fix(proxy): sequence websocket health after settlement

app/modules/proxy/_service/api_key_usage.py

@@ -314,6 +314,26 @@ async def _settle_stream_api_key_usage(
         model_name = api_key_reservation.model or settlement.model or ""
         proxy = cast(_ApiKeyUsageServiceProtocol, self)
 
+        async def _release_ordering_sensitive_fallback() -> bool:
+            fallback_task = asyncio.create_task(
+                self._release_unsettled_stream_api_key_usage(
+                    api_key=api_key,
+                    api_key_reservation=api_key_reservation,
+                    request_id=request_id,
+                ),
+                name=f"proxy-stream-api-key-fallback-{request_id}",
+            )
+            cancellation_pending = False
+            while not fallback_task.done():
+                try:
+                    await asyncio.shield(fallback_task)
+                except asyncio.CancelledError:
+                    cancellation_pending = True
+            settled = fallback_task.result()
+            if cancellation_pending:
+                return False
+            return settled
+
         async def _settle_once() -> bool:
             try:
                 async with proxy._repo_factory() as repos:
@@ -334,18 +354,28 @@ async def _settle_once() -> bool:
                     else:
                         await api_keys_service.release_usage_reservation(reservation_id)
                 return True
+            except asyncio.CancelledError:
+                if wait_for_settlement:
+                    await _release_ordering_sensitive_fallback()
+                    return False
+                raise
             except Exception:
                 logger.warning(
                     "Failed to settle stream API key reservation key_id=%s request_id=%s",
                     api_key.id,
                     request_id,
                     exc_info=True,
                 )
+                if wait_for_settlement:
+                    return await _release_ordering_sensitive_fallback()
                 return False
 
-        # Detach unconditionally instead of shield-awaiting: the tracking
-        # callback already schedules a release when settlement fails or is
-        # cancelled, the caller's finally-net skips via
+        # Detach unconditionally instead of shield-awaiting: for ordinary
+        # callers the tracking callback schedules a release when settlement
+        # fails or is cancelled; an ordering-sensitive settlement task runs
+        # that fallback before the tracked task completes once started, while
+        # the tracker still owns cancellation before coroutine startup. The
+        # caller's finally-net skips via
         # usage_settlement_transferred, and reservations keep counting toward
         # limits until finalized/released, so a briefly-lagging settlement can
         # only over-restrict, never over-admit. Awaiting the ~5+2N-statement
@@ -358,16 +388,26 @@ async def _settle_once() -> bool:
             api_key=api_key,
             api_key_reservation=api_key_reservation,
             request_id=request_id,
+            release_on_failure=not wait_for_settlement,
         )
         if wait_for_settlement:
             # Ordering-sensitive callers (the websocket error path) must
             # commit the settlement before load-balancer health writes; they
             # opt into waiting while everything else stays detached.
+            settlement_committed = False
             with anyio.CancelScope(shield=True):
-                try:
-                    await asyncio.shield(task)
-                except Exception:  # failures release via the tracking callback
-                    pass
+                while True:
+                    try:
+                        settlement_committed = await asyncio.shield(task)
+                        break
+                    except asyncio.CancelledError:
+                        # Caller cancellation must not race fallback release
+                        # against the still-running settlement transaction.
+                        if task.cancelled():
+                            break
+                    except Exception:
+                        break
+            return settlement_committed
         return True
 
     def _track_stream_usage_settlement_task(
@@ -377,10 +417,18 @@ def _track_stream_usage_settlement_task(
         api_key: ApiKeyData,
         api_key_reservation: ApiKeyUsageReservationData,
         request_id: str,
+        release_on_failure: bool = True,
     ) -> None:
         proxy = cast(_ApiKeyUsageServiceProtocol, self)
         proxy._background_cleanup_tasks.add(cast(asyncio.Task[None], task))
 
+        async def _release_after_failed_settlement() -> None:
+            await self._release_unsettled_stream_api_key_usage(
+                api_key=api_key,
+                api_key_reservation=api_key_reservation,
+                request_id=request_id,
+            )
+
         def _settlement_done(done_task: asyncio.Task[bool]) -> None:
             proxy._background_cleanup_tasks.discard(cast(asyncio.Task[None], done_task))
             try:
@@ -391,13 +439,8 @@ def _settlement_done(done_task: asyncio.Task[bool]) -> None:
                     api_key.id,
                     request_id,
                 )
-                release_coro = self._release_unsettled_stream_api_key_usage(
-                    api_key=api_key,
-                    api_key_reservation=api_key_reservation,
-                    request_id=request_id,
-                )
                 self._schedule_cancel_safe_cleanup(
-                    release_coro,
+                    _release_after_failed_settlement(),
                     action="release_stream_api_key_reservation_after_cancelled_settlement",
                     request_id=request_id,
                 )
@@ -409,14 +452,9 @@ def _settlement_done(done_task: asyncio.Task[bool]) -> None:
                     exc_info=(type(exc), exc, exc.__traceback__),
                 )
             else:
-                if not settled:
-                    release_coro = self._release_unsettled_stream_api_key_usage(
-                        api_key=api_key,
-                        api_key_reservation=api_key_reservation,
-                        request_id=request_id,
-                    )
+                if not settled and release_on_failure:
                     self._schedule_cancel_safe_cleanup(
-                        release_coro,
+                        _release_after_failed_settlement(),
                         action="release_stream_api_key_reservation_after_failed_settlement",
                         request_id=request_id,
                     )
@@ -456,7 +494,7 @@ async def _release_unsettled_stream_api_key_usage(
         api_key: ApiKeyData,
         api_key_reservation: ApiKeyUsageReservationData,
         request_id: str,
-    ) -> None:
+    ) -> bool:
         proxy = cast(_ApiKeyUsageServiceProtocol, self)
         with anyio.CancelScope(shield=True):
             try:
@@ -465,10 +503,12 @@ async def _release_unsettled_stream_api_key_usage(
                     await api_keys_service.release_usage_reservation(
                         api_key_reservation.reservation_id,
                     )
+                return True
             except Exception:
                 logger.warning(
                     "Failed to release stream API key reservation key_id=%s request_id=%s",
                     api_key.id,
                     request_id,
                     exc_info=True,
                 )
+                return False

app/modules/proxy/_service/streaming/retry.py

@@ -405,6 +405,8 @@ async def _settle_stream_usage_before_pending_penalty(
                 )
                 pending_penalties = list(pending_post_refresh_transient_penalties)
                 pending_post_refresh_transient_penalties.clear()
+                if not settled_result:
+                    return False
                 for pending_penalty in pending_penalties:
                     (
                         failed_account,
@@ -421,7 +423,7 @@ async def _settle_stream_usage_before_pending_penalty(
                     )
                     if transient_retry_count > 1:
                         await proxy._load_balancer.record_errors(failed_account, transient_retry_count - 1)
-                return settled_result
+                return True
             return await proxy._settle_stream_api_key_usage(
                 api_key,
                 api_key_reservation,
@@ -431,14 +433,16 @@ async def _settle_stream_usage_before_pending_penalty(
 
         async def _drain_pending_post_refresh_penalty_on_terminal(
             current_settlement: _StreamSettlement,
-        ) -> None:
+        ) -> bool:
             nonlocal post_refresh_transient_replacement_selected, settled
             if pending_post_refresh_transient_penalties:
                 # A failed replacement selection still ends the request. Mark
                 # it as terminal so the deferred failure is settled and
                 # recorded before this path returns or re-raises.
                 post_refresh_transient_replacement_selected = True
                 settled = await _settle_stream_usage_before_pending_penalty(current_settlement)
+                return settled
+            return True
 
         async def _wait_for_process_network_recovery(
             account: Account,
@@ -1771,7 +1775,7 @@ async def _retry_account_model_rejection(
                                     settlement.error = tex.error
                                 settlement.account_health_error = _facade()._should_penalize_stream_error(error_code)
                                 settled = await _settle_stream_usage_before_pending_penalty(settlement)
-                                if settlement.account_health_error:
+                                if settled and settlement.account_health_error:
                                     await proxy._handle_stream_error(
                                         account,
                                         _stream_settlement_error_payload(settlement),
@@ -1991,13 +1995,13 @@ async def _retry_account_model_rejection(
                         finally:
                             pop_stream_timeout_overrides(stream_timeout_tokens)
                         settled = await _settle_stream_usage_before_pending_penalty(settlement)
-                        if settlement.account_health_error:
+                        if settled and settlement.account_health_error:
                             await proxy._handle_stream_error(
                                 account,
                                 _stream_settlement_error_payload(settlement),
                                 settlement.error_code or "upstream_error",
                             )
-                        elif settlement.record_success:
+                        elif settled and settlement.record_success:
                             await proxy._load_balancer.record_success(account)
                         network_recovery.log_recovered()
                         upstream_transport_metric_status = settlement.status
@@ -2051,8 +2055,8 @@ async def _retry_account_model_rejection(
                     )
                     continue
                 except _TerminalStreamError as exc:
-                    await _drain_pending_post_refresh_penalty_on_terminal(settlement)
-                    if _facade()._should_penalize_stream_error(exc.code):
+                    health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement)
+                    if health_write_allowed and _facade()._should_penalize_stream_error(exc.code):
                         await proxy._handle_stream_error(account, exc.error, exc.code)
                     return
                 except ProxyResponseError as exc:
@@ -2349,7 +2353,7 @@ async def _retry_account_model_rejection(
                                 settlement.error = _upstream_error_from_openai(error)
                                 settlement.account_health_error = _facade()._should_penalize_stream_error(error_code)
                                 settled = await _settle_stream_usage_before_pending_penalty(settlement)
-                                if settlement.account_health_error:
+                                if settled and settlement.account_health_error:
                                     await proxy._handle_stream_error(
                                         account,
                                         _stream_settlement_error_payload(settlement),
@@ -2451,13 +2455,14 @@ async def _retry_account_model_rejection(
                                 )
                                 excluded_account_ids.add(account.id)
                                 continue
-                            await _drain_pending_post_refresh_penalty_on_terminal(settlement)
-                            await proxy._handle_stream_error(
-                                account,
-                                current_error_payload,
-                                current_error_code,
-                                http_status=retry_exc.status_code,
-                            )
+                            health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement)
+                            if health_write_allowed:
+                                await proxy._handle_stream_error(
+                                    account,
+                                    current_error_payload,
+                                    current_error_code,
+                                    http_status=retry_exc.status_code,
+                                )
                             if propagate_http_errors:
                                 raise
                             error_message = error.message if error else None
@@ -2475,17 +2480,20 @@ async def _retry_account_model_rejection(
                             failed_account is account
                             for failed_account, *_rest in pending_post_refresh_transient_penalties
                         )
-                        if pending_post_refresh_transient_penalties:
-                            await _drain_pending_post_refresh_penalty_on_terminal(settlement)
-                        if settlement.account_health_error and not current_account_penalty_queued:
+                        health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement)
+                        if (
+                            health_write_allowed
+                            and settlement.account_health_error
+                            and not current_account_penalty_queued
+                        ):
                             await proxy._handle_stream_error(
                                 account,
                                 _stream_settlement_error_payload(settlement),
                                 settlement.error_code or "upstream_error",
                             )
-                        elif settlement.record_success:
+                        elif health_write_allowed and settlement.record_success:
                             await proxy._load_balancer.record_success(account)
-                        if not settled:
+                        if not settled and not settlement.usage_settlement_transferred:
                             settled = await _settle_stream_usage_before_pending_penalty(settlement)
                         upstream_transport_metric_status = settlement.status
                         _record_upstream_transport_metric_once(settlement.status)
@@ -2521,8 +2529,8 @@ async def _retry_account_model_rejection(
                             excluded_account_ids.add(account.id)
                             require_security_work_authorized = True
                             continue
-                    await _drain_pending_post_refresh_penalty_on_terminal(settlement)
-                    if _facade()._should_penalize_stream_error(error_code):
+                    health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement)
+                    if health_write_allowed and _facade()._should_penalize_stream_error(error_code):
                         await proxy._handle_stream_error(
                             account,
                             _upstream_error_from_openai(error),

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

@@ -4432,7 +4432,12 @@ async def _finalize_websocket_request_state(
             settlement.account_health_error = False
         proxy._cancel_request_state_api_key_reservation_heartbeat(request_state)
         await _release_websocket_response_create_gate(request_state, response_create_gate)
-        await proxy._settle_stream_api_key_usage(
+        if settlement.account_health_error:
+            # Connection safety must not wait on settlement or health
+            # persistence. The health write remains ordered below.
+            upstream_control.reconnect_requested = True
+            upstream_control.retire_after_drain = True
+        settlement_confirmed = await proxy._settle_stream_api_key_usage(
             api_key,
             request_state.api_key_reservation,
             settlement,
@@ -4442,13 +4447,12 @@ async def _finalize_websocket_request_state(
             wait_for_settlement=settlement.account_health_error,
         )
         if settlement.account_health_error:
-            await proxy._handle_stream_error(
-                account,
-                _stream_settlement_error_payload(settlement),
-                settlement.error_code or "upstream_error",
-            )
-            upstream_control.reconnect_requested = True
-            upstream_control.retire_after_drain = True
+            if settlement_confirmed:
+                await proxy._handle_stream_error(
+                    account,
+                    _stream_settlement_error_payload(settlement),
+                    settlement.error_code or "upstream_error",
+                )
         elif settlement.record_success:
             await proxy._load_balancer.record_success(account)
             for remembered_response_id in _websocket_continuity_response_ids(request_state, response_id):

openspec/changes/sequence-websocket-health-after-settlement/.openspec.yaml

@@ -0,0 +1,2 @@
+schema: spec-driven
+created: 2026-07-31