diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 4ae8afb6dd..501284461f 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -1432,6 +1432,35 @@ def interruptible_api_call(agent, api_kwargs: dict): # a network bug and surfaced to the caller. (PR #6600 — cascading interrupt # hang.) _request_cancelled = {"value": False} + # Codex Responses retirement token (codex_responses only). The worker + # thread reads it through ``agent._active_codex_stream_request_token`` to + # tell whether it still owns the turn. When a watchdog below force-closes + # the connection it clears the agent-level token, so a worker still + # draining SSE frames raises instead of returning its partial output as a + # "completed" response (see run_codex_stream's _request_is_current). + # ``_codex_request_retired`` is the request-local mirror, used to swallow + # the transport error our own force-close causes — same split as + # ``_request_cancelled`` above. + _codex_request_token = object() if agent.api_mode == "codex_responses" else None + _codex_request_retired = {"value": False} + + def _install_codex_request_token() -> None: + if _codex_request_token is None: + return + if _codex_request_retired["value"]: + # Already retired before the worker got going — do not re-publish. + return + agent._active_codex_stream_request_token = _codex_request_token + + def _retire_codex_request_token() -> None: + if _codex_request_token is None: + return + _codex_request_retired["value"] = True + if ( + getattr(agent, "_active_codex_stream_request_token", None) + is _codex_request_token + ): + agent._active_codex_stream_request_token = None def _set_request_client(client, *, kind: str = "openai"): with request_client_lock: @@ -1491,6 +1520,7 @@ def interruptible_api_call(agent, api_kwargs: dict): def _call(): try: + _install_codex_request_token() # _set_request_client registers each per-request client with the # stranger-thread abort machinery above; the shared dispatch helper # builds it via this callback (openai- or anthropic-kind) so the @@ -1513,15 +1543,34 @@ def interruptible_api_call(agent, api_kwargs: dict): # handler, the transport error is the expected consequence of our # own force-close, NOT a network bug. Swallow it instead of # surfacing — the main thread raises InterruptedError. (#6600) - if _request_cancelled["value"]: - logger.debug( - "Non-streaming worker caught %s after request cancellation — " - "exiting without surfacing a network error.", - type(e).__name__, - ) + if _request_cancelled["value"] or _codex_request_retired["value"]: + # Retirement is logged at info: it means a watchdog discarded + # output the provider had already sent, which is exactly the + # event an operator debugging a truncated reply needs to see. + # Cancellation stays at debug — a user interrupt is a normal, + # high-frequency outcome and the caller already surfaces it. + if _codex_request_retired["value"]: + logger.info( + "Codex worker caught %s after request retirement — " + "discarding the stale partial instead of surfacing it " + "as a completed response. %s", + type(e).__name__, + agent._client_log_context(), + ) + else: + logger.debug( + "Non-streaming worker caught %s after request " + "cancellation — exiting without surfacing a network " + "error.", + type(e).__name__, + ) return result["error"] = e finally: + # Retire first: _close_request_client_once can raise (every other + # call site wraps it in try/except), and a leaked token would let a + # later worker mistake itself for the owning attempt. + _retire_codex_request_token() # Reuse reason only on a clean response; any other outcome — # error, or the cancel-swallow return above (which leaves both # result slots None) — really closes so the next attempt builds @@ -1734,6 +1783,7 @@ def interruptible_api_call(agent, api_kwargs: dict): _close_request_client_once("codex_ttfb_kill") except Exception: pass + _retire_codex_request_token() agent._emit_wait_notice( f"⚠ no response from provider in {int(_elapsed)}s — " f"reconnecting..." @@ -1784,6 +1834,7 @@ def interruptible_api_call(agent, api_kwargs: dict): _close_request_client_once("codex_stream_idle_kill") except Exception: pass + _retire_codex_request_token() agent._touch_activity( f"codex stream killed after {int(_event_stale_elapsed)}s with no SSE events" ) @@ -1815,6 +1866,7 @@ def interruptible_api_call(agent, api_kwargs: dict): _close_request_client_once("stale_call_kill") except Exception: pass + _retire_codex_request_token() # Circuit breaker (#58962): count the stale kill. See the # canonical comment block above ``_stale_streak()``. _bump_stale_streak(agent) @@ -1862,6 +1914,7 @@ def interruptible_api_call(agent, api_kwargs: dict): _close_request_client_once("interrupt_abort") except Exception: pass + _retire_codex_request_token() # #81521 (sibling of the streaming-path fix): wait for the worker # to unwind Relay-managed scopes before surfacing # InterruptedError, so turn teardown cannot race a still-open diff --git a/agent/codex_runtime.py b/agent/codex_runtime.py index 1fdd2d153e..30b86c6171 100644 --- a/agent/codex_runtime.py +++ b/agent/codex_runtime.py @@ -1101,7 +1101,9 @@ def _consume_codex_event_stream( * ``on_first_delta()`` — one-shot, fires on the first text delta only. * ``on_event(event)`` — fires for every event before any other processing. Used for watchdog activity, debug logging, anything wire-shape-agnostic. - * ``interrupt_check()`` — returns True to break the loop early. + * ``interrupt_check()`` — returns True to break the loop early, or raises + ``TimeoutError`` / ``InterruptedError`` for request-retirement control + flow that must not be converted into a partial final response. """ collected_output_items: List[Any] = [] # output_index of each collected_output_items entry, appended in lockstep @@ -1606,18 +1608,39 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta max_stream_retries = 1 # Accumulate streamed text so callers / compat shims can read it. agent._codex_streamed_text_parts: list = [] + # Retirement token for THIS request, installed by + # ``interruptible_api_call`` before it hands off to the worker thread. When + # a watchdog (TTFB / stream-idle / stale-call) kills the connection it + # clears the agent-level token, so a worker that is still draining frames + # can tell it has been retired. ``None`` means no watchdog owns this call + # (auxiliary callers drive this function directly) — then every check + # passes and behavior is unchanged. + request_token = getattr(agent, "_active_codex_stream_request_token", None) + + def _request_is_current() -> bool: + if request_token is None: + return True + return getattr(agent, "_active_codex_stream_request_token", None) is request_token def _on_text_delta(text: str) -> None: + if not _request_is_current(): + return agent._codex_streamed_text_parts.append(text) agent._fire_stream_delta(text) def _on_reasoning_delta(text: str) -> None: + if not _request_is_current(): + return agent._fire_reasoning_delta(text) def _on_commentary_message(text: str) -> None: + if not _request_is_current(): + return agent._fire_streamed_codex_commentary(text) def _on_event(event: Any) -> None: + if not _request_is_current(): + return # TTFB watchdog and activity touch — runs once per SSE event. agent._codex_stream_last_event_ts = time.time() agent._touch_activity("receiving stream response") @@ -1721,6 +1744,14 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta raise def _interrupt_or_superseded() -> bool: + # A retired request must NOT break out of the consume loop: breaking + # returns the partial `final` (status defaults to "completed"), which + # the caller persists as a finished assistant turn. Raise so the + # watchdog's own TimeoutError is what the retry path sees. + if not _request_is_current(): + raise TimeoutError( + "Codex Responses stream request retired before terminal response" + ) return bool(agent._interrupt_requested) try: diff --git a/tests/agent/test_codex_ttfb_watchdog.py b/tests/agent/test_codex_ttfb_watchdog.py index 66208a8e1a..bdf53061f3 100644 --- a/tests/agent/test_codex_ttfb_watchdog.py +++ b/tests/agent/test_codex_ttfb_watchdog.py @@ -108,6 +108,89 @@ def test_ttfb_includes_silent_hang_hint_for_gpt_5_5(tmp_path, monkeypatch): stop["flag"] = True +def test_ttfb_installs_and_retires_the_codex_request_token(tmp_path, monkeypatch): + """The watchdog must publish a per-request token and clear it on the kill. + + ``run_codex_stream`` reads ``agent._active_codex_stream_request_token`` to + tell whether it is still the owning attempt. Without an install here the + whole retirement guard would be inert, and without the clear on kill a + retired worker would keep normalizing partial deltas into a "completed" + response. + + The worker also unwinds with its own local error after the force-close; + that error must not replace the watchdog's retryable ``TimeoutError``. + """ + from agent import chat_completion_helpers as h + + agent = _make_codex_agent(tmp_path, monkeypatch) + monkeypatch.setenv("HERMES_CODEX_TTFB_TIMEOUT_SECONDS", "1") + + closes: list = [] + seen = {"token_while_running": None} + dummy_client = SimpleNamespace() + monkeypatch.setattr(agent, "_create_request_openai_client", lambda **k: dummy_client) + monkeypatch.setattr( + agent, + "_abort_request_openai_client", + lambda c, reason=None: closes.append(reason), + ) + monkeypatch.setattr( + agent, + "_close_request_openai_client", + lambda c, reason=None: closes.append(reason), + ) + + def fake_stream(api_kwargs, client=None, on_first_delta=None): + seen["token_while_running"] = getattr( + agent, "_active_codex_stream_request_token", None + ) + deadline = time.time() + 30 + while time.time() < deadline: + if getattr(agent, "_active_codex_stream_request_token", None) is None: + # Retired by the watchdog — mimic the transport unwinding. + raise RuntimeError("retired worker stream ended without terminal") + time.sleep(0.02) + raise RuntimeError("test timed out waiting for retirement") + + monkeypatch.setattr(agent, "_run_codex_stream", fake_stream) + + with pytest.raises(TimeoutError) as excinfo: + h.interruptible_api_call(agent, {"model": "gpt-5.5", "input": "hi"}) + + assert seen["token_while_running"] is not None, ( + "interruptible_api_call must install a request token before the worker runs" + ) + assert "TTFB" in str(excinfo.value) + assert "retired worker" not in str(excinfo.value) + assert "codex_ttfb_kill" in closes + assert getattr(agent, "_active_codex_stream_request_token", None) is None + + +def test_non_codex_api_mode_installs_no_request_token(tmp_path, monkeypatch): + """The token is codex_responses-only — other api_modes stay untouched.""" + from agent import chat_completion_helpers as h + + agent = _make_codex_agent(tmp_path, monkeypatch) + agent.api_mode = "chat_completions" + + seen = {"token": "unset"} + dummy_client = SimpleNamespace() + monkeypatch.setattr(agent, "_create_request_openai_client", lambda **k: dummy_client) + + def fake_dispatch(_agent, _api_kwargs, *, make_client): + make_client("test") + seen["token"] = getattr( + _agent, "_active_codex_stream_request_token", "absent" + ) + return SimpleNamespace(choices=[]) + + monkeypatch.setattr(h, "_dispatch_nonstreaming_api_request", fake_dispatch) + + h.interruptible_api_call(agent, {"model": "gpt-5.5", "messages": []}) + + assert seen["token"] in (None, "absent") + + def test_ttfb_does_not_kill_when_events_flow(tmp_path, monkeypatch): diff --git a/tests/run_agent/test_run_agent_codex_responses.py b/tests/run_agent/test_run_agent_codex_responses.py index 9c0165e066..26228f0b21 100644 --- a/tests/run_agent/test_run_agent_codex_responses.py +++ b/tests/run_agent/test_run_agent_codex_responses.py @@ -2502,3 +2502,126 @@ def test_codex_first_compaction_continuation_is_still_a_bare_retry(monkeypatch): if m.get("role") == "user" and m.get("content") == _CODEX_INCOMPLETE_NUDGE ] + + +class _LazyCreateStream: + """Lazy iterable fake — events are produced during consumption, not upfront. + + ``_FakeCreateStream`` materializes its events with ``list(events)`` in + __init__, which would run any side effect a generator encodes (such as + retiring the request token) before consumption starts. Retirement tests + need the side effect to land *between* two consumed frames. + """ + + def __init__(self, event_factory): + self._event_factory = event_factory + self.closed = False + + def __iter__(self): + return iter(self._event_factory()) + + def close(self): + self.closed = True + + +def _retiring_stream(agent, deltas, *, retire_after): + """Yield ``deltas`` lazily, clearing the request token mid-stream. + + The token is cleared just before yielding delta index ``retire_after``, + mimicking a watchdog (TTFB / stream-idle / stale-call) retiring the + in-flight request while the worker thread is still draining SSE frames. + """ + + def _events(): + yield SimpleNamespace(type="response.created") + for index, delta in enumerate(deltas): + if index == retire_after: + agent._active_codex_stream_request_token = None + yield SimpleNamespace(type="response.output_text.delta", delta=delta) + # A retired stream never reaches a terminal frame on the wire; the + # connection is force-closed under it. + + return _LazyCreateStream(_events) + + +def test_run_codex_stream_retired_request_raises_instead_of_partial_final(monkeypatch): + """A retired request must not be normalized into a completed response. + + ``_consume_codex_event_stream`` returns ``status=terminal_status`` which + defaults to ``"completed"``, and its only guard is + ``if not saw_terminal and not output``. A watchdog kill mid-stream leaves + ``saw_terminal=False`` but ``output``/text non-empty, so the partial text + used to come back as a ``finish_reason=stop`` response and get persisted as + a complete assistant turn (a long reply would just stop mid-sentence). + + Retirement must surface as a retryable ``TimeoutError`` instead. + """ + agent = _build_agent(monkeypatch) + token = object() + agent._active_codex_stream_request_token = token + + def _fake_create(**kwargs): + assert kwargs.get("stream") is True + return _retiring_stream( + agent, ["1. Create ", "(6/6)", " [END-BILLING"], retire_after=2 + ) + + agent.client = SimpleNamespace(responses=SimpleNamespace(create=_fake_create)) + + with pytest.raises(TimeoutError, match="retired"): + agent._run_codex_stream(_codex_request_kwargs()) + + +def test_run_codex_stream_without_token_keeps_partial_tolerance(monkeypatch): + """No token installed (non-watchdog callers) keeps the existing behavior. + + ``_active_codex_stream_request_token`` is only set by + ``interruptible_api_call``. Auxiliary callers (compression summaries, + title generation) drive ``_run_codex_stream`` directly with no token and + must keep tolerating a stream that ends without a terminal frame. + """ + agent = _build_agent(monkeypatch) + agent._active_codex_stream_request_token = None + output_item = SimpleNamespace( + type="message", + status="completed", + content=[SimpleNamespace(type="output_text", text="no terminal frame")], + ) + + def _fake_create(**kwargs): + return _FakeCreateStream([ + SimpleNamespace(type="response.created"), + SimpleNamespace(type="response.output_item.done", item=output_item), + ]) + + agent.client = SimpleNamespace(responses=SimpleNamespace(create=_fake_create)) + + response = agent._run_codex_stream(_codex_request_kwargs()) + assert response.status == "completed" + assert response.output == [output_item] + + +def test_run_codex_stream_retired_request_stops_firing_callbacks(monkeypatch): + """Deltas that arrive after retirement must not reach the UI callbacks. + + The gateway caches AIAgent instances per session, so a retired worker that + keeps draining frames would otherwise stream tokens from an abandoned + attempt into the live turn's bubble alongside the retry's output. + """ + agent = _build_agent(monkeypatch) + token = object() + agent._active_codex_stream_request_token = token + + streamed: list[str] = [] + monkeypatch.setattr(agent, "_fire_stream_delta", streamed.append) + + def _fake_create(**kwargs): + return _retiring_stream(agent, ["keep", "DROPPED"], retire_after=1) + + agent.client = SimpleNamespace(responses=SimpleNamespace(create=_fake_create)) + + with pytest.raises(TimeoutError): + agent._run_codex_stream(_codex_request_kwargs()) + + assert streamed == ["keep"] + assert "DROPPED" not in streamed