diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py index 6d6e9adb96..24dc14b72b 100644 --- a/studio/backend/core/inference/codex_provider.py +++ b/studio/backend/core/inference/codex_provider.py @@ -802,6 +802,7 @@ async def _stream_thread_run( # 2. Legacy: thread.run_streaming(prompt) run_streaming = getattr(thread, "run_streaming", None) if run_streaming is not None: + agent_message_texts: list[str] = [] try: stream_obj = run_streaming(prompt) # Same reasoning as the canonical path above: calling the @@ -812,10 +813,27 @@ async def _stream_thread_run( stream_obj = await stream_obj async for event in stream_obj: turn_started = True - text = _coerce_text(event) + # Match the canonical path's payload-extraction so the + # event-vs-payload shape difference between SDK versions + # is handled the same way in both branches. + payload = getattr(event, "payload", event) + text = _coerce_text(payload) if text: emitted_any = True yield text + else: + # Legacy SDK variants can complete a turn purely via + # ``item.completed`` / agentMessage events with no + # streaming deltas. Capture them so we still emit + # SOMETHING when the stream ends; otherwise the + # request returns 200 with an empty assistant reply + # even though Codex produced a final answer. + final_text = _completed_agent_message_text(payload) + if final_text: + agent_message_texts.append(final_text) + if not emitted_any and agent_message_texts: + yield agent_message_texts[-1] + emitted_any = True return except Exception as exc: logger.warning( @@ -826,6 +844,13 @@ async def _stream_thread_run( turn_started = turn_started, ) if turn_started: + # Flush any agent-message text we collected before the + # crash. The turn already executed on the SDK side, so + # there is no replay risk -- we just want the user to + # see the final answer that was completed before the + # stream broke. + if not emitted_any and agent_message_texts: + yield agent_message_texts[-1] return # 3. Buffered fallback: await the full TurnResult, emit one chunk. @@ -1142,11 +1167,19 @@ async def _stream_codex_parallel( async def _worker(tab_id: int) -> str: """Run one Codex turn, push every chunk into the queue, and return the full accumulated text so the synthesis step can - consume it. Errors are surfaced as a ``codex_tab_error`` - tool-event so the tab strip shows which lane failed without - aborting the whole fan-out. + consume it. Per-turn errors (model rejection, timeout, mid- + stream SDK crash) are surfaced as a ``codex_tab_error`` tool + event so the tab strip shows which lane failed without + aborting the whole fan-out. Setup errors that doom EVERY + worker (SDK not importable, safety enums missing) are re- + raised so the route layer can translate them into a proper + 503 instead of returning a 200 stream with only tool events + and an empty synthesis -- OpenAI-compatible clients that do + not consume ``_toolEvent`` would otherwise see a successful + empty reply. """ collected: list[str] = [] + emit_close = True try: sdk = _import_codex() async_codex_cls = getattr(sdk, "AsyncCodex") @@ -1166,46 +1199,48 @@ async def _stream_codex_parallel( }, ) ) + except CodexUnavailableError: + # Setup-level failure. Same root cause hits every worker, so + # surfacing it as a per-tab error is misleading: every tab + # would emit the same message and the synthesis would be + # blank. Skip the close event too -- the outer fan-out gives + # up before any tabs can render. + emit_close = False + raise except Exception as exc: # CodeQL: never echo str(exc) in client-facing SSE events. # Log full reason server-side; surface a generic message plus # an exception_type discriminator so the UI can still group # failures without leaking file paths / env vars from the - # SDK traceback. CodexUnavailableError is the one exception - # we DO surface verbatim because it's a user-actionable - # install hint with no sensitive content. + # SDK traceback. logger.warning( "codex_provider.parallel_tab_failed", tab_id = tab_id, exc_type = type(exc).__name__, error = str(exc), ) - public_error = ( - str(exc) - if isinstance(exc, CodexUnavailableError) - else "Codex tab failed" - ) await queue.put( _chunk_tool_event( completion_id, { "type": "codex_tab_error", "tab_id": tab_id, - "error": public_error, + "error": "Codex tab failed", "exception_type": type(exc).__name__, }, ) ) finally: - await queue.put( - _chunk_tool_event( - completion_id, - { - "type": "codex_tab_close", - "tab_id": tab_id, - }, + if emit_close: + await queue.put( + _chunk_tool_event( + completion_id, + { + "type": "codex_tab_close", + "tab_id": tab_id, + }, + ) ) - ) return "".join(collected) workers = [asyncio.create_task(_worker(i + 1)) for i in range(n)] @@ -1220,6 +1255,18 @@ async def _stream_codex_parallel( async def _await_workers() -> None: results = await asyncio.gather(*workers, return_exceptions = True) + # If a setup-level failure took out every worker + # (CodexUnavailableError -- SDK not importable, safety enums + # missing, etc.), re-raise so the route layer turns it into a + # 503 instead of letting an empty 200 stream close. Per-tab + # runtime failures stay swallowed (they're already surfaced as + # codex_tab_error events) so a single bad model in the fan-out + # does not kill the others. + setup_errors = [ + r for r in results if isinstance(r, CodexUnavailableError) + ] + if setup_errors and len(setup_errors) == len(results): + raise setup_errors[0] for r in results: if isinstance(r, BaseException): per_tab_texts.append("") @@ -1265,6 +1312,12 @@ async def _stream_codex_parallel( if not cancelled: try: await drain_task + except CodexUnavailableError: + # Setup-level failure took out every worker. Re-raise so + # the route layer translates it into a 503 instead of + # silently continuing into the synthesis step (which + # would itself fail) and returning an empty 200 stream. + raise except Exception as exc: logger.warning( "codex_provider.parallel_drain_failed", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b173a60e2e..629ac56bf6 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1938,7 +1938,15 @@ async def _proxy_to_external_provider( sent_done = False async for line in gen: yield f"{line}\n\n" - if "[DONE]" in line: + # Match the SSE sentinel exactly. The earlier + # substring check (`"[DONE]" in line`) would flip + # the flag when a normal `delta.content` carried + # the literal text "[DONE]" (e.g. an explanation + # of OpenAI's stream terminator), and suppress the + # real `data: [DONE]` frame. OpenAI-compatible + # clients that finalise on the sentinel would + # then hang on stream close. + if line.strip() == "data: [DONE]": sent_done = True if not sent_done: yield "data: [DONE]\n\n" diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py index 8084917fad..55350e3986 100644 --- a/studio/backend/tests/test_codex_provider.py +++ b/studio/backend/tests/test_codex_provider.py @@ -2213,3 +2213,146 @@ class TestStreamReplayProtection: chunks = asyncio.run(collect()) assert thread.run_calls == 1 assert chunks == ["answer"] + + +# ── Round 9: P2 fixes from latest Codex bot review ────────────────── + + +class TestRunStreamingCompletionFallback: + """Round 9 fix: legacy SDK exposes ``thread.run_streaming`` but the + stream only emits completion-style events (no message deltas). We + must still emit the agentMessage text, otherwise the request + returns 200 with an empty assistant reply. + """ + + def test_run_streaming_only_completion_emits_final_text(self): + from core.inference.codex_provider import _stream_thread_run + + # Dict shape matching _completed_agent_message_text's accepted + # form: type=thread.item.completed, item.type=agentMessage, + # item.text=. The legacy run_streaming path now + # extracts ``.payload`` first (matching the canonical path), so + # a plain dict event is the simplest faithful fixture. + completed_event = { + "type": "thread.item.completed", + "item": {"type": "agentMessage", "text": "FINAL_ANSWER"}, + } + + class _Stream: + def __init__(self): + self._sent = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._sent: + raise StopAsyncIteration + self._sent = True + return completed_event + + class _Thread: + # No .turn so the canonical path is skipped; only legacy + # run_streaming exists, and it yields a completion event + # with no streaming deltas. + def run_streaming(self_inner, prompt): + return _Stream() + + async def run(self_inner, prompt): # pragma: no cover + # Should never be called -- the legacy stream + # completed cleanly via the completion event. + raise AssertionError("buffered run must not fire") + + thread = _Thread() + + async def collect(): + return [c async for c in _stream_thread_run(thread, "hi")] + + chunks = asyncio.run(collect()) + assert chunks == ["FINAL_ANSWER"], chunks + + +class TestParallelSetupErrorPropagation: + """Round 9 fix: when CodexUnavailableError takes out every worker + in a parallel-calls fan-out, the function re-raises so the route + layer can return a proper 503. Per-tab runtime failures (timeout + etc.) still get swallowed into codex_tab_error events as before. + """ + + def test_unavailable_in_every_worker_reraises(self, monkeypatch): + from core.inference import codex_provider as cp + + # Force _import_codex to raise CodexUnavailableError. Every + # worker hits this on entry so per_tab_texts stays empty and + # the function MUST re-raise. + def boom(): + raise cp.CodexUnavailableError("SDK not installed (test)") + + monkeypatch.setattr(cp, "_import_codex", boom) + + async def collect_lines(): + lines = [] + try: + async for line in cp._stream_codex_parallel( + model = "gpt-5.4-mini", + system = "", + prompt = "hello", + n = 3, + completion_id = "test-completion", + ): + lines.append(line) + except cp.CodexUnavailableError as exc: + return lines, exc + return lines, None + + lines, exc = asyncio.run(collect_lines()) + assert exc is not None, ( + "CodexUnavailableError did not propagate -- the stream " + "returned a 200 with no assistant content" + ) + assert "SDK not installed (test)" in str(exc) + + +class TestCodexDoneSentinelExactMatch: + """Round 9 fix: the Codex SSE wrapper's `sent_done` detection now + requires an EXACT `data: [DONE]` line match. The substring check + was firing on `delta.content` payloads that happened to contain + the literal text `[DONE]`. + + The route source is the canonical reference -- this test asserts + the source uses an anchored comparison, not a substring `in` + check, so the fix is locked in even if the route is restructured. + """ + + def test_route_uses_exact_done_match(self): + with open( + _backend_file("routes/inference.py"), "r", encoding = "utf-8" + ) as f: + src = f.read() + # The Codex SSE wrapper is the only place we expect this + # comparison style; allow either single or double quotes + # around the canonical line for forward compatibility. + assert ( + 'line.strip() == "data: [DONE]"' in src + or "line.strip() == 'data: [DONE]'" in src + ), ( + "Codex SSE wrapper must terminate on an exact `data: [DONE]` " + "line, not on a substring containing `[DONE]`." + ) + # Inspect the Codex stream block specifically. The old + # substring check `if "[DONE]" in line: sent_done = True` + # must NOT appear as an active comparison. Ignore matches + # inside comments (lines starting with `#` or inside string + # literals describing the old behavior) by scanning for the + # exact statement form. + codex_block_start = src.find('async def _codex_stream():') + if codex_block_start != -1: + window = src[codex_block_start : codex_block_start + 4000] + for line in window.splitlines(): + stripped = line.strip() + if stripped.startswith("#"): + continue + assert 'if "[DONE]" in line' not in stripped, ( + "Codex SSE wrapper still uses substring [DONE] check: " + + stripped + )