fix(quota): release reservations on header failure

Makefile

@@ -12,6 +12,7 @@ POSTGRES_PYTEST_TARGETS := \
 	tests/integration/test_repositories.py::test_accounts_upsert_with_merge_enabled_serializes_concurrent_same_email \
 	tests/integration/test_sticky_sessions_api.py::test_durable_bridge_owned_alias_registration_is_epoch_fenced \
 	tests/integration/test_proxy_api_extended.py::test_proxy_stream_usage_limit_returns_http_error \
+	tests/integration/test_api_keys_api.py::test_rate_limit_header_failure_releases_reservation_once \
 	tests/integration/test_codex_usage_api.py::test_codex_usage_aggregates_windows \
 	tests/integration/test_proxy_compact.py::test_proxy_compact_headers_include_monthly_only_credits \
 	tests/integration/test_repositories.py::test_accounts_upsert_with_merge_disabled_uses_identity_lock_on_postgresql \

app/modules/proxy/api.py

@@ -12,6 +12,7 @@
 from typing import Any, Final, Literal, Protocol, cast
 from uuid import uuid4
 
+import anyio
 from fastapi import (
     APIRouter,
     Body,
@@ -1958,6 +1959,39 @@ async def _rate_limit_headers_for_request(
     return await context.service.rate_limit_headers()
 
 
+async def _release_reservation_deferring_cancellation(
+    reservation: ApiKeyUsageReservationData,
+) -> None:
+    with anyio.CancelScope(shield=True):
+        task = asyncio.create_task(_release_reservation(reservation))
+        while True:
+            try:
+                await asyncio.shield(task)
+                return
+            except asyncio.CancelledError:
+                if task.cancelled():
+                    raise
+
+
+async def _rate_limit_headers_with_reservation_cleanup(
+    context: ProxyContext,
+    api_key: ApiKeyData | None,
+    owned_reservation: ApiKeyUsageReservationData | None,
+) -> dict[str, str]:
+    try:
+        return await _rate_limit_headers_for_request(context, api_key)
+    except BaseException:
+        if owned_reservation is not None:
+            try:
+                await _release_reservation_deferring_cancellation(owned_reservation)
+            except (Exception, asyncio.CancelledError):
+                logger.warning(
+                    "Failed to release API key reservation after rate-limit header failure",
+                    exc_info=True,
+                )
+        raise
+
+
 def _select_codex_usage_limit(
     limits: list[V1UsageLimitResponse],
     window: str,
@@ -4875,7 +4909,15 @@ async def _stream_responses(
         )
     )
 
-    rate_limit_headers = await _rate_limit_headers_for_request(context, api_key) if include_rate_limit_headers else {}
+    rate_limit_headers = (
+        await _rate_limit_headers_with_reservation_cleanup(
+            context,
+            api_key,
+            reservation if owns_reservation else None,
+        )
+        if include_rate_limit_headers
+        else {}
+    )
     bridge_active = prefer_http_bridge and proxy_service_module.get_settings().http_responses_session_bridge_enabled
     effective_headers = forwarded_headers or request.headers
     client_ip = forwarded_client_ip if forwarded_request else resolve_request_client_host(request)
@@ -5080,7 +5122,7 @@ async def _collect_responses(
         request_usage_budget=estimate_api_key_request_usage(payload),
     )
 
-    rate_limit_headers = await _rate_limit_headers_for_request(context, api_key)
+    rate_limit_headers = await _rate_limit_headers_with_reservation_cleanup(context, api_key, reservation)
     bridge_active = prefer_http_bridge and proxy_service_module.get_settings().http_responses_session_bridge_enabled
     downstream_turn_state = (
         proxy_affinity_module.ensure_http_downstream_turn_state(request.headers) if bridge_active else None
@@ -5240,7 +5282,7 @@ async def _compact_responses(
         request_usage_budget=request_usage_budget,
     )
 
-    rate_limit_headers = await _rate_limit_headers_for_request(context, api_key)
+    rate_limit_headers = await _rate_limit_headers_with_reservation_cleanup(context, api_key, reservation)
     try:
         result = await context.service.compact_responses(
             payload,
@@ -5399,7 +5441,7 @@ async def _transcribe_request(
         request_model=_TRANSCRIPTION_MODEL,
         request_service_tier=None,
     )
-    rate_limit_headers = await _rate_limit_headers_for_request(context, api_key)
+    rate_limit_headers = await _rate_limit_headers_with_reservation_cleanup(context, api_key, reservation)
     try:
         result = await context.service.transcribe(
             audio_bytes=multipart.audio_bytes,

openspec/changes/release-quota-reservations-on-header-failure/.openspec.yaml

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

openspec/changes/release-quota-reservations-on-header-failure/context.md

@@ -0,0 +1,30 @@
+# Context: quota cleanup during response-header preparation
+
+## Purpose and scope
+
+This change closes the short ownership gap between API-key quota admission and
+route-specific settlement ownership. Normative behavior lives in
+[`specs/api-keys/spec.md`](./specs/api-keys/spec.md).
+
+## Decisions and constraints
+
+Header calculation remains after admission so successful responses continue to
+reflect the committed reservation. The route keeps cleanup ownership only
+until headers are ready; stream or service settlement then continues unchanged.
+Borrowed reservations remain owned by their origin.
+
+## Failure modes
+
+A database, cache, or calculation error while building rate-limit headers can
+occur after quota has been reserved but before upstream work begins. The owned
+reservation must be released once before that error propagates. A separate
+failure of the release persistence itself is logged without replacing the
+header error and continues to use the repository's existing stale-recovery
+contract.
+
+## Concrete example
+
+For a limited `POST /v1/responses` request with `stream: false`, admission first
+commits reservation `R`. If rate-limit header construction then raises, the
+route releases `R` exactly once, starts no upstream stream, and preserves the
+header failure for normal error handling.

openspec/changes/release-quota-reservations-on-header-failure/design.md

@@ -0,0 +1,88 @@
+## Context
+
+`ApiKeysService.enforce_limits_for_request()` persists a `reserved` usage row
+and commits its limit deltas before returning. The subscription-backed stream,
+collect, compact, and transcription paths then calculate upstream rate-limit
+response headers before their existing stream or `try`/`finally` settlement
+owner is installed. An exception from that header calculation therefore
+escapes with no component responsible for releasing the committed reservation.
+
+Source-routed Responses and transcription requests, image requests, and chat
+completions calculate headers before reservation or use a different ownership
+sequence, so they do not share this gap.
+
+## Goals / Non-Goals
+
+**Goals:**
+
+- Keep cleanup ownership from reservation commit through rate-limit header
+  preparation for stream, collect, compact, and transcription requests.
+- Release each owned reservation exactly once when header preparation fails,
+  then re-raise the original failure without starting upstream work.
+- Prove the behavior at the real HTTP route and persistence seam for all four
+  request shapes.
+- Preserve successful header values and all existing downstream settlement
+  ownership.
+
+**Non-Goals:**
+
+- Making rate-limit header failures best-effort or returning a fallback header
+  set.
+- Changing rate-limit queries, caching, serialization, or response schemas.
+- Adding another retry or detached-settlement mechanism; persistence failure
+  during release remains governed by its existing recovery contracts.
+- Changing source, image, chat, WebSocket, or borrowed/forwarded reservation
+  ownership.
+
+## Decisions
+
+### Centralize the narrow ownership handoff
+
+Add one route helper that calculates rate-limit headers while the caller still
+owns the reservation. If calculation exits unsuccessfully, the helper releases
+an owned reservation and re-raises. Once headers return successfully, the
+existing route-specific stream or `try`/`finally` logic retains responsibility.
+
+Duplicating a `try`/`except` block at each call site was rejected because the
+four paths share the same transition and a later route could easily omit one
+half of the invariant. Expanding each route's downstream finalizer around all
+setup was rejected because streaming paths deliberately transfer reservation
+ownership and must not release it when returning a live stream.
+
+### Keep header calculation after reservation admission
+
+The calculation stays after `enforce_limits_for_request()`. Moving it before
+admission would avoid the leak, but could return quota metadata that does not
+reflect the request's newly committed reservation and would silently change
+successful response semantics.
+
+### Exercise the real commit and release path once per transport shape
+
+Use one parameterized ASGI regression covering streaming Responses, collected
+Responses, compact Responses, and audio transcription. Each case creates a
+limited API key, injects a header-calculation exception after real reservation
+admission, wraps the production release helper, and asserts one release call,
+one released reservation row, and restored limit usage.
+
+A helper-only unit test as the sole proof was rejected because it would not
+prove that every route calls the helper after reservation commit and before
+upstream work.
+
+## Risks / Trade-offs
+
+- **A route bypasses the helper later** → Keep all four cases in one
+  parameterized route-level regression.
+- **Cleanup is accidentally duplicated by a downstream owner** → Inject the
+  failure before downstream construction and require exactly one release call.
+- **Release persistence fails independently** → Log that cleanup failure,
+  preserve the original header failure, and rely on existing stale recovery
+  rather than broadening this fix into a second settlement-retry mechanism.
+
+## Migration Plan
+
+This is a code-only ownership repair with no migration or setting. Deploy it
+through the normal release train; rollback is a code revert.
+
+## Open Questions
+
+None.

openspec/changes/release-quota-reservations-on-header-failure/proposal.md

@@ -0,0 +1,38 @@
+## Why
+
+API-key quota reservation is committed before upstream rate-limit response
+headers are calculated. If that calculation fails, four subscription-backed
+request paths currently propagate the error before any downstream component
+owns cleanup, leaving quota reserved until stale recovery runs.
+
+## What Changes
+
+- Retain reservation cleanup ownership while rate-limit headers are prepared
+  for streaming Responses, collected Responses, Responses compaction, and
+  audio transcription requests.
+- Release an owned reservation exactly once when header preparation fails,
+  then preserve the original failure instead of starting upstream work.
+- Add one parameterized, route-level failure-injection regression that proves
+  all four request shapes restore quota and perform exactly one release.
+- Preserve successful header construction, downstream settlement ownership,
+  and borrowed reservation behavior.
+
+## Capabilities
+
+### New Capabilities
+
+None.
+
+### Modified Capabilities
+
+- `api-keys`: Extend the early-exit reservation cleanup contract to failures
+  while rate-limit response headers are calculated after admission and before
+  upstream ownership begins.
+
+## Impact
+
+- Backend: `app/modules/proxy/api.py`
+- Tests: `tests/integration/test_api_keys_api.py`
+- Contract: API-key quota cleanup on internal response-header failure only
+- No API schema, database migration, dependency, setting, dashboard, or
+  successful-response change

openspec/changes/release-quota-reservations-on-header-failure/specs/api-keys/spec.md

@@ -0,0 +1,37 @@
+## MODIFIED Requirements
+
+### Requirement: 조기 종료 경로에서 reservation release 보장
+
+Reservation 생성 후 upstream API 호출에 진입하지 않고 종료되는 모든 경로에서 reservation이 release되어야 한다. `reserved` 상태로 남는 reservation이 존재하면 안 된다. 시스템은 이 동작을 SHALL 보장해야 한다.
+
+After admission commits an owned reservation, rate-limit response-header
+calculation before upstream work remains part of the early-exit cleanup window.
+If that calculation fails, the system MUST attempt to release the owned
+reservation exactly once before propagating the original header failure.
+
+#### Scenario: no_accounts 즉시 종료 시 release
+
+- **WHEN** reservation 생성 후 `_stream_with_retry()`가 사용 가능한 계정 없음(`no_accounts`)으로 즉시 종료되면
+- **THEN** `release_usage_reservation()`이 호출되어 reservation이 `released` 상태로 전이되어야 한다 (SHALL)
+- **AND** pre-reserved quota가 원복되어야 한다 (SHALL)
+
+#### Scenario: 재시도 소진 후 no_accounts 종료 시 release
+
+- **WHEN** 재시도 루프가 모든 attempt를 소진한 후 `no_accounts`로 종료되면
+- **THEN** `release_usage_reservation()`이 호출되어야 한다 (SHALL)
+
+#### Scenario: reservation 미생성 시 정산 스킵
+
+- **WHEN** API key auth가 비활성이거나 reservation이 생성되지 않은 상태에서 요청이 종료되면
+- **THEN** 정산 로직이 안전하게 스킵되어야 하며 에러가 발생하지 않아야 한다 (SHALL)
+
+#### Scenario: Rate-limit header preparation fails after admission
+
+- **GIVEN** a limited API key has committed an owned reservation for a
+  streaming Responses, collected Responses, compact Responses, or audio
+  transcription request
+- **WHEN** rate-limit response-header calculation fails before upstream work
+  begins
+- **THEN** the reservation is released exactly once
+- **AND** its reserved quota is restored
+- **AND** the header failure propagates without starting upstream work

openspec/changes/release-quota-reservations-on-header-failure/tasks.md

@@ -0,0 +1,35 @@
+## 1. Regression Coverage
+
+- [x] 1.1 Add focused unit coverage that proves header failure or cancellation
+  releases an owned reservation once across a cancellation checkpoint,
+  preserves the original failure if release persistence fails, and does not
+  release a borrowed reservation.
+- [x] 1.2 Add one parameterized ASGI regression for stream, collect, compact,
+  and transcribe that injects header failure after real quota admission and
+  requires exactly one release, a `released` row, restored quota, and no
+  upstream call.
+- [x] 1.3 Run the new regressions against the original implementation and
+  confirm they fail at the missing cleanup handoff.
+
+## 2. Reservation Ownership Fix
+
+- [x] 2.1 Add one narrow helper that retains owned-reservation cleanup through
+  rate-limit header preparation, releases on any unsuccessful exit, and
+  re-raises the original failure.
+- [x] 2.2 Route streaming Responses, collected Responses, compact Responses,
+  and subscription-backed transcription through the helper without changing
+  source, image, chat, forwarded, or downstream settlement behavior.
+- [x] 2.3 Include the route regression in the focused PostgreSQL test target so
+  the commit-to-separate-session-release path runs on asyncpg in required CI.
+
+## 3. Verification
+
+- [x] 3.1 Run the new unit and ASGI regressions plus focused existing API-key
+  reservation, forwarded-owner, compact, and transcription lifecycle tests.
+- [x] 3.2 Run the applicable PostgreSQL proof when locally available, affected
+  Ruff/format/type checks, and the proxy architecture checker.
+  - Ruff, format, scoped `ty`, and proxy architecture checks passed. No local
+    PostgreSQL listener was available on TCP 5432; the regression is included
+    in the required PostgreSQL CI target instead.
+- [x] 3.3 Run strict scoped OpenSpec validation, all main-spec validation,
+  OpenSpec verification, final diff review, and worktree-status inspection.