diff --git a/codex_backend.py b/codex_backend.py index 9365ebd..b8bac33 100644 --- a/codex_backend.py +++ b/codex_backend.py @@ -10,10 +10,17 @@ """ from __future__ import annotations +import asyncio import json +from contextlib import AsyncExitStack from typing import Any, Iterable from llm_router_host import _cached_tokens +from provider_adapters.common import ( + before_first_output, + first_token_timeout_err, + first_token_timeout_s, +) CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex" @@ -267,6 +274,20 @@ def aggregate_codex_sse(lines: Iterable[str], latency_ms: int) -> dict: } +def _codex_line_has_output_delta(line: str) -> bool: + line = line.strip() + if not line or not line.startswith("data:"): + return False + payload = line[len("data:"):].strip() + if payload == "[DONE]": + return False + try: + ev = json.loads(payload) + except ValueError: + return False + return ev.get("type") == "response.output_text.delta" and bool(ev.get("delta")) + + def make_codex_async_call_provider( auth, base_url: str = CODEX_BASE_URL, @@ -309,11 +330,34 @@ async def call(request: dict) -> dict: timeout = (request.get("timeout_ms") or int(timeout_s * 1000)) / 1000.0 t0 = time.monotonic() + saw_output = False + status_seen = False + first_timeout_s = first_token_timeout_s(request) + + def _latency() -> int: + return int((time.monotonic() - t0) * 1000) + + def _saw_output() -> bool: + return saw_output + + def _timeout_err() -> dict: + return first_token_timeout_err(first_timeout_s, _latency()) + try: async with httpx.AsyncClient(timeout=timeout) as c: - async with c.stream("POST", url, json=body, headers=headers) as resp: + async with AsyncExitStack() as stack: + try: + resp = await before_first_output(stack.enter_async_context( + c.stream("POST", url, json=body, headers=headers)), + first_timeout_s, t0, _saw_output) + except (asyncio.TimeoutError, TimeoutError): + if not status_seen: + _notify(0) + return _timeout_err() + _notify(resp.status_code, resp.headers) - latency = int((time.monotonic() - t0) * 1000) + status_seen = True + latency = _latency() if resp.status_code == 401: return _err("auth_error", 401, latency, "codex token rejected") if resp.status_code == 429: @@ -321,13 +365,27 @@ async def call(request: dict) -> dict: if resp.status_code >= 400: detail = (await resp.aread()).decode("utf-8", "replace")[:500] return _err("server_error", resp.status_code, latency, detail) - lines = [line async for line in resp.aiter_lines()] - return aggregate_codex_sse(lines, int((time.monotonic() - t0) * 1000)) + lines = [] + stream_lines = resp.aiter_lines().__aiter__() + while True: + try: + line = await before_first_output( + stream_lines.__anext__(), first_timeout_s, t0, _saw_output) + except StopAsyncIteration: + break + except (asyncio.TimeoutError, TimeoutError): + if not saw_output: + return _timeout_err() + raise + lines.append(line) + if _codex_line_has_output_delta(line): + saw_output = True + return aggregate_codex_sse(lines, _latency()) except httpx.TimeoutException: _notify(0) - return _err("timeout", 0, int((time.monotonic() - t0) * 1000), "codex request timed out") + return _err("timeout", 0, _latency(), "codex request timed out") except (httpx.NetworkError, httpx.RequestError) as e: _notify(0) - return _err("network_error", 0, int((time.monotonic() - t0) * 1000), str(e)) + return _err("network_error", 0, _latency(), str(e)) return call diff --git a/provider_adapters/common.py b/provider_adapters/common.py index 9e0a548..1f4d1c3 100644 --- a/provider_adapters/common.py +++ b/provider_adapters/common.py @@ -2,6 +2,7 @@ from __future__ import annotations import json +import asyncio import time from typing import Awaitable, Callable, Any @@ -159,10 +160,41 @@ def elapsed_ms(t0: float) -> int: return int((time.monotonic() - t0) * 1000) +def first_token_timeout_s(request: dict) -> "float | None": + raw = request.get("first_token_timeout_ms") + try: + seconds = float(raw) / 1000.0 + except (TypeError, ValueError): + return None + return seconds if seconds > 0 else None + + +async def before_first_output( + awaitable, + timeout_s: "float | None", + t0: float, + saw_output: Callable[[], bool], +): + if timeout_s is None or saw_output(): + return await awaitable + remaining = timeout_s - (time.monotonic() - t0) + if remaining <= 0: + raise asyncio.TimeoutError + return await asyncio.wait_for(awaitable, timeout=remaining) + + +def first_token_timeout_err(timeout_s: float, latency_ms: int) -> dict: + return err("timeout", 0, latency_ms, + f"first token timed out after {int(timeout_s * 1000)}ms") + + # Back-compat names used by older tests/scripts that import through # llm_router_host's re-export layer. _cached_tokens = cached_tokens _classify_status = classify_status _err = err _elapsed_ms = elapsed_ms +_first_token_timeout_s = first_token_timeout_s +_before_first_output = before_first_output +_first_token_timeout_err = first_token_timeout_err _provider_error_message = provider_error_message diff --git a/provider_adapters/openai_compatible.py b/provider_adapters/openai_compatible.py index 76476a5..06ba581 100644 --- a/provider_adapters/openai_compatible.py +++ b/provider_adapters/openai_compatible.py @@ -5,12 +5,16 @@ import json import os import time +from contextlib import AsyncExitStack from typing import Awaitable, Callable, Any from provider_adapters.common import ( CallProviderHook, AsyncCallProviderHook, TokenProvider, + before_first_output, + first_token_timeout_err, + first_token_timeout_s, _cached_tokens, _classify_status, _elapsed_ms, @@ -308,26 +312,25 @@ async def stream_openai_compatible( usage: dict = {} raw_model = None saw_output = False - first_token_timeout_ms = request.get("first_token_timeout_ms") - try: - first_token_timeout_s = ( - float(first_token_timeout_ms) / 1000.0 - if first_token_timeout_ms is not None and float(first_token_timeout_ms) > 0 - else None - ) - except (TypeError, ValueError): - first_token_timeout_s = None + first_timeout_s = first_token_timeout_s(request) def _latency() -> int: return int((time.monotonic() - t0) * 1000) - def _first_token_timeout_error() -> dict: - return _err("timeout", 0, _latency(), - f"first token timed out after {int(first_token_timeout_s * 1000)}ms") + def _saw_output() -> bool: + return saw_output + + def _timeout_err() -> dict: + return first_token_timeout_err(first_timeout_s, _latency()) try: - async with client.stream("POST", url, json=body, headers=headers, - timeout=timeout) as resp: + async with AsyncExitStack() as stack: + try: + resp = await before_first_output(stack.enter_async_context( + client.stream("POST", url, json=body, headers=headers, + timeout=timeout)), first_timeout_s, t0, _saw_output) + except (asyncio.TimeoutError, TimeoutError): + return _timeout_err() if not (200 <= resp.status_code < 300): raw = (await resp.aread()).decode("utf-8", "replace")[:500] kind = _classify_from_map(raw, rules.get("error_map")) \ @@ -337,18 +340,13 @@ def _first_token_timeout_error() -> dict: lines = resp.aiter_lines().__aiter__() while True: try: - if first_token_timeout_s is not None and not saw_output: - remaining = first_token_timeout_s - (time.monotonic() - t0) - if remaining <= 0: - return _first_token_timeout_error() - line = await asyncio.wait_for(lines.__anext__(), timeout=remaining) - else: - line = await lines.__anext__() + line = await before_first_output( + lines.__anext__(), first_timeout_s, t0, _saw_output) except StopAsyncIteration: break except (asyncio.TimeoutError, TimeoutError): if not saw_output: - return _first_token_timeout_error() + return _timeout_err() raise if not line or not line.startswith("data:"): continue diff --git a/streaming.py b/streaming.py index fb97657..31d0bca 100644 --- a/streaming.py +++ b/streaming.py @@ -15,11 +15,18 @@ import json import re +import asyncio import time import uuid +from contextlib import AsyncExitStack from typing import Any, Awaitable, Callable -from provider_adapters.common import _err +from provider_adapters.common import ( + _err, + before_first_output, + first_token_timeout_err, + first_token_timeout_s, +) # The openai-compatible STREAM backend lives with its non-stream sibling in the # adapter leaf (it shares _prepare_openai_call); re-exported here so the dispatcher # and shim keep their import site. This module → provider_adapters (the allowed @@ -76,13 +83,25 @@ def _notify(status: int, headers=None) -> None: t0 = time.monotonic() emitted = False lines: list[str] = [] + first_timeout_s = first_token_timeout_s(request) def _latency() -> int: return int((time.monotonic() - t0) * 1000) + def _emitted() -> bool: + return emitted + + def _timeout_err() -> dict: + return first_token_timeout_err(first_timeout_s, _latency()) + try: - async with client.stream("POST", url, json=body, headers=headers, - timeout=timeout) as resp: + async with AsyncExitStack() as stack: + try: + resp = await before_first_output(stack.enter_async_context( + client.stream("POST", url, json=body, headers=headers, + timeout=timeout)), first_timeout_s, t0, _emitted) + except (asyncio.TimeoutError, TimeoutError): + return _timeout_err() _notify(resp.status_code, resp.headers) if resp.status_code == 401: return _err("auth_error", 401, _latency(), "codex token rejected") @@ -91,7 +110,17 @@ def _latency() -> int: if resp.status_code >= 400: detail = (await resp.aread()).decode("utf-8", "replace")[:500] return _err("server_error", resp.status_code, _latency(), detail) - async for line in resp.aiter_lines(): + stream_lines = resp.aiter_lines().__aiter__() + while True: + try: + line = await before_first_output( + stream_lines.__anext__(), first_timeout_s, t0, _emitted) + except StopAsyncIteration: + break + except (asyncio.TimeoutError, TimeoutError): + if not emitted: + return _timeout_err() + raise lines.append(line) if line.startswith("data:"): data = line[len("data:"):].strip() diff --git a/tests/test_codex.py b/tests/test_codex.py index 55b2f40..09d3c93 100644 --- a/tests/test_codex.py +++ b/tests/test_codex.py @@ -239,6 +239,102 @@ def account_id(self): assert seen and seen[0]["status"] == 0 # auth short-circuit observed +class _FakeCodexCallAuth: + def access_token(self): + return "tok" + + def account_id(self): + return "acct" + + +class _SlowCodexStreamResponse: + status_code = 200 + headers = {} + + def __init__(self, lines, open_delay=0, first_line_delay=0): + self._lines = lines + self._open_delay = open_delay + self._first_line_delay = first_line_delay + + async def __aenter__(self): + if self._open_delay: + import asyncio + await asyncio.sleep(self._open_delay) + return self + + async def __aexit__(self, *exc): + return False + + async def aiter_lines(self): + if self._first_line_delay: + import asyncio + await asyncio.sleep(self._first_line_delay) + for line in self._lines: + yield line + + async def aread(self): + return b"" + + +def _patch_codex_async_client(monkeypatch, response): + import httpx + + class FakeAsyncClient: + def __init__(self, timeout=None): + self.timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def stream(self, method, url, json=None, headers=None): + return response + + monkeypatch.setattr(httpx, "AsyncClient", FakeAsyncClient) + + +def test_codex_call_first_token_timeout_opening_stream(monkeypatch): + import asyncio + + response = _SlowCodexStreamResponse( + ['data: {"type": "response.output_text.delta", "delta": "late"}'], + open_delay=0.05, + ) + _patch_codex_async_client(monkeypatch, response) + + call = cb.make_codex_async_call_provider(_FakeCodexCallAuth()) + r = asyncio.run(call({ + "served_model_id": "gpt-5.5", + "messages": [], + "first_token_timeout_ms": 10, + })) + assert r["ok"] is False + assert r["error_kind"] == "timeout" + assert "first token timed out" in r["error_message"] + + +def test_codex_call_first_token_timeout_before_delta(monkeypatch): + import asyncio + + response = _SlowCodexStreamResponse( + ['data: {"type": "response.output_text.delta", "delta": "late"}'], + first_line_delay=0.05, + ) + _patch_codex_async_client(monkeypatch, response) + + call = cb.make_codex_async_call_provider(_FakeCodexCallAuth()) + r = asyncio.run(call({ + "served_model_id": "gpt-5.5", + "messages": [], + "first_token_timeout_ms": 10, + })) + assert r["ok"] is False + assert r["error_kind"] == "timeout" + assert "first token timed out" in r["error_message"] + + def test_messages_to_input_translates_tool_call_history(): # the exact conversation shape that 400d live: every turn after the first # of a tool-using agent carries assistant tool_calls + tool results diff --git a/tests/test_streaming.py b/tests/test_streaming.py index b30d3d7..731d0c1 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -51,6 +51,16 @@ async def aiter_lines(self): yield line +class SlowOpenStreamResponse(FakeStreamResponse): + def __init__(self, delay, lines): + super().__init__(200, lines=lines) + self.delay = delay + + async def __aenter__(self): + await asyncio.sleep(self.delay) + return self + + class FakeStreamClient: def __init__(self, response): self._response = response @@ -163,6 +173,22 @@ async def emit(d): assert "first token timed out" in resp["error_message"] +def test_stream_openai_compatible_first_token_timeout_opening_stream(): + deltas = [] + + async def emit(d): + deltas.append(d) + + req = dict(OPENAI_REQ, first_token_timeout_ms=10) + client = FakeStreamClient(SlowOpenStreamResponse( + 0.05, _openai_lines("late"))) + resp = _collect(streaming.stream_openai_compatible(req, emit, client=client)) + assert deltas == [] + assert resp["ok"] is False + assert resp["error_kind"] == "timeout" + assert "first token timed out" in resp["error_message"] + + # ---- codex stream backend ----------------------------------------------------- CODEX_LINES = [ @@ -213,6 +239,40 @@ async def emit(d): assert resp["error_kind"] == "rate_limit" +def test_stream_codex_first_token_timeout_opening_stream(): + deltas = [] + + async def emit(d): + deltas.append(d) + + client = FakeStreamClient(SlowOpenStreamResponse(0.05, CODEX_LINES)) + resp = _collect(streaming.stream_codex( + {"served_model_id": "gpt-5.5", "messages": [], + "first_token_timeout_ms": 10}, + emit, auth=FakeCodexAuth(), client=client)) + assert deltas == [] + assert resp["ok"] is False + assert resp["error_kind"] == "timeout" + assert "first token timed out" in resp["error_message"] + + +def test_stream_codex_first_token_timeout_before_delta(): + deltas = [] + + async def emit(d): + deltas.append(d) + + client = FakeStreamClient(SlowFirstLineStreamResponse(0.05, CODEX_LINES)) + resp = _collect(streaming.stream_codex( + {"served_model_id": "gpt-5.5", "messages": [], + "first_token_timeout_ms": 10}, + emit, auth=FakeCodexAuth(), client=client)) + assert deltas == [] + assert resp["ok"] is False + assert resp["error_kind"] == "timeout" + assert "first token timed out" in resp["error_message"] + + # ---- SSE chunk encoding ------------------------------------------------------- def test_sse_chunk_encoding():