Merge pull request #80770 from NousResearch/bb/desktop-session-integrity

fix: preserve session history when a turn crashes

agent/chat_completion_helpers.py

@@ -2119,7 +2119,9 @@ def try_activate_fallback(agent, reason: "FailoverReason | None" = None) -> bool
 
 def handle_max_iterations(agent, messages: list, api_call_count: int) -> str:
     """Request a summary when max iterations are reached. Returns the final response text."""
-    print(f"⚠️  Reached maximum iterations ({agent.max_iterations}). Requesting summary...")
+    agent._safe_print(
+        f"⚠️  Reached maximum iterations ({agent.max_iterations}). Requesting summary..."
+    )
 
     summary_api_request_id = f"iteration-summary:{uuid.uuid4()}"
     summary_call_outcome = "failed"

agent/thread_scoped_output.py

@@ -104,18 +104,18 @@ class _ThreadRoutingStream:
         return getattr(self._target(), name)
 
 
-def _ensure_installed(attr: str, sink: TextIO) -> "_ThreadRoutingStream":
+def _ensure_installed(attr: str, passthrough: TextIO) -> "_ThreadRoutingStream":
     """Install (idempotently) a routing proxy as ``sys.<attr>`` and return it."""
     with _install_lock:
         proxy = _installed.get(attr)
         current = getattr(sys, attr, None)
         if proxy is not None and current is proxy:
             return proxy
-        # Capture whatever is currently bound as the passthrough.  If a prior
-        # global redirect_stdout is active we deliberately route non-silenced
-        # threads to *that* (matching prior behaviour) rather than guessing at
-        # the "real" stream.
-        passthrough = current if current is not None else sink
+        # Capture whatever is currently bound as the passthrough. If a prior
+        # global redirect_stdout is active, route non-silenced threads to that
+        # stream to preserve the old behavior.
+        passthrough = current if current is not None else passthrough
+        sink = open(os.devnull, "w", encoding="utf-8")
         proxy = _ThreadRoutingStream(passthrough, sink)
         setattr(sys, attr, proxy)
         _installed[attr] = proxy
@@ -130,10 +130,9 @@ def thread_scoped_silence() -> Iterator[None]:
     thread's body instead of ``contextlib.redirect_stdout(devnull)`` when the
     process is multi-threaded and another thread must keep its console output.
     """
-    sink = open(os.devnull, "w", encoding="utf-8")
     ident = threading.get_ident()
-    out_proxy = _ensure_installed("stdout", sink)
-    err_proxy = _ensure_installed("stderr", sink)
+    out_proxy = _ensure_installed("stdout", sys.__stdout__ or sys.stdout)
+    err_proxy = _ensure_installed("stderr", sys.__stderr__ or sys.stderr)
     out_proxy.silence(ident)
     err_proxy.silence(ident)
     try:
@@ -141,7 +140,3 @@ def thread_scoped_silence() -> Iterator[None]:
     finally:
         out_proxy.unsilence(ident)
         err_proxy.unsilence(ident)
-        try:
-            sink.close()
-        except Exception:
-            pass

agent/turn_finalizer.py

@@ -412,6 +412,14 @@ def finalize_turn(
         _cleanup_errors.append(f"persist_session: {_persist_err}")
         logger.error("finalize_turn: _persist_session failed: %s", _persist_err, exc_info=True)
 
+    # The gateway owns a separate in-memory history snapshot. Keep it current
+    # even when finalization reports a cleanup error: a later prompt must not be
+    # sent with the pre-turn snapshot while the durable DB already has this turn.
+    try:
+        agent._session_messages = messages
+    except Exception:
+        pass
+
     # ── Turn-exit diagnostic log ─────────────────────────────────────
     # Always logged at INFO so agent.log captures WHY every turn ended.
     # When the last message is a tool result (agent was mid-work), log

tests/agent/test_thread_scoped_output.py

@@ -82,3 +82,15 @@ def test_many_concurrent_silenced_and_loud_threads():
     for i in range(5):
         assert f"S{i}" not in captured, f"silenced S{i} leaked"
         assert f"L{i}" in captured, f"loud L{i} swallowed"
+
+
+def test_repeated_contexts_never_write_to_a_closed_sink():
+    """The installed proxy must survive later silenced workers."""
+    original = sys.stdout
+    try:
+        for _ in range(3):
+            with thread_scoped_silence():
+                sys.stdout.write("hidden\n")
+            sys.stdout.fileno()
+    finally:
+        sys.stdout = original

tests/run_agent/test_run_agent.py

@@ -2454,6 +2454,13 @@ class TestMcpParallelToolBatch:
 
 
 class TestHandleMaxIterations:
+    def test_summary_notice_uses_safe_print(self, agent):
+        agent._print_fn = lambda *_args, **_kwargs: (_ for _ in ()).throw(ValueError("closed"))
+        agent.client.chat.completions.create.return_value = _mock_response(content="Summary")
+        agent._cached_system_prompt = "You are helpful."
+
+        assert agent._handle_max_iterations([{"role": "user", "content": "do stuff"}], 60) == "Summary"
+
     def test_returns_summary(self, agent):
         resp = _mock_response(content="Here is a summary of what I did.")
         agent.client.chat.completions.create.return_value = resp

tests/test_tui_gateway_server_crash_history.py

@@ -0,0 +1,37 @@
+"""Regression coverage for crashed TUI gateway turns."""
+
+import threading
+import types
+
+from tui_gateway import server
+
+
+def test_turn_error_restores_agent_transcript_to_gateway_history():
+    agent = types.SimpleNamespace(
+        _session_messages=[
+            {"role": "user", "content": "first"},
+            {"role": "assistant", "content": "completed work"},
+        ]
+    )
+    session = {
+        "history": [{"role": "user", "content": "first"}],
+        "history_lock": threading.Lock(),
+        "history_version": 0,
+    }
+
+    assert server._restore_agent_history_after_turn_error(session, agent) is True
+    assert session["history"] == agent._session_messages
+    assert session["history_version"] == 1
+
+
+def test_turn_error_with_no_agent_transcript_does_not_overwrite_history():
+    agent = types.SimpleNamespace(_session_messages=None)
+    session = {
+        "history": [{"role": "user", "content": "first"}],
+        "history_lock": threading.Lock(),
+        "history_version": 4,
+    }
+
+    assert server._restore_agent_history_after_turn_error(session, agent) is False
+    assert session["history"] == [{"role": "user", "content": "first"}]
+    assert session["history_version"] == 4

tools/code_execution_tool.py

@@ -49,6 +49,7 @@ _IS_WINDOWS = platform.system() == "Windows"
 from typing import Any, Dict, List, Optional, Tuple
 
 from tools.thread_context import propagate_context_to_thread
+from agent.thread_scoped_output import thread_scoped_silence
 
 # Availability gate.  On Windows we fall back to loopback TCP for the
 # sandbox RPC transport (AF_UNIX is unreliable on Windows Python) — see
@@ -741,17 +742,10 @@ def _rpc_server_loop(
                 # Suppress stdout/stderr from internal tool handlers so
                 # their status prints don't leak into the CLI spinner.
                 try:
-                    _real_stdout, _real_stderr = sys.stdout, sys.stderr
-                    devnull = open(os.devnull, "w", encoding="utf-8")
-                    try:
-                        sys.stdout = devnull
-                        sys.stderr = devnull
+                    with thread_scoped_silence():
                         result = handle_function_call(
                             tool_name, tool_args, task_id=task_id
                         )
-                    finally:
-                        sys.stdout, sys.stderr = _real_stdout, _real_stderr
-                        devnull.close()
                 except Exception as exc:
                     logger.error("Tool call failed in sandbox: %s", exc, exc_info=True)
                     result = tool_error(str(exc))
@@ -1023,17 +1017,10 @@ def _rpc_poll_loop(
 
                     # Dispatch through the standard tool handler
                     try:
-                        _real_stdout, _real_stderr = sys.stdout, sys.stderr
-                        devnull = open(os.devnull, "w", encoding="utf-8")
-                        try:
-                            sys.stdout = devnull
-                            sys.stderr = devnull
+                        with thread_scoped_silence():
                             tool_result = handle_function_call(
                                 tool_name, tool_args, task_id=task_id
                             )
-                        finally:
-                            sys.stdout, sys.stderr = _real_stdout, _real_stderr
-                            devnull.close()
                     except Exception as exc:
                         logger.error("Tool call failed in remote sandbox: %s",
                                      exc, exc_info=True)

tui_gateway/server.py

@@ -7683,6 +7683,22 @@ def _emit_terminal_turn_error(sid: str, session: dict, error: Any) -> None:
     _emit("message.complete", sid, payload)
 
 
+def _restore_agent_history_after_turn_error(session: dict, agent) -> bool:
+    """Keep a failed turn's working transcript in the gateway session.
+
+    ``AIAgent`` persists its working messages independently of the gateway's
+    history snapshot. If the turn raises after that persistence, the next
+    prompt must see the working transcript instead of the pre-turn snapshot.
+    """
+    agent_messages = getattr(agent, "_session_messages", None)
+    if not isinstance(agent_messages, list):
+        return False
+    with session["history_lock"]:
+        session["history"] = list(agent_messages)
+        session["history_version"] = int(session.get("history_version", 0)) + 1
+    return True
+
+
 def _queued_prompt_snapshot(session: dict) -> dict | None:
     """Return the accepted next-turn prompt without its transport handle.
 
@@ -10091,6 +10107,12 @@ def _run_prompt_submit(
             print(
                 f"[gateway-turn] {type(e).__name__}: {e}", file=sys.stderr, flush=True
             )
+            # The agent persists its working transcript on normal finalization,
+            # but an exception in that finalizer can otherwise leave the
+            # gateway's separate in-memory history at the turn-start snapshot.
+            # Keep the partial turn available to the next prompt; the durable
+            # inflight record still carries the recoverable error state.
+            _restore_agent_history_after_turn_error(session, agent)
             try:
                 # Close the turn with the same terminal error frame shape as
                 # the returned-error path (uniform client handling), retaining