diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 11a8c86535e98..56fc1dac71f50 100644 --- a/agent/chat_completion_helpers.py +++ b/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" diff --git a/agent/thread_scoped_output.py b/agent/thread_scoped_output.py index e9e494ab83030..e46608e492752 100644 --- a/agent/thread_scoped_output.py +++ b/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.`` 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 diff --git a/agent/turn_finalizer.py b/agent/turn_finalizer.py index d4d6a23e78650..d140545d4cca4 100644 --- a/agent/turn_finalizer.py +++ b/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 diff --git a/tests/agent/test_thread_scoped_output.py b/tests/agent/test_thread_scoped_output.py index 263d25b898fc0..7f85e7d5f8684 100644 --- a/tests/agent/test_thread_scoped_output.py +++ b/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 diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index db68ee4f712e3..2813c52343f3e 100644 --- a/tests/run_agent/test_run_agent.py +++ b/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 diff --git a/tests/test_tui_gateway_server_crash_history.py b/tests/test_tui_gateway_server_crash_history.py new file mode 100644 index 0000000000000..f7a40982977dd --- /dev/null +++ b/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 diff --git a/tools/code_execution_tool.py b/tools/code_execution_tool.py index 3df99f48388fe..0abde608e4f3c 100644 --- a/tools/code_execution_tool.py +++ b/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) diff --git a/tui_gateway/server.py b/tui_gateway/server.py index cc015d3600b77..ed8f1dcc0ce05 100644 --- a/tui_gateway/server.py +++ b/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