fix(proxy): retry missing response.created once

app/core/clients/proxy_websocket.py

@@ -21,6 +21,7 @@
     InvalidProxy,
     InvalidStatus,
 )
+from websockets.protocol import State
 from websockets.typing import Origin, Subprotocol
 
 from app.core.clients.codex import (
@@ -157,6 +158,46 @@ def __init__(self, message: str, *, error_code: str) -> None:
         self.error_code = error_code
 
 
+_SEND_NOT_DISPATCHED_CONSTRUCTOR_TOKEN = object()
+
+
+class UpstreamWebSocketSendNotDispatchedError(UpstreamWebSocketTransportError):
+    """Proof that a websocket frame was rejected before dispatch began.
+
+    The constructor token is module-private, subclassing is disabled, and the
+    adapters are the only production construction path. Python reflection can
+    still bypass these boundaries, so this is a runtime application contract,
+    not a security boundary against hostile in-process code.
+    """
+
+    def __init_subclass__(cls, **kwargs: object) -> None:
+        del cls, kwargs
+        raise TypeError("UpstreamWebSocketSendNotDispatchedError cannot be subclassed")
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        error_code: str,
+        _constructor_token: object,
+    ) -> None:
+        if _constructor_token is not _SEND_NOT_DISPATCHED_CONSTRUCTOR_TOKEN:
+            raise TypeError("Use the websocket adapter's closed-before-send check")
+        super().__init__(message, error_code=error_code)
+
+
+def _websocket_send_not_dispatched_error(
+    *,
+    endpoint_id: str | None = None,
+) -> UpstreamWebSocketSendNotDispatchedError:
+    endpoint_suffix = f" via endpoint {endpoint_id}" if endpoint_id else ""
+    return UpstreamWebSocketSendNotDispatchedError(
+        f"Upstream websocket was already closed before send{endpoint_suffix}",
+        error_code="upstream_websocket_closed_before_send",
+        _constructor_token=_SEND_NOT_DISPATCHED_CONSTRUCTOR_TOKEN,
+    )
+
+
 def _websocket_transport_error_code(exc: BaseException, *, uses_proxy: bool) -> str:
     return process_network_error_code(
         exc,
@@ -225,12 +266,16 @@ def __init__(
         self._preserve_close_semantics = preserve_close_semantics
 
     async def send_text(self, text: str) -> None:
+        if getattr(self._connection, "state", State.OPEN) is not State.OPEN:
+            raise _websocket_send_not_dispatched_error()
         try:
             await self._connection.send(text)
         except Exception as exc:
             await _raise_websocket_send_error(exc, uses_proxy=self._uses_proxy)
 
     async def send_bytes(self, data: bytes) -> None:
+        if getattr(self._connection, "state", State.OPEN) is not State.OPEN:
+            raise _websocket_send_not_dispatched_error()
         try:
             await self._connection.send(data)
         except Exception as exc:
@@ -319,6 +364,8 @@ def __init__(
         self._response_headers = _normalize_response_headers(response_headers)
 
     async def send_text(self, text: str) -> None:
+        if bool(getattr(self._websocket, "closed", False)):
+            raise _websocket_send_not_dispatched_error(endpoint_id=self._endpoint_id)
         try:
             result = self._websocket.send_str(text)
             if asyncio.iscoroutine(result):
@@ -327,6 +374,8 @@ async def send_text(self, text: str) -> None:
             await _raise_websocket_send_error(exc, endpoint_id=self._endpoint_id, uses_proxy=True)
 
     async def send_bytes(self, data: bytes) -> None:
+        if bool(getattr(self._websocket, "closed", False)):
+            raise _websocket_send_not_dispatched_error(endpoint_id=self._endpoint_id)
         try:
             result = self._websocket.send_bytes(data)
             if asyncio.iscoroutine(result):

app/modules/proxy/_service/http_bridge/helpers.py

@@ -188,7 +188,7 @@
 
 logger = logging.getLogger("app.modules.proxy.service")
 _HTTP_BRIDGE_BACKGROUND_CLOSE_TIMEOUT_SECONDS = 5.0
-_HTTP_BRIDGE_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS = 240.0
+_HTTP_BRIDGE_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS = 30.0
 _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL = "missing_response_created_timeout"
 T = TypeVar("T")
 

app/modules/proxy/_service/http_bridge/request_submit.py

@@ -36,7 +36,10 @@
 from app.core.clients.proxy import codex_control_request as core_codex_control_request  # noqa: F401
 from app.core.clients.proxy import compact_responses as core_compact_responses  # noqa: F401
 from app.core.clients.proxy import transcribe_audio as core_transcribe_audio  # noqa: F401
-from app.core.clients.proxy_websocket import UpstreamWebSocketTransportError
+from app.core.clients.proxy_websocket import (
+    UpstreamWebSocketSendNotDispatchedError,
+    UpstreamWebSocketTransportError,
+)
 from app.core.errors import (
     openai_error,
 )
@@ -876,6 +879,25 @@ async def _submit_http_bridge_request_with_handoff(
                     upstream_send_started = True
                     try:
                         await _send_http_bridge_request_text_with_archive_id(session, request_state, text_data)
+                    except UpstreamWebSocketSendNotDispatchedError:
+                        # The adapter observed the warm socket closed before
+                        # dispatch began, so the exact request remains safe to
+                        # send once on a fresh socket with the same leased
+                        # account. Mid-send failures retain the fail-closed path
+                        # below because their dispatch state is ambiguous.
+                        session.closed = True
+                        session.upstream_control.reconnect_requested = True
+                        session.upstream_control.retire_after_drain = True
+                        recovered = await self._retry_http_bridge_request_on_fresh_upstream(
+                            session,
+                            request_state=request_state,
+                            text_data=text_data,
+                            require_same_account=True,
+                            require_preferred_account=True,
+                            dispatch_proven_absent=True,
+                        )
+                        if not recovered:
+                            raise
                     except BaseException:
                         # Publish retirement while lifecycle ownership is still
                         # held; a gate waiter must never reuse an ambiguously sent
@@ -1456,14 +1478,16 @@ async def _retry_http_bridge_request_on_fresh_upstream(
         text_data: str,
         send_request: bool = True,
         require_same_account: bool = False,
+        require_preferred_account: bool = False,
+        dispatch_proven_absent: bool = False,
     ) -> bool:
         require_same_account = require_same_account or is_http_bridge_account_neutral_replay(
             kind=session.key.affinity_kind,
             key=session.key.affinity_key,
         )
         retry_text_data = text_data
         using_fresh_replay = False
-        if request_state.previous_response_id is not None and send_request:
+        if request_state.previous_response_id is not None and send_request and not dispatch_proven_absent:
             # After an ambiguous websocket send failure we cannot prove whether
             # upstream already accepted the continuation. Re-sending the same
             # previous_response_id request can fork continuity with duplicate
@@ -1486,6 +1510,8 @@ async def _retry_http_bridge_request_on_fresh_upstream(
         if request_state.response_event_count > 0:
             return False
         request_state.replay_count += 1
+        if require_preferred_account:
+            request_state.preferred_account_id = session.account.id
         _log_http_bridge_event(
             "retry_fresh_upstream",
             session.key,
@@ -1501,6 +1527,7 @@ async def _retry_http_bridge_request_on_fresh_upstream(
                 request_state=request_state,
                 restart_reader=True,
                 require_same_account=require_same_account,
+                require_preferred_account=require_preferred_account,
             )
             if send_request:
                 retry_text_data = self._http_bridge_text_with_account_installation_id(
@@ -1530,12 +1557,13 @@ async def _retry_http_bridge_precreated_request(
         session: "_HTTPBridgeSession",
         *,
         request_state: _WebSocketRequestState | None = None,
+        require_current_account: bool = False,
     ) -> bool:
         account_neutral_recovery = is_http_bridge_account_neutral_replay(
             kind=session.key.affinity_kind,
             key=session.key.affinity_key,
         )
-        hard_owner_bound = _http_bridge_key_strength(session.key) == "hard"
+        hard_owner_bound = _http_bridge_key_strength(session.key) == "hard" or require_current_account
         async with session.pending_lock:
             if request_state is not None:
                 if (
@@ -1556,6 +1584,8 @@ async def _retry_http_bridge_precreated_request(
                 if len(retryable_requests) != 1:
                     return False
                 request_state = retryable_requests[0]
+            if require_current_account:
+                request_state.preferred_account_id = session.account.id
             if request_state.previous_response_id is not None and not (
                 request_state.proxy_injected_previous_response_id
                 and request_state.fresh_upstream_request_is_retry_safe
@@ -1611,7 +1641,7 @@ async def _retry_http_bridge_precreated_request(
                 request_text = _prepare_websocket_request_state_for_visible_output_replay(request_state)
                 if request_text is None:
                     return False
-                if account_neutral_recovery:
+                if account_neutral_recovery or require_current_account:
                     request_state.preferred_account_id = session.account.id
                 elif not request_state.file_required_preferred_account and not hard_owner_bound:
                     request_state.preferred_account_id = None
@@ -1637,6 +1667,7 @@ async def _retry_http_bridge_precreated_request(
                     session,
                     request_state=request_state,
                     require_same_account=True,
+                    require_preferred_account=require_current_account,
                 )
             elif require_preferred_reconnect:
                 await self._reconnect_http_bridge_session(

app/modules/proxy/_service/http_bridge/upstream_events.py

@@ -610,37 +610,73 @@ async def _relay_http_bridge_upstream_messages(
                             async with session.pending_lock:
                                 if receive_task is not None and receive_task.done():
                                     continue
-                                expired_owner = any(
-                                    deadline is not None and deadline <= now
-                                    for request_state in session.pending_requests
-                                    if (
-                                        deadline := _http_bridge_eventless_precreated_deadline(
-                                            request_state,
-                                            stuck_gate_retire_after_seconds=stuck_gate_retire_after_seconds,
+                                eventless_owner = next(
+                                    (
+                                        request_state
+                                        for request_state in session.pending_requests
+                                        if (
+                                            deadline := _http_bridge_eventless_precreated_deadline(
+                                                request_state,
+                                                stuck_gate_retire_after_seconds=stuck_gate_retire_after_seconds,
+                                            )
                                         )
-                                    )
-                                    is not None
+                                        is not None
+                                        and deadline <= now
+                                    ),
+                                    None,
                                 )
-                                if not expired_owner:
+                                if eventless_owner is None:
                                     continue
                                 pending_count = len(session.pending_requests)
+                                can_retry_eventless_owner = pending_count == 1
+                            receive_cancelled = True
+                            if receive_task is not None:
+                                cancel_requested = receive_task.cancel()
+                                if not cancel_requested:
+                                    continue
+                                try:
+                                    await receive_task
+                                except asyncio.CancelledError:
+                                    relay_task = asyncio.current_task()
+                                    if relay_task is not None and relay_task.cancelling():
+                                        raise
+                                    receive_task = None
+                                except Exception:
+                                    # Preserve the completed task so the next
+                                    # loop iteration raises it outside the
+                                    # lifecycle lock and uses normal cleanup.
+                                    continue
+                                else:
+                                    # A response may win the cancellation race.
+                                    # Leave the completed task in place so the
+                                    # next loop iteration processes its result.
+                                    continue
+                            retried = False
+                            if not session.closed and can_retry_eventless_owner and receive_cancelled:
+                                try:
+                                    retried = await self._retry_http_bridge_precreated_request(
+                                        session,
+                                        request_state=eventless_owner,
+                                        require_current_account=True,
+                                    )
+                                except UpstreamWebSocketTransportError:
+                                    logger.warning(
+                                        "HTTP bridge missing response.created retry transport failed",
+                                        exc_info=True,
+                                    )
+                            if retried:
+                                continue
+                            async with session.pending_lock:
                                 for request_state in session.pending_requests:
                                     if request_state.failure_phase_override is None:
                                         request_state.failure_phase_override = "upstream"
                                     if request_state.failure_detail_override is None:
                                         request_state.failure_detail_override = (
                                             _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL
                                         )
-                            # Claim the session before cancelling receive so a
+                            # Claim the session before terminal settlement so a
                             # gate waiter cannot reopen this ambiguous socket.
                             session.closed = True
-                            if receive_task is not None:
-                                receive_cancelled = await _cancel_http_bridge_reader_child(
-                                    receive_task,
-                                    label="HTTP bridge upstream receive after missing response.created",
-                                )
-                                if receive_cancelled:
-                                    receive_task = None
                             _record_http_bridge_stuck_retire(
                                 reason=_HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL,
                                 session=session,

app/modules/proxy/_service/support.py

@@ -819,6 +819,11 @@ class _WebSocketRequestState:
     fresh_upstream_request_responses_lite_model: str | None = None
     request_stage: str = "first_turn"
     preferred_account_id: str | None = None
+    # A closed-before-send transport proof permits one exact replay, but the
+    # request's response-create lease belongs to the account selected for the
+    # undispatched attempt. Keep that account mandatory until the resend
+    # succeeds.
+    dispatch_absent_replay_account_id: str | None = None
     require_security_work_authorized: bool = False
     file_required_preferred_account: bool = False
     bridge_soft_capacity_reroute_allowed: bool = False