diff --git a/app/db/alembic/versions/20260728_000000_add_file_account_pins.py b/app/db/alembic/versions/20260728_000000_add_file_account_pins.py new file mode 100644 index 0000000000..a12c752f11 --- /dev/null +++ b/app/db/alembic/versions/20260728_000000_add_file_account_pins.py @@ -0,0 +1,40 @@ +"""add durable file account pins + +Revision ID: 20260728_000000_add_file_account_pins +Revises: 20260725_000000_add_http_bridge_pending_tool_calls +Create Date: 2026-07-28 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "20260728_000000_add_file_account_pins" +down_revision = "20260725_000000_add_http_bridge_pending_tool_calls" +branch_labels = None +depends_on = None + +_TABLE = "file_account_pins" + + +def upgrade() -> None: + bind = op.get_bind() + if sa.inspect(bind).has_table(_TABLE): + return + op.create_table( + _TABLE, + sa.Column("file_id", sa.String(), nullable=False), + sa.Column("account_id", sa.String(), nullable=False), + sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("file_id"), + ) + op.create_index("ix_file_account_pins_expires_at", _TABLE, ["expires_at"], unique=False) + + +def downgrade() -> None: + bind = op.get_bind() + if not sa.inspect(bind).has_table(_TABLE): + return + op.drop_index("ix_file_account_pins_expires_at", table_name=_TABLE) + op.drop_table(_TABLE) diff --git a/app/db/models.py b/app/db/models.py index 6f0cb0df9b..fca0e14b7c 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -61,6 +61,16 @@ class StickySessionKind(str, Enum): PROMPT_CACHE = "prompt_cache" +class FileAccountPin(Base): + __tablename__ = "file_account_pins" + + file_id: Mapped[str] = mapped_column(String, primary_key=True) + account_id: Mapped[str] = mapped_column(String, nullable=False) + expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) + + __table_args__ = (Index("ix_file_account_pins_expires_at", "expires_at"),) + + class RequestKind(str, Enum): NORMAL = "normal" WARMUP = "warmup" diff --git a/app/modules/proxy/_service/api_key_usage.py b/app/modules/proxy/_service/api_key_usage.py index 2159e3802d..727c023ed6 100644 --- a/app/modules/proxy/_service/api_key_usage.py +++ b/app/modules/proxy/_service/api_key_usage.py @@ -23,6 +23,7 @@ from app.modules.proxy._service.support import ( _ApiKeyReservationTouchState, _consume_api_key_reservation_heartbeat_result, + _signal_propagated_responses_service_cleanup_ready, _StreamSettlement, _WebSocketRequestState, ) @@ -274,28 +275,34 @@ async def _settle_compact_api_key_usage( ) proxy = cast(_ApiKeyUsageServiceProtocol, self) - with anyio.CancelScope(shield=True): - try: - async with proxy._repo_factory() as repos: - api_keys_service = _service_api_keys_service()(repos.api_keys) - if response is not None and input_tokens is not None and output_tokens is not None: - await api_keys_service.finalize_usage_reservation( - reservation_id, - model=model_name, - input_tokens=input_tokens, - output_tokens=output_tokens, - cached_input_tokens=cached_input_tokens or 0, - service_tier=service_tier, - ) - else: - await api_keys_service.release_usage_reservation(reservation_id) - except Exception: - logger.warning( - "Failed to settle compact API key reservation key_id=%s request_id=%s", - api_key.id, - get_request_id(), - exc_info=True, - ) + try: + with anyio.CancelScope(shield=True): + try: + async with proxy._repo_factory() as repos: + api_keys_service = _service_api_keys_service()(repos.api_keys) + if response is not None and input_tokens is not None and output_tokens is not None: + await api_keys_service.finalize_usage_reservation( + reservation_id, + model=model_name, + input_tokens=input_tokens, + output_tokens=output_tokens, + cached_input_tokens=cached_input_tokens or 0, + service_tier=service_tier, + ) + else: + await api_keys_service.release_usage_reservation(reservation_id) + except Exception: + logger.warning( + "Failed to settle compact API key reservation key_id=%s request_id=%s", + api_key.id, + get_request_id(), + exc_info=True, + ) + finally: + # The compact service has made its one cancellation-safe settlement + # attempt. A caller that created the reservation must not issue a + # second release as a fallback after this boundary. + _signal_propagated_responses_service_cleanup_ready() async def _settle_stream_api_key_usage( self, @@ -429,7 +436,7 @@ def _schedule_cancel_safe_cleanup( *, action: str, request_id: str, - ) -> None: + ) -> asyncio.Task[None]: task = asyncio.create_task(coro, name=f"proxy-{action}-{request_id}") proxy = cast(_ApiKeyUsageServiceProtocol, self) proxy._background_cleanup_tasks.add(task) @@ -449,6 +456,7 @@ def _cleanup_done(done_task: asyncio.Task[None]) -> None: ) task.add_done_callback(_cleanup_done) + return task async def _release_unsettled_stream_api_key_usage( self, diff --git a/app/modules/proxy/_service/compact.py b/app/modules/proxy/_service/compact.py index 4623f61c46..b234b9a0a8 100644 --- a/app/modules/proxy/_service/compact.py +++ b/app/modules/proxy/_service/compact.py @@ -96,6 +96,15 @@ async def _resolve_file_account_for_responses( self, payload: ResponsesCompactRequest, headers: Mapping[str, str] ) -> str | None: ... + async def _resolve_forwarded_file_account_for_responses( + self, + payload: ResponsesCompactRequest, + headers: Mapping[str, str], + *, + forwarded_file_owner_account_id: str | None, + require_forwarded_file_owner: bool = False, + ) -> str | None: ... + async def _acquire_account_response_create_lease_or_overload( self, *, account_id: str, request_id: str, surface: str, concurrency_caps: AccountConcurrencyCaps ) -> AccountLease: ... @@ -573,6 +582,8 @@ async def compact_responses( api_key: ApiKeyData | None = None, api_key_reservation: ApiKeyUsageReservationData | None = None, client_ip: str | None = None, + forwarded_request: bool = False, + forwarded_file_owner_account_id: str | None = None, ) -> CompactResponsePayload: proxy = cast(_CompactServiceProtocol, self) _maybe_log_proxy_request_payload("compact", payload, headers) @@ -596,8 +607,57 @@ async def compact_responses( route_endpoint_id: str | None = None route_fallback_used: bool | None = None route_fail_closed_reason: str | None = None + settlement_attempted = False + + async def settle_compact_usage( + *, + api_key: ApiKeyData | None, + api_key_reservation: ApiKeyUsageReservationData | None, + response: CompactResponsePayload | None, + request_service_tier: str | None, + ) -> None: + nonlocal settlement_attempted + if settlement_attempted: + return + if forwarded_request and response is None: + # A forwarded receiver has not transferred cleanup ownership + # until its successful HTTP 200. Every error before that + # acknowledgement remains the origin's single release path. + return + settlement_attempted = True + await proxy._settle_compact_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + response=response, + request_service_tier=request_service_tier, + ) + proxy._raise_for_unsupported_input_image_references(payload) - rewritten_file_account_id = await proxy._resolve_file_account_for_responses(payload, headers) + try: + rewritten_file_account_id = await proxy._resolve_forwarded_file_account_for_responses( + payload, + headers, + forwarded_file_owner_account_id=forwarded_file_owner_account_id, + require_forwarded_file_owner=forwarded_request, + ) + except ProxyResponseError: + if not forwarded_request and api_key is not None and api_key_reservation is not None: + await settle_compact_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + response=None, + request_service_tier=_service_tier_from_compact_payload(payload), + ) + raise + except asyncio.CancelledError: + if not forwarded_request and api_key is not None and api_key_reservation is not None: + await settle_compact_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + response=None, + request_service_tier=_service_tier_from_compact_payload(payload), + ) + raise settings = await _service_get_settings_cache().get() concurrency_caps = effective_account_concurrency_caps(settings) prefer_earlier_reset = settings.prefer_earlier_reset_accounts @@ -913,13 +973,9 @@ async def _call_compact( if remaining_budget <= 0: logger.warning("Compact request budget exhausted before freshness check request_id=%s", request_id) await proxy._load_balancer.release_account_lease(selected_account_response_create_lease) - # This budget-exhausted terminal exits compact_responses before - # reaching the retry loop's settle sites, so on the HTTP bridge / - # forwarded path (``owns_reservation`` false, ``compact_responses`` - # is the sole settler) the API-key reservation would leak held - # quota. Settle BEFORE raising, mirroring the transport/permanent - # preflight branches above. - await proxy._settle_compact_api_key_usage( + # Direct requests transfer cleanup to the compact service at + # this terminal boundary. Forwarded failures remain origin-owned. + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -937,7 +993,7 @@ async def _call_compact( await proxy._load_balancer.release_account_lease(selected_account_response_create_lease) # Sole-settler leak guard (see above): settle the reservation # before this budget-exhausted terminal raise. - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -965,7 +1021,7 @@ async def _call_compact( # ensure_fresh_with_budget translates terminal process-network # recovery outcomes before the compact upstream settlement # branches run, so this boundary owns reservation cleanup. - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -978,15 +1034,9 @@ async def _call_compact( if isinstance(exc, RefreshError): if exc.is_permanent: # Permanent refresh failures keep their prior - # escalation (they propagate to the caller). On the - # HTTP bridge / forwarded path the caller passes an - # ``api_key_reservation_override`` with - # ``owns_reservation`` false, so ``compact_responses`` - # is the sole settler; settle BEFORE raising so the - # reservation is finalized instead of leaking held - # API-key quota (matching the post-401 permanent - # branch, which settles before re-raising). - await proxy._settle_compact_api_key_usage( + # escalation. Direct requests settle before raising; + # forwarded failures remain origin-owned until HTTP 200. + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1026,14 +1076,9 @@ async def _call_compact( ) if preferred_account_id is not None: # File/previous-response-pinned requests cannot - # fail over. On the HTTP bridge / forwarded path - # the caller passes an ``api_key_reservation_override`` - # with ``owns_reservation`` false, making - # ``compact_responses`` responsible for settling the - # reservation. Settle it BEFORE raising so the - # API-key reservation is finalized instead of leaking - # held quota when the pinned refresh claim times out. - await proxy._settle_compact_api_key_usage( + # fail over. Settle a direct request before raising; + # a forwarded rejection leaves cleanup at its origin. + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1058,15 +1103,11 @@ async def _call_compact( account.id, exc_info=True, ) - # Both terminal (non-failover) transport-failure raises below - # exit compact_responses without reaching the retry loop's - # settle sites, so on the HTTP bridge / forwarded path - # (owns_reservation false, compact_responses is the sole - # settler) the API-key reservation would leak held quota. - # Settle BEFORE raising, mirroring the claim-contention and - # post-401 transport branches. + # Both terminal transport failures exit before the retry + # loop's normal settle sites. Settle direct requests here; + # forwarded rejection cleanup remains at the origin. if not _should_retry_transient_stream_error("upstream_unavailable", message): - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1074,7 +1115,7 @@ async def _call_compact( ) _raise_proxy_unavailable(message) if preferred_account_id is not None: - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1103,7 +1144,7 @@ async def _call_compact( await proxy._load_balancer.release_account_lease(selected_account_response_create_lease) # Sole-settler leak guard (see above): settle the reservation # before this budget-exhausted terminal raise. - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1124,7 +1165,7 @@ async def _call_compact( network_recovery.log_recovered() actual_service_tier = _service_tier_from_response(response) await proxy._load_balancer.record_success(account) - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=response, @@ -1135,7 +1176,7 @@ async def _call_compact( except ProxyResponseError as exc: compact_continuity_error = _compact_previous_response_not_found_error(exc) if compact_continuity_error is not None: - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1154,7 +1195,7 @@ async def _call_compact( try: await proxy._handle_proxy_error(account, exc) except Exception: - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1174,13 +1215,10 @@ async def _call_compact( request_id, account.id, ) - # Sole-settler leak guard (see above): this - # budget-exhausted terminal exits the retry loop - # to the outer handler without settling, so on - # the bridge/forwarded path (``owns_reservation`` - # false) the reservation would leak held quota. - # Settle BEFORE raising. - await proxy._settle_compact_api_key_usage( + # This terminal exits before normal settlement. + # Direct requests settle here; forwarded rejection + # cleanup remains at the origin. + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1196,7 +1234,7 @@ async def _call_compact( # A translated refresh-recovery error escapes the # current upstream-error handler, so settle before # handing it to the request-level error boundary. - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1207,7 +1245,7 @@ async def _call_compact( if isinstance(refresh_exc, RefreshError): if refresh_exc.is_permanent: await proxy._load_balancer.mark_permanent_failure(account, refresh_exc.code) - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1250,7 +1288,7 @@ async def _call_compact( exc_info=True, ) if preferred_account_id is not None: - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1267,7 +1305,7 @@ async def _call_compact( # Non-transport, non-permanent RefreshError # keeps its prior escalation: re-raise the # original 401 to the caller. - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1291,7 +1329,7 @@ async def _call_compact( exc_info=True, ) if not _should_retry_transient_stream_error("upstream_unavailable", message): - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1299,7 +1337,7 @@ async def _call_compact( ) _raise_proxy_unavailable(message) if preferred_account_id is not None: - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1368,7 +1406,7 @@ async def _call_compact( if recovery_decision == "retry": continue if recovery_decision == "exhausted": - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1389,7 +1427,7 @@ async def _call_compact( require_security_work_authorized = True transient_exhausted = True break - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1402,7 +1440,7 @@ async def _call_compact( transient_exhausted = True break if _is_account_neutral_error_code(code): - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1410,7 +1448,7 @@ async def _call_compact( ) raise if code == "upstream_request_timeout": - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1488,7 +1526,7 @@ async def _call_compact( excluded_account_ids.add(account.id) transient_exhausted = True break - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1498,7 +1536,7 @@ async def _call_compact( if transient_exhausted: continue # outer loop: try different account # All account attempts exhausted — raise last error - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, @@ -1523,7 +1561,7 @@ async def _call_compact( route_fail_closed_reason = exc.reason log_error_code = "upstream_proxy_unavailable" log_error_message = exc.reason - await proxy._settle_compact_api_key_usage( + await settle_compact_usage( api_key=api_key, api_key_reservation=api_key_reservation, response=None, diff --git a/app/modules/proxy/_service/file_ops.py b/app/modules/proxy/_service/file_ops.py index bc7cbf233b..6c746312c8 100644 --- a/app/modules/proxy/_service/file_ops.py +++ b/app/modules/proxy/_service/file_ops.py @@ -29,10 +29,11 @@ from app.core.utils.request_id import ensure_request_id, get_request_id from app.db.models import Account from app.modules.api_keys.service import ApiKeyData -from app.modules.proxy._service.support import ( - _FilePinEntry, - _request_log_client_fields, - _RequestLogFailureMetadata, +from app.modules.proxy._service.support import _request_log_client_fields, _RequestLogFailureMetadata +from app.modules.proxy.continuity import resolve_required_account_id +from app.modules.proxy.file_pin_repository import ( + FileAccountPinOwnershipConflict, + FileAccountPinRepository, ) from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection @@ -44,10 +45,9 @@ class _FileOpsServiceProtocol(Protocol): _encryptor: Any - _file_account_pin_lock: asyncio.Lock - _file_account_pins: dict[str, _FilePinEntry] + _file_pin_session_factory: Any _load_balancer: Any - _FILE_ACCOUNT_PIN_TTL_SECONDS: float + _FILE_ACCOUNT_PIN_TTL_SECONDS: int async def _select_account_with_budget_compatible(self, deadline: float, **kwargs: object) -> AccountSelection: ... async def _select_account_with_budget(self, deadline: float, **kwargs: Any) -> AccountSelection: ... @@ -64,8 +64,11 @@ async def _resolve_upstream_route_for_account( async def _proxy_files_call(self, **kwargs: Any) -> tuple[dict[str, JsonValue], str | None]: ... async def _pin_file_account(self, file_id: str, account_id: str) -> None: ... async def _resolve_file_account(self, file_id: str) -> str | None: ... - async def _lookup_file_pin(self, file_id: str) -> _FilePinEntry | None: ... - def _evict_expired_file_pins_locked(self) -> None: ... + async def _resolve_file_account_for_responses( + self, + payload: ResponsesRequest | ResponsesCompactRequest, + headers: Mapping[str, str], + ) -> str | None: ... def _service_core_create_file() -> Callable[..., Awaitable[dict[str, JsonValue]]]: @@ -156,17 +159,32 @@ def _routing_strategy(settings: Any) -> RoutingStrategy: _REQUEST_TRANSPORT_HTTP = "http" +class _FileOwnerPostSuccessError(RuntimeError): + def __init__(self, proxy_error: ProxyResponseError) -> None: + super().__init__("File owner persistence failed after a successful upstream call") + self.proxy_error = proxy_error + + +def _file_owner_unavailable_error() -> ProxyResponseError: + return ProxyResponseError( + 502, + openai_error( + "file_owner_unavailable", + "Input file owner metadata is unavailable; upload the file again and retry.", + error_type="server_error", + ), + ) + + class _FileOpsMixin: # File-account pin TTL: long enough to cover a slow client-side # PUT of a 512 MiB upload (the upstream limit) plus the finalize # poll loop and a follow-up ``/responses`` that references the - # file_id, while still bounding how long stale pins can sit in - # memory on long-lived workers. 30 minutes covers a 512 MiB - # upload at ~280 KiB/s -- well below typical broadband uplink -- - # while keeping the table size negligible (each pin is a short - # string tuple). Eviction runs opportunistically on every write, - # so this acts as an upper bound, not a fixed retention. - _FILE_ACCOUNT_PIN_TTL_SECONDS: float = 30 * 60.0 + # file_id, while still bounding how long stale pins remain in + # shared storage. 30 minutes covers a 512 MiB + # upload at ~280 KiB/s -- well below typical broadband uplink. + # The database clock defines both expiry and opportunistic cleanup. + _FILE_ACCOUNT_PIN_TTL_SECONDS: int = 30 * 60 async def _pin_file_account( self, @@ -176,48 +194,41 @@ async def _pin_file_account( """Remember that ``file_id`` was registered through ``account_id``. Used so a subsequent ``finalize_file`` can be routed to the same - account that created the file. Cross-instance handoff is - best-effort: if the finalize lands on a different replica with - no pin, we fall back to a fresh load-balancer selection. + account that created the file, including when another replica + handles the follow-up request. """ proxy = cast(_FileOpsServiceProtocol, self) if not file_id or not account_id: return - expires_at = time.monotonic() + proxy._FILE_ACCOUNT_PIN_TTL_SECONDS - async with proxy._file_account_pin_lock: - proxy._file_account_pins[file_id] = _FilePinEntry( - account_id=account_id, - expires_at=expires_at, - ) - proxy._evict_expired_file_pins_locked() + try: + async with proxy._file_pin_session_factory() as session: + await FileAccountPinRepository(session).claim( + file_id, + account_id, + ttl_seconds=proxy._FILE_ACCOUNT_PIN_TTL_SECONDS, + ) + except FileAccountPinOwnershipConflict as exc: + raise ProxyResponseError( + 502, + openai_error( + "continuity_owner_conflict", + "File ownership conflicts with an existing live upload.", + error_type="server_error", + ), + ) from exc + except Exception as exc: + raise _file_owner_unavailable_error() from exc async def _resolve_file_account(self, file_id: str) -> str | None: """Return the pinned account_id for ``file_id`` if still live.""" - proxy = cast(_FileOpsServiceProtocol, self) - entry = await proxy._lookup_file_pin(file_id) - return entry.account_id if entry is not None else None - - async def _lookup_file_pin(self, file_id: str) -> _FilePinEntry | None: proxy = cast(_FileOpsServiceProtocol, self) if not file_id: return None - async with proxy._file_account_pin_lock: - proxy._evict_expired_file_pins_locked() - entry = proxy._file_account_pins.get(file_id) - if entry is None: - return None - if entry.expires_at <= time.monotonic(): - proxy._file_account_pins.pop(file_id, None) - return None - return entry - - def _evict_expired_file_pins_locked(self) -> None: - """Drop pins past their TTL. Called under ``_file_account_pin_lock``.""" - proxy = cast(_FileOpsServiceProtocol, self) - now = time.monotonic() - expired = [file_id for file_id, entry in proxy._file_account_pins.items() if entry.expires_at <= now] - for file_id in expired: - proxy._file_account_pins.pop(file_id, None) + try: + async with proxy._file_pin_session_factory() as session: + return await FileAccountPinRepository(session).get_live_account_id(file_id) + except Exception as exc: + raise _file_owner_unavailable_error() from exc async def _resolve_file_account_for_responses( self, @@ -226,8 +237,8 @@ async def _resolve_file_account_for_responses( ) -> str | None: """Resolve a ``preferred_account_id`` from ``input_file.file_id`` pins. - Looks up the in-memory ``file_id -> account_id`` pin table built - by ``create_file``. Used by ``/responses`` flows so a request + Looks up the durable ``file_id -> account_id`` pin table built by + ``create_file``. Used by ``/responses`` flows so a request carrying an ``{type: "input_file", file_id: "file_xxx"}`` part is routed to the same upstream account that registered the upload (the upstream contract is account-scoped via @@ -251,18 +262,21 @@ async def _resolve_file_account_for_responses( if not file_ids: return None - async with proxy._file_account_pin_lock: - proxy._evict_expired_file_pins_locked() - entries = [proxy._file_account_pins.get(file_id) for file_id in file_ids] - - pinned_entries = [entry for entry in entries if entry is not None] - if not pinned_entries: + try: + async with proxy._file_pin_session_factory() as session: + account_ids_by_file_id = await FileAccountPinRepository(session).get_live_account_ids(file_ids) + except Exception as exc: + raise _file_owner_unavailable_error() from exc + resolved_account_ids = [account_ids_by_file_id.get(file_id) for file_id in file_ids] + + pinned_account_ids = [account_id for account_id in resolved_account_ids if account_id is not None] + if not pinned_account_ids: # A raw file_id may have been registered directly with upstream or - # before this replica observed the upload. With zero local proof, + # before durable ownership was observed. With zero shared proof, # it remains an opaque compatibility reference rather than a hard # owner; callers forward it verbatim under ordinary routing. return None - if len(pinned_entries) != len(entries): + if len(pinned_account_ids) != len(resolved_account_ids): raise ProxyResponseError( 502, openai_error( @@ -271,7 +285,7 @@ async def _resolve_file_account_for_responses( error_type="server_error", ), ) - owner_account_ids = {entry.account_id for entry in pinned_entries} + owner_account_ids = set(pinned_account_ids) if len(owner_account_ids) != 1: raise ProxyResponseError( 502, @@ -283,6 +297,30 @@ async def _resolve_file_account_for_responses( ) return next(iter(owner_account_ids)) + async def _resolve_forwarded_file_account_for_responses( + self, + payload: ResponsesRequest | ResponsesCompactRequest, + headers: Mapping[str, str], + *, + forwarded_file_owner_account_id: str | None, + require_forwarded_file_owner: bool = False, + ) -> str | None: + """Revalidate signed bridge ownership against the shared database.""" + proxy = cast(_FileOpsServiceProtocol, self) + durable_owner_account_id = await proxy._resolve_file_account_for_responses(payload, headers) + if ( + require_forwarded_file_owner + and durable_owner_account_id is not None + and forwarded_file_owner_account_id is None + ): + raise _file_owner_unavailable_error() + if forwarded_file_owner_account_id is not None and durable_owner_account_id is None: + raise _file_owner_unavailable_error() + return resolve_required_account_id( + ("signed forwarding context", forwarded_file_owner_account_id), + ("durable file pin", durable_owner_account_id), + ) + def _raise_for_unsupported_input_image_references(self, payload: _ResponsesPayloadT) -> None: references = extract_input_image_file_references(payload.input) if not references: @@ -321,7 +359,13 @@ async def create_file( fail with not-found / unauthorized. """ proxy = cast(_FileOpsServiceProtocol, self) - result, account_id = await proxy._proxy_files_call( + + async def persist_file_owner(result: dict[str, JsonValue], account_id: str) -> None: + file_id = result.get("file_id") + if isinstance(file_id, str) and file_id: + await proxy._pin_file_account(file_id, account_id) + + result, _account_id = await proxy._proxy_files_call( log_model="files-create", kind="files-create", api_key=api_key, @@ -337,12 +381,8 @@ async def create_file( route_trace=route_trace, ) ), + on_success=persist_file_owner, ) - # Best-effort pin so finalize lands on the same account. - if isinstance(result, dict) and account_id: - file_id = result.get("file_id") - if isinstance(file_id, str) and file_id: - await proxy._pin_file_account(file_id, account_id) return result async def finalize_file( @@ -360,20 +400,26 @@ async def finalize_file( verbatim. Routes to the account that handled the matching ``create_file`` - (via the in-memory pin table) so the upstream finalize call + (via the durable pin table) so the upstream finalize call carries the same ``chatgpt-account-id`` that registered the file. Falls back to a fresh load-balancer selection when no - pin is found (unknown ``file_id`` or pin expired / missed across - a replica boundary). + pin is found (unknown ``file_id`` or an expired pin). """ proxy = cast(_FileOpsServiceProtocol, self) - pinned_account_id = await proxy._resolve_file_account(file_id) - result, account_id = await proxy._proxy_files_call( + + async def resolve_file_owner() -> str | None: + return await proxy._resolve_file_account(file_id) + + async def persist_file_owner(result: dict[str, JsonValue], account_id: str) -> None: + if result.get("status") == "success": + await proxy._pin_file_account(file_id, account_id) + + result, _account_id = await proxy._proxy_files_call( log_model="files-finalize", kind="files-finalize", api_key=api_key, headers=headers, - preferred_account_id=pinned_account_id, + resolve_preferred_account_id=resolve_file_owner, invoke=lambda access_token, upstream_account_id, filtered_headers, route, route_trace: ( _service_core_finalize_file()( file_id=file_id, @@ -385,11 +431,8 @@ async def finalize_file( route_trace=route_trace, ) ), + on_success=persist_file_owner, ) - if isinstance(result, dict) and account_id: - status = result.get("status") - if status == "success": - await proxy._pin_file_account(file_id, account_id) return result async def _proxy_files_call( @@ -404,6 +447,8 @@ async def _proxy_files_call( Awaitable[dict[str, JsonValue]], ], preferred_account_id: str | None = None, + resolve_preferred_account_id: Callable[[], Awaitable[str | None]] | None = None, + on_success: Callable[[dict[str, JsonValue], str], Awaitable[None]] | None = None, ) -> tuple[dict[str, JsonValue], str | None]: """Shared account-selection / refresh / 401-retry plumbing for `/files` calls. @@ -411,9 +456,11 @@ async def _proxy_files_call( ensure freshness, invoke upstream, on 401 force-refresh and retry once, translate ``FileProxyError`` -> ``ProxyResponseError``, and always write a request-log entry on the way out. When - ``preferred_account_id`` is provided (e.g. from the file_id pin - for ``finalize_file``), the call is strict to that account and - fails closed when the owner account is unavailable. + ``preferred_account_id`` is provided or resolved (e.g. from the file_id + pin for ``finalize_file``), the call is strict to that account and + fails closed when the owner account is unavailable. ``on_success`` runs + before the request is logged or returned so durable owner persistence + remains part of the route's success contract. """ proxy = cast(_FileOpsServiceProtocol, self) filtered = filter_inbound_headers(headers) @@ -433,10 +480,23 @@ async def _proxy_files_call( route_fallback_used: bool | None = None route_fail_closed_reason: str | None = None - settings = await _service_get_settings_cache().get() - prefer_earlier_reset = settings.prefer_earlier_reset_accounts - routing_strategy = _routing_strategy(settings) try: + if resolve_preferred_account_id is not None: + preferred_account_id = await resolve_preferred_account_id() + settings = await _service_get_settings_cache().get() + prefer_earlier_reset = settings.prefer_earlier_reset_accounts + routing_strategy = _routing_strategy(settings) + + async def _persist_success(result: dict[str, JsonValue], account_id: str) -> None: + if on_success is None: + return + try: + await on_success(result, account_id) + except ProxyResponseError as exc: + raise _FileOwnerPostSuccessError(exc) from exc + except Exception as exc: + raise _FileOwnerPostSuccessError(_file_owner_unavailable_error()) from exc + selection = await proxy._select_account_with_budget_compatible( deadline, request_id=request_id, @@ -528,6 +588,7 @@ async def _select_files_failover(excluded_account_ids: set[str]) -> AccountSelec account_id_value = account.id result = await _call(account) await proxy._load_balancer.record_success(account) + await _persist_success(result, account.id) log_status = "success" return result, account_id_value except RefreshError as refresh_exc: @@ -556,6 +617,7 @@ async def _select_files_failover(excluded_account_ids: set[str]) -> AccountSelec if failover is not None: account, result = failover account_id_value = account.id + await _persist_success(result, account.id) log_status = "success" return result, account_id_value failed_account = _proxy_response_failed_account(exc, account) @@ -610,6 +672,7 @@ async def _select_files_failover(excluded_account_ids: set[str]) -> AccountSelec # caller's pin is consistent with the upstream call. account_id_value = account.id await proxy._load_balancer.record_success(account) + await _persist_success(result, account.id) log_status = "success" return result, account_id_value except ProxyResponseError as retry_exc: @@ -637,12 +700,23 @@ async def _select_files_failover(excluded_account_ids: set[str]) -> AccountSelec try: result = await _call(account) await proxy._load_balancer.record_success(account) + await _persist_success(result, account.id) log_status = "success" return result, account_id_value except ProxyResponseError as failover_exc: await proxy._handle_proxy_error(account, failover_exc) raise raise + except _FileOwnerPostSuccessError as exc: + proxy_error = exc.proxy_error + failure_metadata = _request_log_failure_metadata(proxy_error) + error = _parse_openai_error(proxy_error.payload) + log_error_code = _normalize_error_code( + error.code if error else None, + error.type if error else None, + ) + log_error_message = error.message if error else None + raise proxy_error from exc except ProxyResponseError as exc: failed_account = getattr(exc, _FAILED_ACCOUNT_ATTR, None) if isinstance(failed_account, Account): diff --git a/app/modules/proxy/_service/http_bridge/owner_forwarding.py b/app/modules/proxy/_service/http_bridge/owner_forwarding.py index 51cdbda989..a87976a86f 100644 --- a/app/modules/proxy/_service/http_bridge/owner_forwarding.py +++ b/app/modules/proxy/_service/http_bridge/owner_forwarding.py @@ -2,6 +2,7 @@ import asyncio import logging +from enum import StrEnum from typing import Any, AsyncIterator, Mapping, TypeVar import aiohttp @@ -100,6 +101,9 @@ _HTTPBridgeSessionKey, _signal_propagated_capacity_startup_ready, _signal_propagated_capacity_startup_wait, + _signal_propagated_responses_owner_forward_dispatched, + _signal_propagated_responses_owner_forward_rejected, + _signal_propagated_responses_service_cleanup_ready, ) from app.modules.proxy._service.support import ( _websocket_route_log_kwargs as _websocket_route_log_kwargs, @@ -154,6 +158,43 @@ T = TypeVar("T") +class _OwnerForwardOutcome(StrEnum): + NOT_DISPATCHED = "not_dispatched" + DISPATCH_AMBIGUOUS = "dispatch_ambiguous" + RECEIVER_ACKNOWLEDGED = "receiver_acknowledged" + RECEIVER_REJECTED = "receiver_rejected" + + +class _OwnerForwardRequestError(ProxyResponseError): + def __init__( + self, + source: ProxyResponseError, + *, + outcome: _OwnerForwardOutcome, + ) -> None: + super().__init__( + source.status_code, + source.payload, + failure_phase=source.failure_phase, + retryable_same_contract=source.retryable_same_contract, + failure_detail=source.failure_detail, + failure_exception_type=source.failure_exception_type, + upstream_status_code=source.upstream_status_code, + upstream_error_code=source.upstream_error_code, + failed_session=source.failed_session, + ) + self.outcome = outcome + + +def _owner_forward_failure_allows_local_recovery(exc: ProxyResponseError) -> bool: + if not isinstance(exc, _OwnerForwardRequestError): + return True + return exc.outcome in { + _OwnerForwardOutcome.NOT_DISPATCHED, + _OwnerForwardOutcome.RECEIVER_REJECTED, + } + + def _durable_recovery_supersedes_local_session( durable_lookup: DurableBridgeLookup | None, session: _HTTPBridgeSession, @@ -357,8 +398,8 @@ async def _forward_http_bridge_request_to_owner( incoming_turn_state = None forwarded_turn_state = incoming_turn_state or downstream_turn_state # file_owner_account_id is an origin-side ownership proof, not a route - # hint. The forwarding signer binds it before another replica may skip - # its own process-local file-pin lookup. + # hint. The forwarding signer binds it, and the receiving replica + # corroborates it against the current durable pin before routing. forward_context = HTTPBridgeForwardContext( origin_instance=current_instance, target_instance=owner_forward.owner_instance, @@ -394,8 +435,32 @@ async def _forward_http_bridge_request_to_owner( owner_check_applied=True, ) + forward_outcome = _OwnerForwardOutcome.NOT_DISPATCHED forwarded_any = False forwarded_response_id: str | None = None + + def owner_response_ready() -> None: + nonlocal forward_outcome + forward_outcome = _OwnerForwardOutcome.RECEIVER_ACKNOWLEDGED + _signal_propagated_capacity_startup_ready() + if api_key_reservation is not None: + # A receiver carrying the origin reservation delays its 200 + # response until its settlement finalizer is active. Mirror + # that explicit handoff into the origin's startup guard. + _signal_propagated_responses_service_cleanup_ready() + + def owner_request_dispatched() -> None: + nonlocal forward_outcome + forward_outcome = _OwnerForwardOutcome.DISPATCH_AMBIGUOUS + if api_key_reservation is not None: + _signal_propagated_responses_owner_forward_dispatched() + + def owner_response_rejected() -> None: + nonlocal forward_outcome + forward_outcome = _OwnerForwardOutcome.RECEIVER_REJECTED + if api_key_reservation is not None: + _signal_propagated_responses_owner_forward_rejected() + try: async for event_block in self._http_bridge_owner_client.stream_responses( owner_endpoint=owner_forward.owner_endpoint, @@ -404,7 +469,9 @@ async def _forward_http_bridge_request_to_owner( context=forward_context, request_started_at=request_started_at, on_response_wait=_signal_propagated_capacity_startup_wait, - on_response_ready=_signal_propagated_capacity_startup_ready, + on_request_dispatched=owner_request_dispatched, + on_response_rejected=owner_response_rejected, + on_response_ready=owner_response_ready, ): forwarded_any = True event_payload = parse_sse_data_json(event_block) @@ -432,7 +499,7 @@ async def _forward_http_bridge_request_to_owner( if forwarded_any: yield exc.event_block return - raise ProxyResponseError( + error = ProxyResponseError( 503, openai_error( "bridge_owner_unreachable", @@ -442,7 +509,8 @@ async def _forward_http_bridge_request_to_owner( failure_phase="owner_forward", failure_detail="relay_timeout", failure_exception_type=type(exc).__name__, - ) from exc + ) + raise _OwnerForwardRequestError(error, outcome=forward_outcome) from exc except ProxyResponseError as exc: if PROMETHEUS_AVAILABLE and bridge_owner_forward_total is not None: bridge_owner_forward_total.labels(outcome="fail").inc() @@ -467,7 +535,7 @@ async def _forward_http_bridge_request_to_owner( default_message="HTTP bridge owner request failed", ) return - raise + raise _OwnerForwardRequestError(exc, outcome=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() @@ -493,7 +561,7 @@ async def _forward_http_bridge_request_to_owner( ) ) return - raise ProxyResponseError( + error = ProxyResponseError( 503, openai_error( "bridge_owner_unreachable", @@ -503,7 +571,8 @@ async def _forward_http_bridge_request_to_owner( failure_phase="owner_forward", failure_detail=str(exc) or "owner_forward_request_failed", failure_exception_type=type(exc).__name__, - ) from exc + ) + raise _OwnerForwardRequestError(error, outcome=forward_outcome) from exc else: if PROMETHEUS_AVAILABLE and bridge_owner_forward_total is not None: bridge_owner_forward_total.labels(outcome="success").inc() diff --git a/app/modules/proxy/_service/http_bridge/protocol.py b/app/modules/proxy/_service/http_bridge/protocol.py index b88edce03e..f6f43900c6 100644 --- a/app/modules/proxy/_service/http_bridge/protocol.py +++ b/app/modules/proxy/_service/http_bridge/protocol.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from collections.abc import Mapping from typing import Any, Protocol @@ -102,4 +103,4 @@ async def _maybe_touch_request_state_api_key_reservation(self, *args: Any, **kwa async def _reserve_websocket_api_key_usage(self, *args: Any, **kwargs: Any) -> Any: ... async def _release_websocket_reservation(self, *args: Any, **kwargs: Any) -> None: ... async def _release_websocket_request_state_reservation(self, *args: Any, **kwargs: Any) -> None: ... - def _schedule_cancel_safe_cleanup(self, *args: Any, **kwargs: Any) -> None: ... + def _schedule_cancel_safe_cleanup(self, *args: Any, **kwargs: Any) -> asyncio.Task[None]: ... diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index 1fce2706ea..028ee7dc66 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -94,6 +94,9 @@ _reserve_http_bridge_unanchored_handoff, _trim_http_bridge_previous_response_input_items, ) +from app.modules.proxy._service.http_bridge.owner_forwarding import ( + _owner_forward_failure_allows_local_recovery, +) from app.modules.proxy._service.http_bridge.service_stubs import ( _build_rewritten_stream_response_failed_event, _codex_keepalive_frame, @@ -151,6 +154,7 @@ _is_local_account_cap_code, _signal_propagated_capacity_startup_ready, _signal_propagated_capacity_startup_wait, + _signal_propagated_responses_service_cleanup_ready, _ttft_event_visible_at, _WebSocketRequestState, ) @@ -604,17 +608,11 @@ async def _stream_http_bridge_or_retry( payload_size_estimate_bytes = len( json.dumps(payload.to_payload(), ensure_ascii=True, separators=(",", ":")).encode("utf-8") ) - # File pins are process-local. A remote owner must trust only the - # origin-resolved value carried by the authenticated forward context; - # re-looking it up here would turn a valid cross-replica pin into a miss. - local_file_owner_account_id = ( - None - if forwarded_file_owner_account_id is not None - else await self._resolve_file_account_for_responses(payload, headers) - ) - rewritten_file_account_id = resolve_required_account_id( - ("signed forwarding context", forwarded_file_owner_account_id), - ("local file pin", local_file_owner_account_id), + rewritten_file_account_id = await self._resolve_forwarded_file_account_for_responses( + payload, + headers, + forwarded_file_owner_account_id=forwarded_file_owner_account_id, + require_forwarded_file_owner=forwarded_request, ) ws_payload_budget_bytes = _ws_transport_payload_budget_bytes(_service_get_settings()) if runtime_config.enabled and payload_size_estimate_bytes > ws_payload_budget_bytes: @@ -650,6 +648,7 @@ async def _stream_http_bridge_or_retry( suppress_text_done_events=suppress_text_done_events, request_transport=_REQUEST_TRANSPORT_HTTP, rewritten_file_account_id=rewritten_file_account_id, + file_account_resolution_complete=True, upstream_stream_transport_override=force_upstream_stream_transport, client_ip=client_ip, enforce_openai_sdk_contract=enforce_openai_sdk_contract, @@ -1398,6 +1397,18 @@ def switch_to_account_neutral_replay() -> None: continue break if isinstance(session_or_forward, _HTTPBridgeOwnerForward): + if forwarded_request: + raise ProxyResponseError( + 503, + openai_error( + "bridge_forward_loop_prevented", + ( + "HTTP bridge request was forwarded back to a non-owner instance; " + "refusing takeover to avoid a forward loop" + ), + error_type="server_error", + ), + ) await _current_origin_legacy_owner_anchor_lookup( durable_bridge=self._durable_bridge, bridge_session_key=session_or_forward.key, @@ -1434,6 +1445,8 @@ def switch_to_account_neutral_replay() -> None: default_message="HTTP bridge owner request failed", ) return + if not _owner_forward_failure_allows_local_recovery(exc): + raise owner_forward_fresh_replay = owner_unavailable_allows_account_neutral_replay(exc) if owner_forward_fresh_replay: switch_to_account_neutral_replay() @@ -1718,17 +1731,6 @@ def switch_to_account_neutral_replay() -> None: retry_request_state: _WebSocketRequestState | None = None try: retry_api_key_reservation = api_key_reservation - retry_reservation_reacquired = False - if api_key is not None and api_key_reservation is not None: - retry_api_key_reservation = await self._reserve_websocket_api_key_usage( - api_key, - request_model=recovery_payload.model, - request_service_tier=_normalize_service_tier_value( - dict(recovery_payload.to_payload()).get("service_tier"), - ), - request_usage_budget=estimate_api_key_request_usage(recovery_payload), - ) - retry_reservation_reacquired = True retry_request_state, retry_text_data = prepare_bridge_request( recovery_payload, @@ -1760,10 +1762,6 @@ def switch_to_account_neutral_replay() -> None: request_deadline=request_deadline, ): yield event_block - except BaseException: - if retry_reservation_reacquired and retry_api_key_reservation is not None: - await self._release_websocket_reservation(retry_api_key_reservation) - raise finally: if owner_recovery_scope_id is not None: _release_http_bridge_unanchored_handoff( @@ -2246,7 +2244,7 @@ def switch_to_account_neutral_replay() -> None: retry_payload = _http_bridge_payload_without_previous_response_id(untrimmed_effective_payload) retry_previous_response_id = None retry_request_stage = "context_overflow_recover" - retry_preferred_account_id = None + retry_preferred_account_id = rewritten_file_account_id allow_previous_response_recovery_rebind = False elif should_rollover_after_context_overflow: _log_http_bridge_event( @@ -2395,6 +2393,7 @@ def switch_to_account_neutral_replay() -> None: retry_request_state.transport = _REQUEST_TRANSPORT_HTTP retry_request_state.request_stage = retry_request_stage retry_request_state.preferred_account_id = retry_preferred_account_id + retry_request_state.file_required_preferred_account = file_required_preferred_account retry_request_state.excluded_account_ids.update(request_state.excluded_account_ids) retry_events: AsyncGenerator[str, None] = self._stream_http_bridge_session_events( @@ -2555,6 +2554,11 @@ async def _stream_http_bridge_session_events( continue break try: + # A successful submit installs the request-state finalizer that + # settles or releases the reservation. Signal before the next + # cancellation point so the API startup guard cannot race that + # finalizer while no upstream event has arrived yet. + _signal_propagated_responses_service_cleanup_ready() if downstream_turn_state is not None and not account_neutral_recovery: await self._register_http_bridge_turn_state(session, downstream_turn_state) _signal_propagated_capacity_startup_ready() diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 35e12de908..8dfe2fc13a 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -42,6 +42,7 @@ _request_log_client_fields, _RetryableStreamError, _signal_propagated_capacity_startup_wait, + _signal_propagated_responses_service_cleanup_ready, _stream_settlement_error_payload, _StreamSettlement, _TerminalStreamError, @@ -241,6 +242,7 @@ async def _stream_with_retry( suppress_text_done_events: bool, request_transport: str, rewritten_file_account_id: str | None = None, + file_account_resolution_complete: bool = False, upstream_stream_transport_override: str | None = None, client_ip: str | None = None, enforce_openai_sdk_contract: bool = True, @@ -305,7 +307,7 @@ async def _stream_with_retry( upstream_stream_transport, request_id, ) - if rewritten_file_account_id is None: + if not file_account_resolution_complete: proxy._raise_for_unsupported_input_image_references(payload) rewritten_file_account_id = await proxy._resolve_file_account_for_responses(payload, headers) had_prompt_cache_key = _prompt_cache_key_from_request_model(payload) is not None @@ -812,6 +814,10 @@ async def _retry_account_model_rejection( return True try: + # From this exact point the service finalizer below owns reservation + # settlement/release. Preflight failures before this boundary are + # still owned by the originating API startup guard. + _signal_propagated_responses_service_cleanup_ready() if payload.previous_response_id is not None: previous_response_lookup_session_id = _owner_lookup_session_id_from_headers(headers) preferred_account_id = await proxy._resolve_websocket_previous_response_owner( diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index 0d842e78f9..d714b0eead 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -91,6 +91,18 @@ "propagated_capacity_startup_ready", default=None, ) +_PROPAGATED_RESPONSES_SERVICE_CLEANUP_READY: ContextVar[asyncio.Event | None] = ContextVar( + "propagated_responses_service_cleanup_ready", + default=None, +) +_PROPAGATED_RESPONSES_OWNER_FORWARD_DISPATCHED: ContextVar[asyncio.Event | None] = ContextVar( + "propagated_responses_owner_forward_dispatched", + default=None, +) +_PROPAGATED_RESPONSES_OWNER_FORWARD_REJECTED: ContextVar[asyncio.Event | None] = ContextVar( + "propagated_responses_owner_forward_rejected", + default=None, +) def _strip_blank_html_comment_lines(text: str) -> str: @@ -283,6 +295,48 @@ def _signal_propagated_capacity_startup_ready() -> None: event.set() +def _bind_propagated_responses_service_cleanup_ready(event: asyncio.Event) -> Token[asyncio.Event | None]: + return _PROPAGATED_RESPONSES_SERVICE_CLEANUP_READY.set(event) + + +def _reset_propagated_responses_service_cleanup_ready(token: Token[asyncio.Event | None]) -> None: + _PROPAGATED_RESPONSES_SERVICE_CLEANUP_READY.reset(token) + + +def _signal_propagated_responses_service_cleanup_ready() -> None: + event = _PROPAGATED_RESPONSES_SERVICE_CLEANUP_READY.get() + if event is not None: + event.set() + + +def _bind_propagated_responses_owner_forward_dispatched(event: asyncio.Event) -> Token[asyncio.Event | None]: + return _PROPAGATED_RESPONSES_OWNER_FORWARD_DISPATCHED.set(event) + + +def _reset_propagated_responses_owner_forward_dispatched(token: Token[asyncio.Event | None]) -> None: + _PROPAGATED_RESPONSES_OWNER_FORWARD_DISPATCHED.reset(token) + + +def _signal_propagated_responses_owner_forward_dispatched() -> None: + event = _PROPAGATED_RESPONSES_OWNER_FORWARD_DISPATCHED.get() + if event is not None: + event.set() + + +def _bind_propagated_responses_owner_forward_rejected(event: asyncio.Event) -> Token[asyncio.Event | None]: + return _PROPAGATED_RESPONSES_OWNER_FORWARD_REJECTED.set(event) + + +def _reset_propagated_responses_owner_forward_rejected(token: Token[asyncio.Event | None]) -> None: + _PROPAGATED_RESPONSES_OWNER_FORWARD_REJECTED.reset(token) + + +def _signal_propagated_responses_owner_forward_rejected() -> None: + event = _PROPAGATED_RESPONSES_OWNER_FORWARD_REJECTED.get() + if event is not None: + event.set() + + def _account_selection_recovery_sleep_seconds_from_message( message: str | None, *, @@ -723,12 +777,6 @@ def _consume_api_key_reservation_heartbeat_result(task: asyncio.Task[None]) -> N logger.warning("API key reservation heartbeat task failed during cancellation", exc_info=True) -@dataclass(frozen=True, slots=True) -class _FilePinEntry: - account_id: str - expires_at: float - - @dataclass(frozen=True, slots=True) class _RequestLogFailureMetadata: failure_phase: str | None = None diff --git a/app/modules/proxy/_support.py b/app/modules/proxy/_support.py index b9d401a16f..44a28f5ca0 100644 --- a/app/modules/proxy/_support.py +++ b/app/modules/proxy/_support.py @@ -36,9 +36,6 @@ from app.modules.proxy._service.support import ( _event_type_from_payload as _event_type_from_payload, ) -from app.modules.proxy._service.support import ( - _FilePinEntry as _FilePinEntry, -) from app.modules.proxy._service.support import ( _HTTPBridgeOwnerForward as _HTTPBridgeOwnerForward, ) diff --git a/app/modules/proxy/api.py b/app/modules/proxy/api.py index fcbc012a8a..9432d699b3 100644 --- a/app/modules/proxy/api.py +++ b/app/modules/proxy/api.py @@ -4,7 +4,7 @@ import json import logging import time -from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Mapping from contextlib import asynccontextmanager from dataclasses import dataclass, replace from datetime import datetime, timezone @@ -12,6 +12,7 @@ from typing import Any, Final, Literal, Protocol, cast from uuid import uuid4 +import anyio from fastapi import ( APIRouter, Body, @@ -199,12 +200,18 @@ from app.modules.proxy._service.support import ( _bind_propagated_capacity_startup_ready, _bind_propagated_capacity_startup_wait, + _bind_propagated_responses_owner_forward_dispatched, + _bind_propagated_responses_owner_forward_rejected, + _bind_propagated_responses_service_cleanup_ready, _could_be_blank_html_comment_line, _is_reasoning_summary_interleavable_event, _reasoning_summary_delta_key, _request_log_client_fields, _reset_propagated_capacity_startup_ready, _reset_propagated_capacity_startup_wait, + _reset_propagated_responses_owner_forward_dispatched, + _reset_propagated_responses_owner_forward_rejected, + _reset_propagated_responses_service_cleanup_ready, _strip_blank_html_comment_lines, ) from app.modules.proxy.account_cache import get_account_selection_cache @@ -4781,6 +4788,65 @@ async def _source_chat_stream_with_settlement( ) +@dataclass(slots=True) +class _ResponsesReservationCleanup: + owns_reservation: bool + reservation: ApiKeyUsageReservationData | None + scheduler: _ResponsesCleanupScheduler | None + request_id: str + released: bool = False + + async def release(self, *, action: str) -> None: + if not self.owns_reservation or self.released: + return + self.released = True + await _release_reservation_best_effort( + self.reservation, + action=action, + scheduler=self.scheduler, + request_id=self.request_id, + ) + + +class _ResponsesCleanupScheduler(Protocol): + def _schedule_cancel_safe_cleanup( + self, + coro: Coroutine[Any, Any, None], + *, + action: str, + request_id: str, + ) -> asyncio.Task[None]: ... + + +def _responses_origin_may_release_reservation( + *, + service_cleanup_ready_event: asyncio.Event, + owner_forward_dispatched_event: asyncio.Event | None = None, + owner_forward_rejected_event: asyncio.Event | None = None, +) -> bool: + if service_cleanup_ready_event.is_set(): + return False + if owner_forward_dispatched_event is None or not owner_forward_dispatched_event.is_set(): + return True + return owner_forward_rejected_event is not None and owner_forward_rejected_event.is_set() + + +async def _rate_limit_headers_after_reservation( + context: ProxyContext, + api_key: ApiKeyData | None, + *, + reservation_cleanup: _ResponsesReservationCleanup, + include_headers: bool = True, +) -> dict[str, str]: + if not include_headers: + return {} + try: + return await _rate_limit_headers_for_request(context, api_key) + except BaseException: + await reservation_cleanup.release(action="rate-limit headers") + raise + + async def _stream_responses( request: Request, payload: ResponsesRequest, @@ -4874,8 +4940,26 @@ async def _stream_responses( request_usage_budget=estimate_api_key_request_usage(payload), ) ) + reservation_cleanup = _ResponsesReservationCleanup( + owns_reservation=owns_reservation, + reservation=reservation, + scheduler=( + cast(_ResponsesCleanupScheduler, context.service) + if callable(getattr(context.service, "_schedule_cancel_safe_cleanup", None)) + else None + ), + request_id=ensure_request_id(), + ) + responses_service_cleanup_ready_event = asyncio.Event() + responses_owner_forward_dispatched_event = asyncio.Event() + responses_owner_forward_rejected_event = asyncio.Event() - 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_after_reservation( + context, + api_key, + reservation_cleanup=reservation_cleanup, + include_headers=include_rate_limit_headers, + ) 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) @@ -4892,6 +4976,9 @@ async def _stream_responses( else {} ) if compact_payload is not None: + responses_cleanup_ready_token = _bind_propagated_responses_service_cleanup_ready( + responses_service_cleanup_ready_event + ) try: try: compact_result = await context.service.compact_responses( @@ -4902,6 +4989,8 @@ async def _stream_responses( api_key=api_key, api_key_reservation=reservation, client_ip=client_ip, + forwarded_request=forwarded_request, + forwarded_file_owner_account_id=forwarded_file_owner_account_id, ) except NotImplementedError: error = OpenAIErrorEnvelopeModel( @@ -4925,18 +5014,24 @@ async def _stream_responses( ) compact_item = _compact_response_output_item(compact_result) if compact_item is None: - error = openai_error( - "upstream_error", - "Compact response did not include a compaction output item", - error_type="server_error", + if forwarded_request and responses_service_cleanup_ready_event.is_set(): + # The compact service already settled the forwarded + # reservation, so HTTP 200 must remain the handoff + # acknowledgement even when its payload is malformed. + stream = _synthetic_compaction_failure_stream(response_id=_compact_response_id(compact_result)) + else: + error = openai_error( + "upstream_error", + "Compact response did not include a compaction output item", + error_type="server_error", + ) + return _logged_error_json_response(request, 502, error, headers=rate_limit_headers) + else: + stream = _synthetic_compaction_response_stream( + compact_item, + response_id=_compact_response_id(compact_result), + usage=compact_result.usage, ) - return _logged_error_json_response(request, 502, error, headers=rate_limit_headers) - response_id = _compact_response_id(compact_result) - stream = _synthetic_compaction_response_stream( - compact_item, - response_id=response_id, - usage=compact_result.usage, - ) return StreamingResponse( stream, media_type="text/event-stream", @@ -4948,8 +5043,11 @@ async def _stream_responses( }, ) finally: - if owns_reservation: - await _release_reservation(reservation) + _reset_propagated_responses_service_cleanup_ready(responses_cleanup_ready_token) + if _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event + ): + await reservation_cleanup.release(action="terminal compaction response") capacity_wait_event = asyncio.Event() capacity_ready_event = _CapacityStartupReadyEvent() payload.stream = True @@ -4988,37 +5086,80 @@ async def _stream_responses( client_ip=client_ip, enforce_openai_sdk_contract=enforce_openai_sdk_contract, ) + service_stream = stream + startup_handoff_tasks: list[asyncio.Task[str]] = [] capacity_wait_token = _bind_propagated_capacity_startup_wait(capacity_wait_event) capacity_ready_token = _bind_propagated_capacity_startup_ready(capacity_ready_event) + responses_owner_forward_dispatched_token = _bind_propagated_responses_owner_forward_dispatched( + responses_owner_forward_dispatched_event + ) + responses_owner_forward_rejected_token = _bind_propagated_responses_owner_forward_rejected( + responses_owner_forward_rejected_event + ) + responses_cleanup_ready_token = _bind_propagated_responses_service_cleanup_ready( + responses_service_cleanup_ready_event + ) try: - stream, startup_error = await _probe_stream_startup_error( - stream, - convert_event_errors=bridge_active and enforce_openai_sdk_contract, - timeout_seconds=( - _HTTP_BRIDGE_STARTUP_ERROR_PROBE_SECONDS if prefer_http_bridge else _STREAM_STARTUP_ERROR_PROBE_SECONDS - ), - capacity_wait_event=capacity_wait_event, - capacity_ready_event=capacity_ready_event, - ) - finally: - _reset_propagated_capacity_startup_ready(capacity_ready_token) - _reset_propagated_capacity_startup_wait(capacity_wait_token) + try: + stream, startup_error = await _probe_stream_startup_error( + stream, + convert_event_errors=bridge_active and enforce_openai_sdk_contract, + timeout_seconds=( + _HTTP_BRIDGE_STARTUP_ERROR_PROBE_SECONDS + if prefer_http_bridge + else _STREAM_STARTUP_ERROR_PROBE_SECONDS + ), + capacity_wait_event=capacity_wait_event, + capacity_ready_event=capacity_ready_event, + handoff_task_sink=startup_handoff_tasks, + service_cleanup_ready_event=( + responses_service_cleanup_ready_event if forwarded_request and reservation is not None else None + ), + ) + finally: + _reset_propagated_responses_service_cleanup_ready(responses_cleanup_ready_token) + _reset_propagated_responses_owner_forward_rejected(responses_owner_forward_rejected_token) + _reset_propagated_responses_owner_forward_dispatched(responses_owner_forward_dispatched_token) + _reset_propagated_capacity_startup_ready(capacity_ready_token) + _reset_propagated_capacity_startup_wait(capacity_wait_token) + except BaseException: + # Until the startup probe hands the iterator to StreamingResponse, the + # route owns cancellation/error cleanup for reservations it created. + if _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + await reservation_cleanup.release(action="responses startup") + raise if startup_error is not None: - if owns_reservation: - await _release_reservation(reservation) + if _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + await reservation_cleanup.release(action="responses startup error") return _stream_startup_error_response( request, startup_error, headers=rate_limit_headers, ) + startup_handoff_streams: list[AsyncIterator[str]] = [service_stream, stream] + stream = _stream_response_error_events( + stream, + owns_reservation=owns_reservation, + reservation=reservation, + reservation_cleanup=reservation_cleanup, + responses_service_cleanup_ready_event=responses_service_cleanup_ready_event, + responses_owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + responses_owner_forward_rejected_event=responses_owner_forward_rejected_event, + ) + startup_handoff_streams.append(stream) stream = _normalize_public_responses_stream( - _stream_response_error_events( - stream, - owns_reservation=owns_reservation, - reservation=reservation, - ), + stream, enforce_openai_sdk_contract=enforce_openai_sdk_contract, ) + startup_handoff_streams.append(stream) use_codex_keepalive = native_codex_heartbeat or not enforce_openai_sdk_contract keepalive_frame = CODEX_KEEPALIVE_FRAME if use_codex_keepalive else SSE_KEEPALIVE_FRAME if use_codex_keepalive: @@ -5028,12 +5169,25 @@ async def _stream_responses( request_id=get_request_id(), route_family="responses", ) - return StreamingResponse( - inject_sse_keepalives( + startup_handoff_streams.append(stream) + stream = inject_sse_keepalives( + stream, + get_settings().sse_keepalive_interval_seconds, + keepalive_frame=keepalive_frame, + ) + startup_handoff_streams.append(stream) + if startup_handoff_tasks: + stream = _guard_responses_startup_handoff( stream, - get_settings().sse_keepalive_interval_seconds, - keepalive_frame=keepalive_frame, - ), + startup_task=startup_handoff_tasks[0], + streams_to_close=tuple(startup_handoff_streams), + reservation_cleanup=reservation_cleanup, + responses_service_cleanup_ready_event=responses_service_cleanup_ready_event, + responses_owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + responses_owner_forward_rejected_event=responses_owner_forward_rejected_event, + ) + return StreamingResponse( + stream, media_type="text/event-stream", headers={ "Cache-Control": "no-cache, no-transform", @@ -5079,8 +5233,22 @@ async def _collect_responses( request_service_tier=payload.service_tier, request_usage_budget=estimate_api_key_request_usage(payload), ) + reservation_cleanup = _ResponsesReservationCleanup( + owns_reservation=True, + reservation=reservation, + scheduler=( + cast(_ResponsesCleanupScheduler, context.service) + if callable(getattr(context.service, "_schedule_cancel_safe_cleanup", None)) + else None + ), + request_id=ensure_request_id(), + ) - rate_limit_headers = await _rate_limit_headers_for_request(context, api_key) + rate_limit_headers = await _rate_limit_headers_after_reservation( + context, + api_key, + reservation_cleanup=reservation_cleanup, + ) 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 @@ -5118,13 +5286,39 @@ async def _collect_responses( client_ip=client_ip, ) captured_turn_state_headers: dict[str, str] = {} + responses_service_cleanup_ready_event = asyncio.Event() + responses_owner_forward_dispatched_event = asyncio.Event() + responses_owner_forward_rejected_event = asyncio.Event() + + responses_owner_forward_dispatched_token = _bind_propagated_responses_owner_forward_dispatched( + responses_owner_forward_dispatched_event + ) + responses_owner_forward_rejected_token = _bind_propagated_responses_owner_forward_rejected( + responses_owner_forward_rejected_event + ) + responses_cleanup_ready_token = _bind_propagated_responses_service_cleanup_ready( + responses_service_cleanup_ready_event + ) try: response_payload = await _collect_responses_payload( stream, captured_turn_state_headers=captured_turn_state_headers, ) + except asyncio.CancelledError: + if _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + await reservation_cleanup.release(action="responses collection cancellation") + raise except ProxyResponseError as exc: - await _release_reservation(reservation) + if _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + await reservation_cleanup.release(action="responses collection error") error = _parse_error_envelope(exc.payload) status_code, error = _mask_previous_response_not_found_error(error, default_status=exc.status_code) return _logged_error_json_response( @@ -5133,6 +5327,18 @@ async def _collect_responses( error.model_dump(mode="json", exclude_none=True), headers={**captured_turn_state_headers, **rate_limit_headers}, ) + except BaseException: + if _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + await reservation_cleanup.release(action="responses collection") + raise + finally: + _reset_propagated_responses_service_cleanup_ready(responses_cleanup_ready_token) + _reset_propagated_responses_owner_forward_rejected(responses_owner_forward_rejected_token) + _reset_propagated_responses_owner_forward_dispatched(responses_owner_forward_dispatched_token) if isinstance(response_payload, OpenAIResponsePayload): if response_payload.status == "failed": error_payload = _error_envelope_from_response(response_payload.error) @@ -5239,8 +5445,26 @@ async def _compact_responses( request_service_tier=_compact_request_service_tier(payload), request_usage_budget=request_usage_budget, ) + reservation_cleanup = _ResponsesReservationCleanup( + owns_reservation=True, + reservation=reservation, + scheduler=( + cast(_ResponsesCleanupScheduler, context.service) + if callable(getattr(context.service, "_schedule_cancel_safe_cleanup", None)) + else None + ), + request_id=ensure_request_id(), + ) + responses_service_cleanup_ready_event = asyncio.Event() - rate_limit_headers = await _rate_limit_headers_for_request(context, api_key) + rate_limit_headers = await _rate_limit_headers_after_reservation( + context, + api_key, + reservation_cleanup=reservation_cleanup, + ) + responses_cleanup_ready_token = _bind_propagated_responses_service_cleanup_ready( + responses_service_cleanup_ready_event + ) try: result = await context.service.compact_responses( payload, @@ -5275,7 +5499,9 @@ async def _compact_responses( headers=rate_limit_headers, ) finally: - await _release_reservation(reservation) + _reset_propagated_responses_service_cleanup_ready(responses_cleanup_ready_token) + if _responses_origin_may_release_reservation(service_cleanup_ready_event=responses_service_cleanup_ready_event): + await reservation_cleanup.release(action="compact response") result_payload = result.model_dump(mode="json", exclude_none=True) if codex_session_affinity: result_payload = _normalize_codex_remote_compaction_v2_result(result, result_payload) @@ -5386,6 +5612,17 @@ async def _synthetic_compaction_response_stream( yield "data: [DONE]\n\n" +async def _synthetic_compaction_failure_stream(*, response_id: str) -> AsyncIterator[str]: + yield format_sse_event( + response_failed_event( + "upstream_error", + "Compact response did not include a compaction output item", + response_id=response_id, + ) + ) + yield "data: [DONE]\n\n" + + async def _transcribe_request( *, request: Request, @@ -5677,8 +5914,9 @@ async def _wait_for_first_stream_probe( return_exceptions=True, ) except asyncio.CancelledError: - first_task.cancel() - await asyncio.gather(first_task, return_exceptions=True) + with anyio.CancelScope(shield=True): + first_task.cancel() + await asyncio.gather(first_task, return_exceptions=True) raise @@ -5689,6 +5927,8 @@ async def _probe_stream_startup_error( timeout_seconds: float | None = None, capacity_wait_event: asyncio.Event | None = None, capacity_ready_event: asyncio.Event | None = None, + handoff_task_sink: list[asyncio.Task[str]] | None = None, + service_cleanup_ready_event: asyncio.Event | None = None, ) -> tuple[AsyncIterator[str], ProxyResponseError | OpenAIErrorEnvelopeModel | None]: if timeout_seconds is None: timeout_seconds = _STREAM_STARTUP_ERROR_PROBE_SECONDS @@ -5699,12 +5939,78 @@ async def _probe_stream_startup_error( capacity_wait_event=capacity_wait_event, capacity_ready_event=capacity_ready_event, ) + if service_cleanup_ready_event is not None: + buffered_before_cleanup_ready: list[str] = [] + while not service_cleanup_ready_event.is_set(): + if not first_task.done(): + cleanup_ready_task = asyncio.create_task(service_cleanup_ready_event.wait()) + try: + await asyncio.wait( + {first_task, cleanup_ready_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + first_task.cancel() + cleanup_ready_task.cancel() + await asyncio.gather(first_task, cleanup_ready_task, return_exceptions=True) + raise + finally: + if not cleanup_ready_task.done(): + cleanup_ready_task.cancel() + await asyncio.gather(cleanup_ready_task, return_exceptions=True) + if service_cleanup_ready_event.is_set(): + break + try: + first = first_task.result() + except StopAsyncIteration: + return ( + _prepend_first(None, stream), + ProxyResponseError( + 502, + openai_error( + "stream_incomplete", + "Upstream stream ended before reservation cleanup handoff", + error_type="server_error", + ), + ), + ) + except ProxyResponseError as exc: + return _prepend_first(None, stream), exc + if convert_event_errors: + first_error = _stream_event_error_envelope(first) + if first_error is not None: + aclose = getattr(stream, "aclose", None) + if callable(aclose): + await aclose() + return _prepend_first(None, stream), first_error + # Capacity/keepalive frames can be produced before a bridge submit + # installs the receiver finalizer. Preserve them, but do not let + # one turn the HTTP 200 into a false cleanup acknowledgement. + buffered_before_cleanup_ready.append(first) + first_task = _create_first_stream_probe_task(stream) + + # The receiver has entered its settlement-guarded finalizer. Keep the + # current probe task in the response body so the successful 200 is the + # explicit cross-replica cleanup handoff; returning a startup error here + # would make the origin release concurrently. + if handoff_task_sink is not None: + handoff_task_sink.append(first_task) + return ( + _prepend_items( + buffered_before_cleanup_ready, + _prepend_first_task(first_task, stream), + ), + None, + ) if not probe_done: # Probe window elapsed before the first item arrived. Hand the still- # running task off to be consumed by the streamed response. asyncio.wait # (rather than wait_for + shield) never cancels the task on timeout, # avoiding the Python 3.14 "exception in shielded future" log when the # upstream later returns an error such as a 429 from the admission gate. + if handoff_task_sink is not None: + handoff_task_sink.append(first_task) return _prepend_first_task(first_task, stream), None try: first = first_task.result() @@ -6086,6 +6392,62 @@ async def _prepend_initial_sse_heartbeat( yield line +async def _guard_responses_startup_handoff( + stream: AsyncIterator[str], + *, + startup_task: asyncio.Task[str], + streams_to_close: tuple[AsyncIterator[str], ...], + reservation_cleanup: _ResponsesReservationCleanup, + responses_service_cleanup_ready_event: asyncio.Event, + responses_owner_forward_dispatched_event: asyncio.Event, + responses_owner_forward_rejected_event: asyncio.Event, +) -> AsyncIterator[str]: + try: + async for line in stream: + yield line + finally: + with anyio.CancelScope(shield=True): + release_candidate = False + if startup_task.done(): + release_candidate = startup_task.cancelled() or startup_task.exception() is not None + else: + release_candidate = True + startup_task.cancel() + await asyncio.gather(startup_task, return_exceptions=True) + closed_stream_ids: set[int] = set() + for stream_index, stream_to_close in enumerate(reversed(streams_to_close)): + stream_id = id(stream_to_close) + if stream_id in closed_stream_ids: + continue + closed_stream_ids.add(stream_id) + await _close_responses_stream_best_effort( + stream_to_close, + action=f"startup wrapper {stream_index}", + ) + if release_candidate and _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + await reservation_cleanup.release(action="responses startup handoff") + + +async def _close_responses_stream_best_effort( + stream: AsyncIterator[str], + *, + action: str, +) -> None: + aclose = getattr(stream, "aclose", None) + if not callable(aclose): + return + try: + await aclose() + except asyncio.CancelledError: + logger.debug("Responses %s stream close was cancelled", action) + except Exception: + logger.warning("Failed to close Responses %s stream", action, exc_info=True) + + async def _stream_proxy_errors_as_response_failed(stream: AsyncIterator[str]) -> AsyncIterator[str]: async for line in _stream_response_error_events(stream, owns_reservation=False, reservation=None): yield line @@ -6096,16 +6458,34 @@ async def _stream_response_error_events( *, owns_reservation: bool, reservation: ApiKeyUsageReservationData | None, + reservation_cleanup: _ResponsesReservationCleanup | None = None, + responses_service_cleanup_ready_event: asyncio.Event | None = None, + responses_owner_forward_dispatched_event: asyncio.Event | None = None, + responses_owner_forward_rejected_event: asyncio.Event | None = None, ) -> AsyncIterator[str]: + cleanup = reservation_cleanup or _ResponsesReservationCleanup( + owns_reservation=owns_reservation, + reservation=reservation, + scheduler=None, + request_id=ensure_request_id(), + ) + stream_completed = False + + async def release_owned_reservation() -> None: + if responses_service_cleanup_ready_event is not None and not _responses_origin_may_release_reservation( + service_cleanup_ready_event=responses_service_cleanup_ready_event, + owner_forward_dispatched_event=responses_owner_forward_dispatched_event, + owner_forward_rejected_event=responses_owner_forward_rejected_event, + ): + return + await cleanup.release(action="responses stream cleanup") + try: async for line in stream: yield line + stream_completed = True except ProxyResponseError as exc: - if owns_reservation: - try: - await _release_reservation(reservation) - except Exception: - logger.warning("Failed to release stream reservation after upstream proxy error", exc_info=True) + await release_owned_reservation() envelope = _parse_error_envelope(exc.payload) _, envelope = _mask_previous_response_not_found_error(envelope, default_status=exc.status_code) error = envelope.error @@ -6117,6 +6497,9 @@ async def _stream_response_error_events( error_param=error.param if error else None, ) ) + finally: + if not stream_completed: + await release_owned_reservation() def _stream_startup_error_response( @@ -6376,6 +6759,42 @@ async def _release_reservation(reservation: ApiKeyUsageReservationData | None) - await service.release_usage_reservation(reservation.reservation_id) +async def _release_reservation_best_effort( + reservation: ApiKeyUsageReservationData | None, + *, + action: str, + scheduler: _ResponsesCleanupScheduler | None, + request_id: str, +) -> None: + if reservation is None: + return + + async def release() -> None: + try: + await _release_reservation(reservation) + except Exception: + logger.warning("Failed to release API key reservation during %s", action, exc_info=True) + + if scheduler is None: + # Lightweight service doubles do not own the production background-task + # registry. Keep their synchronous cleanup cancellation-safe without + # manufacturing an untracked detached task. + with anyio.CancelScope(shield=True): + await release() + return + + task = scheduler._schedule_cancel_safe_cleanup( + release(), + action="release_stream_api_key_reservation", + request_id=request_id, + ) + # A raw asyncio task is outside the caller's AnyIO cancel scope and remains + # tracked for graceful shutdown. Shielding the await lets ordinary paths + # observe completion while cancellation propagates without cancelling the + # only persistence attempt. + await asyncio.shield(task) + + async def _finalize_image_reservation( reservation: ApiKeyUsageReservationData | None, *, diff --git a/app/modules/proxy/file_pin_repository.py b/app/modules/proxy/file_pin_repository.py new file mode 100644 index 0000000000..26750826ed --- /dev/null +++ b/app/modules/proxy/file_pin_repository.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +from collections.abc import Collection + +from sqlalchemy import Integer, bindparam, text +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.sql.elements import TextClause + +from app.db.session import sqlite_writer_section + +_TABLE = "file_account_pins" + +# Ownership TTLs stay entirely in the database clock domain. PostgreSQL's +# clock_timestamp() is evaluated when each clause executes. A successful claim +# is followed by a guarded refresh in the same transaction so even an INSERT +# that waited behind an ultimately rolled-back unique contender receives its +# full TTL after the wait. SQLite's padded strftime form matches SQLAlchemy +# DateTime's six-digit fractional width, preserving exact lexicographic expiry +# checks. +_POSTGRES_NOW = "clock_timestamp()" +_POSTGRES_NOW_PLUS_TTL = "clock_timestamp() + make_interval(secs => :ttl)" +_POSTGRES_STATEMENT_NOW = "statement_timestamp()" +_SQLITE_NOW = "(strftime('%Y-%m-%d %H:%M:%f', 'now') || '000')" +_SQLITE_NOW_PLUS_TTL = "(strftime('%Y-%m-%d %H:%M:%f', 'now', '+' || :ttl || ' seconds') || '000')" + +_POSTGRES_CLAIM = text( + f""" + INSERT INTO {_TABLE} (file_id, account_id, expires_at) + VALUES (:file_id, :account_id, {_POSTGRES_NOW_PLUS_TTL}) + ON CONFLICT (file_id) DO UPDATE SET + account_id = excluded.account_id, + expires_at = {_POSTGRES_NOW_PLUS_TTL} + WHERE {_TABLE}.account_id = :account_id + OR {_TABLE}.expires_at <= {_POSTGRES_NOW} + RETURNING account_id + """ +).bindparams(bindparam("ttl", type_=Integer)) + +_SQLITE_CLAIM = text( + f""" + INSERT INTO {_TABLE} (file_id, account_id, expires_at) + VALUES (:file_id, :account_id, {_SQLITE_NOW_PLUS_TTL}) + ON CONFLICT (file_id) DO UPDATE SET + account_id = excluded.account_id, + expires_at = {_SQLITE_NOW_PLUS_TTL} + WHERE {_TABLE}.account_id = :account_id + OR {_TABLE}.expires_at <= {_SQLITE_NOW} + RETURNING account_id + """ +).bindparams(bindparam("ttl", type_=Integer)) + +_POSTGRES_CLEANUP = text(f"DELETE FROM {_TABLE} WHERE expires_at <= {_POSTGRES_STATEMENT_NOW}") +_SQLITE_CLEANUP = text(f"DELETE FROM {_TABLE} WHERE expires_at <= {_SQLITE_NOW}") + +_POSTGRES_REFRESH = text( + f""" + UPDATE {_TABLE} + SET expires_at = {_POSTGRES_NOW_PLUS_TTL} + WHERE file_id = :file_id + AND account_id = :account_id + RETURNING account_id + """ +).bindparams(bindparam("ttl", type_=Integer)) +_SQLITE_REFRESH = text( + f""" + UPDATE {_TABLE} + SET expires_at = {_SQLITE_NOW_PLUS_TTL} + WHERE file_id = :file_id + AND account_id = :account_id + RETURNING account_id + """ +).bindparams(bindparam("ttl", type_=Integer)) + +_POSTGRES_GET_LIVE = text(f"SELECT account_id FROM {_TABLE} WHERE file_id = :file_id AND expires_at > {_POSTGRES_NOW}") +_SQLITE_GET_LIVE = text(f"SELECT account_id FROM {_TABLE} WHERE file_id = :file_id AND expires_at > {_SQLITE_NOW}") + +_POSTGRES_GET_LIVE_MANY = text( + f""" + SELECT file_id, account_id + FROM {_TABLE} + WHERE file_id IN :file_ids + AND expires_at > {_POSTGRES_NOW} + """ +).bindparams(bindparam("file_ids", expanding=True)) +_SQLITE_GET_LIVE_MANY = text( + f""" + SELECT file_id, account_id + FROM {_TABLE} + WHERE file_id IN :file_ids + AND expires_at > {_SQLITE_NOW} + """ +).bindparams(bindparam("file_ids", expanding=True)) + +_GET_ACCOUNT = text(f"SELECT account_id FROM {_TABLE} WHERE file_id = :file_id") + + +class FileAccountPinOwnershipConflict(RuntimeError): + def __init__(self, file_id: str, persisted_account_id: str, requested_account_id: str) -> None: + super().__init__( + f"Live file ownership conflict for {file_id!r}: " + f"persisted={persisted_account_id!r} requested={requested_account_id!r}" + ) + self.file_id = file_id + self.persisted_account_id = persisted_account_id + self.requested_account_id = requested_account_id + + +def build_file_account_pin_claim(*, dialect_name: str) -> TextClause: + if dialect_name == "postgresql": + return _POSTGRES_CLAIM + if dialect_name == "sqlite": + return _SQLITE_CLAIM + raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}") + + +def build_file_account_pin_cleanup(*, dialect_name: str) -> TextClause: + if dialect_name == "postgresql": + return _POSTGRES_CLEANUP + if dialect_name == "sqlite": + return _SQLITE_CLEANUP + raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}") + + +def build_file_account_pin_refresh(*, dialect_name: str) -> TextClause: + if dialect_name == "postgresql": + return _POSTGRES_REFRESH + if dialect_name == "sqlite": + return _SQLITE_REFRESH + raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}") + + +def build_file_account_pin_live_lookup(*, dialect_name: str, many: bool = False) -> TextClause: + if dialect_name == "postgresql": + return _POSTGRES_GET_LIVE_MANY if many else _POSTGRES_GET_LIVE + if dialect_name == "sqlite": + return _SQLITE_GET_LIVE_MANY if many else _SQLITE_GET_LIVE + raise RuntimeError(f"Unsupported database dialect for file account pins: {dialect_name}") + + +class FileAccountPinRepository: + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def claim(self, file_id: str, account_id: str, *, ttl_seconds: int) -> None: + if ttl_seconds <= 0: + raise ValueError("File account pin TTL must be positive") + dialect_name = self._dialect_name() + params = { + "file_id": file_id, + "account_id": account_id, + "ttl": ttl_seconds, + } + async with sqlite_writer_section(): + await self._session.execute(build_file_account_pin_cleanup(dialect_name=dialect_name)) + persisted_account_id = ( + await self._session.execute( + build_file_account_pin_claim(dialect_name=dialect_name), + params, + ) + ).scalar_one_or_none() + if persisted_account_id is None: + persisted_account_id = await self._session.scalar( + _GET_ACCOUNT, + {"file_id": file_id}, + ) + if persisted_account_id != account_id: + await self._session.rollback() + raise FileAccountPinOwnershipConflict( + file_id, + persisted_account_id or "", + account_id, + ) + refreshed_account_id = ( + await self._session.execute( + build_file_account_pin_refresh(dialect_name=dialect_name), + params, + ) + ).scalar_one_or_none() + if refreshed_account_id != account_id: + await self._session.rollback() + raise RuntimeError(f"Failed to refresh file account pin after claim: {file_id!r}") + await self._session.commit() + + async def get_live_account_id(self, file_id: str) -> str | None: + return await self._session.scalar( + build_file_account_pin_live_lookup(dialect_name=self._dialect_name()), + {"file_id": file_id}, + ) + + async def get_live_account_ids(self, file_ids: Collection[str]) -> dict[str, str]: + unique_file_ids = tuple(dict.fromkeys(file_ids)) + if not unique_file_ids: + return {} + rows = ( + ( + await self._session.execute( + build_file_account_pin_live_lookup( + dialect_name=self._dialect_name(), + many=True, + ), + {"file_ids": unique_file_ids}, + ) + ) + .tuples() + .all() + ) + return dict(rows) + + def _dialect_name(self) -> str: + return self._session.get_bind().dialect.name diff --git a/app/modules/proxy/http_bridge_forwarding.py b/app/modules/proxy/http_bridge_forwarding.py index 7f8969da47..c75ed461ac 100644 --- a/app/modules/proxy/http_bridge_forwarding.py +++ b/app/modules/proxy/http_bridge_forwarding.py @@ -124,6 +124,8 @@ async def stream_responses( context: HTTPBridgeForwardContext, request_started_at: float, on_response_wait: Callable[[], None] | None = None, + on_request_dispatched: Callable[[], None] | None = None, + on_response_rejected: Callable[[], None] | None = None, on_response_ready: Callable[[], None] | None = None, ) -> AsyncIterator[str]: settings = get_settings() @@ -134,13 +136,26 @@ async def stream_responses( if on_response_wait is not None: on_response_wait() async with aiohttp.ClientSession(timeout=timeout, trust_env=False) as session: - async with session.post( - f"{owner_endpoint}{HTTP_BRIDGE_INTERNAL_FORWARD_PATH}", - json=payload.model_dump_for_forwarding(), - headers=build_owner_forward_headers(headers=headers, payload=payload, context=context), + request_url = f"{owner_endpoint}{HTTP_BRIDGE_INTERNAL_FORWARD_PATH}" + request_payload = payload.model_dump_for_forwarding() + request_headers = build_owner_forward_headers(headers=headers, payload=payload, context=context) + request_context = session.post( + request_url, + json=request_payload, + headers=request_headers, skip_auto_headers=_OWNER_FORWARD_SKIP_AUTO_HEADERS, - ) as response: + ) + if on_request_dispatched is not None: + # Payload/header construction is still a local pre-dispatch + # failure. Once the request context starts, a transport failure + # cannot prove whether the owner received the signed reservation. + on_request_dispatched() + async with request_context as response: if response.status != 200: + if on_response_rejected is not None: + # The receiver contract never transfers cleanup on a + # non-200 response, so the origin may safely release. + on_response_rejected() payload_text = await response.text() raise ProxyResponseError( response.status, diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index 3b6e75b4f1..9dd9a9422b 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -546,7 +546,6 @@ _clear_websocket_request_error_overrides, # noqa: F401 _DownstreamWebSocketActivity, # noqa: F401 _event_type_from_payload, # noqa: F401 - _FilePinEntry, _finalize_ttft_reasoning_deltas, # noqa: F401 _http_error_status_from_payload, # noqa: F401 _HTTPBridgeSession, @@ -945,15 +944,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() - # 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 - # routing is best-effort: if the finalize request lands on a - # different replica with no pin, we fall back to a fresh load- - # balancer selection. The TTL is short enough (5 min) that we - # never hold stale pins after the upstream upload window closes. - self._file_account_pins: dict[str, _FilePinEntry] = {} - self._file_account_pin_lock = asyncio.Lock() + self._file_pin_session_factory = SessionLocal self._http_bridge_lock = anyio.Lock() self._work_admission: WorkAdmissionController | None = None self._request_log_tasks: set[asyncio.Task[None]] = set() diff --git a/openspec/changes/persist-file-account-pins/.openspec.yaml b/openspec/changes/persist-file-account-pins/.openspec.yaml new file mode 100644 index 0000000000..e8209ffaac --- /dev/null +++ b/openspec/changes/persist-file-account-pins/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-28 diff --git a/openspec/changes/persist-file-account-pins/design.md b/openspec/changes/persist-file-account-pins/design.md new file mode 100644 index 0000000000..2f0fc15738 --- /dev/null +++ b/openspec/changes/persist-file-account-pins/design.md @@ -0,0 +1,45 @@ +## Context + +`ProxyService` records upstream-issued file ownership in a process-local dictionary. File finalize and Responses input-file routing already consume ownership through `_pin_file_account`, `_resolve_file_account`, and `_lookup_file_pin`, but another replica cannot observe that state. The application database and `SessionLocal` are already the shared coordination substrate. + +## Goals / Non-Goals + +**Goals:** + +- Make live file ownership visible to every replica. +- Preserve the existing 30-minute expiry and opaque unknown-file compatibility. +- Keep account-owner routing fail-closed and the change behind existing service boundaries. +- Support the repository's PostgreSQL production path and SQLite test/development path. + +**Non-Goals:** + +- Change upload/finalize API payloads, routing precedence, or retry policy. +- Add operator configuration or background cleanup infrastructure. +- Backfill pins that existed only in memory before migration. + +## Decisions + +1. Add a `file_account_pins` table keyed by `file_id`, with an account identifier and absolute UTC expiry. A dedicated repository owns idempotent ownership claims and live lookup. A live claim is immutable across accounts; same-owner claims renew it, while an expired ID can be claimed again. This is smaller and more explicit than overloading sticky-session namespaces. +2. Use one short-lived `SessionLocal` session per pin/read operation. `ProxyService` is process-scoped, so retaining a request-scoped session would be unsafe; the established durable bridge/ring pattern already uses a session factory. +3. Do not cache durable owner decisions in process memory. `_pin_file_account` writes through the repository, while every `_resolve_file_account` call reads the shared database. Multi-file resolution uses one database query so all referenced IDs are classified from the same repository operation. An authenticated inter-replica forwarding value only corroborates the receiver's fresh database result; it cannot replace that read. +4. Evaluate expiry inside each database statement. PostgreSQL claim, reclaim, and live lookup use `clock_timestamp()`, and every successful claim performs an owner-guarded expiry refresh in the same transaction. The refresh gives a full post-wait TTL even when a new-row insert blocked behind an uncommitted unique contender that later rolled back. PostgreSQL cleanup uses the DB-authoritative, stable `statement_timestamp()` cutoff so the expiry index remains usable. SQLite uses its statement-native UTC clock with the same fractional width as stored `DateTime` values and the same guarded post-claim refresh. All expiry decisions therefore stay in the database clock domain without replica-clock skew or exact-expiry ambiguity. +5. Translate persistence failures at the ownership boundary into the stable fail-closed proxy error. Run finalize lookup and post-upstream pin persistence inside the existing file request-log lifecycle, and keep Responses lookup errors inside the existing startup-error lifecycle, so failures neither leak API-key reservations nor record a failed request as successful. +6. Keep exactly one owner for a Responses usage reservation across API startup, direct or compact service settlement, and HTTP-bridge forwarding. Within one replica, the API layer owns cleanup through the durable file-owner lookup and any following preflight outside the service settlement guard. The direct stream service signals when it enters its settlement-guarded `try/finally`; the local HTTP-bridge service signals only after a successful request submit installs the request-state finalizer; and compact service settlement signals after its one cancellation-safe settlement attempt. From those exact boundaries the service finalizer or settlement attempt owns cleanup even if no upstream event has arrived. For an authenticated cross-replica forward carrying the origin reservation, the receiver delays its successful HTTP 200 until its own service has reached one of those settlement-owned boundaries. The 200 response is the receiver's cleanup-handoff acknowledgement. The origin records dispatch only after local payload and header construction and immediately before entering the request transport. A definitive non-200 response leaves cleanup at the origin. If dispatch occurred but no HTTP status can be observed, the origin must not actively release or replay because the receiver may already own settlement; the receiver finalizer or the existing stale-reservation reaper resolves the ambiguity. The reaper releases reservations older than six hours on its hourly leader loop, so this fail-closed choice can hold quota for up to roughly seven hours but prevents a premature origin release from discarding receiver-recorded usage. A receiver-side owner-revalidation failure or cancellation before dispatch or after a definitive non-200 propagates with origin cleanup intact. An initial client-facing SSE heartbeat or any other frame does not transfer ownership. Cancellation or owner-lookup failure while cleanup remains at the API layer schedules one tracked release, including when a bounded startup probe has handed pending preflight work to the response body; the origin cancels and awaits that pending startup task before releasing. Each cleanup owner makes one cancellation-safe, best-effort attempt without allowing a cleanup-database failure to mask the original `file_owner_unavailable` error or cancellation. + Once compact settlement has transferred ownership, later receiver-side output validation cannot safely turn the response back into a non-200 rejection. The receiver therefore preserves HTTP 200 and emits a terminal `response.failed` SSE event if the settled compact payload lacks a valid compaction output item; the origin keeps the handoff and does not release or replay. + +## Risks / Trade-offs + +- [Every hard ownership decision adds a database read] → use one indexed lookup for one file and one batched indexed lookup for multi-file requests; correctness across replicas takes precedence over process-local locality. +- [A database outage can make file create/finalize unavailable] → fail closed because silently selecting another account can disclose or corrupt account-scoped operations. +- [An idle installation can retain expired rows until the next upload] → the rows are inert and every new claim performs indexed opportunistic cleanup. +- [Migration overlaps another branch] → base the revision on the current single head and report any later head conflict rather than editing another track. + +## Migration Plan + +Upgrade creates the empty ownership table and index; new uploads populate it immediately. Downgrade drops only the new table. Existing in-memory pins cannot be backfilled. + +The behavior change is not safe under an ordinary mixed-version rolling rollout: a legacy replica cannot read pins written by a new replica, and a new replica cannot read a legacy replica's process-local pins. Operators must migrate the database first, stop legacy replicas from accepting new file registrations, drain the legacy upload/finalize window for up to the 30-minute pin TTL (or explicitly accept retrying those in-flight uploads), and then cut all file-serving replicas over without mixed-version file traffic. Deployment automation is intentionally outside this code change. + +## Open Questions + +None. diff --git a/openspec/changes/persist-file-account-pins/proposal.md b/openspec/changes/persist-file-account-pins/proposal.md new file mode 100644 index 0000000000..98e8c3d7da --- /dev/null +++ b/openspec/changes/persist-file-account-pins/proposal.md @@ -0,0 +1,27 @@ +## Why + +File ownership pins currently live only in a replica-local dictionary, so a file finalize or Responses request handled by another replica can select an account that does not own the upstream file. The confirmed P1 bug must be fixed by making the existing ownership boundary durable and shared. + +## What Changes + +- Persist live `file_id -> account_id` ownership pins in the application database with the existing 30-minute lifetime. +- Resolve file ownership through the durable store so finalize and input-file routing remain account-bound across replicas. +- Keep the existing `_pin_file_account` and `_resolve_file_account` service boundaries and fail closed when durable ownership cannot be established safely. +- Add migration and targeted repository/service regression coverage. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `responses-api-compat`: Replace process-local best-effort file ownership with durable, replica-shared ownership for file finalize and Responses input-file routing. +- `files-upload-protocol`: Require file finalization to resolve its durable owner through the shared database and fail closed when that lookup is unavailable. +- `sticky-session-operations`: Require a remote HTTP-bridge owner to corroborate signed file-owner metadata with its own fresh durable lookup instead of trusting origin process memory. +- `sticky-session-operations`: Require every receiving HTTP-bridge transport, including terminal compaction, to revalidate forwarded file ownership against the durable database. + +## Impact + +The change affects the proxy file operations mixin, a small proxy persistence repository, the database model and Alembic graph, and focused proxy/database tests. It adds no dependency or public API surface. diff --git a/openspec/changes/persist-file-account-pins/specs/files-upload-protocol/spec.md b/openspec/changes/persist-file-account-pins/specs/files-upload-protocol/spec.md new file mode 100644 index 0000000000..500215080c --- /dev/null +++ b/openspec/changes/persist-file-account-pins/specs/files-upload-protocol/spec.md @@ -0,0 +1,25 @@ +## ADDED Requirements + +### Requirement: File finalize uses durable replica-shared ownership + +When `POST /backend-api/files/{file_id}/uploaded` references a live durable file pin, the service MUST resolve that pin from the shared database and route finalization only through the owning account. The owner decision MUST NOT use a process-local cache. Expiry, reclaim, and cleanup MUST use database-authoritative time. If durable owner resolution fails, the service MUST fail closed before selecting or invoking an unpinned fallback account. + +#### Scenario: another replica finalizes through the durable owner + +- **GIVEN** one replica registered `file_xyz` through `account_a` +- **WHEN** another replica handles `POST /backend-api/files/file_xyz/uploaded` +- **THEN** it MUST resolve the shared durable pin +- **AND** it MUST finalize only through `account_a` + +#### Scenario: finalize owner lookup failure does not fall back + +- **GIVEN** `file_xyz` requires a durable owner decision +- **WHEN** the shared database lookup fails +- **THEN** finalization MUST fail before any unpinned account selection or upstream invocation + +#### Scenario: an expired identifier can be reclaimed using database time + +- **GIVEN** the durable pin for `file_xyz` has expired according to the database clock +- **WHEN** a later upload claims `file_xyz` through `account_b` +- **THEN** the durable owner MUST become `account_b` +- **AND** every replica's next finalize decision MUST resolve `account_b` diff --git a/openspec/changes/persist-file-account-pins/specs/responses-api-compat/spec.md b/openspec/changes/persist-file-account-pins/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..be36f0c2d6 --- /dev/null +++ b/openspec/changes/persist-file-account-pins/specs/responses-api-compat/spec.md @@ -0,0 +1,124 @@ +## MODIFIED Requirements + +### Requirement: Responses requests with input_file.file_id route to the upload's account + +A `/v1/responses`, `/backend-api/codex/responses`, or `/responses/compact` request that references an `{type: "input_file", file_id}` content item SHALL be routed to the upstream account that registered the file via `POST /backend-api/files` when a durable, unexpired pin for that `file_id` exists. The pin MUST be visible to every replica that shares the application database. A live file pin is hard ownership evidence: it MUST override prompt-cache or bare process-session locality and MUST agree with independently resolved turn-state, previous-response, bridge, or other hard ownership. + +When multiple `file_id`s are referenced, all live pins MUST resolve to the same account. If at least one ID has a live pin and another ID has no live pin, the request MUST fail with `file_owner_unavailable`; if live pins resolve to different accounts, it MUST fail with `continuity_owner_conflict`. If none of the referenced IDs has a live pin, the proxy MUST preserve compatibility with files registered directly upstream or before durable ownership was observed by forwarding the opaque IDs verbatim under ordinary unpinned routing. + +A live durable pin MUST NOT be reassigned to another account. Repeating the claim for the same account MUST be idempotent and MAY renew its expiry; an expired identifier MAY be claimed by a later upload. + +Every hard file-owner decision MUST read the shared database and MUST NOT rely on a process-local owner cache. Authenticated inter-replica forwarding metadata MAY corroborate the freshly resolved durable owner but MUST NOT replace the receiver's database read. A missing or conflicting receiver-side durable owner MUST fail closed before account selection or upstream invocation. Pin expiry, reclaim, and cleanup MUST use database-authoritative statement time rather than a replica's application clock. + +For a streaming Responses request whose durable file-owner lookup runs in the stream service, any API-key usage reservation acquired before that lookup MUST have exactly one cleanup owner if resolution fails or the request is cancelled. Within one replica, the API layer MUST own cleanup until the direct stream service enters its settlement-guarded `try/finally` or the local HTTP-bridge service successfully submits the request and installs its request-state finalizer. The service finalizer MUST own cleanup after that explicit boundary so those layers cannot both release the reservation. Merely completing the durable lookup MUST NOT transfer cleanup before a service finalizer is active, and an initial SSE heartbeat MUST NOT transfer ownership to the client. + +When an authenticated HTTP-bridge origin forwards that reservation to another replica, the receiver MUST delay its successful HTTP 200 response until its service finalizer is active. That 200 response MUST be the cleanup-handoff acknowledgement that transfers ownership from the origin to the receiver. The origin MUST distinguish a request that has not been dispatched, a dispatch with no observed response status, a successful HTTP 200 acknowledgement, and a definitive non-200 rejection. Before dispatch or after a definitive non-200, receiver-side owner-revalidation failure or cancellation MUST propagate with cleanup remaining at the origin. After dispatch when no response status can be observed, the origin MUST NOT actively release or replay the reservation because the receiver may already own settlement; receiver settlement or bounded stale-reservation cleanup MUST resolve that ambiguity. After the acknowledgement, the receiver service finalizer MUST remain authoritative even if no upstream event has arrived. If a bounded startup probe hands pending preflight work to the response body and the body closes first, the active owner MUST cancel and await that work before scheduling one cancellation-safe release attempt. An SSE heartbeat or another frame MUST NOT transfer cleanup ownership. Compact service settlement MUST likewise suppress a second API-layer release after its single settlement attempt. Once a forwarded compact service has made that settlement attempt, a later receiver-side output validation failure MUST preserve HTTP 200 as the cleanup-handoff acknowledgement and surface a terminal `response.failed` event; it MUST NOT become a non-200 rejection that permits origin release or replay. A cleanup-store failure MUST NOT mask the original stable owner error or cancellation. Owner-lookup failure or cancellation MUST NOT trigger account failover or another upstream attempt. + +#### Scenario: file_id pin drives routing for an input_file response + +- **GIVEN** a `POST /backend-api/files` registered `file_xyz` through `account_a` on one replica +- **WHEN** a `/v1/responses` request references `{"type": "input_file", "file_id": "file_xyz"}` on another replica +- **THEN** the proxy MUST route the request to `account_a` + +#### Scenario: file_id pin overrides prompt-cache locality + +- **GIVEN** a pinned `file_xyz -> account_a` +- **WHEN** a `/v1/responses` request references `file_xyz` AND sets an explicit `prompt_cache_key` +- **THEN** the proxy MUST route to `account_a` and MUST NOT send the account-scoped file to the prompt-cache account + +#### Scenario: opaque file_id without a live pin remains compatible + +- **GIVEN** a request references a `file_id` registered directly upstream or before the system durably observed its upload +- **AND** no referenced file has a live durable pin +- **WHEN** the request is routed +- **THEN** the proxy MUST forward the `file_id` verbatim under ordinary unpinned routing +- **AND** it MUST NOT reject the request solely because owner metadata is absent + +#### Scenario: file finalize resolves ownership across replicas + +- **GIVEN** one replica registered `file_xyz` through `account_a` +- **WHEN** another replica handles `POST /backend-api/files/file_xyz/uploaded` +- **THEN** the proxy MUST finalize the file through `account_a` +- **AND** it MUST NOT fall back to a different eligible account + +#### Scenario: concurrent live ownership claims do not overwrite + +- **GIVEN** `file_xyz` has a live durable pin to `account_a` +- **WHEN** another replica attempts to pin `file_xyz` to `account_b` +- **THEN** the claim MUST fail with `continuity_owner_conflict` +- **AND** subsequent routing MUST still resolve `file_xyz` to `account_a` + +#### Scenario: a replica observes an expired pin reclaimed by another replica + +- **GIVEN** a replica previously resolved `file_xyz` to `account_a` +- **AND** the durable pin expires and another replica claims `file_xyz` for `account_b` +- **WHEN** the first replica resolves `file_xyz` again +- **THEN** it MUST read the durable owner and return `account_b` +- **AND** it MUST NOT return `account_a` from process-local state + +#### Scenario: durable owner lookup failure fails closed + +- **GIVEN** a request references a file whose owner decision requires the shared database +- **WHEN** the durable owner lookup fails +- **THEN** the request MUST fail before selecting or invoking an unpinned fallback account + +#### Scenario: cancellation during owner lookup releases admission state + +- **GIVEN** a request has acquired an API-key usage reservation before durable file-owner resolution completes +- **WHEN** the request is cancelled while the owner lookup is pending +- **THEN** exactly one cleanup owner MUST attempt to release or settle the reservation +- **AND** no account selection, upstream invocation, retry, or failover may occur + +#### Scenario: delayed owner failure after stream handoff releases admission state + +- **GIVEN** the streaming startup probe expires while durable file-owner resolution is still pending +- **WHEN** the lookup later fails or the response body is closed +- **THEN** the origin API MUST cancel and await any still-pending lookup +- **AND** the origin API MUST make exactly one release attempt +- **AND** a lookup failure MUST be represented by the stable `file_owner_unavailable` error + +#### Scenario: forwarded owner metadata is revalidated against durable ownership + +- **GIVEN** a replica receives authenticated forwarding metadata that identifies `account_a` as a referenced file's owner +- **WHEN** the receiver's fresh durable lookup has no live owner or identifies a different owner +- **THEN** the receiver MUST fail closed +- **AND** it MUST NOT route using the forwarded value alone +- **AND** it MUST propagate the preflight failure to the origin without releasing the origin reservation +- **AND** the originating request path MUST remain the sole cleanup owner because no successful handoff acknowledgement was sent + +#### Scenario: forwarded stream acknowledges cleanup ownership before HTTP 200 + +- **GIVEN** the origin forwards a file-pinned streaming request and its API-key reservation to the authenticated owner replica +- **WHEN** the receiver completes durable owner revalidation and installs its service settlement finalizer +- **THEN** the receiver MAY return HTTP 200 as the cleanup-handoff acknowledgement +- **AND** the origin MUST stop releasing the reservation after receiving that acknowledgement +- **AND** cancellation before the first upstream event MUST invoke only the receiver's service finalizer + +#### Scenario: ambiguous owner dispatch defers active origin cleanup + +- **GIVEN** the origin has begun dispatching a signed forwarded request carrying its reservation +- **WHEN** the transport fails before the origin can observe an HTTP status +- **THEN** the origin MUST NOT actively release or replay the reservation +- **AND** receiver settlement or stale-reservation cleanup MUST remain the only recovery paths + +#### Scenario: definitive owner rejection retains origin cleanup + +- **GIVEN** the origin dispatches a signed forwarded request carrying its reservation +- **WHEN** the receiver returns a non-200 response without acknowledging cleanup handoff +- **THEN** the origin MUST make exactly one cancellation-safe release attempt +- **AND** the receiver MUST NOT settle the origin reservation + +#### Scenario: compact service settlement is not released twice + +- **GIVEN** terminal or direct compaction receives an API-key usage reservation +- **WHEN** the compact service makes its single settlement or release attempt +- **THEN** the API layer MUST NOT issue another release for that reservation +- **AND** a pre-service failure MUST still leave exactly one release attempt at the API layer + +#### Scenario: malformed compact output after settlement preserves handoff + +- **GIVEN** a forwarded terminal compact request whose receiver service has made its single settlement attempt +- **WHEN** the settled response lacks a valid compaction output item +- **THEN** the receiver MUST return HTTP 200 as the cleanup-handoff acknowledgement +- **AND** it MUST emit a terminal `response.failed` event +- **AND** the origin MUST NOT release or replay the reservation diff --git a/openspec/changes/persist-file-account-pins/specs/sticky-session-operations/spec.md b/openspec/changes/persist-file-account-pins/specs/sticky-session-operations/spec.md new file mode 100644 index 0000000000..9d0ce463f8 --- /dev/null +++ b/openspec/changes/persist-file-account-pins/specs/sticky-session-operations/spec.md @@ -0,0 +1,145 @@ +## MODIFIED Requirements + +### Requirement: Hard continuity remains owner-bound and bounded + +Requests that depend on `previous_response_id`, hard turn-state, nonblank `conversation`, account-scoped `input_file.file_id` pins, live or durable bridge ownership, replay/reattach state, or another required owner continuity source MUST NOT silently reroute to an account that cannot preserve continuity. A resolved required owner MUST override bare process-session locality and MUST be selected without consulting or rewriting that soft mapping. A `previous_response_id` is a stored-object continuation reference and remains owner-bound even when the same request also carries a session header, `prompt_cache_key`, or another soft locality key. If independently resolved hard sources identify different accounts, if live durable referenced-file pins identify different accounts, or if a request has partial live durable file-pin coverage, the service MUST fail closed before upstream dispatch. A request for which no referenced file has a live durable pin MUST preserve opaque `file_id` compatibility and proceed without inventing ownership evidence. If the owner account/session is unavailable or saturated, the service MUST fail closed with an explicit retryable continuity/local overload reason instead of flooding the owner queue indefinitely. + +Every HTTP, compact, direct WebSocket, and HTTP-bridge transport MUST resolve explicit turn state against both live and durable bridge aliases. Live, durable, previous-response, file, and explicit turn-state evidence MUST be compared independently; source ordering MUST NOT choose the first match when distinct sessions or accounts resolve. A reused direct WebSocket MUST repeat nonblank `conversation` ownership validation for each response-create frame because the existing socket account proves only the current route. Single-account routing MUST constrain effective routing without narrowing the ownership-candidate pool used by that validation. + +When an HTTP-bridge owner is on another replica, the origin MUST forward its resolved durable file owner in authenticated full-context metadata. The receiving owner MUST perform its own fresh shared-database lookup and MUST require that durable result to match the forwarded owner. A missing or conflicting receiver-side durable owner MUST fail closed before account selection or upstream invocation. A retired direct WebSocket's upstream turn-state token MUST NOT be sent to a different account selected for a later movable bare-session request. + +A nonblank `conversation` without a dedicated resolved owner MUST proceed only when an explicit hard Codex mapping proves ownership or exactly one account remains in the model/API-key/security-scoped selection pool before transient additional-quota availability, retry exclusions, runtime health, budget, or account-cap filtering. A temporarily quota-filtered, excluded, unhealthy, or capped candidate MUST remain part of this ambiguity check because it may be the actual owner. A bare process-session mapping MUST NOT prove conversation ownership. + +#### Scenario: Previous-response owner queue is saturated + +- **WHEN** a `/v1/responses` follow-up requires a previous-response owner +- **AND** the owner session queue or account cap is saturated +- **THEN** the service fails closed with `hard_affinity_saturated`, `previous_response_owner_unavailable`, or the applicable stable `account_stream_cap` / `account_response_create_cap` code +- **AND** it does not route to an unrelated account that lacks continuity state + +#### Scenario: File-pinned request owner is capped + +- **WHEN** a `/v1/responses` request references an `input_file.file_id` pinned to an owner account +- **AND** the owner account is at its account stream or response-create cap +- **THEN** the service returns a local account-cap overload for the owner +- **AND** it does not route the file reference to another account + +#### Scenario: File-pinned request owner overrides process-session locality + +- **GIVEN** a request carries a bare process-session header mapped to account A +- **AND** its `input_file.file_id` is durably pinned to account B +- **WHEN** the request is routed +- **THEN** account B is treated as the required owner +- **AND** the process-session mapping is neither consulted as an owner nor rewritten + +#### Scenario: Conflicting hard owners fail closed + +- **GIVEN** a turn state, previous response, bridge, or input file resolves to account A +- **AND** another hard source on the same request resolves to account B +- **WHEN** the request is routed +- **THEN** the service fails with `continuity_owner_conflict` before upstream dispatch +- **AND** source ordering does not choose either owner + +#### Scenario: Partial or cross-account file pins fail closed + +- **GIVEN** a request references multiple account-scoped input files +- **AND** at least one file has a live durable owner pin +- **AND** another file has no live durable owner pin or the live pins resolve to different accounts +- **WHEN** the request is routed +- **THEN** the service fails with `file_owner_unavailable` or `continuity_owner_conflict` +- **AND** it does not route the files using a soft affinity account + +#### Scenario: Opaque file IDs with no live durable pins preserve compatibility + +- **GIVEN** a request references one or more `input_file.file_id` values +- **AND** none of those IDs has a live durable owner pin +- **WHEN** the request is routed +- **THEN** the service forwards the opaque file references under ordinary unpinned routing +- **AND** it does not invent a hard owner or fail solely because durable pin metadata is absent + +#### Scenario: Ambiguous conversation fails closed + +- **GIVEN** a request carries nonblank `conversation` continuity and only bare process-session affinity +- **AND** more than one account is eligible +- **WHEN** no dedicated or hard-mapping owner can be resolved +- **THEN** the request fails with a stable owner-unavailable error before upstream dispatch + +#### Scenario: Account-cap pressure does not manufacture a conversation owner + +- **GIVEN** two accounts remain in the model/API-key/security-scoped selection pool +- **AND** one account is temporarily at its local account cap +- **WHEN** a request carries nonblank `conversation` continuity without a dedicated or hard-mapping owner +- **THEN** the request still fails with a stable owner-unavailable error +- **AND** the uncapped account is not treated as the unique owner + +#### Scenario: Retry or additional-quota filtering does not manufacture a conversation owner + +- **GIVEN** two accounts remain in the model/API-key/security-scoped selection pool +- **AND** retry exclusion or transient additional-quota availability removes one from the effective routing pool +- **WHEN** a request carries nonblank `conversation` continuity without a dedicated or hard-mapping owner +- **THEN** the request still fails with a stable owner-unavailable error +- **AND** the remaining effective account is not treated as the unique owner + +#### Scenario: Account status does not manufacture a conversation owner + +- **GIVEN** two accounts are in the model/API-key/security ownership pool +- **AND** one account is paused, requires reauthentication, deactivated, or otherwise unavailable for routing +- **WHEN** a request carries nonblank `conversation` continuity without a dedicated or hard-mapping owner +- **THEN** the request still fails with a stable owner-unavailable error +- **AND** the active account is not treated as the unique owner + +#### Scenario: Preferred file owner does not manufacture a conversation owner + +- **GIVEN** a request carries nonblank `conversation` continuity and a file durably pinned to account B +- **AND** another account remains in the model/API-key/security ownership pool +- **WHEN** no dedicated conversation owner can be resolved +- **THEN** file ownership does not narrow the conversation ambiguity check to account B +- **AND** the request fails closed before upstream dispatch + +#### Scenario: Bridge turn state is owner-bound across transports + +- **GIVEN** an HTTP bridge registered a turn-state alias for account A +- **WHEN** the alias is reused through compact, plain HTTP streaming, or direct WebSocket transport +- **THEN** each transport treats account A as the required owner +- **AND** it does not fall back to unrelated sticky affinity + +#### Scenario: Independent bridge aliases conflict + +- **GIVEN** a live or durable turn-state alias resolves to one bridge session +- **AND** a previous-response alias on the same request resolves to a distinct session or account +- **WHEN** the request is routed +- **THEN** the service fails with `continuity_owner_conflict` +- **AND** alias lookup order does not select either session + +#### Scenario: Reused WebSocket revalidates conversation ownership + +- **GIVEN** a direct upstream WebSocket is already open on account A +- **AND** a later response-create frame carries nonblank `conversation` +- **WHEN** more than one account remains in the ownership-candidate pool +- **THEN** the later frame fails with a stable owner-unavailable error before upstream send +- **AND** the existing socket account is not treated as ownership proof + +#### Scenario: Single-account routing does not manufacture conversation ownership + +- **GIVEN** single-account routing selects account A +- **AND** multiple accounts remain in the model/API-key/security ownership pool +- **WHEN** a request carries nonblank `conversation` without dedicated owner evidence +- **THEN** the request remains ambiguous and fails closed +- **AND** only the effective routing states are constrained to account A + +#### Scenario: Remote bridge owner revalidates forwarded file ownership + +- **GIVEN** origin replica A durably resolves an input file to account A +- **AND** the request's HTTP bridge owner runs on replica B +- **WHEN** replica A forwards the request to replica B with authenticated file-owner metadata +- **THEN** replica B MUST freshly resolve the shared durable pin +- **AND** it MUST accept the forwarded owner only when both owner values match +- **AND** a missing, conflicting, tampered, or legacy-unbound proof MUST be rejected before upstream invocation + +#### Scenario: Retired WebSocket turn state does not cross accounts + +- **GIVEN** a closed upstream WebSocket on account A supplied an account-scoped turn-state token +- **AND** a later movable bare-session frame selects account B +- **WHEN** the proxy opens the replacement WebSocket +- **THEN** it removes account A's stale turn-state token before connect +- **AND** account B never receives that token diff --git a/openspec/changes/persist-file-account-pins/tasks.md b/openspec/changes/persist-file-account-pins/tasks.md new file mode 100644 index 0000000000..0ae0b2f25d --- /dev/null +++ b/openspec/changes/persist-file-account-pins/tasks.md @@ -0,0 +1,22 @@ +## 1. Durable ownership storage + +- [x] 1.1 Add the file-account pin ORM model and forward Alembic migration on the current head. +- [x] 1.2 Implement a focused repository for durable upsert and unexpired owner lookup. + +## 2. Proxy integration + +- [x] 2.1 Wire the repository through the existing file pin/resolve boundaries. +- [x] 2.2 Make multi-file ownership resolution use the durable lookup boundary and preserve fail-closed conflict behavior. + +## 3. Verification + +- [x] 3.1 Add targeted repository and cross-replica service regression tests, including expiry and multi-file behavior. +- [x] 3.2 Run focused tests, migration checks, OpenSpec validation, and inspect the final diff/status. + +## 4. DB-authoritative ownership repair + +- [x] 4.1 Remove process-local caching from hard file-owner decisions and batch multi-file resolution through the repository. +- [x] 4.2 Use database-authoritative time for claim expiry, reclaim, live lookup, and opportunistic cleanup. +- [x] 4.3 Add the file-finalize ownership contract to `files-upload-protocol` and replace the stale process-local forwarding contract in `sticky-session-operations`. +- [x] 4.4 Add hermetic race and fail-closed regression coverage, then rerun focused validation. +- [x] 4.5 Make compact settlement and ambiguous owner-forward dispatch single-owner, and add stream/collect/non-200/lost-status regression coverage. diff --git a/tests/integration/test_migrations.py b/tests/integration/test_migrations.py index 2d82cf4d01..2eb3a0ef30 100644 --- a/tests/integration/test_migrations.py +++ b/tests/integration/test_migrations.py @@ -709,6 +709,58 @@ async def test_run_startup_migrations_drops_accounts_email_unique_with_non_casca await engine.dispose() +@pytest.mark.asyncio +async def test_file_account_pins_migration_upgrade_and_downgrade(tmp_path): + from alembic import command + from sqlalchemy import inspect as sa_inspect + + from app.db.migrate import _build_alembic_config + + db_url = f"sqlite+aiosqlite:///{tmp_path / 'file-account-pins.sqlite'}" + parent_revision = "20260725_000000_add_http_bridge_pending_tool_calls" + pin_revision = "20260728_000000_add_file_account_pins" + + def _schema_state(sync_conn): + inspector = sa_inspect(sync_conn) + if not inspector.has_table("file_account_pins"): + return None + columns = inspector.get_columns("file_account_pins") + file_id_column = next(column for column in columns if column["name"] == "file_id") + return { + "columns": {column["name"] for column in columns}, + "file_id_length": file_id_column["type"].length, + "primary_key": inspector.get_pk_constraint("file_account_pins")["constrained_columns"], + "indexes": {index["name"] for index in inspector.get_indexes("file_account_pins")}, + } + + await to_thread.run_sync(lambda: run_upgrade(db_url, parent_revision, bootstrap_legacy=False)) + engine = create_async_engine(db_url, future=True) + try: + async with engine.connect() as conn: + assert await conn.run_sync(_schema_state) is None + + await to_thread.run_sync(lambda: run_upgrade(db_url, pin_revision, bootstrap_legacy=False)) + async with engine.connect() as conn: + state = await conn.run_sync(_schema_state) + assert state == { + "columns": {"file_id", "account_id", "expires_at"}, + "file_id_length": None, + "primary_key": ["file_id"], + "indexes": {"ix_file_account_pins_expires_at"}, + } + + await to_thread.run_sync(lambda: command.downgrade(_build_alembic_config(db_url), parent_revision)) + async with engine.connect() as conn: + assert await conn.run_sync(_schema_state) is None + + result = await to_thread.run_sync(lambda: run_upgrade(db_url, "head", bootstrap_legacy=False)) + assert result.current_revision == _HEAD_REVISION + async with engine.connect() as conn: + assert await conn.run_sync(_schema_state) is not None + finally: + await engine.dispose() + + @pytest.mark.asyncio async def test_dashboard_settings_default_flip_migration_does_not_infer_intent_from_updated_at(tmp_path): db_url = f"sqlite+aiosqlite:///{tmp_path / 'dashboard-settings-defaults.sqlite'}" diff --git a/tests/integration/test_proxy_compact.py b/tests/integration/test_proxy_compact.py index 77fce895c7..ca03b1ed21 100644 --- a/tests/integration/test_proxy_compact.py +++ b/tests/integration/test_proxy_compact.py @@ -1118,13 +1118,8 @@ async def fake_ensure_fresh(self, account, *, force: bool = False, timeout_secon @pytest.mark.asyncio -async def test_proxy_compact_forwarded_bridge_preflight_budget_exhausted_settles_reservation(async_client, monkeypatch): - """Regression (route-level, forwarded bridge path): a compact request that - reaches the OWNER instance via the internal bridge forward — where - ``owns_reservation`` is false so ``compact_responses`` is the SOLE settler — - and whose preflight budget is exhausted MUST settle (release) the API-key - usage reservation before raising the ``502 upstream_request_timeout``, so - held API-key quota is not leaked. +async def test_proxy_compact_forwarded_bridge_rejection_leaves_reservation_for_origin(async_client, monkeypatch): + """A forwarded compact rejection must leave cleanup at the origin. This drives the REAL external surface, not a handcrafted service call: it POSTs a signed forwarded request to the internal bridge endpoint @@ -1133,16 +1128,9 @@ async def test_proxy_compact_forwarded_bridge_preflight_budget_exhausted_settles reproduced here through the api-keys service). ``internal_bridge_responses`` parses the forward, sets ``skip_limit_enforcement`` + the ``api_key_reservation_override``, and ``_stream_responses`` extracts the - terminal ``compaction_trigger`` and calls ``compact_responses`` with - ``owns_reservation`` false — so ``_compact_or_stream_responses``'s ``finally`` - does NOT release the reservation and ``compact_responses`` alone must settle - it. Pre-fix the budget-exhausted terminal raised via - ``_raise_proxy_budget_exhausted`` without settling (through the outer - ``except ProxyResponseError`` handler and the log-only ``finally``), leaving - the reservation row ``reserved`` (leaked held quota); post-fix the row is - ``released``. PR #1254 fixed the sibling transport-failure / permanent-refresh - preflight raises but left the budget-exhausted terminal out of scope; this - completes that invariant. + terminal ``compaction_trigger`` and returns a definitive non-200 before an + HTTP-200 cleanup handoff. The receiver must not settle the origin's + reservation; the origin remains responsible for its one release attempt. """ import app.modules.proxy._service.compact as compact_module from app.core.config.settings import get_settings @@ -1248,14 +1236,13 @@ async def test_proxy_compact_forwarded_bridge_preflight_budget_exhausted_settles assert response.status_code == 502, response.text assert response.json()["error"]["code"] == "upstream_request_timeout" - # The forwarded reservation row was RELEASED by compact_responses (sole - # settler) before the terminal raised (the fix). Pre-fix it stayed "reserved" - # — leaked held API-key quota — because owns_reservation is false on the - # forwarded path so the route's finally does not release it. + # The receiver returned a definitive non-200 and therefore did not acquire + # cleanup ownership. A real origin observes that rejection and releases the + # still-reserved row exactly once. async with SessionLocal() as session: row = await session.get(ApiKeyUsageReservation, reservation.reservation_id) assert row is not None - assert row.status == "released", f"forwarded reservation leaked held quota; status={row.status!r}" + assert row.status == "reserved" @pytest.mark.asyncio diff --git a/tests/integration/test_proxy_files.py b/tests/integration/test_proxy_files.py index db71df0ebf..fc9cb926eb 100644 --- a/tests/integration/test_proxy_files.py +++ b/tests/integration/test_proxy_files.py @@ -10,16 +10,23 @@ from __future__ import annotations +import asyncio import base64 import json +from datetime import datetime, timedelta from typing import cast import pytest +from sqlalchemy import func, select, text +import app.modules.proxy.file_pin_repository as file_pin_repository_module import app.modules.proxy.service as proxy_module from app.core.auth.refresh import RefreshError from app.core.clients.files import FileProxyError from app.core.clients.proxy import ProxyResponseError +from app.db.models import FileAccountPin +from app.db.session import SessionLocal +from app.modules.proxy.file_pin_repository import FileAccountPinRepository pytestmark = pytest.mark.integration @@ -517,6 +524,243 @@ async def test_resolve_file_account_for_responses_returns_pin_when_no_other_affi assert prepared.affinity_policy.codex_session_source == "session_header" +@pytest.mark.asyncio +async def test_file_account_pin_is_resolved_by_another_proxy_replica(async_client, monkeypatch): + await _import_account(async_client, "acc_cross_replica", "cross-replica@example.com") + + from app.core.openai.requests import ResponsesRequest + from app.dependencies import get_proxy_service_for_app + + origin = get_proxy_service_for_app(async_client._transport.app) + other_replica = proxy_module.ProxyService(origin._repo_factory) + await origin._pin_file_account("file_cross_replica", "acc_cross_replica") + + assert await other_replica._resolve_file_account("file_cross_replica") == "acc_cross_replica" + + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.2", + "instructions": "Summarize the uploaded file.", + "input": [ + { + "role": "user", + "content": [{"type": "input_file", "file_id": "file_cross_replica"}], + } + ], + } + ) + assert await other_replica._resolve_file_account_for_responses(payload, {}) == "acc_cross_replica" + + preferred_accounts: list[str | None] = [] + + async def fake_proxy_files_call(**kwargs): + preferred_accounts.append(await kwargs["resolve_preferred_account_id"]()) + result = {"status": "success"} + await kwargs["on_success"](result, "acc_cross_replica") + return result, "acc_cross_replica" + + monkeypatch.setattr(other_replica, "_proxy_files_call", fake_proxy_files_call) + assert await other_replica.finalize_file("file_cross_replica", {}) == {"status": "success"} + assert preferred_accounts == ["acc_cross_replica"] + + +@pytest.mark.asyncio +async def test_file_account_pin_repository_ignores_expired_rows(async_client): + del async_client + async with SessionLocal() as session: + database_now = await session.scalar(select(func.now())) + assert isinstance(database_now, datetime) + session.add( + FileAccountPin( + file_id="file_expired_repository", + account_id="acc_expired_repository", + expires_at=database_now, + ) + ) + await session.commit() + + async with SessionLocal() as session: + repository = FileAccountPinRepository(session) + assert await repository.get_live_account_id("file_expired_repository") is None + + +@pytest.mark.asyncio +async def test_file_account_pin_same_owner_claim_is_idempotent_and_renews_expiry(async_client): + del async_client + async with SessionLocal() as session: + await FileAccountPinRepository(session).claim( + "file_same_owner", + "acc_same_owner", + ttl_seconds=60, + ) + async with SessionLocal() as session: + first_pin = await session.get(FileAccountPin, "file_same_owner") + assert first_pin is not None + first_expiry = first_pin.expires_at + + async with SessionLocal() as session: + await FileAccountPinRepository(session).claim( + "file_same_owner", + "acc_same_owner", + ttl_seconds=120, + ) + async with SessionLocal() as session: + renewed_pin = await session.get(FileAccountPin, "file_same_owner") + assert renewed_pin is not None + assert renewed_pin.account_id == "acc_same_owner" + assert renewed_pin.expires_at > first_expiry + + +@pytest.mark.asyncio +async def test_file_account_pin_claim_refreshes_a_stale_successful_insert_expiry( + async_client, + monkeypatch, +): + del async_client + stale_insert = text( + """ + INSERT INTO file_account_pins (file_id, account_id, expires_at) + VALUES ( + :file_id, + :account_id, + (strftime('%Y-%m-%d %H:%M:%f', 'now', '-' || :ttl || ' seconds') || '000') + ) + RETURNING account_id + """ + ) + monkeypatch.setattr(file_pin_repository_module, "_SQLITE_CLAIM", stale_insert) + + async with SessionLocal() as session: + await FileAccountPinRepository(session).claim( + "file_stale_insert_candidate", + "acc_stale_insert_candidate", + ttl_seconds=120, + ) + + async with SessionLocal() as session: + database_now = await session.scalar(select(func.now())) + pin = await session.get(FileAccountPin, "file_stale_insert_candidate") + assert isinstance(database_now, datetime) + assert pin is not None + assert pin.expires_at > database_now + timedelta(seconds=100) + + +@pytest.mark.asyncio +async def test_file_account_pin_claim_rolls_back_when_post_claim_refresh_does_not_match( + async_client, + monkeypatch, +): + del async_client + monkeypatch.setattr( + file_pin_repository_module, + "_SQLITE_REFRESH", + text( + """ + UPDATE file_account_pins + SET expires_at = expires_at + WHERE file_id = :file_id + AND account_id = :account_id + AND 0 = 1 + RETURNING account_id + """ + ), + ) + + async with SessionLocal() as session: + with pytest.raises(RuntimeError, match="Failed to refresh file account pin after claim"): + await FileAccountPinRepository(session).claim( + "file_failed_post_claim_refresh", + "acc_failed_post_claim_refresh", + ttl_seconds=120, + ) + + async with SessionLocal() as session: + assert await session.get(FileAccountPin, "file_failed_post_claim_refresh") is None + + +@pytest.mark.asyncio +async def test_reclaimed_file_account_pin_is_observed_without_local_cache(async_client): + from app.dependencies import get_proxy_service_for_app + + first_replica = get_proxy_service_for_app(async_client._transport.app) + second_replica = proxy_module.ProxyService(first_replica._repo_factory) + await first_replica._pin_file_account("file_claim_lifecycle", "acc_claim_a") + assert await first_replica._resolve_file_account("file_claim_lifecycle") == "acc_claim_a" + + async with SessionLocal() as session: + database_now = await session.scalar(select(func.now())) + assert isinstance(database_now, datetime) + pin = await session.get(FileAccountPin, "file_claim_lifecycle") + assert pin is not None + pin.expires_at = database_now - timedelta(seconds=1) + session.add( + FileAccountPin( + file_id="file_cleanup_expired", + account_id="acc_cleanup_expired", + expires_at=database_now - timedelta(seconds=1), + ) + ) + await session.commit() + + await second_replica._pin_file_account("file_claim_lifecycle", "acc_claim_b") + + async with SessionLocal() as session: + assert await session.get(FileAccountPin, "file_cleanup_expired") is None + + assert await first_replica._resolve_file_account("file_claim_lifecycle") == "acc_claim_b" + + +@pytest.mark.asyncio +async def test_concurrent_file_account_pin_claims_choose_one_durable_owner(async_client): + from app.dependencies import get_proxy_service_for_app + + app_service = get_proxy_service_for_app(async_client._transport.app) + start = asyncio.Event() + + async def claim(account_id: str) -> tuple[str, str]: + replica = proxy_module.ProxyService(app_service._repo_factory) + await start.wait() + try: + await replica._pin_file_account("file_concurrent_claim", account_id) + except ProxyResponseError as exc: + return account_id, str(exc.payload["error"]["code"]) + return account_id, "claimed" + + claims = [ + asyncio.create_task(claim("acc_race_a")), + asyncio.create_task(claim("acc_race_b")), + ] + start.set() + results = await asyncio.gather(*claims) + + winners = [account_id for account_id, outcome in results if outcome == "claimed"] + conflicts = [account_id for account_id, outcome in results if outcome == "continuity_owner_conflict"] + assert len(winners) == 1 + assert len(conflicts) == 1 + + observer = proxy_module.ProxyService(app_service._repo_factory) + assert await observer._resolve_file_account("file_concurrent_claim") == winners[0] + + +@pytest.mark.asyncio +async def test_live_file_account_pin_cannot_be_reassigned_by_another_replica(async_client): + await _import_account(async_client, "acc_pin_owner_a", "pin-owner-a@example.com") + await _import_account(async_client, "acc_pin_owner_b", "pin-owner-b@example.com") + + from app.dependencies import get_proxy_service_for_app + + origin = get_proxy_service_for_app(async_client._transport.app) + other_replica = proxy_module.ProxyService(origin._repo_factory) + await origin._pin_file_account("file_immutable_owner", "acc_pin_owner_a") + + with pytest.raises(ProxyResponseError) as exc_info: + await other_replica._pin_file_account("file_immutable_owner", "acc_pin_owner_b") + + assert exc_info.value.status_code == 502 + assert exc_info.value.payload["error"]["code"] == "continuity_owner_conflict" + assert await other_replica._resolve_file_account("file_immutable_owner") == "acc_pin_owner_a" + + @pytest.mark.asyncio async def test_v1_responses_file_id_pin_overrides_prompt_cache_key(async_client, monkeypatch): """A prompt-cache key is locality; an account-scoped file is ownership.""" diff --git a/tests/unit/test_file_pin_repository.py b/tests/unit/test_file_pin_repository.py new file mode 100644 index 0000000000..cc79dc13a6 --- /dev/null +++ b/tests/unit/test_file_pin_repository.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import pytest +from sqlalchemy import String + +from app.db.models import FileAccountPin +from app.modules.proxy.file_pin_repository import ( + build_file_account_pin_claim, + build_file_account_pin_cleanup, + build_file_account_pin_live_lookup, + build_file_account_pin_refresh, +) + +pytestmark = pytest.mark.unit + + +def test_file_account_pin_keeps_upstream_file_id_opaque() -> None: + file_id_type = FileAccountPin.__table__.c.file_id.type + assert isinstance(file_id_type, String) + assert file_id_type.length is None + + +def test_postgresql_file_pin_statements_use_current_database_clock() -> None: + claim_sql = str(build_file_account_pin_claim(dialect_name="postgresql")) + conflict_clause = claim_sql.split("DO UPDATE SET", 1)[1] + + assert claim_sql.count("clock_timestamp() + make_interval(secs => :ttl)") == 2 + assert "excluded.expires_at" not in conflict_clause + assert "expires_at = clock_timestamp() + make_interval(secs => :ttl)" in conflict_clause + assert "file_account_pins.expires_at <= clock_timestamp()" in conflict_clause + cleanup_sql = str(build_file_account_pin_cleanup(dialect_name="postgresql")) + assert "expires_at <= statement_timestamp()" in cleanup_sql + assert "clock_timestamp()" not in cleanup_sql + refresh_sql = str(build_file_account_pin_refresh(dialect_name="postgresql")) + assert "expires_at = clock_timestamp() + make_interval(secs => :ttl)" in refresh_sql + assert "file_id = :file_id" in refresh_sql + assert "account_id = :account_id" in refresh_sql + assert "RETURNING account_id" in refresh_sql + assert "expires_at > clock_timestamp()" in str(build_file_account_pin_live_lookup(dialect_name="postgresql")) + assert "expires_at > clock_timestamp()" in str( + build_file_account_pin_live_lookup(dialect_name="postgresql", many=True) + ) + + +def test_sqlite_file_pin_statements_use_padded_statement_clock() -> None: + claim_sql = str(build_file_account_pin_claim(dialect_name="sqlite")) + sqlite_now = "(strftime('%Y-%m-%d %H:%M:%f', 'now') || '000')" + sqlite_now_plus_ttl = "(strftime('%Y-%m-%d %H:%M:%f', 'now', '+' || :ttl || ' seconds') || '000')" + conflict_clause = claim_sql.split("DO UPDATE SET", 1)[1] + + assert claim_sql.count(sqlite_now_plus_ttl) == 2 + assert "excluded.expires_at" not in conflict_clause + assert f"expires_at = {sqlite_now_plus_ttl}" in conflict_clause + assert f"file_account_pins.expires_at <= {sqlite_now}" in conflict_clause + assert sqlite_now in str(build_file_account_pin_cleanup(dialect_name="sqlite")) + refresh_sql = str(build_file_account_pin_refresh(dialect_name="sqlite")) + assert sqlite_now_plus_ttl in refresh_sql + assert "file_id = :file_id" in refresh_sql + assert "account_id = :account_id" in refresh_sql + assert "RETURNING account_id" in refresh_sql + assert sqlite_now in str(build_file_account_pin_live_lookup(dialect_name="sqlite")) + assert sqlite_now in str(build_file_account_pin_live_lookup(dialect_name="sqlite", many=True)) + assert "CURRENT_TIMESTAMP" not in claim_sql + + +@pytest.mark.parametrize( + "builder", + [ + build_file_account_pin_claim, + build_file_account_pin_cleanup, + build_file_account_pin_refresh, + build_file_account_pin_live_lookup, + ], +) +def test_file_pin_statement_builders_reject_unknown_dialect(builder) -> None: + with pytest.raises(RuntimeError, match="Unsupported database dialect"): + builder(dialect_name="mysql") diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 00dd232e8e..c38658d1e4 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -10026,18 +10026,19 @@ async def test_stream_via_http_bridge_fails_closed_on_forward_loop_prevented( ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) payload = proxy_service.ResponsesRequest.model_validate({"model": "gpt-5.4", "instructions": "hi", "input": "hi"}) + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv-forwarded", + key_id="key-1", + model="gpt-5.4", + ) owner_forward = proxy_service._HTTPBridgeOwnerForward( owner_instance="instance-b", owner_endpoint="http://instance-b", key=proxy_service._HTTPBridgeSessionKey("session_header", "sid-123", None), ) - async def fake_forward(**kwargs: object): - del kwargs - raise ProxyResponseError(503, proxy_service.openai_error("bridge_forward_loop_prevented", "loop")) - yield "" - get_or_create = AsyncMock(return_value=owner_forward) + forward = Mock() monkeypatch.setattr( proxy_service, "get_settings_cache", @@ -10058,7 +10059,7 @@ async def fake_forward(**kwargs: object): monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) monkeypatch.setattr(service._durable_bridge, "lookup_request_targets", AsyncMock(return_value=None)) monkeypatch.setattr(service, "_get_or_create_http_bridge_session", get_or_create) - monkeypatch.setattr(service, "_forward_http_bridge_request_to_owner", fake_forward) + monkeypatch.setattr(service, "_forward_http_bridge_request_to_owner", forward) with pytest.raises(ProxyResponseError) as exc_info: async for _ in service._stream_via_http_bridge( @@ -10067,18 +10068,20 @@ async def fake_forward(**kwargs: object): codex_session_affinity=True, openai_cache_affinity=False, api_key=None, - api_key_reservation=None, + api_key_reservation=reservation, propagate_http_errors=False, suppress_text_done_events=False, idle_ttl_seconds=120.0, codex_idle_ttl_seconds=900.0, max_sessions=8, queue_limit=4, + forwarded_request=True, ): pass assert exc_info.value.payload["error"]["code"] == "bridge_forward_loop_prevented" get_or_create.assert_awaited_once() + forward.assert_not_called() @pytest.mark.asyncio @@ -10497,7 +10500,7 @@ async def test_http_bridge_local_owner_rejects_aliases_for_distinct_live_session @pytest.mark.asyncio -async def test_stream_via_http_bridge_reacquires_api_key_reservation_after_owner_forward_failure( +async def test_stream_via_http_bridge_reuses_origin_reservation_after_safe_owner_forward_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) @@ -10508,11 +10511,6 @@ async def test_stream_via_http_bridge_reacquires_api_key_reservation_after_owner key_id=api_key.id, model="gpt-5.4", ) - retried_reservation = proxy_service.ApiKeyUsageReservationData( - reservation_id="resv-retry", - key_id=api_key.id, - model="gpt-5.4", - ) payload = proxy_service.ResponsesRequest.model_validate( { "model": "gpt-5.4", @@ -10540,7 +10538,7 @@ async def test_stream_via_http_bridge_reacquires_api_key_reservation_after_owner model="gpt-5.4", service_tier=None, reasoning_effort=None, - api_key_reservation=retried_reservation, + api_key_reservation=initial_reservation, started_at=started_at, event_queue=asyncio.Queue(), transport="http", @@ -10614,7 +10612,7 @@ async def produce_after_reattach_delay() -> None: asyncio.create_task(produce_after_reattach_delay()) - reserve_retry = AsyncMock(return_value=retried_reservation) + reserve_retry = AsyncMock() capacity_unavailable = ProxyResponseError( 503, proxy_service.openai_error("no_accounts", "Rate limit exceeded. Try again in 120s"), @@ -10678,9 +10676,145 @@ async def produce_after_reattach_delay() -> None: assert http_bridge_streaming_module._codex_keepalive_frame() in chunks assert chunks[-1] == 'data: {"type":"response.completed"}\n\n' assert get_or_create.await_count == 3 - assert prepare_reservations == [initial_reservation, retried_reservation] - assert submitted_reservations == [retried_reservation] - reserve_retry.assert_awaited_once() + assert prepare_reservations == [initial_reservation, initial_reservation] + assert submitted_reservations == [initial_reservation] + reserve_retry.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("forward_outcome", "expected_recovery"), + [ + pytest.param("not_dispatched", True, id="pre-dispatch"), + pytest.param("dispatch_ambiguous", False, id="ambiguous-dispatch"), + pytest.param("receiver_acknowledged", False, id="acknowledged-then-lost"), + pytest.param("receiver_rejected", True, id="definitive-non-200"), + ], +) +async def test_stream_via_http_bridge_replays_only_definitive_safe_forward_outcomes( + monkeypatch: pytest.MonkeyPatch, + forward_outcome: str, + expected_recovery: bool, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + payload = proxy_service.ResponsesRequest.model_validate( + { + "model": "gpt-5.4", + "instructions": "continue", + "input": "hello", + "previous_response_id": "resp_forward_outcome", + } + ) + request_state = proxy_service._WebSocketRequestState( + request_id="req-forward-outcome", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + event_queue=asyncio.Queue(), + previous_response_id=payload.previous_response_id, + transport="http", + ) + + def fake_prepare( + _payload: proxy_service.ResponsesRequest, + _headers: dict[str, str] | Any, + *, + api_key: proxy_service.ApiKeyData | None, + api_key_reservation: proxy_service.ApiKeyUsageReservationData | None, + request_id: str, + client_ip: str | None = None, + ) -> tuple[proxy_service._WebSocketRequestState, str]: + del api_key, api_key_reservation, request_id, client_ip + return request_state, '{"type":"response.create"}' + + owner_forward = proxy_service._HTTPBridgeOwnerForward( + owner_instance="instance-b", + owner_endpoint="http://instance-b", + key=proxy_service._HTTPBridgeSessionKey("session_header", "sid-forward-outcome", None), + ) + local_replay_attempted = ProxyResponseError( + 500, + openai_error("local_replay_attempted", "sentinel", error_type="server_error"), + ) + get_or_create = AsyncMock(side_effect=[owner_forward, local_replay_attempted]) + + class _OutcomeOwnerClient: + async def stream_responses(self, **kwargs: object): + on_request_dispatched = cast(Any, kwargs["on_request_dispatched"]) + on_response_rejected = cast(Any, kwargs["on_response_rejected"]) + on_response_ready = cast(Any, kwargs["on_response_ready"]) + if forward_outcome != "not_dispatched": + on_request_dispatched() + if forward_outcome == "receiver_rejected": + on_response_rejected() + raise ProxyResponseError( + 503, + openai_error( + "bridge_owner_unreachable", + "Owner rejected the forwarded request", + error_type="server_error", + ), + ) + if forward_outcome == "receiver_acknowledged": + on_response_ready() + raise aiohttp.ClientConnectionError("owner response status unavailable") + yield "" + + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: cast( + Any, + SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + sticky_threads_enabled=False, + openai_cache_affinity_max_age_seconds=1800, + http_responses_session_bridge_prompt_cache_idle_ttl_seconds=3600, + http_responses_session_bridge_gateway_safe_mode=False, + ) + ) + ), + ), + ) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: Settings( + http_responses_session_bridge_enabled=True, + http_responses_session_bridge_instance_id="instance-a", + ), + ) + monkeypatch.setattr(service._durable_bridge, "lookup_request_targets", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_resolve_websocket_previous_response_owner", AsyncMock(return_value="acc-1")) + monkeypatch.setattr(service, "_prepare_http_bridge_request", fake_prepare) + monkeypatch.setattr(service, "_get_or_create_http_bridge_session", get_or_create) + service._http_bridge_owner_client = cast(Any, _OutcomeOwnerClient()) + + with pytest.raises(ProxyResponseError) as exc_info: + _ = [ + chunk + async for chunk in service._stream_via_http_bridge( + payload, + headers={"x-codex-session-id": "sid-forward-outcome"}, + codex_session_affinity=True, + propagate_http_errors=False, + openai_cache_affinity=False, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=1800.0, + max_sessions=8, + queue_limit=4, + ) + ] + + expected_code = "local_replay_attempted" if expected_recovery else "bridge_owner_unreachable" + assert exc_info.value.payload["error"]["code"] == expected_code + assert get_or_create.await_count == (2 if expected_recovery else 1) async def _run_owner_forward_recovery_with_session( @@ -11378,6 +11512,116 @@ async def fake_stream_http_bridge_session_events( assert session.queued_request_count == 0 +@pytest.mark.asyncio +async def test_stream_via_http_bridge_context_overflow_keeps_file_owner_on_soft_affinity_recovery( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + payload = proxy_service.ResponsesRequest.model_validate( + { + "model": "gpt-5.4", + "instructions": "hi", + "input": [ + { + "role": "user", + "content": [{"type": "input_file", "file_id": "file_doc"}], + } + ], + "previous_response_id": "resp_file_context", + "prompt_cache_key": "bridge-file-context-overflow", + } + ) + key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-file-context-overflow", None) + initial_session = _make_bridge_session(key=key) + retry_session = _make_bridge_session(key=key) + initial_session.account.id = "acc-file" + retry_session.account.id = "acc-file" + + observed_states: list[proxy_service._WebSocketRequestState] = [] + observed_frames: list[dict[str, Any]] = [] + + async def fake_stream_http_bridge_session_events( + _session: proxy_service._HTTPBridgeSession, + *, + request_state: proxy_service._WebSocketRequestState, + text_data: str, + queue_limit: int, + propagate_http_errors: bool, + downstream_turn_state: str | None, + request_deadline: float | None = None, + ): + del queue_limit, propagate_http_errors, downstream_turn_state, request_deadline + observed_states.append(request_state) + observed_frames.append(json.loads(text_data)) + if len(observed_states) == 1: + raise ProxyResponseError( + 400, + proxy_service.openai_error( + "context_length_exceeded", + "Your input exceeds the context window of this model.", + error_type="invalid_request_error", + ), + ) + yield 'data: {"type":"response.completed"}\n\n' + + get_or_create = AsyncMock(side_effect=[initial_session, retry_session]) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: cast( + Any, + SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + sticky_threads_enabled=False, + openai_cache_affinity_max_age_seconds=1800, + http_responses_session_bridge_prompt_cache_idle_ttl_seconds=3600, + http_responses_session_bridge_gateway_safe_mode=False, + ) + ) + ), + ), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr(service._durable_bridge, "lookup_request_targets", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_resolve_websocket_previous_response_owner", AsyncMock(return_value="acc-file")) + monkeypatch.setattr(service, "_get_or_create_http_bridge_session", get_or_create) + monkeypatch.setattr(service, "_stream_http_bridge_session_events", fake_stream_http_bridge_session_events) + monkeypatch.setattr(service, "_close_http_bridge_session", AsyncMock()) + + chunks = [ + chunk + async for chunk in service._stream_via_http_bridge( + payload, + headers={}, + codex_session_affinity=False, + propagate_http_errors=True, + openai_cache_affinity=True, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=900.0, + max_sessions=8, + queue_limit=4, + rewritten_file_account_id="acc-file", + ) + ] + + assert chunks == ['data: {"type":"response.completed"}\n\n'] + assert get_or_create.await_count == 2 + initial_call, retry_call = get_or_create.await_args_list + assert initial_call.kwargs["preferred_account_id"] == "acc-file" + assert initial_call.kwargs["fallback_on_preferred_account_unavailable"] is False + assert retry_call.kwargs["previous_response_id"] is None + assert retry_call.kwargs["preferred_account_id"] == "acc-file" + assert retry_call.kwargs["fallback_on_preferred_account_unavailable"] is False + assert "previous_response_id" not in observed_frames[1] + assert observed_frames[1]["input"] == payload.to_payload()["input"] + assert observed_states[1].preferred_account_id == "acc-file" + assert observed_states[1].file_required_preferred_account is True + + @pytest.mark.asyncio async def test_stream_via_http_bridge_context_overflow_keeps_hard_affinity_session( monkeypatch: pytest.MonkeyPatch, @@ -17469,7 +17713,6 @@ async def test_stream_via_http_bridge_fails_closed_before_file_affinity_when_pre monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) - await service._pin_file_account("file_from_other_account", "acc-file") payload = proxy_service.ResponsesRequest.model_validate( { "model": "gpt-5.4", @@ -17525,6 +17768,7 @@ async def test_stream_via_http_bridge_fails_closed_before_file_affinity_when_pre codex_idle_ttl_seconds=1800.0, max_sessions=8, queue_limit=4, + rewritten_file_account_id="acc-file", ): pass diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 31910f68ca..ebd8efffc7 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextvars import errno import gc import hashlib @@ -10,7 +11,7 @@ import ssl import time from collections import deque -from collections.abc import Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from copy import deepcopy from datetime import timedelta @@ -61,9 +62,11 @@ from app.modules.api_keys.service import ApiKeyData from app.modules.proxy import affinity as proxy_affinity from app.modules.proxy import api as proxy_api +from app.modules.proxy import http_bridge_forwarding as proxy_http_bridge_forwarding from app.modules.proxy import request_policy as proxy_request_policy from app.modules.proxy import service as proxy_service from app.modules.proxy._service import compact as proxy_compact_service +from app.modules.proxy._service import file_ops as proxy_file_ops from app.modules.proxy._service import support as proxy_support from app.modules.proxy._service import warmup as proxy_warmup_service from app.modules.proxy._service.http_bridge import request_submit as proxy_http_bridge_request_submit @@ -2351,6 +2354,1897 @@ async def stream_responses(*args, **kwargs): assert "rate_limit_exceeded" in body +@pytest.mark.asyncio +async def test_forwarded_terminal_compaction_passes_signed_file_owner_to_compact_service( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = False + reservation = SimpleNamespace(reservation_id="reservation_forwarded_compaction") + release_reservation = AsyncMock() + compact_calls: list[dict[str, object]] = [] + + async def compact_responses(*_args: object, **kwargs: object) -> CompactResponsePayload: + compact_calls.append(kwargs) + proxy_support._signal_propagated_responses_service_cleanup_ready() + return CompactResponsePayload.model_validate( + { + "id": "resp_forwarded_compaction", + "object": "response.compaction", + "compaction_summary": {"encrypted_content": "enc_forwarded"}, + } + ) + + context = SimpleNamespace( + service=SimpleNamespace( + compact_responses=compact_responses, + rate_limit_headers=AsyncMock(return_value={}), + ) + ) + + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/internal/bridge/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [ + {"type": "input_file", "file_id": "file_forwarded_compaction"}, + {"type": "compaction_trigger"}, + ], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, context), + api_key=None, + codex_session_affinity=True, + prefer_http_bridge=True, + forwarded_request=True, + forwarded_headers={}, + forwarded_file_owner_account_id="acc_forwarded_compaction", + ) + + assert isinstance(response, StreamingResponse) + assert compact_calls[0]["forwarded_file_owner_account_id"] == "acc_forwarded_compaction" + _ = [chunk async for chunk in response.body_iterator] + release_reservation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forwarded_terminal_compaction_malformed_output_keeps_receiver_settlement_authoritative( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = False + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_malformed_compaction", + key_id="key_forwarded_malformed_compaction", + model="gpt-5.1", + ) + origin_release = AsyncMock() + receiver_settle = AsyncMock() + + async def compact_responses(*_args: object, **kwargs: object) -> CompactResponsePayload: + await receiver_settle(kwargs["api_key_reservation"]) + proxy_support._signal_propagated_responses_service_cleanup_ready() + return CompactResponsePayload.model_validate( + { + "id": "resp_forwarded_malformed_compaction", + "object": "response.compaction", + "output": [], + } + ) + + receiver_context = cast( + proxy_api.ProxyContext, + SimpleNamespace( + service=SimpleNamespace( + compact_responses=compact_responses, + rate_limit_headers=AsyncMock(return_value={}), + ) + ), + ) + receiver_request = Request( + { + "type": "http", + "method": "POST", + "path": "/internal/bridge/responses", + "headers": [], + } + ) + forwarded_payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "compaction_trigger"}], + "stream": True, + } + ) + + async def forward_to_receiver(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + proxy_support._signal_propagated_responses_owner_forward_dispatched() + receiver_response = await proxy_api._stream_responses( + receiver_request, + forwarded_payload, + context=receiver_context, + api_key=None, + codex_session_affinity=True, + prefer_http_bridge=True, + skip_limit_enforcement=True, + api_key_reservation_override=reservation, + include_rate_limit_headers=False, + forwarded_request=True, + forwarded_headers={}, + enforce_openai_sdk_contract=False, + ) + assert isinstance(receiver_response, StreamingResponse) + assert receiver_response.status_code == 200 + proxy_support._signal_propagated_responses_service_cleanup_ready() + async for chunk in receiver_response.body_iterator: + yield chunk.decode() if isinstance(chunk, bytes) else str(chunk) + + origin_context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_http_responses=forward_to_receiver)), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", origin_release) + origin_request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + origin_payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_forwarded_malformed_compaction"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + origin_request, + origin_payload, + context=origin_context, + api_key=None, + prefer_http_bridge=True, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 200 + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunk.decode() if isinstance(chunk, bytes) else str(chunk) for chunk in chunks) + assert '"type":"response.failed"' in body + assert '"code":"upstream_error"' in body + receiver_settle.assert_awaited_once_with(reservation) + origin_release.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_terminal_compaction_pre_handoff_failure_releases_at_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = False + reservation = SimpleNamespace(reservation_id="reservation_terminal_compaction_pre_handoff") + release_reservation = AsyncMock() + + async def compact_responses(*_args: object, **_kwargs: object) -> CompactResponsePayload: + raise proxy_module.ProxyResponseError( + 502, + openai_error("file_owner_unavailable", "Durable file owner unavailable"), + ) + + context = SimpleNamespace( + service=SimpleNamespace( + compact_responses=compact_responses, + rate_limit_headers=AsyncMock(return_value={}), + ) + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/backend-api/codex/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "compaction_trigger"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, context), + api_key=None, + codex_session_affinity=True, + ) + + assert response.status_code == 502 + release_reservation.assert_awaited_once_with(reservation) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["stream", "collect", "compact"]) +@pytest.mark.parametrize( + "failure_type", + [RuntimeError, asyncio.CancelledError], + ids=["error", "cancellation"], +) +async def test_responses_rate_limit_header_failure_releases_pre_service_reservation( + monkeypatch: pytest.MonkeyPatch, + surface: str, + failure_type: type[BaseException], +) -> None: + reservation = SimpleNamespace(reservation_id=f"reservation_{surface}_rate_limit_headers") + release_reservation = AsyncMock() + service_call = AsyncMock() + context = cast( + proxy_api.ProxyContext, + SimpleNamespace( + service=SimpleNamespace( + stream_responses=service_call, + stream_http_responses=service_call, + compact_responses=service_call, + ) + ), + ) + + monkeypatch.setattr(proxy_api, "_opportunistic_admission_denial", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr( + proxy_api, + "_rate_limit_headers_for_request", + AsyncMock(side_effect=failure_type("rate-limit headers unavailable")), + ) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": f"/test/{surface}", "headers": []}) + + with pytest.raises(failure_type): + if surface == "compact": + await proxy_api._compact_responses( + request, + ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": "hello", + } + ), + context, + api_key=None, + ) + else: + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "respond", + "input": "hello", + "stream": surface == "stream", + } + ) + if surface == "stream": + await proxy_api._stream_responses( + request, + payload, + context, + api_key=None, + ) + else: + await proxy_api._collect_responses( + request, + payload, + context, + api_key=None, + ) + + service_call.assert_not_awaited() + release_reservation.assert_awaited_once_with(reservation) + + +@pytest.mark.asyncio +async def test_terminal_compaction_cancellation_after_service_handoff_does_not_release_at_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = False + reservation = SimpleNamespace(reservation_id="reservation_terminal_compaction_cancelled") + release_reservation = AsyncMock() + + async def compact_responses(*_args: object, **_kwargs: object) -> CompactResponsePayload: + proxy_support._signal_propagated_responses_service_cleanup_ready() + raise asyncio.CancelledError + + context = cast( + proxy_api.ProxyContext, + SimpleNamespace( + service=SimpleNamespace( + compact_responses=compact_responses, + rate_limit_headers=AsyncMock(return_value={}), + ) + ), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/backend-api/codex/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "compaction_trigger"}], + "stream": True, + } + ) + + with pytest.raises(asyncio.CancelledError): + await proxy_api._stream_responses( + request, + payload, + context=context, + api_key=None, + codex_session_affinity=True, + ) + + release_reservation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_compact_api_does_not_release_after_service_settlement_handoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + reservation = SimpleNamespace(reservation_id="reservation_compact_api_handoff") + release_reservation = AsyncMock() + + async def compact_responses(*_args: object, **_kwargs: object) -> CompactResponsePayload: + proxy_support._signal_propagated_responses_service_cleanup_ready() + return CompactResponsePayload.model_validate( + { + "id": "resp_compact_api_handoff", + "object": "response.compaction", + "output": [{"type": "compaction", "encrypted_content": "enc_compact_api"}], + } + ) + + context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(compact_responses=compact_responses)), + ) + monkeypatch.setattr(proxy_api, "_opportunistic_admission_denial", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/responses/compact", "headers": []}) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": "hello", + } + ) + + response = await proxy_api._compact_responses(request, payload, context, api_key=None) + + assert response.status_code == 200 + release_reservation.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_stage", ["session_open", "owner_query"]) +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_responses_file_owner_database_failure_releases_reservation_before_selection( + monkeypatch: pytest.MonkeyPatch, + failure_stage: str, + release_fails: bool, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_file_owner_failure") + release_reservation = AsyncMock( + side_effect=RuntimeError("reservation database unavailable") if release_fails else None + ) + selection = AsyncMock() + + async def fail_lookup(_repository, _file_ids) -> dict[str, str]: + raise RuntimeError("file pin database unavailable") + + class _FailingSessionContext: + async def __aenter__(self) -> None: + raise RuntimeError("file pin database session unavailable") + + async def __aexit__(self, *_args: object) -> None: + return None + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + if failure_stage == "session_open": + monkeypatch.setattr(service, "_file_pin_session_factory", _FailingSessionContext) + else: + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + fail_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_db_failure"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + ) + + assert not isinstance(response, StreamingResponse) + assert response.status_code == 502 + assert json.loads(bytes(response.body))["error"]["code"] == "file_owner_unavailable" + release_reservation.assert_awaited_once_with(reservation) + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_responses_success_does_not_release_service_owned_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + api_key = _make_api_key_data("key_stream_service_owned_reservation") + reservation = SimpleNamespace(reservation_id="reservation_stream_service_owned") + release_reservation = AsyncMock() + + async def stream_responses(*_args: object, **_kwargs: object): + yield 'data: {"type":"response.completed","response":{"id":"resp_service_owned","status":"completed"}}\n\n' + + context = SimpleNamespace(service=SimpleNamespace(stream_responses=stream_responses)) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest(model="gpt-5.1", instructions="read", input="hello", stream=True) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, context), + api_key=api_key, + ) + + assert isinstance(response, StreamingResponse) + _ = [chunk async for chunk in response.body_iterator] + release_reservation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_collect_responses_success_does_not_release_service_owned_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + api_key = _make_api_key_data("key_collect_service_owned_reservation") + reservation = SimpleNamespace(reservation_id="reservation_collect_service_owned") + release_reservation = AsyncMock() + + async def stream_responses(*_args: object, **_kwargs: object): + yield 'data: {"type":"response.completed","response":{"id":"resp_collect_owned","status":"completed"}}\n\n' + + context = SimpleNamespace(service=SimpleNamespace(stream_responses=stream_responses)) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest(model="gpt-5.1", instructions="read", input="hello", stream=False) + + response = await proxy_api._collect_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, context), + api_key=api_key, + ) + + assert response.status_code == 200 + release_reservation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_collect_responses_cancellation_after_cleanup_handoff_releases_only_in_service( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_collect_cleanup_handoff") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_collect_cleanup_handoff", + key_id=api_key.id, + model="gpt-5.1", + ) + api_release = AsyncMock() + service_release = AsyncMock() + cleanup_ready = asyncio.Event() + stream_cancelled = asyncio.Event() + stream_blocker = asyncio.Event() + + async def stream_http_responses(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + proxy_support._signal_propagated_responses_service_cleanup_ready() + cleanup_ready.set() + try: + await stream_blocker.wait() + raise AssertionError("cancelled collector unexpectedly resumed") + except asyncio.CancelledError: + stream_cancelled.set() + raise + finally: + await service_release(reservation) + yield "" + + monkeypatch.setattr(service, "stream_http_responses", stream_http_responses) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", api_release) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest(model="gpt-5.1", instructions="read", input="hello", stream=False) + + collect_task = asyncio.create_task( + proxy_api._collect_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=api_key, + prefer_http_bridge=True, + ) + ) + try: + await asyncio.wait_for(cleanup_ready.wait(), timeout=1) + collect_task.cancel() + with pytest.raises(asyncio.CancelledError): + await collect_task + + await asyncio.wait_for(stream_cancelled.wait(), timeout=1) + api_release.assert_not_awaited() + service_release.assert_awaited_once_with(reservation) + finally: + stream_blocker.set() + if not collect_task.done(): + collect_task.cancel() + await asyncio.gather(collect_task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_responses_file_owner_lookup_cancellation_releases_reservation_once_before_selection( + monkeypatch: pytest.MonkeyPatch, + release_fails: bool, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_file_owner_cancel") + release_reservation = AsyncMock( + side_effect=RuntimeError("reservation database unavailable") if release_fails else None + ) + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_cancelled_lookup"}], + "stream": True, + } + ) + + request_task = asyncio.create_task( + proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + assert await service.drain_persistence_tasks(1.0) is True + release_reservation.assert_awaited_once_with(reservation) + selection.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_responses_reservation_release_is_tracked_without_swallowing_cancellation( + monkeypatch: pytest.MonkeyPatch, + release_fails: bool, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_cancel_scope_cleanup", + key_id="key_cancel_scope_cleanup", + model="gpt-5.1", + ) + release_count = 0 + reached_after_helper = False + + async def release_reservation(candidate: object) -> None: + nonlocal release_count + assert candidate is reservation + await anyio.sleep(0) + release_count += 1 + if release_fails: + raise RuntimeError("reservation database unavailable") + + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + with anyio.CancelScope() as cancel_scope: + cancel_scope.cancel() + await proxy_api._release_reservation_best_effort( + reservation, + action="cancel-scope regression", + scheduler=service, + request_id="request-cancel-scope-cleanup", + ) + reached_after_helper = True + + assert cancel_scope.cancelled_caught + assert reached_after_helper is False + assert await service.drain_persistence_tasks(1.0) is True + assert release_count == 1 + + +@pytest.mark.asyncio +async def test_direct_file_owner_lookup_cancellation_releases_only_api_owned_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_direct_file_owner_cancel") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_direct_file_owner_cancel", + key_id=api_key.id, + model="gpt-5.1", + ) + api_release_reservation = AsyncMock() + service_release_reservation = AsyncMock() + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr( + service, + "_release_unsettled_stream_api_key_usage", + service_release_reservation, + ) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", api_release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_direct_cancelled_lookup"}], + "stream": True, + } + ) + + request_task = asyncio.create_task( + proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=api_key, + prefer_http_bridge=False, + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + assert await service.drain_persistence_tasks(1.0) is True + api_release_reservation.assert_awaited_once_with(reservation) + service_release_reservation.assert_not_awaited() + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_origin_http_bridge_file_owner_lookup_cancellation_releases_only_api_owned_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_origin_bridge_file_owner_cancel") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_origin_bridge_file_owner_cancel", + key_id=api_key.id, + model="gpt-5.1", + ) + api_release_reservation = AsyncMock() + service_release_reservation = AsyncMock() + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_service, + "_http_bridge_runtime_config", + lambda _dashboard_settings, _app_settings: proxy_service._HTTPBridgeRuntimeConfig( + enabled=False, + idle_ttl_seconds=30.0, + codex_idle_ttl_seconds=30.0, + max_sessions=8, + queue_limit=16, + prompt_cache_idle_ttl_seconds=30.0, + gateway_safe_mode=False, + ), + ) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr( + service, + "_release_unsettled_stream_api_key_usage", + service_release_reservation, + ) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", api_release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_origin_bridge_cancelled_lookup"}], + "stream": True, + } + ) + + request_task = asyncio.create_task( + proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=api_key, + prefer_http_bridge=True, + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + assert await service.drain_persistence_tasks(1.0) is True + api_release_reservation.assert_awaited_once_with(reservation) + service_release_reservation.assert_not_awaited() + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_origin_releases_forwarded_file_owner_preflight_failure_exactly_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_forwarded_owner_preflight") + release_reservation = AsyncMock() + + async def receiver_preflight_failure(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise proxy_module.ProxyResponseError( + 502, + openai_error( + "file_owner_unavailable", + "Durable file ownership is temporarily unavailable", + error_type="server_error", + ), + ) + yield "" + + monkeypatch.setattr(service, "stream_http_responses", receiver_preflight_failure) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_forwarded_owner_preflight"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + prefer_http_bridge=True, + ) + + assert not isinstance(response, StreamingResponse) + assert response.status_code == 502 + assert json.loads(bytes(response.body))["error"]["code"] == "file_owner_unavailable" + release_reservation.assert_awaited_once_with(reservation) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("collect_response", [False, True], ids=["stream", "collect"]) +@pytest.mark.parametrize( + ("receiver_rejected", "expected_release_count"), + [(False, 0), (True, 1)], + ids=["ambiguous-dispatch", "definitive-non-200"], +) +async def test_origin_release_follows_owner_forward_dispatch_outcome( + monkeypatch: pytest.MonkeyPatch, + collect_response: bool, + receiver_rejected: bool, + expected_release_count: int, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + reservation = SimpleNamespace(reservation_id="reservation_owner_dispatch_outcome") + release_reservation = AsyncMock() + + async def forwarded_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + proxy_support._signal_propagated_responses_owner_forward_dispatched() + if receiver_rejected: + proxy_support._signal_propagated_responses_owner_forward_rejected() + raise proxy_module.ProxyResponseError( + 503, + openai_error("upstream_unavailable", "Owner forward failed before acknowledgement"), + ) + yield "" + + context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_http_responses=forwarded_stream)), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_owner_dispatch_outcome"}], + "stream": not collect_response, + } + ) + + if collect_response: + response = await proxy_api._collect_responses( + request, + payload, + context, + api_key=None, + prefer_http_bridge=True, + ) + else: + response = await proxy_api._stream_responses( + request, + payload, + context, + api_key=None, + prefer_http_bridge=True, + ) + + assert response.status_code == 503 + assert release_reservation.await_count == expected_release_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("collect_response", [False, True], ids=["stream", "collect"]) +@pytest.mark.parametrize( + "forward_outcome", + ["dispatch_ambiguous", "receiver_acknowledged"], + ids=["ambiguous-dispatch", "acknowledged-then-lost"], +) +async def test_responses_api_owner_forward_unsafe_outcome_neither_replays_nor_releases( + monkeypatch: pytest.MonkeyPatch, + collect_response: bool, + forward_outcome: str, +) -> None: + dashboard_settings = _make_proxy_settings() + dashboard_settings.http_responses_session_bridge_prompt_cache_idle_ttl_seconds = 3600 + app_settings = Settings( + http_responses_session_bridge_enabled=True, + http_responses_session_bridge_instance_id="instance-origin", + ) + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_public_owner_forward_outcome") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id=f"reservation_public_{forward_outcome}", + key_id=api_key.id, + model="gpt-5.1", + ) + origin_release = AsyncMock() + local_submit = AsyncMock(side_effect=AssertionError("local bridge submit attempted")) + seen_forward_reservations: list[proxy_service.ApiKeyUsageReservationData | None] = [] + owner_forward = proxy_service._HTTPBridgeOwnerForward( + owner_instance="instance-owner", + owner_endpoint="http://owner.invalid", + key=proxy_service._HTTPBridgeSessionKey( + "session_header", + "session-public-owner-forward", + api_key.id, + ), + ) + get_or_create = AsyncMock( + side_effect=[ + owner_forward, + AssertionError("owner-forward failure triggered local session creation"), + ] + ) + + class _OutcomeOwnerClient: + async def stream_responses(self, **kwargs: object) -> AsyncIterator[str]: + seen_forward_reservations.append( + cast(proxy_http_bridge_forwarding.HTTPBridgeForwardContext, kwargs["context"]).reservation + ) + cast(Callable[[], None], kwargs["on_request_dispatched"])() + if forward_outcome == "receiver_acknowledged": + cast(Callable[[], None], kwargs["on_response_ready"])() + raise aiohttp.ClientConnectionError("owner response status lost") + yield "" + + async def forbidden_direct_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise AssertionError("direct local upstream attempted") + yield "" + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(dashboard_settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: app_settings) + monkeypatch.setattr(service._durable_bridge, "lookup_request_targets", AsyncMock(return_value=None)) + monkeypatch.setattr( + service, + "_resolve_forwarded_file_account_for_responses", + AsyncMock(return_value="acc-file-owner"), + ) + monkeypatch.setattr(service, "_get_or_create_http_bridge_session", get_or_create) + monkeypatch.setattr(service, "_submit_http_bridge_request", local_submit) + monkeypatch.setattr(service, "_stream_with_retry", forbidden_direct_stream) + service._http_bridge_owner_client = cast(Any, _OutcomeOwnerClient()) + monkeypatch.setattr(proxy_api, "_opportunistic_admission_denial", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", origin_release) + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"x-codex-session-id", b"session-public-owner-forward")], + } + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_public_owner_forward"}], + "stream": not collect_response, + } + ) + + context = cast(proxy_api.ProxyContext, SimpleNamespace(service=service)) + if collect_response: + response = await proxy_api._collect_responses( + request, + payload, + context, + api_key=api_key, + codex_session_affinity=True, + prefer_http_bridge=True, + ) + else: + response = await proxy_api._stream_responses( + request, + payload, + context=context, + api_key=api_key, + codex_session_affinity=True, + prefer_http_bridge=True, + ) + + assert response.status_code == 503 + assert json.loads(bytes(response.body))["error"]["code"] == "bridge_owner_unreachable" + get_or_create.assert_awaited_once() + get_or_create_call = get_or_create.await_args + assert get_or_create_call is not None + assert get_or_create_call.kwargs["preferred_account_id"] == "acc-file-owner" + local_submit.assert_not_awaited() + assert seen_forward_reservations == [reservation] + origin_release.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_collect_responses_unexpected_pre_handoff_failure_releases_origin_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + reservation = SimpleNamespace(reservation_id="reservation_collect_unexpected_pre_handoff") + release_reservation = AsyncMock() + + async def failing_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise RuntimeError("settings database unavailable") + yield "" + + context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_http_responses=failing_stream)), + ) + monkeypatch.setattr(proxy_api, "_opportunistic_admission_denial", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": "hello", + "stream": False, + } + ) + + with pytest.raises(RuntimeError, match="settings database unavailable"): + await proxy_api._collect_responses( + request, + payload, + context, + api_key=None, + prefer_http_bridge=True, + ) + + release_reservation.assert_awaited_once_with(reservation) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("collect_response", [False, True], ids=["stream", "collect"]) +async def test_lost_owner_ack_keeps_origin_release_disabled_after_receiver_settlement( + monkeypatch: pytest.MonkeyPatch, + collect_response: bool, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_lost_owner_ack", + key_id="key_lost_owner_ack", + model="gpt-5.1", + ) + origin_release = AsyncMock() + receiver_settle = AsyncMock() + + class _LostAckOwnerClient: + async def stream_responses( + self, + *, + on_request_dispatched: Callable[[], None], + **_kwargs: object, + ) -> AsyncIterator[str]: + on_request_dispatched() + await receiver_settle(reservation) + raise aiohttp.ClientConnectionError("receiver 200 acknowledgement lost") + yield "" + + origin_service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + origin_service._http_bridge_owner_client = _LostAckOwnerClient() + owner_forward = proxy_service._HTTPBridgeOwnerForward( + owner_instance="owner-instance", + owner_endpoint="http://owner.invalid", + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-lost-owner-ack", None), + ) + + def origin_stream( + stream_payload: ResponsesRequest, + headers: Mapping[str, str], + **kwargs: object, + ) -> AsyncIterator[str]: + return origin_service._forward_http_bridge_request_to_owner( + owner_forward=owner_forward, + payload=stream_payload, + headers=headers, + api_key_reservation=cast( + proxy_service.ApiKeyUsageReservationData | None, + kwargs["api_key_reservation"], + ), + codex_session_affinity=False, + downstream_turn_state=None, + request_started_at=time.monotonic(), + proxy_api_authorization=None, + file_owner_account_id="acc_lost_owner_ack", + ) + + context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_http_responses=origin_stream)), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", origin_release) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_lost_owner_ack"}], + "stream": not collect_response, + } + ) + + if collect_response: + response = await proxy_api._collect_responses( + request, + payload, + context, + api_key=None, + prefer_http_bridge=True, + ) + else: + response = await proxy_api._stream_responses( + request, + payload, + context, + api_key=None, + prefer_http_bridge=True, + ) + + assert response.status_code == 503 + receiver_settle.assert_awaited_once_with(reservation) + origin_release.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_collect_responses_frame_before_cleanup_handoff_does_not_transfer_ownership( + monkeypatch: pytest.MonkeyPatch, +) -> None: + reservation = SimpleNamespace(reservation_id="reservation_collect_pre_handoff_frame") + release_reservation = AsyncMock() + + async def stream_responses(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + yield ": keepalive\n\n" + raise proxy_module.ProxyResponseError( + 502, + openai_error("file_owner_unavailable", "Durable file owner unavailable"), + ) + + context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_responses=stream_responses)), + ) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": "hello", + "stream": False, + } + ) + + response = await proxy_api._collect_responses( + request, + payload, + context, + api_key=None, + ) + + assert response.status_code == 502 + release_reservation.assert_awaited_once_with(reservation) + + +@pytest.mark.asyncio +async def test_responses_delayed_file_owner_database_failure_releases_reservation_after_stream_handoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_file_owner_delayed_failure") + release_reservation = AsyncMock() + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_release = asyncio.Event() + + async def fail_lookup_after_handoff(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_release.wait() + raise RuntimeError("file pin database unavailable") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + fail_lookup_after_handoff, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_delayed_db_failure"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 200 + assert lookup_started.is_set() + lookup_release.set() + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunk.decode() if isinstance(chunk, bytes) else str(chunk) for chunk in chunks) + + assert "response.failed" in body + assert "file_owner_unavailable" in body + release_reservation.assert_awaited_once_with(reservation) + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_responses_file_owner_lookup_cancellation_after_stream_handoff_releases_reservation_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_file_owner_handoff_cancel") + release_reservation = AsyncMock() + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_cancelled = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + try: + await lookup_blocker.wait() + except asyncio.CancelledError: + lookup_cancelled.set() + raise + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_handoff_cancelled_lookup"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 200 + assert lookup_started.is_set() + body_iterator = cast(AsyncGenerator[Any, None], response.body_iterator) + body_task = asyncio.ensure_future(anext(body_iterator)) + await asyncio.sleep(0) + body_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await body_task + + await asyncio.wait_for(lookup_cancelled.wait(), timeout=1) + assert await service.drain_persistence_tasks(1.0) is True + release_reservation.assert_awaited_once_with(reservation) + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_native_heartbeat_disconnect_before_file_owner_resolution_cancels_lookup_and_releases_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_native_heartbeat_file_owner") + release_reservation = AsyncMock() + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_cancelled = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + try: + await lookup_blocker.wait() + except asyncio.CancelledError: + lookup_cancelled.set() + raise + return {} + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/backend-api/codex/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_native_heartbeat_cancel"}], + "stream": True, + } + ) + baseline_tasks = set(asyncio.all_tasks()) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + native_codex_heartbeat=True, + ) + + try: + assert isinstance(response, StreamingResponse) + assert response.status_code == 200 + assert lookup_started.is_set() + body_iterator = cast(AsyncGenerator[Any, None], response.body_iterator) + first = await anext(body_iterator) + assert first == proxy_service.CODEX_KEEPALIVE_FRAME + + await body_iterator.aclose() + + await asyncio.wait_for(lookup_cancelled.wait(), timeout=1) + assert await service.drain_persistence_tasks(1.0) is True + release_reservation.assert_awaited_once_with(reservation) + selection.assert_not_awaited() + assert not [ + task + for task in asyncio.all_tasks() + if task not in baseline_tasks and task is not asyncio.current_task() and not task.done() + ] + finally: + lookup_blocker.set() + leaked_tasks = [ + task + for task in asyncio.all_tasks() + if task not in baseline_tasks and task is not asyncio.current_task() and not task.done() + ] + for task in leaked_tasks: + task.cancel() + await asyncio.gather(*leaked_tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_native_heartbeat_close_after_file_lookup_before_service_cleanup_guard_releases_at_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_pre_cleanup_guard") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_pre_cleanup_guard", + key_id=api_key.id, + model="gpt-5.1", + ) + api_release_reservation = AsyncMock() + service_release_reservation = AsyncMock() + selection = AsyncMock() + lookup_complete = asyncio.Event() + preflight_started = asyncio.Event() + preflight_cancelled = asyncio.Event() + preflight_blocker = asyncio.Event() + + async def resolve_owner(_repository, file_ids) -> dict[str, str]: + lookup_complete.set() + return {file_id: "acc_file_owner" for file_id in file_ids} + + async def block_turn_state_preflight(*_args: object, **_kwargs: object) -> str | None: + preflight_started.set() + try: + await preflight_blocker.wait() + except asyncio.CancelledError: + preflight_cancelled.set() + raise + raise AssertionError("cancelled turn-state preflight unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + resolve_owner, + ) + monkeypatch.setattr(service, "_resolve_compact_turn_state_owner", block_turn_state_preflight) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr( + service, + "_release_unsettled_stream_api_key_usage", + service_release_reservation, + ) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", api_release_reservation) + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/backend-api/codex/responses", + "headers": [(b"x-codex-turn-state", b"turn_pre_cleanup_guard")], + } + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_pre_cleanup_guard"}], + "stream": True, + } + ) + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=api_key, + native_codex_heartbeat=True, + ) + + try: + assert isinstance(response, StreamingResponse) + await asyncio.wait_for(lookup_complete.wait(), timeout=1) + await asyncio.wait_for(preflight_started.wait(), timeout=1) + body_iterator = cast(AsyncGenerator[Any, None], response.body_iterator) + first = await anext(body_iterator) + assert first == proxy_service.CODEX_KEEPALIVE_FRAME + + await body_iterator.aclose() + + await asyncio.wait_for(preflight_cancelled.wait(), timeout=1) + assert await service.drain_persistence_tasks(1.0) is True + api_release_reservation.assert_awaited_once_with(reservation) + service_release_reservation.assert_not_awaited() + selection.assert_not_awaited() + finally: + preflight_blocker.set() + + +@pytest.mark.asyncio +async def test_native_heartbeat_close_after_file_owner_lookup_leaves_release_to_service_finalizer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_native_heartbeat_service_cleanup") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_native_heartbeat_service_cleanup", + key_id=api_key.id, + model="gpt-5.1", + ) + api_release_reservation = AsyncMock() + service_release_reservation = AsyncMock() + lookup_complete = asyncio.Event() + selection_started = asyncio.Event() + selection_cancelled = asyncio.Event() + selection_blocker = asyncio.Event() + + async def resolve_owner(_repository, file_ids) -> dict[str, str]: + lookup_complete.set() + return {file_id: "acc_file_owner" for file_id in file_ids} + + async def block_selection(*_args: object, **_kwargs: object) -> AccountSelection: + selection_started.set() + try: + await selection_blocker.wait() + except asyncio.CancelledError: + selection_cancelled.set() + raise + raise AssertionError("cancelled selection unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + resolve_owner, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", block_selection) + monkeypatch.setattr( + service, + "_release_unsettled_stream_api_key_usage", + service_release_reservation, + ) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", api_release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/backend-api/codex/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_service_cleanup"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=api_key, + native_codex_heartbeat=True, + ) + try: + assert isinstance(response, StreamingResponse) + await asyncio.wait_for(lookup_complete.wait(), timeout=1) + await asyncio.wait_for(selection_started.wait(), timeout=1) + body_iterator = cast(AsyncGenerator[Any, None], response.body_iterator) + first = await anext(body_iterator) + assert first == proxy_service.CODEX_KEEPALIVE_FRAME + + await body_iterator.aclose() + + await asyncio.wait_for(selection_cancelled.wait(), timeout=1) + assert await service.drain_persistence_tasks(1.0) is True + api_release_reservation.assert_not_awaited() + service_release_reservation.assert_awaited_once_with( + api_key=api_key, + api_key_reservation=reservation, + request_id=ANY, + ) + finally: + selection_blocker.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_native_heartbeat_anyio_cancellation_tracks_one_release_and_closes_lookup( + monkeypatch: pytest.MonkeyPatch, + release_fails: bool, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_native_heartbeat_anyio_cancel") + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_cancelled = asyncio.Event() + lookup_finalized = asyncio.Event() + lookup_blocker = asyncio.Event() + release_attempts = 0 + cancellation_observed = False + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + try: + await lookup_blocker.wait() + except asyncio.CancelledError: + lookup_cancelled.set() + raise + finally: + lookup_finalized.set() + return {} + + async def release_reservation(candidate: object) -> None: + nonlocal release_attempts + assert candidate is reservation + await anyio.sleep(0) + release_attempts += 1 + if release_fails: + raise RuntimeError("reservation database unavailable") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/backend-api/codex/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_native_heartbeat_anyio_cancel"}], + "stream": True, + } + ) + baseline_tasks = set(asyncio.all_tasks()) + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + native_codex_heartbeat=True, + ) + assert isinstance(response, StreamingResponse) + body_iterator = cast(AsyncGenerator[Any, None], response.body_iterator) + heartbeat_seen = asyncio.Event() + + async def consume_body() -> None: + nonlocal cancellation_observed + try: + first = await anext(body_iterator) + assert first == proxy_service.CODEX_KEEPALIVE_FRAME + heartbeat_seen.set() + await anext(body_iterator) + except anyio.get_cancelled_exc_class(): + cancellation_observed = True + raise + + try: + async with anyio.create_task_group() as task_group: + task_group.start_soon(consume_body) + await heartbeat_seen.wait() + task_group.cancel_scope.cancel() + + await asyncio.wait_for(lookup_cancelled.wait(), timeout=1) + await asyncio.wait_for(lookup_finalized.wait(), timeout=1) + assert await service.drain_persistence_tasks(1.0) is True + assert cancellation_observed is True + assert release_attempts == 1 + selection.assert_not_awaited() + assert not [ + task + for task in asyncio.all_tasks() + if task not in baseline_tasks and task is not asyncio.current_task() and not task.done() + ] + finally: + lookup_blocker.set() + leaked_tasks = [ + task + for task in asyncio.all_tasks() + if task not in baseline_tasks and task is not asyncio.current_task() and not task.done() + ] + for task in leaked_tasks: + task.cancel() + await asyncio.gather(*leaked_tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_native_heartbeat_close_after_buffered_service_event_closes_service_without_api_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_buffered_service_event") + release_reservation = AsyncMock() + service_started = asyncio.Event() + release_first_event = asyncio.Event() + first_event_yielded = asyncio.Event() + service_finalized = asyncio.Event() + hold_service_open = asyncio.Event() + + async def delayed_service_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + try: + service_started.set() + await release_first_event.wait() + first_event_yielded.set() + yield 'data: {"type":"response.created","response":{"id":"resp_buffered"}}\n\n' + await hold_service_open.wait() + finally: + service_finalized.set() + + monkeypatch.setattr(service, "stream_responses", delayed_service_stream) + monkeypatch.setattr(proxy_api, "_STREAM_STARTUP_ERROR_PROBE_SECONDS", 0.001) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/backend-api/codex/responses", "headers": []}) + payload = ResponsesRequest(model="gpt-5.1", instructions="read", input="hello", stream=True) + baseline_tasks = set(asyncio.all_tasks()) + + response = await proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + native_codex_heartbeat=True, + ) + try: + assert isinstance(response, StreamingResponse) + await asyncio.wait_for(service_started.wait(), timeout=1) + release_first_event.set() + await asyncio.wait_for(first_event_yielded.wait(), timeout=1) + await asyncio.sleep(0) + + body_iterator = cast(AsyncGenerator[Any, None], response.body_iterator) + first = await anext(body_iterator) + assert first == proxy_service.CODEX_KEEPALIVE_FRAME + await body_iterator.aclose() + + await asyncio.wait_for(service_finalized.wait(), timeout=1) + assert await service.drain_persistence_tasks(1.0) is True + release_reservation.assert_not_awaited() + assert not [ + task + for task in asyncio.all_tasks() + if task not in baseline_tasks and task is not asyncio.current_task() and not task.done() + ] + finally: + hold_service_open.set() + leaked_tasks = [ + task + for task in asyncio.all_tasks() + if task not in baseline_tasks and task is not asyncio.current_task() and not task.done() + ] + for task in leaked_tasks: + task.cancel() + await asyncio.gather(*leaked_tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_collect_responses_file_owner_lookup_cancellation_releases_reservation_once_before_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + reservation = SimpleNamespace(reservation_id="reservation_collect_file_owner_cancel") + release_reservation = AsyncMock() + selection = AsyncMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", release_reservation) + + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_collect_cancelled_lookup"}], + "stream": False, + } + ) + + request_task = asyncio.create_task( + proxy_api._collect_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + assert await service.drain_persistence_tasks(1.0) is True + release_reservation.assert_awaited_once_with(reservation) + selection.assert_not_awaited() + + @pytest.mark.asyncio @pytest.mark.parametrize("startup_surface", ["responses_route", "prime_helper"]) @pytest.mark.parametrize("capacity_recovers", [True, False]) @@ -3885,6 +5779,7 @@ async def select_account(_deadline: float, **kwargs: object) -> AccountSelection {"session_id": "soft-process-session"}, codex_session_affinity=True, openai_cache_affinity=True, + forwarded_file_owner_account_id=account.id, ) assert result.model_extra == {"output": []} @@ -3895,6 +5790,405 @@ async def select_account(_deadline: float, **kwargs: object) -> AccountSelection assert await service.drain_persistence_tasks(timeout_seconds=1) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("durable_owner", "expected_error_code"), + [ + (None, "file_owner_unavailable"), + ("acc_reclaimed_owner", "continuity_owner_conflict"), + ], +) +async def test_compact_revalidates_forwarded_file_owner_before_selection( + monkeypatch: pytest.MonkeyPatch, + durable_owner: str | None, + expected_error_code: str, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + selection = AsyncMock() + upstream_compact = AsyncMock() + resolve_owner = AsyncMock(return_value=durable_owner) + monkeypatch.setattr(service, "_resolve_file_account_for_responses", resolve_owner) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_service, "core_compact_responses", upstream_compact) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "input_file", "file_id": "file_forwarded_compaction"}], + } + ) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.compact_responses( + payload, + {}, + forwarded_file_owner_account_id="acc_forwarded_owner", + ) + + assert _proxy_error_code(exc_info.value) == expected_error_code + resolve_owner.assert_awaited_once_with(payload, {}) + selection.assert_not_awaited() + upstream_compact.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forwarded_compact_rejects_missing_durable_owner_proof_before_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + selection = AsyncMock() + upstream_compact = AsyncMock() + resolve_owner = AsyncMock(return_value="acc_durable_owner") + monkeypatch.setattr(service, "_resolve_file_account_for_responses", resolve_owner) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(proxy_service, "core_compact_responses", upstream_compact) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "input_file", "file_id": "file_missing_forwarded_proof"}], + } + ) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.compact_responses( + payload, + {}, + forwarded_request=True, + forwarded_file_owner_account_id=None, + ) + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + resolve_owner.assert_awaited_once_with(payload, {}) + selection.assert_not_awaited() + upstream_compact.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_compact_file_owner_database_failure_stops_before_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_compact_file_owner_failure") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_compact_file_owner_failure", + key_id=api_key.id, + model="gpt-5.1", + ) + selection = AsyncMock() + upstream_compact = AsyncMock() + settle_compact = AsyncMock() + + async def fail_lookup(_repository, _file_ids) -> dict[str, str]: + raise RuntimeError("file pin database unavailable") + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + fail_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(service, "_settle_compact_api_key_usage", settle_compact) + monkeypatch.setattr(proxy_service, "core_compact_responses", upstream_compact) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "input_file", "file_id": "file_db_failure"}], + } + ) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.compact_responses( + payload, + {}, + api_key=api_key, + api_key_reservation=reservation, + forwarded_file_owner_account_id="acc_forwarded_owner", + ) + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + settle_compact.assert_awaited_once_with( + api_key=api_key, + api_key_reservation=reservation, + response=None, + request_service_tier=None, + ) + selection.assert_not_awaited() + upstream_compact.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_compact_file_owner_lookup_cancellation_settles_reservation_before_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_compact_file_owner_cancel") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_compact_file_owner_cancel", + key_id=api_key.id, + model="gpt-5.1", + ) + selection = AsyncMock() + upstream_compact = AsyncMock() + settle_compact = AsyncMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(service, "_settle_compact_api_key_usage", settle_compact) + monkeypatch.setattr(proxy_service, "core_compact_responses", upstream_compact) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "input_file", "file_id": "file_compact_cancelled_lookup"}], + } + ) + + request_task = asyncio.create_task( + service.compact_responses( + payload, + {}, + api_key=api_key, + api_key_reservation=reservation, + forwarded_file_owner_account_id="acc_forwarded_owner", + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + settle_compact.assert_awaited_once_with( + api_key=api_key, + api_key_reservation=reservation, + response=None, + request_service_tier=None, + ) + selection.assert_not_awaited() + upstream_compact.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forwarded_compact_file_owner_failure_leaves_settlement_to_origin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_forwarded_compact_file_owner_failure") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_compact_file_owner_failure", + key_id=api_key.id, + model="gpt-5.1", + ) + selection = AsyncMock() + settle_compact = AsyncMock() + + async def fail_lookup(_repository, _file_ids) -> dict[str, str]: + raise RuntimeError("file pin database unavailable") + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + fail_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(service, "_settle_compact_api_key_usage", settle_compact) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "input_file", "file_id": "file_forwarded_compact_failure"}], + } + ) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.compact_responses( + payload, + {}, + api_key=api_key, + api_key_reservation=reservation, + forwarded_request=True, + forwarded_file_owner_account_id="acc_forwarded_owner", + ) + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + settle_compact.assert_not_awaited() + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forwarded_compact_file_owner_cancellation_leaves_settlement_to_origin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + api_key = _make_api_key_data("key_forwarded_compact_file_owner_cancel") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_compact_file_owner_cancel", + key_id=api_key.id, + model="gpt-5.1", + ) + selection = AsyncMock() + settle_compact = AsyncMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(service, "_settle_compact_api_key_usage", settle_compact) + payload = ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [{"type": "input_file", "file_id": "file_forwarded_compact_cancel"}], + } + ) + + request_task = asyncio.create_task( + service.compact_responses( + payload, + {}, + api_key=api_key, + api_key_reservation=reservation, + forwarded_request=True, + forwarded_file_owner_account_id="acc_forwarded_owner", + ) + ) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + settle_compact.assert_not_awaited() + selection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forwarded_terminal_compact_rejection_settles_only_at_origin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + api_key = _make_api_key_data("key_forwarded_terminal_compact_rejection") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_terminal_compact_rejection", + key_id=api_key.id, + model="gpt-5.1", + ) + origin_release = AsyncMock() + receiver_settle = AsyncMock() + receiver_service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + receiver_owner_failure = proxy_module.ProxyResponseError( + 502, + openai_error( + "file_owner_unavailable", + "Durable file ownership is temporarily unavailable", + error_type="server_error", + ), + ) + monkeypatch.setattr( + receiver_service, + "_resolve_forwarded_file_account_for_responses", + AsyncMock(side_effect=receiver_owner_failure), + ) + monkeypatch.setattr(receiver_service, "_settle_compact_api_key_usage", receiver_settle) + receiver_payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "compact", + "input": [ + {"type": "input_file", "file_id": "file_forwarded_terminal_compact"}, + {"type": "compaction_trigger"}, + ], + "stream": True, + } + ) + receiver_request = Request( + { + "type": "http", + "method": "POST", + "path": "/internal/bridge/responses", + "headers": [], + } + ) + + async def forward_to_rejecting_receiver(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + proxy_support._signal_propagated_responses_owner_forward_dispatched() + receiver_response = await proxy_api._stream_responses( + receiver_request, + receiver_payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=receiver_service)), + api_key=api_key, + codex_session_affinity=True, + prefer_http_bridge=True, + skip_limit_enforcement=True, + api_key_reservation_override=reservation, + include_rate_limit_headers=False, + forwarded_request=True, + forwarded_headers={}, + forwarded_file_owner_account_id="acc_forwarded_owner", + enforce_openai_sdk_contract=False, + ) + assert receiver_response.status_code == 502 + proxy_support._signal_propagated_responses_owner_forward_rejected() + raise proxy_module.ProxyResponseError( + receiver_response.status_code, + json.loads(bytes(receiver_response.body)), + ) + yield "" + + origin_context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_http_responses=forward_to_rejecting_receiver)), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", origin_release) + origin_request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + origin_payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_origin_forward"}], + "stream": True, + } + ) + + response = await proxy_api._stream_responses( + origin_request, + origin_payload, + context=origin_context, + api_key=api_key, + prefer_http_bridge=True, + ) + + assert response.status_code == 502 + receiver_settle.assert_not_awaited() + origin_release.assert_awaited_once_with(reservation) + + @pytest.mark.asyncio async def test_compact_fails_closed_when_turn_state_and_file_owners_conflict( monkeypatch: pytest.MonkeyPatch, @@ -18853,6 +21147,48 @@ class Settings: release_usage.assert_awaited_once_with(reservation) +@pytest.mark.asyncio +async def test_prepare_websocket_file_owner_database_failure_stops_before_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + settings = _make_proxy_settings() + api_key = _make_api_key_data("key_websocket_file_owner_failure") + reserve_usage = AsyncMock() + + async def fail_lookup(_repository, _file_ids) -> dict[str, str]: + raise RuntimeError("file pin database unavailable") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + fail_lookup, + ) + monkeypatch.setattr(service, "_reserve_websocket_api_key_usage", reserve_usage) + monkeypatch.setattr(service, "_refresh_websocket_api_key_policy", AsyncMock(return_value=api_key)) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service._prepare_websocket_response_create_request( + { + "type": "response.create", + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_websocket_db_failure"}], + }, + headers={}, + codex_session_affinity=False, + openai_cache_affinity=False, + sticky_threads_enabled=False, + openai_cache_affinity_max_age_seconds=300, + api_key=api_key, + ) + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + reserve_usage.assert_not_awaited() + + @pytest.mark.asyncio @pytest.mark.parametrize( "incremental_input", @@ -33037,25 +35373,179 @@ async def __aexit__(self, exc_type, exc, tb) -> None: @pytest.mark.asyncio -async def test_lookup_file_pin_returns_live_entry_and_evicts_expired(monkeypatch): +async def test_resolve_file_account_reads_durable_repository_every_time(monkeypatch): request_logs = _RequestLogsRecorder() service = proxy_service.ProxyService(_repo_factory(request_logs)) - fake_now = [100.0] + resolved_owners = iter(("acc_owner_a", "acc_owner_b")) + looked_up_file_ids: list[str] = [] - monkeypatch.setattr(proxy_service.time, "monotonic", lambda: fake_now[0]) + async def get_live_account_id(_repository, file_id: str) -> str: + looked_up_file_ids.append(file_id) + return next(resolved_owners) + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_id", + get_live_account_id, + ) + + assert await service._resolve_file_account("file_live") == "acc_owner_a" + assert await service._resolve_file_account("file_live") == "acc_owner_b" + assert looked_up_file_ids == ["file_live", "file_live"] + + +@pytest.mark.asyncio +async def test_responses_file_owner_resolution_batches_one_durable_lookup(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + lookup_batches: list[set[str]] = [] + + async def get_live_account_ids(_repository, file_ids) -> dict[str, str]: + lookup_batches.append(set(file_ids)) + return { + "file_batch_a": "acc_batch_owner", + "file_batch_b": "acc_batch_owner", + } + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + get_live_account_ids, + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "Read both files.", + "input": [ + {"type": "input_file", "file_id": "file_batch_a"}, + {"type": "input_file", "file_id": "file_batch_b"}, + {"type": "input_file", "file_id": "file_batch_a"}, + ], + } + ) + + assert await service._resolve_file_account_for_responses(payload, {}) == "acc_batch_owner" + assert lookup_batches == [{"file_batch_a", "file_batch_b"}] + + +@pytest.mark.asyncio +async def test_finalize_file_database_lookup_failure_stops_before_upstream(monkeypatch): + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + selection = AsyncMock() + + async def fail_lookup(_repository, _file_id: str) -> str | None: + raise RuntimeError("file pin database unavailable") + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_id", + fail_lookup, + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.finalize_file("file_db_failure", {}) + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + selection.assert_not_awaited() + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert request_logs.calls[0]["status"] == "error" + assert request_logs.calls[0]["error_code"] == "file_owner_unavailable" + assert request_logs.calls[0]["account_id"] is None + + +@pytest.mark.asyncio +async def test_create_file_database_write_failure_does_not_return_upstream_result(monkeypatch): + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + account = _make_account("acc_unpersisted") + selection = AsyncMock(return_value=AccountSelection(account=account, error_message=None)) + upstream_create = AsyncMock(return_value={"file_id": "file_unpersisted", "upload_url": "https://upload.invalid"}) + claim_calls: list[tuple[str, str, int]] = [] + + async def fail_claim( + _repository, + file_id: str, + account_id: str, + *, + ttl_seconds: int, + ) -> None: + claim_calls.append((file_id, account_id, ttl_seconds)) + raise RuntimeError("file pin database unavailable") + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_file_ops.FileAccountPinRepository, "claim", fail_claim) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(service, "_ensure_previsible_unary_fresh_with_failover", AsyncMock(return_value=account)) + monkeypatch.setattr(service._load_balancer, "record_success", AsyncMock()) + monkeypatch.setattr(proxy_service, "core_create_file", upstream_create) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.create_file({"file_name": "document.txt"}, {}) + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + upstream_create.assert_awaited_once() + assert claim_calls == [ + ( + "file_unpersisted", + "acc_unpersisted", + service._FILE_ACCOUNT_PIN_TTL_SECONDS, + ) + ] + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert request_logs.calls[0]["status"] == "error" + assert request_logs.calls[0]["error_code"] == "file_owner_unavailable" + assert request_logs.calls[0]["account_id"] == account.id - await service._pin_file_account("file_live", "acc_live") - entry = await service._lookup_file_pin("file_live") +@pytest.mark.asyncio +async def test_finalize_file_database_renewal_failure_logs_error_without_upstream_retry(monkeypatch): + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + account = _make_account("acc_finalize_unrenewed") + selection = AsyncMock(return_value=AccountSelection(account=account, error_message=None)) + upstream_finalize = AsyncMock(return_value={"status": "success"}) + unexpected_failover = AsyncMock( + side_effect=AssertionError("post-success pin persistence failure must not retry upstream") + ) + + async def resolve_owner(_repository, _file_id: str) -> str: + return account.id - assert entry is not None - assert entry.account_id == "acc_live" + async def fail_claim( + _repository, + _file_id: str, + _account_id: str, + *, + ttl_seconds: int, + ) -> None: + assert ttl_seconds == service._FILE_ACCOUNT_PIN_TTL_SECONDS + raise RuntimeError("file pin database unavailable") - fake_now[0] += service._FILE_ACCOUNT_PIN_TTL_SECONDS + 1 + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_file_ops.FileAccountPinRepository, "get_live_account_id", resolve_owner) + monkeypatch.setattr(proxy_file_ops.FileAccountPinRepository, "claim", fail_claim) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", selection) + monkeypatch.setattr(service, "_ensure_previsible_unary_fresh_with_failover", AsyncMock(return_value=account)) + monkeypatch.setattr(service, "_retry_previsible_unary_call_failover", unexpected_failover) + monkeypatch.setattr(service._load_balancer, "record_success", AsyncMock()) + monkeypatch.setattr(proxy_service, "core_finalize_file", upstream_finalize) - expired = await service._lookup_file_pin("file_live") + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service.finalize_file("file_finalize_unrenewed", {}) - assert expired is None + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + upstream_finalize.assert_awaited_once() + unexpected_failover.assert_not_awaited() + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert len(request_logs.calls) == 1 + assert request_logs.calls[0]["status"] == "error" + assert request_logs.calls[0]["error_code"] == "file_owner_unavailable" + assert request_logs.calls[0]["account_id"] == account.id @pytest.mark.asyncio @@ -33224,7 +35714,685 @@ async def test_stream_http_bridge_or_retry_rejects_input_image_sediment_url(monk @pytest.mark.asyncio -async def test_stream_http_bridge_or_retry_routes_input_file_file_id_without_rejecting(monkeypatch): +@pytest.mark.parametrize("forwarded_owner", [None, "acc_forwarded_owner"]) +async def test_stream_http_bridge_file_owner_database_failure_stops_before_upstream( + monkeypatch, + forwarded_owner, +): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + settings = _make_proxy_settings() + api_key = _make_api_key_data("key_forwarded_file_owner_failure") if forwarded_owner is not None else None + reservation = ( + proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_file_owner_failure", + key_id=api_key.id, + model="gpt-5.1", + ) + if api_key is not None + else None + ) + release_reservation = AsyncMock() + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_service, + "_http_bridge_runtime_config", + lambda _dashboard_settings, _app_settings: proxy_service._HTTPBridgeRuntimeConfig( + enabled=False, + idle_ttl_seconds=30.0, + codex_idle_ttl_seconds=30.0, + max_sessions=8, + queue_limit=16, + prompt_cache_idle_ttl_seconds=30.0, + gateway_safe_mode=False, + ), + ) + + async def fail_lookup(_repository, _file_ids) -> dict[str, str]: + raise RuntimeError("file pin database unavailable") + + upstream_called = False + + async def unexpected_stream(*_args, **_kwargs): + nonlocal upstream_called + upstream_called = True + yield "data: unexpected\n\n" + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + fail_lookup, + ) + monkeypatch.setattr(service, "_stream_with_retry", unexpected_stream) + monkeypatch.setattr( + service, + "_release_unsettled_stream_api_key_usage", + release_reservation, + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "hi", + "input": [{"type": "input_file", "file_id": "file_db_failure"}], + } + ) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + async for _ in service._stream_http_bridge_or_retry( + payload=payload, + headers={}, + codex_session_affinity=False, + propagate_http_errors=False, + openai_cache_affinity=False, + api_key=api_key, + api_key_reservation=reservation, + suppress_text_done_events=False, + forwarded_request=forwarded_owner is not None, + forwarded_file_owner_account_id=forwarded_owner, + ): + pass + + assert _proxy_error_code(exc_info.value) == "file_owner_unavailable" + assert upstream_called is False + release_reservation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_forwarded_http_bridge_file_owner_lookup_cancellation_leaves_cleanup_to_origin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + settings = _make_proxy_settings() + api_key = _make_api_key_data("key_forwarded_file_owner_cancel") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_file_owner_cancel", + key_id=api_key.id, + model="gpt-5.1", + ) + release_reservation = AsyncMock() + upstream_stream = MagicMock() + lookup_started = asyncio.Event() + lookup_blocker = asyncio.Event() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_service, + "_http_bridge_runtime_config", + lambda _dashboard_settings, _app_settings: proxy_service._HTTPBridgeRuntimeConfig( + enabled=False, + idle_ttl_seconds=30.0, + codex_idle_ttl_seconds=30.0, + max_sessions=8, + queue_limit=16, + prompt_cache_idle_ttl_seconds=30.0, + gateway_safe_mode=False, + ), + ) + + async def block_lookup(_repository, _file_ids) -> dict[str, str]: + lookup_started.set() + await lookup_blocker.wait() + raise AssertionError("cancelled lookup unexpectedly resumed") + + monkeypatch.setattr( + proxy_file_ops.FileAccountPinRepository, + "get_live_account_ids", + block_lookup, + ) + monkeypatch.setattr(service, "_stream_with_retry", upstream_stream) + monkeypatch.setattr( + service, + "_release_unsettled_stream_api_key_usage", + release_reservation, + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_forwarded_cancelled_lookup"}], + } + ) + + stream = service._stream_http_bridge_or_retry( + payload=payload, + headers={}, + codex_session_affinity=False, + propagate_http_errors=True, + openai_cache_affinity=False, + api_key=api_key, + api_key_reservation=reservation, + suppress_text_done_events=False, + forwarded_request=True, + forwarded_file_owner_account_id="acc_forwarded_owner", + ) + request_task = asyncio.ensure_future(anext(stream)) + await asyncio.wait_for(lookup_started.wait(), timeout=1) + request_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await request_task + + release_reservation.assert_not_awaited() + upstream_stream.assert_not_called() + + +@pytest.mark.asyncio +async def test_forwarded_receiver_waits_for_service_cleanup_handoff_before_streaming_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_forwarded_cleanup_handoff", + key_id="key_forwarded_cleanup_handoff", + model="gpt-5.1", + ) + stream_started = asyncio.Event() + pre_handoff_event_sent = asyncio.Event() + allow_cleanup_handoff = asyncio.Event() + allow_first_event = asyncio.Event() + + async def forwarded_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + stream_started.set() + pre_handoff_event_sent.set() + yield proxy_service.CODEX_KEEPALIVE_FRAME + await allow_cleanup_handoff.wait() + proxy_support._signal_propagated_responses_service_cleanup_ready() + await allow_first_event.wait() + yield ( + 'data: {"type":"response.completed","response":' + '{"id":"resp_forwarded_cleanup_handoff","status":"completed"}}\n\n' + ) + + service = SimpleNamespace(stream_http_responses=forwarded_stream) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_HTTP_BRIDGE_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_CAPACITY_STARTUP_SIGNAL_DISCOVERY_SECONDS", 0.01) + request = Request( + { + "type": "http", + "method": "POST", + "path": "/internal/bridge/responses", + "headers": [], + } + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_forwarded_cleanup_handoff"}], + "stream": True, + } + ) + + response_task = asyncio.create_task( + proxy_api._stream_responses( + request, + payload, + context=cast(proxy_api.ProxyContext, SimpleNamespace(service=service)), + api_key=None, + prefer_http_bridge=True, + skip_limit_enforcement=True, + api_key_reservation_override=reservation, + include_rate_limit_headers=False, + forwarded_request=True, + forwarded_headers={}, + forwarded_file_owner_account_id="acc_forwarded_cleanup_handoff", + enforce_openai_sdk_contract=False, + ) + ) + try: + await asyncio.wait_for(stream_started.wait(), timeout=1) + await asyncio.wait_for(pre_handoff_event_sent.wait(), timeout=1) + await asyncio.sleep(0.05) + assert not response_task.done() + + allow_cleanup_handoff.set() + response = await asyncio.wait_for(response_task, timeout=1) + assert isinstance(response, StreamingResponse) + + allow_first_event.set() + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunk.decode() if isinstance(chunk, bytes) else str(chunk) for chunk in chunks) + assert proxy_service.CODEX_KEEPALIVE_FRAME in body + assert "resp_forwarded_cleanup_handoff" in body + finally: + allow_cleanup_handoff.set() + allow_first_event.set() + if not response_task.done(): + response_task.cancel() + await asyncio.gather(response_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_http_bridge_session_submit_signals_service_cleanup_handoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + settings = _make_proxy_settings() + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_cleanup_handoff", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + event_queue=asyncio.Queue(), + request_text='{"type":"response.create"}', + transport="http", + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-cleanup-handoff", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_bridge_cleanup_handoff"), + upstream=AsyncMock(), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque([request_state]), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=1, + last_used_at=0.0, + idle_ttl_seconds=30.0, + ) + cleanup_ready = asyncio.Event() + cleanup_token = proxy_support._bind_propagated_responses_service_cleanup_ready(cleanup_ready) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) + detach = AsyncMock() + monkeypatch.setattr(service, "_detach_http_bridge_request", detach) + events = service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data='{"type":"response.create"}', + queue_limit=10, + propagate_http_errors=True, + downstream_turn_state=None, + ) + next_event = asyncio.create_task(anext(events)) + try: + await asyncio.wait_for(cleanup_ready.wait(), timeout=0.1) + finally: + next_event.cancel() + await asyncio.gather(next_event, return_exceptions=True) + await events.aclose() + proxy_support._reset_propagated_responses_service_cleanup_ready(cleanup_token) + + detach.assert_awaited_once_with(session, request_state=request_state) + + +@pytest.mark.asyncio +async def test_forwarded_owner_response_ready_signals_origin_cleanup_handoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + settings = _make_proxy_settings() + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_owner_response_cleanup_handoff", + key_id="key_owner_response_cleanup_handoff", + model="gpt-5.1", + ) + response_ready = asyncio.Event() + allow_event = asyncio.Event() + + async def owner_stream(*_args: object, **kwargs: object) -> AsyncIterator[str]: + on_response_ready = cast(Callable[[], None], kwargs["on_response_ready"]) + on_response_ready() + response_ready.set() + await allow_event.wait() + yield ( + 'data: {"type":"response.completed","response":' + '{"id":"resp_owner_response_cleanup_handoff","status":"completed"}}\n\n' + ) + + service._http_bridge_owner_client = SimpleNamespace(stream_responses=owner_stream) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + cleanup_ready = asyncio.Event() + cleanup_token = proxy_support._bind_propagated_responses_service_cleanup_ready(cleanup_ready) + owner_forward = proxy_service._HTTPBridgeOwnerForward( + owner_instance="owner-instance", + owner_endpoint="http://owner.invalid", + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-owner-cleanup-handoff", None), + ) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_owner_cleanup_handoff"}], + "stream": True, + } + ) + stream = cast( + AsyncGenerator[str, None], + service._forward_http_bridge_request_to_owner( + owner_forward=owner_forward, + payload=payload, + headers={}, + api_key_reservation=reservation, + codex_session_affinity=False, + downstream_turn_state=None, + request_started_at=time.monotonic(), + proxy_api_authorization=None, + file_owner_account_id="acc_owner_cleanup_handoff", + ), + ) + next_event = asyncio.ensure_future(anext(stream)) + try: + await asyncio.wait_for(response_ready.wait(), timeout=1) + await asyncio.wait_for(cleanup_ready.wait(), timeout=0.1) + allow_event.set() + assert "resp_owner_response_cleanup_handoff" in await asyncio.wait_for(next_event, timeout=1) + finally: + allow_event.set() + if not next_event.done(): + next_event.cancel() + await asyncio.gather(next_event, return_exceptions=True) + await stream.aclose() + proxy_support._reset_propagated_responses_service_cleanup_ready(cleanup_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_status", [None, 503], ids=["ambiguous-transport", "definitive-non-200"]) +async def test_owner_forward_dispatch_callbacks_classify_response_outcome( + monkeypatch: pytest.MonkeyPatch, + response_status: int | None, +) -> None: + settings = SimpleNamespace( + upstream_connect_timeout_seconds=1.0, + stream_idle_timeout_seconds=1.0, + ) + events: list[str] = [] + + class _Response: + status = 503 + + async def text(self) -> str: + events.append("response-text") + return "owner unavailable" + + class _RequestContext: + async def __aenter__(self) -> _Response: + events.append("request-enter") + if response_status is None: + raise aiohttp.ClientConnectionError("lost owner response") + return _Response() + + async def __aexit__(self, *_args: object) -> None: + events.append("request-exit") + + class _Session: + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + def post(self, *_args: object, **_kwargs: object) -> _RequestContext: + events.append("request-built") + return _RequestContext() + + monkeypatch.setattr(proxy_http_bridge_forwarding, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_http_bridge_forwarding, "build_owner_forward_headers", lambda **_kwargs: {}) + monkeypatch.setattr(proxy_http_bridge_forwarding.aiohttp, "ClientSession", lambda **_kwargs: _Session()) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": "hello", + "stream": True, + } + ) + context = proxy_http_bridge_forwarding.HTTPBridgeForwardContext( + origin_instance="origin", + target_instance="owner", + codex_session_affinity=False, + downstream_turn_state=None, + ) + stream = proxy_http_bridge_forwarding.HTTPBridgeOwnerClient().stream_responses( + owner_endpoint="http://owner.invalid", + payload=payload, + headers={}, + context=context, + request_started_at=time.monotonic(), + on_request_dispatched=lambda: events.append("dispatched"), + on_response_rejected=lambda: events.append("rejected"), + on_response_ready=lambda: events.append("ready"), + ) + + expected_error = aiohttp.ClientConnectionError if response_status is None else proxy_module.ProxyResponseError + with pytest.raises(expected_error): + await anext(stream) + + assert events[:3] == ["request-built", "dispatched", "request-enter"] + assert ("rejected" in events) is (response_status == 503) + assert "ready" not in events + + +@pytest.mark.asyncio +async def test_forwarded_receiver_cancellation_after_cleanup_handoff_releases_only_on_receiver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + settings.http_responses_session_bridge_enabled = True + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="reservation_cross_replica_cleanup_owner", + key_id="key_cross_replica_cleanup_owner", + model="gpt-5.1", + ) + origin_release = AsyncMock() + receiver_release = AsyncMock() + receiver_cleanup_ready = asyncio.Event() + receiver_first_event = asyncio.Event() + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "read", + "input": [{"type": "input_file", "file_id": "file_cross_replica_cleanup_owner"}], + "stream": True, + } + ) + + async def receiver_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + try: + proxy_support._signal_propagated_responses_service_cleanup_ready() + receiver_cleanup_ready.set() + await receiver_first_event.wait() + yield ( + 'data: {"type":"response.completed","response":' + '{"id":"resp_cross_replica_cleanup_owner","status":"completed"}}\n\n' + ) + finally: + await receiver_release(reservation) + + receiver_context = cast( + proxy_api.ProxyContext, + SimpleNamespace(service=SimpleNamespace(stream_http_responses=receiver_stream)), + ) + receiver_request = Request( + { + "type": "http", + "method": "POST", + "path": "/internal/bridge/responses", + "headers": [], + } + ) + + class _HermeticOwnerClient: + async def stream_responses( + self, + *, + on_response_ready: Callable[[], None], + **_kwargs: object, + ) -> AsyncIterator[str]: + receiver_task = asyncio.create_task( + proxy_api._stream_responses( + receiver_request, + payload.model_copy(deep=True), + context=receiver_context, + api_key=None, + prefer_http_bridge=True, + skip_limit_enforcement=True, + api_key_reservation_override=reservation, + include_rate_limit_headers=False, + forwarded_request=True, + forwarded_headers={}, + forwarded_file_owner_account_id="acc_cross_replica_cleanup_owner", + enforce_openai_sdk_contract=False, + ), + context=contextvars.Context(), + ) + receiver_response = await receiver_task + assert isinstance(receiver_response, StreamingResponse) + on_response_ready() + receiver_body = cast(AsyncGenerator[Any, None], receiver_response.body_iterator) + try: + async for chunk in receiver_body: + yield chunk.decode() if isinstance(chunk, bytes) else str(chunk) + finally: + await receiver_body.aclose() + + origin_service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + origin_service._http_bridge_owner_client = _HermeticOwnerClient() + owner_forward = proxy_service._HTTPBridgeOwnerForward( + owner_instance="owner-instance", + owner_endpoint="http://owner.invalid", + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-cross-replica-cleanup", None), + ) + + def origin_stream( + stream_payload: ResponsesRequest, + headers: Mapping[str, str], + **kwargs: object, + ) -> AsyncIterator[str]: + return origin_service._forward_http_bridge_request_to_owner( + owner_forward=owner_forward, + payload=stream_payload, + headers=headers, + api_key_reservation=cast( + proxy_service.ApiKeyUsageReservationData | None, + kwargs["api_key_reservation"], + ), + codex_session_affinity=False, + downstream_turn_state=None, + request_started_at=time.monotonic(), + proxy_api_authorization=None, + file_owner_account_id="acc_cross_replica_cleanup_owner", + ) + + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_api, "_HTTP_BRIDGE_STARTUP_ERROR_PROBE_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_CAPACITY_STARTUP_SIGNAL_DISCOVERY_SECONDS", 0.01) + monkeypatch.setattr(proxy_api, "_enforce_request_limits", AsyncMock(return_value=reservation)) + monkeypatch.setattr(proxy_api, "_rate_limit_headers_for_request", AsyncMock(return_value={})) + monkeypatch.setattr(proxy_api, "_release_reservation", origin_release) + origin_request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) + response = await proxy_api._stream_responses( + origin_request, + payload.model_copy(deep=True), + context=cast( + proxy_api.ProxyContext, SimpleNamespace(service=SimpleNamespace(stream_http_responses=origin_stream)) + ), + api_key=None, + prefer_http_bridge=True, + ) + assert isinstance(response, StreamingResponse) + await asyncio.wait_for(receiver_cleanup_ready.wait(), timeout=1) + body = cast(AsyncGenerator[Any, None], response.body_iterator) + body_task = asyncio.create_task(anext(body)) + await asyncio.sleep(0) + body_task.cancel() + try: + with pytest.raises(asyncio.CancelledError): + await body_task + await body.aclose() + receiver_release.assert_awaited_once_with(reservation) + origin_release.assert_not_awaited() + finally: + receiver_first_event.set() + if not body_task.done(): + body_task.cancel() + await asyncio.gather(body_task, return_exceptions=True) + await body.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("durable_owner", "forwarded_owner", "forwarded_request", "expected_error_code"), + [ + (None, "acc_forwarded_owner", False, "file_owner_unavailable"), + ("acc_reclaimed_owner", "acc_forwarded_owner", False, "continuity_owner_conflict"), + ("acc_forwarded_owner", None, True, "file_owner_unavailable"), + ], +) +async def test_stream_http_bridge_revalidates_forwarded_file_owner( + monkeypatch, + durable_owner, + forwarded_owner, + forwarded_request, + expected_error_code, +): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + settings = _make_proxy_settings() + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_service, + "_http_bridge_runtime_config", + lambda _dashboard_settings, _app_settings: proxy_service._HTTPBridgeRuntimeConfig( + enabled=False, + idle_ttl_seconds=30.0, + codex_idle_ttl_seconds=30.0, + max_sessions=8, + queue_limit=16, + prompt_cache_idle_ttl_seconds=30.0, + gateway_safe_mode=False, + ), + ) + resolve_owner = AsyncMock(return_value=durable_owner) + monkeypatch.setattr(service, "_resolve_file_account_for_responses", resolve_owner) + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "hi", + "input": [{"type": "input_file", "file_id": "file_forwarded_owner"}], + } + ) + + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + async for _ in service._stream_http_bridge_or_retry( + payload=payload, + headers={}, + codex_session_affinity=False, + propagate_http_errors=False, + openai_cache_affinity=False, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + forwarded_request=forwarded_request, + forwarded_file_owner_account_id=forwarded_owner, + ): + pass + + assert _proxy_error_code(exc_info.value) == expected_error_code + resolve_owner.assert_awaited_once_with(payload, {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("durable_owner", "forwarded_owner"), + [ + ("acc_doc", "acc_doc"), + (None, None), + ], + ids=["matching-proof", "opaque-no-pin-no-proof"], +) +async def test_stream_http_bridge_or_retry_routes_valid_forwarded_input_file( + monkeypatch, + durable_owner, + forwarded_owner, +): request_logs = _RequestLogsRecorder() service = proxy_service.ProxyService(_repo_factory(request_logs)) settings = _make_proxy_settings() @@ -33243,7 +36411,8 @@ async def test_stream_http_bridge_or_retry_routes_input_file_file_id_without_rej gateway_safe_mode=False, ), ) - await service._pin_file_account("file_doc", "acc_doc") + resolve_file_owner = AsyncMock(return_value=durable_owner) + monkeypatch.setattr(service, "_resolve_file_account_for_responses", resolve_file_owner) payload = ResponsesRequest.model_validate( { "model": "gpt-5.1", @@ -33252,17 +36421,18 @@ async def test_stream_http_bridge_or_retry_routes_input_file_file_id_without_rej } ) - calls: list[tuple[object, str | None]] = [] + calls: list[tuple[object, str | None, bool]] = [] async def fake_stream_with_retry( payload, headers, *, rewritten_file_account_id: str | None = None, + file_account_resolution_complete: bool = False, **kwargs, ): del headers, kwargs - calls.append((payload, rewritten_file_account_id)) + calls.append((payload, rewritten_file_account_id, file_account_resolution_complete)) yield "data: retry\n\n" monkeypatch.setattr(service, "_stream_with_retry", fake_stream_with_retry) @@ -33278,11 +36448,14 @@ async def fake_stream_with_retry( api_key=None, api_key_reservation=None, suppress_text_done_events=False, + forwarded_request=True, + forwarded_file_owner_account_id=forwarded_owner, ) ] assert output == ["data: retry\n\n"] - assert calls == [(payload, "acc_doc")] + assert calls == [(payload, durable_owner, True)] + resolve_file_owner.assert_awaited_once_with(payload, {}) def test_classify_upstream_close_rejected_only_for_clean_close_before_any_response_event(): @@ -34813,11 +37986,13 @@ async def test_files_create_persists_conversation_id_on_refresh_connection_reset account_b = _make_account("acc_files_create_refresh_b") record_error = AsyncMock() record_success = AsyncMock() + pin_file_account = AsyncMock() seen_excluded_account_ids: list[set[str]] = [] _install_two_account_selection(monkeypatch, service, account_a, account_b, seen_excluded_account_ids) monkeypatch.setattr(service._load_balancer, "record_error", record_error) monkeypatch.setattr(service._load_balancer, "record_success", record_success) + monkeypatch.setattr(service, "_pin_file_account", pin_file_account) monkeypatch.setattr( service, "_ensure_fresh", @@ -34847,6 +38022,7 @@ async def fake_create_file( assert seen_excluded_account_ids == [set(), {account_a.id}] record_error.assert_awaited_once_with(account_a) record_success.assert_awaited_once_with(account_b) + pin_file_account.assert_awaited_once_with("file_ok", account_b.id) assert await service.drain_persistence_tasks(timeout_seconds=1) assert request_logs.calls[0]["status"] == "success" assert request_logs.calls[0]["account_id"] == account_b.id @@ -34927,7 +38103,7 @@ async def select_account(**kwargs: object) -> AccountSelection: "core_finalize_file", AsyncMock(side_effect=AssertionError("strict file owner must not fail over or invoke upstream")), ) - await service._pin_file_account("file_pinned", account.id) + monkeypatch.setattr(service, "_resolve_file_account", AsyncMock(return_value=account.id)) with pytest.raises(proxy_module.ProxyResponseError) as exc_info: await service.finalize_file("file_pinned", {"session_id": "sid-files-finalize"}) @@ -34965,7 +38141,7 @@ async def select_account_with_budget(_deadline: float, **kwargs: object) -> Acco monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=fallback_account)) finalize_file = AsyncMock(side_effect=AssertionError("pinned finalize must not invoke a fallback account")) monkeypatch.setattr(proxy_service, "core_finalize_file", finalize_file) - await service._pin_file_account("file_pinned_initial", pinned_account.id) + monkeypatch.setattr(service, "_resolve_file_account", AsyncMock(return_value=pinned_account.id)) with pytest.raises(proxy_module.ProxyResponseError) as exc_info: await service.finalize_file("file_pinned_initial", {"session_id": "sid-files-finalize"})