fix(security): deactivate inherited log captures

Author: erichare

src/lfx/src/lfx/cli/common.py

@@ -63,6 +63,19 @@
 _GITHUB_TOKEN_ENV = "GITHUB_TOKEN"
 
 
+class _StreamCaptureBinding:
+    """Shared capture state inherited by child execution contexts."""
+
+    def __init__(self, target) -> None:
+        self.target = target
+        self.active = True
+
+    def deactivate(self) -> None:
+        """Stop routing writes and release the completed capture buffer."""
+        self.active = False
+        self.target = None
+
+
 class _ContextualTextStream:
     """Route writes to a task-local stream while preserving a process fallback.
 
@@ -75,17 +88,25 @@ class _ContextualTextStream:
 
     def __init__(self, fallback, *, context_name: str) -> None:
         self._fallback = fallback
-        self._target: contextvars.ContextVar[Any | None] = contextvars.ContextVar(context_name, default=None)
+        self._binding: contextvars.ContextVar[_StreamCaptureBinding | None] = contextvars.ContextVar(
+            context_name, default=None
+        )
 
     def _current(self):
-        target = self._target.get()
-        return self._fallback if target is None else target
-
-    def activate(self, target) -> contextvars.Token[Any | None]:
-        return self._target.set(target)
-
-    def reset(self, token: contextvars.Token[Any | None]) -> None:
-        self._target.reset(token)
+        binding = self._binding.get()
+        return self._fallback if binding is None or not binding.active else binding.target
+
+    def activate(self, target) -> tuple[contextvars.Token[_StreamCaptureBinding | None], _StreamCaptureBinding]:
+        binding = _StreamCaptureBinding(target)
+        return self._binding.set(binding), binding
+
+    def reset(
+        self,
+        activation: tuple[contextvars.Token[_StreamCaptureBinding | None], _StreamCaptureBinding],
+    ) -> None:
+        token, binding = activation
+        binding.deactivate()
+        self._binding.reset(token)
 
     def write(self, data):
         return self._current().write(data)

src/lfx/tests/unit/cli/test_common.py

@@ -279,6 +279,43 @@ async def mock_async_start(inputs, **kwargs):  # noqa: ARG001
         assert process_stdout.getvalue() == "outside stdout\n"
         assert process_stderr.getvalue() == "outside stderr\n"
 
+    @pytest.mark.asyncio
+    async def test_execute_graph_with_capture_deactivates_outliving_child_capture(self):
+        """A child task that outlives its graph writes to the process streams."""
+        child_started = asyncio.Event()
+        release_child = asyncio.Event()
+        child_task = None
+
+        async def write_after_graph_returns():
+            child_started.set()
+            await release_child.wait()
+            sys.stdout.write("late child stdout\n")
+            sys.stderr.write("late child stderr\n")
+
+        async def mock_async_start(inputs, **kwargs):  # noqa: ARG001
+            nonlocal child_task
+            child_task = asyncio.create_task(write_after_graph_returns())
+            await child_started.wait()
+            sys.stdout.write("captured stdout\n")
+            sys.stderr.write("captured stderr\n")
+            yield MagicMock(results={"text": "ok"})
+
+        mock_graph = MagicMock()
+        mock_graph.context = {}
+        mock_graph.async_start = mock_async_start
+
+        process_stdout = StringIO()
+        process_stderr = StringIO()
+        with patch.object(sys, "stdout", process_stdout), patch.object(sys, "stderr", process_stderr):
+            _, logs = await execute_graph_with_capture(mock_graph, "test input")
+            release_child.set()
+            assert child_task is not None
+            await child_task
+
+        assert logs == "captured stdout\ncaptured stderr\n"
+        assert process_stdout.getvalue() == "late child stdout\n"
+        assert process_stderr.getvalue() == "late child stderr\n"
+
     @pytest.mark.asyncio
     async def test_execute_graph_with_capture_success(self):
         """Test successful graph execution with output capture."""