fix(core): correct context persistence in virtual environments (#61)

Co-authored-by: Alon Yeshurun <alonyeshurun+microsoft@microsoft.com>

Author: ayeshurun

.changes/unreleased/fixed-20251111-053940.yaml

@@ -0,0 +1,3 @@
+kind: fixed
+body: Fix context persistence in virtual environment
+time: 2025-11-11T05:39:40.62257109Z

src/fabric_cli/core/fab_context.py

@@ -5,6 +5,7 @@
 import json
 import os
 import platform
+import sys
 
 import psutil
 
@@ -247,6 +248,43 @@ def _process_exists(self, pid):
             # If we can't determine process existence, assume it exists to be safe
             return True
 
+    def _is_in_venv(self):
+        """Check if running inside a virtual environment."""
+        return sys.prefix != sys.base_prefix
+
+    def _get_session_process(self, parent_process: psutil.Process) -> psutil.Process:
+        """
+        Get the session process, handling virtual environments and fallbacks.
+
+        If a virtual environment is detected, it traverses one additional level up to
+        find the true shell process.
+
+        Args:
+            parent_process (psutil.Process): The parent process to check.
+
+        Returns:
+            psutil.Process: The actual session process.
+        """
+        grandparent_process = parent_process.parent()
+
+        if not grandparent_process:
+            fab_logger.log_debug(
+                "No grandparent process was found. Falling back to parent process for session ID resolution"
+            )
+            return parent_process
+
+        if self._is_in_venv():
+            great_grandparent_process = grandparent_process.parent()
+
+            if not great_grandparent_process:
+                fab_logger.log_debug(
+                    "No great-grandparent process found in virtual environment. Falling back to grandparent process for session ID resolution"
+                )
+            else:
+                return great_grandparent_process
+
+        return grandparent_process
+
     def _get_context_session_id(self):
         """Get the session ID for the current shell session.
 
@@ -255,35 +293,34 @@ def _get_context_session_id(self):
         process and anchors to the actual shell session, ensuring multiple shell sessions
         don't interfere with each other.
 
+        If a virtual environment is detected, it traverses one additional level up to
+        find the true shell process.
+
         Fallback hierarchy:
-        1. Try to get grandparent process (session/shell)
-        2. If grandparent fails, fall back to parent process
+        1. Try to get session process (grandparent or great-grandparent in venv)
+        2. If that fails, fall back to parent process
         3. If parent fails, fall back to current process
         """
+        parent_process = None
         try:
-            parent_process = psutil.Process().parent()
-            if parent_process is None:
+            current_process = psutil.Process()
+            parent_process = current_process.parent()
+            if not parent_process:
                 fab_logger.log_debug(
-                    "No parent process was found. Falling back to the current process for session ID resolution."
+                    "No parent process found. Falling back to current process for session ID."
                 )
                 return os.getpid()
-        except Exception as e:
-            fab_logger.log_debug(
-                f"Failed to get parent process: {e}. Falling back to current process for session ID resolution."
-            )
-            return os.getpid()
 
-        try:
-            session_process = parent_process.parent()
-            if session_process is not None:
-                return session_process.pid
-            else:
+            session_process = self._get_session_process(parent_process)
+            return session_process.pid
+        except Exception as e:
+            if parent_process:
                 fab_logger.log_debug(
-                    "No grandparent process was found. Falling back to parent process for session ID resolution."
+                    f"Failed to get session process: {e}. Falling back to parent process for session ID resolution."
                 )
                 return parent_process.pid
-        except Exception as e:
-            fab_logger.log_debug(
-                f"Failed to get grandparent process: {e}. Falling back to parent process for session ID resolution."
-            )
-            return parent_process.pid
+            else:
+                fab_logger.log_debug(
+                    f"Failed to get parent process: {e}. Falling back to current process for session ID resolution."
+                )
+                return os.getpid()

tests/test_core/test_fab_context.py

@@ -132,19 +132,57 @@ def mock_process():
     assert session_id == 1234
 
 
-def test_get_context_session_id_no_grandparent_process(monkeypatch):
-    parent_process = MockProcess(pid=1234, parent=None)
-    current_process = MockProcess(pid=9999, parent=parent_process)
+def test_get_context_session_id_in_venv(monkeypatch):
+    great_grandparent_process = MockProcess(pid=1111, parent=None)
+    grandparent_process = MockProcess(pid=2222, parent=great_grandparent_process)
+    parent_process = MockProcess(pid=3333, parent=grandparent_process)
+    current_process = MockProcess(pid=4444, parent=parent_process)
 
     def mock_process():
         return current_process
 
     monkeypatch.setattr("fabric_cli.core.fab_context.psutil.Process", mock_process)
+    context = Context()
+    monkeypatch.setattr(context, "_is_in_venv", lambda: True)
+    session_id = context._get_context_session_id()
+
+    assert session_id == 1111
+
+
+def test_get_context_session_id_in_venv_no_great_grandparent(monkeypatch):
+    grandparent_process = MockProcess(pid=2222, parent=None)
+    parent_process = MockProcess(pid=3333, parent=grandparent_process)
+    current_process = MockProcess(pid=4444, parent=parent_process)
 
+    def mock_process():
+        return current_process
+
+    monkeypatch.setattr("fabric_cli.core.fab_context.psutil.Process", mock_process)
     context = Context()
+    monkeypatch.setattr(context, "_is_in_venv", lambda: True)
     session_id = context._get_context_session_id()
 
-    assert session_id == 1234
+    # Falls back to grandparent PID
+    assert session_id == 2222
+
+
+def test_get_context_session_id_not_in_venv(monkeypatch):
+    great_grandparent_process = MockProcess(pid=1111, parent=None)
+    grandparent_process = MockProcess(pid=2222, parent=great_grandparent_process)
+    parent_process = MockProcess(pid=3333, parent=grandparent_process)
+    current_process = MockProcess(pid=4444, parent=parent_process)
+
+    def mock_process():
+        return current_process
+
+    monkeypatch.setattr("fabric_cli.core.fab_context.psutil.Process", mock_process)
+
+    context = Context()
+    monkeypatch.setattr(context, "_is_in_venv", lambda: False)
+    session_id = context._get_context_session_id()
+
+    # Should return grandparent PID, not great-grandparent
+    assert session_id == 2222
 
 
 def test_get_context_session_id_no_parent_process(monkeypatch):
@@ -172,9 +210,9 @@ def mock_log_debug(message):
 
     assert session_id == 9999
     assert len(log_calls) == 1
-    assert "No parent process was found" in log_calls[0]
     assert (
-        "Falling back to the current process for session ID resolution" in log_calls[0]
+        "No parent process found. Falling back to current process for session ID"
+        in log_calls[0]
     )
 
 
@@ -227,11 +265,11 @@ def mock_log_debug(message):
 
     assert session_id == 5678  # Falls back to parent process PID
     assert len(log_calls) == 1
-    assert "Failed to get grandparent process:" in log_calls[0]
+    assert "Failed to get session process:" in log_calls[0]
     assert "Falling back to parent process for session ID resolution" in log_calls[0]
 
 
-def test_get_context_session_id_grandparent_process_none(monkeypatch):
+def test_get_context_session_id_no_grandparent_process(monkeypatch):
     parent_process = MockProcess(pid=5678, parent=None)
     current_process = MockProcess(pid=9999, parent=parent_process)
 
@@ -253,8 +291,10 @@ def mock_log_debug(message):
 
     assert session_id == 5678  # Falls back to parent process PID
     assert len(log_calls) == 1
-    assert "No grandparent process was found" in log_calls[0]
-    assert "Falling back to parent process for session ID resolution" in log_calls[0]
+    assert (
+        "No grandparent process was found. Falling back to parent process for session ID resolution"
+        in log_calls[0]
+    )
 
 
 class MockProcess: