fix(proxy): retry detached API key release

app/modules/proxy/_service/api_key_usage.py

@@ -33,6 +33,9 @@
 logger = logging.getLogger("app.modules.proxy.service")
 
 _API_KEY_RESERVATION_HEARTBEAT_SECONDS = 300.0
+_STREAM_API_KEY_RELEASE_RETRY_BASE_SECONDS = 0.1
+_STREAM_API_KEY_RELEASE_RETRY_MAX_SECONDS = 5.0
+_STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY = 4
 
 
 def _service_api_keys_service() -> type[ApiKeysService]:
@@ -60,6 +63,7 @@ def _api_key_reservation_heartbeat_seconds() -> float:
 class _ApiKeyUsageServiceProtocol(Protocol):
     _repo_factory: ProxyRepoFactory
     _background_cleanup_tasks: set[asyncio.Task[None]]
+    _stream_api_key_release_retry_semaphore: asyncio.Semaphore
 
 
 def _normalize_service_tier_value(value: Any) -> str | None:
@@ -452,6 +456,7 @@ async def _release_after_failed_settlement() -> None:
                 api_key=api_key,
                 api_key_reservation=api_key_reservation,
                 request_id=request_id,
+                retry_persistence_failures=True,
             )
 
         def _settlement_done(done_task: asyncio.Task[bool]) -> None:
@@ -519,21 +524,48 @@ async def _release_unsettled_stream_api_key_usage(
         api_key: ApiKeyData,
         api_key_reservation: ApiKeyUsageReservationData,
         request_id: str,
+        retry_persistence_failures: bool = False,
     ) -> bool:
         proxy = cast(_ApiKeyUsageServiceProtocol, self)
-        with anyio.CancelScope(shield=True):
+        retry_attempt = 1
+        retry_delay_seconds = _STREAM_API_KEY_RELEASE_RETRY_BASE_SECONDS
+        while True:
+            retry_slot_acquired = False
             try:
-                async with proxy._repo_factory() as repos:
-                    api_keys_service = _service_api_keys_service()(repos.api_keys)
-                    await api_keys_service.release_usage_reservation(
-                        api_key_reservation.reservation_id,
-                    )
+                if retry_persistence_failures:
+                    await proxy._stream_api_key_release_retry_semaphore.acquire()
+                    retry_slot_acquired = True
+                with anyio.CancelScope(shield=True):
+                    async with proxy._repo_factory() as repos:
+                        api_keys_service = _service_api_keys_service()(repos.api_keys)
+                        await api_keys_service.release_usage_reservation(
+                            api_key_reservation.reservation_id,
+                        )
                 return True
             except Exception:
+                if not retry_persistence_failures:
+                    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
                 logger.warning(
-                    "Failed to release stream API key reservation key_id=%s request_id=%s",
+                    "Failed to release stream API key reservation key_id=%s request_id=%s "
+                    "retry_attempt=%d retry_delay_seconds=%.2f",
                     api_key.id,
                     request_id,
+                    retry_attempt,
+                    retry_delay_seconds,
                     exc_info=True,
                 )
-                return False
+            finally:
+                if retry_slot_acquired:
+                    proxy._stream_api_key_release_retry_semaphore.release()
+            await asyncio.sleep(retry_delay_seconds)
+            retry_attempt += 1
+            retry_delay_seconds = min(
+                _STREAM_API_KEY_RELEASE_RETRY_MAX_SECONDS,
+                retry_delay_seconds * 2,
+            )

app/modules/proxy/service.py

@@ -122,6 +122,9 @@
 from app.modules.proxy._service.api_key_usage import (
     _API_KEY_RESERVATION_HEARTBEAT_SECONDS as _API_KEY_RESERVATION_HEARTBEAT_SECONDS,
 )
+from app.modules.proxy._service.api_key_usage import (
+    _STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY as _STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY,
+)
 from app.modules.proxy._service.api_key_usage import _ApiKeyUsageMixin
 from app.modules.proxy._service.codex_control import _CodexControlMixin
 from app.modules.proxy._service.compact import _CompactMixin
@@ -940,6 +943,7 @@ def __init__(
         self._websocket_previous_response_account_index: dict[tuple[str, str | None, str | None], str] = {}
         self._websocket_continuity_index: dict[tuple[str, str | None], _WebSocketContinuityState] = {}
         self._background_cleanup_tasks: set[asyncio.Task[None]] = set()
+        self._stream_api_key_release_retry_semaphore = asyncio.Semaphore(_STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY)
         # In-memory pin from upstream-issued file_id -> codex-lb account_id.
         # Used so ``finalize_file`` for a given ``file_id`` is routed to
         # the same account that handled ``create_file``. Cross-instance

openspec/changes/retry-detached-api-key-release/.openspec.yaml

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

openspec/changes/retry-detached-api-key-release/design.md

@@ -0,0 +1,66 @@
+## Context
+
+Stream reservation settlement is detached from the response path. A failed
+settlement schedules one release task in the existing tracked background-task
+set, but that release currently catches its own exception and returns normally.
+The done callback consequently removes the task, so the persistence drain can
+report success while the reservation remains active.
+
+Reservation release is already transactional and idempotent: once another
+settler has changed the reservation from `reserved`, a later release is a
+no-op. The existing stale sweep remains a last-resort repair, but its six-hour
+age threshold is too slow for a known live cleanup chain.
+
+## Goals / Non-Goals
+
+**Goals:**
+
+- Keep transiently failing fallback release work visible to the existing task
+  drain.
+- Retry until the idempotent release succeeds, with bounded retry pressure.
+- Preserve detached response latency, settlement-before-health ordering, and
+  exactly-once accounting.
+
+**Non-Goals:**
+
+- Changing reservation amounts, quota admission, or stale-sweep timing.
+- Adding a durable job queue, setting, migration, or new public API.
+- Refactoring request-log persistence or unrelated cleanup ownership.
+
+## Decisions
+
+1. **Retry inside the already tracked release task.** The release coroutine
+   stays pending between attempts, so the current task registry and recursive
+   drain remain the single source of cleanup ownership. Creating a second
+   registry or a durable retry row would duplicate state for a narrow failure.
+
+2. **Use capped exponential delay plus a shared retry gate for every persistence
+   exception.** The outer retry covers transient PostgreSQL/session failures
+   that the API-key service's SQLite-lock-specific retry does not classify. A
+   fixed delay cap prevents each task from retrying rapidly, while a per-service
+   concurrency gate prevents many failed streams from opening repository
+   sessions simultaneously. Waiting tasks stay tracked without holding a
+   database connection.
+
+3. **Rely on reservation transition idempotency.** A retry cannot double
+   decrement quota: release only claims a reservation still in `reserved`
+   state, and a concurrent finalizer or release makes subsequent attempts
+   no-ops.
+
+4. **Let the existing drain deadline bound shutdown waiting.** A recovered
+   release completes normally. A release still retrying at the deadline remains
+   pending, so `drain_persistence_tasks` returns `False` instead of claiming
+   durability. No separate retry-count terminal state is introduced.
+
+## Risks / Trade-offs
+
+- **A permanent persistence error leaves a task alive during normal runtime.**
+  → Retries use capped backoff; the task accurately represents unfinished
+  cleanup, and stale recovery remains the final repair path.
+- **Many simultaneous failures could retry together after an outage.**
+  → A shared four-attempt gate bounds aggregate repository pressure in each
+  service instance; exponential delay also bounds each task's retry frequency,
+  and the change adds no inline request-path work.
+- **Cancellation can stop a retry after shutdown has already timed out.**
+  → The drain first reports incomplete, so process termination cannot be
+  mistaken for successful settlement.

openspec/changes/retry-detached-api-key-release/proposal.md

@@ -0,0 +1,34 @@
+## Why
+
+A detached stream settlement can fail, enqueue its reservation-release fallback,
+and then lose the reservation when that fallback also hits a transient
+persistence failure. The task drain reports success even though the reservation
+still consumes quota until stale recovery runs hours later.
+
+## What Changes
+
+- Keep a failed detached reservation release tracked and retry it after
+  transient persistence failures, with a shared concurrency bound on repository
+  attempts.
+- Make the persistence drain report completion only after the tracked
+  settlement/release chain has actually terminated.
+- Add deterministic regression coverage for a finalize failure followed by one
+  failed release attempt, while preserving successful settlement, cancellation,
+  and SQLite-lock behavior.
+
+## Capabilities
+
+### New Capabilities
+
+None.
+
+### Modified Capabilities
+
+- `api-keys`: Clarify that a detached settlement fallback which itself fails
+  transiently remains tracked and retries before persistence drain can succeed.
+
+## Impact
+
+The change is limited to detached API-key reservation cleanup in the proxy
+service, its focused persistence tests, and the existing API-key settlement
+contract. It adds no API, setting, dependency, migration, or dashboard change.

openspec/changes/retry-detached-api-key-release/specs/api-keys/spec.md

@@ -0,0 +1,45 @@
+## MODIFIED Requirements
+
+### Requirement: Stream reservation settlement is detached from the response path
+
+Settling a stream API-key reservation MUST NOT block the response/stream close, with one deliberate exception: when a keyed websocket stream terminates with an account-health error, the finalizer MUST wait for the settlement to commit before the load-balancer health write (the settlement-ordering invariant), so that error path intentionally blocks on settlement. In all other cases the settlement MUST run as a tracked background task; when it fails or is cancelled, the reservation MUST still be released by the tracking fallback, and the request's finalization path MUST NOT double-release a transferred settlement. If the tracking fallback itself encounters a persistence failure, it MUST remain tracked and retry the idempotent release; no more than four retry-enabled detached fallback repository attempts may run concurrently per proxy service instance, waiting fallbacks MUST NOT open repository sessions until admitted, and persistence drain MUST NOT report completion while any retry remains unfinished. Reservations MUST continue to count toward key limits until finalized or released, so deferred settlement can never admit usage a synchronous settlement would have rejected.
+
+#### Scenario: Response close precedes settlement completion
+
+- **GIVEN** a keyed stream whose settlement transaction is still running
+- **WHEN** the stream closes
+- **THEN** the close does not wait for the settlement
+- **AND** the settlement finalizes the reservation exactly once in the background
+
+#### Scenario: Failed detached settlement still releases the reservation
+
+- **GIVEN** a detached settlement whose finalize raises
+- **WHEN** the settlement task completes
+- **THEN** the tracking fallback releases the reservation
+
+#### Scenario: Failed fallback release remains tracked
+
+- **GIVEN** a detached settlement whose finalize raises
+- **AND** the first tracking-fallback release attempt also raises
+- **WHEN** persistence recovers before the drain deadline
+- **THEN** the tracked fallback retries and releases the reservation exactly once
+- **AND** persistence drain does not report completion before that release
+
+#### Scenario: Concurrent fallback release retries are bounded to four
+
+- **GIVEN** five failed detached settlements in one proxy service instance
+- **WHEN** their tracking fallbacks attempt repository persistence concurrently
+- **THEN** no more than four release attempts open repository sessions
+- **AND** the waiting fallbacks remain tracked until they can retry
+
+#### Scenario: Websocket health-error settlement precedes the health write
+
+- **GIVEN** a keyed websocket stream that terminates with an account-health error
+- **WHEN** the finalizer settles the reservation
+- **THEN** it waits for the settlement to commit before recording the account-health error
+
+#### Scenario: Shutdown drains pending settlements
+
+- **WHEN** the service shuts down gracefully with settlements in flight
+- **THEN** shutdown waits for them up to the configured drain timeout
+- **AND** reports an incomplete drain if a tracked settlement or release remains unfinished at that timeout

openspec/changes/retry-detached-api-key-release/tasks.md

@@ -0,0 +1,17 @@
+## 1. Regression
+
+- [x] 1.1 Add a real-repository regression that injects one finalize failure and one fallback-release failure.
+- [x] 1.2 Confirm the regression fails deterministically twice on baseline `3fe0d6f286019a0505783d803db9a1d8cdf6b307`.
+
+## 2. Implementation
+
+- [x] 2.1 Keep the fallback release tracked while retrying persistence failures with capped backoff.
+- [x] 2.2 Preserve idempotent settlement, cancellation ownership, and truthful persistence-drain behavior.
+- [x] 2.3 Bound concurrent fallback repository attempts to four with one shared per-service gate.
+
+## 3. Verification
+
+- [x] 3.1 Run the focused detached-settlement and API-key reservation tests.
+- [x] 3.2 Run changed-file Ruff, format, type, proxy-architecture, and strict OpenSpec checks.
+- [x] 3.3 Inspect the final diff and worktree status for scope and unrelated changes.
+- [x] 3.4 Add deterministic fan-out coverage for the shared retry concurrency bound.

tests/integration/test_detached_persistence.py

@@ -3,9 +3,17 @@
 import pytest
 from httpx import ASGITransport, AsyncClient
 from sqlalchemy import select
+from sqlalchemy.exc import OperationalError
 
 from app.db.models import RequestLog
 from app.db.session import SessionLocal
+from app.modules.api_keys.repository import ApiKeysRepository, UsageReservationData
+from app.modules.api_keys.service import (
+    ApiKeyCreateData,
+    ApiKeyRequestUsageBudget,
+    ApiKeysService,
+    LimitRuleInput,
+)
 from app.modules.proxy import service as proxy_service_module
 
 pytestmark = pytest.mark.integration
@@ -105,6 +113,105 @@ async def never_finishes() -> None:
         service._request_log_tasks.discard(task)
 
 
+@pytest.mark.asyncio
+async def test_failed_detached_settlement_retries_failed_release_until_persisted(raw_client, monkeypatch):
+    import asyncio
+
+    _, app = raw_client
+
+    async with SessionLocal() as session:
+        api_keys = ApiKeysService(ApiKeysRepository(session))
+        created = await api_keys.create_key(
+            ApiKeyCreateData(
+                name="detached-release-retry",
+                allowed_models=None,
+                expires_at=None,
+                limits=[
+                    LimitRuleInput(
+                        limit_type="total_tokens",
+                        limit_window="weekly",
+                        max_value=100,
+                    )
+                ],
+            )
+        )
+        api_key = await api_keys.get_key_by_id(created.id)
+        reservation = await api_keys.enforce_limits_for_request(
+            created.id,
+            request_model="gpt-5.5",
+            request_usage_budget=ApiKeyRequestUsageBudget(
+                input_tokens=4,
+                output_tokens=6,
+            ),
+        )
+
+    original_get_reservation = ApiKeysRepository.get_usage_reservation
+    reservation_read_attempts = 0
+    retry_started = asyncio.Event()
+    allow_retry = asyncio.Event()
+
+    async def fail_first_two_reservation_reads(
+        self: ApiKeysRepository,
+        reservation_id: str,
+    ) -> UsageReservationData | None:
+        nonlocal reservation_read_attempts
+        if reservation_id == reservation.reservation_id:
+            reservation_read_attempts += 1
+            if reservation_read_attempts <= 2:
+                raise OperationalError(
+                    "read usage reservation",
+                    {},
+                    Exception("transient persistence connection failure"),
+                )
+            if reservation_read_attempts == 3:
+                retry_started.set()
+                await allow_retry.wait()
+        return await original_get_reservation(self, reservation_id)
+
+    monkeypatch.setattr(ApiKeysRepository, "get_usage_reservation", fail_first_two_reservation_reads)
+
+    settlement = proxy_service_module._StreamSettlement(
+        status="success",
+        model="gpt-5.5",
+        input_tokens=4,
+        output_tokens=6,
+    )
+    from app.dependencies import get_proxy_service_for_app
+
+    service = get_proxy_service_for_app(app)
+    assert await service._settle_stream_api_key_usage(
+        api_key,
+        reservation,
+        settlement,
+        request_id="req_detached_release_retry",
+    )
+    drain_task = asyncio.create_task(service.drain_persistence_tasks(timeout_seconds=2))
+    retry_wait_task = asyncio.create_task(retry_started.wait())
+    done, _ = await asyncio.wait(
+        {retry_wait_task, drain_task},
+        timeout=1,
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+    retry_was_tracked = retry_wait_task in done and not drain_task.done()
+    allow_retry.set()
+    if not retry_wait_task.done():
+        retry_wait_task.cancel()
+    await asyncio.gather(retry_wait_task, return_exceptions=True)
+    assert await drain_task
+
+    async with SessionLocal() as session:
+        repo = ApiKeysRepository(session)
+        stored = await original_get_reservation(repo, reservation.reservation_id)
+        limits = await repo.get_limits_by_key(created.id)
+
+    assert stored is not None
+    assert stored.status == "released"
+    assert len(limits) == 1
+    assert limits[0].current_value == 0
+    assert reservation_read_attempts == 3
+    assert retry_was_tracked is True
+
+
 @pytest.mark.asyncio
 async def test_drain_ignores_stuck_non_persistence_cleanup_tasks():
     """A stuck bridge-close cleanup in _background_cleanup_tasks must not