fix: bridge backend HTTP responses through websocket sessions (#236)

* fix: bridge backend HTTP responses through websocket sessions

* fix: defer bridge turn-state alias registration

Author: Soju06

app/modules/proxy/api.py

@@ -125,7 +125,13 @@ async def responses(
     api_key: ApiKeyData | None = Security(validate_proxy_api_key),
 ) -> Response:
     return await _stream_responses(
-        request, payload, context, api_key, codex_session_affinity=True, openai_cache_affinity=True
+        request,
+        payload,
+        context,
+        api_key,
+        codex_session_affinity=True,
+        openai_cache_affinity=True,
+        prefer_http_bridge=True,
     )
 
 
@@ -437,6 +443,15 @@ async def _stream_responses(
     )
 
     rate_limit_headers = await context.service.rate_limit_headers()
+    bridge_active = prefer_http_bridge and proxy_service_module.get_settings().http_responses_session_bridge_enabled
+    downstream_turn_state = (
+        proxy_service_module.ensure_http_downstream_turn_state(request.headers) if bridge_active else None
+    )
+    turn_state_headers = (
+        proxy_service_module.build_downstream_turn_state_response_headers(downstream_turn_state)
+        if downstream_turn_state is not None
+        else {}
+    )
     payload.stream = True
     if prefer_http_bridge:
         stream = context.service.stream_http_responses(
@@ -448,6 +463,7 @@ async def _stream_responses(
             api_key=api_key,
             api_key_reservation=reservation,
             suppress_text_done_events=suppress_text_done_events,
+            downstream_turn_state=downstream_turn_state,
         )
     else:
         stream = context.service.stream_responses(
@@ -470,11 +486,16 @@ async def _stream_responses(
         )
     except ProxyResponseError as exc:
         await _release_reservation(reservation)
-        return _logged_error_json_response(request, exc.status_code, exc.payload, headers=rate_limit_headers)
+        return _logged_error_json_response(
+            request,
+            exc.status_code,
+            exc.payload,
+            headers=rate_limit_headers,
+        )
     return StreamingResponse(
         _prepend_first(first, stream),
         media_type="text/event-stream",
-        headers={"Cache-Control": "no-cache", **rate_limit_headers},
+        headers={"Cache-Control": "no-cache", **turn_state_headers, **rate_limit_headers},
     )
 
 
@@ -498,6 +519,15 @@ async def _collect_responses(
     )
 
     rate_limit_headers = await context.service.rate_limit_headers()
+    bridge_active = prefer_http_bridge and proxy_service_module.get_settings().http_responses_session_bridge_enabled
+    downstream_turn_state = (
+        proxy_service_module.ensure_http_downstream_turn_state(request.headers) if bridge_active else None
+    )
+    turn_state_headers = (
+        proxy_service_module.build_downstream_turn_state_response_headers(downstream_turn_state)
+        if downstream_turn_state is not None
+        else {}
+    )
     payload.stream = True
     if prefer_http_bridge:
         stream = context.service.stream_http_responses(
@@ -509,6 +539,7 @@ async def _collect_responses(
             api_key=api_key,
             api_key_reservation=reservation,
             suppress_text_done_events=suppress_text_done_events,
+            downstream_turn_state=downstream_turn_state,
         )
     else:
         stream = context.service.stream_responses(
@@ -540,18 +571,18 @@ async def _collect_responses(
                 request,
                 status_code,
                 error_payload.model_dump(mode="json", exclude_none=True),
-                headers=rate_limit_headers,
+                headers={**turn_state_headers, **rate_limit_headers},
             )
         return JSONResponse(
             content=response_payload.model_dump(mode="json", exclude_none=True),
-            headers=rate_limit_headers,
+            headers={**turn_state_headers, **rate_limit_headers},
         )
     status_code = _status_for_error(response_payload.error)
     return _logged_error_json_response(
         request,
         status_code,
         response_payload.model_dump(mode="json", exclude_none=True),
-        headers=rate_limit_headers,
+        headers={**turn_state_headers, **rate_limit_headers},
     )
 
 

app/modules/proxy/service.py

@@ -7,7 +7,7 @@
 import time
 from collections import deque
 from collections.abc import Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from hashlib import sha256
 from typing import AsyncIterator, Mapping, NoReturn
 from uuid import uuid4
@@ -143,6 +143,7 @@ def __init__(self, repo_factory: ProxyRepoFactory) -> None:
         self._encryptor = TokenEncryptor()
         self._load_balancer = LoadBalancer(repo_factory)
         self._http_bridge_sessions: dict[_HTTPBridgeSessionKey, _HTTPBridgeSession] = {}
+        self._http_bridge_turn_state_index: dict[tuple[str, str | None], _HTTPBridgeSessionKey] = {}
         self._http_bridge_lock = anyio.Lock()
 
     def stream_responses(
@@ -183,6 +184,7 @@ def stream_http_responses(
         api_key: ApiKeyData | None = None,
         api_key_reservation: ApiKeyUsageReservationData | None = None,
         suppress_text_done_events: bool = False,
+        downstream_turn_state: str | None = None,
     ) -> AsyncIterator[str]:
         _maybe_log_proxy_request_payload("stream_http", payload, headers)
         filtered = filter_inbound_headers(headers)
@@ -195,6 +197,7 @@ def stream_http_responses(
             api_key=api_key,
             api_key_reservation=api_key_reservation,
             suppress_text_done_events=suppress_text_done_events,
+            downstream_turn_state=downstream_turn_state,
         )
 
     async def _stream_http_bridge_or_retry(
@@ -208,6 +211,7 @@ async def _stream_http_bridge_or_retry(
         api_key: ApiKeyData | None,
         api_key_reservation: ApiKeyUsageReservationData | None,
         suppress_text_done_events: bool,
+        downstream_turn_state: str | None = None,
     ) -> AsyncIterator[str]:
         settings = await get_settings_cache().get()
         if not _http_responses_session_bridge_enabled(settings):
@@ -238,6 +242,7 @@ async def _stream_http_bridge_or_retry(
             codex_idle_ttl_seconds=getattr(settings, "http_responses_session_bridge_codex_idle_ttl_seconds", 900.0),
             max_sessions=getattr(settings, "http_responses_session_bridge_max_sessions", 256),
             queue_limit=getattr(settings, "http_responses_session_bridge_queue_limit", 8),
+            downstream_turn_state=downstream_turn_state,
         ):
             yield line
 
@@ -256,6 +261,7 @@ async def _stream_via_http_bridge(
         codex_idle_ttl_seconds: float,
         max_sessions: int,
         queue_limit: int,
+        downstream_turn_state: str | None = None,
     ) -> AsyncIterator[str]:
         del propagate_http_errors, suppress_text_done_events
         request_id = ensure_request_id()
@@ -306,7 +312,6 @@ async def _stream_via_http_bridge(
             ),
             max_sessions=max_sessions,
         )
-
         request_state, text_data = self._prepare_http_bridge_request(
             payload,
             api_key=api_key,
@@ -320,6 +325,8 @@ async def _stream_via_http_bridge(
             text_data=text_data,
             queue_limit=queue_limit,
         )
+        if downstream_turn_state is not None:
+            await self._register_http_bridge_turn_state(session, downstream_turn_state)
 
         try:
             event_queue = request_state.event_queue
@@ -1398,8 +1405,8 @@ async def _get_or_create_http_bridge_session(
         idle_ttl_seconds: float,
         max_sessions: int,
     ) -> "_HTTPBridgeSession":
-        del api_key
         settings = get_settings()
+        api_key_id = api_key.id if api_key is not None else None
         effective_idle_ttl_seconds = _effective_http_bridge_idle_ttl_seconds(
             affinity=affinity,
             idle_ttl_seconds=idle_ttl_seconds,
@@ -1409,36 +1416,64 @@ async def _get_or_create_http_bridge_session(
                 900.0,
             ),
         )
-        owner_instance = _http_bridge_owner_instance(key, settings)
-        current_instance, ring = _normalized_http_bridge_instance_ring(settings)
-        if (
-            key.affinity_kind != "request"
-            and owner_instance is not None
-            and len(ring) > 1
-            and owner_instance != current_instance
-        ):
-            _log_http_bridge_event(
-                "owner_mismatch",
-                key,
-                account_id=None,
-                model=request_model,
-                detail=f"expected_instance={owner_instance}, current_instance={current_instance}",
-            )
-            raise ProxyResponseError(
-                409,
-                openai_error(
-                    "bridge_instance_mismatch",
-                    (
-                        "HTTP responses session bridge request reached the wrong instance "
-                        f"(expected {owner_instance}, got {current_instance})"
-                    ),
-                    error_type="server_error",
-                ),
-            )
         async with self._http_bridge_lock:
+            incoming_turn_state = _sticky_key_from_turn_state_header(headers)
+            if incoming_turn_state is not None:
+                alias_key = self._http_bridge_turn_state_index.get(
+                    _http_bridge_turn_state_alias_key(incoming_turn_state, api_key_id)
+                )
+                if alias_key is not None:
+                    key = alias_key
+                elif incoming_turn_state.startswith("http_turn_"):
+                    raise ProxyResponseError(
+                        409,
+                        openai_error(
+                            "bridge_instance_mismatch",
+                            "HTTP bridge turn-state reached an instance that does not own the live session",
+                            error_type="server_error",
+                        ),
+                    )
+            owner_instance = _http_bridge_owner_instance(key, settings)
+            current_instance, ring = _normalized_http_bridge_instance_ring(settings)
+            if (
+                key.affinity_kind != "request"
+                and owner_instance is not None
+                and len(ring) > 1
+                and owner_instance != current_instance
+            ):
+                _log_http_bridge_event(
+                    "owner_mismatch",
+                    key,
+                    account_id=None,
+                    model=request_model,
+                    detail=f"expected_instance={owner_instance}, current_instance={current_instance}",
+                )
+                raise ProxyResponseError(
+                    409,
+                    openai_error(
+                        "bridge_instance_mismatch",
+                        (
+                            "HTTP responses session bridge request reached the wrong instance "
+                            f"(expected {owner_instance}, got {current_instance})"
+                        ),
+                        error_type="server_error",
+                    ),
+                )
             await self._prune_http_bridge_sessions_locked()
             existing = self._http_bridge_sessions.get(key)
             if existing is not None and not existing.closed and existing.account.status == AccountStatus.ACTIVE:
+                if (
+                    incoming_turn_state is not None
+                    and self._http_bridge_turn_state_index.get(
+                        _http_bridge_turn_state_alias_key(incoming_turn_state, api_key_id)
+                    )
+                    == key
+                ):
+                    self._promote_http_bridge_session_to_codex_affinity(
+                        existing,
+                        turn_state=incoming_turn_state,
+                        settings=settings,
+                    )
                 existing.request_model = request_model
                 existing.last_used_at = time.monotonic()
                 _log_http_bridge_event(
@@ -1538,6 +1573,11 @@ async def _prune_http_bridge_sessions_locked(self) -> None:
 
     async def _close_http_bridge_session(self, session: "_HTTPBridgeSession") -> None:
         session.closed = True
+        for alias in session.downstream_turn_state_aliases:
+            self._http_bridge_turn_state_index.pop(
+                _http_bridge_turn_state_alias_key(alias, session.key.api_key_id),
+                None,
+            )
         if session.upstream_reader is not None:
             session.upstream_reader.cancel()
             try:
@@ -1555,6 +1595,34 @@ async def _close_http_bridge_session(self, session: "_HTTPBridgeSession") -> Non
             model=session.request_model,
         )
 
+    async def _register_http_bridge_turn_state(self, session: "_HTTPBridgeSession", turn_state: str) -> None:
+        async with self._http_bridge_lock:
+            if session.closed:
+                return
+            session.downstream_turn_state_aliases.add(turn_state)
+            if session.downstream_turn_state is None:
+                session.downstream_turn_state = turn_state
+            self._http_bridge_turn_state_index[
+                _http_bridge_turn_state_alias_key(turn_state, session.key.api_key_id)
+            ] = session.key
+
+    def _promote_http_bridge_session_to_codex_affinity(
+        self,
+        session: "_HTTPBridgeSession",
+        *,
+        turn_state: str,
+        settings: object,
+    ) -> None:
+        session.affinity = _AffinityPolicy(key=turn_state, kind=StickySessionKind.CODEX_SESSION)
+        session.codex_session = True
+        session.downstream_turn_state = turn_state
+        session.downstream_turn_state_aliases.add(turn_state)
+        session.idle_ttl_seconds = max(
+            session.idle_ttl_seconds,
+            float(getattr(settings, "http_responses_session_bridge_codex_idle_ttl_seconds", 900.0)),
+        )
+        session.headers = _headers_with_turn_state(session.headers, turn_state)
+
     async def _create_http_bridge_session(
         self,
         key: "_HTTPBridgeSessionKey",
@@ -1621,6 +1689,7 @@ async def _create_http_bridge_session(
             codex_session=affinity.kind == StickySessionKind.CODEX_SESSION,
             prewarm_lock=anyio.Lock(),
             upstream_turn_state=_upstream_turn_state_from_socket(upstream),
+            downstream_turn_state=None,
         )
         session.upstream_reader = asyncio.create_task(self._relay_http_bridge_upstream_messages(session))
         return session
@@ -2008,7 +2077,10 @@ async def _reconnect_http_bridge_session(
                 ),
             )
         account = await self._ensure_fresh_with_budget(account, timeout_seconds=_remaining_budget_seconds(deadline))
-        connect_headers = _headers_with_turn_state(session.headers, session.upstream_turn_state)
+        connect_headers = _headers_with_turn_state(
+            session.headers,
+            _preferred_http_bridge_reconnect_turn_state(session),
+        )
         upstream = await self._open_upstream_websocket_with_budget(
             account,
             connect_headers,
@@ -3983,6 +4055,8 @@ class _HTTPBridgeSession:
     prewarmed: bool = False
     prewarm_lock: anyio.Lock | None = None
     upstream_turn_state: str | None = None
+    downstream_turn_state: str | None = None
+    downstream_turn_state_aliases: set[str] = field(default_factory=set)
     upstream_reader: asyncio.Task[None] | None = None
     closed: bool = False
 
@@ -4507,10 +4581,21 @@ def ensure_downstream_turn_state(headers: Mapping[str, str]) -> str:
     return f"turn_{uuid4().hex}"
 
 
+def ensure_http_downstream_turn_state(headers: Mapping[str, str]) -> str:
+    existing = _sticky_key_from_turn_state_header(headers)
+    if existing is not None:
+        return existing
+    return f"http_turn_{uuid4().hex}"
+
+
 def build_downstream_turn_state_accept_headers(turn_state: str) -> list[tuple[bytes, bytes]]:
     return [(b"x-codex-turn-state", turn_state.encode("utf-8"))]
 
 
+def build_downstream_turn_state_response_headers(turn_state: str) -> dict[str, str]:
+    return {"x-codex-turn-state": turn_state}
+
+
 def _upstream_turn_state_from_socket(upstream: UpstreamResponsesWebSocket | None) -> str | None:
     if upstream is None:
         return None
@@ -4531,6 +4616,21 @@ def _headers_with_turn_state(headers: Mapping[str, str], turn_state: str | None)
     return forwarded
 
 
+def _preferred_http_bridge_reconnect_turn_state(session: "_HTTPBridgeSession") -> str | None:
+    if (
+        session.codex_session
+        and session.downstream_turn_state is not None
+        and session.affinity.kind == StickySessionKind.CODEX_SESSION
+        and session.affinity.key == session.downstream_turn_state
+    ):
+        return session.downstream_turn_state
+    return session.upstream_turn_state
+
+
+def _http_bridge_turn_state_alias_key(turn_state: str, api_key_id: str | None) -> tuple[str, str | None]:
+    return (turn_state, api_key_id)
+
+
 def _resolve_prompt_cache_key(
     payload: ResponsesRequest | ResponsesCompactRequest,
     *,