fix(http-bridge): wait for aborted bridge owner

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

@@ -68,26 +68,29 @@
 from app.modules.proxy._service.http_bridge.activity import _HTTPBridgeActivityMixin
 from app.modules.proxy._service.http_bridge.helpers import (
     _HTTP_BRIDGE_BACKGROUND_CLOSE_TIMEOUT_SECONDS,
-    _HTTP_BRIDGE_INFLIGHT_STARTED_AT_ATTR,
     _active_http_bridge_instance_ring,
     _alias_fallback_key,
     _durable_bridge_lookup_active_owner,
     _durable_bridge_lookup_allows_local_reuse,
+    _evict_http_bridge_retained_capacity_waiter_after_error,
     _forwarded_http_bridge_session_key,
     _http_bridge_alias_target_is_stale,
     _http_bridge_allow_durable_takeover,
     _http_bridge_can_local_recover_without_ring,
     _http_bridge_can_recover_during_drain,
     _http_bridge_can_single_instance_owner_takeover_without_anchor,
     _http_bridge_can_single_instance_prompt_cache_takeover_without_anchor,
+    _http_bridge_canonical_inflight_key_locked,
     _http_bridge_capacity_after_planned_closes,
     _http_bridge_claim_allows_takeover,
     _http_bridge_compatible,
     _http_bridge_continuity_lost_error_envelope,
     _http_bridge_endpoint_matches_current_instance,
     _http_bridge_has_durable_recovery_anchor,
     _http_bridge_incompatible_model_fork_key,
+    _http_bridge_inflight_creation_can_register,
     _http_bridge_inflight_creation_count,
+    _http_bridge_key_is_synthesized_turn_state,
     _http_bridge_key_strength,
     _http_bridge_locally_owned_fork_key,
     _http_bridge_models_compatible,
@@ -113,8 +116,10 @@
     _http_bridge_should_wait_for_registration,
     _http_bridge_startup_wait_timeout_error,
     _http_bridge_turn_state_alias_key,
+    _http_bridge_turn_state_key_from,
     _log_http_bridge_event,
     _log_http_bridge_startup_wait_timeout,
+    _mark_http_bridge_inflight_creation_owner,
     _mark_http_bridge_reader_handoff_reconnect_failed,
     _persist_http_bridge_replacement_account,
     _persistent_http_bridge_affinity,
@@ -126,8 +131,9 @@
     _register_http_bridge_turn_state_aliases_locked,
     _require_http_bridge_bound_account_not_excluded,
     _reserve_http_bridge_unanchored_handoff,
-    _settle_failed_http_bridge_creation,
+    _settle_and_close_failed_http_bridge_creation,
     _turn_keys,
+    _wait_for_http_bridge_aborted_owner_within_budget,
 )
 from app.modules.proxy._service.http_bridge.helpers import (
     _close_http_bridge_session as _helpers_close_http_bridge_session,
@@ -288,51 +294,6 @@ async def _drain_http_bridge_background_cleanup_tasks(self, *, reason: str) -> b
         self._http_bridge_background_cleanup_failed |= any(isinstance(result, BaseException) for result in results)
         return not self._http_bridge_background_cleanup_failed
 
-    async def _fail_http_bridge_inflight_session_creation(
-        self,
-        key: "_HTTPBridgeSessionKey",
-        inflight_future: asyncio.Future["_HTTPBridgeSession"] | None,
-        exc: BaseException,
-    ) -> bool:
-        if inflight_future is None:
-            return False
-        async with self._http_bridge_lock:
-            current_future = self._http_bridge_inflight_sessions.get(key)
-            if current_future is not inflight_future:
-                return False
-            if getattr(inflight_future, "_http_bridge_handoff", False):
-                return False
-            self._http_bridge_inflight_sessions.pop(key, None)
-            if inflight_future.done():
-                return True
-            if isinstance(exc, asyncio.CancelledError):
-                inflight_future.cancel()
-            else:
-                inflight_future.set_exception(exc)
-                inflight_future.exception()
-            return True
-
-    async def _evict_http_bridge_inflight_waiter(
-        self,
-        inflight_future: asyncio.Future["_HTTPBridgeSession"],
-        exc: BaseException,
-    ) -> "_HTTPBridgeSessionKey | None":
-        async with self._http_bridge_lock:
-            stale_key = None
-            for candidate_key, candidate_future in self._http_bridge_inflight_sessions.items():
-                if candidate_future is inflight_future:
-                    stale_key = candidate_key
-                    break
-            if stale_key is None:
-                return None
-            if getattr(inflight_future, "_http_bridge_handoff", False):
-                return None
-            self._http_bridge_inflight_sessions.pop(stale_key, None)
-            if not inflight_future.done():
-                inflight_future.set_exception(exc)
-                inflight_future.exception()
-            return stale_key
-
     @overload
     async def _get_or_create_http_bridge_session(
         self,
@@ -434,6 +395,7 @@ async def _get_or_create_http_bridge_session(
         settings = _service_get_settings()
         request_scope_id = ensure_request_scope_id()
         api_key_id = api_key.id if api_key is not None else None
+        requested_key = key
         incoming_turn_state = _sticky_key_from_turn_state_header(headers)
         incoming_session_key, initial_session_key = _turn_keys(headers, api_key, key, session_header_fallback_key)
         original_request_unanchored = _http_bridge_request_needs_unanchored_handoff(
@@ -574,7 +536,7 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                                 bind_account_neutral_recovery_owner(alias_session)
                                 continue
                             self._http_bridge_turn_state_index.pop(alias_index_key, None)
-                            key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id)
+                            key = _http_bridge_turn_state_key_from(requested_key, incoming_turn_state, api_key_id)
                         elif not _http_bridge_models_compatible(alias_session.request_model, request_model):
                             model_transition_rebind, model_transition_parent_key = True, alias_key
                             if is_http_bridge_account_neutral_replay(
@@ -592,7 +554,7 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                                 key = recovery_fork_key
                                 continue
                             else:
-                                key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id)
+                                key = _http_bridge_turn_state_key_from(requested_key, incoming_turn_state, api_key_id)
                         elif not _http_bridge_compatible(
                             alias_session, request_model, request_service_tier, True
                         ) or not _http_bridge_session_matches_preferred_account(
@@ -676,21 +638,25 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                                 kind=key.affinity_kind,
                                 key=key.affinity_key,
                             ):
-                                key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id)
+                                key = _http_bridge_turn_state_key_from(requested_key, incoming_turn_state, api_key_id)
                         elif (
                             fallback_key := _alias_fallback_key(incoming_session_key, initial_session_key, api_key_id)
                         ) is not None:
                             key = fallback_key
                             used_session_header_fallback = True
                         else:
-                            key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id)
+                            key = _http_bridge_turn_state_key_from(requested_key, incoming_turn_state, api_key_id)
                             missing_turn_state_alias = True
                 pruned_sessions = self._prune_http_bridge_sessions_locked()
                 if pruned_sessions:
                     if any(session.key == key for session in pruned_sessions):
                         force_durable_takeover = True
                     self._schedule_http_bridge_session_closes(pruned_sessions, reason="registry_detach")
                 existing = self._http_bridge_sessions.get(key)
+                if existing is not None:
+                    key = existing.key
+                else:
+                    key = _http_bridge_canonical_inflight_key_locked(self, key)
                 retained_handoff = bool(
                     existing and existing.closed and _http_bridge_session_has_admission_waiter(existing)
                 )
@@ -1332,10 +1298,9 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                                 capacity_error_after_planned_closes = capacity_error
                         else:
                             inflight_future = asyncio.get_running_loop().create_future()
-                            setattr(
+                            _mark_http_bridge_inflight_creation_owner(
                                 inflight_future,
-                                _HTTP_BRIDGE_INFLIGHT_STARTED_AT_ATTR,
-                                clock_for(self).monotonic(),
+                                started_at=clock_for(self).monotonic(),
                             )
                             self._http_bridge_inflight_sessions[key] = inflight_future
                             owns_creation = True
@@ -1385,11 +1350,20 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                         pending_count=_http_bridge_session_generation_count(self),
                         inflight_count=len(self._http_bridge_inflight_sessions),
                     )
+                    if await _wait_for_http_bridge_aborted_owner_within_budget(
+                        self, capacity_wait_future, wait_timeout_seconds, request_deadline
+                    ):
+                        continue
                     raise timeout_error from exc
                 except ProxyResponseError:
                     raise
                 except Exception:
-                    pass
+                    await _evict_http_bridge_retained_capacity_waiter_after_error(
+                        self,
+                        capacity_wait_future,
+                        timeout=wait_timeout_seconds,
+                        request_deadline=request_deadline,
+                    )
                 continue
             if inflight_future is not None and not owns_creation:
                 wait_timeout_seconds = _proxy_admission_wait_timeout_seconds()
@@ -1413,6 +1387,12 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                         pending_count=_http_bridge_session_generation_count(self),
                         inflight_count=len(self._http_bridge_inflight_sessions),
                     )
+                    if _http_bridge_key_is_synthesized_turn_state(key):
+                        raise timeout_error from exc
+                    if await _wait_for_http_bridge_aborted_owner_within_budget(
+                        self, inflight_future, wait_timeout_seconds, request_deadline
+                    ):
+                        continue
                     raise timeout_error from exc
                 except Exception:
                     raise
@@ -1553,7 +1533,9 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                 await self._claim_durable_http_bridge_session(created_session, **claim_kwargs)
                 async with self._http_bridge_lock:
                     current_future = self._http_bridge_inflight_sessions.get(key)
-                    if current_future is inflight_future:
+                    if current_future is inflight_future and _http_bridge_inflight_creation_can_register(
+                        inflight_future
+                    ):
                         self._http_bridge_inflight_sessions.pop(key, None)
                         if original_request_unanchored:
                             _reserve_http_bridge_unanchored_handoff(created_session, request_scope_id=request_scope_id)
@@ -1567,18 +1549,14 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None:
                         code="capacity_exhausted_active_sessions",
                     )
             except BaseException as exc:
-                superseded = await _settle_failed_http_bridge_creation(
+                await _settle_and_close_failed_http_bridge_creation(
                     self,
                     key,
                     inflight_future=inflight_future,
                     created_session=created_session,
+                    session_registered=session_registered,
                     exc=exc,
                 )
-                if created_session is not None and not session_registered:
-                    await self._close_http_bridge_session(
-                        created_session,
-                        release_durable_session=not superseded,
-                    )
                 raise
             assert created_session is not None
             _log_http_bridge_event(

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

@@ -198,6 +198,28 @@ def _owner_forward_failure_allows_local_recovery(exc: ProxyResponseError) -> boo
     }
 
 
+def _owner_forward_failure_was_pre_dispatch(exc: ProxyResponseError) -> bool:
+    """Return whether the owner failed before accepting the request upstream."""
+
+    return isinstance(exc, _OwnerForwardRequestError) and exc.outcome in {
+        _OwnerForwardOutcome.NOT_DISPATCHED,
+        _OwnerForwardOutcome.RECEIVER_REJECTED,
+    }
+
+
+def _owner_forward_outcome_for_proxy_error(
+    exc: ProxyResponseError,
+    *,
+    default: _OwnerForwardOutcome,
+) -> _OwnerForwardOutcome:
+    payload = exc.payload
+    if isinstance(payload, dict):
+        error = payload.get("error")
+        if isinstance(error, dict) and error.get("code") == "bridge_drain_active":
+            return _OwnerForwardOutcome.RECEIVER_REJECTED
+    return default
+
+
 def _durable_recovery_supersedes_local_session(
     durable_lookup: DurableBridgeLookup | None,
     session: _HTTPBridgeSession,
@@ -397,6 +419,7 @@ async def _forward_http_bridge_request_to_owner(
         downstream_turn_state: str | None,
         request_started_at: float,
         proxy_api_authorization: str | None,
+        downstream_turn_state_synthesized: bool = False,
         file_owner_account_id: str | None = None,
         client_ip: str | None = None,
     ) -> AsyncIterator[str]:
@@ -419,6 +442,9 @@ async def _forward_http_bridge_request_to_owner(
             reservation=api_key_reservation,
             codex_session_affinity=codex_session_affinity,
             downstream_turn_state=forwarded_turn_state,
+            downstream_turn_state_synthesized=(
+                downstream_turn_state_synthesized and forwarded_turn_state == downstream_turn_state
+            ),
             original_request_unanchored=(
                 recovery_forward
                 or (
@@ -552,7 +578,10 @@ def owner_response_rejected() -> None:
                     default_message="HTTP bridge owner request failed",
                 )
                 return
-            raise _OwnerForwardRequestError(exc, outcome=forward_outcome) from exc
+            raise _OwnerForwardRequestError(
+                exc,
+                outcome=_owner_forward_outcome_for_proxy_error(exc, default=forward_outcome),
+            ) from exc
         except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
             if PROMETHEUS_AVAILABLE and bridge_owner_forward_total is not None:
                 bridge_owner_forward_total.labels(outcome="fail").inc()

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

@@ -115,13 +115,16 @@ async def _register_http_bridge_turn_state_impl(
         self,
         session: _HTTPBridgeSession,
         turn_state: str,
+        *,
+        synthesized: bool = False,
     ) -> bool: ...
     async def _register_http_bridge_turn_state_core(
         self,
         session: _HTTPBridgeSession,
         turn_state: str,
         *,
         reversible: bool,
+        synthesized: bool,
     ) -> tuple[bool, DurableBridgeAliasRegistrationReceipt | None]: ...
     async def _register_http_bridge_previous_response_id_impl(
         self,