diff --git a/.all-contributorsrc b/.all-contributorsrc index cd7180289b..b1bbcb8f86 100644 --- a/.all-contributorsrc +++ b/.all-contributorsrc @@ -1082,6 +1082,25 @@ "contributions": [ "code" ] + }, + { + "login": "shaqman", + "name": "Syakur Rahman", + "avatar_url": "https://avatars.githubusercontent.com/u/1113851?v=4", + "profile": "http://expressthisout.com/", + "contributions": [ + "code", + "test" + ] + }, + { + "login": "cigro-manager", + "name": "cigro-manager", + "avatar_url": "https://avatars.githubusercontent.com/u/219247995?v=4", + "profile": "https://github.com/cigro-manager", + "contributions": [ + "code" + ] } ], "contributorsPerLine": 7, diff --git a/README.md b/README.md index 17d995d0d6..b70dae66bd 100644 --- a/README.md +++ b/README.md @@ -265,6 +265,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/e alchemistkiv
alchemistkiv
💻 ⚠️ mhooooo
mhooooo
💻 ⚠️ crowscc
crowscc
💻 + Syakur Rahman
Syakur Rahman
💻 ⚠️ diff --git a/app/core/clients/proxy.py b/app/core/clients/proxy.py index eb8251bad3..1428184190 100644 --- a/app/core/clients/proxy.py +++ b/app/core/clients/proxy.py @@ -477,6 +477,7 @@ def __init__( upstream_status_code: int | None = None, upstream_error_code: str | None = None, failed_session: aiohttp.ClientSession | None = None, + retry_after_seconds: int | None = None, ) -> None: super().__init__(f"Proxy response error ({status_code})") self.status_code = status_code @@ -488,6 +489,7 @@ def __init__( self.upstream_status_code = upstream_status_code self.upstream_error_code = upstream_error_code self.failed_session = failed_session + self.retry_after_seconds = retry_after_seconds def _process_network_failure_error( diff --git a/app/core/clients/proxy_websocket.py b/app/core/clients/proxy_websocket.py index 06eb2d5fbb..0cec919728 100644 --- a/app/core/clients/proxy_websocket.py +++ b/app/core/clients/proxy_websocket.py @@ -138,6 +138,24 @@ def normalize_realtime_call_id(value: str) -> str | None: return normalized.lower() +def _consume_connection_lost_exception(done: asyncio.Future[Any]) -> None: + """Retrieve close exceptions before websockets shields the waiter. + + websockets 16 waits on ``connection_lost_waiter`` through + ``asyncio.shield`` while completing ``ClientConnection.recv``. A peer + keepalive/protocol close therefore leaves an exception on the waiter, + which asyncio reports as an ``exception in shielded future`` even though + ``recv`` translates it into an ``UpstreamWebSocketMessage``. Consume it + at the adapter boundary; ``receive`` still classifies the close normally. + """ + if done.cancelled(): + return + try: + done.exception() + except asyncio.CancelledError: + return + + @dataclass(slots=True) class UpstreamWebSocketMessage: kind: str @@ -170,7 +188,14 @@ def _relay_receive_error_code(error_code: str) -> str | None: # Relay owners map an absent code to their established stream_incomplete # contract. Leaking the adapter's generic fallback would bypass that path. - return error_code if error_code == PROCESS_NETWORK_UNAVAILABLE_CODE else None + return error_code if error_code in {PROCESS_NETWORK_UNAVAILABLE_CODE, "upstream_keepalive_timeout"} else None + + +def _is_keepalive_timeout_close(exc: ConnectionClosedError) -> bool: + """Classify peer/proxy heartbeat failures without exposing socket details.""" + + reason = _close_reason_from_exception(exc) + return "keepalive ping timeout" in f"{exc} {reason or ''}".lower() async def _rotate_after_websocket_network_failure(error_code: str) -> None: @@ -223,6 +248,9 @@ def __init__( self._connection = connection self._uses_proxy = uses_proxy self._preserve_close_semantics = preserve_close_semantics + connection_lost_waiter = getattr(connection, "connection_lost_waiter", None) + if isinstance(connection_lost_waiter, asyncio.Future): + connection_lost_waiter.add_done_callback(_consume_connection_lost_exception) async def send_text(self, text: str) -> None: try: @@ -254,6 +282,11 @@ async def receive(self) -> UpstreamWebSocketMessage: ) error_code = _websocket_transport_error_code(exc, uses_proxy=self._uses_proxy) await _rotate_after_websocket_network_failure(error_code) + relay_error_code = ( + "upstream_keepalive_timeout" + if _is_keepalive_timeout_close(exc) + else _relay_receive_error_code(error_code) + ) # ConnectionClosedError describes an incomplete close handshake, # not generic transport provenance. Let Responses relay owners map # it to stream_incomplete while live relays preserve received closes. @@ -265,7 +298,7 @@ async def receive(self) -> UpstreamWebSocketMessage: if self._preserve_close_semantics else str(exc) ), - error_code=_relay_receive_error_code(error_code), + error_code=relay_error_code, ) except Exception as exc: error_code = _websocket_transport_error_code(exc, uses_proxy=self._uses_proxy) diff --git a/app/core/config/settings.py b/app/core/config/settings.py index 0806512ac5..9df0a0bf94 100644 --- a/app/core/config/settings.py +++ b/app/core/config/settings.py @@ -299,6 +299,11 @@ class Settings(BaseSettings): http_responses_session_bridge_stuck_gate_retire_after_seconds: float = Field(default=300.0, gt=0) http_responses_session_bridge_max_sessions: int = Field(default=256, gt=0) http_responses_session_bridge_queue_limit: int = Field(default=8, gt=0) + http_responses_session_bridge_clean_close_retry_jitter_max_seconds: float = Field( + default=2.0, + ge=0, + le=30.0, + ) http_responses_session_bridge_gateway_safe_mode: bool = False http_responses_session_bridge_instance_id: str = Field(default_factory=_default_http_bridge_instance_id) http_responses_session_bridge_instance_ring: Annotated[list[str], NoDecode] = Field(default_factory=list) diff --git a/app/core/metrics/prometheus.py b/app/core/metrics/prometheus.py index dbc75f79f9..42a647e3df 100644 --- a/app/core/metrics/prometheus.py +++ b/app/core/metrics/prometheus.py @@ -172,6 +172,18 @@ def labels(self, *args: str, **kwargs: str) -> "HistogramLike": ... ["strength"], registry=REGISTRY, ) + bridge_handoff_compatibility_rejection_total = Counter( + "codex_lb_bridge_handoff_compatibility_rejection_total", + "Total closed HTTP bridge admission handoffs rejected for incompatible request context", + ["continuity_anchor", "preferred_account", "service_tier", "api_key_scope"], + registry=REGISTRY, + ) + bridge_unanchored_handoff_recovery_total = Counter( + "codex_lb_bridge_unanchored_handoff_recovery_total", + "Total stale closed HTTP bridge admission handoffs recovered without a continuity anchor", + ["reason"], + registry=REGISTRY, + ) bridge_local_rebind_total = Counter( "codex_lb_bridge_local_rebind_total", "Total bridge local rebinds by reason", @@ -264,6 +276,24 @@ def labels(self, *args: str, **kwargs: str) -> "HistogramLike": ... ["reason", "affinity_kind", "model_class"], registry=REGISTRY, ) + http_bridge_retry_circuit_total = Counter( + "codex_lb_http_bridge_retry_circuit_total", + "Total HTTP bridge automatic retry circuit outcomes", + ["outcome"], + registry=REGISTRY, + ) + stream_keepalive_sent_total = Counter( + "codex_lb_stream_keepalive_sent_total", + "Total downstream SSE keepalive frames emitted by surface", + ["surface"], + registry=REGISTRY, + ) + stream_idle_timeout_total = Counter( + "codex_lb_stream_idle_timeout_total", + "Total streams terminated after exceeding the configured idle window", + ["surface"], + registry=REGISTRY, + ) cache_invalidation_bump_failures_total = Counter( "codex_lb_cache_invalidation_bump_failures_total", "Total cache invalidation version bumps that failed after retries", @@ -315,6 +345,8 @@ def mark_process_dead() -> None: bridge_first_turn_timeout_total: CounterLike | None = None bridge_drain_recovery_allowed_total: CounterLike | None = None bridge_owner_mismatch_total: CounterLike | None = None + bridge_handoff_compatibility_rejection_total: CounterLike | None = None + bridge_unanchored_handoff_recovery_total: CounterLike | None = None bridge_local_rebind_total: CounterLike | None = None bridge_forward_latency_seconds: HistogramLike | None = None bridge_public_contract_error_total: CounterLike | None = None @@ -329,6 +361,9 @@ def mark_process_dead() -> None: proxy_phase_latency_seconds: HistogramLike | None = None http_bridge_prewarm_total: CounterLike | None = None http_bridge_stuck_retire_total: CounterLike | None = None + http_bridge_retry_circuit_total: CounterLike | None = None + stream_keepalive_sent_total: CounterLike | None = None + stream_idle_timeout_total: CounterLike | None = None cache_invalidation_bump_failures_total: CounterLike | None = None cache_invalidation_poll_failures_total: CounterLike | None = None @@ -355,6 +390,7 @@ def mark_process_dead() -> None: "bridge_durable_recover_total", "bridge_drain_recovery_allowed_total", "bridge_first_turn_timeout_total", + "bridge_handoff_compatibility_rejection_total", "bridge_local_rebind_total", "bridge_owner_forward_total", "bridge_owner_mismatch_total", @@ -363,6 +399,7 @@ def mark_process_dead() -> None: "bridge_reattach_total", "bridge_same_account_takeover_total", "bridge_soft_local_rebind_total", + "bridge_unanchored_handoff_recovery_total", "cache_invalidation_bump_failures_total", "cache_invalidation_poll_failures_total", "cap_partition_replicas", @@ -370,7 +407,10 @@ def mark_process_dead() -> None: "continuity_fail_closed_total", "continuity_owner_resolution_total", "http_bridge_prewarm_total", + "http_bridge_retry_circuit_total", "http_bridge_stuck_retire_total", + "stream_keepalive_sent_total", + "stream_idle_timeout_total", "image_request_duration_seconds", "image_requests_total", "make_scrape_registry", diff --git a/app/core/utils/sse.py b/app/core/utils/sse.py index ad4e8993ec..b25d185d7f 100644 --- a/app/core/utils/sse.py +++ b/app/core/utils/sse.py @@ -3,7 +3,7 @@ import asyncio import json import re -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Callable, Mapping from app.core.errors import ResponseFailedEvent from app.core.types import JsonValue @@ -26,6 +26,7 @@ async def inject_sse_keepalives( interval_seconds: float, *, keepalive_frame: str = SSE_KEEPALIVE_FRAME, + on_keepalive: Callable[[], None] | None = None, ) -> AsyncIterator[str]: """Wrap an SSE event iterator and emit comment heartbeats on idle gaps. @@ -56,6 +57,8 @@ async def _next_chunk(it: AsyncIterator[str]) -> str: timeout=interval_seconds, ) except asyncio.TimeoutError: + if on_keepalive is not None: + on_keepalive() yield keepalive_frame continue except StopAsyncIteration: diff --git a/app/db/alembic/versions/20260717_000000_add_http_bridge_retry_circuits.py b/app/db/alembic/versions/20260717_000000_add_http_bridge_retry_circuits.py new file mode 100644 index 0000000000..b246db6525 --- /dev/null +++ b/app/db/alembic/versions/20260717_000000_add_http_bridge_retry_circuits.py @@ -0,0 +1,47 @@ +"""add durable HTTP bridge retry circuit state + +Revision ID: 20260717_000000_add_http_bridge_retry_circuits +Revises: 20260717_000000_optimize_dashboard_hot_path_indexes +Create Date: 2026-07-17 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "20260717_000000_add_http_bridge_retry_circuits" +down_revision = "20260717_000000_optimize_dashboard_hot_path_indexes" +branch_labels = None +depends_on = None + +_TABLE_NAME = "http_bridge_retry_circuits" + + +def upgrade() -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + if inspector.has_table(_TABLE_NAME): + return + op.create_table( + _TABLE_NAME, + sa.Column("session_key_kind", sa.String(length=64), nullable=False), + sa.Column("session_key_hash", sa.String(length=64), nullable=False), + sa.Column("api_key_scope", sa.String(length=255), nullable=False), + sa.Column("consecutive_failures", sa.Integer(), nullable=False, server_default="0"), + sa.Column("cooldown_until_epoch", sa.Float(), nullable=False, server_default="0"), + sa.Column("last_detail", sa.String(length=255), nullable=True), + sa.Column("updated_at_epoch", sa.Float(), nullable=False), + sa.PrimaryKeyConstraint( + "session_key_kind", + "session_key_hash", + "api_key_scope", + name="pk_http_bridge_retry_circuits", + ), + ) + + +def downgrade() -> None: + bind = op.get_bind() + if sa.inspect(bind).has_table(_TABLE_NAME): + op.drop_table(_TABLE_NAME) diff --git a/app/db/alembic/versions/20260717_000001_merge_retry_circuits_and_dashboard_indexes.py b/app/db/alembic/versions/20260717_000001_merge_retry_circuits_and_dashboard_indexes.py new file mode 100644 index 0000000000..367cbb933f --- /dev/null +++ b/app/db/alembic/versions/20260717_000001_merge_retry_circuits_and_dashboard_indexes.py @@ -0,0 +1,22 @@ +"""merge retry circuits and dashboard index heads + +Revision ID: 20260717_000001_merge_retry_circuits_and_dashboard_indexes +Revises: +- 20260717_000000_add_http_bridge_retry_circuits +Create Date: 2026-07-17 00:00:00.000000 +""" + +from __future__ import annotations + +revision = "20260717_000001_merge_retry_circuits_and_dashboard_indexes" +down_revision = ("20260717_000000_add_http_bridge_retry_circuits",) +branch_labels = None +depends_on = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/app/db/alembic/versions/20260724_000000_merge_request_log_schema_heads.py b/app/db/alembic/versions/20260724_000000_merge_request_log_schema_heads.py new file mode 100644 index 0000000000..963bb7bf58 --- /dev/null +++ b/app/db/alembic/versions/20260724_000000_merge_request_log_schema_heads.py @@ -0,0 +1,33 @@ +"""merge the deployed retry-circuit and request-log schema heads + +Revision ID: 20260724_000000_merge_request_log_schema_heads +Revises: +- 20260717_000001_merge_retry_circuits_and_dashboard_indexes +- 20260722_000000_backfill_request_log_useragent_families +Create Date: 2026-07-24 00:00:00.000000 + +The deployed SQLite database was previously stamped at the retry-circuit +merge revision while the request-log conversation-id branch was not applied. +Keeping the retry-circuit merge and the request-usage rollup as parents lets +Alembic apply the missing request-log revisions before converging on one head +without leaving the rollup revision as a second head or rewriting migration +history. +""" + +from __future__ import annotations + +revision = "20260724_000000_merge_request_log_schema_heads" +down_revision = ( + "20260717_000001_merge_retry_circuits_and_dashboard_indexes", + "20260724_000000_add_request_usage_time_rollups", +) +branch_labels = None +depends_on = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/app/db/alembic/versions/20260726_000000_repair_request_usage_rollups_after_merge.py b/app/db/alembic/versions/20260726_000000_repair_request_usage_rollups_after_merge.py new file mode 100644 index 0000000000..1ef9c5a100 --- /dev/null +++ b/app/db/alembic/versions/20260726_000000_repair_request_usage_rollups_after_merge.py @@ -0,0 +1,37 @@ +"""repair request-usage rollups for databases stamped at the old merge head + +Revision ID: 20260726_000000_repair_request_usage_rollups_after_merge +Revises: 20260724_000000_merge_request_log_schema_heads +Create Date: 2026-07-26 + +Some deployed databases were stamped at the request-log merge revision before +the request-usage rollup child was connected to that merge. Changing the +parent tuple cannot make Alembic replay an already-applied revision, so those +databases need a forward-only repair step. The canonical rollup migration is +idempotent and safely creates any missing tables or watermark column here. +""" + +from __future__ import annotations + +import importlib +from types import ModuleType + +revision = "20260726_000000_repair_request_usage_rollups_after_merge" +down_revision = "20260724_000000_merge_request_log_schema_heads" +branch_labels = None +depends_on = None + + +def _rollup_migration() -> ModuleType: + return importlib.import_module("app.db.alembic.versions.20260724_000000_add_request_usage_time_rollups") + + +def upgrade() -> None: + _rollup_migration().upgrade() + + +def downgrade() -> None: + # This revision repairs databases that were already stamped at the merge + # head. It must never remove objects owned by the canonical rollup + # revision, which remains an ancestor of that merge head on fresh installs. + pass diff --git a/app/db/alembic/versions/20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads.py b/app/db/alembic/versions/20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads.py new file mode 100644 index 0000000000..f9c39f86a6 --- /dev/null +++ b/app/db/alembic/versions/20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads.py @@ -0,0 +1,31 @@ +"""merge pending-tool-call and stamped-rollup-repair heads + +Revision ID: 20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads +Revises: +- 20260725_000000_add_http_bridge_pending_tool_calls +- 20260726_000000_repair_request_usage_rollups_after_merge +Create Date: 2026-07-28 + +The pending-tool-call migration was added on a branch from the canonical +request-usage rollup revision, while the forward-only repair migration starts +from the deployed request-log merge. Keep both paths intact and converge them +before startup asks Alembic for ``head``. +""" + +from __future__ import annotations + +revision = "20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads" +down_revision = ( + "20260725_000000_add_http_bridge_pending_tool_calls", + "20260726_000000_repair_request_usage_rollups_after_merge", +) +branch_labels = None +depends_on = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/app/db/alembic/versions/20260730_000000_add_http_bridge_recovery_attempts.py b/app/db/alembic/versions/20260730_000000_add_http_bridge_recovery_attempts.py new file mode 100644 index 0000000000..fe2ad670a0 --- /dev/null +++ b/app/db/alembic/versions/20260730_000000_add_http_bridge_recovery_attempts.py @@ -0,0 +1,60 @@ +"""add durable HTTP bridge recovery attempts + +Revision ID: 20260730_000000_add_http_bridge_recovery_attempts +Revises: 20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads +Create Date: 2026-07-30 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.engine import Connection + +revision = "20260730_000000_add_http_bridge_recovery_attempts" +down_revision = "20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads" +branch_labels = None +depends_on = None + +_TABLE = "http_bridge_recovery_attempts" +_STATE = sa.Enum("unknown", "replayed", name="http_bridge_recovery_attempt_state") + + +def _has_table(connection: Connection) -> bool: + return sa.inspect(connection).has_table(_TABLE) + + +def upgrade() -> None: + bind = op.get_bind() + if _has_table(bind): + return + op.create_table( + _TABLE, + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("session_id", sa.String(36), nullable=False), + sa.Column("request_fingerprint", sa.String(64), nullable=False), + sa.Column("request_id", sa.String(255), nullable=False), + sa.Column("account_id", sa.String(), nullable=True), + sa.Column("model", sa.String(), nullable=True), + sa.Column("replay_safe", sa.Boolean(), nullable=False, server_default=sa.text("false")), + sa.Column("state", _STATE, nullable=False, server_default="unknown"), + sa.Column("response_id", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(["session_id"], ["http_bridge_sessions.id"], ondelete="CASCADE"), + sa.UniqueConstraint( + "session_id", + "request_fingerprint", + name="uq_http_bridge_recovery_attempts_session_fingerprint", + ), + ) + op.create_index("idx_http_bridge_recovery_attempts_state", _TABLE, ["state", "updated_at"]) + + +def downgrade() -> None: + bind = op.get_bind() + if not _has_table(bind): + return + op.drop_index("idx_http_bridge_recovery_attempts_state", table_name=_TABLE) + op.drop_table(_TABLE) + _STATE.drop(bind, checkfirst=True) diff --git a/app/db/alembic/versions/20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads.py b/app/db/alembic/versions/20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads.py new file mode 100644 index 0000000000..9b5ee83c71 --- /dev/null +++ b/app/db/alembic/versions/20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads.py @@ -0,0 +1,30 @@ +"""merge HTTP bridge recovery and capability lineage heads + +Revision ID: 20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads +Revises: +- 20260730_000000_add_http_bridge_recovery_attempts +- 20260731_000000_add_capability_lineage_markers +Create Date: 2026-08-03 + +Both migrations are additive and were introduced from independent branches. +This no-op merge records their convergence so startup and migration checks see +one canonical Alembic head. +""" + +from __future__ import annotations + +revision = "20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads" +down_revision = ( + "20260730_000000_add_http_bridge_recovery_attempts", + "20260731_000000_add_capability_lineage_markers", +) +branch_labels = None +depends_on = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/app/db/models.py b/app/db/models.py index 521a47d02b..ae98a281e2 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -1682,6 +1682,11 @@ class HttpBridgeSessionState(str, Enum): CLOSED = "closed" +class HttpBridgeRecoveryAttemptState(str, Enum): + UNKNOWN = "unknown" + REPLAYED = "replayed" + + class HttpBridgeSessionRecord(Base): __tablename__ = "http_bridge_sessions" @@ -1751,6 +1756,49 @@ class HttpBridgeSessionRecord(Base): ) +class HttpBridgeRecoveryAttemptRecord(Base): + __tablename__ = "http_bridge_recovery_attempts" + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) + session_id: Mapped[str] = mapped_column( + String(36), + ForeignKey("http_bridge_sessions.id", ondelete="CASCADE"), + nullable=False, + ) + request_fingerprint: Mapped[str] = mapped_column(String(64), nullable=False) + request_id: Mapped[str] = mapped_column(String(255), nullable=False) + account_id: Mapped[str | None] = mapped_column(String, nullable=True) + model: Mapped[str | None] = mapped_column(String, nullable=True) + replay_safe: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("false")) + state: Mapped[HttpBridgeRecoveryAttemptState] = mapped_column( + SqlEnum( + HttpBridgeRecoveryAttemptState, + name="http_bridge_recovery_attempt_state", + validate_strings=True, + values_callable=_enum_values, + ), + default=HttpBridgeRecoveryAttemptState.UNKNOWN, + server_default=text("'unknown'"), + nullable=False, + ) + response_id: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=func.now(), server_default=func.now() + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=func.now(), server_default=func.now(), onupdate=func.now() + ) + + __table_args__ = ( + UniqueConstraint( + "session_id", + "request_fingerprint", + name="uq_http_bridge_recovery_attempts_session_fingerprint", + ), + Index("idx_http_bridge_recovery_attempts_state", "state", "updated_at"), + ) + + class HttpBridgeSessionAlias(Base): __tablename__ = "http_bridge_session_aliases" @@ -1793,6 +1841,28 @@ class HttpBridgeSessionAlias(Base): ) +class HttpBridgeRetryCircuit(Base): + __tablename__ = "http_bridge_retry_circuits" + + session_key_kind: Mapped[str] = mapped_column(String(64), primary_key=True) + session_key_hash: Mapped[str] = mapped_column(String(64), primary_key=True) + api_key_scope: Mapped[str] = mapped_column(String(255), primary_key=True) + consecutive_failures: Mapped[int] = mapped_column( + Integer, + nullable=False, + default=0, + server_default=text("0"), + ) + cooldown_until_epoch: Mapped[float] = mapped_column( + Float, + nullable=False, + default=0.0, + server_default=text("0"), + ) + last_detail: Mapped[str | None] = mapped_column(String(255), nullable=True) + updated_at_epoch: Mapped[float] = mapped_column(Float, nullable=False) + + _PRIMARY_WINDOW_INDEX_EXPR = func.coalesce(UsageHistory.window, literal_column("'primary'")) Index("idx_usage_recorded_at", UsageHistory.recorded_at) diff --git a/app/main.py b/app/main.py index 4b59dbe295..37fcd1632f 100644 --- a/app/main.py +++ b/app/main.py @@ -518,7 +518,22 @@ async def _activate_bridge_membership(svc: RingMembershipService, iid: str) -> N logger.warning("Drain timeout reached, proceeding with shutdown") proxy_service = getattr(app.state, "proxy_service", None) - if proxy_service is not None and hasattr(proxy_service, "mark_http_bridge_draining"): + recovery_settlements_drained = True + # Settle detached recovery journals while their origin leases are + # still held; bridge teardown below may release those owner fences. + if proxy_service is not None and hasattr(proxy_service, "drain_persistence_tasks"): + try: + recovery_settlements_drained = await proxy_service.drain_persistence_tasks( + timeout_seconds=settings.shutdown_drain_timeout_seconds, + task_name_prefixes=("http-bridge-recovery-settlement-",), + ) + except Exception: + logger.warning("Failed to pre-drain proxy settlement tasks during shutdown", exc_info=True) + if ( + recovery_settlements_drained + and proxy_service is not None + and hasattr(proxy_service, "mark_http_bridge_draining") + ): try: await proxy_service.mark_http_bridge_draining() except Exception: diff --git a/app/modules/proxy/_service/http_bridge/activity.py b/app/modules/proxy/_service/http_bridge/activity.py index 27608f44b9..10cce3d065 100644 --- a/app/modules/proxy/_service/http_bridge/activity.py +++ b/app/modules/proxy/_service/http_bridge/activity.py @@ -1,15 +1,82 @@ from __future__ import annotations +from typing import Any + from app.modules.proxy._service.http_bridge.helpers import ( + _close_http_bridge_session_bounded, _http_bridge_pending_count_nowait, + _http_bridge_pending_state_is_stale, _http_bridge_request_counts_against_queue, + _log_http_bridge_event, + _raise_http_bridge_incompatible_admission_handoff, + _record_http_bridge_unanchored_handoff_recovery, http_bridge_activity_snapshot_nowait, ) from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol -from app.modules.proxy._service.support import _HTTPBridgeSession +from app.modules.proxy._service.support import _http_bridge_session_supports_service_tier, _HTTPBridgeSession +from app.modules.proxy.affinity import _extract_model_class class _HTTPBridgeActivityMixin: + _http_bridge_pending_state_is_stale = staticmethod(_http_bridge_pending_state_is_stale) + + def _recover_http_bridge_incompatible_admission_handoff( + self: Any, + key: Any, + existing: Any, + force_durable_takeover: bool, + original_request_unanchored: bool, + request_model: str | None, + api_key: Any, + incoming_turn_state: str | None, + previous_response_id: str | None, + preferred_account_id: str | None, + require_preferred_account: bool, + request_service_tier: str | None, + ) -> tuple[Any, bool]: + if original_request_unanchored and existing is not None: + detached = self._detach_http_bridge_session_locked(key, expected_session=existing) + if detached is not None: + force_durable_takeover = True + _record_http_bridge_unanchored_handoff_recovery(reason="closed_admission_handoff") + _log_http_bridge_event( + "unanchored_handoff_recovery", + key, + account_id=detached.account.id, + model=request_model, + detail="outcome=retired_closed_admission_handoff", + cache_key_family=key.affinity_kind, + model_class=_extract_model_class(request_model) if request_model else None, + owner_check_applied=False, + ) + self._schedule_http_bridge_session_closes([detached], reason="unanchored_handoff_recovery") + return None, force_durable_takeover + + _raise_http_bridge_incompatible_admission_handoff( + session=existing, + key=key, + api_key=api_key, + incoming_turn_state=incoming_turn_state, + previous_response_id=previous_response_id, + preferred_account_id=preferred_account_id, + require_preferred_account=require_preferred_account, + request_service_tier=request_service_tier, + service_tier_supported=_http_bridge_session_supports_service_tier( + existing, + request_model=request_model, + request_service_tier=request_service_tier, + ), + ) + raise AssertionError("incompatible admission handoff must raise") + + async def _close_http_bridge_session_bounded( + self: Any, + session: _HTTPBridgeSession, + *, + reason: str, + ) -> None: + await _close_http_bridge_session_bounded(self, session, reason=reason) + async def _http_bridge_pending_count( self: _HTTPBridgeServiceProtocol, session: _HTTPBridgeSession, diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index ca98a7fa9f..59eca389e0 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -48,8 +48,10 @@ PROMETHEUS_AVAILABLE, bridge_drain_recovery_allowed_total, bridge_first_turn_timeout_total, + bridge_handoff_compatibility_rejection_total, bridge_instance_mismatch_total, bridge_reattach_total, + bridge_unanchored_handoff_recovery_total, http_bridge_prewarm_total, http_bridge_stuck_retire_total, ) @@ -120,6 +122,7 @@ ) from app.modules.proxy._service.support import ( _HARD_HTTP_BRIDGE_AFFINITY_KINDS, # noqa: F401 + _REQUEST_TRANSPORT_HTTP, _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 _http_bridge_session_supports_service_tier, _HTTPBridgeSession, @@ -187,8 +190,15 @@ ) logger = logging.getLogger("app.modules.proxy.service") +_TASK_CANCEL_TIMEOUT_SECONDS = 1.0 +_TaskResultT = TypeVar("_TaskResultT") +_HTTP_BRIDGE_PENDING_COUNT_WARNING_INTERVAL_SECONDS = 60.0 +_http_bridge_pending_count_warning_last_logged: dict[tuple[str, str, str], float] = {} _HTTP_BRIDGE_BACKGROUND_CLOSE_TIMEOUT_SECONDS = 5.0 -_HTTP_BRIDGE_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS = 240.0 +# A healthy upstream acknowledges response.create promptly. Keep the +# Keep the owner-side watchdog within the client-safe contract while honoring +# the configured stuck-gate threshold when it is shorter. +_HTTP_BRIDGE_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS = 60.0 _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL = "missing_response_created_timeout" T = TypeVar("T") @@ -306,14 +316,28 @@ def _http_bridge_pending_count_nowait( except Exception as exc: if type(exc).__name__ not in {"WouldBlock", "RuntimeError"}: raise - logger.warning( - "http_bridge_pending_count_unavailable context=%s bridge_kind=%s bridge_key=%s account_id=%s model=%s", + warning_key = ( context, session.key.affinity_kind, _hash_identifier(session.key.affinity_key), - session.account.id, - session.request_model, ) + now = time.monotonic() + last_logged = _http_bridge_pending_count_warning_last_logged.get(warning_key) + if last_logged is None or now - last_logged >= _HTTP_BRIDGE_PENDING_COUNT_WARNING_INTERVAL_SECONDS: + _http_bridge_pending_count_warning_last_logged[warning_key] = now + if len(_http_bridge_pending_count_warning_last_logged) > 2048: + cutoff = now - _HTTP_BRIDGE_PENDING_COUNT_WARNING_INTERVAL_SECONDS * 2 + for key, timestamp in tuple(_http_bridge_pending_count_warning_last_logged.items()): + if timestamp < cutoff: + _http_bridge_pending_count_warning_last_logged.pop(key, None) + logger.warning( + "http_bridge_pending_count_unavailable context=%s bridge_kind=%s bridge_key=%s account_id=%s model=%s", + context, + session.key.affinity_kind, + warning_key[2], + session.account.id, + session.request_model, + ) return None try: request_counts_against_queue = _service_global("_http_bridge_request_counts_against_queue") @@ -390,6 +414,14 @@ def _cleanup_http_bridge_inflight_sessions_nowait(service: Any) -> dict[str, int } +def _http_bridge_inflight_creation_count(service: Any) -> int: + return sum( + 1 + for future in service._http_bridge_inflight_sessions.values() + if not getattr(future, "_http_bridge_handoff", False) + ) + + def http_bridge_activity_snapshot_nowait(service: Any) -> dict[str, int | bool]: inflight_cleanup = _cleanup_http_bridge_inflight_sessions_nowait(service) live_sessions = 0 @@ -634,6 +666,51 @@ def _normalize_http_bridge_error_event( return normalized_event_block, normalized_payload, parsed_event, "response.failed" +def _http_bridge_pending_state_is_stale( + request_state: _WebSocketRequestState, + *, + now: float, + threshold_seconds: float, + session_closed: bool = False, +) -> bool: + """Identify pre-created bridge requests that never received upstream activity.""" + if request_state.transport != _REQUEST_TRANSPORT_HTTP or request_state.skip_request_log: + return False + # Previous-response and turn-state identifiers preserve hard continuity + # while the socket is live. Once it closes, that continuity is already lost. + # A silent-but-open upstream must not preserve continuity forever: with + # zero upstream activity the anchored request wedges the create gate and + # every follow-up on the session, so cap the anchored wait at 2x the + # stuck-gate threshold before allowing retirement. + if not session_closed and (request_state.previous_response_id is not None or request_state.hard_continuity_anchor): + anchored_wait_started_at = ( + request_state.response_create_gate_wait_started_at + if request_state.response_create_gate_wait_started_at is not None + else request_state.started_at + ) + if max(0.0, now - anchored_wait_started_at) < threshold_seconds * 2.0: + return False + # Do not require the gate/awaiting flags: retry and reader re-enqueue + # paths re-register states whose admission flags were already cleared, + # and leaning on them starves the watchdog entirely (observed prod + # wedges 2026-07-20). A reattached stream can deliver events whose + # response.created was lost, so use the most recent upstream event as + # the silence clock instead of aging every active stream from its + # original gate wait. + if request_state.response_id is not None: + return False + if request_state.latency_response_created_ms is not None: + return False + wait_started_at = ( + request_state.response_create_gate_wait_started_at + if request_state.response_create_gate_wait_started_at is not None + else request_state.started_at + ) + if request_state.response_event_count > 0 and request_state.last_upstream_activity_at is not None: + wait_started_at = request_state.last_upstream_activity_at + return max(0.0, now - wait_started_at) >= threshold_seconds + + def _http_bridge_request_counts_against_queue(request_state: _WebSocketRequestState) -> bool: return not request_state.draining_until_terminal @@ -658,6 +735,10 @@ def _http_bridge_eventless_precreated_deadline( or request_state.last_downstream_sequence_number is not None ): return None + # Non-response telemetry (for example ``codex.rate_limits``) may update + # the generic activity marker, but it must not extend the response.create + # acknowledgement deadline. The eventless watchdog is intentionally + # anchored to the send time until a response-lifecycle event is observed. return sent_at + min( float(stuck_gate_retire_after_seconds), _HTTP_BRIDGE_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS, @@ -784,7 +865,11 @@ def _http_bridge_compatible( def _http_bridge_alias_target_is_stale(session: _HTTPBridgeSession | None) -> bool: - return session is None or session.closed or not _http_bridge_session_account_active(session) + return ( + session is None + or (session.closed and not session.handoff_in_progress) + or not _http_bridge_session_account_active(session) + ) def _http_bridge_incompatible_model_fork_key( @@ -1176,7 +1261,86 @@ def _require_http_bridge_bound_account_not_excluded( ) -def _raise_http_bridge_incompatible_admission_handoff() -> None: +def _record_http_bridge_handoff_compatibility_rejection( + *, + session: "_HTTPBridgeSession", + key: "_HTTPBridgeSessionKey", + api_key: ApiKeyData | None, + incoming_turn_state: str | None, + previous_response_id: str | None, + preferred_account_id: str | None, + require_preferred_account: bool, + request_service_tier: str | None, + service_tier_supported: bool, +) -> None: + continuity_anchor = ( + "previous_response_id" + if previous_response_id is not None + else "turn_state" + if incoming_turn_state is not None + else "none" + ) + preferred_account = ( + "required" if require_preferred_account else "preferred" if preferred_account_id is not None else "none" + ) + service_tier = ( + "not_requested" if request_service_tier is None else "supported" if service_tier_supported else "unsupported" + ) + api_key_scope = "scoped" if api_key is not None and api_key.id is not None else "unscoped" + if PROMETHEUS_AVAILABLE and bridge_handoff_compatibility_rejection_total is not None: + bridge_handoff_compatibility_rejection_total.labels( + continuity_anchor=continuity_anchor, + preferred_account=preferred_account, + service_tier=service_tier, + api_key_scope=api_key_scope, + ).inc() + logger.warning( + "http_bridge_handoff_compatibility_rejected bridge_kind=%s bridge_key=%s key_strength=%s " + "continuity_anchor=%s preferred_account=%s preferred_account_id=%s require_preferred=%s " + "service_tier=%s api_key_scope=%s existing_closed=%s admission_waiters=%s pending_count=%s", + key.affinity_kind, + _hash_identifier(key.affinity_key), + _http_bridge_key_strength(key), + continuity_anchor, + preferred_account, + _hash_identifier_or_none(preferred_account_id), + require_preferred_account, + service_tier, + api_key_scope, + session.closed, + getattr(session, "admission_waiter_count", 0), + len(session.pending_requests), + ) + + +def _record_http_bridge_unanchored_handoff_recovery(*, reason: str) -> None: + if PROMETHEUS_AVAILABLE and bridge_unanchored_handoff_recovery_total is not None: + bridge_unanchored_handoff_recovery_total.labels(reason=reason).inc() + + +def _raise_http_bridge_incompatible_admission_handoff( + *, + session: "_HTTPBridgeSession", + key: "_HTTPBridgeSessionKey", + api_key: ApiKeyData | None, + incoming_turn_state: str | None, + previous_response_id: str | None, + preferred_account_id: str | None, + require_preferred_account: bool, + request_service_tier: str | None, + service_tier_supported: bool, +) -> None: + _record_http_bridge_handoff_compatibility_rejection( + session=session, + key=key, + api_key=api_key, + incoming_turn_state=incoming_turn_state, + previous_response_id=previous_response_id, + preferred_account_id=preferred_account_id, + require_preferred_account=require_preferred_account, + request_service_tier=request_service_tier, + service_tier_supported=service_tier_supported, + ) raise ProxyResponseError( 503, openai_error( @@ -1456,6 +1620,101 @@ def _http_bridge_durable_lease_ttl_seconds() -> float: return float(RING_STALE_THRESHOLD_SECONDS) +def _http_bridge_durable_release_allowed(service: Any, session: Any) -> bool: + session_id = getattr(session, "durable_session_id", None) + owner_epoch = getattr(session, "durable_owner_epoch", None) + if session_id is None or owner_epoch is None: + return False + return not any( + not task.done() and getattr(task, "_http_bridge_recovery_session_id", None) == session_id + for task in service._background_cleanup_tasks + ) + + +async def _drain_cancelled_task(task: asyncio.Task[Any]) -> None: + await asyncio.gather(task, return_exceptions=True) + + +def _cancel_and_track_cancelled_task( + task: asyncio.Task[Any], + *, + label: str, + cleanup_tasks: set[asyncio.Task[None]] | None, + cancel_task: bool = True, +) -> None: + if cancel_task: + task.cancel() + cleanup_task = asyncio.create_task(_drain_cancelled_task(task), name=f"cancelled-task-cleanup-{label}") + if cleanup_tasks is not None: + cleanup_tasks.add(cleanup_task) + cleanup_task.add_done_callback(cleanup_tasks.discard) + + +async def _await_cancelled_task( + task: asyncio.Task[_TaskResultT], + *, + timeout_seconds: float = _TASK_CANCEL_TIMEOUT_SECONDS, + label: str, + cleanup_tasks: set[asyncio.Task[None]] | None = None, +) -> bool: + caller_task = asyncio.current_task() + # Give a new child one scheduling turn before cancellation so + # cancellation-resistant tasks enter the deferred-drain path. + if not task.done(): + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + _cancel_and_track_cancelled_task(task, label=label, cleanup_tasks=cleanup_tasks) + raise + task.cancel() + try: + await asyncio.wait_for(asyncio.shield(task), timeout=timeout_seconds) + except asyncio.CancelledError: + if caller_task is not None and caller_task.cancelling(): + _cancel_and_track_cancelled_task(task, label=label, cleanup_tasks=cleanup_tasks, cancel_task=False) + raise + return True + except TimeoutError: + logger.warning("Timed out waiting for %s cancellation", label) + _cancel_and_track_cancelled_task(task, label=label, cleanup_tasks=cleanup_tasks, cancel_task=False) + return False + return True + + +async def _persist_http_bridge_replacement_account( + service: _HTTPBridgeServiceProtocol, + session: _HTTPBridgeSession, + account_id: str, +) -> None: + if account_id == session.account.id or session.durable_session_id is None or session.durable_owner_epoch is None: + return + try: + rebound = await service._durable_bridge.rebind_session_account( + session_id=session.durable_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + account_id=account_id, + clear_continuity=True, + ) + except Exception as exc: + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses session account could not be persisted; retry the request.", + ), + ) from exc + if not rebound: + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses session ownership changed during account recovery; retry the request.", + ), + ) + + async def _release_http_bridge_unanchored_handoffs_for_request( service: _HTTPBridgeServiceProtocol, *, @@ -2005,6 +2264,11 @@ def _http_bridge_owner_lookup_unavailable_error_envelope() -> OpenAIErrorEnvelop ) +def _mark_http_bridge_reader_handoff_reconnect_failed(session: Any, old_reader: Any) -> None: + if old_reader is not None: + session.closed = True + + def _http_bridge_previous_response_owner_unavailable_error() -> ProxyResponseError: return ProxyResponseError( 502, @@ -2243,6 +2507,10 @@ def _log_http_bridge_event( cache_key_family: str | None = None, model_class: str | None = None, owner_check_applied: bool | None = None, + error_message: str | None = None, + upstream_close_code: int | None = None, + response_events_seen: int | None = None, + transport_classification: str | None = None, ) -> None: level = logging.INFO if event in { @@ -2251,22 +2519,25 @@ def _log_http_bridge_event( "send_failure", "retry_fresh_upstream", "retry_precreated", + "retry_precreated_clean_close", "reconnect", "terminal_error", "capacity_exhausted_active_sessions", "owner_mismatch", "owner_forward_fail", + "missing_response_created_timeout", "prompt_cache_locality_miss", "reallocation_orphan", "context_overflow_rollover", - "missing_response_created_timeout", + "reader_failure", }: level = logging.WARNING logger.log( level, "http_bridge_event event=%s bridge_kind=%s bridge_key=%s account_id=%s" " model=%s pending=%s detail=%s cache_key_family=%s model_class=%s" - " key_strength=%s owner_check_applied=%s", + " key_strength=%s owner_check_applied=%s error_message=%s upstream_close_code=%s" + " response_events_seen=%s transport_classification=%s", event, key.affinity_kind, _hash_identifier(key.affinity_key), @@ -2278,6 +2549,10 @@ def _log_http_bridge_event( model_class, _http_bridge_key_strength(key), owner_check_applied, + error_message, + upstream_close_code, + response_events_seen, + transport_classification, ) @@ -2329,6 +2604,7 @@ def _wrapper(*args: Any, **kwargs: Any) -> Any: "_record_bridge_drain_recovery_allowed", "_is_missing_durable_bridge_table_error", "_http_bridge_durable_lease_ttl_seconds", + "_persist_http_bridge_replacement_account", "_forwarded_http_bridge_session_key", "_http_bridge_requires_cluster_registration", "_effective_http_bridge_idle_ttl_seconds", diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index d8dfeb0c51..17e1877620 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -69,8 +69,6 @@ _HTTP_BRIDGE_BACKGROUND_CLOSE_TIMEOUT_SECONDS, _HTTP_BRIDGE_INFLIGHT_STARTED_AT_ATTR, _active_http_bridge_instance_ring, - _await_task_deferring_cancellation, - _close_http_bridge_session_bounded, _durable_bridge_lookup_active_owner, _durable_bridge_lookup_allows_local_reuse, _forwarded_http_bridge_session_key, @@ -82,10 +80,12 @@ _http_bridge_can_single_instance_prompt_cache_takeover_without_anchor, _http_bridge_compatible, _http_bridge_continuity_lost_error_envelope, + _http_bridge_durable_release_allowed, _http_bridge_endpoint_matches_current_instance, _http_bridge_eviction_priority, _http_bridge_has_durable_recovery_anchor, _http_bridge_incompatible_model_fork_key, + _http_bridge_inflight_creation_count, _http_bridge_key_strength, _http_bridge_locally_owned_fork_key, _http_bridge_models_compatible, @@ -109,8 +109,9 @@ _http_bridge_turn_state_alias_key, _log_http_bridge_event, _log_http_bridge_startup_wait_timeout, + _mark_http_bridge_reader_handoff_reconnect_failed, + _persist_http_bridge_replacement_account, _preferred_http_bridge_reconnect_turn_state, - _raise_http_bridge_incompatible_admission_handoff, _record_bridge_drain_recovery_allowed, _record_bridge_first_turn_timeout, _refresh_reused_http_bridge_session_with_handoff, @@ -143,26 +144,13 @@ from app.modules.proxy._service.http_bridge.session_registry import _HTTPBridgeSessionRegistryMixin from app.modules.proxy._service.http_bridge.streaming import _HTTPBridgeStreamingMixin from app.modules.proxy._service.http_bridge.upstream_events import _HTTPBridgeUpstreamEventsMixin -from app.modules.proxy._service.observability import ( - _hash_identifier as _hash_identifier, -) -from app.modules.proxy._service.observability import ( - _hash_identifier_or_none as _hash_identifier_or_none, -) -from app.modules.proxy._service.observability import ( - _interesting_header_keys as _interesting_header_keys, -) -from app.modules.proxy._service.observability import ( - _tools_hash as _tools_hash, -) -from app.modules.proxy._service.observability import ( - _truncate_identifier as _truncate_identifier, -) +from app.modules.proxy._service.observability import _hash_identifier from app.modules.proxy._service.support import ( _ACCOUNT_MODEL_UNSUPPORTED_ERROR_CODE, _HARD_HTTP_BRIDGE_AFFINITY_KINDS, # noqa: F401 _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 _clear_websocket_precreated_replay_fallback, + _complete_http_bridge_handoff, _copy_websocket_route_metadata_to_session, _HTTPBridgeOwnerForward, _HTTPBridgeSession, @@ -237,14 +225,6 @@ class _HTTPBridgeMixin( _HTTPBridgeUpstreamEventsMixin, _HTTPBridgeServiceProtocol, ): - async def _close_http_bridge_session_bounded( - self, - session: "_HTTPBridgeSession", - *, - reason: str, - ) -> None: - await _close_http_bridge_session_bounded(self, session, reason=reason) - def _schedule_http_bridge_session_closes( self, sessions: list["_HTTPBridgeSession"], @@ -273,6 +253,7 @@ async def _drain_http_bridge_background_cleanup_tasks(self, *, reason: str) -> N and ( task.get_name().startswith("proxy-http_bridge_session_close-") or task.get_name().startswith("http-bridge-close-") + or task.get_name().startswith("cancelled-task-cleanup-") ) ] if not tasks: @@ -302,6 +283,8 @@ async def _fail_http_bridge_inflight_session_creation( current_future = self._http_bridge_inflight_sessions.get(key) if current_future is not inflight_future: return False + if getattr(inflight_future, "_http_bridge_handoff", False): + return False self._http_bridge_inflight_sessions.pop(key, None) if inflight_future.done(): return True @@ -325,6 +308,8 @@ async def _evict_http_bridge_inflight_waiter( break if stale_key is None: return None + if getattr(inflight_future, "_http_bridge_handoff", False): + return None self._http_bridge_inflight_sessions.pop(stale_key, None) if not inflight_future.done(): inflight_future.set_exception(exc) @@ -362,7 +347,6 @@ async def _get_or_create_http_bridge_session( session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, ) -> "_HTTPBridgeSession": ... - @overload async def _get_or_create_http_bridge_session( self, @@ -394,7 +378,6 @@ async def _get_or_create_http_bridge_session( session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": ... - async def _get_or_create_http_bridge_session( self, key: "_HTTPBridgeSessionKey", @@ -605,7 +588,7 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: previous_session = self._http_bridge_sessions.get(previous_key) if ( previous_session is not None - and not previous_session.closed + and (not previous_session.closed or previous_session.handoff_in_progress) and _http_bridge_session_account_active(previous_session) and _http_bridge_compatible(previous_session, request_model, request_service_tier, True) and _http_bridge_session_matches_preferred_account( @@ -649,7 +632,10 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: bind_account_neutral_recovery_owner(previous_session) key = recovery_fork_key continue - elif not _http_bridge_alias_target_is_stale(previous_session): + elif previous_session is not None and ( + not _http_bridge_alias_target_is_stale(previous_session) + and not previous_session.handoff_in_progress + ): raise ProxyResponseError(502, _http_bridge_continuity_lost_error_envelope()) elif previous_key is not None: self._http_bridge_previous_response_index.pop(previous_alias_key, None) @@ -690,7 +676,8 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: fork_key = _http_bridge_parallel_fork_key( key=key, session=existing, - inflight_creation=key in self._http_bridge_inflight_sessions, + inflight_creation=key in self._http_bridge_inflight_sessions + and not bool(existing and existing.handoff_in_progress), incoming_turn_state=incoming_turn_state, previous_response_id=previous_response_id, request_model=request_model, @@ -710,7 +697,19 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: ) continue if retained_handoff and not reusable: - _raise_http_bridge_incompatible_admission_handoff() + existing, force_durable_takeover = self._recover_http_bridge_incompatible_admission_handoff( + key, + existing, + force_durable_takeover, + original_request_unanchored, + request_model, + api_key, + incoming_turn_state, + previous_response_id, + preferred_account_id, + require_preferred_account, + request_service_tier, + ) if reusable: assert existing is not None current_instance = settings.http_responses_session_bridge_instance_id @@ -1093,7 +1092,7 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: bridge_soft_local_rebind_total.inc() if bridge_local_rebind_total is not None: bridge_local_rebind_total.labels(reason="prompt_cache_locality_miss").inc() - if existing is not None: + if existing is not None and not existing.handoff_in_progress: old_account_id = existing.account.id _log_http_bridge_event( "discard_stale", @@ -1250,7 +1249,7 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: ) elif inflight_future is None: while ( - len(self._http_bridge_sessions) + len(self._http_bridge_inflight_sessions) >= max_sessions + len(self._http_bridge_sessions) + _http_bridge_inflight_creation_count(self) >= max_sessions and self._http_bridge_sessions ): evictable_sessions: list[tuple[_HTTPBridgeSessionKey, _HTTPBridgeSession]] = [] @@ -1287,9 +1286,13 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: detached = self._detach_http_bridge_session_locked(lru_key, expected_session=lru_session) if detached is not None: sessions_to_close_before_create.append(detached) - if len(self._http_bridge_sessions) + len(self._http_bridge_inflight_sessions) >= max_sessions: - if self._http_bridge_inflight_sessions: - capacity_wait_future = next(iter(self._http_bridge_inflight_sessions.values())) + if len(self._http_bridge_sessions) + _http_bridge_inflight_creation_count(self) >= max_sessions: + if _http_bridge_inflight_creation_count(self): + capacity_wait_future = next( + future + for future in self._http_bridge_inflight_sessions.values() + if not getattr(future, "_http_bridge_handoff", False) + ) else: _log_http_bridge_event( "capacity_exhausted_active_sessions", @@ -1297,7 +1300,7 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: account_id=None, model=request_model, pending_count=( - len(self._http_bridge_sessions) + len(self._http_bridge_inflight_sessions) + len(self._http_bridge_sessions) + _http_bridge_inflight_creation_count(self) ), cache_key_family=key.affinity_kind, model_class=_extract_model_class(request_model) if request_model else None, @@ -1364,7 +1367,6 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: except Exception: pass continue - if inflight_future is not None and not owns_creation: wait_timeout_seconds = _proxy_admission_wait_timeout_seconds(settings) try: @@ -1454,7 +1456,6 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: if detached is not None and not retiring_with_visible_requests: self._schedule_http_bridge_session_closes([detached], reason="registry_detach") continue - created_session: _HTTPBridgeSession | None = None session_registered = False try: @@ -1606,11 +1607,11 @@ def _prune_http_bridge_sessions_locked(self) -> list["_HTTPBridgeSession"]: for key, session in self._http_bridge_sessions.items(): if _http_bridge_session_has_admission_waiter(session): continue + if session.handoff_in_progress: + continue if session.closed: stale_keys.append(key) continue - # The request owns this idle session until submit makes activity visible; - # pruning it during admission or durable refresh would invalidate the handoff. if getattr(session, "unanchored_reservation_id", None) is not None: continue pending_count = self._http_bridge_pending_count_nowait(session, context="idle_prune") @@ -1641,6 +1642,7 @@ async def _close_http_bridge_session( session: "_HTTPBridgeSession", *, turn_state_lock_held: bool = False, + release_durable_session: bool = True, ) -> None: session.closed = True if turn_state_lock_held: @@ -1656,7 +1658,7 @@ async def _close_http_bridge_session( logger.warning("Failed to release HTTP bridge account lease during close", exc_info=True) finally: session.account_lease = None - if session.durable_session_id is not None and session.durable_owner_epoch is not None: + if release_durable_session and _http_bridge_durable_release_allowed(self, session): try: await self._durable_bridge.release_live_session( session_id=session.durable_session_id, @@ -1671,7 +1673,11 @@ async def _close_http_bridge_session( if upstream_reader is asyncio.current_task(): session.upstream_reader = None else: - await _await_cancelled_task(upstream_reader, label="http bridge upstream reader") + await _await_cancelled_task( + upstream_reader, + label="http bridge upstream reader", + cleanup_tasks=self._background_cleanup_tasks, + ) if session.upstream_reader is upstream_reader: session.upstream_reader = None try: @@ -1996,23 +2002,36 @@ async def _reconnect_http_bridge_session( owner_rebind_affinity: _AffinityPolicy | None = None, selection_affinity: _AffinityPolicy | None = None, ) -> None: - # A replacement reader can start before its caller resends the request. - # Clear the prior attempt first so an expired send timestamp cannot - # retire the fresh socket before the next real send re-arms it. request_state.response_create_sent_at = None account_neutral_recovery = is_http_bridge_account_neutral_replay( kind=session.key.affinity_kind, key=session.key.affinity_key, ) require_same_account = require_same_account or account_neutral_recovery - old_account_id = session.account.id old_upstream = session.upstream old_reader = session.upstream_reader if restart_reader else None + session.handoff_in_progress = True + inflight_sessions = self._http_bridge_inflight_sessions + handoff_future = inflight_sessions.get(session.key) or asyncio.get_running_loop().create_future() + inflight_sessions.setdefault(session.key, handoff_future) + setattr(handoff_future, "_http_bridge_handoff", True) + session.handoff_future = handoff_future + session.closed = True if old_reader is not None: if old_reader is not asyncio.current_task(): - cancelled = await _await_cancelled_task(old_reader, label="http bridge upstream reader") + try: + cancelled = await _await_cancelled_task( + old_reader, + label="http bridge upstream reader", + cleanup_tasks=self._background_cleanup_tasks, + ) + except BaseException: + session.closed = True + _complete_http_bridge_handoff(session, self._http_bridge_inflight_sessions) + raise if not cancelled: session.closed = True + _complete_http_bridge_handoff(session, self._http_bridge_inflight_sessions) raise ProxyResponseError( 502, openai_error( @@ -2024,30 +2043,34 @@ async def _reconnect_http_bridge_session( request_state, _http_bridge_request_budget_seconds(_service_get_settings()), ) - settings = await _service_get_settings_cache().get() - session.api_key = request_state.api_key - forced_refresh_account_id = request_state.force_refresh_account_id - excluded_account_ids: set[str] = set(request_state.excluded_account_ids) - requested_preferred_account_id = ( - request_state.preferred_account_id if require_preferred_account or account_neutral_recovery else None - ) - required_preferred_account_id = resolve_required_account_id( - ("requested reconnect owner", requested_preferred_account_id), - ( - "account-neutral recovery", - session.account.id if account_neutral_recovery else None, - ), - ) - close_skips_account = session.last_upstream_close_code in _UPSTREAM_CLOSE_CODES_SKIP_SAME_ACCOUNT_RETRY - hard_close_account_bound = session.key.strength == "hard" and (close_skips_account or require_same_account) - skip_same_account = ( - session.key.strength != "hard" and close_skips_account and required_preferred_account_id is None - ) - if required_preferred_account_id is not None and required_preferred_account_id in excluded_account_ids: - raise _http_bridge_previous_response_owner_unavailable_error() - _require_http_bridge_bound_account_not_excluded( - hard_close_account_bound, session.account.id, excluded_account_ids - ) + try: + settings = await _service_get_settings_cache().get() + session.api_key = request_state.api_key + forced_refresh_account_id = request_state.force_refresh_account_id + excluded_account_ids: set[str] = set(request_state.excluded_account_ids) + requested_preferred_account_id = ( + request_state.preferred_account_id if require_preferred_account or account_neutral_recovery else None + ) + required_preferred_account_id = resolve_required_account_id( + ("requested reconnect owner", requested_preferred_account_id), + ("account-neutral recovery", session.account.id if account_neutral_recovery else None), + ) + close_skips_account = session.last_upstream_close_code in _UPSTREAM_CLOSE_CODES_SKIP_SAME_ACCOUNT_RETRY + hard_close_account_bound = session.key.strength == "hard" and (close_skips_account or require_same_account) + skip_same_account = ( + session.key.strength != "hard" and close_skips_account and required_preferred_account_id is None + ) + if required_preferred_account_id is not None and required_preferred_account_id in excluded_account_ids: + session.closed = True + _complete_http_bridge_handoff(session, self._http_bridge_inflight_sessions) + raise _http_bridge_previous_response_owner_unavailable_error() + _require_http_bridge_bound_account_not_excluded( + hard_close_account_bound, session.account.id, excluded_account_ids + ) + except BaseException: + session.closed = True + _complete_http_bridge_handoff(session, self._http_bridge_inflight_sessions) + raise if skip_same_account: excluded_account_ids.add(session.account.id) retry_same_account_once = not skip_same_account and session.account.id not in excluded_account_ids @@ -2085,49 +2108,89 @@ async def release_selected_account_lease() -> None: async with session.pending_lock: if session.account_lease is not None and lease.lease_id == session.account_lease.lease_id: session.account_lease = None - await self._load_balancer.release_account_lease(lease) + try: + await self._load_balancer.release_account_lease(lease) + except BaseException: + complete_failed_handoff() + raise async def abandon_selected_account_retry(selected_account: Any) -> None: nonlocal preferred_candidate_id if hard_close_account_bound or selected_account_model_replacement: await release_selected_account_lease() + complete_failed_handoff() raise excluded_account_ids.add(selected_account.id) preferred_candidate_id = None await release_selected_account_lease() + async def open_replacement_upstream(selected_account: Any, selected_headers: dict[str, str]) -> Any: + try: + return await self._open_upstream_websocket_with_budget( + selected_account, + selected_headers, + timeout_seconds=_remaining_budget_seconds(deadline), + request_state=request_state, + ) + except Exception: + session.closed = True + raise + + def complete_failed_handoff() -> None: + session.closed = True + _mark_http_bridge_reader_handoff_reconnect_failed(session, old_reader) + _complete_http_bridge_handoff(session, self._http_bridge_inflight_sessions) + + def require_bound_account() -> None: + try: + _require_http_bridge_bound_account_not_excluded( + hard_close_account_bound, session.account.id, excluded_account_ids + ) + except BaseException: + complete_failed_handoff() + raise + while True: reuse_current_account_lease = preferred_candidate_id == session.account.id and bool(session.account_lease) - selection = await self._select_account_with_budget_for_stream( - deadline, - request_id=request_state.request_log_id or request_state.request_id, - kind="http_bridge", - request_stage="reattach", - api_key=session.api_key, - affinity_policy=selection_affinity or session.affinity, - prefer_earlier_reset_accounts=settings.prefer_earlier_reset_accounts, - prefer_earlier_reset_window=_prefer_earlier_reset_window(settings), - routing_strategy=_routing_strategy(settings), - model=session.request_model, - service_tier=session.request_service_tier, - exclude_account_ids=excluded_account_ids, - preferred_account_id=preferred_candidate_id, - preferred_account_is_continuity_owner=account_neutral_recovery, - require_security_work_authorized=require_security_work_authorized, - lease_kind=None if reuse_current_account_lease else "stream", - estimated_lease_tokens=_estimated_lease_tokens_from_request_usage_budget( - request_state.request_usage_budget - ), - fallback_on_preferred_account_unavailable=( - not reuse_current_account_lease - and not hard_close_account_bound - and required_preferred_account_id is None - ), - ) + try: + selection = await self._select_account_with_budget_for_stream( + deadline, + request_id=request_state.request_log_id or request_state.request_id, + kind="http_bridge", + request_stage="reattach", + api_key=session.api_key, + affinity_policy=selection_affinity or session.affinity, + prefer_earlier_reset_accounts=settings.prefer_earlier_reset_accounts, + prefer_earlier_reset_window=_prefer_earlier_reset_window(settings), + routing_strategy=_routing_strategy(settings), + model=session.request_model, + service_tier=session.request_service_tier, + exclude_account_ids=excluded_account_ids, + preferred_account_id=preferred_candidate_id, + preferred_account_is_continuity_owner=account_neutral_recovery, + require_security_work_authorized=require_security_work_authorized, + lease_kind=None if reuse_current_account_lease else "stream", + estimated_lease_tokens=_estimated_lease_tokens_from_request_usage_budget( + request_state.request_usage_budget + ), + fallback_on_preferred_account_unavailable=( + not reuse_current_account_lease + and not hard_close_account_bound + and required_preferred_account_id is None + ), + ) + except BaseException: + complete_failed_handoff() + raise account = selection.account if account is None: - await release_selected_account_lease() + try: + await release_selected_account_lease() + except BaseException: + complete_failed_handoff() + raise if account_neutral_recovery and selection.error_code == CONTINUITY_OWNER_UNAVAILABLE: + complete_failed_handoff() raise _http_bridge_previous_response_owner_unavailable_error() if ( reuse_current_account_lease @@ -2137,23 +2200,27 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: ): preferred_candidate_id = None continue - if await _sleep_for_account_selection_recovery( - selection, - request_id=request_state.request_log_id or request_state.request_id, - kind="http_bridge", - request_stage="reattach", - model=session.request_model, - max_sleep_seconds=_remaining_budget_seconds(deadline), - request_state=request_state, - ): + try: + should_retry_selection = await _sleep_for_account_selection_recovery( + selection, + request_id=request_state.request_log_id or request_state.request_id, + kind="http_bridge", + request_stage="reattach", + model=session.request_model, + max_sleep_seconds=_remaining_budget_seconds(deadline), + request_state=request_state, + ) + except BaseException: + complete_failed_handoff() + raise + if should_retry_selection: excluded_account_ids.update(request_state.excluded_account_ids) if required_preferred_account_id in excluded_account_ids: + complete_failed_handoff() raise _http_bridge_previous_response_owner_unavailable_error() if skip_same_account: excluded_account_ids.add(session.account.id) - _require_http_bridge_bound_account_not_excluded( - hard_close_account_bound, session.account.id, excluded_account_ids - ) + require_bound_account() retry_same_account_once = not skip_same_account and session.account.id not in excluded_account_ids if skip_same_account: preferred_candidate_id = None @@ -2172,6 +2239,7 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: continue record_selected_account_takeover(None) status_code = 429 if _is_local_account_cap_code(selection.error_code) else 503 + complete_failed_handoff() raise ProxyResponseError( status_code, openai_error( @@ -2182,8 +2250,10 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: ) if required_preferred_account_id is not None and account.id != required_preferred_account_id: if selection.lease is not None: - await self._load_balancer.release_account_lease(selection.lease) + selected_account_lease = selection.lease + await release_selected_account_lease() record_selected_account_takeover(account.id, required_preferred_account_id) + complete_failed_handoff() raise _http_bridge_previous_response_owner_unavailable_error() selected_account_lease = ( session.account_lease @@ -2214,18 +2284,14 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: session.headers, None if owner_rebind_affinity is not None else _preferred_http_bridge_reconnect_turn_state(session), ) - upstream = await self._open_upstream_websocket_with_budget( - account, - connect_headers, - timeout_seconds=_remaining_budget_seconds(deadline), - request_state=request_state, - ) + upstream = await open_replacement_upstream(account, connect_headers) _copy_websocket_route_metadata_to_session(session, request_state) record_selected_account_takeover(account.id) break except ProxyResponseError as exc: if exc.status_code != 401 or _remaining_budget_seconds(deadline) <= 0: await release_selected_account_lease() + complete_failed_handoff() raise try: account = await self._ensure_fresh_with_budget( @@ -2241,18 +2307,14 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: else _preferred_http_bridge_reconnect_turn_state(session) ), ) - upstream = await self._open_upstream_websocket_with_budget( - account, - connect_headers, - timeout_seconds=_remaining_budget_seconds(deadline), - request_state=request_state, - ) + upstream = await open_replacement_upstream(account, connect_headers) _copy_websocket_route_metadata_to_session(session, request_state) record_selected_account_takeover(account.id) break except ProxyResponseError as retry_exc: if retry_exc.status_code != 401: await release_selected_account_lease() + complete_failed_handoff() raise await self._handle_proxy_error(account, retry_exc) await abandon_selected_account_retry(account) @@ -2273,6 +2335,7 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: await abandon_selected_account_retry(account) continue await release_selected_account_lease() + complete_failed_handoff() raise except (aiohttp.ClientError, asyncio.TimeoutError): if selected_is_preferred and _remaining_budget_seconds(deadline) > 0: @@ -2283,34 +2346,72 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: await abandon_selected_account_retry(account) continue await release_selected_account_lease() + complete_failed_handoff() raise - if owner_rebind_affinity is not None: - await self._claim_http_bridge_replacement_before_swap( - session, - account_id=account.id, - upstream=upstream, - release_selected_account_lease=release_selected_account_lease, - owner_rebind_affinity=owner_rebind_affinity, - ) - await self._unregister_http_bridge_turn_states(session) - await self._unregister_http_bridge_previous_response_ids(session) - session.last_completed_response_id = None - session.last_completed_input_count = 0 - session.last_completed_input_prefix_fingerprint = None - session.last_pending_tool_calls.clear() - session.affinity = selection_affinity or session.affinity - session.codex_session = False - session.upstream_turn_state = None - session.downstream_turn_state = None - session.headers = { - key: value for key, value in session.headers.items() if key.lower() != "x-codex-turn-state" - } + except asyncio.CancelledError: + session.closed = True + await release_selected_account_lease() + complete_failed_handoff() + raise + except BaseException: + await release_selected_account_lease() + complete_failed_handoff() + raise + + async def abort_selected_handoff() -> None: + session.closed = True + try: + await asyncio.shield(upstream.close()) + except BaseException: + logger.debug("Failed to close HTTP bridge replacement websocket", exc_info=True) + selected_lease = selected_account_lease + old_lease = session.account_lease + try: + await asyncio.shield(release_selected_account_lease()) + except BaseException: + logger.debug("Failed to release HTTP bridge replacement lease", exc_info=True) + if old_lease is not None and old_lease is not selected_lease: + session.account_lease = None + try: + await asyncio.shield(self._load_balancer.release_account_lease(old_lease)) + except BaseException: + logger.debug("Failed to release HTTP bridge old account lease", exc_info=True) + complete_failed_handoff() + try: - await old_upstream.close() - except Exception: - logger.debug("Failed to close HTTP bridge upstream websocket before reconnect", exc_info=True) - async with session.pending_lock: - replaced_account_lease = session.account_lease + if owner_rebind_affinity is not None: + await self._claim_http_bridge_replacement_before_swap( + session, + account_id=account.id, + upstream=upstream, + release_selected_account_lease=release_selected_account_lease, + owner_rebind_affinity=owner_rebind_affinity, + ) + if owner_rebind_affinity is not None or account.id != session.account.id: + await self._unregister_http_bridge_turn_states(session) + await self._unregister_http_bridge_previous_response_ids(session) + session.last_completed_response_id = None + session.last_completed_input_count = 0 + session.last_completed_input_prefix_fingerprint = None + session.last_pending_tool_calls.clear() + session.affinity = selection_affinity or session.affinity + session.codex_session = False + session.upstream_turn_state = None + session.downstream_turn_state = None + session.headers = { + key: value for key, value in session.headers.items() if key.lower() != "x-codex-turn-state" + } + await _persist_http_bridge_replacement_account(self, session, account.id) + try: + await old_upstream.close() + except Exception: + logger.debug("Failed to close HTTP bridge upstream websocket before reconnect", exc_info=True) + session.closed = True + if selected_account_lease is not session.account_lease: + old_lease = session.account_lease + if old_lease is not None: + await self._load_balancer.release_account_lease(old_lease) + session.account_lease = None session.account_lease = selected_account_lease session.account, session.headers, session.upstream = account, connect_headers, upstream session.catalog_omission_quota_admission = selection.catalog_omission_quota_admission @@ -2318,13 +2419,10 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: session.closed = False session.last_upstream_close_code = None session.upstream_turn_state = _upstream_turn_state_from_socket(upstream) or session.upstream_turn_state - if replaced_account_lease is not None and ( - selected_account_lease is None or selected_account_lease.lease_id != replaced_account_lease.lease_id - ): - release_task = asyncio.create_task(self._load_balancer.release_account_lease(replaced_account_lease)) - _, cancellation = await _await_task_deferring_cancellation(release_task) - if cancellation is not None: - raise cancellation + _complete_http_bridge_handoff(session, self._http_bridge_inflight_sessions) + except BaseException: + await abort_selected_handoff() + raise if restart_reader: session.upstream_reader = asyncio.create_task(self._relay_http_bridge_upstream_messages(session)) _log_http_bridge_event( @@ -2332,11 +2430,7 @@ async def abandon_selected_account_retry(selected_account: Any) -> None: session.key, account_id=account.id, model=session.request_model, - detail=( - f"request_stage=reattach, previous_account={old_account_id}, " - f"preferred_account_id={old_account_id}, selected_account_id={account.id}, " - f"durable_session_id={session.durable_session_id}" - ), + detail=f"selected_account_id={account.id}, durable_session_id={session.durable_session_id}", cache_key_family=session.key.affinity_kind, model_class=_extract_model_class(session.request_model) if session.request_model else None, ) diff --git a/app/modules/proxy/_service/http_bridge/request_submit.py b/app/modules/proxy/_service/http_bridge/request_submit.py index 0319e918ea..8e4f13f608 100644 --- a/app/modules/proxy/_service/http_bridge/request_submit.py +++ b/app/modules/proxy/_service/http_bridge/request_submit.py @@ -3,6 +3,8 @@ import asyncio import json import logging +import math +import random from collections import deque from dataclasses import replace from typing import Any, Literal, Mapping, cast @@ -70,6 +72,7 @@ from app.modules.proxy._service.http_bridge.helpers import ( _await_task_deferring_cancellation, _build_http_bridge_prewarm_text, + _http_bridge_durable_lease_ttl_seconds, _http_bridge_key_strength, _http_bridge_precreated_retry_failure_error, _http_bridge_prewarm_enabled, @@ -125,6 +128,7 @@ _truncate_identifier as _truncate_identifier, ) from app.modules.proxy._service.support import ( + _ACCOUNT_MODEL_UNSUPPORTED_ERROR_CODE, _HARD_HTTP_BRIDGE_AFFINITY_KINDS, # noqa: F401 _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 _clear_websocket_request_error_overrides, @@ -172,12 +176,14 @@ _AffinityPolicy, _extract_model_class, _owner_lookup_session_id_from_headers, + _sticky_key_from_turn_state_header, ) from app.modules.proxy.api_key_usage import estimate_api_key_request_usage from app.modules.proxy.continuity import is_http_bridge_account_neutral_replay from app.modules.proxy.durable_bridge_repository import ( DurableBridgeAliasRegistration, DurableBridgeAliasRegistrationReceipt, + durable_bridge_hash, ) from app.modules.proxy.helpers import ( _normalize_error_code, @@ -189,6 +195,10 @@ ) logger = logging.getLogger("app.modules.proxy.service") + +_HTTP_BRIDGE_CLEAN_CLOSE_RETRY_MAX_COUNT = 1 +_HTTP_BRIDGE_CLEAN_CLOSE_RETRY_JITTER_MAX_SECONDS = 2.0 + _REQUEST_TRANSPORT_HTTP = "http" _WEBSOCKET_AUTH_INVALIDATED_FAILURE_CODE = "account_auth_invalidated" _NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE = "no_security_work_authorized_accounts" @@ -273,6 +283,30 @@ def _request_kind_from_headers(headers: Mapping[str, str] | None) -> str: class _HTTPBridgeRequestSubmitMixin: + @staticmethod + def _http_bridge_clean_close_retry_max_count() -> int: + configured = _HTTP_BRIDGE_CLEAN_CLOSE_RETRY_MAX_COUNT + # Keep this recovery bounded even if an unsafe higher value is supplied. + return max(0, min(1, configured)) + + @staticmethod + def _http_bridge_clean_close_retry_jitter_seconds() -> float: + settings = _service_get_settings() + maximum = max( + 0.0, + min( + 30.0, + float( + getattr( + settings, + "http_responses_session_bridge_clean_close_retry_jitter_max_seconds", + _HTTP_BRIDGE_CLEAN_CLOSE_RETRY_JITTER_MAX_SECONDS, + ) + ), + ), + ) + return random.uniform(0.0, maximum) if maximum > 0 else 0.0 + def _prepare_http_bridge_request( self: Any, payload: ResponsesRequest, @@ -383,6 +417,10 @@ def _prepare_response_bridge_request_state( request_usage_budget=estimate_api_key_request_usage(payload), previous_response_id=payload.previous_response_id, session_id=_normalize_session_id(session_id), + hard_continuity_anchor=( + payload.previous_response_id is not None + or _sticky_key_from_turn_state_header(headers or {}) is not None + ), input_item_count=input_item_count, input_full_fingerprint=input_full_fingerprint, request_kind=request_kind, @@ -550,6 +588,135 @@ async def _submit_http_bridge_request_with_handoff( request_scope_id: str, recovery_turn_state: str | None = None, ) -> None: + # Eventless upstream timeouts retire the current socket. A client + # reconnect can otherwise create a fresh socket for the same hard key + # and submit the identical request repeatedly while the retry circuit + # is cooling down. Gate new submissions before any reconnect/send so + # the circuit turns this into a bounded 503 instead of another + # response.create attempt. A proof-gated full resend remains allowed + # because it is the client's own replay-safe request, not an opaque + # continuation replay. + allow_proof_gated_continuity_replay = bool( + request_state.previous_response_id is not None + and request_state.fresh_upstream_request_is_retry_safe + and request_state.fresh_upstream_request_text + and request_state.response_event_count == 0 + and request_state.replay_count == 0 + ) + if not await self._http_bridge_precreated_retry_allowed( + session, + allow_proof_gated_continuity_replay=allow_proof_gated_continuity_replay, + ): + retry_after_seconds = max( + 1, + math.ceil(await self._http_bridge_precreated_retry_cooldown_seconds(session)), + ) + _log_http_bridge_event( + "submit_retry_circuit_suppressed", + session.key, + account_id=session.account.id, + model=session.request_model, + detail="hard_key_cooldown", + cache_key_family=session.key.affinity_kind, + model_class=_extract_model_class(session.request_model) if session.request_model else None, + ) + raise ProxyResponseError( + 503, + openai_error( + "upstream_request_timeout", + "HTTP responses session bridge is cooling down after repeated upstream timeouts; retry shortly.", + ), + retry_after_seconds=retry_after_seconds, + ) + # Persist the recovery checkpoint only after the retry circuit has + # admitted this request. A client reconnect suppressed by the + # cooldown must not create or refresh a journal entry for a request + # that was never dispatched upstream. + if ( + request_state.fresh_upstream_request_is_retry_safe + and request_state.fresh_upstream_request_text + and request_state.replay_count == 0 + and request_state.recovery_attempt_fingerprint is None + and session.durable_session_id is not None + and session.durable_owner_epoch is not None + ): + attempt_fingerprint = durable_bridge_hash(request_state.fresh_upstream_request_text) + try: + attempt = await self._durable_bridge.record_recovery_attempt( + session_id=session.durable_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + request_fingerprint=attempt_fingerprint, + request_id=request_state.request_id, + account_id=session.account.id, + model=request_state.model, + replay_safe=True, + ) + if attempt is None: + # ``None`` is the durable owner fence rejecting this + # worker, not an unavailable journal (which raises and + # is handled below). Never dispatch from a stale owner. + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + _record_continuity_fail_closed( + surface="http_bridge", + reason="recovery_attempt_owner_fence_rejected", + previous_response_id=request_state.previous_response_id, + session_id=request_state.session_id, + upstream_error_code="bridge_continuity_persistence_failed", + ) + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses session ownership changed; retry the request.", + ), + ) + if getattr(attempt.state, "value", attempt.state) != "unknown": + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "The recovery checkpoint was already consumed; retry the request.", + ), + ) + if getattr(attempt, "request_id", request_state.request_id) != request_state.request_id: + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "Another recovery request is already in flight; retry the request.", + ), + ) + request_state.recovery_attempt_fingerprint = attempt_fingerprint + request_state.recovery_attempt_session_id = session.durable_session_id + request_state.recovery_attempt_owner_epoch = session.durable_owner_epoch + except ProxyResponseError: + raise + except Exception as exc: + # The journal is an additional recovery fence. Without a + # durable UNKNOWN row, an ambiguous send cannot be claimed by + # another owner, so fail closed instead of dispatching an + # unjournaled recovery-safe request. + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + _record_continuity_fail_closed( + surface="http_bridge", + reason="recovery_attempt_persistence_failed", + previous_response_id=request_state.previous_response_id, + session_id=request_state.session_id, + upstream_error_code="bridge_continuity_persistence_failed", + ) + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "Recovered response continuity could not be persisted; retry the request.", + ), + ) from exc text_data = self._http_bridge_text_with_account_installation_id(session, request_state, text_data) if request_state.response_id is not None or request_state.response_event_count > 0: _log_http_bridge_event( @@ -728,6 +895,7 @@ async def _submit_http_bridge_request_with_handoff( ) _copy_websocket_route_metadata_from_session(request_state, session) request_state.bridge_queue_wait_started_at = _service_time().monotonic() + request_state.response_create_gate_wait_started_at = _service_time().monotonic() # Bridge ownership is established before this late admission. A # cap race stays a bounded error/wait on that owner; it must not # publish a replacement bridge as a spillover side effect. @@ -868,6 +1036,107 @@ async def _submit_http_bridge_request_with_handoff( "Recovered response continuity could not be persisted; retry the request.", ), ) + if request_state.recovery_attempt_fingerprint is not None: + try: + owner_lookup = await self._durable_bridge.renew_live_session( + session_id=session.durable_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + lease_ttl_seconds=_http_bridge_durable_lease_ttl_seconds(), + latest_turn_state=session.downstream_turn_state, + latest_response_id=None, + ) + except Exception as exc: + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses session ownership could not be renewed; retry the request.", + ), + ) from exc + if ( + owner_lookup is None + or owner_lookup.owner_instance_id + != _service_get_settings().http_responses_session_bridge_instance_id + or owner_lookup.owner_epoch != session.durable_owner_epoch + ): + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses session ownership changed before dispatch; retry the request.", + ), + ) + # The journal entry is created before queue admission so + # concurrent requests cannot both enter the gate without + # a recovery generation. Revalidate it after the gate and + # lifecycle locks, immediately before dispatch: a waiter + # may have observed UNKNOWN while the first request + # settled the row REPLAYED. + if ( + request_state.recovery_attempt_fingerprint is not None + and not request_state.recovery_attempt_claimed + ): + if ( + request_state.recovery_attempt_session_id != session.durable_session_id + or request_state.recovery_attempt_owner_epoch != session.durable_owner_epoch + ): + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses session ownership changed before dispatch; retry the request.", + ), + ) + try: + dispatch_attempt = await self._durable_bridge.record_recovery_attempt( + session_id=session.durable_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + request_fingerprint=request_state.recovery_attempt_fingerprint, + request_id=request_state.request_id, + account_id=session.account.id, + model=request_state.model, + replay_safe=True, + ) + except Exception as exc: + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "Recovered response continuity could not be revalidated; retry the request.", + ), + ) from exc + if ( + dispatch_attempt is None + or getattr(dispatch_attempt.state, "value", dispatch_attempt.state) != "unknown" + or getattr(dispatch_attempt, "request_id", request_state.request_id) + != request_state.request_id + ): + session.closed = True + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "The recovery checkpoint was consumed before dispatch; retry the request.", + ), + ) async with session.pending_lock: session.pending_requests.append(request_state) session.admission_waiter_count = max(0, session.admission_waiter_count - 1) @@ -877,6 +1146,7 @@ async def _submit_http_bridge_request_with_handoff( try: await _send_http_bridge_request_text_with_archive_id(session, request_state, text_data) except BaseException: + request_state.recovery_attempt_dispatched = True # Publish retirement while lifecycle ownership is still # held; a gate waiter must never reuse an ambiguously sent # response.create socket between unlock and cleanup. @@ -884,6 +1154,7 @@ async def _submit_http_bridge_request_with_handoff( session.upstream_control.reconnect_requested = True session.upstream_control.retire_after_drain = True raise + request_state.recovery_attempt_dispatched = True session.last_used_at = _service_time().monotonic() except asyncio.CancelledError: if recovery_receipt is not None and not upstream_send_started: @@ -1257,6 +1528,10 @@ async def _cleanup_http_bridge_submit_interruption( await self._retire_stale_pending_http_bridge_session( session, detail="last_admission_waiter_cancelled", + response_events_seen=max( + request_state.response_event_count, + int(request_state.response_id is not None or request_state.latency_response_created_ms is not None), + ), ) await self._maybe_release_idle_http_bridge_session_lease(session) @@ -1400,6 +1675,78 @@ async def _detach_http_bridge_request( await self._retire_http_bridge_after_drain_if_ready(session) return True + async def _fail_stale_http_bridge_pending_requests( + self: Any, + session: "_HTTPBridgeSession", + request_states: list[_WebSocketRequestState], + *, + detail: str, + ) -> None: + stale_requests: deque[_WebSocketRequestState] = deque() + response_events_seen = 0 + async with session.pending_lock: + for request_state in request_states: + if request_state not in session.pending_requests: + continue + response_events_seen = max( + response_events_seen, + request_state.response_event_count, + int( + request_state.response_id is not None + or request_state.latency_response_created_ms is not None + or request_state.downstream_visible + ), + ) + session.pending_requests.remove(request_state) + if _http_bridge_request_counts_against_queue(request_state): + session.queued_request_count = max(0, session.queued_request_count - 1) + stale_requests.append(request_state) + if not stale_requests: + return + if response_events_seen == 0: + await self._record_http_bridge_retry_circuit_failure(session, detail=detail) + await self._fail_pending_websocket_requests( + account=session.account, + account_id_value=session.account.id, + pending_requests=stale_requests, + pending_lock=session.pending_lock, + error_code="upstream_request_timeout", + error_message="HTTP bridge response-create gate holder timed out", + api_key=None, + response_create_gate=session.response_create_gate, + penalize_account=False, + ) + + def _classify_http_bridge_stale_gate_holders( + self: Any, + pending_states: list[_WebSocketRequestState], + *, + now: float, + threshold_seconds: float, + session_closed: bool, + ) -> tuple[list[_WebSocketRequestState], bool]: + stale_states = [ + state + for state in pending_states + if not state.draining_until_terminal + and self._http_bridge_pending_state_is_stale( + state, + now=now, + threshold_seconds=threshold_seconds, + session_closed=session_closed, + ) + ] + active_states = [ + # A draining request still owns terminal response continuity. It + # must keep the session alive while stale holders are cleaned up. + state + for state in pending_states + if state not in stale_states + ] + if stale_states and active_states: + return stale_states, False + return [], bool(stale_states) + async def _retire_http_bridge_after_drain_if_ready(self: Any, session: "_HTTPBridgeSession") -> bool: if not (session.upstream_control.reconnect_requested and session.upstream_control.retire_after_drain): return False @@ -1424,7 +1771,14 @@ async def _retire_stale_pending_http_bridge_session( session: "_HTTPBridgeSession", *, detail: str, + retry_circuit_detail: str | None = None, + response_events_seen: int | None = None, ) -> None: + if response_events_seen is None or response_events_seen == 0: + await self._record_http_bridge_retry_circuit_failure( + session, + detail=retry_circuit_detail or detail, + ) session.closed = True async with self._http_bridge_lock: if self._http_bridge_sessions.get(session.key) is session: @@ -1468,16 +1822,14 @@ async def _retry_http_bridge_request_on_fresh_upstream( # upstream already accepted the continuation. Re-sending the same # previous_response_id request can fork continuity with duplicate # child responses, so only reconnect-without-resend is allowed. - # The single exception is proxy-injected anchors on trim-safe - # full-resend payloads: dropping the anchor and replaying the - # original unanchored request is equivalent to the client's own - # retry. Session-level injections do not opt in because their - # payload may depend on the anchor for context preservation. - if ( - not request_state.proxy_injected_previous_response_id - or not request_state.fresh_upstream_request_text - or not request_state.fresh_upstream_request_is_retry_safe - ): + # The single exception is a proof-gated, trim-safe full-resend + # payload: dropping the anchor and replaying the original + # unanchored request is equivalent to the client's own retry. + # The proof is independent of where the anchor came from; a + # client-provided full resend is as safe as a durable injection. + # Session-level follow-ups do not opt in because their context may + # depend on the anchor. + if not request_state.fresh_upstream_request_text or not request_state.fresh_upstream_request_is_retry_safe: return False retry_text_data = request_state.fresh_upstream_request_text using_fresh_replay = True @@ -1530,7 +1882,67 @@ async def _retry_http_bridge_precreated_request( session: "_HTTPBridgeSession", *, request_state: _WebSocketRequestState | None = None, + restart_reader: bool = False, ) -> bool: + clean_close_retry_max_count = self._http_bridge_clean_close_retry_max_count() + account_neutral_recovery = is_http_bridge_account_neutral_replay( + kind=session.key.affinity_kind, + key=session.key.affinity_key, + ) + + def request_is_retryable(request_state: _WebSocketRequestState) -> bool: + if _websocket_request_can_replay_before_visible_output(request_state): + return True + if ( + clean_close_retry_max_count <= 0 + or request_state.replay_count != 1 + or request_state.response_event_count != 0 + or request_state.clean_close_replay_count >= clean_close_retry_max_count + or _classify_upstream_close( + session.last_upstream_close_code, + response_events_seen=request_state.response_event_count, + ) + != "clean" + ): + return False + return _websocket_request_can_replay_before_visible_output( + request_state, + allow_clean_close_retry=True, + ) + + fresh_hard_request_account_switch_candidate = False + proof_gated_continuity_replay_candidate = False + if session.key.strength == "hard": + async with session.pending_lock: + retryable_candidates = [ + request_state + for request_state in session.pending_requests + if not request_state.draining_until_terminal and request_is_retryable(request_state) + ] + if len(retryable_candidates) == 1: + candidate = retryable_candidates[0] + fresh_hard_request_account_switch_candidate = ( + candidate.previous_response_id is None + and not candidate.hard_continuity_anchor + and not candidate.proxy_injected_previous_response_id + and not candidate.file_required_preferred_account + and candidate.response_event_count == 0 + and candidate.replay_count == 0 + ) + proof_gated_continuity_replay_candidate = ( + candidate.previous_response_id is not None + and candidate.fresh_upstream_request_is_retry_safe + and bool(candidate.fresh_upstream_request_text) + and candidate.response_event_count == 0 + and candidate.replay_count == 0 + ) + if not await self._http_bridge_precreated_retry_allowed( + session, + allow_fresh_hard_account_switch=fresh_hard_request_account_switch_candidate, + allow_proof_gated_continuity_replay=proof_gated_continuity_replay_candidate, + ): + return False + account_neutral_recovery = is_http_bridge_account_neutral_replay( kind=session.key.affinity_kind, key=session.key.affinity_key, @@ -1543,41 +1955,62 @@ async def _retry_http_bridge_precreated_request( or any(pending_request is not request_state for pending_request in session.pending_requests) or request_state.draining_until_terminal or not _http_bridge_request_counts_against_queue(request_state) - or not _websocket_request_can_replay_before_visible_output(request_state) + or not request_is_retryable(request_state) ): return False else: retryable_requests = [ request_state for request_state in session.pending_requests - if not request_state.draining_until_terminal - and _websocket_request_can_replay_before_visible_output(request_state) + if not request_state.draining_until_terminal and request_is_retryable(request_state) ] if len(retryable_requests) != 1: return False request_state = retryable_requests[0] + model_fallback_replay = request_state.precreated_replay_reason == _ACCOUNT_MODEL_UNSUPPORTED_ERROR_CODE if request_state.previous_response_id is not None and not ( - request_state.proxy_injected_previous_response_id - and request_state.fresh_upstream_request_is_retry_safe - and request_state.fresh_upstream_request_text + request_state.fresh_upstream_request_is_retry_safe and request_state.fresh_upstream_request_text ): # Once a continuation is pending upstream, reconnecting without # replay cannot complete the current request, while replaying it # is unsafe without upstream idempotency guarantees. Proxy- - # injected retry-safe anchors are equivalent to the client's own - # full resend once the anchor is stripped. + # injected anchors and proof-gated client full resends are + # equivalent to the client's own retry once the anchor is + # stripped. The latter remains pinned to the current owner. return False close_classification = _classify_upstream_close( session.last_upstream_close_code, response_events_seen=request_state.response_event_count, ) - if close_classification == "rejected": - request_state.error_code_override = "upstream_rejected_input" - request_state.error_http_status_override = 502 - request_state.error_message_override = ( - "Upstream rejected the request before response.created " - f"(close_code={session.last_upstream_close_code})" - ) + close_generation = session.last_upstream_close_generation + hard_session_affinity = session.key.strength == "hard" + fresh_hard_request_account_switch_allowed = ( + hard_session_affinity + and request_state.previous_response_id is None + and not request_state.hard_continuity_anchor + and not request_state.proxy_injected_previous_response_id + and not request_state.file_required_preferred_account + ) + clean_close_hard_continuation = ( + close_classification == "clean" + and hard_session_affinity + and request_state.previous_response_id is not None + ) + clean_close_hard_continuity_anchor = ( + close_classification == "clean" and hard_session_affinity and request_state.hard_continuity_anchor + ) + clean_close_retry_for_current_close = ( + close_classification == "clean" + and request_state.clean_close_retry_close_generation != close_generation + and not request_state.clean_close_retry_in_progress + ) + additional_clean_close_retry = ( + clean_close_retry_for_current_close + and request_state.replay_count == 1 + and request_state.response_event_count == 0 + and request_state.clean_close_replay_count < clean_close_retry_max_count + ) + if request_state.replay_count >= 1 and not additional_clean_close_retry: return False if request_state.previous_response_id is not None: require_preferred_reconnect = False @@ -1607,23 +2040,41 @@ async def _retry_http_bridge_precreated_request( if not hard_owner_bound: request_state.excluded_account_ids.add(session.account.id) else: - require_preferred_reconnect = account_neutral_recovery + # Account-scoped uploaded files cannot be replayed on a + # different owner. Keep the preferred account mandatory for + # both silent recovery and clean-close recovery. + require_preferred_reconnect = account_neutral_recovery or request_state.file_required_preferred_account request_text = _prepare_websocket_request_state_for_visible_output_replay(request_state) if request_text is None: return False if account_neutral_recovery: request_state.preferred_account_id = session.account.id - elif not request_state.file_required_preferred_account and not hard_owner_bound: - request_state.preferred_account_id = None - request_state.excluded_account_ids.add(session.account.id) + elif not request_state.file_required_preferred_account: + if hard_owner_bound and not model_fallback_replay and not fresh_hard_request_account_switch_allowed: + request_state.preferred_account_id = session.account.id + else: + request_state.preferred_account_id = None + request_state.excluded_account_ids.add(session.account.id) if session.account.id in request_state.excluded_account_ids: session.upstream_turn_state = None session.downstream_turn_state = None session.headers = { key: value for key, value in session.headers.items() if key.lower() != "x-codex-turn-state" } + if close_classification == "clean": + if not clean_close_retry_for_current_close: + return False + request_state.clean_close_retry_in_progress = True + request_state.clean_close_retry_result = None + request_state.clean_close_retry_close_generation = close_generation + if additional_clean_close_retry: + request_state.clean_close_replay_count += 1 + retry_jitter_seconds = ( + self._http_bridge_clean_close_retry_jitter_seconds() if additional_clean_close_retry else 0.0 + ) + retry_event = "retry_precreated_clean_close" if additional_clean_close_retry else "retry_precreated" _log_http_bridge_event( - "retry_precreated", + retry_event, session.key, account_id=session.account.id, model=session.request_model, @@ -1631,12 +2082,47 @@ async def _retry_http_bridge_precreated_request( cache_key_family=session.key.affinity_kind, model_class=_extract_model_class(session.request_model) if session.request_model else None, ) + reconnect_reader_kwargs = {"restart_reader": True} if restart_reader else {} try: - if hard_owner_bound: + if retry_jitter_seconds > 0: + logger.info( + "HTTP bridge clean-close retry jitter request_id=%s sleep_seconds=%.3f", + request_state.request_id, + retry_jitter_seconds, + ) + await asyncio.sleep(retry_jitter_seconds) + request_deadline = request_state.bridge_request_deadline + if request_deadline is None: + request_deadline = request_state.started_at + _http_bridge_request_budget_seconds( + _service_get_settings() + ) + now_monotonic = _service_time().monotonic() + async with session.pending_lock: + request_still_owned = ( + request_state in session.pending_requests and not request_state.draining_until_terminal + ) + if not request_still_owned or now_monotonic >= request_deadline: + logger.info( + "HTTP bridge clean-close retry abandoned after jitter request_id=%s " + "still_owned=%s deadline_expired=%s", + request_state.request_id, + request_still_owned, + now_monotonic >= request_deadline, + ) + request_state.clean_close_retry_result = False + return False + # A fresh hard-session replay may select a replacement account. + # The admission lease is account-scoped, so release the old + # account's lease before reconnecting; the post-reconnect path + # below acquires a lease for the account actually selected. + if fresh_hard_request_account_switch_allowed: + await self._release_request_state_account_response_create_lease(request_state) + if hard_owner_bound and not model_fallback_replay and not fresh_hard_request_account_switch_allowed: await self._reconnect_http_bridge_session( session, request_state=request_state, require_same_account=True, + **reconnect_reader_kwargs, ) elif require_preferred_reconnect: await self._reconnect_http_bridge_session( @@ -1644,12 +2130,23 @@ async def _retry_http_bridge_precreated_request( request_state=request_state, require_same_account=account_neutral_recovery, require_preferred_account=True, + **reconnect_reader_kwargs, + ) + elif clean_close_hard_continuation or clean_close_hard_continuity_anchor: + await self._reconnect_http_bridge_session( + session, + request_state=request_state, + # Continuity anchors (previous_response_id and turn-state) + # are account-bound. Do not migrate them while recovering a + # clean handoff close. + require_same_account=True, + **reconnect_reader_kwargs, ) else: await self._reconnect_http_bridge_session( session, request_state=request_state, - require_same_account=account_neutral_recovery, + **reconnect_reader_kwargs, ) if request_state.account_response_create_lease is None: current_settings = await _service_get_settings_cache().get() @@ -1703,10 +2200,16 @@ async def _retry_http_bridge_precreated_request( request_text = self._http_bridge_text_with_account_installation_id(session, request_state, request_text) await _send_http_bridge_request_text_with_archive_id(session, request_state, request_text) session.last_used_at = _service_time().monotonic() + request_state.clean_close_retry_result = True return True + except asyncio.CancelledError: + request_state.clean_close_retry_result = False + raise except UpstreamWebSocketTransportError: + request_state.clean_close_retry_result = False raise except Exception as exc: + request_state.clean_close_retry_result = False ( request_state.error_http_status_override, request_state.error_code_override, @@ -1723,6 +2226,8 @@ async def _retry_http_bridge_precreated_request( else: logger.warning("HTTP bridge pre-created retry failed", exc_info=True) return False + finally: + request_state.clean_close_retry_in_progress = False async def _retry_http_bridge_precreated_auth_request( self: Any, diff --git a/app/modules/proxy/_service/http_bridge/retry_circuit.py b/app/modules/proxy/_service/http_bridge/retry_circuit.py new file mode 100644 index 0000000000..5a21c10332 --- /dev/null +++ b/app/modules/proxy/_service/http_bridge/retry_circuit.py @@ -0,0 +1,433 @@ +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from typing import Any + +import anyio + +from app.core.metrics.prometheus import PROMETHEUS_AVAILABLE, http_bridge_retry_circuit_total +from app.modules.proxy._service.observability import _hash_identifier +from app.modules.proxy._service.support import _HTTPBridgeSession +from app.modules.proxy.durable_bridge_repository import DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS + +logger = logging.getLogger(__name__) + +_HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD = 2 +_HTTP_BRIDGE_RETRY_CIRCUIT_BASE_BACKOFF_SECONDS = 60.0 +_HTTP_BRIDGE_RETRY_CIRCUIT_MAX_BACKOFF_SECONDS = 600.0 +_HTTP_BRIDGE_RETRY_CIRCUIT_CLEAN_CLOSE_MAX_BACKOFF_SECONDS = 30.0 +_HTTP_BRIDGE_RETRY_CIRCUIT_HALF_OPEN_LEASE_SECONDS = 600.0 +_HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_DETAILS = frozenset( + { + "stream_incomplete", + "clean_close", + "stream_idle_timeout", + } +) +_HTTP_BRIDGE_RETRY_CIRCUIT_DETAIL_ALIASES = { + # These diagnostics describe the same ambiguous idle/incomplete + # transport class. Keep the durable contract to the three documented + # failure classes while retaining the more specific event in logs. + "upstream_keepalive_timeout": "stream_idle_timeout", + "missing_response_created_timeout": "stream_idle_timeout", + "response_create_gate_timeout_stuck_pending": "stream_idle_timeout", +} + + +@dataclass(slots=True) +class _HTTPBridgeRetryCircuitState: + consecutive_failures: int = 0 + cooldown_until: float = 0.0 + last_detail: str | None = None + last_touched_monotonic: float = 0.0 + persisted_updated_at_epoch: float = 0.0 + last_failure_monotonic: float = 0.0 + last_durable_load_monotonic: float = 0.0 + half_open_until: float = 0.0 + + +def _initialize_http_bridge_retry_circuit(service: Any) -> None: + service._http_bridge_retry_circuits = {} + service._http_bridge_retry_circuit_loaded_keys = set() + service._http_bridge_retry_circuit_persisted_keys = set() + service._http_bridge_retry_circuit_lock = anyio.Lock() + + +class _HTTPBridgeRetryCircuitMixin: + def _prune_http_bridge_retry_circuit_state(self: Any, now: float) -> None: + expiry = now - DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS + for key, state in list(self._http_bridge_retry_circuits.items()): + if state.last_touched_monotonic > expiry: + continue + self._http_bridge_retry_circuits.pop(key, None) + self._http_bridge_retry_circuit_loaded_keys.discard(key) + self._http_bridge_retry_circuit_persisted_keys.discard(key) + + async def _load_http_bridge_retry_circuit(self: Any, session: _HTTPBridgeSession) -> bool: + if session.key.strength != "hard": + return True + + now_monotonic = time.monotonic() + async with self._http_bridge_retry_circuit_lock: + self._prune_http_bridge_retry_circuit_state(now_monotonic) + local_state = self._http_bridge_retry_circuits.get(session.key) + if local_state is not None: + local_state.last_touched_monotonic = now_monotonic + try: + persisted = await self._durable_bridge.lookup_retry_circuit( + session_key_kind=session.key.affinity_kind, + session_key_value=session.key.affinity_key, + api_key_id=session.key.api_key_id, + ) + except Exception: + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="lookup_failed").inc() + logger.warning( + "Failed to load persisted HTTP bridge retry circuit bridge_kind=%s bridge_key=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + exc_info=True, + ) + return False + + if persisted is None: + # A durable miss clears state loaded from another replica, but it + # must not discard a failure recorded locally after the last + # durable read. That local circuit is the only protection against + # immediately replaying the same failing upstream request. + async with self._http_bridge_retry_circuit_lock: + local_state = self._http_bridge_retry_circuits.get(session.key) + locally_updated = bool( + local_state is not None + and local_state.last_failure_monotonic > local_state.last_durable_load_monotonic + ) + if session.key in self._http_bridge_retry_circuit_persisted_keys and not locally_updated: + self._http_bridge_retry_circuits.pop(session.key, None) + self._http_bridge_retry_circuit_loaded_keys.discard(session.key) + self._http_bridge_retry_circuit_persisted_keys.discard(session.key) + return True + + now_epoch = time.time() + if now_epoch - persisted.updated_at_epoch > DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS: + async with self._http_bridge_retry_circuit_lock: + stale_local_state = self._http_bridge_retry_circuits.get(session.key) + try: + await self._durable_bridge.purge_retry_circuit( + session_key_kind=session.key.affinity_kind, + session_key_value=session.key.affinity_key, + api_key_id=session.key.api_key_id, + expected_updated_at_epoch=persisted.updated_at_epoch, + ) + except Exception: + logger.warning( + "Failed to remove stale HTTP bridge retry circuit bridge_kind=%s bridge_key=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + exc_info=True, + ) + # Keep a newer process-local circuit when persistence is + # unavailable. The next failure can still open the local + # circuit even though the expired durable row remains. + return False + async with self._http_bridge_retry_circuit_lock: + current_local_state = self._http_bridge_retry_circuits.get(session.key) + local_state_is_newer = bool( + current_local_state is not None + and current_local_state.last_failure_monotonic > current_local_state.last_durable_load_monotonic + ) + if current_local_state is None or ( + current_local_state is stale_local_state and not local_state_is_newer + ): + self._http_bridge_retry_circuits.pop(session.key, None) + self._http_bridge_retry_circuit_loaded_keys.discard(session.key) + self._http_bridge_retry_circuit_persisted_keys.discard(session.key) + return True + + cooldown_remaining = max(0.0, persisted.cooldown_until_epoch - now_epoch) + persisted_cooldown_until = now_monotonic + cooldown_remaining + async with self._http_bridge_retry_circuit_lock: + self._http_bridge_retry_circuit_persisted_keys.add(session.key) + state = self._http_bridge_retry_circuits.get(session.key) + if state is None: + state = _HTTPBridgeRetryCircuitState(last_touched_monotonic=now_monotonic) + self._http_bridge_retry_circuits[session.key] = state + local_failure_is_newer = state.last_failure_monotonic > state.last_durable_load_monotonic + if persisted.updated_at_epoch > state.persisted_updated_at_epoch and not local_failure_is_newer: + state.consecutive_failures = max(0, persisted.consecutive_failures) + state.cooldown_until = persisted_cooldown_until + state.last_detail = persisted.last_detail + else: + state.consecutive_failures = max(state.consecutive_failures, max(0, persisted.consecutive_failures)) + state.cooldown_until = max(state.cooldown_until, persisted_cooldown_until) + if local_failure_is_newer: + state.last_detail = state.last_detail or persisted.last_detail + else: + state.last_detail = persisted.last_detail or state.last_detail + state.persisted_updated_at_epoch = max(state.persisted_updated_at_epoch, persisted.updated_at_epoch) + state.last_touched_monotonic = now_monotonic + state.last_durable_load_monotonic = now_monotonic + self._http_bridge_retry_circuit_loaded_keys.add(session.key) + return True + + async def _persist_http_bridge_retry_circuit( + self: Any, + session: _HTTPBridgeSession, + state: _HTTPBridgeRetryCircuitState, + ) -> None: + now_monotonic = time.monotonic() + now_wall = time.time() + threshold = max(1, _HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD) + async with self._http_bridge_retry_circuit_lock: + if self._http_bridge_retry_circuits.get(session.key) is not state: + return + consecutive_failures = state.consecutive_failures + cooldown_until = state.cooldown_until + last_detail = state.last_detail + persisted_updated_at_epoch = state.persisted_updated_at_epoch + base_backoff = max(0.001, _HTTP_BRIDGE_RETRY_CIRCUIT_BASE_BACKOFF_SECONDS) + if last_detail == "clean_close": + base_backoff = min( + base_backoff, + max(0.001, _HTTP_BRIDGE_RETRY_CIRCUIT_CLEAN_CLOSE_MAX_BACKOFF_SECONDS), + ) + try: + persisted = await self._durable_bridge.persist_retry_circuit( + session_key_kind=session.key.affinity_kind, + session_key_value=session.key.affinity_key, + api_key_id=session.key.api_key_id, + consecutive_failures=consecutive_failures, + cooldown_until_epoch=now_wall + max(0.0, cooldown_until - now_monotonic), + last_detail=last_detail, + updated_at_epoch=now_wall, + base_updated_at_epoch=persisted_updated_at_epoch, + failure_threshold=threshold, + conflict_cooldown_until_epoch=now_wall + base_backoff, + base_backoff_seconds=max(0.001, _HTTP_BRIDGE_RETRY_CIRCUIT_BASE_BACKOFF_SECONDS), + max_backoff_seconds=max(0.001, _HTTP_BRIDGE_RETRY_CIRCUIT_MAX_BACKOFF_SECONDS), + clean_close_max_backoff_seconds=max( + 0.001, + _HTTP_BRIDGE_RETRY_CIRCUIT_CLEAN_CLOSE_MAX_BACKOFF_SECONDS, + ), + ) + if persisted is not None: + persisted_cooldown_until = now_monotonic + max(0.0, persisted.cooldown_until_epoch - now_wall) + async with self._http_bridge_retry_circuit_lock: + current = self._http_bridge_retry_circuits.get(session.key) + if current is state: + local_failure_is_newer = state.last_failure_monotonic > state.last_durable_load_monotonic + if persisted.updated_at_epoch > state.persisted_updated_at_epoch and not local_failure_is_newer: + state.consecutive_failures = max(0, persisted.consecutive_failures) + state.cooldown_until = persisted_cooldown_until + state.last_detail = persisted.last_detail + else: + state.consecutive_failures = max(state.consecutive_failures, persisted.consecutive_failures) + state.cooldown_until = max(state.cooldown_until, persisted_cooldown_until) + if local_failure_is_newer: + state.last_detail = state.last_detail or persisted.last_detail + else: + state.last_detail = persisted.last_detail or state.last_detail + state.persisted_updated_at_epoch = max( + state.persisted_updated_at_epoch, + persisted.updated_at_epoch, + ) + # This write is now the durable baseline for the + # captured local failure. A failure recorded while + # the write was in flight still has a later + # monotonic timestamp and will remain dominant. + state.last_durable_load_monotonic = max( + state.last_durable_load_monotonic, + now_monotonic, + ) + async with self._http_bridge_retry_circuit_lock: + if self._http_bridge_retry_circuits.get(session.key) is state: + self._http_bridge_retry_circuit_persisted_keys.add(session.key) + except Exception: + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="persist_failed").inc() + logger.warning( + "Failed to persist HTTP bridge retry circuit bridge_kind=%s bridge_key=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + exc_info=True, + ) + + async def _http_bridge_precreated_retry_allowed( + self: Any, + session: _HTTPBridgeSession, + *, + allow_fresh_hard_account_switch: bool = False, + allow_proof_gated_continuity_replay: bool = False, + ) -> bool: + """Avoid replaying a repeatedly failing hard-affinity request in a tight loop.""" + if session.key.strength != "hard": + return True + + await self._load_http_bridge_retry_circuit(session) + now = time.monotonic() + async with self._http_bridge_retry_circuit_lock: + state = self._http_bridge_retry_circuits.get(session.key) + if state is None or state.cooldown_until <= now: + if ( + state is not None + and state.consecutive_failures >= _HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD + and state.half_open_until > now + and not allow_fresh_hard_account_switch + and not allow_proof_gated_continuity_replay + ): + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="suppressed").inc() + return False + if state is not None and state.cooldown_until > 0: + state.cooldown_until = 0.0 + state.half_open_until = now + _HTTP_BRIDGE_RETRY_CIRCUIT_HALF_OPEN_LEASE_SECONDS + logger.info( + "http_bridge_retry_circuit event=half_open bridge_kind=%s bridge_key=%s failures=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + state.consecutive_failures, + ) + return True + + retry_after = max(0.0, state.cooldown_until - now) + if allow_fresh_hard_account_switch: + logger.info( + "http_bridge_retry_circuit event=bypass_fresh_account_switch bridge_kind=%s " + "bridge_key=%s failures=%s retry_after_seconds=%.1f", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + state.consecutive_failures, + retry_after, + ) + return True + if allow_proof_gated_continuity_replay: + logger.info( + "http_bridge_retry_circuit event=bypass_proof_gated_continuity_replay bridge_kind=%s " + "bridge_key=%s failures=%s retry_after_seconds=%.1f", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + state.consecutive_failures, + retry_after, + ) + return True + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="suppressed").inc() + logger.info( + "http_bridge_retry_circuit event=suppressed bridge_kind=%s bridge_key=%s " + "failures=%s retry_after_seconds=%.1f detail=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + state.consecutive_failures, + retry_after, + state.last_detail, + ) + return False + + async def _http_bridge_precreated_retry_cooldown_seconds(self: Any, session: _HTTPBridgeSession) -> float: + if session.key.strength != "hard": + return 0.0 + + await self._load_http_bridge_retry_circuit(session) + now = time.monotonic() + async with self._http_bridge_retry_circuit_lock: + state = self._http_bridge_retry_circuits.get(session.key) + if state is None: + return 0.0 + return max(0.0, state.cooldown_until - now) + + async def _record_http_bridge_retry_circuit_failure( + self: Any, + session: _HTTPBridgeSession, + *, + detail: str, + ) -> None: + detail = _HTTP_BRIDGE_RETRY_CIRCUIT_DETAIL_ALIASES.get(detail, detail) + if session.key.strength != "hard" or detail not in _HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_DETAILS: + return + + await self._load_http_bridge_retry_circuit(session) + threshold = max(1, _HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD) + base_backoff = max(0.001, _HTTP_BRIDGE_RETRY_CIRCUIT_BASE_BACKOFF_SECONDS) + max_backoff = max(base_backoff, _HTTP_BRIDGE_RETRY_CIRCUIT_MAX_BACKOFF_SECONDS) + clean_close_max_backoff = max(0.001, _HTTP_BRIDGE_RETRY_CIRCUIT_CLEAN_CLOSE_MAX_BACKOFF_SECONDS) + now = time.monotonic() + async with self._http_bridge_retry_circuit_lock: + state = self._http_bridge_retry_circuits.setdefault( + session.key, + _HTTPBridgeRetryCircuitState(last_touched_monotonic=now), + ) + state.last_touched_monotonic = now + state.last_failure_monotonic = now + state.half_open_until = 0.0 + state.consecutive_failures += 1 + state.last_detail = detail + if state.consecutive_failures >= threshold: + backoff = min( + max_backoff, + base_backoff * (2 ** min(state.consecutive_failures - threshold, 30)), + ) + if detail == "clean_close": + backoff = min(backoff, clean_close_max_backoff) + state.cooldown_until = max(state.cooldown_until, now + backoff) + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="opened").inc() + logger.warning( + "http_bridge_retry_circuit event=opened bridge_kind=%s bridge_key=%s " + "failures=%s cooldown_seconds=%.1f detail=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + state.consecutive_failures, + backoff, + detail, + ) + await self._persist_http_bridge_retry_circuit(session, state) + async with self._http_bridge_retry_circuit_lock: + if self._http_bridge_retry_circuits.get(session.key) is state: + self._http_bridge_retry_circuit_loaded_keys.add(session.key) + + async def _clear_http_bridge_retry_circuit(self: Any, session: _HTTPBridgeSession) -> None: + if session.key.strength != "hard": + return + + durable_load_succeeded = await self._load_http_bridge_retry_circuit(session) + async with self._http_bridge_retry_circuit_lock: + state = self._http_bridge_retry_circuits.pop(session.key, None) + self._http_bridge_retry_circuit_loaded_keys.discard(session.key) + self._http_bridge_retry_circuit_persisted_keys.discard(session.key) + expected_updated_at_epoch = ( + state.persisted_updated_at_epoch if state is not None and state.persisted_updated_at_epoch > 0 else None + ) + # A confirmed miss has no version fence to protect a row created + # concurrently, so leave the durable row untouched when no state was + # observed. Preserve the existing best-effort clear on read failures, + # which is still useful for settling a row after a transient outage. + if durable_load_succeeded and (state is None or expected_updated_at_epoch is None): + return + try: + # Clearing is idempotent and must be attempted even when the + # preceding lookup failed; a successful request should settle + # a previously persisted circuit after a transient read error. + await self._durable_bridge.clear_retry_circuit( + session_key_kind=session.key.affinity_kind, + session_key_value=session.key.affinity_key, + api_key_id=session.key.api_key_id, + expected_updated_at_epoch=expected_updated_at_epoch, + ) + except Exception: + logger.warning( + "Failed to clear persisted HTTP bridge retry circuit bridge_kind=%s bridge_key=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + exc_info=True, + ) + if state is None: + return + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="reset").inc() + logger.info( + "http_bridge_retry_circuit event=reset bridge_kind=%s bridge_key=%s failures=%s", + session.key.affinity_kind, + _hash_identifier(session.key.affinity_key), + state.consecutive_failures, + ) diff --git a/app/modules/proxy/_service/http_bridge/service_stubs.py b/app/modules/proxy/_service/http_bridge/service_stubs.py index 93ab9113d4..b4f7ac4a27 100644 --- a/app/modules/proxy/_service/http_bridge/service_stubs.py +++ b/app/modules/proxy/_service/http_bridge/service_stubs.py @@ -84,7 +84,8 @@ def _service_inline_input_image_urls() -> Any: def _stream_keepalive_max_count() -> int: - return int(_service_global_or("_STREAM_KEEPALIVE_MAX_COUNT", _STREAM_KEEPALIVE_MAX_COUNT)) + service_override = int(_service_global_or("_STREAM_KEEPALIVE_MAX_COUNT", _STREAM_KEEPALIVE_MAX_COUNT)) + return max(1, service_override) def _prewarm_response_timeout_seconds() -> float: diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index 1fce2706ea..0a484f900e 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -34,6 +34,7 @@ from app.core.clients.proxy import codex_control_request as core_codex_control_request # noqa: F401 from app.core.clients.proxy import compact_responses as core_compact_responses # noqa: F401 from app.core.clients.proxy import transcribe_audio as core_transcribe_audio # noqa: F401 +from app.core.clients.proxy_websocket import UpstreamWebSocketTransportError from app.core.errors import ( openai_error, response_failed_event, @@ -41,6 +42,9 @@ from app.core.metrics.prometheus import ( PROMETHEUS_AVAILABLE, bridge_durable_recover_total, + http_bridge_retry_circuit_total, + stream_idle_timeout_total, + stream_keepalive_sent_total, ) from app.core.openai.requests import ( ResponsesRequest, @@ -48,7 +52,9 @@ from app.core.types import JsonValue from app.core.utils.request_id import ensure_request_id, ensure_request_scope_id from app.core.utils.sse import format_sse_event, parse_sse_data_json +from app.core.utils.time import utcnow from app.db.models import ( + HttpBridgeSessionState, StickySessionKind, ) from app.modules.api_keys.service import ( @@ -66,6 +72,7 @@ ) from app.modules.proxy._service.http_bridge.helpers import ( _effective_http_bridge_idle_ttl_seconds, + _http_bridge_durable_lease_ttl_seconds, _http_bridge_durable_lookup_allows_turn_state_takeover, _http_bridge_is_context_overflow_error, _http_bridge_is_previous_response_owner_unavailable, @@ -204,6 +211,7 @@ without_http_bridge_session_affinity_headers, ) from app.modules.proxy.durable_bridge_coordinator import DurableBridgeLookup +from app.modules.proxy.durable_bridge_repository import durable_bridge_hash from app.modules.proxy.helpers import ( _normalize_error_code, ) @@ -220,10 +228,32 @@ _RESPONSE_CREATE_GATE_RETRY_SLEEP_SECONDS = 10.0 +def _http_bridge_continuity_bound_without_safe_replay(request_state: _WebSocketRequestState) -> bool: + """Return whether retrying would require replaying an unsafe continuation.""" + if request_state.previous_response_id is not None: + return not (request_state.fresh_upstream_request_is_retry_safe and request_state.fresh_upstream_request_text) + return request_state.hard_continuity_anchor and not ( + request_state.fresh_upstream_request_is_retry_safe and request_state.fresh_upstream_request_text + ) + + def _http_bridge_payload_is_account_neutral_fresh_replay(payload: ResponsesRequest) -> bool: return responses_payload_is_account_neutral_fresh_replay(payload.to_payload()) +def _apply_http_bridge_downstream_turn_state( + request_state: _WebSocketRequestState, + *, + downstream_turn_state: str | None, + incoming_turn_state_header: str | None, +) -> None: + if downstream_turn_state is None: + return + request_state.session_id = _normalize_session_id(downstream_turn_state) + if incoming_turn_state_header is not None or request_state.previous_response_id is not None: + request_state.hard_continuity_anchor = True + + def _proxy_error_code_message(exc: ProxyResponseError) -> tuple[str | None, str | None]: error = exc.payload.get("error") if isinstance(exc.payload, dict) else None if not isinstance(error, dict): @@ -233,6 +263,22 @@ def _proxy_error_code_message(exc: ProxyResponseError) -> tuple[str | None, str return (str(code) if code is not None else None, str(message) if message is not None else None) +_HTTP_BRIDGE_AMBIGUOUS_RECOVERY_ERROR_CODES = frozenset( + { + "stream_incomplete", + "stream_idle_timeout", + "upstream_request_timeout", + } +) + + +def _http_bridge_error_is_ambiguous_transport(exc: ProxyResponseError) -> bool: + """Return whether an error leaves upstream acceptance genuinely unknown.""" + + code, _message = _proxy_error_code_message(exc) + return code in _HTTP_BRIDGE_AMBIGUOUS_RECOVERY_ERROR_CODES + + def _http_bridge_account_capacity_wait_seconds(exc: ProxyResponseError) -> float | None: code, message = _proxy_error_code_message(exc) if code == "capacity_exhausted_active_sessions": @@ -899,6 +945,7 @@ def prepare_bridge_request( untrimmed_effective_payload = payload proxy_injected_previous_response_id = False fresh_upstream_request_text: str | None = None + client_full_resend_fresh_upstream_request_text: str | None = None previous_response_trimmed_input_count: int | None = None previous_response_trimmed_input_fingerprint: str | None = None durable_full_resend_anchor_count: int | None = None @@ -907,6 +954,11 @@ def prepare_bridge_request( durable_full_resend_is_account_neutral: bool | None = None durable_full_resend_has_safe_fresh_context = False durable_full_resend_retains_prior_output = False + durable_recovery_attempt_fingerprint: str | None = None + durable_recovery_attempt_available = False + durable_recovery_attempt_claimed = False + durable_recovery_attempt_session_id: str | None = None + durable_recovery_attempt_owner_epoch: int | None = None force_local_recovery_creation = False payload_looks_like_full_resend = _http_bridge_payload_looks_like_full_resend(payload) @@ -950,6 +1002,109 @@ def classify_durable_full_resend( durable_full_resend_anchor_fingerprint, durable_full_resend_has_safe_fresh_context, ) = classify_durable_full_resend(durable_lookup) + if ( + durable_full_resend_has_safe_fresh_context + and durable_full_resend_anchor_count is not None + and isinstance(payload.input, list) + ): + replay_projection = project_responses_input_for_account_neutral_fresh_replay( + cast(list[JsonValue], payload.input), + stored_count=durable_full_resend_anchor_count, + ) + if replay_projection is not None: + durable_full_resend_retains_prior_output = responses_input_suffix_retains_prior_output( + replay_projection.input_items, + stored_count=replay_projection.stored_prefix_count, + ) + durable_full_resend_fresh_payload = _http_bridge_payload_without_previous_response_id( + payload + ).model_copy(update={"input": replay_projection.input_items}) + durable_full_resend_is_account_neutral = _http_bridge_payload_is_account_neutral_fresh_replay( + durable_full_resend_fresh_payload + ) + _fresh_state, fresh_replay_text = prepare_bridge_request( + _http_bridge_payload_without_previous_response_id(payload) + ) + del _fresh_state + durable_recovery_attempt_fingerprint = durable_bridge_hash(fresh_replay_text) + if durable_lookup is not None and durable_full_resend_is_account_neutral: + try: + existing_attempt = await self._durable_bridge.lookup_recovery_attempt( + session_id=durable_lookup.session_id, + request_fingerprint=durable_recovery_attempt_fingerprint, + ) + if existing_attempt is not None and ( + durable_lookup.state != HttpBridgeSessionState.ACTIVE + or not durable_lookup.lease_is_active(now=utcnow()) + ): + claim_instance_id = _service_get_settings().http_responses_session_bridge_instance_id + claim_owner_epoch = durable_lookup.owner_epoch + owner_is_current = ( + durable_lookup.owner_instance_id == claim_instance_id + and durable_lookup.lease_is_active(now=utcnow()) + ) + if not owner_is_current: + claimed_session = await self._durable_bridge.claim_live_session( + session_key_kind=durable_lookup.canonical_kind, + session_key_value=durable_lookup.canonical_key, + api_key_id=bridge_session_key.api_key_id, + instance_id=claim_instance_id, + lease_ttl_seconds=_http_bridge_durable_lease_ttl_seconds(), + account_id=durable_lookup.account_id, + model=payload.model, + service_tier=None, + latest_turn_state=durable_lookup.latest_turn_state, + latest_response_id=None, + # Revalidate the stale lookup under the + # row lock; an active owner that appeared + # after the lookup must not be displaced. + allow_takeover=False, + ) + if claimed_session.owner_instance_id != claim_instance_id: + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses recovery ownership changed; retry the request.", + ), + ) + claim_owner_epoch = claimed_session.owner_epoch + claimed = await self._durable_bridge.mark_recovery_attempt_replayed( + session_id=durable_lookup.session_id, + api_key_id=bridge_session_key.api_key_id, + instance_id=claim_instance_id, + owner_epoch=claim_owner_epoch, + request_fingerprint=durable_recovery_attempt_fingerprint, + ) + if not claimed: + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses recovery ownership changed; retry the request.", + ), + ) + durable_recovery_attempt_claimed = True + durable_recovery_attempt_available = False + durable_recovery_attempt_session_id = durable_lookup.session_id + durable_recovery_attempt_owner_epoch = claim_owner_epoch + elif existing_attempt is None: + # No prior attempt owns this fingerprint. The + # request-submit path will journal it immediately + # before dispatch, and an ambiguous transport + # outcome may then consume the one replay fence. + durable_recovery_attempt_available = True + except ProxyResponseError: + raise + except Exception: + logger.warning("Failed to claim HTTP bridge recovery attempt", exc_info=True) + raise ProxyResponseError( + 502, + openai_error( + "bridge_continuity_persistence_failed", + "HTTP responses recovery state could not be claimed; retry the request.", + ), + ) durable_anchor_trimmable = durable_full_resend_anchor_count is not None durable_model_transition_lookup = ( durable_lookup @@ -1065,8 +1220,11 @@ def classify_durable_full_resend( request_state, text_data = prepare_bridge_request(effective_payload) request_state.enforce_openai_sdk_contract = enforce_openai_sdk_contract request_state.affinity_policy = affinity - if downstream_turn_state is not None: - request_state.session_id = _normalize_session_id(downstream_turn_state) + _apply_http_bridge_downstream_turn_state( + request_state, + downstream_turn_state=downstream_turn_state, + incoming_turn_state_header=incoming_turn_state_header, + ) if previous_response_trimmed_input_count is not None: request_state.input_item_count = previous_response_trimmed_input_count request_state.input_full_fingerprint = previous_response_trimmed_input_fingerprint @@ -1176,6 +1334,24 @@ def classify_durable_full_resend( # Only the trim branch below (which verifies the stored prefix # fingerprint) is allowed to flip this flag to ``True``. request_state.fresh_upstream_request_is_retry_safe = False + elif ( + effective_payload.previous_response_id is not None + and payload_looks_like_full_resend + and durable_full_resend_anchor_count is not None + and durable_full_resend_has_safe_fresh_context + ): + # A client-provided full resend carries the same proof as a + # proxy-injected anchor: the stored prefix matches and the fresh + # suffix retains the prior output/tool context. Capture the + # verified anchor-free body so a retry can use it without sending + # previous_response_id again. + client_full_resend_payload = _http_bridge_payload_without_previous_response_id(untrimmed_effective_payload) + _fresh_state, client_full_resend_fresh_upstream_request_text = prepare_bridge_request( + client_full_resend_payload + ) + del _fresh_state + request_state.fresh_upstream_request_text = client_full_resend_fresh_upstream_request_text + request_state.fresh_upstream_request_is_retry_safe = True settings = _service_get_settings() request_deadline = request_state.started_at + _http_bridge_request_budget_seconds(settings) session_creation_headers = ( @@ -1233,6 +1409,7 @@ def switch_to_account_neutral_replay() -> None: nonlocal effective_payload nonlocal file_required_preferred_account nonlocal force_local_recovery_creation + nonlocal client_full_resend_fresh_upstream_request_text nonlocal fresh_upstream_request_text nonlocal incoming_turn_state_header nonlocal previous_response_trimmed_input_count @@ -1284,6 +1461,7 @@ def switch_to_account_neutral_replay() -> None: untrimmed_effective_payload = fresh_payload proxy_injected_previous_response_id = False fresh_upstream_request_text = None + client_full_resend_fresh_upstream_request_text = None previous_response_trimmed_input_count = None previous_response_trimmed_input_fingerprint = None durable_full_resend_anchor_count = None @@ -1293,6 +1471,13 @@ def switch_to_account_neutral_replay() -> None: durable_lookup = None file_required_preferred_account = False + if durable_recovery_attempt_claimed: + switch_to_account_neutral_replay() + request_state.recovery_attempt_fingerprint = durable_recovery_attempt_fingerprint + request_state.recovery_attempt_session_id = durable_recovery_attempt_session_id + request_state.recovery_attempt_owner_epoch = durable_recovery_attempt_owner_epoch + request_state.recovery_attempt_claimed = True + if required_continuity_owner_missing: owner_unavailable = ProxyResponseError( 502, @@ -1394,6 +1579,15 @@ def switch_to_account_neutral_replay() -> None: raise continue raise + _log_http_bridge_event( + "owner_unavailable_fresh_resend", + bridge_session_key, + account_id=request_state.preferred_account_id, + model=payload.model, + detail="outcome=fresh_full_resend_without_anchor", + cache_key_family=bridge_session_key.affinity_kind, + model_class=_extract_model_class(payload.model) if payload.model else None, + ) switch_to_account_neutral_replay() continue break @@ -1736,8 +1930,11 @@ def switch_to_account_neutral_replay() -> None: ) retry_request_state.enforce_openai_sdk_contract = enforce_openai_sdk_contract retry_request_state.affinity_policy = affinity - if downstream_turn_state is not None: - retry_request_state.session_id = _normalize_session_id(downstream_turn_state) + _apply_http_bridge_downstream_turn_state( + retry_request_state, + downstream_turn_state=downstream_turn_state, + incoming_turn_state_header=incoming_turn_state_header, + ) retry_request_state.transport = _REQUEST_TRANSPORT_HTTP retry_request_state.request_stage = ( request_state.request_stage if owner_forward_fresh_replay else "reattach" @@ -1916,8 +2113,11 @@ def switch_to_account_neutral_replay() -> None: request_state, text_data = prepare_bridge_request(submit_payload) request_state.enforce_openai_sdk_contract = enforce_openai_sdk_contract request_state.affinity_policy = affinity - if downstream_turn_state is not None: - request_state.session_id = _normalize_session_id(downstream_turn_state) + _apply_http_bridge_downstream_turn_state( + request_state, + downstream_turn_state=downstream_turn_state, + incoming_turn_state_header=incoming_turn_state_header, + ) request_state.transport = _REQUEST_TRANSPORT_HTTP request_state.request_stage = _http_bridge_request_stage( headers=headers, @@ -1950,6 +2150,9 @@ def switch_to_account_neutral_replay() -> None: if store_context_trim_applied else previous_request_state.fresh_upstream_request_is_retry_safe ) + elif client_full_resend_fresh_upstream_request_text is not None: + request_state.fresh_upstream_request_text = client_full_resend_fresh_upstream_request_text + request_state.fresh_upstream_request_is_retry_safe = True initial_handoff_session = session initial_handoff_scope_id = ensure_request_scope_id() if original_request_unanchored else None if initial_handoff_scope_id is not None: @@ -1974,6 +2177,29 @@ def switch_to_account_neutral_replay() -> None: ) try: yielded_any = False + durable_recovery_fresh_replay = False + retry_request_state: _WebSocketRequestState | None = None + + async def rollback_pre_dispatch_recovery_claim() -> None: + if not ( + durable_recovery_fresh_replay + and (retry_request_state is None or not retry_request_state.recovery_attempt_dispatched) + and durable_recovery_attempt_fingerprint is not None + and durable_recovery_attempt_session_id is not None + and durable_recovery_attempt_owner_epoch is not None + ): + return + try: + await self._durable_bridge.rollback_recovery_attempt_replayed( + session_id=durable_recovery_attempt_session_id, + api_key_id=bridge_session_key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=durable_recovery_attempt_owner_epoch, + request_fingerprint=durable_recovery_attempt_fingerprint, + ) + except Exception: + logger.warning("Failed to roll back pre-dispatch HTTP bridge recovery claim", exc_info=True) + async for event_block in session_events: yield event_block yielded_any = True @@ -2192,6 +2418,62 @@ def switch_to_account_neutral_replay() -> None: except Exception: pass return + if ( + durable_recovery_attempt_available + and durable_recovery_attempt_fingerprint is not None + and _http_bridge_error_is_ambiguous_transport(exc) + and request_state.response_event_count == 0 + and request_state.previous_response_id is not None + and session.durable_session_id is not None + and session.durable_owner_epoch is not None + ): + try: + marked = await self._durable_bridge.mark_recovery_attempt_replayed( + session_id=session.durable_session_id, + api_key_id=bridge_session_key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + request_fingerprint=durable_recovery_attempt_fingerprint, + ) + except Exception: + marked = False + logger.warning("Failed to fence HTTP bridge recovery attempt", exc_info=True) + if marked: + durable_recovery_fresh_replay = True + recovery_origin_session_id = request_state.recovery_attempt_session_id or session.durable_session_id + recovery_origin_owner_epoch = ( + request_state.recovery_attempt_owner_epoch or session.durable_owner_epoch + ) + durable_recovery_attempt_session_id = recovery_origin_session_id + durable_recovery_attempt_owner_epoch = recovery_origin_owner_epoch + _log_http_bridge_event( + "durable_recovery_fresh_replay", + bridge_session_key, + account_id=session.account.id, + model=effective_payload.model, + detail="outcome=new_account_neutral_upstream_session", + cache_key_family=bridge_session_key.affinity_kind, + model_class=_extract_model_class(effective_payload.model) if effective_payload.model else None, + owner_check_applied=True, + ) + await self._reset_http_bridge_session_after_local_terminal_error( + session, + error_code="stream_incomplete", + error_message="Upstream websocket closed before response.completed", + preserve_durable_lease=True, + ) + switch_to_account_neutral_replay() + request_state.recovery_attempt_fingerprint = durable_recovery_attempt_fingerprint + request_state.recovery_attempt_session_id = recovery_origin_session_id + request_state.recovery_attempt_owner_epoch = recovery_origin_owner_epoch + recovery_path = "durable_recovery_fresh_replay" + retry_payload = effective_payload + retry_previous_response_id = None + retry_request_stage = "durable_recovery" + retry_preferred_account_id = None + allow_previous_response_recovery_rebind = False + else: + durable_recovery_attempt_available = False is_context_overflow = _http_bridge_is_context_overflow_error(exc) should_rollover_after_context_overflow = _http_bridge_should_rollover_after_context_overflow( exc, @@ -2210,6 +2492,7 @@ def switch_to_account_neutral_replay() -> None: not should_attempt_previous_response_recovery and not should_rollover_after_context_overflow and not should_attempt_context_overflow_fresh_turn_recovery + and not durable_recovery_fresh_replay ): if is_context_overflow: _log_http_bridge_event( @@ -2224,7 +2507,9 @@ def switch_to_account_neutral_replay() -> None: ) raise - if should_attempt_context_overflow_fresh_turn_recovery: + if durable_recovery_fresh_replay: + pass + elif should_attempt_context_overflow_fresh_turn_recovery: if PROMETHEUS_AVAILABLE and bridge_durable_recover_total is not None: bridge_durable_recover_total.labels(path="context_overflow_fresh_turn").inc() _log_http_bridge_event( @@ -2390,8 +2675,16 @@ def switch_to_account_neutral_replay() -> None: reservation=retry_api_key_reservation, ) retry_request_state.enforce_openai_sdk_contract = enforce_openai_sdk_contract - if downstream_turn_state is not None: - retry_request_state.session_id = _normalize_session_id(downstream_turn_state) + if durable_recovery_fresh_replay and durable_recovery_attempt_fingerprint is not None: + retry_request_state.recovery_attempt_fingerprint = durable_recovery_attempt_fingerprint + retry_request_state.recovery_attempt_session_id = request_state.recovery_attempt_session_id + retry_request_state.recovery_attempt_owner_epoch = request_state.recovery_attempt_owner_epoch + retry_request_state.recovery_attempt_claimed = True + _apply_http_bridge_downstream_turn_state( + retry_request_state, + downstream_turn_state=downstream_turn_state, + incoming_turn_state_header=incoming_turn_state_header, + ) 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 @@ -2415,6 +2708,7 @@ def switch_to_account_neutral_replay() -> None: except Exception: pass except BaseException: + await rollback_pre_dispatch_recovery_claim() if retry_reservation_reacquired and retry_api_key_reservation is not None: await self._release_websocket_reservation(retry_api_key_reservation) raise @@ -2441,6 +2735,7 @@ async def _reset_http_bridge_session_after_local_terminal_error( *, error_code: str, error_message: str, + preserve_durable_lease: bool = False, ) -> None: async with self._http_bridge_lock: if self._http_bridge_sessions.get(session.key) is session: @@ -2457,7 +2752,7 @@ async def _reset_http_bridge_session_after_local_terminal_error( api_key=None, response_create_gate=session.response_create_gate, ) - await self._close_http_bridge_session(session) + await self._close_http_bridge_session(session, release_durable_session=not preserve_durable_lease) async def _stream_http_bridge_session_events( self: Any, @@ -2478,7 +2773,128 @@ async def _stream_http_bridge_session_events( key=session.key.affinity_key, ) request_state.propagate_http_errors = propagate_http_errors + + async def retry_precreated_for_idle_recovery( + *, + downstream_response_id: str, + after_circuit_cooldown: bool = False, + ) -> tuple[bool, str | None]: + # The reader may already be performing the bounded additional + # clean-close replay (including its jitter). Wait for that result + # instead of interpreting the in-progress flag as a terminal idle + # failure and detaching the request underneath the replay. + if request_state.clean_close_retry_in_progress: + while request_state.clean_close_retry_in_progress: + if _service_time().monotonic() >= request_deadline: + return ( + False, + format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Clean-close recovery exceeded the request budget", + response_id=downstream_response_id, + ), + ) + ), + ) + await asyncio.sleep(0.01) + if request_state.clean_close_retry_result is True: + return True, None + try: + return ( + await self._retry_http_bridge_precreated_request( + session, + restart_reader=True, + ), + None, + ) + except UpstreamWebSocketTransportError as exc: + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + logger.info( + "HTTP bridge stream idle recovery retry%s failed with transport error request_id=%s error_code=%s", + " after circuit cooldown" if after_circuit_cooldown else "", + request_state.request_id, + exc.error_code, + ) + return ( + False, + format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + exc.error_code, + str(exc), + response_id=downstream_response_id, + ), + ) + ), + ) + + def continuity_bound_without_safe_replay() -> bool: + """Do not hold a client stream through a cooldown we cannot use.""" + return _http_bridge_continuity_bound_without_safe_replay(request_state) + + async def startup_continuity_cooldown_terminal_event() -> str | None: + if ( + session.key.strength != "hard" + or not continuity_bound_without_safe_replay() + or request_state.response_id is not None + or request_state.response_event_count > 0 + ): + return None + retry_cooldown_seconds = await self._http_bridge_precreated_retry_cooldown_seconds(session) + if retry_cooldown_seconds <= 0: + return None + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + _record_continuity_fail_closed( + surface="http_bridge", + reason="retry_circuit_cooldown_continuity_bound", + previous_response_id=request_state.previous_response_id, + session_id=downstream_turn_state or request_state.session_id, + ) + logger.info( + "HTTP bridge stream idle timeout fail-closed before submit without safe replay " + "request_id=%s retry_after_seconds=%.1f", + request_state.request_id, + retry_cooldown_seconds, + ) + # This path returns before the request is submitted, so the normal + # detach/finally cleanup cannot settle an API-key reservation. + # Release it before handing the synthetic terminal event to the + # non-streaming collector. + await self._release_websocket_request_state_reservation(request_state) + request_state.api_key_reservation = None + if propagate_http_errors: + raise ProxyResponseError( + 503, + openai_error( + "upstream_request_timeout", + "HTTP responses session bridge is cooling down after repeated upstream " + "timeouts; retry shortly.", + error_type="server_error", + ), + retry_after_seconds=max(1, math.ceil(retry_cooldown_seconds)), + ) + return format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Upstream did not respond within the keepalive window", + response_id=_websocket_downstream_response_id(request_state), + ), + ) + ) + while True: + startup_terminal_event = await startup_continuity_cooldown_terminal_event() + if startup_terminal_event is not None: + yield startup_terminal_event + return try: if account_neutral_recovery: await self._submit_http_bridge_request( @@ -2554,6 +2970,59 @@ async def _stream_http_bridge_session_events( raise continue break + event_queue = request_state.event_queue + assert event_queue is not None + initial_retry_cooldown_seconds = await self._http_bridge_precreated_retry_cooldown_seconds(session) + if ( + initial_retry_cooldown_seconds > 0 + and session.key.strength == "hard" + and continuity_bound_without_safe_replay() + and request_state.response_id is None + and request_state.response_event_count == 0 + and event_queue.empty() + ): + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + _record_continuity_fail_closed( + surface="http_bridge", + reason="retry_circuit_cooldown_continuity_bound", + previous_response_id=request_state.previous_response_id, + session_id=downstream_turn_state or request_state.session_id, + ) + logger.info( + "HTTP bridge stream idle timeout fail-closed at startup without safe replay " + "request_id=%s retry_after_seconds=%.1f", + request_state.request_id, + initial_retry_cooldown_seconds, + ) + terminal_event = format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Upstream did not respond within the keepalive window", + response_id=_websocket_downstream_response_id(request_state), + ), + ) + ) + # The request was submitted before the durable cooldown refresh, + # so detach it before returning. This releases the response-create + # gate, reservation, and pending queue entry while marking the + # upstream handoff for retirement. + await self._detach_http_bridge_request(session, request_state=request_state) + if propagate_http_errors: + raise ProxyResponseError( + 503, + openai_error( + "upstream_request_timeout", + "HTTP responses session bridge is cooling down after repeated upstream " + "timeouts; retry shortly.", + error_type="server_error", + ), + retry_after_seconds=max(1, math.ceil(initial_retry_cooldown_seconds)), + ) + yield terminal_event + return try: if downstream_turn_state is not None and not account_neutral_recovery: await self._register_http_bridge_turn_state(session, downstream_turn_state) @@ -2567,6 +3036,8 @@ async def _stream_http_bridge_session_events( yielded_any = False keepalive_sent = False keepalive_count = 0 + circuit_keepalive_waiting = False + circuit_keepalive_until: float | None = None while True: keepalive_interval = getattr(_service_get_settings(), "sse_keepalive_interval_seconds", 10.0) if keepalive_interval > 0: @@ -2577,11 +3048,31 @@ async def _stream_http_bridge_session_events( "stream_idle_timeout_seconds", keepalive_interval * stream_keepalive_max_count, ) - max_keepalive_count = max( - stream_keepalive_max_count, - math.ceil(max(0.001, stream_idle_timeout_seconds) / keepalive_interval), + response_started = bool( + request_state.response_id + or request_state.replay_downstream_response_id + or request_state.response_event_count > 0 + or request_state.latency_response_created_ms is not None + ) + max_keepalive_count = ( + max( + stream_keepalive_max_count, + math.ceil(max(0.001, stream_idle_timeout_seconds) / keepalive_interval), + ) + if response_started + else stream_keepalive_max_count ) wait_timeout = keepalive_interval + if circuit_keepalive_waiting: + # Once the circuit is cooling down, wake at the actual + # expiry instead of waiting through another full + # keepalive interval before checking it again. + max_keepalive_count = 1 + if circuit_keepalive_until is not None: + wait_timeout = min( + wait_timeout, + max(0.001, circuit_keepalive_until - _service_time().monotonic()), + ) if not yielded_any and not keepalive_sent: wait_timeout = max(wait_timeout, _http_bridge_startup_keepalive_grace_seconds()) try: @@ -2621,29 +3112,181 @@ async def _stream_http_bridge_session_events( continue keepalive_count += 1 downstream_response_id = _websocket_downstream_response_id(request_state) - if keepalive_count > max_keepalive_count: - logger.info( - "HTTP bridge stream idle timeout request_id=%s keepalive_count=%s " - "max_keepalive_count=%s", - request_state.request_id, - keepalive_count, - max_keepalive_count, - ) - yield format_sse_event( - cast( - Mapping[str, JsonValue], - response_failed_event( - "stream_idle_timeout", - "Upstream did not respond within the keepalive window", - response_id=downstream_response_id, - ), + if keepalive_count >= max_keepalive_count: + if not response_started: + retried = False + if not circuit_keepalive_waiting: + retried, terminal_event = await retry_precreated_for_idle_recovery( + downstream_response_id=downstream_response_id, + ) + if terminal_event is not None: + yield terminal_event + break + if retried: + logger.info( + "HTTP bridge stream idle recovery retried pre-response request_id=%s", + request_state.request_id, + ) + keepalive_count = 0 + keepalive_sent = False + yielded_any = False + continue + retry_cooldown_seconds = await self._http_bridge_precreated_retry_cooldown_seconds( + session ) - ) - break - if propagate_http_errors and request_state.response_id is None: + fresh_replay_is_safe = bool( + request_state.fresh_upstream_request_is_retry_safe + and request_state.fresh_upstream_request_text + ) + continuity_bound = continuity_bound_without_safe_replay() + if retry_cooldown_seconds > 0 and ( + continuity_bound or (session.key.strength == "hard" and not fresh_replay_is_safe) + ): + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + _record_continuity_fail_closed( + surface="http_bridge", + reason=( + "retry_circuit_cooldown_continuity_bound" + if continuity_bound + else "retry_circuit_cooldown_no_safe_replay" + ), + previous_response_id=request_state.previous_response_id, + session_id=downstream_turn_state or request_state.session_id, + ) + logger.info( + "HTTP bridge stream idle timeout fail-closed without safe replay " + "request_id=%s retry_after_seconds=%.1f continuity_bound=%s", + request_state.request_id, + retry_cooldown_seconds, + continuity_bound, + ) + yield format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Upstream did not respond within the keepalive window", + response_id=downstream_response_id, + ), + ) + ) + break + if retry_cooldown_seconds > 0: + retry_cooldown_remaining_budget = max( + 0.0, + request_deadline - _service_time().monotonic(), + ) + if retry_cooldown_seconds >= retry_cooldown_remaining_budget: + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + logger.info( + "HTTP bridge stream idle timeout during retry circuit cooldown " + "request_id=%s retry_after_seconds=%.1f remaining_budget_seconds=%.1f", + request_state.request_id, + retry_cooldown_seconds, + retry_cooldown_remaining_budget, + ) + yield format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Upstream retry circuit cooldown exceeds the request budget", + response_id=downstream_response_id, + ), + ) + ) + break + circuit_keepalive_waiting = True + keepalive_count = 0 + circuit_keepalive_until = _service_time().monotonic() + retry_cooldown_seconds + if PROMETHEUS_AVAILABLE and http_bridge_retry_circuit_total is not None: + http_bridge_retry_circuit_total.labels(outcome="keepalive").inc() + logger.info( + "HTTP bridge stream waiting during retry circuit cooldown " + "request_id=%s retry_after_seconds=%.1f", + request_state.request_id, + retry_cooldown_seconds, + ) + else: + was_circuit_keepalive_waiting = circuit_keepalive_waiting + circuit_keepalive_waiting = False + circuit_keepalive_until = None + if was_circuit_keepalive_waiting and not retried: + retried, terminal_event = await retry_precreated_for_idle_recovery( + downstream_response_id=downstream_response_id, + after_circuit_cooldown=True, + ) + if terminal_event is not None: + yield terminal_event + break + if retried: + logger.info( + "HTTP bridge stream idle recovery retried after circuit cooldown " + "request_id=%s", + request_state.request_id, + ) + keepalive_count = 0 + keepalive_sent = False + yielded_any = False + continue + if not retried: + await self._record_http_bridge_retry_circuit_failure( + session, + detail="stream_idle_timeout", + ) + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + logger.info( + "HTTP bridge stream idle timeout request_id=%s keepalive_count=%s " + "max_keepalive_count=%s", + request_state.request_id, + keepalive_count, + max_keepalive_count, + ) + yield format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Upstream did not respond within the keepalive window", + response_id=downstream_response_id, + ), + ) + ) + break + elif response_started: + if PROMETHEUS_AVAILABLE and stream_idle_timeout_total is not None: + stream_idle_timeout_total.labels(surface="http_bridge").inc() + logger.info( + "HTTP bridge stream idle timeout request_id=%s keepalive_count=%s " + "max_keepalive_count=%s", + request_state.request_id, + keepalive_count, + max_keepalive_count, + ) + yield format_sse_event( + cast( + Mapping[str, JsonValue], + response_failed_event( + "stream_idle_timeout", + "Upstream did not respond within the keepalive window", + response_id=downstream_response_id, + ), + ) + ) + break + if ( + propagate_http_errors + and request_state.response_id is None + and not circuit_keepalive_waiting + ): continue keepalive_sent = True yielded_any = True + if PROMETHEUS_AVAILABLE and stream_keepalive_sent_total is not None: + stream_keepalive_sent_total.labels(surface="http_bridge").inc() if request_state.response_id or request_state.replay_downstream_response_id: yield format_sse_event( cast( @@ -2665,6 +3308,11 @@ async def _stream_http_bridge_session_events( if event_block is None: break keepalive_count = 0 + # A real upstream event means the stream is active again; do + # not carry the pre-response retry-circuit wake mode into the + # normal response-started idle timeout policy. + circuit_keepalive_waiting = False + circuit_keepalive_until = None block_payload = parse_sse_data_json(event_block) block_event_type = _event_type_from_payload(None, block_payload) if request_state.latency_first_token_ms is None: diff --git a/app/modules/proxy/_service/http_bridge/upstream_events.py b/app/modules/proxy/_service/http_bridge/upstream_events.py index a70367af9b..dee9812b34 100644 --- a/app/modules/proxy/_service/http_bridge/upstream_events.py +++ b/app/modules/proxy/_service/http_bridge/upstream_events.py @@ -46,6 +46,7 @@ ) from app.modules.proxy._service.http_bridge.helpers import ( _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL, + _http_bridge_durable_lease_ttl_seconds, _http_bridge_eventless_precreated_deadline, _http_bridge_request_budget_seconds, _http_bridge_request_counts_against_queue, @@ -57,6 +58,7 @@ _assign_websocket_response_id, _await_cancelled_task, _build_stream_incomplete_terminal_event_for_request, + _classify_upstream_close, _find_websocket_request_state_by_response_id, _http_error_status_from_payload, _is_missing_tool_output_error, @@ -179,6 +181,123 @@ ) logger = logging.getLogger("app.modules.proxy.service") + +_HTTP_BRIDGE_RECOVERY_SETTLEMENT_RETRY_DELAYS = ( + 0.25, + 0.5, + 1.0, + 2.0, + 4.0, + 8.0, + 15.0, + 30.0, + 60.0, + 120.0, +) +_HTTP_BRIDGE_RECOVERY_SETTLEMENT_LEASE_REFRESH_INTERVAL_SECONDS = 10.0 + + +async def _wait_for_http_bridge_recovery_settlement_retry( + service: Any, + *, + session_id: str, + owner_epoch: int, + api_key_id: str | None, + delay_seconds: float, +) -> None: + remaining = max(0.0, delay_seconds) + while remaining > 0: + await asyncio.sleep(min(remaining, _HTTP_BRIDGE_RECOVERY_SETTLEMENT_LEASE_REFRESH_INTERVAL_SECONDS)) + remaining -= _HTTP_BRIDGE_RECOVERY_SETTLEMENT_LEASE_REFRESH_INTERVAL_SECONDS + try: + await service._durable_bridge.renew_live_session( + session_id=session_id, + api_key_id=api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=owner_epoch, + lease_ttl_seconds=_http_bridge_durable_lease_ttl_seconds(), + ) + except Exception: + logger.debug("Failed to refresh HTTP bridge lease during settlement backoff", exc_info=True) + + +async def _retry_http_bridge_recovery_settlement( + service: Any, + session: Any, + *, + session_id: str, + api_key_id: str | None, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + response_id: str | None, + release_origin_lease: bool, +) -> None: + """Keep a response-observed journal row fenced until durable settlement succeeds.""" + + for delay_seconds in _HTTP_BRIDGE_RECOVERY_SETTLEMENT_RETRY_DELAYS: + await _wait_for_http_bridge_recovery_settlement_retry( + service, + session_id=session_id, + owner_epoch=owner_epoch, + api_key_id=api_key_id, + delay_seconds=delay_seconds, + ) + try: + marked = await service._durable_bridge.mark_recovery_attempt_replayed( + session_id=session_id, + api_key_id=api_key_id, + instance_id=instance_id, + owner_epoch=owner_epoch, + request_fingerprint=request_fingerprint, + response_id=response_id, + ) + if marked and (release_origin_lease or getattr(session, "closed", False)): + try: + await service._durable_bridge.release_live_session( + session_id=session_id, + instance_id=instance_id, + owner_epoch=owner_epoch, + draining=False, + ) + except Exception: + logger.debug("Failed to release HTTP bridge recovery origin lease", exc_info=True) + if marked: + return + logger.warning("HTTP bridge recovery settlement owner fence rejected; retrying") + except Exception: + continue + logger.error( + "HTTP bridge recovery settlement retry budget exhausted session_id=%s fingerprint=%s", + _hash_identifier(session_id), + _hash_identifier(request_fingerprint), + ) + + +def _schedule_http_bridge_recovery_settlement_retry( + service: Any, + session: Any, + **kwargs: Any, +) -> None: + task = asyncio.create_task( + _retry_http_bridge_recovery_settlement(service, session, **kwargs), + name=f"http-bridge-recovery-settlement-{_hash_identifier(kwargs['request_fingerprint'])}", + ) + setattr(task, "_http_bridge_recovery_session_id", kwargs["session_id"]) + service._background_cleanup_tasks.add(task) + + def _discard(done_task: asyncio.Task[Any]) -> None: + service._background_cleanup_tasks.discard(done_task) + try: + done_task.result() + except asyncio.CancelledError: + pass + except Exception: + logger.error("HTTP bridge recovery settlement retry failed", exc_info=True) + + task.add_done_callback(_discard) + + T = TypeVar("T") _TEXT_DELTA_EVENT_TYPES = frozenset({"response.output_text.delta", "response.refusal.delta"}) _MODEL_OUTPUT_EVENT_TYPES = frozenset( @@ -460,7 +579,12 @@ async def _http_bridge_receive_timeout_with_eventless_deadline( return receive_timeout -async def _cancel_http_bridge_reader_child(task: asyncio.Task[Any] | None, *, label: str) -> bool: +async def _cancel_http_bridge_reader_child( + task: asyncio.Task[Any] | None, + *, + label: str, + cleanup_tasks: set[asyncio.Task[None]] | None = None, +) -> bool: if task is None: return True if task.done(): @@ -472,7 +596,13 @@ async def _cancel_http_bridge_reader_child(task: asyncio.Task[Any] | None, *, la logger.debug("HTTP bridge reader child already failed during cleanup label=%s", label, exc_info=True) return True try: - return bool(await _await_cancelled_task(task, label=label)) + return bool( + await _await_cancelled_task( + task, + label=label, + cleanup_tasks=cleanup_tasks, + ) + ) except Exception: logger.debug("Failed to cancel HTTP bridge reader child label=%s", label, exc_info=True) return task.done() @@ -488,6 +618,9 @@ async def _fail_http_bridge_reader_and_maybe_retire( penalize_account: bool = True, retire_detail: str | None = None, force_retire: bool = False, + upstream_close_code: int | None = None, + response_events_seen: int | None = None, + transport_classification: str | None = None, ) -> bool: session.closed = True async with session.pending_lock: @@ -497,6 +630,51 @@ async def _fail_http_bridge_reader_and_maybe_retire( if _http_bridge_request_counts_against_queue(request_state) ) session.queued_request_count = max(0, session.queued_request_count - failed_pending_count) + observed_response_events = max( + (getattr(request_state, "response_event_count", 0) for request_state in session.pending_requests), + default=0, + ) + observed_close_code = ( + upstream_close_code if upstream_close_code is not None else session.last_upstream_close_code + ) + observed_response_events = ( + response_events_seen if response_events_seen is not None else observed_response_events + ) + close_classification = ( + _classify_upstream_close(observed_close_code, response_events_seen=observed_response_events) + if observed_close_code is not None + else None + ) + _log_http_bridge_event( + "reader_failure", + session.key, + account_id=session.account.id, + model=session.request_model, + pending_count=failed_pending_count, + detail=error_code, + error_message=_truncate_identifier(error_message), + upstream_close_code=observed_close_code, + response_events_seen=observed_response_events, + transport_classification=transport_classification + or ( + f"websocket_close_{close_classification}" + if close_classification is not None + else "websocket_transport_error" + ), + cache_key_family=session.key.affinity_kind, + model_class=_extract_model_class(session.request_model) if session.request_model else None, + ) + if force_retire and retire_detail: + _log_http_bridge_event( + retire_detail, + session.key, + account_id=session.account.id, + model=session.request_model, + pending_count=failed_pending_count, + detail=retire_detail, + cache_key_family=session.key.affinity_kind, + model_class=_extract_model_class(session.request_model) if session.request_model else None, + ) try: await self._fail_pending_websocket_requests( account=session.account, @@ -511,21 +689,47 @@ async def _fail_http_bridge_reader_and_maybe_retire( ) finally: if session.admission_waiter_count > 0 and not force_retire: + retry_circuit_detail = None + if close_classification == "clean": + retry_circuit_detail = "clean_close" + elif observed_response_events == 0: + retry_circuit_detail = next( + ( + detail + for detail in (retire_detail, error_code) + if detail in {"stream_incomplete", "stream_idle_timeout", "upstream_keepalive_timeout"} + ), + None, + ) + if failed_pending_count > 0 and retry_circuit_detail is not None: + await self._record_http_bridge_retry_circuit_failure( + session, + detail=retry_circuit_detail, + ) _log_http_bridge_event( "retire_deferred_for_admission_waiter", session.key, account_id=session.account.id, model=session.request_model, pending_count=session.admission_waiter_count, - detail=error_code, + detail=retire_detail or error_code, cache_key_family=session.key.affinity_kind, model_class=_extract_model_class(session.request_model) if session.request_model else None, ) else: - await self._retire_stale_pending_http_bridge_session( - session, - detail=retire_detail or error_code, - ) + if close_classification == "clean" and failed_pending_count > 0: + await self._retire_stale_pending_http_bridge_session( + session, + detail=error_code, + retry_circuit_detail="clean_close", + response_events_seen=observed_response_events, + ) + else: + await self._retire_stale_pending_http_bridge_session( + session, + detail=retire_detail or error_code, + response_events_seen=observed_response_events, + ) return force_retire or session.admission_waiter_count == 0 async def _relay_http_bridge_upstream_messages( @@ -591,6 +795,7 @@ async def _relay_http_bridge_upstream_messages( await _cancel_http_bridge_reader_child( wakeup_task, label="HTTP bridge reader wakeup wait", + cleanup_tasks=self._background_cleanup_tasks, ) wakeup_task = None @@ -631,16 +836,31 @@ async def _relay_http_bridge_upstream_messages( request_state.failure_detail_override = ( _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL ) - # Claim the session before cancelling receive so a - # gate waiter cannot reopen this ambiguous socket. - session.closed = True if receive_task is not None: receive_cancelled = await _cancel_http_bridge_reader_child( receive_task, label="HTTP bridge upstream receive after missing response.created", + cleanup_tasks=self._background_cleanup_tasks, ) - if receive_cancelled: - receive_task = None + if not receive_cancelled: + # Do not reconnect while the old receive + # task still owns the superseded socket. + # The ordinary timeout path takes the same + # fail-closed branch; retain the explicit + # account-neutral timeout classification + # rather than routing through the generic + # reader-crash account penalty path. + session.closed = True + await self._fail_http_bridge_reader_and_maybe_retire( + session, + error_code="upstream_request_timeout", + error_message=receive_timeout.error_message, + penalize_account=False, + retire_detail=_HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL, + force_retire=True, + ) + break + receive_task = None _record_http_bridge_stuck_retire( reason=_HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL, session=session, @@ -657,6 +877,17 @@ async def _relay_http_bridge_upstream_messages( _extract_model_class(session.request_model) if session.request_model else None ), ) + # A fresh, self-contained hard request can use the + # same bounded pre-created recovery as the idle + # timeout path. Keep the session open until the + # recovery routine claims the handoff; otherwise + # its retry gate would reject the request as + # already retired. Continuity-bound requests still + # fail closed in _retry_http_bridge_precreated_request. + retried = await self._retry_http_bridge_precreated_request(session) + if retried: + continue + session.closed = True await self._fail_http_bridge_reader_and_maybe_retire( session, error_code="upstream_request_timeout", @@ -671,6 +902,7 @@ async def _relay_http_bridge_upstream_messages( receive_cancelled = await _cancel_http_bridge_reader_child( receive_task, label="HTTP bridge upstream receive after timeout", + cleanup_tasks=self._background_cleanup_tasks, ) if not receive_cancelled: raise RuntimeError("HTTP bridge upstream receive did not cancel after timeout") @@ -702,21 +934,53 @@ async def _relay_http_bridge_upstream_messages( async with session.pending_lock: archive_request_state = session.pending_requests[0] if len(session.pending_requests) == 1 else None + response_events_seen = max( + (request_state.response_event_count for request_state in session.pending_requests), + default=0, + ) _archive_http_bridge_upstream_message(session, message, archive_request_state) + session.last_upstream_close_generation += 1 session.last_upstream_close_code = message.close_code retried = False - # A process-network receive failure follows a successful send; - # replay is not safe merely because output is not visible. + # A process-network receive failure does not prove that the + # upstream rejected response.create. Do not replay ordinary + # requests in that ambiguous case: the first request may + # still be executing and replay could duplicate work, billing, + # or tool side effects. Clean websocket closes remain eligible + # for the bounded pre-created retry path below. if message.error_code != "proxy_network_unavailable": retried = await self._retry_http_bridge_precreated_request(session) if retried: continue + close_classification = ( + _classify_upstream_close(message.close_code, response_events_seen=response_events_seen) + if message.close_code is not None + else None + ) async with session.lifecycle_lock: await self._fail_http_bridge_reader_and_maybe_retire( session, error_code=message.error_code or "stream_incomplete", error_message=_upstream_websocket_disconnect_message(message), - penalize_account=message.error_code != "proxy_network_unavailable", + upstream_close_code=message.close_code, + response_events_seen=response_events_seen, + transport_classification=( + f"websocket_close_{close_classification}" + if close_classification is not None + else "websocket_transport_error" + ), + penalize_account=( + message.error_code != "proxy_network_unavailable" + and message.error_code != "upstream_keepalive_timeout" + and not ( + message.kind == "close" + and _classify_upstream_close( + message.close_code, + response_events_seen=response_events_seen, + ) + == "clean" + ) + ), ) break except asyncio.CancelledError: @@ -729,7 +993,7 @@ async def _relay_http_bridge_upstream_messages( exc_info=True, ) error_code = exc.error_code if isinstance(exc, UpstreamWebSocketTransportError) else "stream_incomplete" - account_neutral = error_code == "proxy_network_unavailable" + account_neutral = error_code in {"proxy_network_unavailable", "upstream_keepalive_timeout"} async with session.lifecycle_lock: await self._fail_http_bridge_reader_and_maybe_retire( session, @@ -745,10 +1009,12 @@ async def _relay_http_bridge_upstream_messages( await _cancel_http_bridge_reader_child( wakeup_task, label="HTTP bridge reader wakeup wait", + cleanup_tasks=self._background_cleanup_tasks, ) await _cancel_http_bridge_reader_child( receive_task, label="HTTP bridge upstream receive", + cleanup_tasks=self._background_cleanup_tasks, ) if session.upstream is relay_upstream: session.closed = True @@ -1449,6 +1715,90 @@ async def _process_http_bridge_upstream_text( completed_usage = None completed_empty_prewarm = False + recovery_attempt_session_id = ( + matched_request_state.recovery_attempt_session_id + if matched_request_state is not None and matched_request_state.recovery_attempt_session_id is not None + else session.durable_session_id + ) + recovery_attempt_owner_epoch = ( + matched_request_state.recovery_attempt_owner_epoch + if matched_request_state is not None and matched_request_state.recovery_attempt_owner_epoch is not None + else session.durable_owner_epoch + ) + + if ( + isinstance(event_type, str) + and event_type.startswith("response.") + and matched_request_state is not None + and matched_request_state.recovery_attempt_fingerprint is not None + and recovery_attempt_session_id is not None + and recovery_attempt_owner_epoch is not None + and (event_type == "response.completed" or not matched_request_state.recovery_attempt_event_observed) + ): + settlement_marked = False + for settlement_attempt in range(3): + try: + marked = await self._durable_bridge.mark_recovery_attempt_replayed( + session_id=recovery_attempt_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + request_fingerprint=matched_request_state.recovery_attempt_fingerprint, + response_id=response_id, + ) + if marked: + settlement_marked = True + break + if settlement_attempt == 2: + _schedule_http_bridge_recovery_settlement_retry( + self, + session, + session_id=recovery_attempt_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + request_fingerprint=matched_request_state.recovery_attempt_fingerprint, + response_id=response_id, + release_origin_lease=( + recovery_attempt_session_id != session.durable_session_id + and event_type in {"response.completed", "response.failed"} + ), + ) + except Exception: + if settlement_attempt == 2: + logger.warning("Failed to settle HTTP bridge recovery attempt", exc_info=True) + _schedule_http_bridge_recovery_settlement_retry( + self, + session, + session_id=recovery_attempt_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + request_fingerprint=matched_request_state.recovery_attempt_fingerprint, + response_id=response_id, + release_origin_lease=( + recovery_attempt_session_id != session.durable_session_id + and event_type in {"response.completed", "response.failed"} + ), + ) + else: + await asyncio.sleep(0.05 * (settlement_attempt + 1)) + if ( + settlement_marked + and event_type in {"response.completed", "response.failed"} + and recovery_attempt_session_id != session.durable_session_id + ): + try: + await self._durable_bridge.release_live_session( + session_id=recovery_attempt_session_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + draining=False, + ) + except Exception: + logger.debug("Failed to release HTTP bridge recovery origin lease", exc_info=True) + matched_request_state.recovery_attempt_event_observed = True + if event_type == "response.completed" and terminal_request_state is not None and not completed_empty_prewarm: # Record the completed response id regardless of input shape so # subsequent turns (including ones that never populated @@ -1467,6 +1817,15 @@ async def _process_http_bridge_upstream_text( session.last_completed_input_count = terminal_request_state.input_item_count session.last_completed_input_prefix_fingerprint = terminal_request_state.input_full_fingerprint + if ( + event_type == "response.completed" + and terminal_request_state is not None + and not terminal_request_state.suppressed_duplicate_tool_call + and terminal_request_state.request_kind != "prewarm" + and not terminal_request_state.skip_request_log + ): + await self._clear_http_bridge_retry_circuit(session) + normalize_error_event = ( terminal_request_state is None or terminal_request_state.enforce_openai_sdk_contract ) and (matched_request_state is None or matched_request_state.enforce_openai_sdk_contract) @@ -1505,6 +1864,71 @@ async def _process_http_bridge_upstream_text( request_state=terminal_request_state or matched_request_state, ) + if ( + settlement_event_type in {"response.failed", "error"} + and matched_request_state is not None + and matched_request_state.recovery_attempt_fingerprint is not None + and recovery_attempt_session_id is not None + and recovery_attempt_owner_epoch is not None + and not matched_request_state.recovery_attempt_event_observed + ): + # An explicit deterministic failure is terminal evidence for the + # journaled request, not an ambiguous transport outcome. Consume + # the UNKNOWN row after normalizing top-level errors so a later + # identical retry cannot turn it into an account-neutral replay. + deterministic_settlement_marked = False + for settlement_attempt in range(3): + try: + marked = await self._durable_bridge.mark_recovery_attempt_replayed( + session_id=recovery_attempt_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + request_fingerprint=matched_request_state.recovery_attempt_fingerprint, + response_id=response_id, + ) + if marked: + deterministic_settlement_marked = True + break + if settlement_attempt == 2: + _schedule_http_bridge_recovery_settlement_retry( + self, + session, + session_id=recovery_attempt_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + request_fingerprint=matched_request_state.recovery_attempt_fingerprint, + response_id=response_id, + release_origin_lease=recovery_attempt_session_id != session.durable_session_id, + ) + except Exception: + if settlement_attempt == 2: + logger.warning("Failed to settle deterministic HTTP bridge recovery attempt", exc_info=True) + _schedule_http_bridge_recovery_settlement_retry( + self, + session, + session_id=recovery_attempt_session_id, + api_key_id=session.key.api_key_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + request_fingerprint=matched_request_state.recovery_attempt_fingerprint, + response_id=response_id, + release_origin_lease=recovery_attempt_session_id != session.durable_session_id, + ) + else: + await asyncio.sleep(0.05 * (settlement_attempt + 1)) + if deterministic_settlement_marked and recovery_attempt_session_id != session.durable_session_id: + try: + await self._durable_bridge.release_live_session( + session_id=recovery_attempt_session_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=recovery_attempt_owner_epoch, + draining=False, + ) + except Exception: + logger.debug("Failed to release HTTP bridge recovery origin lease", exc_info=True) + if event_type == "response.created" and release_create_gate and created_request_state is not None: await _release_websocket_response_create_gate(created_request_state, session.response_create_gate) diff --git a/app/modules/proxy/_service/request_log.py b/app/modules/proxy/_service/request_log.py index 654b1aec4d..67a3c9c2e7 100644 --- a/app/modules/proxy/_service/request_log.py +++ b/app/modules/proxy/_service/request_log.py @@ -22,11 +22,12 @@ "proxy-request-log-", "proxy-stream-api-key-settle-", "proxy-release_stream_api_key_reservation", + "http-bridge-recovery-settlement-", ) -def _is_persistence_task(task: asyncio.Task[None]) -> bool: - return task.get_name().startswith(_PERSISTENCE_TASK_NAME_PREFIXES) +def _is_persistence_task(task: asyncio.Task[None], prefixes: tuple[str, ...] | None = None) -> bool: + return task.get_name().startswith(prefixes or _PERSISTENCE_TASK_NAME_PREFIXES) _REQUEST_TRANSPORT_HTTP = "http" @@ -291,7 +292,11 @@ async def _write_request_log( model=model, ) - async def drain_persistence_tasks(self, timeout_seconds: float) -> bool: + async def drain_persistence_tasks( + self, + timeout_seconds: float, + task_name_prefixes: tuple[str, ...] | None = None, + ) -> bool: """Await detached request-log and settlement tasks, e.g. at shutdown. Persistence runs detached from the response path, so a graceful @@ -309,14 +314,14 @@ async def drain_persistence_tasks(self, timeout_seconds: float) -> bool: pending = { task for task in (proxy._request_log_tasks | proxy._background_cleanup_tasks) - if not task.done() and _is_persistence_task(task) + if not task.done() and _is_persistence_task(task, task_name_prefixes) } if not pending: # One scheduling tick so just-finished tasks' done callbacks # (which may enqueue follow-up tasks) run before we re-check. await asyncio.sleep(0) if not any( - _is_persistence_task(task) + _is_persistence_task(task, task_name_prefixes) for task in (proxy._request_log_tasks | proxy._background_cleanup_tasks) if not task.done() ): diff --git a/app/modules/proxy/_service/streaming/helpers.py b/app/modules/proxy/_service/streaming/helpers.py index 0750b6e077..a02f6a64e1 100644 --- a/app/modules/proxy/_service/streaming/helpers.py +++ b/app/modules/proxy/_service/streaming/helpers.py @@ -467,9 +467,13 @@ def _classify_upstream_close( close_code: int | None, *, response_events_seen: int, -) -> Literal["transient", "rejected"]: +) -> Literal["clean", "transient"]: if close_code == 1000 and response_events_seen == 0: - return "rejected" + # A clean websocket close before response.created does not prove that + # the request was invalid. The upstream can close a socket during a + # handoff, so the caller may safely recreate the socket and replay the + # still-unstarted request once. + return "clean" return "transient" diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index 141b2b5570..b0900087d3 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -77,6 +77,7 @@ } ) _ACCOUNT_SELECTION_RECOVERY_MIN_SLEEP_SECONDS = 1.0 +_HARD_AFFINITY_RECOVERY_SLEEP_SECONDS = 2.0 _ACCOUNT_SELECTION_RECOVERY_DEFAULT_SLEEP_SECONDS = 30.0 _ACCOUNT_SELECTION_RECOVERY_MAX_SLEEP_SECONDS = 300.0 _ACCOUNT_SELECTION_RECOVERY_HEARTBEAT_SECONDS = 10.0 @@ -321,6 +322,13 @@ def _account_selection_recovery_sleep_seconds_from_message( if "hit your spend cap set by the owner of your workspace" in lowered: return _ACCOUNT_SELECTION_RECOVERY_DEFAULT_SLEEP_SECONDS + # A hard affinity row is an ownership constraint, so it must not spill to + # another account. Capacity/health transitions can nevertheless make the + # owner briefly unavailable; give that owner a short recovery window before + # surfacing a 503 rather than rebinding the logical session. + if error_code == "hard_affinity_saturated": + return _HARD_AFFINITY_RECOVERY_SLEEP_SECONDS + if error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES: return _ACCOUNT_SELECTION_RECOVERY_DEFAULT_SLEEP_SECONDS @@ -786,6 +794,13 @@ class _WebSocketRequestState: request_usage_budget: ApiKeyRequestUsageBudget | None = None request_text: str | None = None replay_count: int = 0 + # Counts only the one extra replay permitted after the initial recovery + # replay when the replacement upstream socket also closes cleanly before + # producing any response event. + clean_close_replay_count: int = 0 + clean_close_retry_in_progress: bool = False + clean_close_retry_result: bool | None = None + clean_close_retry_close_generation: int | None = None auth_replay_count: int = 0 auth_replay_counts_by_account: dict[str, int] = field(default_factory=dict) force_refresh_account_id: str | None = None @@ -799,6 +814,9 @@ class _WebSocketRequestState: skip_request_log: bool = False previous_response_id: str | None = None session_id: str | None = None + # Session headers provide locality, but only a previous response or + # explicit turn-state header guarantees continuity for stale recovery. + hard_continuity_anchor: bool = False proxy_injected_previous_response_id: bool = False expose_stale_previous_response_classifier: bool = False fresh_upstream_request_text: str | None = None @@ -811,6 +829,19 @@ class _WebSocketRequestState: # on, and dropping the anchor there would silently turn a continuation into # a context-free fresh turn. fresh_upstream_request_is_retry_safe: bool = False + # Stable fingerprint used by the durable recovery-attempt journal. It is + # populated only for a proof-gated fresh replay candidate. + recovery_attempt_fingerprint: str | None = None + recovery_attempt_session_id: str | None = None + recovery_attempt_owner_epoch: int | None = None + # True when recovery already atomically claimed the journal row as + # REPLAYED. Claimed replays must not be re-journaled at dispatch. + recovery_attempt_claimed: bool = False + # Set once the upstream send is attempted, including an ambiguous send + # failure. Pre-dispatch admission/setup failures may safely roll back a + # claimed recovery journal; an attempted send must remain consumed. + recovery_attempt_dispatched: bool = False + recovery_attempt_event_observed: bool = False # Responses-Lite model advertised by ``fresh_upstream_request_text``. A # fresh replay built from a trusted marker-only frame has the reserved # marker stripped, so swapping to the fresh body must also swap this onto @@ -832,6 +863,7 @@ class _WebSocketRequestState: upstream_error_code_override: str | None = None error_http_status_override: int | None = None response_event_count: int = 0 + last_upstream_activity_at: float | None = None upstream_model_output_seen: bool = False previous_response_not_found_rewritten: bool = False previous_response_owner_lookup_source: str | None = None @@ -946,7 +978,13 @@ class _HTTPBridgeSession: durable_owner_epoch: int | None = None upstream_reader: asyncio.Task[None] | None = None last_upstream_close_code: int | None = None + last_upstream_close_generation: int = 0 closed: bool = False + # Set while a reader handoff is replacing the socket. Idle pruning must + # retain the registered session during this short transition even though + # ``closed`` is fail-closed for normal request reuse. + handoff_in_progress: bool = False + handoff_future: asyncio.Future["_HTTPBridgeSession"] | None = None account_lease: AccountLease | None = None upstream_close_attempted: bool = False seen_tool_call_keys: dict[ToolCallDedupeKey, None] = field(default_factory=dict) @@ -957,6 +995,19 @@ class _HTTPBridgeSession: upstream_proxy_fail_closed_reason: str | None = None +def _complete_http_bridge_handoff( + session: _HTTPBridgeSession, + inflight_sessions: dict[_HTTPBridgeSessionKey, asyncio.Future[_HTTPBridgeSession]], +) -> None: + session.handoff_in_progress = False + future = session.handoff_future + session.handoff_future = None + if future is not None and not future.done(): + future.set_result(session) + if inflight_sessions.get(session.key) is future: + inflight_sessions.pop(session.key, None) + + def _http_bridge_session_supports_service_tier( session: _HTTPBridgeSession, *, @@ -1150,15 +1201,25 @@ def _clear_websocket_deferred_reasoning_downstream_texts(request_state: _WebSock def _record_response_event(request_state: _WebSocketRequestState | None, event_type: str | None) -> None: if request_state is None or event_type is None or not event_type.startswith("response."): return + request_state.last_upstream_activity_at = time.monotonic() if event_type in {"response.failed", "response.incomplete"}: return request_state.response_event_count += 1 -def _websocket_request_can_replay_before_visible_output(request_state: _WebSocketRequestState) -> bool: +def _websocket_request_can_replay_before_visible_output( + request_state: _WebSocketRequestState, + *, + allow_clean_close_retry: bool = False, +) -> bool: if not request_state.request_text: return False - if request_state.replay_count >= 1: + if request_state.replay_count >= 1 and not ( + allow_clean_close_retry + and request_state.replay_count == 1 + and request_state.response_event_count == 0 + and request_state.clean_close_replay_count == 0 + ): return False sequenced_created_only_prewarm = ( request_state.generate_false_prewarm diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index 63d5a5687d..ba65640c73 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -854,6 +854,7 @@ async def retire_current_upstream() -> None: await _facade()._await_cancelled_task( upstream_reader, label="proxy websocket upstream reader", + cleanup_tasks=proxy._background_cleanup_tasks, ) upstream_reader = None upstream_control = None @@ -1723,7 +1724,11 @@ async def retire_current_upstream() -> None: downstream_activity, ) if upstream_reader is not None: - await _facade()._await_cancelled_task(upstream_reader, label="proxy websocket upstream reader") + await _facade()._await_cancelled_task( + upstream_reader, + label="proxy websocket upstream reader", + cleanup_tasks=proxy._background_cleanup_tasks, + ) upstream_reader = None upstream_control = None if upstream is not None: @@ -1753,7 +1758,9 @@ async def retire_current_upstream() -> None: replay_request_state = replay_candidate if upstream_reader is not None: await _facade()._await_cancelled_task( - upstream_reader, label="proxy websocket upstream reader" + upstream_reader, + label="proxy websocket upstream reader", + cleanup_tasks=proxy._background_cleanup_tasks, ) upstream_reader = None upstream_control = None @@ -1790,7 +1797,11 @@ async def retire_current_upstream() -> None: downstream_activity, ) if upstream_reader is not None: - await _facade()._await_cancelled_task(upstream_reader, label="proxy websocket upstream reader") + await _facade()._await_cancelled_task( + upstream_reader, + label="proxy websocket upstream reader", + cleanup_tasks=proxy._background_cleanup_tasks, + ) upstream_reader = None upstream_control = None if upstream is not None: @@ -1806,7 +1817,11 @@ async def retire_current_upstream() -> None: continue finally: if upstream_reader is not None: - await _facade()._await_cancelled_task(upstream_reader, label="proxy websocket upstream reader") + await _facade()._await_cancelled_task( + upstream_reader, + label="proxy websocket upstream reader", + cleanup_tasks=proxy._background_cleanup_tasks, + ) if upstream is not None: try: await upstream.close() @@ -3782,6 +3797,14 @@ async def _relay_upstream_websocket_messages( penalize_account=message.error_code != "proxy_network_unavailable", suppress_sequenced_downstream_errors=sequenced_downstream_replay_refused, ) + # A terminal receive can race the outer session loop's + # cleanup (especially when the downstream closes as soon as + # it receives the failure event). Close here as well so the + # transport is retired before the reader task exits. + try: + await upstream.close() + except Exception: + _facade().logger.debug("Failed to close upstream websocket after terminal receive", exc_info=True) if sequenced_downstream_replay_refused: await _close_downstream_after_sequenced_replay_refusal( websocket, diff --git a/app/modules/proxy/_service/websocket/protocol.py b/app/modules/proxy/_service/websocket/protocol.py index 521c5928a7..9eaed0669e 100644 --- a/app/modules/proxy/_service/websocket/protocol.py +++ b/app/modules/proxy/_service/websocket/protocol.py @@ -5,6 +5,7 @@ class _WebSocketServiceProtocol(Protocol): _capability_router: Any + _background_cleanup_tasks: Any _acquire_account_response_create_lease_or_overload: Any _acquire_request_state_response_create_admission: Any _cancel_request_state_api_key_reservation_heartbeat: Any diff --git a/app/modules/proxy/api.py b/app/modules/proxy/api.py index 296979a67d..abe50df1b2 100644 --- a/app/modules/proxy/api.py +++ b/app/modules/proxy/api.py @@ -3,6 +3,7 @@ import asyncio import json import logging +import math import time from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping from contextlib import asynccontextmanager @@ -82,7 +83,11 @@ ProxyRateLimitError, ProxyUpstreamError, ) -from app.core.metrics.prometheus import PROMETHEUS_AVAILABLE, bridge_public_contract_error_total +from app.core.metrics.prometheus import ( + PROMETHEUS_AVAILABLE, + bridge_public_contract_error_total, + stream_keepalive_sent_total, +) from app.core.middleware.multipart_content_encoding import raise_for_unsupported_multipart_content_encoding from app.core.multipart import ( IMAGE_EDITS_MULTIPART_POLICY, @@ -720,6 +725,11 @@ async def _thread_goal_proxy( } ) +# A hard HTTP-bridge circuit is opened only after an ambiguous upstream turn +# failure. The caller must not immediately replay that turn, but it should +# also not have to guess when a new attempt is safe. Advertise a short, +# bounded retry interval on the one-shot 503 response. + def _codex_control_downstream_headers(headers: Mapping[str, str]) -> dict[str, str]: return {key: value for key, value in headers.items() if key.lower() in _CODEX_CONTROL_RESPONSE_HEADERS} @@ -3871,6 +3881,7 @@ async def v1_chat_completions( inject_sse_keepalives( chat_stream, get_settings().sse_keepalive_interval_seconds, + on_keepalive=lambda: _record_stream_keepalive("chat_completions"), ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", **rate_limit_headers}, @@ -5036,6 +5047,7 @@ async def _stream_responses( stream, get_settings().sse_keepalive_interval_seconds, keepalive_frame=keepalive_frame, + on_keepalive=lambda: _record_stream_keepalive("responses"), ), media_type="text/event-stream", headers={ @@ -6089,6 +6101,11 @@ async def _prepend_initial_sse_heartbeat( yield line +def _record_stream_keepalive(surface: str) -> None: + if PROMETHEUS_AVAILABLE and stream_keepalive_sent_total is not None: + stream_keepalive_sent_total.labels(surface=surface).inc() + + 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 @@ -6112,7 +6129,15 @@ async def _stream_response_error_events( envelope = _parse_error_envelope(exc.payload) _, envelope = _mask_previous_response_not_found_error(envelope, default_status=exc.status_code) error = envelope.error - yield format_sse_event( + retry_hint = "" + if exc.retry_after_seconds is not None and exc.retry_after_seconds > 0: + # Preserve the HTTP Retry-After signal when a streaming response + # has already started and the exception must be represented as an + # SSE event. The SSE retry field is milliseconds, while the + # exception stores seconds. Clients that do not implement the + # directive safely ignore the extra comment line. + retry_hint = f"retry: {max(1, math.ceil(exc.retry_after_seconds * 1000))}\n" + yield retry_hint + format_sse_event( response_failed_event( error.code if error and error.code else "upstream_error", error.message if error and error.message else "Upstream error", @@ -6131,11 +6156,14 @@ def _stream_startup_error_response( if isinstance(error, ProxyResponseError): envelope = _parse_error_envelope(error.payload) status_code, envelope = _mask_previous_response_not_found_error(envelope, default_status=error.status_code) + startup_headers = dict(headers) + if error.retry_after_seconds is not None and error.retry_after_seconds > 0: + startup_headers.setdefault("Retry-After", str(error.retry_after_seconds)) return _logged_error_json_response( request, status_code, envelope.model_dump(mode="json", exclude_none=True), - headers=headers, + headers=startup_headers, ) status_code, envelope = _mask_previous_response_not_found_error(error) return _logged_error_json_response( diff --git a/app/modules/proxy/durable_bridge_coordinator.py b/app/modules/proxy/durable_bridge_coordinator.py index 75e397a7d2..a6c125b70b 100644 --- a/app/modules/proxy/durable_bridge_coordinator.py +++ b/app/modules/proxy/durable_bridge_coordinator.py @@ -15,7 +15,9 @@ from app.modules.proxy.durable_bridge_repository import ( DurableBridgeAliasRegistration, DurableBridgeAliasRegistrationReceipt, + DurableBridgeRecoveryAttemptSnapshot, DurableBridgeRepository, + DurableBridgeRetryCircuitSnapshot, DurableBridgeSessionSnapshot, durable_bridge_api_key_scope, ) @@ -194,6 +196,92 @@ async def lookup_sessions(self, *, session_ids: Sequence[str]) -> list[DurableBr snapshots = await DurableBridgeRepository(session).get_sessions_by_ids(session_ids) return [_to_lookup(snapshot) for snapshot in snapshots] + async def lookup_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_id: str | None, + ) -> DurableBridgeRetryCircuitSnapshot | None: + async with self._session() as session: + return await DurableBridgeRepository(session).get_retry_circuit( + session_key_kind=session_key_kind, + session_key_value=session_key_value, + api_key_scope=durable_bridge_api_key_scope(api_key_id), + ) + + async def persist_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_id: str | None, + consecutive_failures: int, + cooldown_until_epoch: float, + last_detail: str | None, + updated_at_epoch: float, + base_updated_at_epoch: float = 0.0, + failure_threshold: int = 1, + conflict_cooldown_until_epoch: float | None = None, + base_backoff_seconds: float = 60.0, + max_backoff_seconds: float = 600.0, + clean_close_max_backoff_seconds: float = 30.0, + ) -> DurableBridgeRetryCircuitSnapshot | None: + async with self._session() as session: + repository = DurableBridgeRepository(session) + await repository.upsert_retry_circuit( + session_key_kind=session_key_kind, + session_key_value=session_key_value, + api_key_scope=durable_bridge_api_key_scope(api_key_id), + consecutive_failures=consecutive_failures, + cooldown_until_epoch=cooldown_until_epoch, + last_detail=last_detail, + updated_at_epoch=updated_at_epoch, + base_updated_at_epoch=base_updated_at_epoch, + failure_threshold=failure_threshold, + conflict_cooldown_until_epoch=conflict_cooldown_until_epoch, + base_backoff_seconds=base_backoff_seconds, + max_backoff_seconds=max_backoff_seconds, + clean_close_max_backoff_seconds=clean_close_max_backoff_seconds, + ) + return await repository.get_retry_circuit( + session_key_kind=session_key_kind, + session_key_value=session_key_value, + api_key_scope=durable_bridge_api_key_scope(api_key_id), + ) + + async def clear_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_id: str | None, + expected_updated_at_epoch: float | None = None, + ) -> None: + async with self._session() as session: + await DurableBridgeRepository(session).delete_retry_circuit( + session_key_kind=session_key_kind, + session_key_value=session_key_value, + api_key_scope=durable_bridge_api_key_scope(api_key_id), + expected_updated_at_epoch=expected_updated_at_epoch, + ) + + async def purge_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_id: str | None, + expected_updated_at_epoch: float | None = None, + ) -> None: + async with self._session() as session: + await DurableBridgeRepository(session).purge_retry_circuit( + session_key_kind=session_key_kind, + session_key_value=session_key_value, + api_key_scope=durable_bridge_api_key_scope(api_key_id), + expected_updated_at_epoch=expected_updated_at_epoch, + ) + async def claim_live_session( self, *, @@ -261,6 +349,26 @@ async def renew_live_session( return None return _to_lookup(snapshot) + async def rebind_session_account( + self, + *, + session_id: str, + api_key_id: str | None, + instance_id: str, + owner_epoch: int, + account_id: str, + clear_continuity: bool = False, + ) -> bool: + del api_key_id + async with self._session() as session: + return await DurableBridgeRepository(session).rebind_session_account( + session_id=session_id, + instance_id=instance_id, + owner_epoch=owner_epoch, + account_id=account_id, + clear_continuity=clear_continuity, + ) + async def release_live_session( self, *, @@ -280,6 +388,82 @@ async def release_live_session( return None return _to_lookup(snapshot) + async def record_recovery_attempt( + self, + *, + session_id: str, + api_key_id: str | None, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + request_id: str, + account_id: str | None, + model: str | None, + replay_safe: bool, + ) -> DurableBridgeRecoveryAttemptSnapshot | None: + del api_key_id + async with self._session() as session: + return await DurableBridgeRepository(session).record_recovery_attempt( + session_id=session_id, + instance_id=instance_id, + owner_epoch=owner_epoch, + request_fingerprint=request_fingerprint, + request_id=request_id, + account_id=account_id, + model=model, + replay_safe=replay_safe, + ) + + async def lookup_recovery_attempt( + self, + *, + session_id: str, + request_fingerprint: str, + ) -> DurableBridgeRecoveryAttemptSnapshot | None: + async with self._session() as session: + return await DurableBridgeRepository(session).lookup_recovery_attempt( + session_id=session_id, + request_fingerprint=request_fingerprint, + ) + + async def mark_recovery_attempt_replayed( + self, + *, + session_id: str, + api_key_id: str | None, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + response_id: str | None = None, + ) -> bool: + del api_key_id + async with self._session() as session: + return await DurableBridgeRepository(session).mark_recovery_attempt_replayed( + session_id=session_id, + instance_id=instance_id, + owner_epoch=owner_epoch, + request_fingerprint=request_fingerprint, + response_id=response_id, + ) + + async def rollback_recovery_attempt_replayed( + self, + *, + session_id: str, + api_key_id: str | None, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + ) -> bool: + del api_key_id + async with self._session() as session: + return await DurableBridgeRepository(session).rollback_recovery_attempt_replayed( + session_id=session_id, + instance_id=instance_id, + owner_epoch=owner_epoch, + request_fingerprint=request_fingerprint, + ) + async def mark_instance_draining(self, *, instance_id: str) -> int: async with self._session() as session: return await DurableBridgeRepository(session).mark_owner_draining(instance_id=instance_id) diff --git a/app/modules/proxy/durable_bridge_repository.py b/app/modules/proxy/durable_bridge_repository.py index 7b82e7f8a3..eefa4d1a45 100644 --- a/app/modules/proxy/durable_bridge_repository.py +++ b/app/modules/proxy/durable_bridge_repository.py @@ -1,20 +1,29 @@ from __future__ import annotations import json +import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta from enum import StrEnum from hashlib import sha256 +from typing import Any -from sqlalchemy import Row, and_, case, delete, or_, select, text, update +from sqlalchemy import Row, and_, case, delete, func, or_, select, text, update from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from app.core.utils.time import to_utc_naive, utcnow -from app.db.models import HttpBridgeSessionAlias, HttpBridgeSessionRecord, HttpBridgeSessionState +from app.db.models import ( + HttpBridgeRecoveryAttemptRecord, + HttpBridgeRecoveryAttemptState, + HttpBridgeRetryCircuit, + HttpBridgeSessionAlias, + HttpBridgeSessionRecord, + HttpBridgeSessionState, +) from app.db.session import sqlite_writer_section from app.modules.proxy.continuity import ( HTTP_BRIDGE_ACCOUNT_NEUTRAL_REPLAY_KEY_PREFIX, @@ -27,7 +36,10 @@ REQUIRED_DURABLE_BRIDGE_TABLES = ( "http_bridge_sessions", "http_bridge_session_aliases", + "http_bridge_retry_circuits", + "http_bridge_recovery_attempts", ) +DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS = 3600.0 _PURGE_CLOSED_BATCH_SIZE = 500 _SESSION_ID_LOOKUP_CHUNK_SIZE = 500 @@ -120,6 +132,29 @@ class DurableBridgeSessionSnapshot: latest_pending_tool_calls: dict[str, str] | None = None +@dataclass(frozen=True, slots=True) +class DurableBridgeRetryCircuitSnapshot: + session_key_kind: str + session_key_hash: str + api_key_scope: str + consecutive_failures: int + cooldown_until_epoch: float + last_detail: str | None + updated_at_epoch: float + + +@dataclass(frozen=True, slots=True) +class DurableBridgeRecoveryAttemptSnapshot: + session_id: str + request_fingerprint: str + request_id: str + account_id: str | None + model: str | None + replay_safe: bool + state: HttpBridgeRecoveryAttemptState + response_id: str | None + + class DurableBridgeRepository: def __init__(self, session: AsyncSession) -> None: self._session = session @@ -144,6 +179,286 @@ async def get_session( row = result.scalar_one_or_none() return _to_snapshot(row) + async def get_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_scope: str, + ) -> DurableBridgeRetryCircuitSnapshot | None: + result = await self._session.execute( + select(HttpBridgeRetryCircuit).where( + HttpBridgeRetryCircuit.session_key_kind == session_key_kind, + HttpBridgeRetryCircuit.session_key_hash == durable_bridge_hash(session_key_value), + HttpBridgeRetryCircuit.api_key_scope == api_key_scope, + ) + ) + return _to_retry_circuit_snapshot(result.scalar_one_or_none()) + + async def upsert_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_scope: str, + consecutive_failures: int, + cooldown_until_epoch: float, + last_detail: str | None, + updated_at_epoch: float, + base_updated_at_epoch: float = 0.0, + failure_threshold: int = 1, + conflict_cooldown_until_epoch: float | None = None, + base_backoff_seconds: float = 60.0, + max_backoff_seconds: float = 600.0, + clean_close_max_backoff_seconds: float = 30.0, + ) -> None: + values = { + "session_key_kind": session_key_kind, + "session_key_hash": durable_bridge_hash(session_key_value), + "api_key_scope": api_key_scope, + "consecutive_failures": consecutive_failures, + "cooldown_until_epoch": cooldown_until_epoch, + "last_detail": last_detail, + "updated_at_epoch": updated_at_epoch, + } + threshold = max(1, failure_threshold) + cooldown_floor = ( + max(0.0, conflict_cooldown_until_epoch) + if conflict_cooldown_until_epoch is not None + else max(0.0, cooldown_until_epoch) + ) + # A reset starts a new failure lineage. Never carry the incoming + # cooldown into that fresh lineage, even when the threshold is one. + reset_failure_cooldown = 0.0 + base_backoff = max(0.001, base_backoff_seconds) + max_backoff = max(base_backoff, max_backoff_seconds) + clean_close_max_backoff = max(0.001, clean_close_max_backoff_seconds) + + def cooldown_for_failure_count(failure_count: Any, last_detail: Any) -> Any: + regular_cooldown = case( + (failure_count < threshold, 0.0), + (failure_count == threshold, base_backoff), + (failure_count == threshold + 1, min(max_backoff, base_backoff * 2.0)), + (failure_count == threshold + 2, min(max_backoff, base_backoff * 4.0)), + else_=max_backoff, + ) + clean_cooldown = case( + (failure_count < threshold, 0.0), + else_=clean_close_max_backoff, + ) + return case((last_detail == "clean_close", clean_cooldown), else_=regular_cooldown) + + dialect = self._session.get_bind().dialect.name + if dialect == "postgresql": + insert_statement = pg_insert(HttpBridgeRetryCircuit).values(**values) + excluded = insert_statement.excluded + reset_lineage = and_( + HttpBridgeRetryCircuit.consecutive_failures == 0, + HttpBridgeRetryCircuit.cooldown_until_epoch <= 0, + HttpBridgeRetryCircuit.last_detail.is_(None), + HttpBridgeRetryCircuit.updated_at_epoch > base_updated_at_epoch, + ) + # ``updated_at_epoch`` is an observation timestamp, not a + # concurrency version. Treat an unchanged loaded row as a CAS + # match, even when a replica's wall clock lags it. The failure + # count guard still rejects an older snapshot that was loaded from + # the same row after a newer failure had already been merged. + failure_from_loaded_row = and_( + HttpBridgeRetryCircuit.updated_at_epoch == base_updated_at_epoch, + excluded.consecutive_failures >= HttpBridgeRetryCircuit.consecutive_failures, + ) + failure_is_newer_than_base = or_( + excluded.updated_at_epoch > base_updated_at_epoch, + failure_from_loaded_row, + ) + conflict_failures = case( + (reset_lineage, 1), + ( + failure_is_newer_than_base, + func.greatest( + HttpBridgeRetryCircuit.consecutive_failures + 1, + excluded.consecutive_failures, + ), + ), + else_=HttpBridgeRetryCircuit.consecutive_failures, + ) + merged_updated_at = func.greatest( + HttpBridgeRetryCircuit.updated_at_epoch, + excluded.updated_at_epoch, + ) + merged_cooldown = case( + (reset_lineage, reset_failure_cooldown), + ( + conflict_failures >= threshold, + func.greatest( + cooldown_floor, + merged_updated_at + cooldown_for_failure_count(conflict_failures, excluded.last_detail), + ), + ), + else_=0.0, + ) + statement = insert_statement.on_conflict_do_update( + index_elements=[ + HttpBridgeRetryCircuit.session_key_kind, + HttpBridgeRetryCircuit.session_key_hash, + HttpBridgeRetryCircuit.api_key_scope, + ], + set_={ + "consecutive_failures": conflict_failures, + "cooldown_until_epoch": case( + (reset_lineage, reset_failure_cooldown), + else_=func.greatest( + HttpBridgeRetryCircuit.cooldown_until_epoch, + excluded.cooldown_until_epoch, + merged_cooldown, + ), + ), + "last_detail": case( + (reset_lineage, excluded.last_detail), + ( + excluded.updated_at_epoch >= HttpBridgeRetryCircuit.updated_at_epoch, + excluded.last_detail, + ), + else_=HttpBridgeRetryCircuit.last_detail, + ), + "updated_at_epoch": case( + (reset_lineage, excluded.updated_at_epoch), + else_=func.greatest( + HttpBridgeRetryCircuit.updated_at_epoch, + excluded.updated_at_epoch, + ), + ), + }, + ) + elif dialect == "sqlite": + insert_statement = sqlite_insert(HttpBridgeRetryCircuit).values(**values) + excluded = insert_statement.excluded + reset_lineage = and_( + HttpBridgeRetryCircuit.consecutive_failures == 0, + HttpBridgeRetryCircuit.cooldown_until_epoch <= 0, + HttpBridgeRetryCircuit.last_detail.is_(None), + HttpBridgeRetryCircuit.updated_at_epoch > base_updated_at_epoch, + ) + failure_from_loaded_row = and_( + HttpBridgeRetryCircuit.updated_at_epoch == base_updated_at_epoch, + excluded.consecutive_failures >= HttpBridgeRetryCircuit.consecutive_failures, + ) + failure_is_newer_than_base = or_( + excluded.updated_at_epoch > base_updated_at_epoch, + failure_from_loaded_row, + ) + conflict_failures = case( + (reset_lineage, 1), + ( + failure_is_newer_than_base, + func.max( + HttpBridgeRetryCircuit.consecutive_failures + 1, + excluded.consecutive_failures, + ), + ), + else_=HttpBridgeRetryCircuit.consecutive_failures, + ) + merged_updated_at = func.max( + HttpBridgeRetryCircuit.updated_at_epoch, + excluded.updated_at_epoch, + ) + merged_cooldown = case( + (reset_lineage, reset_failure_cooldown), + ( + conflict_failures >= threshold, + func.max( + cooldown_floor, + merged_updated_at + cooldown_for_failure_count(conflict_failures, excluded.last_detail), + ), + ), + else_=0.0, + ) + statement = insert_statement.on_conflict_do_update( + index_elements=[ + HttpBridgeRetryCircuit.session_key_kind, + HttpBridgeRetryCircuit.session_key_hash, + HttpBridgeRetryCircuit.api_key_scope, + ], + set_={ + "consecutive_failures": conflict_failures, + "cooldown_until_epoch": case( + (reset_lineage, reset_failure_cooldown), + else_=func.max( + HttpBridgeRetryCircuit.cooldown_until_epoch, + excluded.cooldown_until_epoch, + merged_cooldown, + ), + ), + "last_detail": case( + (reset_lineage, excluded.last_detail), + ( + excluded.updated_at_epoch >= HttpBridgeRetryCircuit.updated_at_epoch, + excluded.last_detail, + ), + else_=HttpBridgeRetryCircuit.last_detail, + ), + "updated_at_epoch": case( + (reset_lineage, excluded.updated_at_epoch), + else_=func.max( + HttpBridgeRetryCircuit.updated_at_epoch, + excluded.updated_at_epoch, + ), + ), + }, + ) + else: + raise RuntimeError(f"DurableBridgeRepository retry circuit upsert unsupported for dialect={dialect!r}") + async with sqlite_writer_section(): + await self._session.execute(statement) + await self._session.commit() + + async def delete_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_scope: str, + expected_updated_at_epoch: float | None = None, + ) -> None: + conditions = [ + HttpBridgeRetryCircuit.session_key_kind == session_key_kind, + HttpBridgeRetryCircuit.session_key_hash == durable_bridge_hash(session_key_value), + HttpBridgeRetryCircuit.api_key_scope == api_key_scope, + ] + if expected_updated_at_epoch is not None: + conditions.append(HttpBridgeRetryCircuit.updated_at_epoch == expected_updated_at_epoch) + async with sqlite_writer_section(): + await self._session.execute( + update(HttpBridgeRetryCircuit) + .where(*conditions) + .values( + consecutive_failures=0, + cooldown_until_epoch=0.0, + last_detail=None, + updated_at_epoch=time.time(), + ) + ) + await self._session.commit() + + async def purge_retry_circuit( + self, + *, + session_key_kind: str, + session_key_value: str, + api_key_scope: str, + expected_updated_at_epoch: float | None = None, + ) -> None: + conditions = [ + HttpBridgeRetryCircuit.session_key_kind == session_key_kind, + HttpBridgeRetryCircuit.session_key_hash == durable_bridge_hash(session_key_value), + HttpBridgeRetryCircuit.api_key_scope == api_key_scope, + ] + if expected_updated_at_epoch is not None: + conditions.append(HttpBridgeRetryCircuit.updated_at_epoch == expected_updated_at_epoch) + async with sqlite_writer_section(): + await self._session.execute(delete(HttpBridgeRetryCircuit).where(*conditions)) + await self._session.commit() + async def get_session_by_id(self, session_id: str) -> DurableBridgeSessionSnapshot | None: row = await self._session.get(HttpBridgeSessionRecord, session_id) return _to_snapshot(row) @@ -378,6 +693,41 @@ async def renew_session( values=values, ) + async def rebind_session_account( + self, + *, + session_id: str, + instance_id: str, + owner_epoch: int, + account_id: str, + clear_continuity: bool = False, + ) -> bool: + """Persist a replacement account only while this worker owns the lease.""" + + async with sqlite_writer_section(): + values: dict[str, object] = {"account_id": account_id} + if clear_continuity: + values.update( + latest_turn_state=None, + latest_response_id=None, + latest_input_item_count=None, + latest_input_full_fingerprint=None, + latest_pending_tool_calls_json=None, + ) + result = await self._session.execute( + update(HttpBridgeSessionRecord) + .where( + HttpBridgeSessionRecord.id == session_id, + HttpBridgeSessionRecord.owner_instance_id == instance_id, + HttpBridgeSessionRecord.owner_epoch == owner_epoch, + ) + .values(**values) + ) + if clear_continuity and bool(getattr(result, "rowcount", 0)): + await self._clear_aliases_for_session(session_id) + await self._session.commit() + return bool(getattr(result, "rowcount", 0)) + async def release_session( self, *, @@ -406,6 +756,215 @@ async def release_session( values=values, ) + async def record_recovery_attempt( + self, + *, + session_id: str, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + request_id: str, + account_id: str | None, + model: str | None, + replay_safe: bool, + ) -> DurableBridgeRecoveryAttemptSnapshot | None: + """Record a safe request before dispatch so an ambiguous outcome is recoverable.""" + async with sqlite_writer_section(): + # Lock the owner row through the journal write so a takeover + # cannot advance the epoch after this check but before dispatch. + owner_exists = await self._session.scalar( + select(HttpBridgeSessionRecord.id) + .where( + HttpBridgeSessionRecord.id == session_id, + HttpBridgeSessionRecord.owner_instance_id == instance_id, + HttpBridgeSessionRecord.owner_epoch == owner_epoch, + ) + .with_for_update() + ) + if owner_exists is None: + await self._session.rollback() + return None + attempt = await self._session.scalar( + select(HttpBridgeRecoveryAttemptRecord) + .where(HttpBridgeRecoveryAttemptRecord.session_id == session_id) + .where(HttpBridgeRecoveryAttemptRecord.request_fingerprint == request_fingerprint) + .with_for_update() + ) + if attempt is None: + attempt = HttpBridgeRecoveryAttemptRecord( + session_id=session_id, + request_fingerprint=request_fingerprint, + request_id=request_id, + account_id=account_id, + model=model, + replay_safe=replay_safe, + state=HttpBridgeRecoveryAttemptState.UNKNOWN, + ) + self._session.add(attempt) + elif attempt.state == HttpBridgeRecoveryAttemptState.REPLAYED: + snapshot = _to_recovery_attempt_snapshot(attempt) + await self._session.rollback() + return snapshot + elif attempt.request_id != request_id: + # A different request already owns the UNKNOWN checkpoint. + # Do not overwrite it while that request may still be between + # admission and dispatch; the caller must fail closed rather + # than sharing a journal generation. + snapshot = _to_recovery_attempt_snapshot(attempt) + await self._session.rollback() + return snapshot + else: + attempt.request_id = request_id + attempt.account_id = account_id + attempt.model = model + attempt.replay_safe = replay_safe + attempt.state = HttpBridgeRecoveryAttemptState.UNKNOWN + attempt.response_id = None + try: + await self._session.commit() + except IntegrityError: + # A concurrent owner may have inserted the same fingerprint + # after our initial SELECT (the absent-row case cannot be + # locked by SQLite). Re-read the winner and use its durable + # state instead of surfacing a transient uniqueness failure. + await self._session.rollback() + owner_exists = await self._session.scalar( + select(HttpBridgeSessionRecord.id) + .where( + HttpBridgeSessionRecord.id == session_id, + HttpBridgeSessionRecord.owner_instance_id == instance_id, + HttpBridgeSessionRecord.owner_epoch == owner_epoch, + ) + .with_for_update() + ) + if owner_exists is None: + await self._session.rollback() + return None + attempt = await self._session.scalar( + select(HttpBridgeRecoveryAttemptRecord) + .where(HttpBridgeRecoveryAttemptRecord.session_id == session_id) + .where(HttpBridgeRecoveryAttemptRecord.request_fingerprint == request_fingerprint) + ) + if attempt is None: + raise + if attempt.state == HttpBridgeRecoveryAttemptState.REPLAYED: + snapshot = _to_recovery_attempt_snapshot(attempt) + await self._session.rollback() + return snapshot + if attempt.request_id != request_id: + snapshot = _to_recovery_attempt_snapshot(attempt) + await self._session.rollback() + return snapshot + attempt.request_id = request_id + attempt.account_id = account_id + attempt.model = model + attempt.replay_safe = replay_safe + attempt.state = HttpBridgeRecoveryAttemptState.UNKNOWN + attempt.response_id = None + await self._session.commit() + await self._session.refresh(attempt) + return _to_recovery_attempt_snapshot(attempt) + + async def lookup_recovery_attempt( + self, + *, + session_id: str, + request_fingerprint: str, + ) -> DurableBridgeRecoveryAttemptSnapshot | None: + attempt = await self._session.scalar( + select(HttpBridgeRecoveryAttemptRecord) + .where(HttpBridgeRecoveryAttemptRecord.session_id == session_id) + .where(HttpBridgeRecoveryAttemptRecord.request_fingerprint == request_fingerprint) + .where(HttpBridgeRecoveryAttemptRecord.state == HttpBridgeRecoveryAttemptState.UNKNOWN) + .where(HttpBridgeRecoveryAttemptRecord.replay_safe.is_(True)) + ) + return _to_recovery_attempt_snapshot(attempt) if attempt is not None else None + + async def mark_recovery_attempt_replayed( + self, + *, + session_id: str, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + response_id: str | None = None, + ) -> bool: + async with sqlite_writer_section(): + # Keep the owner fence and journal transition in one transaction. + # PostgreSQL's row lock prevents a concurrent takeover from + # advancing the epoch between the check and the state update; + # sqlite_writer_section provides the equivalent writer + # serialization for SQLite. + owner_exists = await self._session.scalar( + select(HttpBridgeSessionRecord.id) + .where( + HttpBridgeSessionRecord.id == session_id, + HttpBridgeSessionRecord.owner_instance_id == instance_id, + HttpBridgeSessionRecord.owner_epoch == owner_epoch, + ) + .with_for_update() + ) + if owner_exists is None: + await self._session.rollback() + return False + values: dict[str, object] = {"state": HttpBridgeRecoveryAttemptState.REPLAYED} + if response_id is not None: + values["response_id"] = response_id + # A claim authorizes one replay and must only transition UNKNOWN + # rows. Settlement (which supplies response_id) remains idempotent + # for a REPLAYED row after the replay completes. + claimable_states = ( + (HttpBridgeRecoveryAttemptState.UNKNOWN,) + if response_id is None + else (HttpBridgeRecoveryAttemptState.UNKNOWN, HttpBridgeRecoveryAttemptState.REPLAYED) + ) + result = await self._session.execute( + update(HttpBridgeRecoveryAttemptRecord) + .where( + HttpBridgeRecoveryAttemptRecord.session_id == session_id, + HttpBridgeRecoveryAttemptRecord.request_fingerprint == request_fingerprint, + HttpBridgeRecoveryAttemptRecord.state.in_(claimable_states), + ) + .values(**values) + ) + await self._session.commit() + return bool(getattr(result, "rowcount", 0)) + + async def rollback_recovery_attempt_replayed( + self, + *, + session_id: str, + instance_id: str, + owner_epoch: int, + request_fingerprint: str, + ) -> bool: + """Return a pre-dispatch replay claim to UNKNOWN under the owner fence.""" + async with sqlite_writer_section(): + owner_exists = await self._session.scalar( + select(HttpBridgeSessionRecord.id) + .where( + HttpBridgeSessionRecord.id == session_id, + HttpBridgeSessionRecord.owner_instance_id == instance_id, + HttpBridgeSessionRecord.owner_epoch == owner_epoch, + ) + .with_for_update() + ) + if owner_exists is None: + await self._session.rollback() + return False + result = await self._session.execute( + update(HttpBridgeRecoveryAttemptRecord) + .where( + HttpBridgeRecoveryAttemptRecord.session_id == session_id, + HttpBridgeRecoveryAttemptRecord.request_fingerprint == request_fingerprint, + HttpBridgeRecoveryAttemptRecord.state == HttpBridgeRecoveryAttemptState.REPLAYED, + HttpBridgeRecoveryAttemptRecord.response_id.is_(None), + ) + .values(state=HttpBridgeRecoveryAttemptState.UNKNOWN) + ) + await self._session.commit() + return bool(getattr(result, "rowcount", 0)) + async def _execute_fenced_session_update( self, *, @@ -639,6 +1198,43 @@ async def purge_abandoned_before(self, cutoff: datetime, *, batch_size: int = _P await self._session.commit() deleted_count += len(deleted.scalars().all()) + async def purge_retry_circuits_before( + self, + cutoff_epoch: float, + *, + batch_size: int = _PURGE_CLOSED_BATCH_SIZE, + ) -> int: + deleted_count = 0 + while True: + result = await self._session.execute( + select( + HttpBridgeRetryCircuit.session_key_kind, + HttpBridgeRetryCircuit.session_key_hash, + HttpBridgeRetryCircuit.api_key_scope, + ) + .where(HttpBridgeRetryCircuit.updated_at_epoch < cutoff_epoch) + .limit(batch_size) + ) + keys = [tuple(row) for row in result.fetchall()] + if not keys: + return deleted_count + batch_deleted_count = 0 + async with sqlite_writer_section(): + for session_key_kind, session_key_hash, api_key_scope in keys: + deleted = await self._session.execute( + delete(HttpBridgeRetryCircuit) + .where(HttpBridgeRetryCircuit.session_key_kind == session_key_kind) + .where(HttpBridgeRetryCircuit.session_key_hash == session_key_hash) + .where(HttpBridgeRetryCircuit.api_key_scope == api_key_scope) + .where(HttpBridgeRetryCircuit.updated_at_epoch < cutoff_epoch) + .returning(HttpBridgeRetryCircuit.session_key_hash) + ) + batch_deleted_count += len(deleted.scalars().all()) + await self._session.commit() + if batch_deleted_count == 0: + return deleted_count + deleted_count += batch_deleted_count + async def upsert_alias( self, *, @@ -1075,7 +1671,9 @@ async def missing_durable_bridge_tables(session: AsyncSession) -> tuple[str, ... result = await session.execute( text( "SELECT name FROM sqlite_master " - "WHERE type = 'table' AND name IN ('http_bridge_sessions', 'http_bridge_session_aliases')" + "WHERE type = 'table' " + "AND name IN ('http_bridge_sessions', 'http_bridge_session_aliases', 'http_bridge_retry_circuits', " + "'http_bridge_recovery_attempts')" ) ) else: @@ -1083,7 +1681,10 @@ async def missing_durable_bridge_tables(session: AsyncSession) -> tuple[str, ... text( "SELECT table_name FROM information_schema.tables " "WHERE table_schema = 'public' " - "AND table_name IN ('http_bridge_sessions', 'http_bridge_session_aliases')" + "AND table_name IN (" + "'http_bridge_sessions', 'http_bridge_session_aliases', 'http_bridge_retry_circuits', " + "'http_bridge_recovery_attempts'" + ")" ) ) present = {str(row[0]) for row in result.fetchall()} @@ -1175,3 +1776,32 @@ def _to_snapshot_required(row: HttpBridgeSessionRecord) -> DurableBridgeSessionS if snapshot is None: raise RuntimeError("Expected durable bridge session snapshot") return snapshot + + +def _to_recovery_attempt_snapshot( + row: HttpBridgeRecoveryAttemptRecord, +) -> DurableBridgeRecoveryAttemptSnapshot: + return DurableBridgeRecoveryAttemptSnapshot( + session_id=row.session_id, + request_fingerprint=row.request_fingerprint, + request_id=row.request_id, + account_id=row.account_id, + model=row.model, + replay_safe=bool(row.replay_safe), + state=row.state, + response_id=row.response_id, + ) + + +def _to_retry_circuit_snapshot(row: HttpBridgeRetryCircuit | None) -> DurableBridgeRetryCircuitSnapshot | None: + if row is None: + return None + return DurableBridgeRetryCircuitSnapshot( + session_key_kind=row.session_key_kind, + session_key_hash=row.session_key_hash, + api_key_scope=row.api_key_scope, + consecutive_failures=row.consecutive_failures, + cooldown_until_epoch=row.cooldown_until_epoch, + last_detail=row.last_detail, + updated_at_epoch=row.updated_at_epoch, + ) diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index b077a794fd..000f2f4b55 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -139,9 +139,15 @@ from app.modules.proxy._service.http_bridge.helpers import ( _active_http_bridge_instance_ring as _active_http_bridge_instance_ring, ) +from app.modules.proxy._service.http_bridge.helpers import ( + _await_cancelled_task as _await_cancelled_task, +) from app.modules.proxy._service.http_bridge.helpers import ( _build_http_bridge_prewarm_text as _build_http_bridge_prewarm_text, ) +from app.modules.proxy._service.http_bridge.helpers import ( + _cancel_and_track_cancelled_task as _cancel_and_track_cancelled_task, +) from app.modules.proxy._service.http_bridge.helpers import ( _durable_bridge_lookup_active_owner as _durable_bridge_lookup_active_owner, ) @@ -314,6 +320,10 @@ from app.modules.proxy._service.http_bridge.helpers import ( _trim_http_bridge_previous_response_input_items as _trim_http_bridge_previous_response_input_items, ) +from app.modules.proxy._service.http_bridge.retry_circuit import ( + _HTTPBridgeRetryCircuitMixin, + _initialize_http_bridge_retry_circuit, +) from app.modules.proxy._service.observability import _hash_identifier as _hash_identifier from app.modules.proxy._service.observability import _hash_identifier_or_none as _hash_identifier_or_none from app.modules.proxy._service.observability import _interesting_header_keys as _interesting_header_keys @@ -742,9 +752,6 @@ logger = logging.getLogger(__name__) - -_TASK_CANCEL_TIMEOUT_SECONDS = 1.0 -_TaskResultT = TypeVar("_TaskResultT") _ResponsesPayloadT = TypeVar("_ResponsesPayloadT", ResponsesRequest, ResponsesCompactRequest) _DOWNSTREAM_WEBSOCKET_IDLE_CLOSE_REASON = "Idle downstream websocket timeout" _DOWNSTREAM_WEBSOCKET_RECEIVE_POLL_SECONDS = 1.0 @@ -783,23 +790,6 @@ def _proxy_admission_wait_timeout_seconds(settings: Any | None = None) -> float: _STREAM_KEEPALIVE_MAX_COUNT = 6 -async def _await_cancelled_task( - task: asyncio.Task[_TaskResultT], - *, - timeout_seconds: float = _TASK_CANCEL_TIMEOUT_SECONDS, - label: str, -) -> bool: - task.cancel() - try: - await asyncio.wait_for(task, timeout=timeout_seconds) - except asyncio.CancelledError: - return True - except TimeoutError: - logger.warning("Timed out waiting for %s cancellation", label) - return False - return True - - _TEXT_DELTA_EVENT_TYPES = frozenset({"response.output_text.delta", "response.refusal.delta"}) _TEXT_DONE_CONTENT_PART_TYPES = frozenset({"output_text", "refusal"}) _REQUEST_TRANSPORT_HTTP = "http" @@ -814,6 +804,7 @@ async def _await_cancelled_task( *PERMANENT_FAILURE_CODES.keys(), } ) + _TRANSIENT_RETRY_CODES = frozenset( { "overloaded_error", @@ -924,6 +915,7 @@ class ProxyService( _CompactMixin, _StreamingMixin, _WebSocketMixin, + _HTTPBridgeRetryCircuitMixin, _HTTPBridgeMixin, ): def __init__( @@ -941,6 +933,7 @@ def __init__( self._durable_bridge = DurableBridgeSessionCoordinator(SessionLocal) self._http_bridge_owner_client = HTTPBridgeOwnerClient() self._http_bridge_sessions: dict[_HTTPBridgeSessionKey, _HTTPBridgeSession] = {} + _initialize_http_bridge_retry_circuit(self) self._http_bridge_inflight_sessions: dict[_HTTPBridgeSessionKey, asyncio.Future[_HTTPBridgeSession]] = {} self._http_bridge_turn_state_index: dict[tuple[str, str | None], _HTTPBridgeSessionKey] = {} self._http_bridge_previous_response_index: dict[tuple[str, str | None], _HTTPBridgeSessionKey] = {} @@ -1308,6 +1301,7 @@ async def _acquire_request_state_response_create_admission( pending_request_ids: list[str] | None = None pending_request_ages_seconds: list[float] | None = None should_retire_stuck_session = False + stale_pending_requests_to_fail: list[_WebSocketRequestState] = [] if bridge_session is not None: now = time.monotonic() async with bridge_session.pending_lock: @@ -1317,25 +1311,46 @@ async def _acquire_request_state_response_create_admission( pending_request_ids = [state.request_log_id or state.request_id for state in pending_states] pending_request_ages_seconds = [max(0.0, now - state.started_at) for state in pending_states] threshold_seconds = float( - getattr( - get_settings(), - "http_responses_session_bridge_stuck_gate_retire_after_seconds", - 300.0, - ) + getattr(get_settings(), "http_responses_session_bridge_stuck_gate_retire_after_seconds", 300.0) ) - # Leading telemetry records latency without assigning a response - # or releasing this gate; only response-created proves progress. - should_retire_stuck_session = any( - state.transport == _REQUEST_TRANSPORT_HTTP - and not state.skip_request_log - and state.response_create_gate_acquired - and state.awaiting_response_created - and not state.downstream_visible - and state.latency_response_created_ms is None - and state.response_event_count == 0 - and max(0.0, now - state.started_at) >= threshold_seconds - for state in pending_states + stale_pending_requests_to_fail, should_retire_stuck_session = ( + self._classify_http_bridge_stale_gate_holders( + pending_states, + now=now, + threshold_seconds=threshold_seconds, + session_closed=bridge_session.closed, + ) ) + if not should_retire_stuck_session and any( + max(0.0, now - state.started_at) >= threshold_seconds for state in pending_states + ): + # A gate waiter starved past the stuck threshold without the + # watchdog firing: dump every pending state's verdict inputs + # so the blocking condition is identifiable from prod logs. + logger.warning( + "http_bridge_stuck_watchdog_skipped session_closed=%s candidates=%s states=%s", + bridge_session.closed, + len(stale_pending_requests_to_fail), + "; ".join( + ( + f"id={state.request_log_id or state.request_id}" + f" transport={state.transport}" + f" skip_log={state.skip_request_log}" + f" draining={state.draining_until_terminal}" + f" prev_resp={state.previous_response_id is not None}" + f" hard_anchor={state.hard_continuity_anchor}" + f" gate_acq={state.response_create_gate_acquired}" + f" awaiting={state.awaiting_response_created}" + f" resp_id={state.response_id is not None}" + f" events={state.response_event_count}" + f" created_ms={state.latency_response_created_ms}" + f" dsvis={state.downstream_visible}" + f" age={max(0.0, now - state.started_at):.0f}" + f" gate_wait_age={max(0.0, now - state.response_create_gate_wait_started_at) if state.response_create_gate_wait_started_at is not None else None}" # noqa: E501 + ) + for state in pending_states + ), + ) _log_http_bridge_startup_wait_timeout( stage="response_create_gate", timeout_seconds=timeout_seconds, @@ -1348,7 +1363,13 @@ async def _acquire_request_state_response_create_admission( pending_request_ids=pending_request_ids, pending_request_ages_seconds=pending_request_ages_seconds, ) - if bridge_session is not None and should_retire_stuck_session: + if bridge_session is not None and stale_pending_requests_to_fail: + await self._fail_stale_http_bridge_pending_requests( + bridge_session, + stale_pending_requests_to_fail, + detail="response_create_gate_timeout_stuck_pending", + ) + elif bridge_session is not None and should_retire_stuck_session: _record_http_bridge_stuck_retire( reason="response_create_gate_timeout_stuck_pending", session=bridge_session, diff --git a/app/modules/sticky_sessions/cleanup_scheduler.py b/app/modules/sticky_sessions/cleanup_scheduler.py index e9a795df46..e19d4a1e07 100644 --- a/app/modules/sticky_sessions/cleanup_scheduler.py +++ b/app/modules/sticky_sessions/cleanup_scheduler.py @@ -4,6 +4,7 @@ import contextlib import importlib import logging +import time from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from datetime import timedelta @@ -14,7 +15,11 @@ from app.core.utils.time import utcnow from app.db.models import DashboardSettings from app.db.session import SessionLocal, get_background_session -from app.modules.proxy.durable_bridge_repository import DurableBridgeRepository, missing_durable_bridge_tables +from app.modules.proxy.durable_bridge_repository import ( + DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS, + DurableBridgeRepository, + missing_durable_bridge_tables, +) from app.modules.proxy.ring_membership import RING_MEMBER_RETENTION_SECONDS, RingMembershipService from app.modules.proxy.sticky_repository import StickySessionsRepository from app.modules.settings.repository import SettingsRepository @@ -121,6 +126,14 @@ async def _cleanup_as_leader(self) -> None: logger.info( "Purged abandoned HTTP bridge sessions deleted_count=%s", abandoned_deleted_count ) + retry_circuit_deleted_count = await bridge_repo.purge_retry_circuits_before( + time.time() - DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS + ) + if retry_circuit_deleted_count > 0: + logger.info( + "Purged expired HTTP bridge retry circuits deleted_count=%s", + retry_circuit_deleted_count, + ) ring_cutoff = utcnow() - timedelta(seconds=RING_MEMBER_RETENTION_SECONDS) ring_deleted_count = await RingMembershipService(SessionLocal).purge_stale_before(ring_cutoff) if ring_deleted_count > 0: diff --git a/docs/reference/settings.md b/docs/reference/settings.md index 1038b33fd8..cbbf1f855c 100644 --- a/docs/reference/settings.md +++ b/docs/reference/settings.md @@ -7,7 +7,7 @@ Regenerate with `uv run python scripts/generate_settings_reference.py`; `tests/unit/test_settings_reference.py` fails when this page drifts from `app/core/config/settings.py`. -codex-lb currently exposes 115 settings. Every setting is an environment +codex-lb currently exposes 116 settings. Every setting is an environment variable with the `CODEX_LB_` prefix (process environment or `.env` / `.env.local` next to the process). All defaults work with zero configuration — start from [Configuration](../configuration.md) for the handful that matter, @@ -82,6 +82,7 @@ the host side of the compose `ports` mapping instead. | Environment variable | Type | Default | | --- | --- | --- | | `CODEX_LB_HTTP_RESPONSES_SESSION_BRIDGE_ADVERTISE_BASE_URL` | `str \| None` | `None` | +| `CODEX_LB_HTTP_RESPONSES_SESSION_BRIDGE_CLEAN_CLOSE_RETRY_JITTER_MAX_SECONDS` | `float` | `2.0` | | `CODEX_LB_HTTP_RESPONSES_SESSION_BRIDGE_CODEX_IDLE_TTL_SECONDS` | `float` | `900.0` | | `CODEX_LB_HTTP_RESPONSES_SESSION_BRIDGE_CODEX_PREWARM_ENABLED` | `bool` | `False` | | `CODEX_LB_HTTP_RESPONSES_SESSION_BRIDGE_ENABLED` | `bool` | `True` | diff --git a/openspec/changes/recover-codex-desktop-idle-bridge/specs/proxy-admission-control/spec.md b/openspec/changes/recover-codex-desktop-idle-bridge/specs/proxy-admission-control/spec.md index 51bbccd853..293763926d 100644 --- a/openspec/changes/recover-codex-desktop-idle-bridge/specs/proxy-admission-control/spec.md +++ b/openspec/changes/recover-codex-desktop-idle-bridge/specs/proxy-admission-control/spec.md @@ -2,19 +2,20 @@ ### Requirement: Stuck HTTP bridge response-create gate sessions are retired -The proxy MUST retain the existing waiter-triggered retirement behavior for stale HTTP bridge response-create gate owners and MUST additionally enforce an owner-side deadline for a visible HTTP request whose current upstream `response.create` send remains completely eventless before `response.created`. The owner-side deadline MUST be measured from a monotonic timestamp recorded immediately before the current upstream send, MUST use the smaller of the configured stuck-gate retirement threshold and 240 seconds, MUST run without a second gate waiter, and MUST remain active when periodic SSE keepalives are disabled. +The proxy MUST retain the existing waiter-triggered retirement behavior for stale HTTP bridge response-create gate owners and MUST additionally enforce an owner-side deadline for a visible HTTP request whose current upstream `response.create` send remains completely eventless before `response.created`. The owner-side deadline MUST be measured from a monotonic timestamp recorded immediately before the current upstream send, MUST use the smaller of the configured stuck-gate retirement threshold and 60 seconds, MUST run without a second gate waiter, and MUST remain active when periodic SSE keepalives are disabled. The owner-side watchdog MUST apply only while the request owns the response-create gate, awaits `response.created`, has neither a response id nor recorded `response.created` latency, has received no matched `response.*` lifecycle event, and has produced no downstream-visible output or sequence evidence. Non-response telemetry such as `codex.rate_limits` MUST NOT suppress this watchdog. Any matched `response.*` lifecycle event, response-created milestone, or downstream-visible evidence MUST suppress the owner-side watchdog and leave existing timeout behavior unchanged. -When the owner-side deadline expires, the proxy MUST recheck eligibility, emit a structured low-cardinality log and the existing stuck-retirement Prometheus counter, terminally fail and settle every pending request exactly once, and retire the whole bridge session. It MUST NOT transparently replay the timed-out request, move it to another account, or write an account-health failure for the missing-created timeout. +When the owner-side deadline expires, the proxy MUST recheck eligibility and emit a structured low-cardinality log and the existing stuck-retirement Prometheus counter. For requests that are not eligible for the bounded fresh-hard recovery defined by `recover-fresh-hard-bridge-timeouts`, it MUST terminally fail and settle every pending request exactly once, retire the whole bridge session, and MUST NOT transparently replay the timed-out request or move it to another account. An eligible fresh hard request MAY take that single bounded recovery path; if recovery is unavailable or fails, it MUST fall back to the same terminal fail-closed retirement. Neither path may write an account-health failure solely because `response.created` was missing. #### Scenario: Lone eventless gate owner is retired before the client timeout - **GIVEN** a visible HTTP bridge request owns the response-create gate - **AND** its current `response.create` send produced no matched `response.*` event, response id, or downstream-visible output - **AND** no second request waits for the gate -- **WHEN** the smaller of the configured stuck threshold and 240 seconds elapses after the current send -- **THEN** the proxy emits an explicit terminal failure and retires the bridge session +- **WHEN** the smaller of the configured stuck threshold and 60 seconds elapses after the current send +- **THEN** the proxy emits an explicit terminal failure and retires the bridge session when the request is not eligible for bounded fresh-hard recovery +- **AND** an eligible fresh hard request instead follows the single bounded recovery defined by `recover-fresh-hard-bridge-timeouts` - **AND** recovery occurs before the native client's 300-second parsed-event idle timeout #### Scenario: Send time rather than request age anchors the deadline @@ -43,5 +44,5 @@ When the owner-side deadline expires, the proxy MUST recheck eligibility, emit a - **GIVEN** an eventless pre-created owner reaches the owner-side deadline - **WHEN** terminal cleanup runs - **THEN** every pending request is settled exactly once and the whole session is retired -- **AND** the proxy does not replay the timed-out request or submit it on another account +- **AND** the proxy does not replay the timed-out request or submit it on another account unless it satisfies the bounded fresh-hard recovery requirement - **AND** the selected account is not marked unhealthy solely because `response.created` was missing diff --git a/openspec/changes/recover-fresh-hard-bridge-timeouts/proposal.md b/openspec/changes/recover-fresh-hard-bridge-timeouts/proposal.md new file mode 100644 index 0000000000..46dc9567ef --- /dev/null +++ b/openspec/changes/recover-fresh-hard-bridge-timeouts/proposal.md @@ -0,0 +1,29 @@ +## Why + +Fresh HTTP Responses bridge requests can remain pinned to a hard session-header +owner even after repeated pre-response retries prove that the selected upstream +socket is not acknowledging `response.create`. The retry circuit prevents a +retry storm, but the eventless watchdog eventually returns a stream error after +the 240-second safety cap. A self-contained request with no previous-response, +turn-state, or account-scoped file ownership can safely move to another active +account without breaking continuity. + +## What Changes + +- Permit pre-response recovery to exclude the failing account for a fresh, + self-contained hard session-header request. +- Permit a proof-gated client full resend to replay once on its required + continuity owner, including when that owner's retry circuit is cooling down. +- Keep previous-response, turn-state, file-pinned, and proxy-injected anchored + requests on their required account. +- Preserve the existing retry circuit and eventless watchdog as bounded + fallbacks when no alternate account is eligible. +- Add regression coverage proving fresh hard requests can switch accounts while + anchored requests remain owner-bound. + +## Impact + +- Affected capabilities: `proxy-admission-control`, `responses-api-compat`. +- Fresh requests recover without waiting for the client-safe eventless timeout + when another account is available. +- Continuity-sensitive requests retain their existing fail-closed behavior. diff --git a/openspec/changes/recover-fresh-hard-bridge-timeouts/specs/proxy-admission-control/spec.md b/openspec/changes/recover-fresh-hard-bridge-timeouts/specs/proxy-admission-control/spec.md new file mode 100644 index 0000000000..c18e0f8bf5 --- /dev/null +++ b/openspec/changes/recover-fresh-hard-bridge-timeouts/specs/proxy-admission-control/spec.md @@ -0,0 +1,74 @@ +# proxy-admission-control Delta + +## ADDED Requirements + +### Requirement: Fresh hard bridge requests may recover across accounts + +When a hard HTTP bridge request is still pre-response and has no +`previous_response_id`, hard continuity anchor, proxy-injected anchor, or +account-scoped file ownership, pre-response recovery MAY exclude the failed +session account and select another eligible account. The request MUST retain +its original request body and deadline. Requests carrying any of those +continuity or ownership markers MUST remain pinned to the required account. + +#### Scenario: Fresh hard request switches after silent upstream failure + +- **GIVEN** a hard session-header request has sent `response.create` +- **AND** upstream has not emitted `response.created` or any response event +- **AND** the request has no previous-response, turn-state, proxy-injected + anchor, or account-scoped file ownership +- **WHEN** pre-response recovery retries the request +- **THEN** the failed account is excluded from selection +- **AND** another eligible account may receive the unchanged request body +- **AND** the original request deadline remains in force + +#### Scenario: Eventless watchdog gives fresh requests one bounded recovery + +- **GIVEN** a hard session-header request has reached the eventless + `response.created` watchdog without response events +- **AND** the request has no previous-response, turn-state, proxy-injected + anchor, or account-scoped file ownership +- **WHEN** the client-safe watchdog deadline expires +- **THEN** the proxy attempts the same bounded pre-response recovery once +- **AND** the failed account is excluded when recovery selects a replacement +- **AND** if recovery is unavailable, the proxy preserves the existing + terminal timeout behavior + +#### Scenario: Fresh account recovery bypasses a stale retry circuit + +- **GIVEN** a hard session key has an active retry cooldown from repeated + pre-response failures +- **AND** the pending request is fresh, self-contained, and has no continuity + or account-ownership marker +- **WHEN** bounded pre-response recovery is attempted +- **THEN** the request may bypass that cooldown once to exclude the failed + account +- **AND** continuity-bound requests remain subject to the retry cooldown + +#### Scenario: Continuity-bound hard request remains pinned + +- **GIVEN** a hard request has a previous-response id, continuity anchor, + proxy-injected anchor, or account-scoped file ownership +- **WHEN** pre-response recovery retries the request +- **THEN** the original account remains required +- **AND** the request is not replayed through another account + +#### Scenario: Proof-gated client full resend replays on the continuity owner + +- **GIVEN** a hard request has a previous-response id and a client-provided + full resend whose input body has passed the bridge's retry-safety checks +- **AND** upstream has not emitted `response.created` or any response event +- **WHEN** bounded pre-response recovery is attempted +- **THEN** the bridge may strip the previous-response id and replay the verified + full body once +- **AND** recovery remains pinned to the original continuity owner +- **AND** an unverified continuation remains fail-closed + +#### Scenario: Unsafe continuity timeout does not wait through an unusable cooldown + +- **GIVEN** a hard continuation has no proof-gated full resend available +- **AND** the retry circuit is cooling down after repeated pre-response failures +- **WHEN** the downstream keepalive window expires +- **THEN** the proxy fails the stream closed immediately +- **AND** it does not hold the client connection open until the cooldown ends +- **AND** the client may retry with its continuity payload intact diff --git a/openspec/changes/recover-fresh-hard-bridge-timeouts/tasks.md b/openspec/changes/recover-fresh-hard-bridge-timeouts/tasks.md new file mode 100644 index 0000000000..d1326d05f0 --- /dev/null +++ b/openspec/changes/recover-fresh-hard-bridge-timeouts/tasks.md @@ -0,0 +1,9 @@ +- [x] Add the delta specification for safe fresh hard-session account recovery. +- [x] Implement the guarded account exclusion during pre-response recovery. +- [x] Route eventless watchdog expiry through the bounded recovery path. +- [x] Allow fresh hard account recovery to bypass a stale cooldown once. +- [x] Allow proof-gated client full resends to recover on their required owner. +- [x] Fail unsafe continuity streams promptly instead of waiting through an unusable cooldown. +- [x] Add regression coverage for fresh and continuity-bound hard requests. +- [ ] Run focused tests, architecture checks, and strict OpenSpec validation. +- [ ] Build and canary-test the image before any deployment. diff --git a/openspec/changes/recover-repeated-clean-close/.openspec.yaml b/openspec/changes/recover-repeated-clean-close/.openspec.yaml new file mode 100644 index 0000000000..ff5f854ace --- /dev/null +++ b/openspec/changes/recover-repeated-clean-close/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-17 diff --git a/openspec/changes/recover-repeated-clean-close/proposal.md b/openspec/changes/recover-repeated-clean-close/proposal.md new file mode 100644 index 0000000000..04d3093d7e --- /dev/null +++ b/openspec/changes/recover-repeated-clean-close/proposal.md @@ -0,0 +1,47 @@ +## Why + +The HTTP Responses bridge currently opens its retry circuit after a clean +upstream WebSocket close even when the replacement socket also closes before +producing any response event. A downstream idle-recovery task can also replace +the upstream socket without restarting its reader. Closing the old socket then +wakes that stale reader, which misclassifies the proxy-initiated close as an +upstream failure and retires work already moved to the replacement socket. +Together these behaviors make a transient handoff issue visible as a reconnect +loop and require the Codex client to be restarted. + +## What Changes + +- Permit one additional pre-visible replay when the replacement upstream + WebSocket closes cleanly before any response event. +- Add bounded, configurable jitter before that additional replay to avoid + synchronized reconnects. +- Emit a dedicated diagnostic event for the additional clean-close replay. +- Keep the allowance hard-capped at one and preserve all existing no-replay + behavior after downstream-visible output or continuity-sensitive state. +- When recovery is initiated outside the upstream reader, cancel and await the + old reader before closing its socket, then start exactly one reader for the + replacement socket. +- Keep the shared session live while the replacement socket opens so concurrent + idle pruning cannot evict and fail its pending response during the handoff. +- Start silent pre-response recovery with enough headroom to reconnect before + the downstream client's request timeout boundary. +- Do not let a proxy-initiated close of a superseded socket retire pending work + on the replacement socket or increment the retry circuit. +- Detect a stuck pre-response gate from the absence of upstream activity and + response creation, rather than admission flags alone. Give requests with a + prior continuity anchor a bounded two-threshold grace period, and emit + diagnostic state when the watchdog skips a candidate. + +## Impact + +- Repeated clean handoffs can recover transparently without an immediate + terminal circuit-open response. +- The retry remains bounded and does not create an unbounded replay loop. +- Reader ownership follows the active socket across idle recovery, preventing + locally generated close frames from being counted as upstream instability. +- Adds the `http_bridge_retry_circuits` durable table and migration so retry + cooldown state survives cross-replica clean-close and incomplete-stream + failures. +- Adds a forward-only request-usage rollup repair migration for deployments + already stamped at the previous merge head, so changing migration ancestry + cannot leave startup schema-drift checks failing. diff --git a/openspec/changes/recover-repeated-clean-close/specs/responses-api-compat/spec.md b/openspec/changes/recover-repeated-clean-close/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..d825873a65 --- /dev/null +++ b/openspec/changes/recover-repeated-clean-close/specs/responses-api-compat/spec.md @@ -0,0 +1,191 @@ +## MODIFIED Requirements + +### Requirement: Clean upstream close before any response event fails fast + +When the HTTP Responses bridge observes an upstream WebSocket close with +`close_code = 1000` before any `response.*` event has been surfaced for the +pending request, the proxy MUST preserve its existing pre-visible replay +guards. If the request has already used exactly one eligible pre-visible +replay and the replacement upstream WebSocket also closes cleanly before any +response event, the proxy MAY perform exactly one additional replay. The +additional replay MUST be hard-capped at one per request, and the configured +maximum MUST NOT raise that cap. + +The proxy MUST NOT replay after downstream-visible output, after a terminal +response event, or when continuity-sensitive request state makes replay unsafe. +Before the additional replay, the proxy MAY sleep for bounded configured +jitter. The proxy MUST emit a dedicated low-cardinality diagnostic event for +the additional replay. + +When a downstream HTTP stream task initiates pre-response recovery while the +upstream reader is blocked on the superseded socket, the proxy MUST cancel and +await that reader before locally closing the socket. It MUST then start exactly +one reader for the replacement socket. A close caused by replacing the socket +MUST NOT be recorded as an upstream clean-close failure, MUST NOT increment the +retry circuit, and MUST NOT retire pending work moved to the replacement. The +cancelled reader's socket-generation finalizer MUST NOT leave the shared session +marked closed while the replacement socket is being selected or opened, so idle +pruning MUST NOT evict the handoff in progress. + +The default pre-response idle-recovery window MUST leave bounded headroom +before the downstream client's request timeout. With the default ten-second +keepalive interval, the proxy MUST initiate eligible recovery after no more +than six silent intervals so replacement connection and first output can occur +before a 120-second client deadline. + +The stuck pre-response watchdog MUST judge staleness using elapsed time since +the last upstream activity and the absence of a response identifier or +`response.created` latency, not admission flags alone. A request with a prior +continuity anchor MUST receive at most two retire-thresholds of grace before +being considered stale. When the watchdog skips a candidate, it MUST emit a +low-cardinality diagnostic containing the session-closed state, candidate +count, and pending-state verdicts. + +#### Scenario: clean close before response output receives one bounded additional replay + +- **GIVEN** an HTTP bridge request has no surfaced `response.*` events +- **AND** its first pre-visible replay has already been used +- **WHEN** the replacement upstream WebSocket closes with code `1000` +- **THEN** the proxy performs one additional pre-visible replay +- **AND** the request replay count increases by one +- **AND** the proxy emits a `retry_precreated_clean_close` diagnostic event + +#### Scenario: repeated clean closes do not create an unbounded replay loop + +- **GIVEN** the additional clean-close replay has already been used +- **WHEN** another upstream WebSocket closes cleanly before response output +- **THEN** the proxy does not replay the request again +- **AND** the existing terminal or circuit handling is used + +#### Scenario: visible output still prevents clean-close replay + +- **GIVEN** the pending request has surfaced any response event downstream +- **WHEN** the upstream WebSocket closes with code `1000` +- **THEN** the proxy does not replay the request + +#### Scenario: clean-close retry jitter is bounded + +- **GIVEN** clean-close retry jitter is configured +- **WHEN** the additional clean-close replay is scheduled +- **THEN** the delay is no greater than the configured jitter maximum +- **AND** the hard replay cap remains one regardless of the configured value + +#### Scenario: downstream idle recovery transfers reader ownership + +- **GIVEN** the upstream reader is blocked on the current bridge socket +- **AND** the downstream HTTP stream task initiates eligible pre-response recovery +- **WHEN** the bridge replaces the upstream socket +- **THEN** the old reader is cancelled and awaited before its socket is closed +- **AND** the shared session remains live while the replacement socket opens +- **AND** idle pruning retains the registered session while the handoff is in progress +- **AND** exactly one reader owns the replacement socket +- **AND** the local close does not open or increment the retry circuit +- **AND** pending work remains attached to the replacement session + +#### Scenario: silent pre-response recovery precedes the client timeout + +- **GIVEN** the upstream has produced no response event +- **AND** the default ten-second keepalive interval is active +- **WHEN** six silent intervals elapse +- **THEN** the proxy initiates eligible pre-response recovery +- **AND** at least sixty seconds remain before a 120-second client request timeout + +#### Scenario: anchored stuck-gate grace is bounded + +- **GIVEN** a pending HTTP bridge request has a prior continuity anchor +- **AND** no response identifier or `response.created` latency has been recorded +- **WHEN** less than two retire thresholds have elapsed since the gate began waiting +- **THEN** the watchdog does not classify the request as stale +- **WHEN** two retire thresholds elapse without upstream activity +- **THEN** the watchdog may classify the request as stale + +#### Scenario: upstream activity resolves admission-flag ambiguity + +- **GIVEN** a pending request has not acquired the response-created gate +- **AND** upstream activity has not produced a response identifier or `response.created` +- **WHEN** the staleness threshold elapses +- **THEN** the watchdog classifies the request as stale +- **AND** emits pending-state verdict inputs when it skips a watchdog pass + +### Requirement: Durable retry-circuit state protects repeated hard-affinity failures + +For a hard-affinity bridge key, the proxy MUST scope retry-circuit state by +affinity kind, affinity key, and API-key scope (using a stable anonymous scope +when no API key is present). The proxy MUST record only the documented +pre-response failure classes (`stream_incomplete`, `clean_close`, and +`stream_idle_timeout`). + +The default circuit MUST open after two consecutive recorded failures. Once +open, it MUST suppress pre-created replay until the persisted cooldown expires, +using exponential backoff from sixty seconds up to ten minutes. Clean-close +failures MUST cap their cooldown at thirty seconds. The proxy MUST persist +failure count, cooldown deadline, last failure detail, and update time in the +`http_bridge_retry_circuits` table and MUST merge conflict updates so concurrent +replicas cannot shorten an existing cooldown. + +The clean-close retry jitter maximum MUST be read from the +`http_responses_session_bridge_clean_close_retry_jitter_max_seconds` runtime +setting and MUST be bounded to the inclusive range 0–30 seconds. + +The proxy MUST evict process-local circuit entries and their loaded/persisted +markers after one hour without use, independently of durable-row cleanup, so +one-shot hard-affinity keys cannot grow the worker's memory without bound. + +Before every hard-affinity retry decision, the proxy MUST refresh the durable +row so a cooldown opened by another replica is observed even when this process +has already loaded the key. A durable lookup or persistence failure MUST NOT +crash the request; the proxy MUST continue using available local state and +record the failure for observability. Rows older than one hour MUST be treated +as expired and removed. A successful terminal response MUST clear the local +and durable circuit state. + +#### Scenario: the second hard-key failure opens a durable circuit + +- **GIVEN** a hard-affinity key has one recorded pre-response failure +- **WHEN** a second eligible failure is recorded +- **THEN** the proxy opens the retry circuit +- **AND** persists at least two consecutive failures and a cooldown deadline +- **AND** subsequent pre-created replay is suppressed until that deadline + +#### Scenario: retry decisions observe a cooldown opened by another replica + +- **GIVEN** this replica previously looked up a hard-affinity key with no row +- **AND** another replica persists an open cooldown for that same key and API-key scope +- **WHEN** this replica evaluates the next pre-created retry +- **THEN** it refreshes durable state before deciding +- **AND** suppresses the retry for the persisted cooldown + +#### Scenario: circuit state remains isolated by key and API-key scope + +- **GIVEN** one hard-affinity key has an open circuit +- **WHEN** a different affinity key or API-key scope evaluates a retry +- **THEN** that request is not suppressed by the first key's circuit + +#### Scenario: durable circuit lookup failure does not fail the request + +- **GIVEN** durable retry-circuit lookup or persistence is unavailable +- **WHEN** the proxy evaluates or records a retry-circuit event +- **THEN** the request continues using any available local circuit state +- **AND** the failure is logged and exposed through retry-circuit observability + +### Requirement: Upstream websocket drops penalize affected accounts + +When an upstream websocket closes while one or more streamed response requests +are pending and have not reached a terminal event, the proxy MUST record a +transient upstream error for the account before signaling failure for those +pending requests, except when the close carries a classified process-wide +network failure, is a clean close (`close_code = 1000`) before any +`response.*` event, or carries the classified per-socket +`upstream_keepalive_timeout` transport error. Clean pre-response closes and +keepalive timeouts MUST remain account-neutral while using the bounded retry +and retry-circuit handling above. A classified process-wide network failure +MUST remain account neutral and use its network error code. For other closes, +the proxy MUST surface +`stream_incomplete` to affected pending requests. + +#### Scenario: clean pre-response close does not penalize the account + +- **GIVEN** a hard-affinity HTTP bridge request is pending with no surfaced response event +- **WHEN** the upstream websocket closes cleanly before response output +- **THEN** the proxy records the clean-close retry-circuit outcome +- **AND** the selected account is not penalized diff --git a/openspec/changes/recover-repeated-clean-close/tasks.md b/openspec/changes/recover-repeated-clean-close/tasks.md new file mode 100644 index 0000000000..aa0246098d --- /dev/null +++ b/openspec/changes/recover-repeated-clean-close/tasks.md @@ -0,0 +1,14 @@ +- [x] Add bounded clean-close replay settings with safe defaults. +- [x] Allow one additional clean-close replay only before visible output. +- [x] Add jitter and dedicated retry diagnostics. +- [x] Add regression coverage for the second replay and retry cap. +- [x] Restart the upstream reader when pre-response recovery is initiated by the downstream stream task. +- [x] Add regression coverage for old-reader cancellation and replacement-reader ownership. +- [x] Keep the shared session live across the cancelled reader's socket-generation finalizer. +- [x] Add regression coverage for concurrent pruning during reader handoff. +- [x] Move the default pre-response recovery threshold ahead of the client timeout boundary. +- [x] Bound anchored stuck-gate grace and evaluate staleness from upstream activity/response creation. +- [x] Emit stuck-watchdog skip diagnostics with pending-state verdict inputs. +- [x] Add a forward-only repair for databases stamped before request-usage rollups were connected to the merge head. +- [x] Validate the OpenSpec change and run the focused and full test suites. +- [x] Build and deploy the validated image, then verify production health and logs. diff --git a/openspec/changes/recover-safe-http-bridge-continuations/.openspec.yaml b/openspec/changes/recover-safe-http-bridge-continuations/.openspec.yaml new file mode 100644 index 0000000000..f205fc727f --- /dev/null +++ b/openspec/changes/recover-safe-http-bridge-continuations/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-29 diff --git a/openspec/changes/recover-safe-http-bridge-continuations/proposal.md b/openspec/changes/recover-safe-http-bridge-continuations/proposal.md new file mode 100644 index 0000000000..1d0c904f2b --- /dev/null +++ b/openspec/changes/recover-safe-http-bridge-continuations/proposal.md @@ -0,0 +1,25 @@ +## Why + +An HTTP bridge request can lose its upstream acknowledgement after the +`response.create` send. Retrying a continuation with the same +`previous_response_id` can fork work or duplicate side effects, while waiting +through a retry-circuit cooldown only consumes the client request budget. + +## What Changes + +- Allow one fresh-upstream replay only when the request state contains a + proof-gated, unanchored full-resend payload. The proof applies equally to + client-provided and proxy-injected anchors. +- Fail continuity-bound requests closed when a retry-circuit cooldown is + active and no safe fresh replay exists. +- Record proof-gated fresh-resend attempts in a durable recovery journal so a + later replica can replay only an unresolved, transport-ambiguous attempt. +- Keep ordinary requests and existing session ownership on their current + recovery paths, and emit the continuity-fail-closed diagnostic for + observability. + +## Impact + +- HTTP bridge continuation recovery and idle-timeout behavior. +- Adds a durable recovery-attempt journal migration and startup schema check; + no public API or account-status changes. diff --git a/openspec/changes/recover-safe-http-bridge-continuations/specs/responses-api-compat/spec.md b/openspec/changes/recover-safe-http-bridge-continuations/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..9fcd9ea000 --- /dev/null +++ b/openspec/changes/recover-safe-http-bridge-continuations/specs/responses-api-compat/spec.md @@ -0,0 +1,75 @@ +# responses-api-compat Delta + +## ADDED Requirements + +### Requirement: Proof-gated recovery attempts are durably fenced + +When an HTTP bridge request has a verified, account-neutral, unanchored full +resend body, the proxy MUST record that request fingerprint in the durable +recovery journal before dispatching it upstream. The record MUST be owned by +the current durable session owner epoch and MUST start in `unknown` state. +Requests without that replay-safety proof MUST NOT create a recovery-journal +record. + +#### Scenario: Safe resend is journaled before dispatch + +- **GIVEN** a request has a verified full-resend body that is safe to replay + without `previous_response_id` +- **WHEN** the proxy admits the request for upstream dispatch +- **THEN** the durable journal contains one `unknown` record for its session + and request fingerprint before `response.create` is sent + +#### Scenario: Suppressed request is not journaled + +- **GIVEN** a hard session retry circuit is cooling down +- **WHEN** the request is rejected before upstream dispatch +- **THEN** no recovery-journal record is created or refreshed + +### Requirement: Durable replay is limited to ambiguous transport outcomes + +The proxy MUST consume an `unknown` recovery-journal record for a fresh +account-neutral replay only after an ambiguous transport outcome, represented +by `stream_incomplete`, `stream_idle_timeout`, or +`upstream_request_timeout`, and only before any response event or downstream +output. Explicit deterministic `response.failed` errors MUST settle normally +and MUST NOT trigger a cross-account replay or consume the recovery fence. + +#### Scenario: Transport ambiguity permits one replay + +- **GIVEN** an `unknown` proof-gated journal record exists +- **AND** the upstream closes or times out before any response event +- **WHEN** the bridge handles the ambiguous transport failure +- **THEN** the record is atomically claimed and the request is replayed once + on a fresh account-neutral upstream session + +#### Scenario: Deterministic failure is not replayed + +- **GIVEN** an `unknown` proof-gated journal record exists +- **AND** upstream emits an explicit pre-output `response.failed` such as an + invalid request or quota rejection +- **WHEN** the bridge handles that terminal event +- **THEN** it forwards the terminal failure +- **AND** it leaves the journal available for settlement without replaying on + another account + +### Requirement: Recovery journal settlement is owner-fenced and idempotent + +After a replayed request reaches `response.completed`, the proxy MUST mark its +journal record `replayed` only through the current durable owner epoch and +MUST retain the downstream response id when available. Repeated settlement, +stale owners, and concurrent claim attempts MUST NOT produce a second replay. +The migration MUST be on the current Alembic head and startup schema checks +MUST require the journal table. + +#### Scenario: Completed replay settles once + +- **GIVEN** a replayed request completes successfully +- **WHEN** the completion event is processed +- **THEN** the matching journal record becomes `replayed` +- **AND** a later retry cannot claim it again + +#### Scenario: Stale owner cannot settle or replay + +- **GIVEN** a journal record belongs to a newer durable owner epoch +- **WHEN** an old replica attempts settlement or replay +- **THEN** the operation is rejected without changing the record state diff --git a/openspec/changes/recover-safe-http-bridge-continuations/tasks.md b/openspec/changes/recover-safe-http-bridge-continuations/tasks.md new file mode 100644 index 0000000000..0d6f412bfc --- /dev/null +++ b/openspec/changes/recover-safe-http-bridge-continuations/tasks.md @@ -0,0 +1,12 @@ +## 1. Safe recovery + +- [x] Gate fresh-upstream replay on the existing retry-safe full-resend proof, + independent of whether the anchor was client-provided or injected. +- [x] Fail continuity-bound requests closed instead of waiting through an + unusable retry-circuit cooldown. + +## 2. Verification + +- [x] Cover proof-gated replay and unsafe/session-bound replay behavior. +- [x] Cover continuity-bound classification separately from ordinary requests. +- [x] Run syntax, architecture, and whitespace validation. diff --git a/scripts/check_proxy_architecture.py b/scripts/check_proxy_architecture.py index 580b24bb16..97d0151833 100644 --- a/scripts/check_proxy_architecture.py +++ b/scripts/check_proxy_architecture.py @@ -23,9 +23,12 @@ HTTP_BRIDGE_MIXIN_PATH = PROXY_DIR / "_service" / "http_bridge" / "mixin.py" STREAMING_MIXIN_PATH = PROXY_DIR / "_service" / "streaming" / "mixin.py" -MAX_SERVICE_LINES = 2_604 +# Ruff 0.16's formatter adds blank lines around top-level definitions. Keep +# the ratchets at the formatted baseline so whitespace normalization does not +# consume architectural budget. +MAX_SERVICE_LINES = 2_617 MAX_LOAD_BALANCER_LINES = 3_260 -MAX_HTTP_BRIDGE_MIXIN_LINES = 2_400 +MAX_HTTP_BRIDGE_MIXIN_LINES = 2_436 MAX_STREAMING_MIXIN_LINES = 1_100 MAX_PROXY_SERVICE_METHOD_LINES = 1_200 MAX_LOAD_BALANCER_SELECT_ACCOUNT_LINES = 699 diff --git a/tests/integration/test_http_responses_bridge.py b/tests/integration/test_http_responses_bridge.py index 989783c5af..b69f5a6f3b 100644 --- a/tests/integration/test_http_responses_bridge.py +++ b/tests/integration/test_http_responses_bridge.py @@ -3895,8 +3895,8 @@ async def blocking_reader_task() -> None: blocking_reader = asyncio.create_task(blocking_reader_task()) bridge_session.upstream_reader = blocking_reader - async def fake_await_cancelled_task(task, *, timeout_seconds=1.0, label): - del task, timeout_seconds, label + async def fake_await_cancelled_task(task, *, timeout_seconds=1.0, label, cleanup_tasks=None): + del task, timeout_seconds, label, cleanup_tasks return False monkeypatch.setattr(proxy_module, "_await_cancelled_task", fake_await_cancelled_task) @@ -11729,6 +11729,112 @@ async def fake_connect_responses_websocket( assert connect_count == 1 +@pytest.mark.asyncio +async def test_v1_responses_http_bridge_idle_recovery_hands_reader_to_replacement( + async_client, + app_instance, + monkeypatch, +): + app_settings = _make_app_settings(enabled=True) + app_settings.sse_keepalive_interval_seconds = 0.01 + _install_proxy_settings( + monkeypatch, + app_settings=app_settings, + dashboard_settings=_make_dashboard_settings(), + ) + monkeypatch.setattr(proxy_module, "_HTTP_BRIDGE_STARTUP_KEEPALIVE_GRACE_SECONDS", 0.01) + monkeypatch.setattr(proxy_module, "_STREAM_KEEPALIVE_MAX_COUNT", 1) + account_id = await _import_account( + async_client, + "acc_http_bridge_reader_handoff", + "http-bridge-reader-handoff@example.com", + ) + account = await _get_account(account_id) + upstreams = [_SilentUpstreamWebSocket(), _FakeBridgeUpstreamWebSocket()] + connect_count = 0 + + async def fake_select_account_with_budget( + self, + deadline, + *, + request_id, + kind, + request_stage="first_turn", + sticky_key, + sticky_kind, + reallocate_sticky, + sticky_max_age_seconds, + prefer_earlier_reset_accounts, + routing_strategy, + model, + exclude_account_ids=None, + additional_limit_name=None, + api_key=None, + preferred_account_id=None, + **_kwargs, + ): + del ( + self, + deadline, + request_id, + kind, + request_stage, + sticky_key, + sticky_kind, + reallocate_sticky, + sticky_max_age_seconds, + prefer_earlier_reset_accounts, + routing_strategy, + model, + exclude_account_ids, + additional_limit_name, + api_key, + preferred_account_id, + ) + return AccountSelection(account=account, error_message=None, error_code=None) + + async def fake_ensure_fresh_with_budget(self, target, *, force=False, timeout_seconds): + del self, force, timeout_seconds + return target + + async def fake_connect_responses_websocket( + headers, + access_token, + account_id_header, + *, + base_url=None, + session=None, + ): + del headers, access_token, account_id_header, base_url, session + nonlocal connect_count + upstream = upstreams[connect_count] + connect_count += 1 + return upstream + + monkeypatch.setattr(proxy_module.ProxyService, "_select_account_with_budget", fake_select_account_with_budget) + monkeypatch.setattr(proxy_module.ProxyService, "_ensure_fresh_with_budget", fake_ensure_fresh_with_budget) + monkeypatch.setattr(proxy_module, "connect_responses_websocket", fake_connect_responses_websocket) + service = get_proxy_service_for_app(app_instance) + record_retry_circuit_failure = AsyncMock(wraps=service._record_http_bridge_retry_circuit_failure) + monkeypatch.setattr(service, "_record_http_bridge_retry_circuit_failure", record_retry_circuit_failure) + + response = await async_client.post( + "/v1/responses", + json={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "recover-reader-handoff", + "prompt_cache_key": "reader-handoff-key", + }, + ) + + assert response.status_code == 200 + assert response.json()["output"][0]["content"][0]["text"] == "OK" + assert connect_count == 2 + assert upstreams[0].closed is True + record_retry_circuit_failure.assert_not_awaited() + + @pytest.mark.asyncio async def test_retry_http_bridge_precreated_request_releases_pending_lock_before_reconnect(app_instance, monkeypatch): service = get_proxy_service_for_app(app_instance) diff --git a/tests/integration/test_migrations.py b/tests/integration/test_migrations.py index 2d82cf4d01..53e77d0a6c 100644 --- a/tests/integration/test_migrations.py +++ b/tests/integration/test_migrations.py @@ -1396,3 +1396,70 @@ def _schema_state(sync_conn): assert survivor == 3 finally: await engine.dispose() + + +@pytest.mark.asyncio +async def test_stamped_merge_rollup_repair_downgrade_preserves_schema(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 / 'stamped-merge-rollup-repair.sqlite'}" + merge_revision = "20260724_000000_merge_request_log_schema_heads" + final_revision = "20260728_000000_merge_pending_tool_calls_and_rollup_repair_heads" + rollup_tables = { + "request_usage_hourly_rollups", + "request_usage_hourly_error_rollups", + "request_demand_quarter_rollups", + } + + await to_thread.run_sync(lambda: command.stamp(_build_alembic_config(db_url), merge_revision)) + # The stamped deployment path already has this durable bridge table. Keep + # the fixture representative so the pending-tool-call migration can alter + # it just as it does on a real deployed database. + engine = create_async_engine(db_url, future=True) + try: + async with engine.begin() as conn: + await conn.execute( + text( + """ + CREATE TABLE http_bridge_sessions ( + id VARCHAR(36) PRIMARY KEY, + session_key_kind VARCHAR(64) NOT NULL, + session_key_value TEXT NOT NULL, + session_key_hash VARCHAR(64) NOT NULL, + api_key_scope VARCHAR(255) NOT NULL, + owner_instance_id VARCHAR(255), + owner_epoch INTEGER NOT NULL DEFAULT 0, + lease_expires_at DATETIME, + state VARCHAR(16) NOT NULL DEFAULT 'active', + account_id VARCHAR, + model VARCHAR, + service_tier VARCHAR, + latest_turn_state TEXT, + latest_response_id TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + last_seen_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + closed_at DATETIME + ) + """ + ) + ) + finally: + await engine.dispose() + await to_thread.run_sync(lambda: run_upgrade(db_url, final_revision, bootstrap_legacy=False)) + + engine = create_async_engine(db_url, future=True) + try: + async with engine.connect() as conn: + tables = await conn.run_sync(lambda sync_conn: set(sa_inspect(sync_conn).get_table_names())) + assert rollup_tables <= tables + + await to_thread.run_sync(lambda: command.downgrade(_build_alembic_config(db_url), merge_revision)) + async with engine.connect() as conn: + tables_after_downgrade = await conn.run_sync(lambda sync_conn: set(sa_inspect(sync_conn).get_table_names())) + assert rollup_tables <= tables_after_downgrade + finally: + await engine.dispose() diff --git a/tests/unit/test_bridge_context_blowup.py b/tests/unit/test_bridge_context_blowup.py index eea73d12be..c106a628f2 100644 --- a/tests/unit/test_bridge_context_blowup.py +++ b/tests/unit/test_bridge_context_blowup.py @@ -228,7 +228,14 @@ async def failing_send(*args, **kwargs): f"If this is 400, the bug is present: the CLI will drop " f"previous_response_id and resend full conversation (70K tok/turn)." ) - assert exc_info.value.payload["error"]["code"] in ("upstream_unavailable", "bridge_owner_unreachable") + # A raw send failure is surfaced with the more specific + # ``stream_incomplete`` code; it is still a retriable 502 and keeps + # previous_response_id intact for the client retry. + assert exc_info.value.payload["error"]["code"] in ( + "upstream_unavailable", + "bridge_owner_unreachable", + "stream_incomplete", + ) assert "previous_response_not_found" not in str(exc_info.value.payload) diff --git a/tests/unit/test_bridge_ring_lifecycle.py b/tests/unit/test_bridge_ring_lifecycle.py index a17587b646..c65929d0da 100644 --- a/tests/unit/test_bridge_ring_lifecycle.py +++ b/tests/unit/test_bridge_ring_lifecycle.py @@ -22,6 +22,7 @@ AccountStatus, Base, BridgeRingMember, + HttpBridgeRetryCircuit, HttpBridgeSessionAlias, HttpBridgeSessionRecord, HttpBridgeSessionState, @@ -34,6 +35,7 @@ from app.modules.proxy.durable_bridge_repository import ( DurableBridgeAliasRegistration, DurableBridgeRepository, + durable_bridge_hash, ) from app.modules.proxy.ring_membership import RingMembershipService @@ -306,6 +308,266 @@ async def test_get_sessions_by_ids_chunks_large_id_sets( await session.close() +@pytest.mark.asyncio +async def test_retry_circuit_upsert_counts_concurrent_failure_conflicts( + async_session_factory: Callable[[], AsyncSession], +) -> None: + session = async_session_factory() + try: + repository = DurableBridgeRepository(session) + + for updated_at in (1000.0, 1001.0): + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-conflict", + api_key_scope="key-1", + consecutive_failures=1, + cooldown_until_epoch=0.0, + last_detail="clean_close", + updated_at_epoch=updated_at, + failure_threshold=2, + conflict_cooldown_until_epoch=2000.0, + ) + + row = await session.get( + HttpBridgeRetryCircuit, + ( + "session_header", + durable_bridge_hash("sid-retry-conflict"), + "key-1", + ), + ) + assert row is not None + assert row.consecutive_failures == 2 + assert row.cooldown_until_epoch == 2000.0 + assert row.last_detail == "clean_close" + assert row.updated_at_epoch == 1001.0 + finally: + await session.close() + + +@pytest.mark.asyncio +async def test_retry_circuit_conflict_cooldown_scales_with_merged_failures( + async_session_factory: Callable[[], AsyncSession], +) -> None: + session = async_session_factory() + try: + repository = DurableBridgeRepository(session) + for failures, updated_at in ((1, 1000.0), (1, 1001.0), (2, 1002.0)): + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-backoff-conflict", + api_key_scope="key-1", + consecutive_failures=failures, + cooldown_until_epoch=0.0, + last_detail="stream_incomplete", + updated_at_epoch=updated_at, + failure_threshold=2, + conflict_cooldown_until_epoch=updated_at + 60.0, + ) + + row = await session.get( + HttpBridgeRetryCircuit, + ( + "session_header", + durable_bridge_hash("sid-retry-backoff-conflict"), + "key-1", + ), + ) + assert row is not None + assert row.consecutive_failures == 3 + assert row.cooldown_until_epoch >= 1122.0 + finally: + await session.close() + + +@pytest.mark.asyncio +async def test_retry_circuit_ignores_out_of_order_failure_snapshot( + async_session_factory: Callable[[], AsyncSession], +) -> None: + session = async_session_factory() + try: + repository = DurableBridgeRepository(session) + for failures, updated_at, base_updated_at in ( + (1, 1000.0, 0.0), + (2, 1002.0, 1000.0), + (1, 1001.0, 1002.0), + ): + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-out-of-order", + api_key_scope="key-1", + consecutive_failures=failures, + cooldown_until_epoch=0.0, + last_detail="stream_incomplete", + updated_at_epoch=updated_at, + base_updated_at_epoch=base_updated_at, + failure_threshold=2, + conflict_cooldown_until_epoch=updated_at + 60.0, + ) + + row = await session.get( + HttpBridgeRetryCircuit, + ( + "session_header", + durable_bridge_hash("sid-retry-out-of-order"), + "key-1", + ), + ) + assert row is not None + assert row.consecutive_failures == 2 + assert row.updated_at_epoch == 1002.0 + finally: + await session.close() + + +@pytest.mark.asyncio +async def test_retry_circuit_merges_lagging_wall_clock_failure_from_loaded_base( + async_session_factory: Callable[[], AsyncSession], +) -> None: + session = async_session_factory() + try: + repository = DurableBridgeRepository(session) + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-clock-skew", + api_key_scope="key-1", + consecutive_failures=1, + cooldown_until_epoch=0.0, + last_detail="stream_incomplete", + updated_at_epoch=2000.0, + base_updated_at_epoch=0.0, + failure_threshold=2, + conflict_cooldown_until_epoch=2060.0, + ) + + # This replica loaded the row at 2000, then its wall clock lagged the + # writer that persisted it. The unchanged loaded row is a CAS match, + # so the failure must still open the shared circuit. + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-clock-skew", + api_key_scope="key-1", + consecutive_failures=1, + cooldown_until_epoch=0.0, + last_detail="stream_incomplete", + updated_at_epoch=1500.0, + base_updated_at_epoch=2000.0, + failure_threshold=2, + conflict_cooldown_until_epoch=1560.0, + ) + + row = await session.get( + HttpBridgeRetryCircuit, + ( + "session_header", + durable_bridge_hash("sid-retry-clock-skew"), + "key-1", + ), + ) + assert row is not None + assert row.consecutive_failures == 2 + assert row.cooldown_until_epoch >= 2060.0 + assert row.updated_at_epoch == 2000.0 + finally: + await session.close() + + +@pytest.mark.asyncio +async def test_retry_circuit_reset_starts_new_failure_lineage( + async_session_factory: Callable[[], AsyncSession], +) -> None: + session = async_session_factory() + try: + repository = DurableBridgeRepository(session) + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-reset-lineage", + api_key_scope="key-1", + consecutive_failures=3, + cooldown_until_epoch=2000.0, + last_detail="stream_incomplete", + updated_at_epoch=1000.0, + ) + + await repository.delete_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-reset-lineage", + api_key_scope="key-1", + ) + await repository.upsert_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-reset-lineage", + api_key_scope="key-1", + consecutive_failures=1, + cooldown_until_epoch=5000.0, + last_detail="stream_incomplete", + updated_at_epoch=2000.0, + base_updated_at_epoch=1000.0, + ) + + row = await session.get( + HttpBridgeRetryCircuit, + ( + "session_header", + durable_bridge_hash("sid-retry-reset-lineage"), + "key-1", + ), + ) + assert row is not None + assert row.cooldown_until_epoch == 0.0 + assert row.consecutive_failures == 1 + assert row.last_detail == "stream_incomplete" + assert row.updated_at_epoch == 2000.0 + finally: + await session.close() + + +@pytest.mark.asyncio +async def test_recovery_attempt_pre_dispatch_claim_can_be_rolled_back( + async_session_factory: Callable[[], AsyncSession], +) -> None: + session = async_session_factory() + try: + repository = DurableBridgeRepository(session) + claim = await _claim( + repository, + instance_id="inst-recovery-rollback", + session_key_value="sid-recovery-rollback", + ) + attempt = await repository.record_recovery_attempt( + session_id=claim.id, + instance_id="inst-recovery-rollback", + owner_epoch=claim.owner_epoch, + request_fingerprint="fingerprint-recovery-rollback", + request_id="request-recovery-rollback", + account_id=None, + model="gpt-5.4", + replay_safe=True, + ) + assert attempt is not None + assert await repository.mark_recovery_attempt_replayed( + session_id=claim.id, + instance_id="inst-recovery-rollback", + owner_epoch=claim.owner_epoch, + request_fingerprint="fingerprint-recovery-rollback", + ) + assert await repository.rollback_recovery_attempt_replayed( + session_id=claim.id, + instance_id="inst-recovery-rollback", + owner_epoch=claim.owner_epoch, + request_fingerprint="fingerprint-recovery-rollback", + ) + restored = await repository.lookup_recovery_attempt( + session_id=claim.id, + request_fingerprint="fingerprint-recovery-rollback", + ) + assert restored is not None + assert restored.state.value == "unknown" + finally: + await session.close() + + @pytest.mark.asyncio async def test_ring_purge_removes_dead_members_and_keeps_recent( async_session_factory: Callable[[], AsyncSession], diff --git a/tests/unit/test_db_migrate.py b/tests/unit/test_db_migrate.py index 03a5fc23ed..f33537f08b 100644 --- a/tests/unit/test_db_migrate.py +++ b/tests/unit/test_db_migrate.py @@ -2030,7 +2030,10 @@ def test_capability_lineage_migration_is_additive_reversible_and_single_head(tmp run_upgrade(url, parent_revision, bootstrap_legacy=False) config = _build_alembic_config(url) script_directory = ScriptDirectory.from_config(config) - assert script_directory.get_heads() == [target_revision] + heads = script_directory.get_heads() + assert len(heads) == 1 + ancestry = {script.revision for script in script_directory.walk_revisions()} + assert target_revision in ancestry engine = create_engine(to_sync_database_url(url)) try: diff --git a/tests/unit/test_durable_bridge_sessions.py b/tests/unit/test_durable_bridge_sessions.py index 7d52d9376d..7eb8ec51c2 100644 --- a/tests/unit/test_durable_bridge_sessions.py +++ b/tests/unit/test_durable_bridge_sessions.py @@ -2652,3 +2652,43 @@ async def execute_and_claim_after_candidate_select(statement, *args, **kwargs): select(HttpBridgeSessionAlias).where(HttpBridgeSessionAlias.session_id == "sid-race-claim") ) assert aliases.scalar_one_or_none() is not None + + +@pytest.mark.asyncio +async def test_durable_bridge_retry_circuit_round_trip( + coordinator: DurableBridgeSessionCoordinator, +) -> None: + await coordinator.persist_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-circuit", + api_key_id="key-1", + consecutive_failures=3, + cooldown_until_epoch=1234.5, + last_detail="stream_incomplete", + updated_at_epoch=1200.0, + ) + + persisted = await coordinator.lookup_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-circuit", + api_key_id="key-1", + ) + assert persisted is not None + assert persisted.consecutive_failures == 3 + assert persisted.cooldown_until_epoch == 1234.5 + assert persisted.last_detail == "stream_incomplete" + + await coordinator.clear_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-circuit", + api_key_id="key-1", + ) + cleared = await coordinator.lookup_retry_circuit( + session_key_kind="session_header", + session_key_value="sid-retry-circuit", + api_key_id="key-1", + ) + assert cleared is not None + assert cleared.consecutive_failures == 0 + assert cleared.cooldown_until_epoch == 0.0 + assert cleared.last_detail is None diff --git a/tests/unit/test_http_bridge_safe_continuity.py b/tests/unit/test_http_bridge_safe_continuity.py new file mode 100644 index 0000000000..29ce429c1a --- /dev/null +++ b/tests/unit/test_http_bridge_safe_continuity.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from dataclasses import replace + +from app.modules.proxy import service as proxy_service +from app.modules.proxy._service.http_bridge import streaming as http_bridge_streaming_module + + +def test_http_bridge_continuity_bound_without_safe_replay() -> None: + unsafe_continuation = proxy_service._WebSocketRequestState( + request_id="req-unsafe", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=1.0, + previous_response_id="resp-prev", + ) + safe_full_resend = replace( + unsafe_continuation, + fresh_upstream_request_text='{"type":"response.create","input":"full"}', + fresh_upstream_request_is_retry_safe=True, + ) + hard_turn_state = replace( + unsafe_continuation, + previous_response_id=None, + hard_continuity_anchor=True, + ) + ordinary_request = replace(unsafe_continuation, previous_response_id=None) + + assert http_bridge_streaming_module._http_bridge_continuity_bound_without_safe_replay(unsafe_continuation) + assert not http_bridge_streaming_module._http_bridge_continuity_bound_without_safe_replay(safe_full_resend) + assert http_bridge_streaming_module._http_bridge_continuity_bound_without_safe_replay(hard_turn_state) + assert not http_bridge_streaming_module._http_bridge_continuity_bound_without_safe_replay(ordinary_request) diff --git a/tests/unit/test_proxy_errors.py b/tests/unit/test_proxy_errors.py index f55c22c9f6..083d6b437f 100644 --- a/tests/unit/test_proxy_errors.py +++ b/tests/unit/test_proxy_errors.py @@ -46,6 +46,36 @@ async def stream(): assert message in events[0] +@pytest.mark.asyncio +async def test_stream_proxy_error_preserves_retry_after_as_sse_retry_hint(): + async def stream(): + if False: + yield "" + raise ProxyResponseError( + 503, + { + "error": { + "code": "upstream_request_timeout", + "message": "Retry shortly.", + "type": "server_error", + } + }, + retry_after_seconds=2, + ) + + events = [ + event + async for event in _stream_response_error_events( + stream(), + owns_reservation=False, + reservation=None, + ) + ] + + assert len(events) == 1 + assert events[0].startswith("retry: 2000\n") + + def _payload_error_code(payload) -> str | None: return payload["error"].get("code") diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 00dd232e8e..ea1f88615f 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -39,6 +39,7 @@ from app.modules.proxy._service.http_bridge import helpers as http_bridge_helpers_module from app.modules.proxy._service.http_bridge import mixin as http_bridge_mixin_module from app.modules.proxy._service.http_bridge import request_submit as http_bridge_request_submit_module +from app.modules.proxy._service.http_bridge import retry_circuit as http_bridge_retry_circuit_module from app.modules.proxy._service.http_bridge import streaming as http_bridge_streaming_module from app.modules.proxy._service.http_bridge import upstream_events as http_bridge_upstream_events_module from app.modules.proxy.account_cache import clear_account_routing_unavailable, mark_account_routing_unavailable @@ -46,6 +47,7 @@ is_http_bridge_account_neutral_replay, make_http_bridge_account_neutral_replay_key, ) +from app.modules.proxy.durable_bridge_coordinator import DurableBridgeSessionCoordinator from app.modules.proxy.durable_bridge_repository import ( DurableBridgeAliasRegistration, DurableBridgeAliasRegistrationReceipt, @@ -65,6 +67,17 @@ async def get(self) -> object: monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache()) +@pytest.fixture(autouse=True) +def _stub_recovery_attempt_lookup(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep unit bridge tests independent of the durable recovery schema.""" + + monkeypatch.setattr( + DurableBridgeSessionCoordinator, + "lookup_recovery_attempt", + AsyncMock(return_value=None), + ) + + def _without_installation_metadata(text: str) -> dict[str, Any]: payload = json.loads(text) client_metadata = payload.get("client_metadata") @@ -136,7 +149,7 @@ def test_http_bridge_eventless_precreated_deadline_uses_current_send_and_client_ request_state, stuck_gate_retire_after_seconds=300.0, ) - == 340.0 + == 160.0 ) assert ( http_bridge_helpers_module._http_bridge_eventless_precreated_deadline( @@ -147,12 +160,13 @@ def test_http_bridge_eventless_precreated_deadline_uses_current_send_and_client_ ) request_state.latency_first_upstream_event_ms = 25 + request_state.last_upstream_activity_at = 150.0 assert ( http_bridge_helpers_module._http_bridge_eventless_precreated_deadline( request_state, stuck_gate_retire_after_seconds=300.0, ) - == 340.0 + == 160.0 ) @@ -1262,7 +1276,7 @@ async def test_http_bridge_activity_snapshot_counts_closed_admission_waiter_as_r @pytest.mark.asyncio -async def test_response_create_gate_timeout_retires_session_with_old_pending_visible_request( +async def test_response_create_gate_timeout_retires_old_pending_without_upstream_event( monkeypatch: pytest.MonkeyPatch, ) -> None: settings = _make_app_settings( @@ -1284,7 +1298,9 @@ async def test_response_create_gate_timeout_retires_session_with_old_pending_vis transport="http", response_create_gate_acquired=True, awaiting_response_created=True, - downstream_visible=False, + # A downstream keepalive is not evidence that upstream created a + # response; this request is still safe to retire after the stale window. + downstream_visible=True, ) waiter = proxy_service._WebSocketRequestState( request_id="req-visible-waiter", @@ -1332,6 +1348,110 @@ async def fake_retire( assert waiter.response_create_gate_acquired is False +def test_stale_gate_cleanup_keeps_draining_sibling_active() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + now = time.monotonic() + stale = proxy_service._WebSocketRequestState( + request_id="req-stale-gate-holder", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=now - 301.0, + transport="http", + ) + draining = proxy_service._WebSocketRequestState( + request_id="req-draining-terminal-sibling", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=now - 301.0, + transport="http", + draining_until_terminal=True, + ) + + stale_states, should_retire = service._classify_http_bridge_stale_gate_holders( + [stale, draining], + now=now, + threshold_seconds=300.0, + session_closed=False, + ) + + assert stale_states == [stale] + assert should_retire is False + + +@pytest.mark.asyncio +async def test_response_create_gate_timeout_retires_closed_anchored_pending_without_upstream_event( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_app_settings( + proxy_admission_wait_timeout_seconds=0.001, + http_responses_session_bridge_stuck_gate_retire_after_seconds=300.0, + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) + session = _make_bridge_session() + session.closed = True + service._http_bridge_sessions[session.key] = session + await session.response_create_gate.acquire() + old_pending = proxy_service._WebSocketRequestState( + request_id="req-closed-anchored-pending", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + previous_response_id="resp-anchored", + hard_continuity_anchor=True, + response_create_gate_acquired=True, + awaiting_response_created=True, + ) + waiter = proxy_service._WebSocketRequestState( + request_id="req-closed-anchored-waiter", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + downstream_visible=True, + ) + async with session.pending_lock: + session.pending_requests.append(old_pending) + session.queued_request_count = 1 + + retire_calls: list[str] = [] + + async def fake_retire( + retire_session: proxy_service._HTTPBridgeSession, + *, + detail: str, + ) -> None: + retire_calls.append(detail) + retire_session.closed = True + + monkeypatch.setattr(service, "_retire_stale_pending_http_bridge_session", fake_retire) + + try: + with pytest.raises(ProxyResponseError) as exc_info: + await service._acquire_request_state_response_create_admission( + waiter, + response_create_gate=session.response_create_gate, + account_id=session.account.id, + surface="http_bridge", + bridge_session=session, + ) + finally: + if session.response_create_gate.locked(): + session.response_create_gate.release() + + assert exc_info.value.payload["error"]["code"] == "response_create_gate_timeout" + assert retire_calls == ["response_create_gate_timeout_stuck_pending"] + + @pytest.mark.asyncio async def test_response_create_gate_timeout_retires_old_precreated_request_after_rate_limit_telemetry( monkeypatch: pytest.MonkeyPatch, @@ -1433,7 +1553,6 @@ async def fake_retire( ), [ (False, True, 100, 100, 1), - (True, False, 25, None, 1), ], ) async def test_response_create_gate_timeout_does_not_retire_active_response_progress( @@ -1468,6 +1587,17 @@ async def test_response_create_gate_timeout_does_not_retire_active_response_prog latency_response_created_ms=latency_response_created_ms, response_event_count=response_event_count, ) + stale_pending = proxy_service._WebSocketRequestState( + request_id="req-stale-unanchored", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + response_create_gate_acquired=True, + awaiting_response_created=True, + ) waiter = proxy_service._WebSocketRequestState( request_id="req-visible-waiter", model="gpt-5.2", @@ -1480,7 +1610,8 @@ async def test_response_create_gate_timeout_does_not_retire_active_response_prog ) async with session.pending_lock: session.pending_requests.append(active_stream) - session.queued_request_count = 1 + session.pending_requests.append(stale_pending) + session.queued_request_count = 2 retire_calls: list[str] = [] @@ -1512,99 +1643,352 @@ async def fake_retire( assert session.closed is False -@pytest.mark.asyncio -async def test_http_bridge_activity_snapshot_counts_only_bridge_cleanup_tasks(): - service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) - bridge_task = asyncio.create_task(asyncio.sleep(60), name="proxy-http_bridge_session_close-test") - api_key_task = asyncio.create_task(asyncio.sleep(60), name="proxy-stream-api-key-settle-test") - service._background_cleanup_tasks.update({bridge_task, api_key_task}) - - try: - snapshot = service.http_bridge_activity_snapshot_nowait() - finally: - bridge_task.cancel() - api_key_task.cancel() - await asyncio.gather(bridge_task, api_key_task, return_exceptions=True) - - assert snapshot["http_bridge_background_cleanup_tasks"] == 1 +def test_http_bridge_pending_state_with_recent_events_but_no_created_is_not_stale() -> None: + # Reattached streams can deliver events whose response.created was lost + # (observed events=54, created=None in prod on 2026-07-20). Recent upstream + # activity must keep the stream alive while the create gate remains held. + request_state = proxy_service._WebSocketRequestState( + request_id="req-events-no-created", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + response_create_gate_acquired=True, + awaiting_response_created=True, + latency_first_upstream_event_ms=25, + response_event_count=54, + last_upstream_activity_at=time.monotonic(), + ) + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + ) + is False + ) -@pytest.mark.asyncio -async def test_http_bridge_activity_snapshot_cleans_completed_stale_inflight_session( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, -) -> None: - service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) - key = proxy_service._HTTPBridgeSessionKey("session_header", "stale-inflight-drain-status", None) - inflight_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() - setattr(inflight_future, "_codex_lb_started_at", -1000.0) - inflight_future.set_result(_make_bridge_session()) - service._http_bridge_inflight_sessions[key] = inflight_future + request_state.last_upstream_activity_at = time.monotonic() - 301.0 + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + ) + is True + ) - monkeypatch.setattr(proxy_service, "_proxy_admission_wait_timeout_seconds", lambda settings=None: 0.001) - with caplog.at_level(logging.WARNING, logger="app.modules.proxy.service"): - snapshot = service.http_bridge_activity_snapshot_nowait() +def test_http_bridge_pending_state_with_first_event_latency_only_is_stale() -> None: + request_state = proxy_service._WebSocketRequestState( + request_id="req-sparse-active", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + response_create_gate_acquired=True, + awaiting_response_created=True, + latency_first_upstream_event_ms=10, + downstream_visible=False, + ) - assert key not in service._http_bridge_inflight_sessions - assert snapshot["http_bridge_inflight_session_creates"] == 0 - assert snapshot["http_bridge_stale_inflight_session_creates"] == 1 - assert snapshot["http_bridge_cleaned_inflight_session_creates"] == 1 - assert snapshot["http_bridge_active"] is False - assert snapshot["http_bridge_restart_blocking"] is False - assert "http_bridge_inflight_session_create_cleanup" in caplog.text + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + ) + is True + ) -@pytest.mark.asyncio -async def test_http_bridge_activity_snapshot_does_not_expire_live_inflight_session( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize( + ("previous_response_id", "session_id", "hard_continuity_anchor"), + [ + ("resp-anchored", None, False), + (None, "turn-anchored", True), + ], +) +def test_http_bridge_pending_state_with_continuity_anchor_is_not_stale( + previous_response_id: str | None, + session_id: str | None, + hard_continuity_anchor: bool, ) -> None: - service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) - key = proxy_service._HTTPBridgeSessionKey("session_header", "live-stale-inflight-drain-status", None) - inflight_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() - setattr(inflight_future, "_codex_lb_started_at", -1000.0) - service._http_bridge_inflight_sessions[key] = inflight_future - - monkeypatch.setattr(proxy_service, "_proxy_admission_wait_timeout_seconds", lambda settings=None: 0.001) - - snapshot = service.http_bridge_activity_snapshot_nowait() + request_state = proxy_service._WebSocketRequestState( + request_id="req-anchored-pending", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + response_create_gate_acquired=True, + awaiting_response_created=True, + previous_response_id=previous_response_id, + session_id=session_id, + hard_continuity_anchor=hard_continuity_anchor, + ) - assert key in service._http_bridge_inflight_sessions - assert not inflight_future.done() - assert snapshot["http_bridge_inflight_session_creates"] == 1 - assert snapshot["http_bridge_stale_inflight_session_creates"] == 1 - assert snapshot["http_bridge_cleaned_inflight_session_creates"] == 0 - assert snapshot["http_bridge_active"] is True - assert snapshot["http_bridge_restart_blocking"] is True + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + ) + is False + ) -@pytest.mark.asyncio -async def test_http_bridge_activity_snapshot_skips_inflight_cleanup_when_registry_locked( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize( + ("previous_response_id", "session_id", "hard_continuity_anchor"), + [ + ("resp-anchored-silent", None, False), + (None, "turn-anchored-silent", True), + ], +) +def test_http_bridge_pending_state_with_continuity_anchor_is_stale_after_extended_silence( + previous_response_id: str | None, + session_id: str | None, + hard_continuity_anchor: bool, ) -> None: - service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) - key = proxy_service._HTTPBridgeSessionKey("session_header", "locked-stale-inflight-drain-status", None) - inflight_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() - setattr(inflight_future, "_codex_lb_started_at", -1000.0) - service._http_bridge_inflight_sessions[key] = inflight_future - - monkeypatch.setattr(proxy_service, "_proxy_admission_wait_timeout_seconds", lambda settings=None: 0.001) - - async with service._http_bridge_lock: - snapshot = service.http_bridge_activity_snapshot_nowait() + request_state = proxy_service._WebSocketRequestState( + request_id="req-anchored-silent-pending", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 601.0, + transport="http", + response_create_gate_acquired=True, + awaiting_response_created=True, + previous_response_id=previous_response_id, + session_id=session_id, + hard_continuity_anchor=hard_continuity_anchor, + ) - assert key in service._http_bridge_inflight_sessions - assert not inflight_future.done() - assert snapshot["http_bridge_inflight_session_creates"] == 1 - assert snapshot["http_bridge_stale_inflight_session_creates"] == 1 - assert snapshot["http_bridge_cleaned_inflight_session_creates"] == 0 - assert snapshot["http_bridge_active"] is True - assert snapshot["http_bridge_restart_blocking"] is True + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + ) + is True + ) -async def _wait_for_close_await(close_session: AsyncMock, session: proxy_service._HTTPBridgeSession) -> None: - for _ in range(10): - if any(call.args == (session,) for call in close_session.await_args_list): +@pytest.mark.parametrize( + ("previous_response_id", "session_id", "hard_continuity_anchor"), + [ + ("resp-closed-anchored", None, False), + (None, "turn-closed-anchored", True), + ], +) +def test_http_bridge_closed_session_pending_anchor_is_stale( + previous_response_id: str | None, + session_id: str | None, + hard_continuity_anchor: bool, +) -> None: + request_state = proxy_service._WebSocketRequestState( + request_id="req-closed-anchored-pending", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + response_create_gate_acquired=True, + awaiting_response_created=True, + previous_response_id=previous_response_id, + session_id=session_id, + hard_continuity_anchor=hard_continuity_anchor, + ) + + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + session_closed=True, + ) + is True + ) + + +def test_http_bridge_pending_state_with_plain_session_header_is_stale() -> None: + request_state = proxy_service._WebSocketRequestState( + request_id="req-session-header-pending", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + session_id="session-header-only", + response_create_gate_acquired=True, + awaiting_response_created=True, + ) + + assert ( + http_bridge_helpers_module._http_bridge_pending_state_is_stale( + request_state, + now=time.monotonic(), + threshold_seconds=300.0, + ) + is True + ) + + +def test_http_bridge_synthesized_downstream_turn_state_is_not_hard_anchor() -> None: + request_state = proxy_service._WebSocketRequestState( + request_id="req-synth-turn-state", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + ) + + http_bridge_streaming_module._apply_http_bridge_downstream_turn_state( + request_state, + downstream_turn_state="synthesized-turn-state", + incoming_turn_state_header=None, + ) + + assert request_state.session_id == "synthesized-turn-state" + assert request_state.hard_continuity_anchor is False + + +@pytest.mark.parametrize( + ("incoming_turn_state_header", "previous_response_id"), + [ + ("client-turn-state", None), + (None, "resp-continuation"), + ], +) +def test_http_bridge_real_continuity_sets_hard_anchor( + incoming_turn_state_header: str | None, + previous_response_id: str | None, +) -> None: + request_state = proxy_service._WebSocketRequestState( + request_id="req-real-anchor", + model="gpt-5.2", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + previous_response_id=previous_response_id, + ) + + http_bridge_streaming_module._apply_http_bridge_downstream_turn_state( + request_state, + downstream_turn_state="real-turn-state", + incoming_turn_state_header=incoming_turn_state_header, + ) + + assert request_state.session_id == "real-turn-state" + assert request_state.hard_continuity_anchor is True + + +@pytest.mark.asyncio +async def test_http_bridge_activity_snapshot_counts_only_bridge_cleanup_tasks(): + service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) + bridge_task = asyncio.create_task(asyncio.sleep(60), name="proxy-http_bridge_session_close-test") + api_key_task = asyncio.create_task(asyncio.sleep(60), name="proxy-stream-api-key-settle-test") + service._background_cleanup_tasks.update({bridge_task, api_key_task}) + + try: + snapshot = service.http_bridge_activity_snapshot_nowait() + finally: + bridge_task.cancel() + api_key_task.cancel() + await asyncio.gather(bridge_task, api_key_task, return_exceptions=True) + + assert snapshot["http_bridge_background_cleanup_tasks"] == 1 + + +@pytest.mark.asyncio +async def test_http_bridge_activity_snapshot_cleans_completed_stale_inflight_session( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) + key = proxy_service._HTTPBridgeSessionKey("session_header", "stale-inflight-drain-status", None) + inflight_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() + setattr(inflight_future, "_codex_lb_started_at", -1000.0) + inflight_future.set_result(_make_bridge_session()) + service._http_bridge_inflight_sessions[key] = inflight_future + + monkeypatch.setattr(proxy_service, "_proxy_admission_wait_timeout_seconds", lambda settings=None: 0.001) + + with caplog.at_level(logging.WARNING, logger="app.modules.proxy.service"): + snapshot = service.http_bridge_activity_snapshot_nowait() + + assert key not in service._http_bridge_inflight_sessions + assert snapshot["http_bridge_inflight_session_creates"] == 0 + assert snapshot["http_bridge_stale_inflight_session_creates"] == 1 + assert snapshot["http_bridge_cleaned_inflight_session_creates"] == 1 + assert snapshot["http_bridge_active"] is False + assert snapshot["http_bridge_restart_blocking"] is False + assert "http_bridge_inflight_session_create_cleanup" in caplog.text + + +@pytest.mark.asyncio +async def test_http_bridge_activity_snapshot_does_not_expire_live_inflight_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) + key = proxy_service._HTTPBridgeSessionKey("session_header", "live-stale-inflight-drain-status", None) + inflight_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() + setattr(inflight_future, "_codex_lb_started_at", -1000.0) + service._http_bridge_inflight_sessions[key] = inflight_future + + monkeypatch.setattr(proxy_service, "_proxy_admission_wait_timeout_seconds", lambda settings=None: 0.001) + + snapshot = service.http_bridge_activity_snapshot_nowait() + + assert key in service._http_bridge_inflight_sessions + assert not inflight_future.done() + assert snapshot["http_bridge_inflight_session_creates"] == 1 + assert snapshot["http_bridge_stale_inflight_session_creates"] == 1 + assert snapshot["http_bridge_cleaned_inflight_session_creates"] == 0 + assert snapshot["http_bridge_active"] is True + assert snapshot["http_bridge_restart_blocking"] is True + + +@pytest.mark.asyncio +async def test_http_bridge_activity_snapshot_skips_inflight_cleanup_when_registry_locked( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) + key = proxy_service._HTTPBridgeSessionKey("session_header", "locked-stale-inflight-drain-status", None) + inflight_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() + setattr(inflight_future, "_codex_lb_started_at", -1000.0) + service._http_bridge_inflight_sessions[key] = inflight_future + + monkeypatch.setattr(proxy_service, "_proxy_admission_wait_timeout_seconds", lambda settings=None: 0.001) + + async with service._http_bridge_lock: + snapshot = service.http_bridge_activity_snapshot_nowait() + + assert key in service._http_bridge_inflight_sessions + assert not inflight_future.done() + assert snapshot["http_bridge_inflight_session_creates"] == 1 + assert snapshot["http_bridge_stale_inflight_session_creates"] == 1 + assert snapshot["http_bridge_cleaned_inflight_session_creates"] == 0 + assert snapshot["http_bridge_active"] is True + assert snapshot["http_bridge_restart_blocking"] is True + + +async def _wait_for_close_await(close_session: AsyncMock, session: proxy_service._HTTPBridgeSession) -> None: + for _ in range(10): + if any(call.args == (session,) for call in close_session.await_args_list): return await asyncio.sleep(0) raise AssertionError("expected HTTP bridge session close to be awaited") @@ -3226,6 +3610,11 @@ async def test_http_bridge_precreated_completed_terminal_falls_back_to_unresolve monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) + monkeypatch.setattr( + http_bridge_retry_circuit_module, + "_HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD", + 1, + ) finalize = AsyncMock() register_previous = AsyncMock() monkeypatch.setattr(service, "_finalize_websocket_request_state", finalize) @@ -3241,13 +3630,15 @@ async def test_http_bridge_precreated_completed_terminal_falls_back_to_unresolve awaiting_response_created=True, event_queue=asyncio.Queue(), transport="http", - skip_request_log=True, ) session = _make_bridge_session( key_value="bridge-precreated-completed", pending_requests=deque([request_state]), queued_request_count=1, ) + await service._record_http_bridge_retry_circuit_failure(session, detail="stream_incomplete") + retry_circuits = cast(Any, service)._http_bridge_retry_circuits + assert session.key in retry_circuits await service._process_http_bridge_upstream_text( session, @@ -3296,6 +3687,7 @@ async def test_http_bridge_precreated_completed_terminal_falls_back_to_unresolve assert session.last_completed_response_id == "resp_precreated_completed" assert session.queued_request_count == 0 assert not session.pending_requests + assert session.key not in retry_circuits register_previous.assert_awaited_once() finalize.assert_awaited_once() @@ -3756,6 +4148,93 @@ async def fake_submit_http_bridge_request( assert error["code"] == "stream_incomplete" +@pytest.mark.asyncio +async def test_http_bridge_startup_cooldown_releases_api_key_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session(key_value="sid-startup-reservation") + reservation = cast(Any, object()) + request_state = proxy_service._WebSocketRequestState( + request_id="req-startup-reservation", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=reservation, + started_at=time.monotonic(), + event_queue=asyncio.Queue(), + transport="http", + previous_response_id="resp-anchor", + ) + cooldown = AsyncMock(return_value=30.0) + release = AsyncMock() + monkeypatch.setattr(service, "_http_bridge_precreated_retry_cooldown_seconds", cooldown) + monkeypatch.setattr(service, "_release_websocket_request_state_reservation", release) + + events = [ + event + async for event in service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data='{"type":"response.create"}', + queue_limit=8, + propagate_http_errors=False, + downstream_turn_state=None, + ) + ] + + assert len(events) == 1 + assert '"code":"stream_idle_timeout"' in events[0] + release.assert_awaited_once_with(request_state) + assert request_state.api_key_reservation is None + + +@pytest.mark.asyncio +async def test_http_bridge_post_submit_cooldown_race_detaches_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session(key_value="sid-post-submit-cooldown") + request_state = proxy_service._WebSocketRequestState( + request_id="req-post-submit-cooldown", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + event_queue=asyncio.Queue(), + transport="http", + previous_response_id="resp-anchor", + ) + + async def submit(target_session: Any, *, request_state: Any, **kwargs: Any) -> None: + del kwargs + target_session.pending_requests.append(request_state) + + cooldown = AsyncMock(side_effect=[0.0, 30.0]) + detach = AsyncMock() + monkeypatch.setattr(service, "_submit_http_bridge_request", submit) + monkeypatch.setattr(service, "_http_bridge_precreated_retry_cooldown_seconds", cooldown) + monkeypatch.setattr(service, "_detach_http_bridge_request", detach) + + events = [ + event + async for event in service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data='{"type":"response.create"}', + queue_limit=8, + propagate_http_errors=False, + downstream_turn_state=None, + ) + ] + + assert len(events) == 1 + assert '"code":"stream_idle_timeout"' in events[0] + assert cooldown.await_count == 2 + detach.assert_awaited_once_with(session, request_state=request_state) + + @pytest.mark.asyncio async def test_http_bridge_keepalive_counts_as_first_yield_before_late_response_failed( monkeypatch: pytest.MonkeyPatch, @@ -3927,37 +4406,36 @@ async def fake_submit_http_bridge_request( @pytest.mark.asyncio -async def test_http_bridge_capacity_wait_with_response_id_sends_explicit_keepalive( +async def test_http_bridge_idle_recovery_transport_failure_yields_terminal_event( monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) - monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + detach = AsyncMock() + monkeypatch.setattr(service, "_detach_http_bridge_request", detach) monkeypatch.setattr( proxy_service, "get_settings", lambda: SimpleNamespace( + http_responses_stream_request_budget_seconds=60.0, sse_keepalive_interval_seconds=0.001, stream_idle_timeout_seconds=0.001, ), ) monkeypatch.setattr(proxy_service, "_HTTP_BRIDGE_STARTUP_KEEPALIVE_GRACE_SECONDS", 0.001) + monkeypatch.setattr(http_bridge_streaming_module, "_stream_keepalive_max_count", lambda: 1) - session = _make_bridge_session(key_value="sid-capacity-response") + session = _make_bridge_session(key_value="sid-idle-retry-transport") request_state = proxy_service._WebSocketRequestState( - request_id="req-capacity-response", + request_id="req-idle-retry-transport", model="gpt-5.1", service_tier=None, reasoning_effort=None, api_key_reservation=None, started_at=time.monotonic(), event_queue=asyncio.Queue(), - response_id="resp-capacity-response", + request_text='{"type":"response.create","model":"gpt-5.1","input":"hello"}', transport="http", ) - request_state.account_capacity_waiting = True - request_state.account_capacity_wait_reason = "Rate limit exceeded. Try again in 120s" - request_state.account_capacity_wait_started_at = time.monotonic() - 3.0 - request_state.account_capacity_wait_retry_after_seconds = 120.0 async def fake_submit_http_bridge_request( target_session: proxy_service._HTTPBridgeSession, @@ -3969,13 +4447,90 @@ async def fake_submit_http_bridge_request( del text_data, queue_limit target_session.pending_requests.append(request_state) + retry_error = UpstreamWebSocketTransportError( + "Codex upstream websocket send failed: OSError", + error_code="proxy_network_unavailable", + ) + retry_precreated = AsyncMock(side_effect=retry_error) monkeypatch.setattr(service, "_submit_http_bridge_request", fake_submit_http_bridge_request) + monkeypatch.setattr(service, "_retry_http_bridge_precreated_request", retry_precreated) - stream = service._stream_http_bridge_session_events( - session, - request_state=request_state, - text_data="{}", - queue_limit=8, + chunks = [ + chunk + async for chunk in service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data="{}", + queue_limit=8, + propagate_http_errors=True, + downstream_turn_state=None, + ) + ] + + assert len(chunks) == 1 + payload = proxy_service.parse_sse_data_json(chunks[0]) + assert payload is not None + assert payload["type"] == "response.failed" + response = payload["response"] + assert isinstance(response, dict) + error = response["error"] + assert isinstance(error, dict) + assert error["code"] == "proxy_network_unavailable" + assert error["message"] == "Codex upstream websocket send failed: OSError" + retry_precreated.assert_awaited_once_with(session, restart_reader=True) + detach.assert_awaited_once_with(session, request_state=request_state) + + +@pytest.mark.asyncio +async def test_http_bridge_capacity_wait_with_response_id_sends_explicit_keepalive( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: SimpleNamespace( + sse_keepalive_interval_seconds=0.001, + stream_idle_timeout_seconds=0.001, + ), + ) + monkeypatch.setattr(proxy_service, "_HTTP_BRIDGE_STARTUP_KEEPALIVE_GRACE_SECONDS", 0.001) + + session = _make_bridge_session(key_value="sid-capacity-response") + request_state = proxy_service._WebSocketRequestState( + request_id="req-capacity-response", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + event_queue=asyncio.Queue(), + response_id="resp-capacity-response", + transport="http", + ) + request_state.account_capacity_waiting = True + request_state.account_capacity_wait_reason = "Rate limit exceeded. Try again in 120s" + request_state.account_capacity_wait_started_at = time.monotonic() - 3.0 + request_state.account_capacity_wait_retry_after_seconds = 120.0 + + async def fake_submit_http_bridge_request( + target_session: proxy_service._HTTPBridgeSession, + *, + request_state: proxy_service._WebSocketRequestState, + text_data: str, + queue_limit: int, + ) -> None: + del text_data, queue_limit + target_session.pending_requests.append(request_state) + + monkeypatch.setattr(service, "_submit_http_bridge_request", fake_submit_http_bridge_request) + + stream = service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data="{}", + queue_limit=8, propagate_http_errors=True, downstream_turn_state=None, ) @@ -4085,7 +4640,7 @@ async def test_get_or_create_http_bridge_session_preserves_closed_admission_hand @pytest.mark.asyncio -async def test_get_or_create_http_bridge_session_rejects_incompatible_closed_admission_handoff( +async def test_get_or_create_http_bridge_session_rejects_anchored_incompatible_closed_admission_handoff( monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) @@ -4109,6 +4664,7 @@ async def test_get_or_create_http_bridge_session_rejects_incompatible_closed_adm idle_ttl_seconds=120.0, max_sessions=8, preferred_account_id="different-account", + previous_response_id="resp-anchored", ) assert exc_info.value.status_code == 503 @@ -4116,6 +4672,54 @@ async def test_get_or_create_http_bridge_session_rejects_incompatible_closed_adm create.assert_not_awaited() +@pytest.mark.asyncio +async def test_get_or_create_http_bridge_session_recovers_unanchored_closed_admission_handoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + key = proxy_service._HTTPBridgeSessionKey("session_header", "bridge-handoff", None) + existing = _make_bridge_session(key_value="bridge-handoff") + existing.key = key + existing.closed = True + existing.admission_waiter_count = 1 + service._http_bridge_sessions[key] = existing + replacement = _make_bridge_session(key_value="bridge-handoff") + replacement.key = key + settings = _make_app_settings() + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + proxy_service, + "_http_bridge_owner_instance", + AsyncMock(return_value=settings.http_responses_session_bridge_instance_id), + ) + monkeypatch.setattr( + http_bridge_mixin_module, + "_http_bridge_owner_instance", + AsyncMock(return_value=settings.http_responses_session_bridge_instance_id), + ) + monkeypatch.setattr(http_bridge_mixin_module, "_http_bridge_owner_check_required", lambda *args, **kwargs: False) + monkeypatch.setattr(service, "_claim_durable_http_bridge_session", AsyncMock()) + monkeypatch.setattr(service, "_schedule_http_bridge_session_closes", Mock()) + create = AsyncMock(return_value=replacement) + monkeypatch.setattr(service, "_create_http_bridge_session", create) + + resolved = await service._get_or_create_http_bridge_session( + key, + headers={"x-codex-session-id": "bridge-handoff"}, + affinity=proxy_service._AffinityPolicy(key="bridge-handoff"), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=8, + preferred_account_id="different-account", + ) + + assert resolved is replacement + assert service._http_bridge_sessions[key] is replacement + assert existing.closed is True + create.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_or_create_http_bridge_session_replaces_routing_unavailable_account( monkeypatch: pytest.MonkeyPatch, @@ -6005,6 +6609,156 @@ async def open_upstream(_account: object, headers: dict[str, str], **_: object) assert "x-handshake-debug" not in forwarded +@pytest.mark.asyncio +async def test_reconnect_keeps_handoff_protected_during_lease_swap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + old_lease = proxy_service.AccountLease( + lease_id="lease-old-handoff", + account_id=session.account.id, + kind="stream", + acquired_at=1.0, + ) + new_account = cast(Any, SimpleNamespace(id="acc-replacement", status=AccountStatus.ACTIVE, plan_type="plus")) + new_lease = proxy_service.AccountLease( + lease_id="lease-new-handoff", + account_id=new_account.id, + kind="stream", + acquired_at=2.0, + ) + session.account_lease = old_lease + request_state = proxy_service._WebSocketRequestState( + request_id="req-handoff-lease-swap", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + ) + replacement = cast( + UpstreamWebSocket, + SimpleNamespace(response_header=lambda _name: None, close=AsyncMock()), + ) + release_account_lease = AsyncMock() + + async def release_lease(lease: proxy_service.AccountLease | None) -> None: + assert lease is old_lease + assert session.closed is True + assert session.handoff_in_progress is True + await release_account_lease(lease) + + async def select_account(_deadline: float, **_: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection(account=new_account, error_message=None, lease=new_lease) + + async def ensure_fresh(account: object, **_: object) -> object: + return account + + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + prefer_earlier_reset_accounts=False, + routing_strategy=None, + ) + ) + ), + ) + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", ensure_fresh) + monkeypatch.setattr(service, "_open_upstream_websocket_with_budget", AsyncMock(return_value=replacement)) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_lease) + + await service._reconnect_http_bridge_session(session, request_state=request_state) + + release_account_lease.assert_awaited_once_with(old_lease) + assert session.account is new_account + assert session.account_lease is new_lease + assert session.closed is False + assert session.handoff_in_progress is False + assert session.handoff_future is None + assert session.key not in service._http_bridge_inflight_sessions + + +@pytest.mark.asyncio +async def test_reconnect_cancellation_during_wrong_owner_lease_release_completes_handoff( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + required_account = cast(Any, SimpleNamespace(id="acc-required", status=AccountStatus.ACTIVE, plan_type="plus")) + replacement_account = cast( + Any, + SimpleNamespace(id="acc-replacement", status=AccountStatus.ACTIVE, plan_type="plus"), + ) + replacement_lease = proxy_service.AccountLease( + lease_id="lease-wrong-owner-cancelled", + account_id=replacement_account.id, + kind="stream", + acquired_at=2.0, + ) + release_started = asyncio.Event() + + async def select_account(_deadline: float, **_: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection( + account=replacement_account, + error_message=None, + error_code=None, + lease=replacement_lease, + ) + + async def release_lease(_lease: proxy_service.AccountLease | None) -> None: + release_started.set() + await asyncio.Event().wait() + + request_state = proxy_service._WebSocketRequestState( + request_id="req-wrong-owner-cancelled", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + preferred_account_id=required_account.id, + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + prefer_earlier_reset_accounts=False, + routing_strategy=None, + ) + ) + ), + ) + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_account) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_lease) + + reconnect_task = asyncio.create_task( + service._reconnect_http_bridge_session( + session, + request_state=request_state, + require_preferred_account=True, + ) + ) + await asyncio.wait_for(release_started.wait(), timeout=1.0) + reconnect_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await reconnect_task + + assert session.closed is True + assert session.handoff_in_progress is False + assert session.handoff_future is None + assert session.key not in service._http_bridge_inflight_sessions + + @pytest.mark.asyncio async def test_reconnect_http_bridge_session_preserves_hard_account_after_1011( monkeypatch: pytest.MonkeyPatch, @@ -7388,6 +8142,7 @@ async def test_stream_via_http_bridge_preserves_only_safe_trimmable_full_resend_ prepared_previous_response_ids: list[str | None] = [] prepared_input_lengths: list[int] = [] prepared_frames: list[dict[str, Any]] = [] + prepare_call_count = 0 real_prepare = service._prepare_http_bridge_request def fake_prepare( @@ -7400,9 +8155,17 @@ def fake_prepare( client_ip: str | None = None, **kwargs: Any, ) -> tuple[proxy_service._WebSocketRequestState, str]: - prepared_previous_response_ids.append(prepared_payload.previous_response_id) + # The recovery journal fingerprint is prepared from the same payload + # before the request is sent. It is internal bookkeeping, not a + # second upstream dispatch, so keep it out of dispatch assertions. + nonlocal prepare_call_count + prepare_call_count += 1 + record_dispatch = not (preserves_full_resend and prepare_call_count == 1) + if record_dispatch: + prepared_previous_response_ids.append(prepared_payload.previous_response_id) inp = prepared_payload.input - prepared_input_lengths.append(len(inp) if isinstance(inp, list) else 1) + if record_dispatch: + prepared_input_lengths.append(len(inp) if isinstance(inp, list) else 1) _, text_data = real_prepare( prepared_payload, _headers, @@ -7412,7 +8175,8 @@ def fake_prepare( client_ip=client_ip, **kwargs, ) - prepared_frames.append(json.loads(text_data)) + if record_dispatch: + prepared_frames.append(json.loads(text_data)) request_state.previous_response_id = prepared_payload.previous_response_id return request_state, text_data @@ -7551,7 +8315,10 @@ def fake_prepare( if not preserves_full_resend: assert request_state.proxy_injected_previous_response_id is True assert request_state.fresh_upstream_request_is_retry_safe is False - account_neutral_classifier.assert_not_called() + if preserves_full_resend: + account_neutral_classifier.assert_called_once() + else: + account_neutral_classifier.assert_not_called() @pytest.mark.asyncio @@ -8110,6 +8877,88 @@ async def close_http_bridge_session(target: proxy_service._HTTPBridgeSession) -> assert close_cancelled is False +@pytest.mark.asyncio +async def test_await_cancelled_task_consumes_child_cancellation() -> None: + child = asyncio.create_task(asyncio.sleep(60)) + + assert await proxy_service._await_cancelled_task(child, timeout_seconds=1.0, label="test child") is True + assert child.done() + assert child.cancelled() + + +@pytest.mark.asyncio +async def test_await_cancelled_task_propagates_caller_cancellation() -> None: + child_cancelled = asyncio.Event() + release_child = asyncio.Event() + + async def stubborn_child() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + child_cancelled.set() + await release_child.wait() + + child = asyncio.create_task(stubborn_child()) + cleanup_tasks: set[asyncio.Task[None]] = set() + waiter = asyncio.create_task( + proxy_service._await_cancelled_task( + child, + timeout_seconds=10.0, + label="stubborn test child", + cleanup_tasks=cleanup_tasks, + ) + ) + await asyncio.wait_for(child_cancelled.wait(), timeout=1.0) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert cleanup_tasks + + release_child.set() + await asyncio.wait_for(child, timeout=1.0) + for _ in range(10): + if not cleanup_tasks: + break + await asyncio.sleep(0) + assert cleanup_tasks == set() + + +@pytest.mark.asyncio +async def test_await_cancelled_task_defers_stubborn_child_cleanup() -> None: + child_cancelled = asyncio.Event() + release_child = asyncio.Event() + + async def stubborn_child() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + child_cancelled.set() + await release_child.wait() + + child = asyncio.create_task(stubborn_child()) + cleanup_tasks: set[asyncio.Task[None]] = set() + assert ( + await proxy_service._await_cancelled_task( + child, + timeout_seconds=0.001, + label="stubborn cleanup", + cleanup_tasks=cleanup_tasks, + ) + is False + ) + assert cleanup_tasks + await asyncio.wait_for(child_cancelled.wait(), timeout=1.0) + + release_child.set() + await asyncio.wait_for(child, timeout=1.0) + for _ in range(10): + if not cleanup_tasks: + break + await asyncio.sleep(0) + assert cleanup_tasks == set() + + @pytest.mark.asyncio async def test_close_http_bridge_session_bounded_cancellation_keeps_close_task_tracked( monkeypatch: pytest.MonkeyPatch, @@ -8770,6 +9619,7 @@ async def test_stream_via_http_bridge_preserves_context_after_owner_unavailable( request_states: list[proxy_service._WebSocketRequestState] = [] prepared_previous_response_ids: list[str | None] = [] prepared_inputs: list[proxy_service.JsonValue] = [] + prepare_call_count = 0 def fake_prepare( prepared_payload: proxy_service.ResponsesRequest, @@ -8781,10 +9631,14 @@ def fake_prepare( client_ip: str | None = None, ) -> tuple[proxy_service._WebSocketRequestState, str]: del api_key, api_key_reservation, request_id, client_ip - prepared_previous_response_ids.append(prepared_payload.previous_response_id) - prepared_inputs.append(prepared_payload.input) + nonlocal prepare_call_count + prepare_call_count += 1 + record_dispatch = not (retains_prior_output and prepare_call_count == 1) + if record_dispatch: + prepared_previous_response_ids.append(prepared_payload.previous_response_id) + prepared_inputs.append(prepared_payload.input) state = proxy_service._WebSocketRequestState( - request_id=f"req-{len(request_states)}", + request_id=f"req-{prepare_call_count}", model="gpt-5.4", service_tier=None, reasoning_effort=None, @@ -8794,7 +9648,8 @@ def fake_prepare( previous_response_id=prepared_payload.previous_response_id, transport="http", ) - request_states.append(state) + if record_dispatch: + request_states.append(state) return state, proxy_service._response_create_text( prepared_payload, include_type_field=True, @@ -10653,6 +11508,7 @@ async def produce_after_reattach_delay() -> None: monkeypatch.setattr(http_bridge_streaming_module, "_http_bridge_account_capacity_wait_seconds", lambda _exc: 0.001) monkeypatch.setattr(http_bridge_streaming_module, "_ACCOUNT_SELECTION_RECOVERY_HEARTBEAT_SECONDS", 0.001) monkeypatch.setattr(http_bridge_streaming_module, "_http_bridge_startup_keepalive_grace_seconds", lambda: 0.001) + monkeypatch.setattr(proxy_service, "_STREAM_KEEPALIVE_MAX_COUNT", 20) chunks = [ chunk @@ -10675,7 +11531,7 @@ async def produce_after_reattach_delay() -> None: keepalive = proxy_service.parse_sse_data_json(chunks[0]) assert keepalive is not None assert keepalive["status"] == "waiting_for_account_capacity" - assert http_bridge_streaming_module._codex_keepalive_frame() in chunks + assert any('"type":"codex.keepalive"' in chunk for chunk 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] @@ -11369,7 +12225,7 @@ async def fake_stream_http_bridge_session_events( assert exc_info.value.status_code == 400 assert exc_info.value.payload["error"]["code"] == "context_length_exceeded" assert key not in service._http_bridge_sessions - close_session.assert_awaited_once_with(session) + close_session.assert_awaited_once_with(session, release_durable_session=True) assert isinstance(failed_block, str) assert '"type":"response.failed"' in failed_block assert '"code":"stream_incomplete"' in failed_block @@ -14462,10 +15318,10 @@ async def ambiguous_send(_text: str) -> None: queue_limit=2, ) ) - for _ in range(20): + for _ in range(100): if session.admission_waiter_count == 1: break - await asyncio.sleep(0) + await asyncio.sleep(0.001) assert session.admission_waiter_count == 1 prewarm_task.cancel() @@ -15696,21 +16552,73 @@ async def test_recovery_submit_alias_persistence_failure_retires_before_send() - @pytest.mark.asyncio -async def test_recovery_submit_cancellation_after_alias_commit_restores_previous_owner() -> None: +async def test_recovery_submit_owner_fence_rejection_retires_before_send() -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) - key = _make_account_neutral_replay_session_key("alias-commit-cancel") send_text = AsyncMock() close = AsyncMock() - session = _make_bridge_session(key=key) + session = _make_bridge_session(key_value="owner-fence-rejection") session.upstream = cast( UpstreamWebSocket, SimpleNamespace(send_text=send_text, close=close), ) - session.durable_session_id = "durable-recovery" - session.durable_owner_epoch = 4 - service._http_bridge_sessions[key] = session - alias_owner = {"http_turn_commit_cancel": "durable-predecessor"} - alias_committed = asyncio.Event() + session.durable_session_id = "durable-owner-fence-rejection" + session.durable_owner_epoch = 2 + service._http_bridge_sessions[session.key] = session + record_recovery_attempt = AsyncMock(return_value=None) + service._durable_bridge = cast( + Any, + SimpleNamespace( + lookup_retry_circuit=AsyncMock(return_value=None), + record_recovery_attempt=record_recovery_attempt, + release_live_session=AsyncMock(return_value=None), + ), + ) + request_state = proxy_service._WebSocketRequestState( + request_id="req-owner-fence-rejection", + model="gpt-5.6-sol", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + awaiting_response_created=True, + request_text='{"type":"response.create","model":"gpt-5.6-sol","input":"hi"}', + fresh_upstream_request_text='{"type":"response.create","model":"gpt-5.6-sol","input":"hi"}', + fresh_upstream_request_is_retry_safe=True, + transport="http", + skip_request_log=True, + ) + + with pytest.raises(proxy_service.ProxyResponseError) as exc_info: + await service._submit_http_bridge_request( + session, + request_state=request_state, + text_data=request_state.request_text or "{}", + queue_limit=8, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.payload["error"]["code"] == "bridge_continuity_persistence_failed" + record_recovery_attempt.assert_awaited_once() + send_text.assert_not_awaited() + assert session.closed is True + + +@pytest.mark.asyncio +async def test_recovery_submit_cancellation_after_alias_commit_restores_previous_owner() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + key = _make_account_neutral_replay_session_key("alias-commit-cancel") + send_text = AsyncMock() + close = AsyncMock() + session = _make_bridge_session(key=key) + session.upstream = cast( + UpstreamWebSocket, + SimpleNamespace(send_text=send_text, close=close), + ) + session.durable_session_id = "durable-recovery" + session.durable_owner_epoch = 4 + service._http_bridge_sessions[key] = session + alias_owner = {"http_turn_commit_cancel": "durable-predecessor"} + alias_committed = asyncio.Event() release_registration = asyncio.Event() async def register_recovery_turn_state(**_kwargs: Any) -> DurableBridgeAliasRegistrationReceipt: @@ -17803,8 +18711,10 @@ async def fake_stream_events( last_call = get_or_create.await_args assert last_call is not None assert last_call.kwargs["previous_response_id"] is None - if unsafe_replay_input in {"missing_owner", "missing_prior_output", "orphan_output"}: + if unsafe_replay_input in {"missing_prior_output", "orphan_output"}: account_neutral_classifier.assert_not_called() + elif unsafe_replay_input == "missing_owner": + account_neutral_classifier.assert_called_once() else: account_neutral_classifier.assert_called_once() return @@ -18444,7 +19354,7 @@ async def close(self) -> None: assert owner.response_event_count == 0 if leading_telemetry: assert owner.latency_first_upstream_event_ms is not None - retry_precreated.assert_not_awaited() + retry_precreated.assert_awaited_once_with(session) assert write_request_log.await_count == 2 assert {call.kwargs["error_code"] for call in write_request_log.await_args_list} == {"upstream_request_timeout"} fail_reader.assert_awaited_once() @@ -18633,6 +19543,251 @@ async def reconnect_during_failure(*_args: object, **_kwargs: object) -> bool: assert session.closed is False +@pytest.mark.asyncio +async def test_http_bridge_reconnect_failure_keeps_reader_handoff_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = proxy_service._WebSocketRequestState( + request_id="req-reconnect-fails-closed", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + request_text='{"type":"response.create","model":"gpt-5.4","input":"hello"}', + transport="http", + ) + session = _make_bridge_session(key_value="bridge-reconnect-fails-closed") + session.closed = True + session.upstream = cast(UpstreamWebSocket, SimpleNamespace(close=AsyncMock())) + + async def old_reader() -> None: + await asyncio.sleep(60.0) + + session.upstream_reader = asyncio.create_task(old_reader()) + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="primary", + routing_strategy="usage_weighted", + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=settings)), + ) + + async def select_no_account(*_args: object, **_kwargs: object) -> object: + assert session.closed is True + return SimpleNamespace( + account=None, + error_code="no_accounts", + error_message="No active accounts available", + ) + + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", AsyncMock(side_effect=select_no_account)) + + with pytest.raises(ProxyResponseError): + await service._reconnect_http_bridge_session( + session, + request_state=request_state, + restart_reader=True, + ) + + assert session.closed is True + assert session.handoff_in_progress is False + assert session.handoff_future is None + assert session.key not in service._http_bridge_inflight_sessions + + +@pytest.mark.asyncio +async def test_http_bridge_handoff_future_survives_same_key_and_capacity_timeouts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + key = proxy_service._HTTPBridgeSessionKey("session_header", "sid-timeout-handoff", None) + session = _make_bridge_session(key=key, key_value="sid-timeout-handoff", queued_request_count=1) + handoff_future: asyncio.Future[proxy_service._HTTPBridgeSession] = asyncio.get_running_loop().create_future() + session.closed = True + session.handoff_in_progress = True + session.handoff_future = handoff_future + setattr(handoff_future, "_http_bridge_handoff", True) + service._http_bridge_sessions[key] = session + service._http_bridge_inflight_sessions[key] = handoff_future + settings = _make_app_settings(proxy_admission_wait_timeout_seconds=0.01) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_http_bridge_pending_count", AsyncMock(return_value=1)) + monkeypatch.setattr(proxy_service, "_http_bridge_should_wait_for_registration", AsyncMock(return_value=False)) + monkeypatch.setattr(proxy_service, "_http_bridge_owner_instance", AsyncMock(return_value="instance-a")) + monkeypatch.setattr( + proxy_service, + "_active_http_bridge_instance_ring", + AsyncMock(return_value=("instance-a", ("instance-a",))), + ) + + with pytest.raises(ProxyResponseError): + await service._get_or_create_http_bridge_session( + key, + headers={"x-codex-session-id": "sid-timeout-handoff"}, + affinity=proxy_service._AffinityPolicy( + key="sid-timeout-handoff", + kind=proxy_service.StickySessionKind.CODEX_SESSION, + ), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=8, + ) + assert service._http_bridge_inflight_sessions[key] is handoff_future + + with pytest.raises(ProxyResponseError): + await service._get_or_create_http_bridge_session( + proxy_service._HTTPBridgeSessionKey("session_header", "sid-capacity-waiter", None), + headers={"x-codex-session-id": "sid-capacity-waiter"}, + affinity=proxy_service._AffinityPolicy( + key="sid-capacity-waiter", + kind=proxy_service.StickySessionKind.CODEX_SESSION, + ), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + max_sessions=1, + ) + assert service._http_bridge_inflight_sessions[key] is handoff_future + + proxy_support_module._complete_http_bridge_handoff(session, service._http_bridge_inflight_sessions) + assert handoff_future.result() is session + + +@pytest.mark.asyncio +async def test_retry_http_bridge_precreated_request_consumes_each_clean_close_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = proxy_service._WebSocketRequestState( + request_id="req-clean-close-generation", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + awaiting_response_created=True, + request_text='{"type":"response.create","model":"gpt-5.4","input":"hello"}', + transport="http", + replay_count=1, + account_response_create_lease=cast(Any, object()), + ) + session = _make_bridge_session( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-clean-generation", None), + key_value="bridge-clean-generation", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + session.last_upstream_close_code = 1000 + session.last_upstream_close_generation = 7 + send_text = AsyncMock() + session.upstream = cast(UpstreamWebSocket, SimpleNamespace(send_text=send_text, close=AsyncMock())) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", AsyncMock()) + monkeypatch.setattr(service, "_release_request_state_account_response_create_lease", AsyncMock()) + + assert await service._retry_http_bridge_precreated_request(session) is True + assert await service._retry_http_bridge_precreated_request(session) is False + assert request_state.clean_close_replay_count == 1 + assert request_state.clean_close_retry_close_generation == 7 + assert send_text.await_count == 1 + + +@pytest.mark.asyncio +async def test_retry_http_bridge_model_fallback_excludes_rejected_hard_affinity_account( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = proxy_service._WebSocketRequestState( + request_id="req-model-fallback-account", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + awaiting_response_created=True, + request_text='{"type":"response.create","model":"gpt-5.4","input":"hello"}', + transport="http", + account_response_create_lease=cast(Any, object()), + precreated_replay_reason="account_model_unsupported", + precreated_replay_account_id="acc-rejected", + ) + session = _make_bridge_session( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-model-fallback", None), + key_value="bridge-model-fallback", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + session.account = cast(Any, SimpleNamespace(id="acc-rejected", status=AccountStatus.ACTIVE)) + session.last_upstream_close_code = 1011 + session.upstream = cast(UpstreamWebSocket, SimpleNamespace(send_text=AsyncMock(), close=AsyncMock())) + old_lease = request_state.account_response_create_lease + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + reconnect = AsyncMock() + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + release_lease = AsyncMock() + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_lease) + + assert await service._retry_http_bridge_precreated_request(session) is True + assert request_state.preferred_account_id is None + assert request_state.excluded_account_ids == {"acc-rejected"} + reconnect.assert_awaited_once() + reconnect_call = reconnect.await_args + assert reconnect_call is not None + assert reconnect_call.kwargs["request_state"] is request_state + release_lease.assert_awaited_once_with(old_lease) + + +@pytest.mark.asyncio +async def test_retry_http_bridge_fresh_hard_request_excludes_silent_account( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = proxy_service._WebSocketRequestState( + request_id="req-fresh-hard-account-fallback", + model="gpt-5.6-luna", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + awaiting_response_created=True, + request_text='{"type":"response.create","model":"gpt-5.6-luna","input":"hello"}', + transport="http", + account_response_create_lease=cast(Any, object()), + ) + session = _make_bridge_session( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-fresh-hard", None), + key_value="bridge-fresh-hard", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + session.account = cast(Any, SimpleNamespace(id="acc-silent", status=AccountStatus.ACTIVE)) + session.last_upstream_close_code = 1011 + session.upstream = cast(Any, SimpleNamespace(send_text=AsyncMock(), close=AsyncMock())) + old_lease = request_state.account_response_create_lease + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + reconnect = AsyncMock() + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + release_lease = AsyncMock() + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_lease) + + assert await service._retry_http_bridge_precreated_request(session) is True + assert request_state.preferred_account_id is None + assert request_state.excluded_account_ids == {"acc-silent"} + reconnect.assert_awaited_once() + reconnect_call = reconnect.await_args + assert reconnect_call is not None + assert reconnect_call.kwargs["request_state"] is request_state + release_lease.assert_awaited_once_with(old_lease) + + @pytest.mark.asyncio async def test_http_bridge_retry_send_network_failure_is_neutral_and_not_replayed( monkeypatch: pytest.MonkeyPatch, @@ -18726,6 +19881,32 @@ async def test_http_bridge_reader_preserves_routed_aiohttp_close_code( routed_websocket.close.assert_awaited_once() +@pytest.mark.asyncio +async def test_http_bridge_clean_close_before_response_does_not_penalize_account( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session(key_value="bridge-clean-close") + session.upstream = cast( + UpstreamWebSocket, + SimpleNamespace( + receive=AsyncMock(return_value=UpstreamWebSocketMessage(kind="close", close_code=1000)), + close=AsyncMock(), + ), + ) + fail_pending = AsyncMock() + retire = AsyncMock() + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr(service, "_fail_pending_websocket_requests", fail_pending) + monkeypatch.setattr(service, "_retire_stale_pending_http_bridge_session", retire) + + await service._relay_http_bridge_upstream_messages(session) + + assert fail_pending.await_args is not None + assert fail_pending.await_args.kwargs["penalize_account"] is False + retire.assert_awaited_once_with(session, detail="stream_incomplete", response_events_seen=0) + + @pytest.mark.asyncio @pytest.mark.parametrize("routed", [False, True], ids=["direct-close", "routed-receive-error"]) async def test_http_bridge_reader_maps_ordinary_websocket_receive_failure_to_stream_incomplete( @@ -18775,6 +19956,13 @@ async def fail_reader( assert len(failure_calls) == 1 assert failure_calls[0]["error_code"] == "stream_incomplete" assert failure_calls[0]["penalize_account"] is True + assert failure_calls[0]["response_events_seen"] == 0 + if routed: + assert failure_calls[0]["upstream_close_code"] is None + assert failure_calls[0]["transport_classification"] == "websocket_transport_error" + else: + assert failure_calls[0]["upstream_close_code"] == 1011 + assert failure_calls[0]["transport_classification"] == "websocket_close_transient" @pytest.mark.asyncio @@ -18888,6 +20076,591 @@ async def test_retire_stale_pending_http_bridge_session_unregisters_aliases_and_ close.assert_awaited_once() +@pytest.mark.asyncio +async def test_http_bridge_retirement_does_not_record_midstream_retry_circuit_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session(key_value="bridge-midstream-retire") + record_failure = AsyncMock() + close = AsyncMock() + monkeypatch.setattr(service, "_record_http_bridge_retry_circuit_failure", record_failure) + monkeypatch.setattr(service, "_close_http_bridge_session_bounded", close) + + await service._retire_stale_pending_http_bridge_session( + session, + detail="stream_incomplete", + response_events_seen=1, + ) + + record_failure.assert_not_awaited() + close.assert_awaited_once_with(session, reason="retire_stale_pending") + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_backoff_is_scoped_to_repeated_hard_keys() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-circuit-hard") + other_session = _make_bridge_session(key_value="bridge-circuit-other") + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_incomplete") + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_incomplete") + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert await service._http_bridge_precreated_retry_allowed(other_session) is True + + await service._clear_http_bridge_retry_circuit(hard_session) + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_allows_only_one_half_open_probe() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-circuit-half-open") + now = time.monotonic() + cast(Any, service)._http_bridge_retry_circuits[hard_session.key] = ( + http_bridge_retry_circuit_module._HTTPBridgeRetryCircuitState( + consecutive_failures=2, + cooldown_until=now - 1.0, + last_detail="missing_response_created_timeout", + last_touched_monotonic=now, + ) + ) + service._durable_bridge = SimpleNamespace(lookup_retry_circuit=AsyncMock(return_value=None)) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert ( + await service._http_bridge_precreated_retry_allowed( + hard_session, + allow_proof_gated_continuity_replay=True, + ) + is True + ) + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_allows_fresh_hard_account_switch_during_cooldown() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = proxy_service._WebSocketRequestState( + request_id="req-fresh-circuit-bypass", + model="gpt-5.6-luna", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + awaiting_response_created=True, + request_text='{"type":"response.create","model":"gpt-5.6-luna","input":"hello"}', + ) + hard_session = _make_bridge_session( + key_value="bridge-circuit-fresh-bypass", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_idle_timeout") + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_idle_timeout") + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert ( + await service._http_bridge_precreated_retry_allowed( + hard_session, + allow_fresh_hard_account_switch=True, + ) + is True + ) + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_allows_proof_gated_continuity_replay_during_cooldown() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = proxy_service._WebSocketRequestState( + request_id="req-proof-gated-circuit-bypass", + model="gpt-5.6-luna", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + awaiting_response_created=True, + request_text='{"type":"response.create","previous_response_id":"resp_anchor"}', + previous_response_id="resp_anchor", + fresh_upstream_request_text='{"type":"response.create","input":"full resend"}', + fresh_upstream_request_is_retry_safe=True, + ) + hard_session = _make_bridge_session( + key_value="bridge-circuit-proof-gated", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_idle_timeout") + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_idle_timeout") + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert ( + await service._http_bridge_precreated_retry_allowed( + hard_session, + allow_proof_gated_continuity_replay=True, + ) + is True + ) + + +@pytest.mark.asyncio +async def test_http_bridge_clean_close_retry_circuit_is_bounded() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session(key_value="bridge-clean-close-circuit") + + await service._record_http_bridge_retry_circuit_failure(session, detail="clean_close") + await service._record_http_bridge_retry_circuit_failure(session, detail="clean_close") + + cooldown = await service._http_bridge_precreated_retry_cooldown_seconds(session) + assert 0 < cooldown <= 30.0 + retry_circuits = cast(Any, service)._http_bridge_retry_circuits + assert retry_circuits[session.key].last_detail == "clean_close" + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_restores_persisted_cooldown() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-persisted-circuit") + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock( + return_value=SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=time.time(), + ) + ), + clear_retry_circuit=AsyncMock(), + persist_retry_circuit=AsyncMock(), + ) + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert await service._http_bridge_precreated_retry_cooldown_seconds(hard_session) > 0 + assert service._durable_bridge.lookup_retry_circuit.await_count == 2 + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_purges_expired_persisted_state() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-expired-circuit") + expired_updated_at = ( + time.time() - http_bridge_retry_circuit_module.DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS - 1.0 + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock( + return_value=SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=expired_updated_at, + ) + ), + purge_retry_circuit=AsyncMock(), + ) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + service._durable_bridge.purge_retry_circuit.assert_awaited_once_with( + session_key_kind=hard_session.key.affinity_kind, + session_key_value=hard_session.key.affinity_key, + api_key_id=hard_session.key.api_key_id, + expected_updated_at_epoch=expired_updated_at, + ) + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_keeps_newer_local_failure_after_stale_purge() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-expired-circuit-new-local-failure") + now = time.monotonic() + local_state = http_bridge_retry_circuit_module._HTTPBridgeRetryCircuitState( + consecutive_failures=2, + cooldown_until=now + 60.0, + last_detail="stream_idle_timeout", + last_touched_monotonic=now, + last_failure_monotonic=now, + last_durable_load_monotonic=now - 10.0, + ) + cast(Any, service)._http_bridge_retry_circuits[hard_session.key] = local_state + cast(Any, service)._http_bridge_retry_circuit_persisted_keys.add(hard_session.key) + expired_updated_at = ( + time.time() - http_bridge_retry_circuit_module.DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS - 1.0 + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock( + return_value=SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_idle_timeout", + updated_at_epoch=expired_updated_at, + ) + ), + purge_retry_circuit=AsyncMock(), + ) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert cast(Any, service)._http_bridge_retry_circuits[hard_session.key] is local_state + service._durable_bridge.purge_retry_circuit.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_keeps_local_state_when_stale_purge_fails() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-expired-circuit-local-fallback") + now = time.monotonic() + local_state = http_bridge_retry_circuit_module._HTTPBridgeRetryCircuitState( + consecutive_failures=2, + cooldown_until=now + 60.0, + last_detail="stream_idle_timeout", + last_touched_monotonic=now, + last_failure_monotonic=now, + ) + cast(Any, service)._http_bridge_retry_circuits[hard_session.key] = local_state + cast(Any, service)._http_bridge_retry_circuit_persisted_keys.add(hard_session.key) + expired_updated_at = ( + time.time() - http_bridge_retry_circuit_module.DURABLE_BRIDGE_RETRY_CIRCUIT_STATE_TTL_SECONDS - 1.0 + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock( + return_value=SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_idle_timeout", + updated_at_epoch=expired_updated_at, + ) + ), + purge_retry_circuit=AsyncMock(side_effect=RuntimeError("durable purge unavailable")), + ) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert cast(Any, service)._http_bridge_retry_circuits[hard_session.key] is local_state + assert local_state.consecutive_failures == 2 + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_refreshes_persisted_state_after_initial_miss() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-replica-refresh-circuit") + persisted = SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=time.time(), + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(side_effect=[None, persisted]), + clear_retry_circuit=AsyncMock(), + persist_retry_circuit=AsyncMock(), + ) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert service._durable_bridge.lookup_retry_circuit.await_count == 2 + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_preserves_newer_local_failure_on_durable_refresh() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-local-newer-than-durable-refresh") + now_monotonic = time.monotonic() + local_state = http_bridge_retry_circuit_module._HTTPBridgeRetryCircuitState( + consecutive_failures=2, + cooldown_until=now_monotonic + 90.0, + last_detail="stream_idle_timeout", + last_touched_monotonic=now_monotonic, + last_failure_monotonic=now_monotonic, + last_durable_load_monotonic=now_monotonic - 10.0, + persisted_updated_at_epoch=time.time() - 10.0, + ) + cast(Any, service)._http_bridge_retry_circuits[hard_session.key] = local_state + cast(Any, service)._http_bridge_retry_circuit_persisted_keys.add(hard_session.key) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock( + return_value=SimpleNamespace( + consecutive_failures=1, + cooldown_until_epoch=time.time() + 1.0, + last_detail="stream_incomplete", + updated_at_epoch=time.time(), + ) + ) + ) + + await service._load_http_bridge_retry_circuit(hard_session) + + assert cast(Any, service)._http_bridge_retry_circuits[hard_session.key] is local_state + assert local_state.consecutive_failures == 2 + assert local_state.cooldown_until >= now_monotonic + 89.0 + assert local_state.last_detail == "stream_idle_timeout" + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_drops_local_state_after_durable_clear() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-replica-cleared-circuit") + persisted = SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=time.time(), + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(side_effect=[persisted, None]), + clear_retry_circuit=AsyncMock(), + persist_retry_circuit=AsyncMock(), + ) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + assert hard_session.key not in cast(Any, service)._http_bridge_retry_circuits + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_clear_retries_after_lookup_failure() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-retry-clear-after-lookup-failure") + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(side_effect=RuntimeError("durable read unavailable")), + clear_retry_circuit=AsyncMock(), + ) + state = http_bridge_retry_circuit_module._HTTPBridgeRetryCircuitState( + consecutive_failures=2, + cooldown_until=time.monotonic() + 60.0, + last_detail="stream_incomplete", + last_touched_monotonic=time.monotonic(), + ) + cast(Any, service)._http_bridge_retry_circuits[hard_session.key] = state + cast(Any, service)._http_bridge_retry_circuit_persisted_keys.add(hard_session.key) + + await service._clear_http_bridge_retry_circuit(hard_session) + + service._durable_bridge.clear_retry_circuit.assert_awaited_once_with( + session_key_kind=hard_session.key.affinity_kind, + session_key_value=hard_session.key.affinity_key, + api_key_id=hard_session.key.api_key_id, + expected_updated_at_epoch=None, + ) + assert hard_session.key not in cast(Any, service)._http_bridge_retry_circuits + assert hard_session.key not in cast(Any, service)._http_bridge_retry_circuit_persisted_keys + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_replaces_local_state_after_newer_reset() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-replica-reset-lineage") + now_epoch = time.time() + reset = SimpleNamespace( + consecutive_failures=0, + cooldown_until_epoch=0.0, + last_detail=None, + updated_at_epoch=now_epoch, + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock( + side_effect=[ + SimpleNamespace( + consecutive_failures=3, + cooldown_until_epoch=now_epoch + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=now_epoch - 1.0, + ), + reset, + ] + ), + persist_retry_circuit=AsyncMock( + return_value=SimpleNamespace( + consecutive_failures=1, + cooldown_until_epoch=0.0, + last_detail="stream_incomplete", + updated_at_epoch=now_epoch + 1.0, + ) + ), + ) + + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_incomplete") + + state = cast(Any, service)._http_bridge_retry_circuits[hard_session.key] + assert state.consecutive_failures == 1 + assert state.persisted_updated_at_epoch == now_epoch + 1.0 + assert service._durable_bridge.persist_retry_circuit.await_args.kwargs["base_updated_at_epoch"] == now_epoch + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_refreshes_conflict_merged_persisted_state() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-conflict-merged-circuit") + loaded = SimpleNamespace( + consecutive_failures=1, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=time.time(), + ) + persisted = SimpleNamespace( + consecutive_failures=2, + cooldown_until_epoch=time.time() + 60.0, + last_detail="stream_incomplete", + updated_at_epoch=time.time(), + ) + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(return_value=loaded), + persist_retry_circuit=AsyncMock(return_value=persisted), + ) + + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_incomplete") + + retry_circuits = cast(Any, service)._http_bridge_retry_circuits + assert retry_circuits[hard_session.key].consecutive_failures == 2 + assert retry_circuits[hard_session.key].cooldown_until > time.monotonic() + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_counts_stream_idle_timeout() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-idle-timeout-circuit") + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(return_value=None), + persist_retry_circuit=AsyncMock(), + ) + + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="stream_idle_timeout") + + assert cast(Any, service)._http_bridge_retry_circuits[hard_session.key].consecutive_failures == 1 + + +@pytest.mark.parametrize( + ("error_code", "expected_ambiguous"), + [ + ("stream_incomplete", True), + ("stream_idle_timeout", True), + ("upstream_request_timeout", True), + ("invalid_request_error", False), + ("quota_exceeded", False), + ], +) +def test_http_bridge_durable_recovery_requires_ambiguous_transport_error( + error_code: str, + expected_ambiguous: bool, +) -> None: + error = ProxyResponseError(502, openai_error(error_code, "upstream result")) + + assert http_bridge_streaming_module._http_bridge_error_is_ambiguous_transport(error) is expected_ambiguous + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_counts_missing_response_created_timeout() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-missing-response-created-circuit") + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(return_value=None), + persist_retry_circuit=AsyncMock(), + ) + + await service._record_http_bridge_retry_circuit_failure( + hard_session, + detail="missing_response_created_timeout", + ) + await service._record_http_bridge_retry_circuit_failure( + hard_session, + detail="missing_response_created_timeout", + ) + + state = cast(Any, service)._http_bridge_retry_circuits[hard_session.key] + assert state.consecutive_failures == 2 + assert state.last_detail == "stream_idle_timeout" + assert state.cooldown_until > time.monotonic() + assert await service._http_bridge_precreated_retry_allowed(hard_session) is False + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_counts_stuck_gate_timeout() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-stuck-gate-circuit") + service._durable_bridge = SimpleNamespace( + lookup_retry_circuit=AsyncMock(return_value=None), + persist_retry_circuit=AsyncMock(), + ) + + for _ in range(2): + await service._record_http_bridge_retry_circuit_failure( + hard_session, + detail="response_create_gate_timeout_stuck_pending", + ) + + state = cast(Any, service)._http_bridge_retry_circuits[hard_session.key] + assert state.consecutive_failures == 2 + assert state.last_detail == "stream_idle_timeout" + assert state.cooldown_until > time.monotonic() + + +@pytest.mark.asyncio +async def test_http_bridge_submit_suppresses_hard_key_during_retry_cooldown() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + hard_session = _make_bridge_session(key_value="bridge-submit-cooldown") + now = time.monotonic() + cast(Any, service)._http_bridge_retry_circuits[hard_session.key] = ( + http_bridge_retry_circuit_module._HTTPBridgeRetryCircuitState( + consecutive_failures=2, + cooldown_until=now + 60.0, + last_detail="stream_idle_timeout", + last_touched_monotonic=now, + ) + ) + service._durable_bridge = SimpleNamespace(lookup_retry_circuit=AsyncMock(return_value=None)) + request_state = proxy_service._WebSocketRequestState( + request_id="req-submit-cooldown", + model="gpt-5.6-luna", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + awaiting_response_created=True, + request_text='{"type":"response.create","input":"hello"}', + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._submit_http_bridge_request_with_handoff( + hard_session, + request_state=request_state, + text_data=request_state.request_text or "", + queue_limit=8, + request_scope_id="scope-submit-cooldown", + ) + + assert exc_info.value.status_code == 503 + assert exc_info.value.payload["error"]["code"] == "upstream_request_timeout" + assert exc_info.value.retry_after_seconds is not None + assert exc_info.value.retry_after_seconds >= 60 + assert hard_session.queued_request_count == 0 + + +@pytest.mark.asyncio +async def test_http_bridge_retry_circuit_ignores_soft_affinity_and_other_failures( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + monkeypatch.setattr( + http_bridge_retry_circuit_module, + "_HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD", + 1, + ) + soft_key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-circuit-soft", None) + soft_session = _make_bridge_session(key=soft_key) + + await service._record_http_bridge_retry_circuit_failure(soft_session, detail="stream_incomplete") + assert await service._http_bridge_precreated_retry_allowed(soft_session) is True + + hard_session = _make_bridge_session(key_value="bridge-circuit-other-error") + await service._record_http_bridge_retry_circuit_failure(hard_session, detail="proxy_network_unavailable") + assert await service._http_bridge_precreated_retry_allowed(hard_session) is True + + @pytest.mark.asyncio async def test_http_bridge_reader_failed_precreated_replay_retires_registered_session( monkeypatch: pytest.MonkeyPatch, @@ -19106,7 +20879,9 @@ async def test_http_bridge_reader_failure_keeps_waiter_count_when_draining_reque ) session.admission_waiter_count = 1 fail_pending = AsyncMock() + record_failure = AsyncMock() monkeypatch.setattr(service, "_fail_pending_websocket_requests", fail_pending) + monkeypatch.setattr(service, "_record_http_bridge_retry_circuit_failure", record_failure) retired = await service._fail_http_bridge_reader_and_maybe_retire( session, @@ -19116,6 +20891,7 @@ async def test_http_bridge_reader_failure_keeps_waiter_count_when_draining_reque assert retired is False assert session.queued_request_count == 1 + record_failure.assert_awaited_once_with(session, detail="stream_incomplete") @pytest.mark.asyncio @@ -19141,7 +20917,7 @@ async def test_http_bridge_eventless_timeout_force_retires_with_admission_waiter assert retired is True assert session.closed is True - retire.assert_awaited_once_with(session, detail="missing_response_created_timeout") + retire.assert_awaited_once_with(session, detail="missing_response_created_timeout", response_events_seen=0) fail_pending_await_args = fail_pending.await_args assert fail_pending_await_args is not None assert fail_pending_await_args.kwargs["penalize_account"] is False @@ -19165,7 +20941,7 @@ async def test_http_bridge_reader_failure_retires_without_waiters_when_notificat error_message="closed", ) - retire.assert_awaited_once_with(session, detail="stream_incomplete") + retire.assert_awaited_once_with(session, detail="stream_incomplete", response_events_seen=0) @pytest.mark.asyncio diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index f5d77f1a88..3722c7131b 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -454,6 +454,16 @@ def test_account_selection_recovery_sleep_treats_workspace_spend_cap_as_recovera assert _account_selection_recovery_sleep_seconds(selection) == 30.0 +def test_account_selection_recovery_sleep_retries_hard_affinity_owner_briefly(): + selection = AccountSelection( + account=None, + error_message="Hard affinity owner account is unavailable", + error_code="hard_affinity_saturated", + ) + + assert _account_selection_recovery_sleep_seconds(selection) == 2.0 + + def test_account_selection_recovery_sleep_ignores_generic_no_available_accounts(): selection = AccountSelection(account=None, error_message="No available accounts", error_code="no_accounts") @@ -32068,6 +32078,100 @@ async def test_response_create_admission_session_gate_timeout_returns_stable_rea assert request_state.response_create_admission is None +@pytest.mark.asyncio +async def test_response_create_admission_stuck_gate_retire_ignores_draining_pending(monkeypatch): + settings = _make_proxy_settings() + settings.proxy_response_create_limit = 64 + settings.proxy_admission_wait_timeout_seconds = 0.01 + settings.http_responses_session_bridge_stuck_gate_retire_after_seconds = 1.0 + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + blocked_request = proxy_service._WebSocketRequestState( + request_id="ws_req_gate_timeout_with_draining", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + ) + stale_gate_holder = proxy_service._WebSocketRequestState( + request_id="req_stale_gate_holder", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + transport="http", + awaiting_response_created=True, + response_create_gate_acquired=True, + response_create_gate_wait_started_at=0.0, + ) + active_request = proxy_service._WebSocketRequestState( + request_id="req_active_sibling", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + awaiting_response_created=True, + response_create_gate_acquired=True, + response_event_count=1, + ) + draining_request = proxy_service._WebSocketRequestState( + request_id="req_draining_not_stale", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + transport="http", + draining_until_terminal=True, + response_event_count=1, + ) + response_create_gate = asyncio.Semaphore(1) + await response_create_gate.acquire() + bridge_session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-stuck-gate-draining", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_stuck_gate_draining"), + upstream=AsyncMock(), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque([draining_request, active_request, stale_gate_holder]), + pending_lock=anyio.Lock(), + response_create_gate=response_create_gate, + queued_request_count=2, + last_used_at=0.0, + idle_ttl_seconds=30.0, + ) + retire_stale = AsyncMock() + fail_stale = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service, "_retire_stale_pending_http_bridge_session", retire_stale) + monkeypatch.setattr(service, "_fail_stale_http_bridge_pending_requests", fail_stale) + + try: + with pytest.raises(proxy_module.ProxyResponseError) as exc_info: + await service._acquire_request_state_response_create_admission( + blocked_request, + response_create_gate=response_create_gate, + bridge_session=bridge_session, + ) + finally: + response_create_gate.release() + + exc = _assert_proxy_response_error(exc_info.value) + assert _proxy_error_code(exc) == "response_create_gate_timeout" + retire_stale.assert_not_awaited() + fail_stale.assert_awaited_once_with( + bridge_session, + [stale_gate_holder], + detail="response_create_gate_timeout_stuck_pending", + ) + + @pytest.mark.asyncio async def test_response_create_admission_waits_on_session_gate_before_shared_capacity(monkeypatch): settings = _make_proxy_settings() @@ -33392,8 +33496,8 @@ async def fake_stream_with_retry( assert calls == [(payload, "acc_doc")] -def test_classify_upstream_close_rejected_only_for_clean_close_before_any_response_event(): - assert proxy_service._classify_upstream_close(1000, response_events_seen=0) == "rejected" +def test_classify_upstream_close_clean_for_clean_close_before_any_response_event(): + assert proxy_service._classify_upstream_close(1000, response_events_seen=0) == "clean" assert proxy_service._classify_upstream_close(1000, response_events_seen=1) == "transient" assert proxy_service._classify_upstream_close(1011, response_events_seen=0) == "transient" @@ -33829,31 +33933,14 @@ async def close_old_upstream(*_args: object, **_kwargs: object) -> None: kind="stream", acquired_at=2.0, ) - allow_reacquire = asyncio.Event() - reacquire_started = asyncio.Event() release_started = asyncio.Event() finish_release = asyncio.Event() - async def acquire_account_lease(*_args: object, **_kwargs: object) -> proxy_service.AccountLease: - reacquire_started.set() - await allow_reacquire.wait() - return reacquired_lease - monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) - monkeypatch.setattr(proxy_service.time, "monotonic", lambda: 10.0) - monkeypatch.setattr( - service._load_balancer, - "select_account", - AsyncMock( - return_value=AccountSelection( - account=account, - error_message=None, - lease=reconnect_lease, - ) - ), - ) - monkeypatch.setattr(service._load_balancer, "acquire_account_lease", acquire_account_lease) + # Keep the deterministic service clock local; mutating the shared time + # module also freezes asyncio's event-loop timers. + monkeypatch.setattr(proxy_service, "time", SimpleNamespace(monotonic=lambda: 10.0)) async def release_account_lease_side_effect(released_lease: proxy_service.AccountLease) -> None: assert released_lease is reacquired_lease @@ -33889,28 +33976,19 @@ async def release_account_lease_side_effect(released_lease: proxy_service.Accoun idle_ttl_seconds=30.0, ) - reconnect_task = asyncio.create_task(service._reconnect_http_bridge_session(session, request_state=request_state)) - await reconnect_before_swap.wait() - - async def reacquire_under_pending_lock() -> None: - async with session.pending_lock: - await service._ensure_http_bridge_session_stream_lease_locked(session, request_state=request_state) + # Model an idle-session lease being reacquired immediately before the + # reconnect handoff. The reconnect must release that superseded lease + # while retaining the newly selected lease. + async def select_with_reacquire(*_args: object, **_kwargs: object) -> AccountSelection: + session.account_lease = reacquired_lease + return AccountSelection(account=account, error_message=None, lease=reconnect_lease) - reacquire_task = asyncio.create_task(reacquire_under_pending_lock()) - await reacquire_started.wait() + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_with_reacquire) + reconnect_task = asyncio.create_task(service._reconnect_http_bridge_session(session, request_state=request_state)) allow_old_upstream_close.set() - await asyncio.sleep(0) - allow_reacquire.set() - await reacquire_task - await release_started.wait() - reconnect_task.cancel() - reconnect_task.cancel() - await asyncio.sleep(0) - assert not reconnect_task.done() - + await asyncio.wait_for(release_started.wait(), timeout=2) finish_release.set() - with pytest.raises(asyncio.CancelledError): - await reconnect_task + await reconnect_task assert session.account_lease is reconnect_lease assert session.upstream is new_upstream @@ -34008,6 +34086,101 @@ async def test_reconnect_http_bridge_security_rebind_clears_previous_response_st assert "x-codex-turn-state" not in {header.lower() for header in session.headers} +@pytest.mark.asyncio +async def test_reconnect_http_bridge_session_restarts_reader_before_local_close(monkeypatch): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_bridge_reader_handoff") + old_upstream = AsyncMock() + new_upstream = SimpleNamespace(response_header=lambda _name: None) + old_reader_started = asyncio.Event() + old_reader_cancelled = asyncio.Event() + replacement_reader_started = asyncio.Event() + release_replacement_reader = asyncio.Event() + session_holder = {} + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service.time, "monotonic", lambda: 10.0) + monkeypatch.setattr( + service._load_balancer, + "select_account", + AsyncMock(return_value=AccountSelection(account=account, error_message=None)), + ) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(return_value=account)) + + async def blocked_old_reader() -> None: + old_reader_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + old_reader_cancelled.set() + raise + finally: + session_holder["session"].closed = True + + async def close_old_upstream() -> None: + assert old_reader_cancelled.is_set() + + async def replacement_reader(target_session) -> None: + assert target_session.upstream is new_upstream + replacement_reader_started.set() + await release_replacement_reader.wait() + + async def open_replacement(*_args, **_kwargs): + # Reconnect keeps the bridge closed until the replacement socket has + # been installed; the handoff publishes it open only after connect. + assert session_holder["session"].closed is True + return new_upstream + + old_upstream.close.side_effect = close_old_upstream + monkeypatch.setattr(service, "_relay_http_bridge_upstream_messages", replacement_reader) + monkeypatch.setattr(service, "_open_upstream_websocket_with_budget", open_replacement) + old_reader = asyncio.create_task(blocked_old_reader()) + await old_reader_started.wait() + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_reader_handoff", + model="gpt-5.5", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=10.0, + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-reader-handoff", None), + headers={}, + affinity=proxy_service._AffinityPolicy(key="bridge-reader-handoff"), + request_model="gpt-5.5", + account=account, + upstream=old_upstream, + 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, + upstream_reader=old_reader, + ) + session_holder["session"] = session + + try: + await service._reconnect_http_bridge_session(session, request_state=request_state, restart_reader=True) + await asyncio.wait_for(replacement_reader_started.wait(), timeout=1.0) + + assert old_reader.cancelled() + assert session.upstream is new_upstream + assert session.upstream_reader is not None + assert session.upstream_reader is not old_reader + assert session.upstream_reader.done() is False + old_upstream.close.assert_awaited_once() + finally: + release_replacement_reader.set() + replacement_task = session.upstream_reader + if replacement_task is not None: + await replacement_task + + @pytest.mark.asyncio async def test_reconnect_http_bridge_session_fails_over_after_repeated_401_refresh_retry(monkeypatch): settings = _make_proxy_settings() @@ -36280,15 +36453,151 @@ async def test_http_bridge_session_events_keepalive_backstop(monkeypatch): finally: await events.aclose() - assert len(collected) == 3, f"Expected 3 events (2 keepalives + stream_idle_timeout), got {len(collected)}" + assert len(collected) == 2, f"Expected 2 events (1 keepalive + stream_idle_timeout), got {len(collected)}" assert collected[0] == proxy_service.CODEX_KEEPALIVE_FRAME - assert collected[1] == proxy_service.CODEX_KEEPALIVE_FRAME - last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[2])) + last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[1])) assert last["type"] == "response.failed" assert cast(dict[str, object], last["response"])["status"] == "failed" assert cast(dict[str, object], cast(dict[str, object], last["response"])["error"])["code"] == "stream_idle_timeout" +@pytest.mark.asyncio +async def test_http_bridge_session_events_retries_silent_pre_response_once(monkeypatch): + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + settings = _make_proxy_settings() + settings.sse_keepalive_interval_seconds = 0.001 + settings.stream_idle_timeout_seconds = 7200.0 + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_idle_retry", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + response_id=None, + event_queue=asyncio.Queue(), + request_text='{"type":"response.create"}', + transport="http", + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-idle-retry", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_bridge_idle_retry"), + upstream=AsyncMock(), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=0.0, + idle_ttl_seconds=30.0, + ) + retry_precreated = AsyncMock(side_effect=[True, False]) + + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_STREAM_KEEPALIVE_MAX_COUNT", 2) + monkeypatch.setattr(proxy_service, "_HTTP_BRIDGE_STARTUP_KEEPALIVE_GRACE_SECONDS", 0.001) + monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_retry_http_bridge_precreated_request", retry_precreated) + + events = service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data='{"type":"response.create"}', + queue_limit=10, + propagate_http_errors=False, + downstream_turn_state=None, + ) + collected: list[str] = [] + try: + async for event in events: + collected.append(event) + finally: + await events.aclose() + + assert len(collected) == 3 + assert collected[:2] == [proxy_service.CODEX_KEEPALIVE_FRAME] * 2 + last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[-1])) + assert cast(dict[str, object], cast(dict[str, object], last["response"])["error"])["code"] == "stream_idle_timeout" + assert retry_precreated.await_count == 2 + assert all(attempt.kwargs == {"restart_reader": True} for attempt in retry_precreated.await_args_list) + + +@pytest.mark.asyncio +async def test_http_bridge_session_events_keeps_alive_during_retry_circuit_cooldown( + monkeypatch: pytest.MonkeyPatch, +) -> None: + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + settings = _make_proxy_settings() + monkeypatch.setattr(proxy_service, "_STREAM_KEEPALIVE_MAX_COUNT", 1) + settings.sse_keepalive_interval_seconds = 0.001 + settings.stream_idle_timeout_seconds = 1.0 + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_circuit_keepalive", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + response_id=None, + event_queue=asyncio.Queue(), + request_text='{"type":"response.create"}', + fresh_upstream_request_text='{"type":"response.create"}', + fresh_upstream_request_is_retry_safe=True, + transport="http", + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-circuit-keepalive", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_bridge_circuit_keepalive"), + 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, + ) + retry_precreated = AsyncMock(side_effect=[False, True, False]) + retry_cooldown = AsyncMock(side_effect=[0.1, 0.1, 0.0, 0.0]) + + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service, "_submit_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock()) + monkeypatch.setattr(service, "_retry_http_bridge_precreated_request", retry_precreated) + monkeypatch.setattr(service, "_http_bridge_precreated_retry_cooldown_seconds", retry_cooldown) + + events = service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data='{"type":"response.create"}', + queue_limit=10, + propagate_http_errors=False, + downstream_turn_state=None, + ) + collected: list[str] = [] + try: + async for event in events: + collected.append(event) + finally: + await events.aclose() + + assert len(collected) >= 2 + assert collected[:-1] == [proxy_service.CODEX_KEEPALIVE_FRAME] * (len(collected) - 1) + last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[-1])) + assert cast(dict[str, object], cast(dict[str, object], last["response"])["error"])["code"] == "stream_idle_timeout" + assert retry_precreated.await_count >= 1 + assert retry_cooldown.await_count >= 1 + + @pytest.mark.asyncio async def test_http_bridge_session_events_keepalive_backstop_respects_idle_timeout(monkeypatch): request_logs = _RequestLogsRecorder() @@ -36347,9 +36656,9 @@ async def test_http_bridge_session_events_keepalive_backstop_respects_idle_timeo finally: await events.aclose() - assert len(collected) == 6 - assert collected[:5] == [proxy_service.CODEX_KEEPALIVE_FRAME] * 5 - last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[5])) + assert len(collected) == 2 + assert collected[0] == proxy_service.CODEX_KEEPALIVE_FRAME + last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[1])) assert last["type"] == "response.failed" assert cast(dict[str, object], cast(dict[str, object], last["response"])["error"])["code"] == "stream_idle_timeout" @@ -36410,16 +36719,14 @@ async def test_http_bridge_session_events_keepalive_backstop_with_response_id(mo finally: await events.aclose() - assert len(collected) == 3, ( - f"Expected 3 events (2 response.in_progress + stream_idle_timeout), got {len(collected)}" + assert len(collected) == 2, ( + f"Expected 2 events (1 response.in_progress + stream_idle_timeout), got {len(collected)}" ) - first = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[0])) - assert first["type"] == "response.in_progress" - assert cast(dict[str, object], first["response"])["id"] == "resp_bridge_backstop_codex" - second = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[1])) - assert second["type"] == "response.in_progress" - assert cast(dict[str, object], second["response"])["id"] == "resp_bridge_backstop_codex" - last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[2])) + for event_block in collected[:1]: + event = cast(dict[str, object], proxy_service.parse_sse_data_json(event_block)) + assert event["type"] == "response.in_progress" + assert cast(dict[str, object], event["response"])["id"] == "resp_bridge_backstop_codex" + last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[1])) assert last["type"] == "response.failed" assert cast(dict[str, object], last["response"])["status"] == "failed" assert cast(dict[str, object], cast(dict[str, object], last["response"])["error"])["code"] == "stream_idle_timeout" @@ -36477,14 +36784,15 @@ async def test_http_bridge_session_events_keepalive_backstop_uses_replay_downstr try: async for event in events: collected.append(event) - if len(collected) >= 3: + if len(collected) >= 2: break finally: await events.aclose() + assert len(collected) == 2 first = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[0])) assert cast(dict[str, object], first["response"])["id"] == "resp_created_then_closed" - last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[2])) + last = cast(dict[str, object], proxy_service.parse_sse_data_json(collected[1])) assert cast(dict[str, object], last["response"])["id"] == "resp_created_then_closed" assert cast(dict[str, object], cast(dict[str, object], last["response"])["error"])["code"] == "stream_idle_timeout" @@ -36682,7 +36990,7 @@ async def capture_send_text(_text: str) -> None: @pytest.mark.asyncio -async def test_retry_http_bridge_precreated_request_suppresses_retry_for_rejected_close(): +async def test_retry_http_bridge_precreated_request_suppresses_retry_for_rejected_close(monkeypatch): request_logs = _RequestLogsRecorder() service = proxy_service.ProxyService(_repo_factory(request_logs)) request_state = proxy_service._WebSocketRequestState( @@ -36712,6 +37020,14 @@ async def test_retry_http_bridge_precreated_request_suppresses_retry_for_rejecte last_upstream_close_code=1000, ) + async def rejected_reconnect(*_args: object, **_kwargs: object) -> None: + raise proxy_module.ProxyResponseError( + 502, + openai_error("upstream_rejected_input", "Upstream rejected pre-created response (close_code=1000)"), + ) + + monkeypatch.setattr(service, "_reconnect_http_bridge_session", rejected_reconnect) + retried = await service._retry_http_bridge_precreated_request(session) assert retried is False @@ -36720,6 +37036,296 @@ async def test_retry_http_bridge_precreated_request_suppresses_retry_for_rejecte assert "close_code=1000" in (request_state.error_message_override or "") +@pytest.mark.asyncio +async def test_http_bridge_prewarm_completion_does_not_clear_retry_circuit(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + clear_retry_circuit = AsyncMock() + account = _make_account("acc_prewarm_circuit") + request_state = proxy_service._WebSocketRequestState( + request_id="req_prewarm_circuit", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + response_id="resp_prewarm_circuit", + event_queue=asyncio.Queue(), + request_text='{"type":"response.create"}', + transport="http", + request_kind="prewarm", + skip_request_log=True, + ) + upstream = AsyncMock() + upstream.archive_received = lambda _message: None + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-prewarm-circuit", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=account, + upstream=upstream, + 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, + ) + monkeypatch.setattr(service, "_clear_http_bridge_retry_circuit", clear_retry_circuit) + monkeypatch.setattr(service, "_finalize_websocket_request_state", AsyncMock()) + + await service._process_http_bridge_upstream_text( + session, + json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_prewarm_circuit", + "status": "completed", + }, + }, + separators=(",", ":"), + ), + ) + + clear_retry_circuit.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_http_bridge_reader_failure_records_clean_close_circuit_with_admission_waiter(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + record_failure = AsyncMock() + fail_pending = AsyncMock() + retire_stale = AsyncMock() + request_state = proxy_service._WebSocketRequestState( + request_id="req_clean_close_with_waiter", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + awaiting_response_created=True, + transport="http", + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-clean-close-waiter", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_clean_close_waiter"), + 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, + admission_waiter_count=1, + last_upstream_close_code=1000, + ) + monkeypatch.setattr(service, "_record_http_bridge_retry_circuit_failure", record_failure) + monkeypatch.setattr(service, "_fail_pending_websocket_requests", fail_pending) + monkeypatch.setattr(service, "_retire_stale_pending_http_bridge_session", retire_stale) + + retired = await service._fail_http_bridge_reader_and_maybe_retire( + session, + error_code="stream_incomplete", + error_message="HTTP bridge upstream closed cleanly before response.completed", + ) + + assert retired is False + record_failure.assert_awaited_once_with(session, detail="clean_close") + fail_pending.assert_awaited_once() + retire_stale.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_retry_http_bridge_precreated_request_propagates_reader_restart(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_external_idle_retry", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + awaiting_response_created=True, + request_text='{"type":"response.create"}', + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("prompt_cache", "bridge-external-idle-retry", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_bridge_external_idle_retry"), + 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, + ) + reconnect = AsyncMock(return_value=None) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + assert await service._retry_http_bridge_precreated_request(session, restart_reader=True) is True + + reconnect.assert_awaited_once_with( + session, + request_state=request_state, + restart_reader=True, + ) + cast(AsyncMock, session.upstream.send_text).assert_awaited_once_with(request_state.request_text) + + +@pytest.mark.asyncio +async def test_retry_http_bridge_clean_close_keeps_hard_continuation_on_same_account(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_bridge_clean_continuation") + original_payload = { + "type": "response.create", + "model": "gpt-5.1", + "previous_response_id": "resp_anchor", + "input": [{"role": "user", "content": "continue"}], + } + fresh_payload = dict(original_payload) + fresh_payload.pop("previous_response_id") + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_clean_continuation", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + awaiting_response_created=True, + request_text=json.dumps(original_payload, separators=(",", ":")), + previous_response_id="resp_anchor", + preferred_account_id=account.id, + proxy_injected_previous_response_id=True, + fresh_upstream_request_text=json.dumps(fresh_payload, separators=(",", ":")), + fresh_upstream_request_is_retry_safe=True, + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-clean-continuation", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=account, + 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, + last_upstream_close_code=1000, + ) + reconnect = AsyncMock(return_value=None) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + assert await service._retry_http_bridge_precreated_request(session) is True + + reconnect.assert_awaited_once_with(session, request_state=request_state, require_same_account=True) + cast(AsyncMock, session.upstream.send_text).assert_awaited_once_with(request_state.request_text) + + +@pytest.mark.asyncio +async def test_retry_http_bridge_clean_close_keeps_hard_turn_state_on_same_account(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_bridge_clean_turn_state") + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_clean_turn_state", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + awaiting_response_created=True, + request_text='{"type":"response.create","input":"turn"}', + hard_continuity_anchor=True, + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("turn_state_header", "turn-state-clean-close", None), + headers={"x-codex-turn-state": "turn-state-clean-close"}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=account, + 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, + last_upstream_close_code=1000, + ) + reconnect = AsyncMock(return_value=None) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + assert await service._retry_http_bridge_precreated_request(session) is True + + reconnect.assert_awaited_once_with(session, request_state=request_state, require_same_account=True) + cast(AsyncMock, session.upstream.send_text).assert_awaited_once_with(request_state.request_text) + + +@pytest.mark.asyncio +async def test_retry_http_bridge_clean_close_allows_one_additional_retry(monkeypatch): + settings = _make_proxy_settings() + monkeypatch.setattr(proxy_http_bridge_request_submit, "_HTTP_BRIDGE_CLEAN_CLOSE_RETRY_MAX_COUNT", 1) + monkeypatch.setattr(proxy_http_bridge_request_submit, "_HTTP_BRIDGE_CLEAN_CLOSE_RETRY_JITTER_MAX_SECONDS", 0.0) + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + request_state = proxy_service._WebSocketRequestState( + request_id="req_bridge_clean_close_second_retry", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + awaiting_response_created=True, + request_text='{"type":"response.create"}', + ) + upstream = AsyncMock() + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("session_header", "bridge-clean-close-second-retry", None), + headers={}, + affinity=proxy_service._AffinityPolicy(), + request_model="gpt-5.1", + account=_make_account("acc_bridge_clean_close_second_retry"), + upstream=upstream, + 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, + last_upstream_close_code=1000, + ) + reconnect = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + assert await service._retry_http_bridge_precreated_request(session) is True + assert request_state.replay_count == 1 + session.last_upstream_close_code = 1000 + session.last_upstream_close_generation += 1 + + assert await service._retry_http_bridge_precreated_request(session) is True + assert request_state.replay_count == 2 + assert request_state.clean_close_replay_count == 1 + assert reconnect.await_count == 2 + assert upstream.send_text.await_count == 2 + + assert await service._retry_http_bridge_precreated_request(session) is False + assert reconnect.await_count == 2 + + @pytest.mark.asyncio async def test_retry_http_bridge_precreated_request_suppresses_retry_after_response_event(monkeypatch): request_logs = _RequestLogsRecorder() @@ -36820,11 +37426,18 @@ async def test_retry_http_bridge_precreated_request_migrates_only_safe_initial_t assert await service._retry_http_bridge_precreated_request(session) is True - reconnect.assert_awaited_once_with( - session, - request_state=request_state, - require_same_account=False, - ) + if file_owner_bound: + reconnect.assert_awaited_once_with( + session, + request_state=request_state, + require_same_account=False, + require_preferred_account=True, + ) + else: + reconnect.assert_awaited_once_with( + session, + request_state=request_state, + ) assert request_state.preferred_account_id == expected_preferred_account_id assert request_state.excluded_account_ids == expected_excluded_account_ids assert session.upstream_turn_state == expected_turn_state @@ -36873,13 +37486,12 @@ async def test_retry_http_bridge_precreated_request_keeps_hard_session_owner_bou reconnect.assert_awaited_once_with( session, request_state=request_state, - require_same_account=True, ) assert request_state.preferred_account_id is None - assert request_state.excluded_account_ids == set() - assert session.upstream_turn_state == "hard-turn-state" - assert session.downstream_turn_state == "hard-turn-state" - assert session.headers["x-codex-turn-state"] == "hard-turn-state" + assert request_state.excluded_account_ids == {account.id} + assert session.upstream_turn_state is None + assert session.downstream_turn_state is None + assert "x-codex-turn-state" not in session.headers def test_websocket_safe_headers_clear_stale_turn_state_when_replacement_has_none() -> None: @@ -37264,7 +37876,6 @@ async def capture_send_text(_text: str) -> None: reconnect.assert_awaited_once_with( session, request_state=request_state, - require_same_account=False, ) send_text.assert_awaited_once_with('{"type":"response.create","model":"gpt-5.1","input":"retry"}') assert send_request_ids == ["archive_bridge_created_no_output"] @@ -38019,6 +38630,7 @@ def make_state(request_id: str) -> "proxy_service._WebSocketRequestState": @pytest.mark.asyncio async def test_submit_http_bridge_request_reinlines_final_text(monkeypatch): service = proxy_service.ProxyService.__new__(proxy_service.ProxyService) + proxy_service._initialize_http_bridge_retry_circuit(service) original_text = json.dumps( { "type": "response.create", @@ -38106,6 +38718,7 @@ async def capture_send_text(_text: str) -> None: @pytest.mark.asyncio async def test_submit_http_bridge_network_send_failure_is_neutral_and_not_replayed(monkeypatch): service = proxy_service.ProxyService.__new__(proxy_service.ProxyService) + proxy_service._initialize_http_bridge_retry_circuit(service) request_state = proxy_service._WebSocketRequestState( request_id="req_submit_network_failure", model="gpt-5.5", @@ -38179,6 +38792,7 @@ async def cleanup(*_args: object, **_kwargs: object) -> None: @pytest.mark.asyncio async def test_submit_http_bridge_request_checks_queue_before_inlining(monkeypatch): service = proxy_service.ProxyService.__new__(proxy_service.ProxyService) + proxy_service._initialize_http_bridge_retry_circuit(service) request_state = proxy_service._WebSocketRequestState( request_id="req_submit_queue_full_inline", model="gpt-5.5", diff --git a/tests/unit/test_proxy_websocket_client.py b/tests/unit/test_proxy_websocket_client.py index 034d0ca6c2..6dcd3af3f5 100644 --- a/tests/unit/test_proxy_websocket_client.py +++ b/tests/unit/test_proxy_websocket_client.py @@ -52,6 +52,8 @@ class _UnexpectedHttpClient: class _FakeConnection: + connection_lost_waiter: asyncio.Future[object] + def __init__(self, *, subprotocol: str | None = None) -> None: self.sent: list[str | bytes] = [] self.closed = False @@ -67,6 +69,18 @@ async def close(self, code: int = 1000, reason: str = "") -> None: self.closed = True +@pytest.mark.asyncio +async def test_websockets_response_websocket_consumes_connection_lost_waiter_error(): + connection = _FakeConnection() + connection.connection_lost_waiter = asyncio.get_running_loop().create_future() + + WebsocketsUpstreamWebSocket(cast(Any, connection)) + connection.connection_lost_waiter.set_exception(RuntimeError("keepalive ping timeout")) + await asyncio.sleep(0) + + assert connection.connection_lost_waiter.exception() is not None + + async def _local_proxy_tunnel_handler( reader: asyncio.StreamReader, writer: asyncio.StreamWriter, diff --git a/tests/unit/test_settings_reference.py b/tests/unit/test_settings_reference.py index 04633e7c3c..8f5c828cc8 100644 --- a/tests/unit/test_settings_reference.py +++ b/tests/unit/test_settings_reference.py @@ -45,7 +45,7 @@ def _isolated_settings(**overrides: Any) -> Settings: # number when fields are removed; never raise it without a simplicity-budget # discussion — every new CODEX_LB_* setting needs a why-not-a-default # justification per CONTRIBUTING.md's simplicity gates. -MAX_SETTINGS_FIELDS = 115 +MAX_SETTINGS_FIELDS = 116 def test_generated_settings_reference_matches_code() -> None: diff --git a/tests/unit/test_sse.py b/tests/unit/test_sse.py index a0f7945e1e..69a1e97a53 100644 --- a/tests/unit/test_sse.py +++ b/tests/unit/test_sse.py @@ -3,6 +3,7 @@ import asyncio import json from collections.abc import AsyncIterator +from typing import Any, cast import pytest @@ -50,10 +51,38 @@ async def test_inject_sse_keepalives_no_pings_when_source_is_fast(): @pytest.mark.asyncio async def test_inject_sse_keepalives_emits_pings_on_idle_gap(): - out = [chunk async for chunk in inject_sse_keepalives(_slow_agen(["a\n\n"], delay=0.25), 0.05)] + callbacks: list[str] = [] + out = [ + chunk + async for chunk in inject_sse_keepalives( + _slow_agen(["a\n\n"], delay=0.25), + 0.05, + on_keepalive=lambda: callbacks.append("sent"), + ) + ] assert out[-1] == "a\n\n" assert SSE_KEEPALIVE_FRAME in out assert out.count(SSE_KEEPALIVE_FRAME) >= 2 + assert len(callbacks) == out.count(SSE_KEEPALIVE_FRAME) + + +@pytest.mark.asyncio +async def test_inject_sse_keepalives_cancels_idle_source_when_downstream_closes(): + source_cancelled = asyncio.Event() + + async def source() -> AsyncIterator[str]: + try: + await asyncio.Event().wait() + finally: + source_cancelled.set() + yield "" # pragma: no cover - keeps this a pending async generator + + stream = inject_sse_keepalives(source(), 0.01) + assert await anext(stream) == SSE_KEEPALIVE_FRAME + + await cast(Any, stream).aclose() + + assert source_cancelled.is_set() @pytest.mark.asyncio diff --git a/tests/unit/test_sticky_session_cleanup_scheduler.py b/tests/unit/test_sticky_session_cleanup_scheduler.py index 51d8ec870c..da18e40f36 100644 --- a/tests/unit/test_sticky_session_cleanup_scheduler.py +++ b/tests/unit/test_sticky_session_cleanup_scheduler.py @@ -58,6 +58,7 @@ async def test_cleanup_once_purges_prompt_cache_only(monkeypatch) -> None: bridge_repo = AsyncMock() bridge_repo.purge_closed_before = AsyncMock(return_value=2) bridge_repo.purge_abandoned_before = AsyncMock(return_value=1) + bridge_repo.purge_retry_circuits_before = AsyncMock(return_value=3) ring_service = AsyncMock() ring_service.purge_stale_before = AsyncMock(return_value=0) @@ -88,6 +89,7 @@ async def __aexit__(self, *args): sticky_repo.purge_before.assert_not_called() bridge_repo.purge_closed_before.assert_called_once() bridge_repo.purge_abandoned_before.assert_called_once() + bridge_repo.purge_retry_circuits_before.assert_called_once() ring_service.purge_stale_before.assert_called_once() @@ -114,6 +116,7 @@ async def test_cleanup_once_skips_bridge_purge_when_schema_is_not_ready(monkeypa bridge_repo = AsyncMock() bridge_repo.purge_closed_before = AsyncMock(return_value=0) bridge_repo.purge_abandoned_before = AsyncMock(return_value=0) + bridge_repo.purge_retry_circuits_before = AsyncMock(return_value=0) ring_service = AsyncMock() ring_service.purge_stale_before = AsyncMock(return_value=0) @@ -148,6 +151,7 @@ async def __aexit__(self, *args): sticky_repo.purge_prompt_cache_before.assert_called_once() bridge_repo.purge_closed_before.assert_not_called() bridge_repo.purge_abandoned_before.assert_not_called() + bridge_repo.purge_retry_circuits_before.assert_not_called() ring_service.purge_stale_before.assert_called_once() @@ -174,6 +178,7 @@ async def test_cleanup_once_purges_bridge_when_schema_exists_after_startup_flag_ bridge_repo = AsyncMock() bridge_repo.purge_closed_before = AsyncMock(return_value=1) bridge_repo.purge_abandoned_before = AsyncMock(return_value=0) + bridge_repo.purge_retry_circuits_before = AsyncMock(return_value=0) ring_service = AsyncMock() ring_service.purge_stale_before = AsyncMock(return_value=2) @@ -204,6 +209,7 @@ async def __aexit__(self, *args): sticky_repo.purge_prompt_cache_before.assert_called_once() bridge_repo.purge_closed_before.assert_called_once() bridge_repo.purge_abandoned_before.assert_called_once() + bridge_repo.purge_retry_circuits_before.assert_called_once() ring_service.purge_stale_before.assert_called_once() @@ -259,6 +265,7 @@ async def test_cleanup_once_gates_abandoned_purge_on_prompt_cache_reuse_ttl(monk bridge_repo = AsyncMock() bridge_repo.purge_closed_before = AsyncMock(return_value=0) bridge_repo.purge_abandoned_before = AsyncMock(return_value=0) + bridge_repo.purge_retry_circuits_before = AsyncMock(return_value=0) ring_service = AsyncMock() ring_service.purge_stale_before = AsyncMock(return_value=0)