diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index b52bd6a1fb17..33f5036e36d9 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -64,6 +64,38 @@ logger = logging.getLogger(__name__) +def _api_payload_for_stale_timeout(api_kwargs: dict) -> list: + """Return the request payload used to estimate non-stream stale timeout. + + Chat Completions requests carry ``messages``. Responses/Codex requests + carry ``input`` plus optional ``instructions``. The stale detector only + looked at ``messages`` before, so Codex prompts were estimated as ~0 tokens + and never received the larger-context timeout bump. + """ + messages = api_kwargs.get("messages") + if messages: + return messages if isinstance(messages, list) else [messages] + + payload = [] + instructions = api_kwargs.get("instructions") + if instructions: + payload.append({"role": "system", "content": instructions}) + + response_input = api_kwargs.get("input") + if isinstance(response_input, list): + payload.extend(response_input) + elif response_input: + payload.append(response_input) + return payload + + +def _estimate_payload_tokens(payload: object) -> int: + """Cheap, conservative token estimate for log messages/timeouts.""" + if not payload: + return 0 + return sum(len(str(v)) for v in (payload if isinstance(payload, list) else [payload])) // 4 + + def _ra(): """Lazy ``run_agent`` reference. @@ -200,9 +232,8 @@ def _call(): # httpx timeout (default 1800s) with zero feedback. The stale # detector kills the connection early so the main retry loop can # apply richer recovery (credential rotation, provider fallback). - _stale_timeout = agent._compute_non_stream_stale_timeout( - api_kwargs.get("messages", []) - ) + _stale_payload = _api_payload_for_stale_timeout(api_kwargs) + _stale_timeout = agent._compute_non_stream_stale_timeout(_stale_payload) _call_start = time.time() agent._touch_activity("waiting for non-streaming API response") @@ -226,7 +257,7 @@ def _call(): # arrives within the configured timeout. _elapsed = time.time() - _call_start if _elapsed > _stale_timeout: - _est_ctx = sum(len(str(v)) for v in api_kwargs.get("messages", [])) // 4 + _est_ctx = _estimate_payload_tokens(_stale_payload) logger.warning( "Non-streaming API call stale for %.0fs (threshold %.0fs). " "model=%s context=~%s tokens. Killing connection.", diff --git a/tests/agent/test_non_stream_stale_timeout.py b/tests/agent/test_non_stream_stale_timeout.py new file mode 100644 index 000000000000..ef379a401506 --- /dev/null +++ b/tests/agent/test_non_stream_stale_timeout.py @@ -0,0 +1,51 @@ +from unittest.mock import patch + +from agent.chat_completion_helpers import ( + _api_payload_for_stale_timeout, + _estimate_payload_tokens, +) +from run_agent import AIAgent + + +def test_stale_payload_uses_responses_input_and_instructions(): + payload = _api_payload_for_stale_timeout( + { + "model": "gpt-5.5", + "instructions": "system prompt", + "input": [{"role": "user", "content": "hello"}], + } + ) + + assert payload == [ + {"role": "system", "content": "system prompt"}, + {"role": "user", "content": "hello"}, + ] + + +def test_stale_payload_prefers_chat_messages_when_present(): + messages = [{"role": "user", "content": "chat path"}] + + assert _api_payload_for_stale_timeout( + {"messages": messages, "input": [{"role": "user", "content": "responses path"}]} + ) is messages + + +def test_codex_responses_payload_gets_large_context_timeout_bump(monkeypatch, tmp_path): + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + with patch.object(AIAgent, "__init__", lambda self, **kw: None): + agent = AIAgent() + + setattr(agent, "provider", "openai-codex") + setattr(agent, "model", "gpt-5.5") + setattr(agent, "base_url", "https://chatgpt.com/backend-api/codex") + setattr(agent, "_base_url", agent.base_url) + + payload = _api_payload_for_stale_timeout( + { + "instructions": "system", + "input": [{"role": "user", "content": "x" * 240_000}], + } + ) + + assert _estimate_payload_tokens(payload) > 50_000 + assert agent._compute_non_stream_stale_timeout(payload) == 450.0