diff --git a/.all-contributorsrc b/.all-contributorsrc index c848e32026..6e631698b4 100644 --- a/.all-contributorsrc +++ b/.all-contributorsrc @@ -1084,50 +1084,10 @@ ] }, { - "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" - ] - }, - { - "login": "lkraider", - "name": "Paul Eipper", - "avatar_url": "https://avatars.githubusercontent.com/u/52256?v=4", - "profile": "https://github.com/lkraider", - "contributions": [ - "code", - "test", - "doc" - ] - }, - { - "login": "ret2basic", - "name": "ret2basic.eth", - "avatar_url": "https://avatars.githubusercontent.com/u/59381775?v=4", - "profile": "https://github.com/ret2basic", - "contributions": [ - "code", - "test" - ] - }, - { - "login": "yeongjun-cigro", - "name": "유영준", - "avatar_url": "https://avatars.githubusercontent.com/u/260819931?v=4", - "profile": "https://github.com/yeongjun-cigro", + "login": "glopyglerky", + "name": "glopyglerky", + "avatar_url": "https://avatars.githubusercontent.com/u/189872235?v=4", + "profile": "https://github.com/glopyglerky", "contributions": [ "code", "test" diff --git a/app/core/balancer/__init__.py b/app/core/balancer/__init__.py index 2e29cad9db..5b69cf3131 100644 --- a/app/core/balancer/__init__.py +++ b/app/core/balancer/__init__.py @@ -11,6 +11,7 @@ ROUTING_POLICY_PRESERVE, TRAFFIC_CLASS_FOREGROUND, TRAFFIC_CLASS_OPPORTUNISTIC, + USAGE_LIMIT_REACHED, AccountState, FailoverAction, ResetPreferenceWindow, @@ -28,6 +29,7 @@ handle_quota_exceeded, handle_rate_limit, plausible_rate_limit_reset_at, + pool_usage_exhaustion, select_account, ) @@ -52,6 +54,7 @@ "RoutingStrategy", "TrafficClass", "SelectionResult", + "USAGE_LIMIT_REACHED", "UsageWeightedOrder", "account_status_for_permanent_failure", "configure_replica_salt", @@ -61,5 +64,6 @@ "handle_quota_exceeded", "handle_rate_limit", "plausible_rate_limit_reset_at", + "pool_usage_exhaustion", "select_account", ] diff --git a/app/core/balancer/logic.py b/app/core/balancer/logic.py index a3caf85afa..63303f19a6 100644 --- a/app/core/balancer/logic.py +++ b/app/core/balancer/logic.py @@ -145,6 +145,90 @@ class AccountState: class SelectionResult: account: AccountState | None error_message: str | None + error_code: str | None = None + resets_at: int | None = None + + +USAGE_LIMIT_REACHED = "usage_limit_reached" + + +def pool_usage_exhaustion( + states: Iterable[AccountState], + *, + current: float, + ignore_standard_quota: bool = False, + ignore_standard_quota_account_ids: Collection[str] | None = None, +) -> SelectionResult | None: + """Describe pool-wide subscription exhaustion without parsing retry text.""" + if ignore_standard_quota: + return None + ignored_account_ids = set(ignore_standard_quota_account_ids or ()) + + def _primary_usage_evidence(state: AccountState) -> float | None: + return state.priority_used_percent if state.priority_used_percent is not None else state.used_percent + + def _secondary_usage_evidence(state: AccountState) -> float | None: + if state.limit_scoped_usage and state.priority_secondary_used_percent is None: + return _primary_usage_evidence(state) + return ( + state.priority_secondary_used_percent + if state.priority_secondary_used_percent is not None + else state.secondary_used_percent + ) + + def _usage_exhausted(state: AccountState) -> bool: + if state.status not in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED): + return False + usage_values = (_primary_usage_evidence(state), _secondary_usage_evidence(state)) + return any(value is not None and float(value) >= 100.0 for value in usage_values) + + def _usage_exhausted_reset_at(state: AccountState) -> float | None: + if state.status not in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED): + return None + candidates: list[float] = [] + primary_evidence = _primary_usage_evidence(state) + secondary_evidence = _secondary_usage_evidence(state) + if primary_evidence is not None and float(primary_evidence) >= 100.0 and state.primary_reset_at is not None: + candidates.append(float(state.primary_reset_at)) + if ( + secondary_evidence is not None + and float(secondary_evidence) >= 100.0 + and state.secondary_reset_at is not None + ): + candidates.append(float(state.secondary_reset_at)) + if not candidates: + return None + return max(candidates) + + eligible = [ + state + for state in states + if not state.ignore_standard_quota + and state.account_id not in ignored_account_ids + and state.status + not in ( + AccountStatus.PAUSED, + AccountStatus.REAUTH_REQUIRED, + AccountStatus.DEACTIVATED, + ) + ] + if not eligible or any(not _usage_exhausted(state) for state in eligible): + return None + + # Only usage-proven exhausted accounts reach this branch; surface the + # earliest reset for an actually exhausted window so the structured 429 + # does not retry a secondary-window exhaustion at the primary reset time. + reset_candidates = [reset_at for state in eligible if (reset_at := _usage_exhausted_reset_at(state)) is not None] + resets_at = int(min(reset_candidates)) if reset_candidates else None + message = "Usage limit reached" + if resets_at is not None: + message = _format_retry_hint(max(0.0, resets_at - current)) + return SelectionResult( + account=None, + error_message=message, + error_code=USAGE_LIMIT_REACHED, + resets_at=resets_at, + ) @dataclass(frozen=True, slots=True) @@ -379,6 +463,8 @@ def select_account( primary_first_usage_weighted: bool = False, routing_costs: RoutingCostsByAccount | None = None, replica_salt: str | None = None, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> SelectionResult: """Select an eligible account by applying availability checks and routing strategy. @@ -446,6 +532,7 @@ def select_account( available: list[AccountState] = [] in_error_backoff: list[AccountState] = [] all_states = list(states) + usage_exhaustion_state_list = list(usage_exhaustion_states) if usage_exhaustion_states is not None else all_states bypass_account_ids = None if bypass_quota_exceeded_account_ids is None else set(bypass_quota_exceeded_account_ids) for state in all_states: @@ -529,6 +616,15 @@ def _backoff_expires_at(s: AccountState) -> float: return SelectionResult(None, f"opportunistic burn window closed: {reason}") available = opportunistic_available else: + if allow_usage_exhaustion_error: + usage_exhaustion = pool_usage_exhaustion( + usage_exhaustion_state_list, + current=current, + ignore_standard_quota=ignore_standard_quota or bypass_quota_exceeded, + ignore_standard_quota_account_ids=bypass_account_ids, + ) + if usage_exhaustion is not None: + return usage_exhaustion reauth_required = [s for s in all_states if s.status == AccountStatus.REAUTH_REQUIRED] deactivated = [s for s in all_states if s.status == AccountStatus.DEACTIVATED] paused = [s for s in all_states if s.status == AccountStatus.PAUSED] diff --git a/app/core/errors.py b/app/core/errors.py index 269d142d4c..395470dbef 100644 --- a/app/core/errors.py +++ b/app/core/errors.py @@ -47,8 +47,17 @@ class ResponseFailedEvent(TypedDict): PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE = "Previous response was not found; retry without previous_response_id." -def openai_error(code: str, message: str, error_type: str = "server_error") -> OpenAIErrorEnvelope: - return {"error": {"message": message, "type": error_type, "code": code}} +def openai_error( + code: str, + message: str, + error_type: str = "server_error", + *, + resets_at: int | float | None = None, +) -> OpenAIErrorEnvelope: + detail: OpenAIErrorDetail = {"message": message, "type": error_type, "code": code} + if resets_at is not None: + detail["resets_at"] = int(resets_at) + return {"error": detail} def dashboard_error(code: str, message: str) -> DashboardErrorEnvelope: @@ -105,9 +114,10 @@ def response_failed_event( response_id: str | None = None, created_at: int | None = None, error_param: str | None = None, + resets_at: int | float | None = None, incomplete_details: dict[str, str] | None = None, ) -> ResponseFailedEvent: - error = openai_error(code, message, error_type)["error"] + error = openai_error(code, message, error_type, resets_at=resets_at)["error"] if error_param: error["param"] = error_param if created_at is None: diff --git a/app/core/openai/requests.py b/app/core/openai/requests.py index 896287fbaa..654b9a9d28 100644 --- a/app/core/openai/requests.py +++ b/app/core/openai/requests.py @@ -42,9 +42,9 @@ _ASSISTANT_TEXT_PART_TYPES = frozenset({"text", "input_text", "output_text"}) _TOOL_TEXT_PART_TYPES = frozenset({"text", "input_text", "output_text", "refusal"}) _COMPACT_STATE_TOOL_NAMES = frozenset({"create_goal", "get_goal", "update_goal", "update_plan"}) -_COMPACT_TOOL_CALL_ITEM_TYPES = frozenset({"function_call", "custom_tool_call", "apply_patch_call"}) +_COMPACT_TOOL_CALL_ITEM_TYPES = frozenset({"function_call", "custom_tool_call", "apply_patch_call", "tool_search_call"}) _COMPACT_TOOL_CALL_OUTPUT_ITEM_TYPES = frozenset( - {"function_call_output", "custom_tool_call_output", "apply_patch_call_output"} + {"function_call_output", "custom_tool_call_output", "apply_patch_call_output", "tool_search_output"} ) _EXPLICIT_PROMPT_CACHE_CONTENT_TYPES = frozenset({"input_text", "input_image", "input_file"}) _GOAL_CONTINUATION_CONTEXT_PREFIX = '' diff --git a/app/db/alembic/versions/20260722_000000_add_security_lineage_persistence.py b/app/db/alembic/versions/20260722_000000_add_security_lineage_persistence.py new file mode 100644 index 0000000000..ee9b541e03 --- /dev/null +++ b/app/db/alembic/versions/20260722_000000_add_security_lineage_persistence.py @@ -0,0 +1,270 @@ +"""Reconcile durable security-lineage persistence without a second head. + +Revision ID: 20260722_000000_add_security_lineage_persistence +Revises: 20260720_000000_add_request_log_conversation_id +Create Date: 2026-07-22 00:00:00.000000 +""" + +from __future__ import annotations + +from collections.abc import Mapping +from hashlib import pbkdf2_hmac + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.engine import Connection + +revision = "20260722_000000_add_security_lineage_persistence" +down_revision = "20260720_000000_add_request_log_conversation_id" +branch_labels = None +depends_on = None + +_MARKER_PREFIX = "@security-work/v2/" +_LEGACY_MARKER_PREFIX = "security-work:" +_CODEX_SESSION_KIND = "codex_session" +_LINEAGE_ALIAS_KINDS = ("session_header", "turn_state") +_ANONYMOUS_SCOPE = "__anonymous__" +_BATCH_NAMING_CONVENTION = { + "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", +} + + +def _identifier_digest(value: str) -> str: + return pbkdf2_hmac("sha256", value.encode(), b"codex-lb-marker-v1", 120_000).hex() + + +def _columns(connection: Connection, table_name: str) -> dict[str, Mapping[str, object]]: + inspector = sa.inspect(connection) + if not inspector.has_table(table_name): + return {} + return {str(column["name"]): column for column in inspector.get_columns(table_name) if column.get("name")} + + +def _account_foreign_key(connection: Connection, table_name: str) -> Mapping[str, object] | None: + inspector = sa.inspect(connection) + if not inspector.has_table(table_name): + return None + for foreign_key in inspector.get_foreign_keys(table_name): + if foreign_key.get("constrained_columns") == ["account_id"] and foreign_key.get("referred_table") == "accounts": + return foreign_key + return None + + +def _marker_key(lineage_id: str, api_key_scope: str | None) -> str: + scope = (api_key_scope or "").strip() or _ANONYMOUS_SCOPE + digest = _identifier_digest(f"{scope}\0{lineage_id}") + return f"{_MARKER_PREFIX}{digest}" + + +def _legacy_marker_key(lineage_id: str) -> str: + return f"{_MARKER_PREFIX}{_identifier_digest(lineage_id)}" + + +def _insert_marker(bind: Connection, marker_key: str) -> None: + bind.execute( + sa.text( + """ + INSERT INTO sticky_sessions ( + key, kind, account_id, requires_security_work_authorized, created_at, updated_at + ) + SELECT :key, :kind, NULL, :required, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP + WHERE NOT EXISTS ( + SELECT 1 FROM sticky_sessions WHERE key = :key AND kind = :kind + ) + """ + ), + {"key": marker_key, "kind": _CODEX_SESSION_KIND, "required": True}, + ) + + +def _backfill_marker(bind: Connection, lineage_id: str, api_key_scope: str | None, *, legacy: bool = False) -> None: + if lineage_id.startswith(_MARKER_PREFIX): + bind.execute( + sa.text( + """ + UPDATE sticky_sessions + SET account_id = NULL, requires_security_work_authorized = :required, updated_at = CURRENT_TIMESTAMP + WHERE key = :key AND kind = :kind + """ + ), + {"key": lineage_id, "kind": _CODEX_SESSION_KIND, "required": True}, + ) + return + if lineage_id.startswith(_LEGACY_MARKER_PREFIX): + lineage_id = lineage_id.removeprefix(_LEGACY_MARKER_PREFIX) + _insert_marker(bind, _marker_key(lineage_id, api_key_scope)) + if legacy: + _insert_marker(bind, _legacy_marker_key(lineage_id)) + + +def _backfill_detached_markers(bind: Connection) -> None: + sticky_columns = _columns(bind, "sticky_sessions") + required_sticky = {"key", "kind", "account_id", "requires_security_work_authorized"} + if required_sticky.issubset(sticky_columns): + rows = bind.execute( + sa.text( + """ + SELECT key FROM sticky_sessions + WHERE kind = :kind AND account_id IS NOT NULL AND requires_security_work_authorized = :required + """ + ), + {"kind": _CODEX_SESSION_KIND, "required": True}, + ).fetchall() + for (lineage_id,) in rows: + if isinstance(lineage_id, str): + _backfill_marker(bind, lineage_id, None, legacy=True) + + bridge_columns = _columns(bind, "http_bridge_sessions") + required_bridge = {"session_key_kind", "session_key_value", "api_key_scope", "requires_security_work_authorized"} + if not required_bridge.issubset(bridge_columns): + return + turn_state = "latest_turn_state" if "latest_turn_state" in bridge_columns else "NULL" + rows = bind.execute( + sa.text( + f""" + SELECT session_key_kind, session_key_value, api_key_scope, {turn_state} AS latest_turn_state + FROM http_bridge_sessions + WHERE requires_security_work_authorized = :required + """ + ), + {"required": True}, + ).fetchall() + for kind, value, scope, latest_turn_state in rows: + if kind in _LINEAGE_ALIAS_KINDS and isinstance(value, str): + _backfill_marker(bind, value, scope if isinstance(scope, str) else None) + if isinstance(latest_turn_state, str): + _backfill_marker(bind, latest_turn_state, scope if isinstance(scope, str) else None) + + alias_columns = _columns(bind, "http_bridge_session_aliases") + required_alias = {"session_id", "alias_kind", "alias_value"} + if not required_alias.issubset(alias_columns) or "id" not in bridge_columns: + return + if "api_key_scope" in alias_columns: + alias_scope = "COALESCE(a.api_key_scope, s.api_key_scope, :anonymous_scope)" + else: + alias_scope = "COALESCE(s.api_key_scope, :anonymous_scope)" + alias_rows = bind.execute( + sa.text( + f""" + SELECT a.alias_value, {alias_scope} AS api_key_scope + FROM http_bridge_session_aliases AS a + JOIN http_bridge_sessions AS s ON s.id = a.session_id + WHERE s.requires_security_work_authorized = :required + AND a.alias_kind IN :alias_kinds + """ + ).bindparams(sa.bindparam("alias_kinds", expanding=True)), + { + "required": True, + "alias_kinds": list(_LINEAGE_ALIAS_KINDS), + "anonymous_scope": _ANONYMOUS_SCOPE, + }, + ).fetchall() + for alias_value, scope in alias_rows: + if isinstance(alias_value, str): + _backfill_marker(bind, alias_value, scope if isinstance(scope, str) else None) + + +def _add_columns(bind: Connection) -> None: + usage = _columns(bind, "usage_history") + if usage: + with op.batch_alter_table("usage_history") as batch: + if "requires_security_work_authorized" not in usage: + batch.add_column( + sa.Column( + "requires_security_work_authorized", sa.Boolean(), nullable=False, server_default=sa.false() + ) + ) + if not bool(usage.get("account_id", {}).get("nullable", False)): + batch.alter_column("account_id", existing_type=sa.String(), nullable=True) + if bind.dialect.name == "sqlite": + # SQLite batch-alter rebuilds the table and does not preserve its + # expression indexes, which are required by the usage hot path. + op.execute(sa.text("DROP INDEX IF EXISTS idx_usage_window_account_latest")) + op.execute(sa.text("DROP INDEX IF EXISTS idx_usage_window_account_time")) + op.execute( + sa.text( + "CREATE INDEX idx_usage_window_account_latest " + "ON usage_history (coalesce(\"window\", 'primary'), account_id, recorded_at DESC, id DESC)" + ) + ) + op.execute( + sa.text( + "CREATE INDEX idx_usage_window_account_time " + "ON usage_history (coalesce(\"window\", 'primary'), account_id, recorded_at DESC)" + ) + ) + + sticky = _columns(bind, "sticky_sessions") + if sticky: + account_foreign_key = _account_foreign_key(bind, "sticky_sessions") + raw_account_foreign_key_options = account_foreign_key.get("options") if account_foreign_key else None + account_foreign_key_options = ( + raw_account_foreign_key_options if isinstance(raw_account_foreign_key_options, Mapping) else {} + ) + replace_account_foreign_key = ( + account_foreign_key is None or str(account_foreign_key_options.get("ondelete", "")).upper() != "SET NULL" + ) + with op.batch_alter_table( + "sticky_sessions", + naming_convention=_BATCH_NAMING_CONVENTION, + ) as batch: + if "requires_security_work_authorized" not in sticky: + batch.add_column( + sa.Column( + "requires_security_work_authorized", sa.Boolean(), nullable=False, server_default=sa.false() + ) + ) + if not bool(sticky.get("account_id", {}).get("nullable", False)): + batch.alter_column("account_id", existing_type=sa.String(), nullable=True) + if replace_account_foreign_key: + if account_foreign_key is not None: + constraint_name = str(account_foreign_key.get("name") or "fk_sticky_sessions_account_id_accounts") + batch.drop_constraint(constraint_name, type_="foreignkey") + batch.create_foreign_key( + "fk_sticky_sessions_account_id_accounts", + "accounts", + ["account_id"], + ["id"], + ondelete="SET NULL", + ) + + bridge = _columns(bind, "http_bridge_sessions") + if bridge: + with op.batch_alter_table("http_bridge_sessions") as batch: + if "requires_security_work_authorized" not in bridge: + batch.add_column( + sa.Column( + "requires_security_work_authorized", sa.Boolean(), nullable=False, server_default=sa.false() + ) + ) + if "latest_pending_function_call_ids" not in bridge: + batch.add_column(sa.Column("latest_pending_function_call_ids", sa.Text(), nullable=True)) + if "latest_pending_custom_tool_call_ids" not in bridge: + batch.add_column(sa.Column("latest_pending_custom_tool_call_ids", sa.Text(), nullable=True)) + + quota = _columns(bind, "quota_planner_settings") + if quota: + with op.batch_alter_table("quota_planner_settings") as batch: + if "auto_redeem_expiring_reset_credits" not in quota: + batch.add_column( + sa.Column( + "auto_redeem_expiring_reset_credits", sa.Boolean(), nullable=False, server_default=sa.false() + ) + ) + if "reset_credit_redeem_lead_minutes" not in quota: + batch.add_column( + sa.Column("reset_credit_redeem_lead_minutes", sa.Integer(), nullable=False, server_default="30") + ) + + +def upgrade() -> None: + bind = op.get_bind() + _add_columns(bind) + _backfill_detached_markers(bind) + + +def downgrade() -> None: + # This revision reconciles columns and detached markers that may have been + # created by a previous aggregate. Their original owner cannot be inferred, + # so dropping them could destroy live lineage data. + return diff --git a/app/db/alembic/versions/20260728_000000_merge_security_lineage_and_pending_tool_calls_heads.py b/app/db/alembic/versions/20260728_000000_merge_security_lineage_and_pending_tool_calls_heads.py new file mode 100644 index 0000000000..eab4dc4ec7 --- /dev/null +++ b/app/db/alembic/versions/20260728_000000_merge_security_lineage_and_pending_tool_calls_heads.py @@ -0,0 +1,22 @@ +"""Merge security-lineage and pending tool call manifest heads. + +Revision ID: 20260728_000000_merge_security_lineage_and_pending_tool_calls_heads +Revises: 20260722_000000_add_security_lineage_persistence, 20260725_000000_add_http_bridge_pending_tool_calls +Create Date: 2026-07-28 +""" + +revision = "20260728_000000_merge_security_lineage_and_pending_tool_calls_heads" +down_revision = ( + "20260722_000000_add_security_lineage_persistence", + "20260725_000000_add_http_bridge_pending_tool_calls", +) +branch_labels = None +depends_on = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/app/db/alembic/versions/20260729_000000_drop_legacy_bridge_pending_tool_columns.py b/app/db/alembic/versions/20260729_000000_drop_legacy_bridge_pending_tool_columns.py new file mode 100644 index 0000000000..4c52540970 --- /dev/null +++ b/app/db/alembic/versions/20260729_000000_drop_legacy_bridge_pending_tool_columns.py @@ -0,0 +1,55 @@ +"""Drop legacy split pending tool call columns. + +Revision ID: 20260729_000000_drop_legacy_bridge_pending_tool_columns +Revises: 20260728_000000_merge_security_lineage_and_pending_tool_calls_heads +Create Date: 2026-07-29 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.engine import Connection + +revision = "20260729_000000_drop_legacy_bridge_pending_tool_columns" +down_revision = "20260728_000000_merge_security_lineage_and_pending_tool_calls_heads" +branch_labels = None +depends_on = None + +_TABLE = "http_bridge_sessions" +_CURRENT_COLUMN = "latest_pending_tool_calls_json" +_LEGACY_COLUMNS = ( + "latest_pending_function_call_ids", + "latest_pending_custom_tool_call_ids", +) + + +def _columns(connection: Connection) -> set[str]: + inspector = sa.inspect(connection) + if not inspector.has_table(_TABLE): + return set() + return {str(column["name"]) for column in inspector.get_columns(_TABLE) if column.get("name") is not None} + + +def upgrade() -> None: + bind = op.get_bind() + columns = _columns(bind) + if not columns: + return + with op.batch_alter_table(_TABLE) as batch_op: + if _CURRENT_COLUMN not in columns: + batch_op.add_column(sa.Column(_CURRENT_COLUMN, sa.Text(), nullable=True)) + for column in _LEGACY_COLUMNS: + if column in columns: + batch_op.drop_column(column) + + +def downgrade() -> None: + bind = op.get_bind() + columns = _columns(bind) + if not columns: + return + with op.batch_alter_table(_TABLE) as batch_op: + for column in _LEGACY_COLUMNS: + if column not in columns: + batch_op.add_column(sa.Column(column, sa.Text(), nullable=True)) diff --git a/app/db/alembic/versions/20260802_000000_merge_bridge_and_capability_lineage_heads.py b/app/db/alembic/versions/20260802_000000_merge_bridge_and_capability_lineage_heads.py new file mode 100644 index 0000000000..5192643059 --- /dev/null +++ b/app/db/alembic/versions/20260802_000000_merge_bridge_and_capability_lineage_heads.py @@ -0,0 +1,24 @@ +"""Merge live bridge and capability lineage migration heads. + +Revision ID: 20260802_000000_merge_bridge_and_capability_lineage_heads +Revises: 20260729_000000_drop_legacy_bridge_pending_tool_columns, 20260731_000000_add_capability_lineage_markers +Create Date: 2026-08-02 +""" + +from __future__ import annotations + +revision = "20260802_000000_merge_bridge_and_capability_lineage_heads" +down_revision = ( + "20260729_000000_drop_legacy_bridge_pending_tool_columns", + "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/alembic/versions/20260805_000000_merge_live_carrier_heads.py b/app/db/alembic/versions/20260805_000000_merge_live_carrier_heads.py new file mode 100644 index 0000000000..4bddd1d8ce --- /dev/null +++ b/app/db/alembic/versions/20260805_000000_merge_live_carrier_heads.py @@ -0,0 +1,24 @@ +"""Merge live carrier migration heads. + +Revision ID: 20260805_000000_merge_live_carrier_heads +Revises: 20260802_000000_merge_bridge_and_capability_lineage_heads, 20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads +Create Date: 2026-08-05 +""" + +from __future__ import annotations + +revision = "20260805_000000_merge_live_carrier_heads" +down_revision = ( + "20260802_000000_merge_bridge_and_capability_lineage_heads", + "20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads", +) +branch_labels = None +depends_on = None + + +def upgrade() -> None: + pass + + +def downgrade() -> None: + pass diff --git a/app/db/migrate.py b/app/db/migrate.py index 1e19350940..a7b8d4e773 100644 --- a/app/db/migrate.py +++ b/app/db/migrate.py @@ -103,7 +103,18 @@ ) _LEGACY_EXTRA_COLUMNS = frozenset( { + ("http_bridge_sessions", "requires_security_work_authorized"), + ("quota_planner_settings", "auto_redeem_expiring_reset_credits"), + ("quota_planner_settings", "reset_credit_redeem_lead_minutes"), ("request_logs", "slim_summary_json"), + ("sticky_sessions", "requires_security_work_authorized"), + ("usage_history", "requires_security_work_authorized"), + } +) +_LEGACY_NULLABLE_COLUMNS = frozenset( + { + ("sticky_sessions", "account_id"), + ("usage_history", "account_id"), } ) @@ -582,6 +593,20 @@ def _is_ignored_schema_drift(connection: Connection, diff: object) -> bool: if (str(diff[2]), str(column_name)) in _LEGACY_EXTRA_COLUMNS: return True + if diff[0] == "modify_nullable" and len(diff) >= 7: + table_name = str(diff[2]) + column_name = str(diff[3]) + if (table_name, column_name) in _LEGACY_NULLABLE_COLUMNS: + return True + + if diff[0] in {"add_fk", "remove_fk"} and len(diff) >= 2: + constraint = diff[1] + table = getattr(constraint, "table", None) + table_name = getattr(table, "name", None) + columns = {str(column.name) for column in getattr(constraint, "columns", ()) if getattr(column, "name", None)} + if table_name == "sticky_sessions" and columns == {"account_id"}: + return True + if connection.dialect.name == "sqlite" and diff[0] == "modify_type" and len(diff) >= 7: table_name = str(diff[2]) column_name = str(diff[3]) diff --git a/app/modules/proxy/_load_balancer/sticky_selection.py b/app/modules/proxy/_load_balancer/sticky_selection.py index 48d80c7b5d..1af2e3c9ca 100644 --- a/app/modules/proxy/_load_balancer/sticky_selection.py +++ b/app/modules/proxy/_load_balancer/sticky_selection.py @@ -170,6 +170,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback: bool, traffic_class: TrafficClass, ignore_standard_quota: bool, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> _StickySelectionOutcome: ... async def release_account_lease(self, lease: AccountLease | None) -> None: ... @@ -204,6 +206,7 @@ class StickySelectionRequest(Generic[SelectionInputsT]): selection_inputs: SelectionInputsT reload_inputs: Callable[[], Awaitable[SelectionInputsT]] record_account_cap_rejection: AccountCapRejectionCallback + allow_usage_exhaustion_error: bool = True @dataclass(frozen=True, slots=True) @@ -226,6 +229,7 @@ class StickySelectionOutcome(Generic[SelectionInputsT]): selected_lease: AccountLease | None error_message: str | None error_code: str | None + resets_at: int | None = None disposition: StickySelectionDisposition = "shared_result" @@ -261,11 +265,13 @@ async def run_sticky_selection_path( redact_sensitive_details = request.redact_sensitive_details load_selection_inputs = request.reload_inputs _record_account_cap_rejection = request.record_account_cap_rejection + allow_usage_exhaustion_error = request.allow_usage_exhaustion_error selected_snapshot: Account | None = None selected_lease: AccountLease | None = None error_message: str | None = None selection_error_code: str | None = None + selection_resets_at: int | None = None def _direct_error( *, @@ -405,9 +411,11 @@ def _direct_error( sticky_outcome = _StickySelectionOutcome(selection=SelectionResult(None, None)) if hard_sticky and not selection_states: selection_error_code = "hard_affinity_saturated" + selection_resets_at = None result = SelectionResult(None, "Hard affinity owner account is unavailable") elif not selection_states and states: selection_error_code = _account_cap_error_code(lease_kind) + selection_resets_at = None result = SelectionResult(None, _account_cap_error_message(lease_kind, caps)) logger.warning( "Account cap exhausted during sticky selection lease_kind=%s reason=%s candidates=%s", @@ -432,17 +440,22 @@ def _direct_error( traffic_class=traffic_class, ignore_standard_quota=False, routing_costs_by_account_id=effective_routing_costs, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=states, ) if result.account is None: selection_error_code = "hard_affinity_saturated" + selection_resets_at = None result = SelectionResult( None, result.error_message or "Hard affinity owner account is unavailable", ) else: selection_error_code = None + selection_resets_at = None else: selection_error_code = None + selection_resets_at = None try: async with owner._repo_factory() as repos: sticky_outcome = await owner._select_with_stickiness( @@ -465,8 +478,25 @@ def _direct_error( traffic_class=traffic_class, ignore_standard_quota=False, routing_costs_by_account_id=effective_routing_costs, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=states, ) result = sticky_outcome.selection + if ( + result.account is None + and result.error_code is None + and lease_kind is not None + and len(selection_states) < len(states) + and any( + state.status == AccountStatus.ACTIVE for state in states if state not in selection_states + ) + ): + selection_error_code = _account_cap_error_code(lease_kind) + result = SelectionResult( + None, + _account_cap_error_message(lease_kind, caps), + error_code=selection_error_code, + ) except BaseException: async with owner._runtime_lock: owner._release_due_probe_reservation_locked(probe_reservation) @@ -525,6 +555,8 @@ def _direct_error( probe_reservation_invalidated = True if result.account is None: error_message = result.error_message + selection_error_code = result.error_code or selection_error_code + selection_resets_at = result.resets_at or selection_resets_at elif probe_reservation_invalidated: selected = None else: @@ -856,6 +888,7 @@ def _direct_error( selected_lease=selected_lease, error_message=error_message, error_code=selection_error_code, + resets_at=selection_resets_at, ) @@ -880,6 +913,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback: bool = False, traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> _StickySelectionOutcome: if not sticky_key or not sticky_repo: return _StickySelectionOutcome( @@ -894,6 +929,8 @@ async def _select_with_stickiness( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs_by_account_id=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) ) if sticky_kind is None: @@ -1019,6 +1056,8 @@ def finish_selection( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs_by_account_id=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) pool_also_exhausted = pool_best.account is not None and ( pool_best.account.account_id == pinned.account_id @@ -1105,6 +1144,8 @@ def finish_selection( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs_by_account_id=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) if persist_fallback and chosen.account is not None and chosen.account.account_id in account_map: return finish_selection(chosen, persist_account_id=chosen.account.account_id) @@ -1308,6 +1349,8 @@ def _select_account_preferring_budget_safe( traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, routing_costs_by_account_id: RoutingCostsByAccount | None = None, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> SelectionResult: state_list = list(states) if routing_strategy not in ("sequential_drain", "reset_drain", "single_account"): @@ -1326,6 +1369,7 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=False, ) if recovery_probe.account is not None: return recovery_probe @@ -1358,6 +1402,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) best_health_states = _best_health_tier_states(state_list) @@ -1375,6 +1421,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) if burn_first.account is not None: return burn_first @@ -1398,6 +1446,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) if preferred.account is not None: return preferred @@ -1415,6 +1465,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) return select_account( state_list, @@ -1428,6 +1480,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) diff --git a/app/modules/proxy/_load_balancer/unbound_selection.py b/app/modules/proxy/_load_balancer/unbound_selection.py index 5731d018ee..11c5a5e7b3 100644 --- a/app/modules/proxy/_load_balancer/unbound_selection.py +++ b/app/modules/proxy/_load_balancer/unbound_selection.py @@ -14,7 +14,7 @@ SelectionResult, TrafficClass, ) -from app.db.models import Account +from app.db.models import Account, AccountStatus from app.modules.proxy._load_balancer.sticky_selection import ( SelectionInputsProtocol, StickySelectionOwner, @@ -70,6 +70,7 @@ class UnboundSelectionRequest(Generic[SelectionInputsT]): selection_inputs: SelectionInputsT reload_inputs: Callable[[], Awaitable[SelectionInputsT]] record_account_cap_rejection: AccountCapRejectionCallback + allow_usage_exhaustion_error: bool = True @dataclass(frozen=True, slots=True) @@ -79,6 +80,7 @@ class UnboundSelectionOutcome(Generic[SelectionInputsT]): selected_lease: AccountLease | None error_message: str | None error_code: str | None + resets_at: int | None = None disposition: str = "shared_result" @@ -105,11 +107,13 @@ async def run_unbound_selection_path( redact_sensitive_details = request.redact_sensitive_details load_selection_inputs = request.reload_inputs _record_account_cap_rejection = request.record_account_cap_rejection + allow_usage_exhaustion_error = request.allow_usage_exhaustion_error selected_snapshot: Account | None = None selected_lease: AccountLease | None = None error_message: str | None = None selection_error_code: str | None = None + selection_resets_at: int | None = None def _direct_error( *, @@ -161,6 +165,7 @@ def _direct_error( ) if not selection_states and states: selection_error_code = _account_cap_error_code(lease_kind) + selection_resets_at = None error_message = _account_cap_error_message(lease_kind, caps) result = SelectionResult(None, error_message) logger.warning( @@ -172,6 +177,7 @@ def _direct_error( _record_account_cap_rejection(lease_kind) else: selection_error_code = None + selection_resets_at = None result = _select_account_preferring_budget_safe( selection_states, prefer_earlier_reset=prefer_earlier_reset_accounts, @@ -184,7 +190,22 @@ def _direct_error( traffic_class=traffic_class, ignore_standard_quota=False, routing_costs_by_account_id=effective_routing_costs, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=states, ) + if ( + result.account is None + and result.error_code is None + and lease_kind is not None + and len(selection_states) < len(states) + and any(state.status == AccountStatus.ACTIVE for state in states if state not in selection_states) + ): + selection_error_code = _account_cap_error_code(lease_kind) + result = SelectionResult( + None, + _account_cap_error_message(lease_kind, caps), + error_code=selection_error_code, + ) probing_result_requires_reservation = _probing_result_requires_recovery_reservation( selection_states, result.account, @@ -266,6 +287,8 @@ def _direct_error( selected_snapshot.reset_at = selected_reset_at elif result.account is None: error_message = result.error_message + selection_error_code = result.error_code or selection_error_code + selection_resets_at = result.resets_at or selection_resets_at if probe_reservation_invalidated: selected_snapshot = None @@ -442,4 +465,5 @@ def _direct_error( selected_lease=selected_lease, error_message=error_message, error_code=selection_error_code, + resets_at=selection_resets_at, ) diff --git a/app/modules/proxy/_service/api_key_usage.py b/app/modules/proxy/_service/api_key_usage.py index 7312cc37e3..e00b13190c 100644 --- a/app/modules/proxy/_service/api_key_usage.py +++ b/app/modules/proxy/_service/api_key_usage.py @@ -33,6 +33,9 @@ logger = logging.getLogger("app.modules.proxy.service") _API_KEY_RESERVATION_HEARTBEAT_SECONDS = 300.0 +_STREAM_API_KEY_RELEASE_RETRY_BASE_SECONDS = 0.1 +_STREAM_API_KEY_RELEASE_RETRY_MAX_SECONDS = 5.0 +_STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY = 4 def _service_api_keys_service() -> type[ApiKeysService]: @@ -60,6 +63,7 @@ def _api_key_reservation_heartbeat_seconds() -> float: class _ApiKeyUsageServiceProtocol(Protocol): _repo_factory: ProxyRepoFactory _background_cleanup_tasks: set[asyncio.Task[None]] + _stream_api_key_release_retry_semaphore: asyncio.Semaphore def _normalize_service_tier_value(value: Any) -> str | None: @@ -464,6 +468,12 @@ def _settlement_done(done_task: asyncio.Task[bool]) -> None: api_key.id, request_id, ) + release_coro = self._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + request_id=request_id, + retry_persistence_failures=True, + ) self._schedule_cancel_safe_cleanup( _release_after_failed_settlement(), action="release_stream_api_key_reservation_after_cancelled_settlement", @@ -477,7 +487,13 @@ def _settlement_done(done_task: asyncio.Task[bool]) -> None: exc_info=(type(exc), exc, exc.__traceback__), ) else: - if not settled and release_on_failure: + if not settled: + release_coro = self._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + request_id=request_id, + retry_persistence_failures=True, + ) self._schedule_cancel_safe_cleanup( _release_after_failed_settlement(), action="release_stream_api_key_reservation_after_failed_settlement", @@ -519,21 +535,48 @@ async def _release_unsettled_stream_api_key_usage( api_key: ApiKeyData, api_key_reservation: ApiKeyUsageReservationData, request_id: str, - ) -> bool: + retry_persistence_failures: bool = False, + ) -> None: proxy = cast(_ApiKeyUsageServiceProtocol, self) - with anyio.CancelScope(shield=True): + retry_attempt = 1 + retry_delay_seconds = _STREAM_API_KEY_RELEASE_RETRY_BASE_SECONDS + while True: + retry_slot_acquired = False try: - async with proxy._repo_factory() as repos: - api_keys_service = _service_api_keys_service()(repos.api_keys) - await api_keys_service.release_usage_reservation( - api_key_reservation.reservation_id, - ) - return True + if retry_persistence_failures: + await proxy._stream_api_key_release_retry_semaphore.acquire() + retry_slot_acquired = True + with anyio.CancelScope(shield=True): + async with proxy._repo_factory() as repos: + api_keys_service = _service_api_keys_service()(repos.api_keys) + await api_keys_service.release_usage_reservation( + api_key_reservation.reservation_id, + ) + return except Exception: + if not retry_persistence_failures: + logger.warning( + "Failed to release stream API key reservation key_id=%s request_id=%s", + api_key.id, + request_id, + exc_info=True, + ) + return logger.warning( - "Failed to release stream API key reservation key_id=%s request_id=%s", + "Failed to release stream API key reservation key_id=%s request_id=%s " + "retry_attempt=%d retry_delay_seconds=%.2f", api_key.id, request_id, + retry_attempt, + retry_delay_seconds, exc_info=True, ) - return False + finally: + if retry_slot_acquired: + proxy._stream_api_key_release_retry_semaphore.release() + await asyncio.sleep(retry_delay_seconds) + retry_attempt += 1 + retry_delay_seconds = min( + _STREAM_API_KEY_RELEASE_RETRY_MAX_SECONDS, + retry_delay_seconds * 2, + ) diff --git a/app/modules/proxy/_service/codex_control.py b/app/modules/proxy/_service/codex_control.py index 1e768bf761..6879c7c3c4 100644 --- a/app/modules/proxy/_service/codex_control.py +++ b/app/modules/proxy/_service/codex_control.py @@ -36,6 +36,7 @@ from app.modules.proxy.affinity import _AffinityPolicy, _sticky_key_for_codex_control_request from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection, effective_account_concurrency_caps +from app.modules.proxy.selection_errors import selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -338,10 +339,8 @@ async def _finalize_success( if account is None: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call_control(target: Account) -> CodexControlResponse: diff --git a/app/modules/proxy/_service/compact.py b/app/modules/proxy/_service/compact.py index f8eea368f6..f6a87d32a7 100644 --- a/app/modules/proxy/_service/compact.py +++ b/app/modules/proxy/_service/compact.py @@ -58,6 +58,7 @@ AccountSelection, effective_account_concurrency_caps, ) +from app.modules.proxy.selection_errors import selection_failure_response from app.modules.proxy.work_admission import AdmissionLease, WorkAdmissionController logger = logging.getLogger("app.modules.proxy.service") @@ -897,14 +898,10 @@ async def _call_compact( else: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - status_code = 429 if log_error_code == "account_response_create_cap" else 503 + status_code, error_payload = selection_failure_response(selection) raise ProxyResponseError( status_code, - openai_error( - log_error_code, - log_error_message, - error_type="rate_limit_error" if status_code == 429 else "server_error", - ), + error_payload, ) assert account is not None account_id_value = account.id diff --git a/app/modules/proxy/_service/file_ops.py b/app/modules/proxy/_service/file_ops.py index b07696ef35..acc0e39a13 100644 --- a/app/modules/proxy/_service/file_ops.py +++ b/app/modules/proxy/_service/file_ops.py @@ -36,6 +36,7 @@ ) from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection +from app.modules.proxy.selection_errors import selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -455,10 +456,8 @@ async def _proxy_files_call( if not account: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call(target: Account) -> dict[str, JsonValue]: diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index 59eca389e0..8fe372f6aa 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -195,10 +195,7 @@ _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 -# 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_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS = 15.0 _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL = "missing_response_created_timeout" T = TypeVar("T") @@ -730,7 +727,6 @@ def _http_bridge_eventless_precreated_deadline( or sent_at is None or request_state.response_id is not None or request_state.latency_response_created_ms is not None - or request_state.response_event_count != 0 or request_state.downstream_visible or request_state.last_downstream_sequence_number is not None ): @@ -761,11 +757,12 @@ async def _close_http_bridge_session_bounded( session: "_HTTPBridgeSession", *, reason: str, + clear_continuity: bool = False, ) -> None: if session.upstream_reader is asyncio.current_task(): session.upstream_reader = None close_task = asyncio.create_task( - service._close_http_bridge_session(session), + service._close_http_bridge_session(session, clear_continuity=clear_continuity), name=f"http-bridge-close-{_hash_identifier(session.key.affinity_key)}", ) diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index 17e1877620..3573ccefa7 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -126,7 +126,6 @@ _await_cancelled_task, _call_with_supported_optional_kwargs, _estimated_lease_tokens_from_request_usage_budget, - _is_local_account_cap_code, _prefer_earlier_reset_window, _proxy_admission_wait_timeout_seconds, _raise_proxy_unavailable, @@ -207,6 +206,7 @@ DurableBridgeLookup, ) from app.modules.proxy.load_balancer import CONTINUITY_OWNER_UNAVAILABLE, AccountLease +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -225,6 +225,20 @@ class _HTTPBridgeMixin( _HTTPBridgeUpstreamEventsMixin, _HTTPBridgeServiceProtocol, ): + async def _close_http_bridge_session_bounded( + self, + session: "_HTTPBridgeSession", + *, + reason: str, + clear_continuity: bool = False, + ) -> None: + await _close_http_bridge_session_bounded( + self, + session, + reason=reason, + clear_continuity=clear_continuity, + ) + def _schedule_http_bridge_session_closes( self, sessions: list["_HTTPBridgeSession"], @@ -1642,7 +1656,7 @@ async def _close_http_bridge_session( session: "_HTTPBridgeSession", *, turn_state_lock_held: bool = False, - release_durable_session: bool = True, + clear_continuity: bool = False, ) -> None: session.closed = True if turn_state_lock_held: @@ -1665,6 +1679,7 @@ async def _close_http_bridge_session( instance_id=_service_get_settings().http_responses_session_bridge_instance_id, owner_epoch=session.durable_owner_epoch, draining=shutdown_state.is_bridge_drain_active(), + clear_continuity=clear_continuity, ) except Exception: logger.warning("Failed to release durable HTTP bridge session", exc_info=True) @@ -1781,7 +1796,6 @@ async def _create_http_bridge_session( preferred_account_id=preferred_account_id, selected_account_id=None, ) - is_local_account_cap = _is_local_account_cap_code(selection.error_code) if ( require_preferred_account and preferred_account_id is not None @@ -1789,16 +1803,8 @@ async def _create_http_bridge_session( and selection.error_code == CONTINUITY_OWNER_UNAVAILABLE ): raise _http_bridge_previous_response_owner_unavailable_error() - status_code = 429 if is_local_account_cap else 503 - error_type = "rate_limit_error" if status_code == 429 else "server_error" - raise ProxyResponseError( - status_code, - openai_error( - selection.error_code or "no_accounts", - selection.error_message or "No active accounts available", - error_type=error_type, - ), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) if require_preferred_account and preferred_account_id is not None and account.id != preferred_account_id: await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None @@ -2200,20 +2206,23 @@ def require_bound_account() -> None: ): preferred_candidate_id = None continue - 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: + if selection.error_code == USAGE_LIMIT_REACHED and ( + required_preferred_account_id is not None or hard_close_account_bound + ): + raise _http_bridge_previous_response_owner_unavailable_error() + if selection.error_code == USAGE_LIMIT_REACHED: + record_selected_account_takeover(None) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) + 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, + ): excluded_account_ids.update(request_state.excluded_account_ids) if required_preferred_account_id in excluded_account_ids: complete_failed_handoff() @@ -2238,16 +2247,8 @@ def require_bound_account() -> None: preferred_candidate_id = 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( - selection.error_code or "no_accounts", - selection.error_message or "No active accounts available", - error_type="rate_limit_error" if status_code == 429 else "server_error", - ), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) if required_preferred_account_id is not None and account.id != required_preferred_account_id: if selection.lease is not None: selected_account_lease = selection.lease diff --git a/app/modules/proxy/_service/http_bridge/request_submit.py b/app/modules/proxy/_service/http_bridge/request_submit.py index 1a211865ac..1e0c1af401 100644 --- a/app/modules/proxy/_service/http_bridge/request_submit.py +++ b/app/modules/proxy/_service/http_bridge/request_submit.py @@ -12,6 +12,7 @@ import anyio +from app.core import shutdown as shutdown_state from app.core.clients.files import create_file as core_create_file # noqa: F401 from app.core.clients.files import finalize_file as core_finalize_file # noqa: F401 from app.core.clients.proxy import CodexControlResponse as CodexControlResponse @@ -131,6 +132,7 @@ _ACCOUNT_MODEL_UNSUPPORTED_ERROR_CODE, _HARD_HTTP_BRIDGE_AFFINITY_KINDS, # noqa: F401 _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 + _WEBSOCKET_TRANSPARENT_CLOSE_MAX_REPLAYS, _clear_websocket_request_error_overrides, _copy_websocket_route_metadata_from_session, _event_type_from_payload, @@ -202,6 +204,7 @@ _REQUEST_TRANSPORT_HTTP = "http" _WEBSOCKET_AUTH_INVALIDATED_FAILURE_CODE = "account_auth_invalidated" _NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE = "no_security_work_authorized_accounts" +_HTTP_BRIDGE_MISSING_RESPONSE_CREATED_MAX_RETRIES = 1 _SECURITY_WORK_NO_AUTHORIZED_ACCOUNTS_MESSAGE = ( "Upstream flagged this request as possible cybersecurity work, but no account is marked as authorized for " "security work. codex-lb is continuing with normal account selection; the upstream request may still fail until " @@ -209,6 +212,29 @@ ) +def _http_bridge_can_replay_same_anchor_before_created(request_state: _WebSocketRequestState) -> bool: + if not request_state.request_text: + return False + if request_state.replay_count >= _WEBSOCKET_TRANSPARENT_CLOSE_MAX_REPLAYS: + return False + return ( + request_state.previous_response_id is not None + and request_state.response_id is None + and request_state.awaiting_response_created + and request_state.response_event_count == 0 + and request_state.last_downstream_sequence_number is None + and not request_state.downstream_visible + and not request_state.upstream_model_output_seen + ) + + +def _http_bridge_can_retry_missing_response_created(request_state: _WebSocketRequestState) -> bool: + return ( + request_state.missing_response_created_retry_count < _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_MAX_RETRIES + and _http_bridge_can_replay_same_anchor_before_created(request_state) + ) + + async def _rollback_http_bridge_recovery_turn_state_registration( service: Any, receipt: DurableBridgeAliasRegistrationReceipt, @@ -1772,11 +1798,7 @@ async def _retire_stale_pending_http_bridge_session( 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, - ) + clear_continuity = detail == "missing_response_created_timeout" session.closed = True async with self._http_bridge_lock: if self._http_bridge_sessions.get(session.key) is session: @@ -1788,7 +1810,22 @@ async def _retire_stale_pending_http_bridge_session( if should_close: session.upstream_close_attempted = True if should_close: - await self._close_http_bridge_session_bounded(session, reason="retire_stale_pending") + await self._close_http_bridge_session_bounded( + session, + reason="retire_stale_pending", + clear_continuity=clear_continuity, + ) + elif clear_continuity and session.durable_session_id is not None and session.durable_owner_epoch is not None: + try: + await self._durable_bridge.release_live_session( + session_id=session.durable_session_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + draining=shutdown_state.is_bridge_drain_active(), + clear_continuity=True, + ) + except Exception: + logger.warning("Failed to clear durable HTTP bridge continuity during stale retire", exc_info=True) _log_http_bridge_event( "retire_stale_pending", session.key, @@ -1880,7 +1917,7 @@ async def _retry_http_bridge_precreated_request( session: "_HTTPBridgeSession", *, request_state: _WebSocketRequestState | None = None, - restart_reader: bool = False, + allow_expired_deadline: 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( @@ -1946,6 +1983,11 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool: key=session.key.affinity_key, ) hard_owner_bound = _http_bridge_key_strength(session.key) == "hard" + now = _service_time().monotonic() + missing_created_retry = False + rebind_missing_created_owner = False + owner_rebind_affinity: _AffinityPolicy | None = None + selection_rebind_affinity: _AffinityPolicy | None = None async with session.pending_lock: if request_state is not None: if ( @@ -1953,21 +1995,55 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool: 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 request_is_retryable(request_state) + or not ( + _websocket_request_can_replay_before_visible_output(request_state) + or ( + allow_expired_deadline + and hard_owner_bound + and _http_bridge_can_retry_missing_response_created(request_state) + ) + ) + or ( + not allow_expired_deadline + and request_state.bridge_request_deadline is not None + and request_state.bridge_request_deadline <= now + ) ): return False else: retryable_requests = [ request_state for request_state in session.pending_requests - if not request_state.draining_until_terminal and request_is_retryable(request_state) + if not request_state.draining_until_terminal + and ( + _websocket_request_can_replay_before_visible_output(request_state) + or ( + allow_expired_deadline + and hard_owner_bound + and _http_bridge_can_retry_missing_response_created(request_state) + ) + ) + and ( + allow_expired_deadline + or request_state.bridge_request_deadline is None + or request_state.bridge_request_deadline > now + ) ] 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.fresh_upstream_request_is_retry_safe and request_state.fresh_upstream_request_text + 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 + ) + and not ( + allow_expired_deadline + and hard_owner_bound + and _http_bridge_can_retry_missing_response_created(request_state) + ) ): # Once a continuation is pending upstream, reconnecting without # replay cannot complete the current request, while replaying it @@ -1976,6 +2052,11 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool: # equivalent to the client's own retry once the anchor is # stripped. The latter remains pinned to the current owner. return False + missing_created_retry = ( + allow_expired_deadline + and hard_owner_bound + and _http_bridge_can_retry_missing_response_created(request_state) + ) close_classification = _classify_upstream_close( session.last_upstream_close_code, response_events_seen=request_state.response_event_count, @@ -2035,7 +2116,25 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool: request_text = _prepare_websocket_request_state_for_visible_output_replay(request_state) if request_text is None: return False - if not hard_owner_bound: + if missing_created_retry and hard_owner_bound: + request_state.preferred_account_id = None + request_state.excluded_account_ids.add(session.account.id) + request_state.affinity_policy = replace( + request_state.affinity_policy, + key=None, + kind=None, + reallocate_sticky=True, + ) + rebind_missing_created_owner = True + owner_rebind_affinity = session.affinity + selection_rebind_affinity = replace( + session.affinity, + key=None, + kind=None, + reallocate_sticky=True, + codex_session_source=None, + ) + elif not hard_owner_bound: request_state.excluded_account_ids.add(session.account.id) else: # Account-scoped uploaded files cannot be replayed on a @@ -2047,12 +2146,11 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool: return False if account_neutral_recovery: request_state.preferred_account_id = 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) + 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) + if missing_created_retry: + request_state.missing_response_created_retry_count += 1 if session.account.id in request_state.excluded_account_ids: session.upstream_turn_state = None session.downstream_turn_state = None @@ -2082,40 +2180,18 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool: ) reconnect_reader_kwargs = {"restart_reader": True} if restart_reader else {} try: - 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, + if rebind_missing_created_owner: + await _call_with_supported_optional_kwargs( + self._reconnect_http_bridge_session, + session, + optional_kwargs={ + "owner_rebind_affinity": owner_rebind_affinity, + "selection_affinity": selection_rebind_affinity, + }, + request_state=request_state, + require_same_account=account_neutral_recovery, ) - 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: + elif hard_owner_bound: await self._reconnect_http_bridge_session( session, request_state=request_state, diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index 9664e77bcb..a06c6092d8 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -75,7 +75,6 @@ _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, _http_bridge_models_compatible, _http_bridge_owner_lookup_unavailable_error_envelope, _http_bridge_payload_looks_like_full_resend, @@ -96,6 +95,7 @@ _proxy_admission_wait_timeout_seconds, _record_bridge_reattach, _record_continuity_fail_closed, + _record_http_bridge_stuck_retire, _release_http_bridge_unanchored_handoff, _release_http_bridge_unanchored_handoffs_for_request, _reserve_http_bridge_unanchored_handoff, @@ -263,20 +263,9 @@ 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.""" - +def _http_bridge_owner_failure_allows_account_neutral_replay(exc: ProxyResponseError) -> bool: code, _message = _proxy_error_code_message(exc) - return code in _HTTP_BRIDGE_AMBIGUOUS_RECOVERY_ERROR_CODES + return code in {"previous_response_owner_unavailable", "continuity_owner_conflict"} def _http_bridge_account_capacity_wait_seconds(exc: ProxyResponseError) -> float | None: @@ -980,6 +969,11 @@ def classify_durable_full_resend( replay_projection = project_responses_input_for_account_neutral_fresh_replay( cast(list[JsonValue], payload.input), stored_count=stored_count, + # Classification only: inline Responses-Lite developer IDs + # must remain visible until the exact-manifest check rejects + # response-owned messages. Cross-account replay uses the + # default ID-stripping projection below. + preserve_developer_message_ids=True, ) safe_fresh_context = False if replay_projection is not None: @@ -1366,7 +1360,7 @@ def owner_unavailable_allows_account_neutral_replay(exc: ProxyResponseError) -> nonlocal durable_full_resend_retains_prior_output if ( - not _http_bridge_is_previous_response_owner_unavailable(exc) + not _http_bridge_owner_failure_allows_account_neutral_replay(exc) or forwarded_request or rewritten_file_account_id is not None or durable_full_resend_anchor_count is None @@ -1471,12 +1465,50 @@ 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 + def switch_model_transition_to_account_neutral_fork(exc: ProxyResponseError) -> bool: + nonlocal account_neutral_recovery + nonlocal affinity + nonlocal bridge_session_key + nonlocal force_local_recovery_creation + nonlocal incoming_turn_state_header + nonlocal preferred_account_has_continuity_provenance + nonlocal request_state + nonlocal session_creation_headers + nonlocal session_header_fallback_key + + if ( + durable_model_transition_lookup is None + or not _http_bridge_owner_failure_allows_account_neutral_replay(exc) + or owner_unavailable_allows_account_neutral_replay(exc) + or request_state.previous_response_id is not None + or rewritten_file_account_id is not None + ): + return False + failed_owner_id = request_state.preferred_account_id + _log_http_bridge_event( + "model_transition_owner_conflict_fork", + bridge_session_key, + account_id=failed_owner_id, + model=effective_payload.model, + detail="outcome=retry_without_previous_model_owner", + 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, + ) + if failed_owner_id is not None: + fresh_replay_excluded_account_ids.add(failed_owner_id) + session_creation_headers = without_http_bridge_session_affinity_headers(session_creation_headers) + incoming_turn_state_header = None + session_header_fallback_key = None + affinity = _AffinityPolicy() + replay_kind, replay_key = make_http_bridge_account_neutral_replay_key(uuid4().hex) + bridge_session_key = _HTTPBridgeSessionKey(replay_kind, replay_key, bridge_session_key.api_key_id) + account_neutral_recovery = True + force_local_recovery_creation = True + request_state.preferred_account_id = None + request_state.excluded_account_ids.update(fresh_replay_excluded_account_ids) + preferred_account_has_continuity_provenance = False + return True if required_continuity_owner_missing: owner_unavailable = ProxyResponseError( @@ -1531,6 +1563,8 @@ def switch_to_account_neutral_replay() -> None: exclude_account_ids=fresh_replay_excluded_account_ids or None, ) except ProxyResponseError as exc: + if switch_model_transition_to_account_neutral_fork(exc): + continue if not owner_unavailable_allows_account_neutral_replay(exc): exc_code, _exc_message = _proxy_error_code_message(exc) if not unanchored_fork_spill_attempted and _http_bridge_unanchored_fork_can_spill_on_cap( @@ -2962,8 +2996,40 @@ async def startup_continuity_cooldown_terminal_event() -> str | None: yield line finally: if gate_contention: + should_retire_stale_gate = False async with session.pending_lock: session.queued_request_count = max(0, session.queued_request_count - 1) + retire_after_seconds = float( + getattr( + _service_get_settings(), + "http_responses_session_bridge_stuck_gate_retire_after_seconds", + 300.0, + ) + ) + now = _service_time().monotonic() + should_retire_stale_gate = any( + pending_request is not request_state + and pending_request.transport == _REQUEST_TRANSPORT_HTTP + and pending_request.response_create_gate_acquired + and pending_request.response_create_gate is session.response_create_gate + and pending_request.response_create_sent_at is None + and pending_request.response_id is None + and pending_request.response_event_count == 0 + and not pending_request.downstream_visible + and pending_request.last_downstream_sequence_number is None + and now - pending_request.started_at >= retire_after_seconds + for pending_request in session.pending_requests + ) + if should_retire_stale_gate and not session.closed: + session.closed = True + _record_http_bridge_stuck_retire( + reason="response_create_gate_timeout_stuck_pending", + session=session, + ) + await self._retire_stale_pending_http_bridge_session( + session, + detail="response_create_gate_timeout_stuck_pending", + ) if _service_time().monotonic() >= request_deadline: raise if gate_contention and session.closed: @@ -3168,42 +3234,17 @@ def stream_idle_keepalive(*, downstream_response_id: str) -> str | None: ) ) continue - downstream_response_id = _websocket_downstream_response_id(request_state) - completed_delivery_in_progress = await completed_delivery_suppresses_idle_timeout( - downstream_response_id=downstream_response_id, - revoke_queue_if_inactive=False, + completed_delivery_scope = request_state.completed_delivery_scope + completed_delivery_in_progress = ( + completed_delivery_scope is not None and completed_delivery_scope.active ) - if not completed_delivery_in_progress: + if completed_delivery_in_progress: + keepalive_count = 0 + else: keepalive_count += 1 - if not completed_delivery_in_progress and keepalive_count >= max_keepalive_count: + downstream_response_id = _websocket_downstream_response_id(request_state) + if not completed_delivery_in_progress and 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: - if await completed_delivery_suppresses_idle_timeout( - downstream_response_id=downstream_response_id, - revoke_queue_if_inactive=True, - ): - keepalive_event = stream_idle_keepalive( - downstream_response_id=downstream_response_id, - ) - if keepalive_event is not None: - yield keepalive_event - continue - 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 ) diff --git a/app/modules/proxy/_service/http_bridge/upstream_events.py b/app/modules/proxy/_service/http_bridge/upstream_events.py index 91e08cf6a2..08c0517097 100644 --- a/app/modules/proxy/_service/http_bridge/upstream_events.py +++ b/app/modules/proxy/_service/http_bridge/upstream_events.py @@ -831,42 +831,14 @@ async def _relay_http_bridge_upstream_messages( if not expired_owner: continue pending_count = len(session.pending_requests) - for request_state in session.pending_requests: - if request_state.failure_phase_override is None: - request_state.failure_phase_override = "upstream" - if request_state.failure_detail_override is None: - request_state.failure_detail_override = ( - _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL - ) 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 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, - ) + if receive_cancelled: + receive_task = None _log_http_bridge_event( "missing_response_created_timeout", session.key, @@ -879,22 +851,34 @@ 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) + retried = await self._retry_http_bridge_precreated_request( + session, + allow_expired_deadline=True, + ) if retried: continue + # Claim the session only after the safe pre-created + # retry path refuses it. Reconnecting first prevents + # a silent upstream websocket from stranding clients + # until the request budget expires. session.closed = True + _record_http_bridge_stuck_retire( + reason=_HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL, + session=session, + ) + async with session.pending_lock: + for request_state in session.pending_requests: + if request_state.failure_phase_override is None: + request_state.failure_phase_override = "upstream" + if request_state.failure_detail_override is None: + request_state.failure_detail_override = ( + _HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL + ) await self._fail_http_bridge_reader_and_maybe_retire( session, error_code="upstream_request_timeout", error_message=receive_timeout.error_message, - penalize_account=False, + penalize_account=True, retire_detail=_HTTP_BRIDGE_MISSING_RESPONSE_CREATED_TIMEOUT_DETAIL, force_retire=True, ) @@ -2049,12 +2033,6 @@ async def _process_parsed_http_bridge_upstream_event( await terminal_event_queue.put(event_block) if terminal_event_queue is not None: await terminal_event_queue.put(None) - if completed_event_queue_claimed and completed_delivery_scope is not None: - async with session.pending_lock: - # Keep the completed claim authoritative after its producer - # returns. A concurrent timeout may still be finishing - # awaited recovery work before it rechecks this scope. - completed_delivery_scope.terminal_enqueued = True if settlement_event_type in {"response.failed", "response.incomplete", "error"}: error_code = None diff --git a/app/modules/proxy/_service/streaming/mixin.py b/app/modules/proxy/_service/streaming/mixin.py index 470e7a9376..bdf81ea8c4 100644 --- a/app/modules/proxy/_service/streaming/mixin.py +++ b/app/modules/proxy/_service/streaming/mixin.py @@ -298,6 +298,7 @@ _RetryableStreamError, _StreamSettlement, _TerminalStreamError, + _TransientStreamError, _ttft_event_latency_ms, _WebSocketUpstreamControl, ) @@ -491,6 +492,7 @@ async def _stream_once( enforce_openai_sdk_contract: bool = True, ) -> AsyncIterator[str]: proxy = cast(_StreamingServiceProtocol, self) + settlement.reset() account_id_value = account.id access_token = proxy._encryptor.decrypt(account.access_token_encrypted) account_id = _header_account_id(account.chatgpt_account_id) @@ -577,6 +579,8 @@ async def _stream_once( settlement.record_success = False settlement.account_health_error = True settlement.error = {"message": error_message} + if allow_transient_retry and payload.previous_response_id is not None: + raise _TransientStreamError(error_code, settlement.error) yield format_sse_event( response_failed_event( error_code, diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 51c044c422..607516b96e 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -71,6 +71,7 @@ is_upstream_model_capacity_error, ) from app.modules.proxy.load_balancer import AccountLease, AccountSelection +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response _REQUEST_TRANSPORT_HTTP = "http" _REQUEST_TRANSPORT_WEBSOCKET = "websocket" @@ -1071,6 +1072,7 @@ async def _retry_account_model_rejection( continue if ( not account + and selection.error_code != USAGE_LIMIT_REACHED and ( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or not (propagate_http_errors and last_transient_exc is not None) @@ -1110,6 +1112,40 @@ async def _retry_account_model_rejection( account_id=last_account_model_rejection_account_id, ) return + if selection.error_code == USAGE_LIMIT_REACHED: + await _drain_pending_post_refresh_penalty_on_terminal(settlement) + no_accounts_msg = selection.error_message or "Usage limit reached" + status_code, error_payload = selection_failure_response(selection) + await proxy._write_request_log( + account_id=None, + api_key=api_key, + request_id=request_id, + model=payload.model, + latency_ms=int((time.monotonic() - start) * 1000), + status="error", + error_code=USAGE_LIMIT_REACHED, + error_message=no_accounts_msg, + reasoning_effort=payload.reasoning.effort if payload.reasoning else None, + transport=request_transport, + upstream_transport=upstream_stream_transport, + service_tier=payload.service_tier, + requested_service_tier=payload.service_tier, + useragent=useragent, + useragent_group=useragent_group, + client_ip=client_ip, + ) + if propagate_http_errors: + raise ProxyResponseError(status_code, error_payload) + yield format_sse_event( + response_failed_event( + USAGE_LIMIT_REACHED, + no_accounts_msg, + error_type=USAGE_LIMIT_REACHED, + response_id=request_id, + resets_at=selection.resets_at, + ) + ) + return if selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES: await _drain_pending_post_refresh_penalty_on_terminal(settlement) no_accounts_msg = selection.error_message or "Local account capacity is exhausted" diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index be299e8d47..09e062c3ed 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -67,6 +67,8 @@ _TTFT_OUTPUT_ITEM_TYPES = _PENDING_TOOL_CALL_ITEM_TYPES - {"function_call"} _WEBSOCKET_FULL_REPLAY_WAIT_MIN_ITEMS = 20 _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS = 0.05 +_WEBSOCKET_CREATED_ONLY_CLOSE_MAX_REPLAYS = 1 +_WEBSOCKET_TRANSPARENT_CLOSE_MAX_REPLAYS = 20 _HARD_HTTP_BRIDGE_AFFINITY_KINDS = frozenset( { "turn_state_header", @@ -750,7 +752,6 @@ class _RequestLogFailureMetadata: @dataclass(slots=True) class _HTTPBridgeCompletedDeliveryScope: active: bool = False - terminal_enqueued: bool = False @dataclass @@ -800,13 +801,7 @@ 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 + missing_response_created_retry_count: int = 0 auth_replay_count: int = 0 auth_replay_counts_by_account: dict[str, int] = field(default_factory=dict) force_refresh_account_id: str | None = None @@ -1221,13 +1216,14 @@ def _websocket_request_can_replay_before_visible_output( ) -> bool: if not request_state.request_text: return False - 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 + if request_state.replay_count >= 1: + unanchored_precreated_pending = ( + request_state.previous_response_id is None + and request_state.response_id is None + and request_state.awaiting_response_created + ) + if not unanchored_precreated_pending or request_state.replay_count >= _WEBSOCKET_TRANSPARENT_CLOSE_MAX_REPLAYS: + return False sequenced_created_only_prewarm = ( request_state.generate_false_prewarm and request_state.last_downstream_sequence_number == 0 @@ -1236,6 +1232,12 @@ def _websocket_request_can_replay_before_visible_output( and request_state.response_event_count == 1 and not request_state.downstream_visible ) + if ( + request_state.response_id is not None + and not sequenced_created_only_prewarm + and request_state.replay_count >= _WEBSOCKET_CREATED_ONLY_CLOSE_MAX_REPLAYS + ): + return False if request_state.last_downstream_sequence_number is not None and not sequenced_created_only_prewarm: return False if request_state.downstream_visible: diff --git a/app/modules/proxy/_service/transcribe.py b/app/modules/proxy/_service/transcribe.py index e2abbe4c1b..f27d08f4e6 100644 --- a/app/modules/proxy/_service/transcribe.py +++ b/app/modules/proxy/_service/transcribe.py @@ -30,6 +30,7 @@ from app.modules.proxy._service.support import _request_log_client_fields, _RequestLogFailureMetadata from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection +from app.modules.proxy.selection_errors import selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -202,10 +203,8 @@ async def transcribe( if not account: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call_transcribe(target: Account) -> dict[str, JsonValue]: diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index 604e00d76d..b476e29bdf 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -77,7 +77,6 @@ ProcessNetworkRecovery, process_network_error_code, ) -from app.core.resilience.overload import is_local_overload_error_code from app.core.types import JsonValue from app.core.upstream_proxy import UpstreamProxyRouteError from app.core.utils.request_id import get_request_id, reset_request_id, set_request_id @@ -471,6 +470,7 @@ openai_validation_error, validate_model_access, ) +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response from app.modules.proxy.tool_call_dedupe import ( mark_duplicate_tool_call_downstream_event, rewrite_parallel_tool_call_text, @@ -2554,6 +2554,8 @@ async def _select_websocket_connect_account( account = selection.account if account is not None: break + if selection.error_code == USAGE_LIMIT_REACHED: + break async def _heartbeat(remaining_seconds: float) -> None: event = _account_capacity_wait_payload( @@ -2660,18 +2662,15 @@ async def _heartbeat(remaining_seconds: float) -> None: return None if require_preferred_account and preferred_account_id is not None: if _facade()._is_local_account_cap_code(error_code): + status_code, error_payload = selection_failure_response(selection) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, account_id=preferred_account_id, api_key=api_key, request_state=request_state, - status_code=429, - payload=openai_error( - error_code, - error_message, - error_type="rate_limit_error", - ), + status_code=status_code, + payload=error_payload, error_code=error_code, error_message=error_message, ) @@ -2712,7 +2711,7 @@ async def _heartbeat(remaining_seconds: float) -> None: len(exclude_account_ids), api_key is not None, ) - status_code = 429 if is_local_overload_error_code(error_code) else 503 + status_code, error_payload = selection_failure_response(selection) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, @@ -2720,11 +2719,7 @@ async def _heartbeat(remaining_seconds: float) -> None: api_key=api_key, request_state=request_state, status_code=status_code, - payload=openai_error( - error_code, - error_message, - error_type="rate_limit_error" if status_code == 429 else "server_error", - ), + payload=error_payload, error_code=error_code, error_message=error_message, ) diff --git a/app/modules/proxy/api.py b/app/modules/proxy/api.py index b410a87f6d..f7b2236c93 100644 --- a/app/modules/proxy/api.py +++ b/app/modules/proxy/api.py @@ -262,6 +262,7 @@ WarmupSkippedAccount, WarmupSubmittedAccount, ) +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED from app.modules.proxy.types import ( CreditStatusDetailsData, RateLimitResetCreditsData, @@ -6400,6 +6401,17 @@ async def _opportunistic_admission_denial( ) if selection.account is not None: return None + if selection.error_code == USAGE_LIMIT_REACHED: + return _logged_error_json_response( + request, + 429, + openai_error( + USAGE_LIMIT_REACHED, + selection.error_message or "Usage limit reached", + error_type=USAGE_LIMIT_REACHED, + resets_at=selection.resets_at, + ), + ) message = selection.error_message or "opportunistic burn window closed" if not message.startswith("opportunistic burn window closed"): message = f"opportunistic burn window closed: {message}" diff --git a/app/modules/proxy/durable_bridge_coordinator.py b/app/modules/proxy/durable_bridge_coordinator.py index a6c125b70b..440a08b228 100644 --- a/app/modules/proxy/durable_bridge_coordinator.py +++ b/app/modules/proxy/durable_bridge_coordinator.py @@ -376,6 +376,7 @@ async def release_live_session( instance_id: str, owner_epoch: int, draining: bool, + clear_continuity: bool = False, ) -> DurableBridgeLookup | None: async with self._session() as session: snapshot = await DurableBridgeRepository(session).release_session( @@ -383,6 +384,7 @@ async def release_live_session( instance_id=instance_id, owner_epoch=owner_epoch, draining=draining, + clear_continuity=clear_continuity, ) if snapshot is None: return None diff --git a/app/modules/proxy/durable_bridge_repository.py b/app/modules/proxy/durable_bridge_repository.py index eefa4d1a45..a983e8b810 100644 --- a/app/modules/proxy/durable_bridge_repository.py +++ b/app/modules/proxy/durable_bridge_repository.py @@ -735,6 +735,7 @@ async def release_session( instance_id: str, owner_epoch: int, draining: bool, + clear_continuity: bool = False, ) -> DurableBridgeSessionSnapshot | None: """Release the lease with a single fenced UPDATE. @@ -749,6 +750,16 @@ async def release_session( "state": HttpBridgeSessionState.DRAINING if draining else HttpBridgeSessionState.CLOSED, "closed_at": None if draining else now, } + 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, + } + ) return await self._execute_fenced_session_update( session_id=session_id, instance_id=instance_id, diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 27244b522a..d4b63e6217 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -22,6 +22,7 @@ ROUTING_POLICY_PRESERVE, TRAFFIC_CLASS_FOREGROUND, TRAFFIC_CLASS_OPPORTUNISTIC, + USAGE_LIMIT_REACHED, AccountState, ResetPreferenceWindow, RoutingCostsByAccount, @@ -183,6 +184,7 @@ class AccountSelection: account: Account | None error_message: str | None error_code: str | None = None + resets_at: int | None = None lease: AccountLease | None = None catalog_omission_quota_admission: CatalogOmissionQuotaAdmission | None = None @@ -440,6 +442,7 @@ async def select_account( traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, concurrency_caps: AccountConcurrencyCaps | None = None, redact_sensitive_details: bool = False, + allow_usage_exhaustion_error: bool = True, ) -> AccountSelection: if (required_account_is_ownership_constraint or required_continuity_owner) and required_account_id is None: raise ValueError("required account ownership flags require required_account_id") @@ -574,6 +577,7 @@ async def load_selection_inputs() -> _SelectionInputs: error_message: str | None = None selected_lease: AccountLease | None = None selection_error_code: str | None = None + selection_resets_at: int | None = None legacy_existing_account_id: str | None = None if sticky_source == "session_header" and legacy_sticky_key is not None: async with self._repo_factory() as repos: @@ -638,6 +642,7 @@ async def load_selection_inputs() -> _SelectionInputs: selection_inputs=selection_inputs, reload_inputs=load_selection_inputs, record_account_cap_rejection=_record_account_cap_rejection, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, ), ) selection_inputs = unbound_outcome.selection_inputs @@ -645,6 +650,7 @@ async def load_selection_inputs() -> _SelectionInputs: selected_lease = unbound_outcome.selected_lease error_message = unbound_outcome.error_message selection_error_code = unbound_outcome.error_code + selection_resets_at = unbound_outcome.resets_at if unbound_outcome.disposition == "direct_error": return AccountSelection( account=None, @@ -682,6 +688,7 @@ async def load_selection_inputs() -> _SelectionInputs: selection_inputs=selection_inputs, reload_inputs=load_selection_inputs, record_account_cap_rejection=_record_account_cap_rejection, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, ), ) selection_inputs = sticky_outcome.selection_inputs @@ -689,6 +696,7 @@ async def load_selection_inputs() -> _SelectionInputs: selected_lease = sticky_outcome.selected_lease error_message = sticky_outcome.error_message selection_error_code = sticky_outcome.error_code + selection_resets_at = sticky_outcome.resets_at if sticky_outcome.disposition == "direct_error": return AccountSelection( account=None, @@ -735,7 +743,12 @@ async def load_selection_inputs() -> _SelectionInputs: and (selection_inputs.accounts or selection_inputs.error_code is not None) ): set_normal() - return AccountSelection(account=None, error_message=error_message, error_code=selection_error_code) + return AccountSelection( + account=None, + error_message=error_message, + error_code=selection_error_code, + resets_at=selection_resets_at, + ) if not circuit_breaker_open: set_normal() logger.info( @@ -1188,8 +1201,16 @@ async def check_opportunistic_admission( deterministic_probe=True, traffic_class=TRAFFIC_CLASS_OPPORTUNISTIC, ignore_standard_quota=False, + usage_exhaustion_states=states, ) if result.account is None: + if result.error_code == USAGE_LIMIT_REACHED: + return AccountSelection( + account=None, + error_message=result.error_message, + error_code=result.error_code, + resets_at=result.resets_at, + ) return AccountSelection( account=None, error_message=result.error_message, @@ -1410,6 +1431,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback: bool = False, traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> _StickySelectionOutcome: return await _run_select_with_stickiness( states=states, @@ -1431,6 +1454,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback=preserve_existing_mapping_on_fallback, traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) _persist_sticky_mutation = staticmethod(_persist_sticky_mutation) @@ -2267,6 +2292,7 @@ def _state_from_account( pressure_pct = inflight_pressure_pct + leased_token_pressure_pct effective_used_percent = None if used_percent is None else min(100.0, used_percent + pressure_pct) effective_secondary_used_percent = None if secondary_used is None else min(100.0, secondary_used + pressure_pct) + usage_exhaustion_evidence_status = status in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED) return AccountState( account_id=account.id, @@ -2286,6 +2312,8 @@ def _state_from_account( plan_type=account.plan_type, capacity_credits=capacity_credits, health_tier=new_tier, + priority_used_percent=used_percent if usage_exhaustion_evidence_status else None, + priority_secondary_used_percent=secondary_used if usage_exhaustion_evidence_status else None, inflight_response_creates=runtime.inflight_response_creates, inflight_streams=runtime.inflight_streams, leased_tokens=runtime.leased_tokens, diff --git a/app/modules/proxy/replay_safety.py b/app/modules/proxy/replay_safety.py index e350b1aa03..471cb5a9a1 100644 --- a/app/modules/proxy/replay_safety.py +++ b/app/modules/proxy/replay_safety.py @@ -156,8 +156,14 @@ def project_responses_input_for_account_neutral_fresh_replay( input_items: list[JsonValue], *, stored_count: int, + preserve_developer_message_ids: bool = False, ) -> AccountNeutralReplayProjection | None: - """Remove known response-owned bookkeeping after durable prefix proof.""" + """Remove known response-owned bookkeeping after durable prefix proof. + + ``preserve_developer_message_ids`` is classification-only evidence for + inline Responses-Lite messages. A projection created with that option must + not be serialized as an account-neutral replay payload. + """ if stored_count <= 0 or stored_count > len(input_items): return None @@ -165,7 +171,10 @@ def project_responses_input_for_account_neutral_fresh_replay( projected_items: list[JsonValue] = [] projected_stored_count = 0 for index, item in enumerate(input_items): - projected_item = _project_account_neutral_replay_item(item) + projected_item = _project_account_neutral_replay_item( + item, + preserve_developer_message_ids=preserve_developer_message_ids, + ) if projected_item is not None: projected_items.append(projected_item) if index + 1 == stored_count: @@ -177,7 +186,11 @@ def project_responses_input_for_account_neutral_fresh_replay( ) -def _project_account_neutral_replay_item(item: JsonValue) -> JsonValue | None: +def _project_account_neutral_replay_item( + item: JsonValue, + *, + preserve_developer_message_ids: bool, +) -> JsonValue | None: if not isinstance(item, dict): return item @@ -191,6 +204,8 @@ def _project_account_neutral_replay_item(item: JsonValue) -> JsonValue | None: if "id" not in item: return item + if preserve_developer_message_ids and item_type in (None, "message") and item.get("role") == "developer": + return item projected_item = dict(item) projected_item.pop("id") return projected_item @@ -265,11 +280,17 @@ def responses_input_suffix_retains_prior_output( return False pending_suffix_calls, seen_suffix_call_ids = prefix_state retained_output_seen = False + retained_output_is_final_answer = False fresh_followup_seen = False + fresh_followup_count = 0 + fresh_followup_is_user_message = False + fresh_developer_followup_seen = False for item in input_items[stored_count:]: - if not isinstance(item, dict): + if fresh_developer_followup_seen or not isinstance(item, dict): return False item_type_value = item.get("type") + if "type" in item and not _is_nonblank_string(item_type_value): + return False item_type = item_type_value if isinstance(item_type_value, str) else None if item_type in _TOOL_CALL_TYPES: if item.get("status") not in (None, "completed"): @@ -283,7 +304,10 @@ def responses_input_suffix_retains_prior_output( # prove that an omitted parallel call was not part of the response. # Require a later completed assistant message as the turn boundary. retained_output_seen = False + retained_output_is_final_answer = False fresh_followup_seen = False + fresh_followup_count = 0 + fresh_followup_is_user_message = False continue call_type = _TOOL_CALL_TYPE_BY_OUTPUT_TYPE.get(item_type or "") if call_type is not None: @@ -300,12 +324,28 @@ def responses_input_suffix_retains_prior_output( if pending_suffix_calls or not _is_retained_response_message(item): return False retained_output_seen = True + retained_output_is_final_answer = item.get("phase") == "final_answer" fresh_followup_seen = False + fresh_followup_count = 0 + fresh_followup_is_user_message = False continue if _is_fresh_followup_input(item): if not retained_output_seen or pending_suffix_calls: return False fresh_followup_seen = True + fresh_followup_count += 1 + fresh_followup_is_user_message = item_type in (None, "message") and item.get("role") == "user" + continue + if _fresh_developer_message_is_transparent(item): + if ( + not fresh_followup_seen + or fresh_followup_count != 1 + or not fresh_followup_is_user_message + or not retained_output_is_final_answer + or pending_suffix_calls + ): + return False + fresh_developer_followup_seen = True continue return False return retained_output_seen and fresh_followup_seen and not pending_suffix_calls @@ -321,10 +361,20 @@ def responses_input_suffix_matches_pending_tool_calls( if stored_count <= 0 or len(input_items) <= stored_count or not pending_tool_calls: return False - prefix_state = _direct_tool_call_prefix_state(input_items[:stored_count]) - if prefix_state is None or prefix_state[1] & pending_tool_calls.keys(): + prefix_state = _direct_tool_call_prefix_state( + input_items[:stored_count], + allow_historical_developer_interleave=True, + ) + if prefix_state is None or prefix_state[0] or prefix_state[1] & pending_tool_calls.keys(): return False suffix = input_items[stored_count:] + if ( + len(suffix) == 3 + and isinstance(suffix[1], dict) + and _fresh_developer_message_is_transparent(suffix[1]) + and _fresh_developer_interleave_is_bounded(suffix, index=1) + ): + suffix = [suffix[0], suffix[2]] if not all( isinstance(item, dict) and isinstance(item.get("type"), str) @@ -349,6 +399,8 @@ def responses_input_suffix_matches_pending_tool_calls( def _direct_tool_call_prefix_state( input_items: list[JsonValue], + *, + allow_historical_developer_interleave: bool = False, ) -> tuple[deque[tuple[str, str]], set[str]] | None: pending_calls: deque[tuple[str, str]] = deque() seen_call_ids: set[str] = set() @@ -379,6 +431,12 @@ def _direct_tool_call_prefix_state( return None pending_calls.popleft() continue + if ( + pending_calls + and allow_historical_developer_interleave + and _historical_pending_developer_message_is_transparent(item, item_type=item_type) + ): + continue if pending_calls and ( (item_type in (None, "message") and item.get("role") in _ACCOUNT_NEUTRAL_MESSAGE_ROLES) or item_type in {"input_file", "input_image", "input_text"} @@ -394,6 +452,73 @@ def _direct_tool_call_prefix_state( return pending_calls, seen_call_ids +def _historical_pending_developer_message_is_transparent( + item: Mapping[str, JsonValue], + *, + item_type: str | None, +) -> bool: + return ( + item_type in (None, "message") + and item.get("role") == "developer" + and item.get("id") is None + and item.get("phase") is None + and item.get("status") in (None, "completed") + and _internal_chat_message_metadata_is_account_neutral(item.get(_INTERNAL_CHAT_MESSAGE_METADATA_FIELD)) + and _input_item_has_only_known_fields(item, item_type) + and _message_has_valid_account_neutral_content(item) + ) + + +def _fresh_developer_interleave_is_bounded( + input_items: list[JsonValue], + *, + index: int, +) -> bool: + if len(input_items) != 3 or index != 1: + return False + preceding_item = input_items[0] + following_item = input_items[2] + if not isinstance(preceding_item, dict) or not isinstance(following_item, dict): + return False + call_type = preceding_item.get("type") + output_type = following_item.get("type") + call_id = preceding_item.get("call_id") + return ( + call_type == "custom_tool_call" + and output_type == "custom_tool_call_output" + and _is_nonblank_string(call_id) + and following_item.get("call_id") == call_id + ) + + +def _fresh_developer_message_is_transparent( + item: Mapping[str, JsonValue], +) -> bool: + item_type_value = item.get("type") + item_type = item_type_value if isinstance(item_type_value, str) else None + metadata = item.get(_INTERNAL_CHAT_MESSAGE_METADATA_FIELD) + content = item.get("content") + return ( + ("type" not in item or _is_nonblank_string(item.get("type"))) + and item_type in (None, "message") + and item.get("role") == "developer" + and item.get("id") in (None, "") + and item.get("phase") is None + and item.get("status") in (None, "completed") + and isinstance(metadata, dict) + and _internal_chat_message_metadata_is_account_neutral(metadata) + and _input_item_has_only_known_fields(item, item_type) + and isinstance(content, list) + and len(content) == 1 + and isinstance(content[0], dict) + and content[0].get("type") == "input_text" + and _input_content_part_is_self_contained( + cast(dict[str, JsonValue], content[0]), + allow_output=False, + ) + ) + + def _is_retained_response_message(item: Mapping[str, JsonValue]) -> bool: item_type = item.get("type") if ( diff --git a/app/modules/proxy/selection_errors.py b/app/modules/proxy/selection_errors.py new file mode 100644 index 0000000000..cdcfcd9e39 --- /dev/null +++ b/app/modules/proxy/selection_errors.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import Protocol + +from app.core.errors import OpenAIErrorEnvelope, openai_error +from app.core.resilience.overload import is_local_overload_error_code + +USAGE_LIMIT_REACHED = "usage_limit_reached" + + +class SelectionFailure(Protocol): + error_message: str | None + error_code: str | None + resets_at: int | None + + +def selection_failure_response(selection: SelectionFailure) -> tuple[int, OpenAIErrorEnvelope]: + """Map an account-selection failure to its externally visible HTTP response. + + The ``usage_limit_reached`` mapping is strictly for upstream usage/quota + exhaustion of the whole eligible pool. Local capacity codes (account caps, + admission gates, fair-share throttles) resolve against the canonical + ``LOCAL_OVERLOAD_CODES`` registry so they keep their stable 429 + ``rate_limit_error`` contract and are never reclassified as upstream usage + exhaustion or collapsed into a generic 503. + """ + code = selection.error_code or "no_accounts" + message = selection.error_message or "No active accounts available" + if code == USAGE_LIMIT_REACHED: + return ( + 429, + openai_error( + code, + message, + error_type=USAGE_LIMIT_REACHED, + resets_at=selection.resets_at, + ), + ) + if is_local_overload_error_code(code): + return 429, openai_error(code, message, error_type="rate_limit_error") + return 503, openai_error(code, message) diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index 689913390d..2cc5d98c97 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -122,6 +122,9 @@ from app.modules.proxy._service.api_key_usage import ( _API_KEY_RESERVATION_HEARTBEAT_SECONDS as _API_KEY_RESERVATION_HEARTBEAT_SECONDS, ) +from app.modules.proxy._service.api_key_usage import ( + _STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY as _STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY, +) from app.modules.proxy._service.api_key_usage import _ApiKeyUsageMixin from app.modules.proxy._service.codex_control import _CodexControlMixin from app.modules.proxy._service.compact import _CompactMixin @@ -748,6 +751,7 @@ from app.modules.proxy.ring_membership import ( RingMembershipService, ) +from app.modules.proxy.selection_errors import selection_failure_response from app.modules.proxy.work_admission import WorkAdmissionController logger = logging.getLogger(__name__) @@ -940,6 +944,7 @@ def __init__( self._websocket_previous_response_account_index: dict[tuple[str, str | None, str | None], str] = {} self._websocket_continuity_index: dict[tuple[str, str | None], _WebSocketContinuityState] = {} self._background_cleanup_tasks: set[asyncio.Task[None]] = set() + self._stream_api_key_release_retry_semaphore = asyncio.Semaphore(_STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY) # In-memory pin from upstream-issued file_id -> codex-lb account_id. # Used so ``finalize_file`` for a given ``file_id`` is routed to # the same account that handled ``create_file``. Cross-instance @@ -1029,10 +1034,8 @@ async def thread_goal_request( if account is None: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call_goal(target: Account) -> dict[str, JsonValue]: @@ -1894,6 +1897,7 @@ def log_account_id(account_id: str | None) -> str | None: traffic_class=effective_traffic_class, concurrency_caps=concurrency_caps, redact_sensitive_details=redact_sensitive_details, + allow_usage_exhaustion_error=not required_preferred_account, ) if preferred_selection.account is not None: logger.info( diff --git a/app/modules/proxy/tool_call_dedupe.py b/app/modules/proxy/tool_call_dedupe.py index a5d2cd7625..730a46b378 100644 --- a/app/modules/proxy/tool_call_dedupe.py +++ b/app/modules/proxy/tool_call_dedupe.py @@ -137,6 +137,14 @@ def mark_duplicate_tool_call_downstream_event( None, argument_key, ) + if item_name is not None and same_response_argument_key in seen_tool_call_keys: + logger.warning( + "Suppressed duplicate downstream side-effect tool call response_id=%s item_type=%s name=%s", + response_id, + item_type, + item_name, + ) + return True code_mode_call = item_name in tool_call_safety.CODE_MODE_DOWNSTREAM_SIDE_EFFECT_TOOL_CALL_NAMES identity_scoped_call = code_mode_call or item_namespace is not None cross_response_call_id = call_id if identity_scoped_call else None @@ -162,6 +170,11 @@ def mark_duplicate_tool_call_downstream_event( item_name, ) return True + _clear_downstream_side_effect_burst_keys( + seen_tool_call_keys, + current_key=key, + current_argument_key=same_response_argument_key, + ) seen_tool_call_keys[key] = None if is_side_effect_tool_call: seen_tool_call_keys[same_response_argument_key] = None @@ -180,6 +193,22 @@ def _clear_legacy_downstream_tool_call_keys(seen_tool_call_keys: dict[ToolCallDe seen_tool_call_keys.pop(key, None) +def _clear_downstream_side_effect_burst_keys( + seen_tool_call_keys: dict[ToolCallDedupeKey, None], + *, + current_key: ToolCallDedupeKey, + current_argument_key: ToolCallDedupeKey, +) -> None: + for key in tuple(seen_tool_call_keys): + if key == current_key or key == current_argument_key: + continue + response_id, _, namespace, _, call_id, _ = key + if response_id == "": + continue + if call_id is None: + seen_tool_call_keys.pop(key, None) + + def _mark_duplicate_parallel_tool_call_downstream_event( item: dict[str, JsonValue], argument_value: str, diff --git a/openspec/changes/allow-developer-interleaved-fresh-resend/.openspec.yaml b/openspec/changes/allow-developer-interleaved-fresh-resend/.openspec.yaml new file mode 100644 index 0000000000..5849c2dbf4 --- /dev/null +++ b/openspec/changes/allow-developer-interleaved-fresh-resend/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-08-01 diff --git a/openspec/changes/allow-developer-interleaved-fresh-resend/proposal.md b/openspec/changes/allow-developer-interleaved-fresh-resend/proposal.md new file mode 100644 index 0000000000..fc090c3846 --- /dev/null +++ b/openspec/changes/allow-developer-interleaved-fresh-resend/proposal.md @@ -0,0 +1,56 @@ +## Why + +A verified durable Responses-Lite input prefix can contain a completed direct +tool call with a Codex `developer` message between the call and its matching +output. Because the Lite `additional_tools` bundle keeps that message inline, +the fresh full-resend classifier encounters it while the historical call is +pending and rejects the otherwise valid shape. + +Two additional Responses-Lite resend shapes are now observed after the stored +prefix: a fresh `developer` message after retained final output plus one user +follow-up, and a fresh `developer` message between a custom tool call and its +matching output. Treating every fresh developer message as unsafe makes these +bounded, account-neutral resends fall back to anchor injection and can trigger +an upstream acknowledgement timeout. + +## What Changes + +- Keep the observed unphased, non-response-owned historical `developer` + message transparent only while proving an exact durable pending-tool + manifest from inline Responses-Lite input. +- Allow a fresh developer message after retained output only when the latest + retained assistant message is `final_answer`, exactly one explicit user + message follows it, and the developer message is terminal. +- Allow a fresh developer message in a tool suffix only when the entire suffix + is exactly `custom_tool_call -> developer -> matching custom_tool_call_output` + and the pair exactly equals the durable pending-tool manifest. +- Require fresh developer messages to contain exactly one account-neutral + `input_text` part, exact `turn_id` metadata, known fields, no response-owned + ID or phase, and no status other than `completed`. +- Keep function calls, apply-patch calls, parallel batches, extra leading or + trailing items, malformed types, account-scoped content, and all unproven + developer positions fail-closed. +- Leave non-Lite `input` and `messages` instruction hoisting unchanged; this + change does not add hoist provenance to the durable proof. +- Add helper, bridge-unit, and public `/v1/responses` regressions for both + observed fresh suffixes and their rejection boundaries. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `responses-api-compat`: verified Responses-Lite developer-interleaved history + and two bounded fresh developer suffixes can preserve the existing safe fresh + full-resend path. + +## Impact + +- Code: replay classification; the existing HTTP bridge projection and owner + selection contracts remain unchanged. +- Tests: focused replay-safety, bridge-unit, and HTTP bridge route coverage. +- Owner forwarding, retry policy, logging, storage, public schemas, and + non-Lite instruction hoisting remain unchanged. diff --git a/openspec/changes/allow-developer-interleaved-fresh-resend/specs/responses-api-compat/spec.md b/openspec/changes/allow-developer-interleaved-fresh-resend/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..40eeacd7b0 --- /dev/null +++ b/openspec/changes/allow-developer-interleaved-fresh-resend/specs/responses-api-compat/spec.md @@ -0,0 +1,73 @@ +# responses-api-compat Delta + +## ADDED Requirements + +### Requirement: Responses-Lite replay proof tolerates only verified developer interleaving + +When a fresh durable HTTP bridge classifies a client-unanchored Responses-Lite +full resend whose `additional_tools` bundle preserves developer messages inline, +the replay proof MUST tolerate a developer message only in the historical and +fresh positions defined below. Every other developer position or shape MUST +remain fail-closed. + +A tolerated fresh developer message MUST have `type` omitted or equal to `message`, +MUST have role `developer`, MUST have no non-empty response-owned ID or phase, +MUST have no status or a `completed` status, MUST contain exact account-neutral +metadata with one nonblank `turn_id`, MUST contain exactly one self-contained +`input_text` content part, and MUST contain no unknown or account-scoped fields. +Explicit null or malformed item types MUST fail closed. + +Classification MUST retain response-owned developer-message ID evidence until +these checks have completed, even when other response-owned IDs are projected +out. Non-Lite `input` or `messages` forms whose instruction-role messages are +normalized into top-level `instructions` remain outside this requirement. + +#### Scenario: Verified historical Responses-Lite developer message is transparent + +- **GIVEN** a Responses-Lite input contains an `additional_tools` bundle +- **AND** its fingerprint-verified stored prefix contains a supported direct call +- **AND** a valid developer message appears before that call's matching output +- **AND** the fresh suffix exactly settles the durable pending-tool manifest +- **WHEN** the HTTP bridge opens a replacement session on the durable owner +- **THEN** it sends the original full input without injecting `previous_response_id` +- **AND** it sends the request once + +#### Scenario: Other historical messages remain fail-closed + +- **GIVEN** a supported direct call is pending in the verified stored prefix +- **WHEN** a user, assistant, system, malformed developer, or response-owned message appears before its output +- **THEN** exact manifest proof fails + +#### Scenario: Historical output remains mandatory + +- **GIVEN** a valid developer message follows a supported historical call +- **WHEN** the matching output is missing or has another call ID or type +- **THEN** exact manifest proof fails + +#### Scenario: Bounded fresh custom-tool developer interleave is transparent + +- **GIVEN** the fingerprint-verified stored prefix is followed by a fresh suffix +- **AND** the durable pending-tool manifest contains exactly one `custom_tool_call` +- **WHEN** the entire suffix is exactly that custom call, one valid developer message, and its matching custom-tool output +- **THEN** exact manifest proof passes +- **AND** the original full input is sent once without injecting `previous_response_id` + +#### Scenario: Other fresh tool-loop developer positions remain fail-closed + +- **GIVEN** a durable pending-tool manifest +- **WHEN** a fresh developer message is used with a function or apply-patch call, appears in a parallel batch, is duplicated, lacks exact metadata, contains malformed or account-scoped content, or has leading or trailing suffix items +- **THEN** exact manifest proof fails + +#### Scenario: Bounded retained-output developer follow-up is transparent + +- **GIVEN** the fingerprint-verified stored prefix is followed by a completed assistant `final_answer` +- **AND** exactly one explicit user message follows that retained output +- **WHEN** one valid developer message is the terminal suffix item +- **THEN** retained-output proof passes +- **AND** the original full input is sent once without injecting `previous_response_id` + +#### Scenario: Unproven retained-output developer follow-up remains fail-closed + +- **GIVEN** a retained-output full resend +- **WHEN** the latest assistant output is not `final_answer`, the developer message is not terminal, the fresh input is raw or contains multiple user items, the developer metadata or content is not account-neutral, or the stored prefix contains historical developer interleaving +- **THEN** retained-output proof fails diff --git a/openspec/changes/allow-developer-interleaved-fresh-resend/tasks.md b/openspec/changes/allow-developer-interleaved-fresh-resend/tasks.md new file mode 100644 index 0000000000..482c6ac99e --- /dev/null +++ b/openspec/changes/allow-developer-interleaved-fresh-resend/tasks.md @@ -0,0 +1,19 @@ +## 1. Replay classification + +- [x] 1.1 Allow a valid historical developer message only in exact manifest proof. +- [x] 1.2 Keep the historical call/output match and fresh suffix checks fail-closed. +- [x] 1.3 Allow a terminal fresh developer message only after `final_answer` and exactly one explicit user follow-up. +- [x] 1.4 Allow fresh tool interleaving only for an exact custom call/developer/matching-output suffix. +- [x] 1.5 Reject malformed, account-scoped, parallel, function/apply-patch, leading, trailing, and repeated variants. + +## 2. Regression coverage + +- [x] 2.1 Add focused positive and negative replay-safety cases. +- [x] 2.2 Exercise historical developer interleaving through `/v1/responses`. +- [x] 2.3 Exercise both bounded fresh developer suffixes through bridge-unit and `/v1/responses` coverage. +- [x] 2.4 Verify the new positive regressions fail before the production fix and pass after it. + +## 3. Validation + +- [x] 3.1 Run the full replay-safety, bridge-unit, and HTTP bridge integration suites. +- [x] 3.2 Run changed-file Ruff, type, diff, and strict OpenSpec checks. diff --git a/openspec/changes/preserve-http-bridge-terminal-delivery/proposal.md b/openspec/changes/preserve-http-bridge-terminal-delivery/proposal.md index 3855f32c14..e795f391d6 100644 --- a/openspec/changes/preserve-http-bridge-terminal-delivery/proposal.md +++ b/openspec/changes/preserve-http-bridge-terminal-delivery/proposal.md @@ -12,14 +12,8 @@ processing already claimed the request. - Capture the downstream queue when completed-event processing removes the request from pending ownership. - Use that captured queue for completed delivery after asynchronous bookkeeping. -- Serialize completed-queue claiming with the terminal idle-timeout decision - under the bridge pending lock. A timeout that wins MUST revoke the mutable - queue before releasing the lock; a completed claim that wins MUST remain - authoritative through queued terminal delivery and suppress that timeout. - Keep emitting liveness frames, without manufacturing an idle timeout, while that completed-delivery operation is actively doing bookkeeping. -- Log the first completed-delivery timeout suppression with bounded request, - response, and elapsed-time context. - Add stream-level regressions for slow and failed completed bookkeeping. ## Capabilities diff --git a/openspec/changes/preserve-http-bridge-terminal-delivery/specs/responses-api-compat/spec.md b/openspec/changes/preserve-http-bridge-terminal-delivery/specs/responses-api-compat/spec.md index bdb0af7bcb..04841e06b8 100644 --- a/openspec/changes/preserve-http-bridge-terminal-delivery/specs/responses-api-compat/spec.md +++ b/openspec/changes/preserve-http-bridge-terminal-delivery/specs/responses-api-compat/spec.md @@ -16,21 +16,6 @@ requirement. While the claimed completed-delivery operation remains active, ordinary stream idle accounting MUST NOT replace the upstream completion with a synthetic idle failure, and the stream MUST continue emitting its existing liveness frames. -The completed-queue claim and the terminal idle-timeout decision MUST be -serialized under the bridge pending lock. If completed processing wins that -serialization and claims a live queue, the timeout MUST be suppressed. If the -terminal event and end-of-stream marker are already queued when a concurrent -timeout finishes awaited recovery work, the completed claim MUST remain -authoritative until the stream consumes that queued delivery. If the -terminal idle timeout wins while no completed delivery is active, it MUST -revoke the request's mutable event queue before releasing the pending lock so a -later completed event cannot claim an orphaned queue. - -The first idle-timeout suppression for one completed-delivery operation MUST -emit one bounded diagnostic containing the request ID, downstream response ID, -and elapsed seconds. Further liveness intervals for that same operation MUST -NOT repeat the diagnostic. - When that operation returns, raises, or is cancelled before delivery, idle timeout behavior MUST resume. @@ -51,24 +36,6 @@ client-disconnect and drain behavior MUST remain unchanged. - **WHEN** later completed bookkeeping exceeds the configured stream idle window - **THEN** the stream continues emitting liveness frames - **AND** it does not emit a synthetic idle failure while that operation remains active -- **AND** it logs the suppression once with request, response, and elapsed-time context - -#### Scenario: Terminal idle timeout wins before completed processing - -- **GIVEN** an HTTP bridge stream has exhausted its configured idle window -- **AND** no completed-delivery operation has claimed its queue -- **WHEN** the stream acquires the bridge pending lock before a concurrent completed event -- **THEN** it revokes the mutable event queue while still holding that lock -- **AND** it emits the existing synthetic idle failure -- **AND** later completed processing does not deliver to the revoked queue - -#### Scenario: Completed delivery finishes during timeout recovery - -- **GIVEN** an HTTP bridge timeout path is awaiting pre-response recovery work -- **AND** completed processing claims the live queue and enqueues its terminal event and end-of-stream marker -- **WHEN** completed processing returns before the timeout path rechecks ownership -- **THEN** the completed claim remains authoritative -- **AND** the stream consumes the queued completion without emitting a synthetic idle failure #### Scenario: Completed bookkeeping aborts diff --git a/openspec/changes/preserve-http-bridge-terminal-delivery/tasks.md b/openspec/changes/preserve-http-bridge-terminal-delivery/tasks.md index f1e1c13d80..4f5080bd19 100644 --- a/openspec/changes/preserve-http-bridge-terminal-delivery/tasks.md +++ b/openspec/changes/preserve-http-bridge-terminal-delivery/tasks.md @@ -4,20 +4,14 @@ - [x] 1.2 Deliver the completed event and end-of-stream marker through that captured queue. - [x] 1.3 Preserve detach-first cancellation and retry behavior. - [x] 1.4 Suppress synthetic idle failure only while completed delivery is actively producing. -- [x] 1.5 Serialize the completed-claim/terminal-timeout decision under the pending lock and atomically revoke the timeout-first queue. -- [x] 1.6 Emit one bounded diagnostic when completed delivery suppresses an idle timeout. -- [x] 1.7 Keep a completed claim authoritative after terminal delivery is queued but before the stream consumes it. ## 2. Regression coverage - [x] 2.1 Add a stream-level regression for slow bookkeeping after completed pending removal. - [x] 2.2 Verify the regression fails on the pre-fix implementation and passes after the fix. - [x] 2.3 Verify failed completed bookkeeping releases timeout suppression. -- [x] 2.4 Add deterministic timeout-first and completed-first race regressions, including one-time suppression logging. -- [x] 2.5 Add a regression where completed delivery finishes while the timeout path awaits pre-response recovery. ## 3. Validation - [x] 3.1 Run focused HTTP bridge tests and changed-file Ruff checks. - [x] 3.2 Run strict OpenSpec validation. -- [x] 3.3 Re-run focused/full HTTP bridge tests, static checks, architecture checks, strict OpenSpec validation, and local Codex review. diff --git a/openspec/changes/recover-codex-desktop-idle-bridge/design.md b/openspec/changes/recover-codex-desktop-idle-bridge/design.md index 45f132d300..b4202a5366 100644 --- a/openspec/changes/recover-codex-desktop-idle-bridge/design.md +++ b/openspec/changes/recover-codex-desktop-idle-bridge/design.md @@ -53,7 +53,7 @@ Leading non-response telemetry such as `codex.rate_limits` does not change those When the deadline expires, reuse the reader-owned terminal failure and whole-session retirement path. Emit a stable `missing_response_created_timeout` detail, increment the existing stuck-retirement metric, settle every pending request exactly once, and close the bridge session. -Do not transparently replay the timed-out request, submit it on another account, or mark the selected account unhealthy. Upstream acceptance is unknown, so duplicate submission and account movement are less safe than an explicit terminal failure. A later client request creates a fresh session through existing behavior. +Do not transparently replay the timed-out request or submit that same request on another account. Upstream acceptance is unknown, so duplicate submission is less safe than an explicit terminal failure. Record a transient health failure for the selected account so a later client request creates a fresh session through existing selection and can avoid repeating an account that accepted `response.create` but emitted no `response.created`. Example: a request sends at monotonic time 1,000 with the default 300-second stuck threshold. With no matched response lifecycle event, it becomes eligible at 1,240 and receives an explicit terminal failure; it does not wait for a second request or the 300-second Desktop idle timeout. @@ -66,7 +66,7 @@ Explicit `x-stainless-*` headers or an OpenAI User-Agent retain comment liveness ## Risks / Trade-offs - **A send fails after the timestamp is set.** Existing send-error cleanup retires or settles the request before the watchdog can act; tests cover that the timestamp alone is not sufficient eligibility. -- **A quiet upstream accepted the request but emitted no event.** The proxy returns an explicit failure rather than risking a duplicate replay. The selected account remains healthy because silence is not proof of account failure. +- **A quiet upstream accepted the request but emitted no event.** The proxy returns an explicit failure rather than risking a duplicate replay. The selected account receives a transient health failure because repeated missing-created timeouts on the same account are a live availability fault. - **A matched lifecycle event arrives just before timeout.** Eligibility is rechecked under the existing request/session synchronization before retirement, and any matched `response.*` event suppresses this watchdog. - **Whole-session retirement interrupts a healthy sibling.** This narrow design chooses fail-closed session cleanup rather than attempting unsafe sibling isolation on current `main`. Existing terminal settlement must cover every pending sibling exactly once. - **A client spoofs native identity.** The only benefit is an ignored vendor liveness event on the authenticated Codex backend route; explicit SDK markers still take precedence. diff --git a/openspec/changes/recover-codex-desktop-idle-bridge/proposal.md b/openspec/changes/recover-codex-desktop-idle-bridge/proposal.md index d9849d4af5..d17fe39d34 100644 --- a/openspec/changes/recover-codex-desktop-idle-bridge/proposal.md +++ b/openspec/changes/recover-codex-desktop-idle-bridge/proposal.md @@ -6,9 +6,9 @@ A production Codex Desktop request on the HTTP-to-WebSocket bridge remained pend - Record the monotonic time of the current upstream `response.create` send. - Proactively expire an eventless request that remains pre-`response.created` for the smaller of the existing stuck-gate threshold and 240 seconds, even when no second gate waiter exists and periodic keepalives are disabled. -- Fail the affected bridge session closed through existing terminal settlement and retirement paths, without transparent replay, account movement, or account-health penalties. +- Fail the affected bridge session closed through existing terminal settlement and retirement paths, without transparent replay or moving the timed-out request to another account, while recording a transient account-health failure so later requests can avoid the eventless account. - Give verified native Codex identity parser-visible `codex.keepalive` frames even when payload-shape heuristics still require OpenAI-compatible event normalization; explicit SDK markers and public `/v1/responses` retain comment liveness. -- Add regressions for the no-waiter deadline, protected created/eventful requests, account-neutral retirement, and contrasting Desktop/SDK/public heartbeat contracts. +- Add regressions for the no-waiter deadline, protected created/eventful requests, health-accounted retirement without replay, and contrasting Desktop/SDK/public heartbeat contracts. ## Capabilities 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 293763926d..6db9f29f1a 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 @@ -6,7 +6,7 @@ The proxy MUST retain the existing waiter-triggered retirement behavior for stal 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 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. +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, write the selected account's transient health failure, and retire the whole bridge session. It MUST NOT transparently replay the timed-out request or move that timed-out request to another account. #### Scenario: Lone eventless gate owner is retired before the client timeout @@ -39,10 +39,10 @@ When the owner-side deadline expires, the proxy MUST recheck eligibility and emi - **THEN** this watchdog does not retire the session - **AND** existing stream, request-budget, and waiter-triggered timeout behavior remains authoritative -#### Scenario: Timeout is fail-closed and account-neutral +#### Scenario: Timeout is fail-closed and health-accounted - **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 unless it satisfies the bounded fresh-hard recovery requirement -- **AND** the selected account is not marked unhealthy solely because `response.created` was missing +- **AND** the proxy does not replay the timed-out request or submit it on another account +- **AND** the selected account records a transient health failure so later requests can avoid repeating the eventless account diff --git a/openspec/changes/recover-codex-desktop-idle-bridge/tasks.md b/openspec/changes/recover-codex-desktop-idle-bridge/tasks.md index 1c356ad8d4..32cc5e81e6 100644 --- a/openspec/changes/recover-codex-desktop-idle-bridge/tasks.md +++ b/openspec/changes/recover-codex-desktop-idle-bridge/tasks.md @@ -3,8 +3,8 @@ - [x] 1.1 Record the current monotonic `response.create` send timestamp in HTTP bridge request state and replace it on every real send. - [x] 1.2 Add a pure client-safe deadline helper that uses the smaller of the existing stuck-gate threshold and 240 seconds. - [x] 1.3 Enforce the deadline from the upstream reader without requiring a second gate waiter or SSE keepalives; recheck narrow eventless eligibility before acting. -- [x] 1.4 Fail and retire the whole bridge session through existing settlement, logging, and Prometheus paths without replay, account movement, or account-health writes. -- [x] 1.5 Add focused regressions for no-waiter expiry, send-time anchoring, leading telemetry, created/eventful/downstream protection, terminal settlement, and account neutrality. +- [x] 1.4 Fail and retire the whole bridge session through existing settlement, logging, Prometheus, and transient account-health paths without replaying or moving the timed-out request. +- [x] 1.5 Add focused regressions for no-waiter expiry, send-time anchoring, leading telemetry, created/eventful/downstream protection, terminal settlement, and health-accounted retirement. ## 2. Native Codex SSE liveness diff --git a/openspec/changes/report-pool-usage-exhaustion/.openspec.yaml b/openspec/changes/report-pool-usage-exhaustion/.openspec.yaml new file mode 100644 index 0000000000..0bd76e6186 --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-18 diff --git a/openspec/changes/report-pool-usage-exhaustion/proposal.md b/openspec/changes/report-pool-usage-exhaustion/proposal.md new file mode 100644 index 0000000000..41367d46a4 --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/proposal.md @@ -0,0 +1,35 @@ +## Why + +When every account eligible for a Responses request is exhausted by known pool +usage windows, codex-lb can currently collapse the selection failure into a +generic no-account/server-unavailable response. That hides the user-actionable +upstream condition from Codex/OpenAI-compatible clients and makes agents treat a +quota window as infrastructure failure. + +## What Changes + +- Preserve the stable `usage_limit_reached` code from account selection when the + whole eligible pool is exhausted by usage windows. +- Return HTTP `429` with an OpenAI-style error envelope whose + `error.code` and `error.type` are both `usage_limit_reached`. +- Preserve the selected reset hint as `error.resets_at` when account selection + has one, and use the same contract across HTTP, streaming, bridge, and + WebSocket selection-failure paths. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `responses-api-compat`: define the externally visible Responses error + contract for pool-wide usage exhaustion. + +## Impact + +- Affected code: account selection failure mapping and Responses proxy surfaces. +- Affected APIs: failure status/body for pool-wide usage exhaustion changes from + generic unavailable/no-account semantics to HTTP 429 `usage_limit_reached`. +- Configuration and schema: no changes. diff --git a/openspec/changes/report-pool-usage-exhaustion/specs/responses-api-compat/spec.md b/openspec/changes/report-pool-usage-exhaustion/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..65482726e5 --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/specs/responses-api-compat/spec.md @@ -0,0 +1,63 @@ +## ADDED Requirements + +### Requirement: Pool usage exhaustion is reported as a usage-limit error + +The proxy MUST report pool-wide Responses usage exhaustion as a usage-limit +error. When every account eligible for a Responses request is exhausted by known +usage windows, the proxy MUST reject the request with HTTP `429` and an +OpenAI-style error envelope whose `error.code` and `error.type` are both +`usage_limit_reached`. If account selection has an authoritative upstream reset +timestamp for the exhausted pool, the response envelope MUST include that +timestamp as `error.resets_at`; the proxy MUST NOT expose the capped +human-facing retry hint or a synthesized fallback as `error.resets_at`. The +proxy MUST NOT collapse this condition into generic `no_accounts`, +`server_error`, or HTTP `503` semantics. Exhaustion classification MUST be +based on structured account state after the same eligibility filtering as +ordinary selection, and MUST NOT reclassify local capacity or overload codes +(account caps, admission gates, fair-share throttles) as usage exhaustion. + +#### Scenario: Public Responses request exhausts the eligible usage pool + +- **WHEN** account selection for a public `/v1/responses` or + `/backend-api/codex/responses` request finds only usage-exhausted eligible + accounts +- **THEN** the response status is HTTP `429` +- **AND** the response body has `error.code = "usage_limit_reached"` +- **AND** the response body has `error.type = "usage_limit_reached"` +- **AND** any selected pool reset timestamp is surfaced as `error.resets_at` + +#### Scenario: Streaming selection failure preserves usage-limit semantics + +- **WHEN** a streaming Responses request cannot select an account because every + eligible account is usage-exhausted before downstream-visible output +- **THEN** the terminal error event uses `usage_limit_reached` +- **AND** clients do not receive a generic no-account/server-unavailable error + +#### Scenario: Usage-limit selection failures are terminal, not waitable + +- **WHEN** account selection fails with `usage_limit_reached` on a streaming, + HTTP-bridge, or WebSocket Responses path +- **THEN** the proxy reports the structured usage-limit failure immediately +- **AND** it does not enter an account-capacity recovery wait for the + remaining request budget before reporting it + +#### Scenario: Local capacity codes keep their rate-limit contract + +- **WHEN** account selection fails with a local capacity or overload code such + as `account_stream_cap` or `account_response_create_cap` +- **THEN** the response keeps HTTP `429` with `error.type = "rate_limit_error"` + and the stable local error code +- **AND** the response is not reported as `usage_limit_reached` + +#### Scenario: Unusable non-exhausted pools keep existing semantics + +- **WHEN** every account is paused, deactivated, or requires re-authentication + and no eligible account is exhausted by a known usage window +- **THEN** the pre-existing `no_accounts` failure semantics are preserved + +#### Scenario: Owner-scoped exhaustion preserves continuity semantics + +- **WHEN** a request is pinned to a previous-response or file owner account and + only that owner is usage-exhausted while the wider eligible pool is usable +- **THEN** the proxy keeps the existing continuity-owner failure semantics +- **AND** it does not report pool-wide `usage_limit_reached` diff --git a/openspec/changes/report-pool-usage-exhaustion/tasks.md b/openspec/changes/report-pool-usage-exhaustion/tasks.md new file mode 100644 index 0000000000..9dd87603ff --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/tasks.md @@ -0,0 +1,26 @@ +## 1. Error contract + +- [x] Preserve `usage_limit_reached` from pool-wide account selection failures. +- [x] Map pool-wide usage exhaustion to HTTP 429 with OpenAI-style + `error.code = "usage_limit_reached"` and + `error.type = "usage_limit_reached"`. +- [x] Preserve `error.resets_at` when account selection provides a reset hint. + +## 2. Proxy surfaces + +- [x] Apply the same selection-failure response helper across HTTP, streaming, + bridge, compact, file, transcription, WebSocket, and Codex-control paths. +- [x] Keep local capacity cap errors as 429 `rate_limit_error` responses rather + than weakening their existing contract. + +## 3. Regression coverage + +- [x] Add unit coverage for pool usage exhaustion selection and response mapping. +- [x] Add externally routed HTTP/streaming regressions for the 429 envelope. + +## 4. Validation + +- [x] Run focused pytest for selection, load balancer, and Responses proxy + regressions. +- [x] Run lint/type checks for touched Python files. +- [x] Validate the OpenSpec change strictly. diff --git a/openspec/changes/retry-detached-api-key-release/.openspec.yaml b/openspec/changes/retry-detached-api-key-release/.openspec.yaml new file mode 100644 index 0000000000..ab39675458 --- /dev/null +++ b/openspec/changes/retry-detached-api-key-release/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-30 diff --git a/openspec/changes/retry-detached-api-key-release/design.md b/openspec/changes/retry-detached-api-key-release/design.md new file mode 100644 index 0000000000..9736d61e2e --- /dev/null +++ b/openspec/changes/retry-detached-api-key-release/design.md @@ -0,0 +1,66 @@ +## Context + +Stream reservation settlement is detached from the response path. A failed +settlement schedules one release task in the existing tracked background-task +set, but that release currently catches its own exception and returns normally. +The done callback consequently removes the task, so the persistence drain can +report success while the reservation remains active. + +Reservation release is already transactional and idempotent: once another +settler has changed the reservation from `reserved`, a later release is a +no-op. The existing stale sweep remains a last-resort repair, but its six-hour +age threshold is too slow for a known live cleanup chain. + +## Goals / Non-Goals + +**Goals:** + +- Keep transiently failing fallback release work visible to the existing task + drain. +- Retry until the idempotent release succeeds, with bounded retry pressure. +- Preserve detached response latency, settlement-before-health ordering, and + exactly-once accounting. + +**Non-Goals:** + +- Changing reservation amounts, quota admission, or stale-sweep timing. +- Adding a durable job queue, setting, migration, or new public API. +- Refactoring request-log persistence or unrelated cleanup ownership. + +## Decisions + +1. **Retry inside the already tracked release task.** The release coroutine + stays pending between attempts, so the current task registry and recursive + drain remain the single source of cleanup ownership. Creating a second + registry or a durable retry row would duplicate state for a narrow failure. + +2. **Use capped exponential delay plus a shared retry gate for every persistence + exception.** The outer retry covers transient PostgreSQL/session failures + that the API-key service's SQLite-lock-specific retry does not classify. A + fixed delay cap prevents each task from retrying rapidly, while a per-service + concurrency gate prevents many failed streams from opening repository + sessions simultaneously. Waiting tasks stay tracked without holding a + database connection. + +3. **Rely on reservation transition idempotency.** A retry cannot double + decrement quota: release only claims a reservation still in `reserved` + state, and a concurrent finalizer or release makes subsequent attempts + no-ops. + +4. **Let the existing drain deadline bound shutdown waiting.** A recovered + release completes normally. A release still retrying at the deadline remains + pending, so `drain_persistence_tasks` returns `False` instead of claiming + durability. No separate retry-count terminal state is introduced. + +## Risks / Trade-offs + +- **A permanent persistence error leaves a task alive during normal runtime.** + → Retries use capped backoff; the task accurately represents unfinished + cleanup, and stale recovery remains the final repair path. +- **Many simultaneous failures could retry together after an outage.** + → A shared four-attempt gate bounds aggregate repository pressure in each + service instance; exponential delay also bounds each task's retry frequency, + and the change adds no inline request-path work. +- **Cancellation can stop a retry after shutdown has already timed out.** + → The drain first reports incomplete, so process termination cannot be + mistaken for successful settlement. diff --git a/openspec/changes/retry-detached-api-key-release/proposal.md b/openspec/changes/retry-detached-api-key-release/proposal.md new file mode 100644 index 0000000000..7a32890cea --- /dev/null +++ b/openspec/changes/retry-detached-api-key-release/proposal.md @@ -0,0 +1,34 @@ +## Why + +A detached stream settlement can fail, enqueue its reservation-release fallback, +and then lose the reservation when that fallback also hits a transient +persistence failure. The task drain reports success even though the reservation +still consumes quota until stale recovery runs hours later. + +## What Changes + +- Keep a failed detached reservation release tracked and retry it after + transient persistence failures, with a shared concurrency bound on repository + attempts. +- Make the persistence drain report completion only after the tracked + settlement/release chain has actually terminated. +- Add deterministic regression coverage for a finalize failure followed by one + failed release attempt, while preserving successful settlement, cancellation, + and SQLite-lock behavior. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `api-keys`: Clarify that a detached settlement fallback which itself fails + transiently remains tracked and retries before persistence drain can succeed. + +## Impact + +The change is limited to detached API-key reservation cleanup in the proxy +service, its focused persistence tests, and the existing API-key settlement +contract. It adds no API, setting, dependency, migration, or dashboard change. diff --git a/openspec/changes/retry-detached-api-key-release/specs/api-keys/spec.md b/openspec/changes/retry-detached-api-key-release/specs/api-keys/spec.md new file mode 100644 index 0000000000..588671d4f0 --- /dev/null +++ b/openspec/changes/retry-detached-api-key-release/specs/api-keys/spec.md @@ -0,0 +1,45 @@ +## MODIFIED Requirements + +### Requirement: Stream reservation settlement is detached from the response path + +Settling a stream API-key reservation MUST NOT block the response/stream close, with one deliberate exception: when a keyed websocket stream terminates with an account-health error, the finalizer MUST wait for the settlement to commit before the load-balancer health write (the settlement-ordering invariant), so that error path intentionally blocks on settlement. In all other cases the settlement MUST run as a tracked background task; when it fails or is cancelled, the reservation MUST still be released by the tracking fallback, and the request's finalization path MUST NOT double-release a transferred settlement. If the tracking fallback itself encounters a persistence failure, it MUST remain tracked and retry the idempotent release; no more than four retry-enabled detached fallback repository attempts may run concurrently per proxy service instance, waiting fallbacks MUST NOT open repository sessions until admitted, and persistence drain MUST NOT report completion while any retry remains unfinished. Reservations MUST continue to count toward key limits until finalized or released, so deferred settlement can never admit usage a synchronous settlement would have rejected. + +#### Scenario: Response close precedes settlement completion + +- **GIVEN** a keyed stream whose settlement transaction is still running +- **WHEN** the stream closes +- **THEN** the close does not wait for the settlement +- **AND** the settlement finalizes the reservation exactly once in the background + +#### Scenario: Failed detached settlement still releases the reservation + +- **GIVEN** a detached settlement whose finalize raises +- **WHEN** the settlement task completes +- **THEN** the tracking fallback releases the reservation + +#### Scenario: Failed fallback release remains tracked + +- **GIVEN** a detached settlement whose finalize raises +- **AND** the first tracking-fallback release attempt also raises +- **WHEN** persistence recovers before the drain deadline +- **THEN** the tracked fallback retries and releases the reservation exactly once +- **AND** persistence drain does not report completion before that release + +#### Scenario: Concurrent fallback release retries are bounded to four + +- **GIVEN** five failed detached settlements in one proxy service instance +- **WHEN** their tracking fallbacks attempt repository persistence concurrently +- **THEN** no more than four release attempts open repository sessions +- **AND** the waiting fallbacks remain tracked until they can retry + +#### Scenario: Websocket health-error settlement precedes the health write + +- **GIVEN** a keyed websocket stream that terminates with an account-health error +- **WHEN** the finalizer settles the reservation +- **THEN** it waits for the settlement to commit before recording the account-health error + +#### Scenario: Shutdown drains pending settlements + +- **WHEN** the service shuts down gracefully with settlements in flight +- **THEN** shutdown waits for them up to the configured drain timeout +- **AND** reports an incomplete drain if a tracked settlement or release remains unfinished at that timeout diff --git a/openspec/changes/retry-detached-api-key-release/tasks.md b/openspec/changes/retry-detached-api-key-release/tasks.md new file mode 100644 index 0000000000..4c04aacc2f --- /dev/null +++ b/openspec/changes/retry-detached-api-key-release/tasks.md @@ -0,0 +1,17 @@ +## 1. Regression + +- [x] 1.1 Add a real-repository regression that injects one finalize failure and one fallback-release failure. +- [x] 1.2 Confirm the regression fails deterministically twice on baseline `3fe0d6f286019a0505783d803db9a1d8cdf6b307`. + +## 2. Implementation + +- [x] 2.1 Keep the fallback release tracked while retrying persistence failures with capped backoff. +- [x] 2.2 Preserve idempotent settlement, cancellation ownership, and truthful persistence-drain behavior. +- [x] 2.3 Bound concurrent fallback repository attempts to four with one shared per-service gate. + +## 3. Verification + +- [x] 3.1 Run the focused detached-settlement and API-key reservation tests. +- [x] 3.2 Run changed-file Ruff, format, type, proxy-architecture, and strict OpenSpec checks. +- [x] 3.3 Inspect the final diff and worktree status for scope and unrelated changes. +- [x] 3.4 Add deterministic fan-out coverage for the shared retry concurrency bound. diff --git a/tests/integration/test_detached_persistence.py b/tests/integration/test_detached_persistence.py index 8f37d03890..cb18ef800b 100644 --- a/tests/integration/test_detached_persistence.py +++ b/tests/integration/test_detached_persistence.py @@ -3,9 +3,17 @@ import pytest from httpx import ASGITransport, AsyncClient from sqlalchemy import select +from sqlalchemy.exc import OperationalError from app.db.models import RequestLog from app.db.session import SessionLocal +from app.modules.api_keys.repository import ApiKeysRepository, UsageReservationData +from app.modules.api_keys.service import ( + ApiKeyCreateData, + ApiKeyRequestUsageBudget, + ApiKeysService, + LimitRuleInput, +) from app.modules.proxy import service as proxy_service_module pytestmark = pytest.mark.integration @@ -105,6 +113,105 @@ async def never_finishes() -> None: service._request_log_tasks.discard(task) +@pytest.mark.asyncio +async def test_failed_detached_settlement_retries_failed_release_until_persisted(raw_client, monkeypatch): + import asyncio + + _, app = raw_client + + async with SessionLocal() as session: + api_keys = ApiKeysService(ApiKeysRepository(session)) + created = await api_keys.create_key( + ApiKeyCreateData( + name="detached-release-retry", + allowed_models=None, + expires_at=None, + limits=[ + LimitRuleInput( + limit_type="total_tokens", + limit_window="weekly", + max_value=100, + ) + ], + ) + ) + api_key = await api_keys.get_key_by_id(created.id) + reservation = await api_keys.enforce_limits_for_request( + created.id, + request_model="gpt-5.5", + request_usage_budget=ApiKeyRequestUsageBudget( + input_tokens=4, + output_tokens=6, + ), + ) + + original_get_reservation = ApiKeysRepository.get_usage_reservation + reservation_read_attempts = 0 + retry_started = asyncio.Event() + allow_retry = asyncio.Event() + + async def fail_first_two_reservation_reads( + self: ApiKeysRepository, + reservation_id: str, + ) -> UsageReservationData | None: + nonlocal reservation_read_attempts + if reservation_id == reservation.reservation_id: + reservation_read_attempts += 1 + if reservation_read_attempts <= 2: + raise OperationalError( + "read usage reservation", + {}, + Exception("transient persistence connection failure"), + ) + if reservation_read_attempts == 3: + retry_started.set() + await allow_retry.wait() + return await original_get_reservation(self, reservation_id) + + monkeypatch.setattr(ApiKeysRepository, "get_usage_reservation", fail_first_two_reservation_reads) + + settlement = proxy_service_module._StreamSettlement( + status="success", + model="gpt-5.5", + input_tokens=4, + output_tokens=6, + ) + from app.dependencies import get_proxy_service_for_app + + service = get_proxy_service_for_app(app) + assert await service._settle_stream_api_key_usage( + api_key, + reservation, + settlement, + request_id="req_detached_release_retry", + ) + drain_task = asyncio.create_task(service.drain_persistence_tasks(timeout_seconds=2)) + retry_wait_task = asyncio.create_task(retry_started.wait()) + done, _ = await asyncio.wait( + {retry_wait_task, drain_task}, + timeout=1, + return_when=asyncio.FIRST_COMPLETED, + ) + retry_was_tracked = retry_wait_task in done and not drain_task.done() + allow_retry.set() + if not retry_wait_task.done(): + retry_wait_task.cancel() + await asyncio.gather(retry_wait_task, return_exceptions=True) + assert await drain_task + + async with SessionLocal() as session: + repo = ApiKeysRepository(session) + stored = await original_get_reservation(repo, reservation.reservation_id) + limits = await repo.get_limits_by_key(created.id) + + assert stored is not None + assert stored.status == "released" + assert len(limits) == 1 + assert limits[0].current_value == 0 + assert reservation_read_attempts == 3 + assert retry_was_tracked is True + + @pytest.mark.asyncio async def test_drain_ignores_stuck_non_persistence_cleanup_tasks(): """A stuck bridge-close cleanup in _background_cleanup_tasks must not diff --git a/tests/integration/test_http_responses_bridge.py b/tests/integration/test_http_responses_bridge.py index b69f5a6f3b..f0c338c95f 100644 --- a/tests/integration/test_http_responses_bridge.py +++ b/tests/integration/test_http_responses_bridge.py @@ -6571,6 +6571,39 @@ async def test_backend_responses_http_bridge_startup_error_omits_turn_state_head assert "x-codex-turn-state" not in response.headers +@pytest.mark.asyncio +async def test_backend_responses_http_bridge_pool_usage_exhaustion_returns_429(async_client, monkeypatch): + _install_bridge_settings(monkeypatch, enabled=True) + + async def fake_select_account_with_budget(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="Usage limit reached", + error_code="usage_limit_reached", + ) + + monkeypatch.setattr( + proxy_module.ProxyService, + "_select_account_with_budget", + fake_select_account_with_budget, + ) + + response = await async_client.post( + "/backend-api/codex/responses", + json={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "hello", + "stream": True, + }, + ) + + assert response.status_code == 429 + assert response.json()["error"]["type"] == "usage_limit_reached" + assert response.json()["error"]["code"] == "usage_limit_reached" + assert "x-codex-turn-state" not in response.headers + + @pytest.mark.asyncio async def test_v1_responses_http_bridge_startup_error_omits_turn_state_header(async_client, monkeypatch): _install_bridge_settings(monkeypatch, enabled=True) @@ -7116,8 +7149,32 @@ async def fake_connect_responses_websocket( assert connect_count == 2 +@pytest.mark.parametrize( + ("developer_message_extra", "fresh_developer_message", "preserves_full_resend"), + [ + pytest.param({}, None, True, id="unowned-developer-message"), + pytest.param({"id": "msg_response_owned"}, None, False, id="response-owned-developer-message"), + pytest.param( + {}, + { + "type": "message", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_fresh"}, + "content": [{"type": "input_text", "text": "fresh control"}], + }, + True, + id="fresh-developer-interleave", + ), + ], +) @pytest.mark.asyncio -async def test_v1_responses_http_bridge_preserves_full_resend_before_fresh_bridge_send(async_client, monkeypatch): +async def test_v1_responses_http_bridge_classifies_responses_lite_developer_interleaved_full_resend( + async_client, + monkeypatch, + developer_message_extra, + fresh_developer_message, + preserves_full_resend, +): _install_bridge_settings(monkeypatch, enabled=True) account_id = await _import_account( async_client, @@ -7156,10 +7213,31 @@ async def fake_connect_responses_websocket( session_headers = {"x-codex-session-id": "fresh-reattach-full-resend"} historical_input = [ + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "custom", "name": "shell"}], + }, { "role": "user", "content": [{"type": "input_text", "text": "first question"}], - } + }, + { + "type": "custom_tool_call", + "call_id": "call_historical_shell", + "name": "shell", + "input": "printf historical", + }, + { + "role": "developer", + "content": [{"type": "input_text", "text": "historical control"}], + **developer_message_extra, + }, + { + "type": "custom_tool_call_output", + "call_id": "call_historical_shell", + "output": "historical", + }, ] first = await asyncio.wait_for( async_client.post( @@ -7183,6 +7261,7 @@ async def fake_connect_responses_websocket( "name": "shell", "input": "pwd", }, + *([fresh_developer_message] if fresh_developer_message is not None else []), { "type": "custom_tool_call_output", "call_id": "call_custom_shell", @@ -7210,8 +7289,11 @@ async def fake_connect_responses_websocket( assert len(first_upstream.sent_text) == 1 assert len(replay_upstream.sent_text) == 1 replay_payload = json.loads(replay_upstream.sent_text[0]) - assert "previous_response_id" not in replay_payload - assert replay_payload["input"] == full_resend + if preserves_full_resend: + assert "previous_response_id" not in replay_payload + assert replay_payload["input"] == full_resend + else: + assert replay_payload["previous_response_id"] == "resp_bridge_custom_1" @pytest.mark.asyncio @@ -7407,9 +7489,16 @@ async def fake_connect_responses_websocket( assert stale_session.closed is True +@pytest.mark.parametrize( + "fresh_developer_followup", + [ + pytest.param(False, id="ordinary-user-followup"), + pytest.param(True, id="fresh-developer-followup"), + ], +) @pytest.mark.asyncio async def test_v1_responses_http_bridge_replays_full_resend_once_then_stays_on_new_owner( - async_client, app_instance, monkeypatch + async_client, app_instance, monkeypatch, fresh_developer_followup ): _install_bridge_settings(monkeypatch, enabled=True) owner_account_id = await _import_account( @@ -7474,10 +7563,21 @@ async def fake_connect_responses_websocket( monkeypatch.setattr(proxy_module, "connect_responses_websocket", fake_connect_responses_websocket) historical_input = [ + *( + [ + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "custom", "name": "shell"}], + } + ] + if fresh_developer_followup + else [] + ), { "role": "user", "content": [{"type": "input_text", "text": "first question"}], - } + }, ] first = await asyncio.wait_for( async_client.post( @@ -7505,19 +7605,46 @@ async def fake_connect_responses_websocket( assert durable_lookup.latest_input_item_count == len(historical_input) assert durable_lookup.latest_input_full_fingerprint is not None - retained_prior_output = { - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "first answer"}], - } + if fresh_developer_followup: + retained_prior_output = { + "type": "message", + "role": "assistant", + "phase": "final_answer", + "status": "completed", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_previous"}, + "content": [{"type": "output_text", "text": "first answer"}], + } + fresh_followup_items = [ + { + "type": "message", + "role": "user", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + "content": [{"type": "input_text", "text": "second question"}], + }, + { + "type": "message", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + "content": [{"type": "input_text", "text": "fresh control"}], + }, + ] + else: + retained_prior_output = { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "first answer"}], + } + fresh_followup_items = [ + { + "role": "user", + "content": [{"type": "input_text", "text": "second question"}], + } + ] full_resend = [ *historical_input, retained_prior_output, - { - "role": "user", - "content": [{"type": "input_text", "text": "second question"}], - }, + *fresh_followup_items, ] second = await asyncio.wait_for( async_client.post( @@ -7595,6 +7722,154 @@ async def fake_connect_responses_websocket( assert owner_miss["preferred_account_is_continuity_owner"] is True +@pytest.mark.asyncio +async def test_v1_responses_http_bridge_replays_full_resend_after_owner_conflict(async_client, monkeypatch): + _install_bridge_settings(monkeypatch, enabled=True) + owner_account_id = await _import_account( + async_client, + "acc_http_bridge_conflict_owner", + "http-bridge-conflict-owner@example.com", + ) + alternate_account_id = await _import_account( + async_client, + "acc_http_bridge_conflict_alternate", + "http-bridge-conflict-alternate@example.com", + ) + owner_account = await _get_account(owner_account_id) + alternate_account = await _get_account(alternate_account_id) + owner_chatgpt_account_id = cast(str, owner_account.chatgpt_account_id) + alternate_chatgpt_account_id = cast(str, alternate_account.chatgpt_account_id) + owner_upstream = _ClosingBridgeUpstreamWebSocket("resp_conflict_owner") + alternate_upstream = _FakeBridgeUpstreamWebSocket("resp_conflict_replay") + selection_calls: list[dict[str, object]] = [] + connected_account_ids: list[str] = [] + connect_headers_by_account: dict[str, dict[str, str]] = {} + + async def fake_select_account_with_budget(self, deadline, **kwargs): + del self, deadline + selection_calls.append(dict(kwargs)) + preferred_account_id = cast(str | None, kwargs.get("preferred_account_id")) + excluded_account_ids = cast(set[str], kwargs.get("exclude_account_ids") or set()) + fallback_enabled = bool(kwargs.get("fallback_on_preferred_account_unavailable", True)) + if preferred_account_id == owner_account.id and not fallback_enabled: + assert kwargs.get("preferred_account_is_continuity_owner") is True + return AccountSelection( + account=None, + error_message="Account-owned continuity sources conflict; retry the logical turn", + error_code="continuity_owner_conflict", + ) + if owner_account.id in excluded_account_ids: + return AccountSelection(account=alternate_account, error_message=None, error_code=None) + return AccountSelection(account=owner_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 access_token, base_url, session + connected_account_ids.append(account_id_header) + connect_headers_by_account[account_id_header] = dict(headers) + if account_id_header == owner_chatgpt_account_id: + return owner_upstream + assert account_id_header == alternate_chatgpt_account_id + return alternate_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) + + historical_input = [ + { + "role": "user", + "content": [{"type": "input_text", "text": "first question"}], + } + ] + first = await asyncio.wait_for( + async_client.post( + "/v1/responses", + json={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": historical_input, + "prompt_cache_key": "http-bridge-conflict-replay", + }, + ), + timeout=_TEST_SYNC_TIMEOUT_SECONDS, + ) + assert first.status_code == 200, first.text + + retained_prior_output = { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "first answer"}], + } + full_resend = [ + *historical_input, + retained_prior_output, + { + "role": "user", + "content": [{"type": "input_text", "text": "second question"}], + }, + ] + second = await asyncio.wait_for( + async_client.post( + "/v1/responses", + json={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": full_resend, + "prompt_cache_key": "http-bridge-conflict-replay", + "previous_response_id": first.json()["id"], + }, + headers={ + "session_id": "stale-session", + "x-codex-session-id": "stale-codex-session", + "x-codex-turn-state": "http_turn_stale", + "x-request-trace": "keep-me", + }, + ), + timeout=_TEST_SYNC_TIMEOUT_SECONDS, + ) + + assert second.status_code == 200, second.text + assert second.json()["id"] == "resp_conflict_replay_1" + assert connected_account_ids == [ + owner_chatgpt_account_id, + alternate_chatgpt_account_id, + ] + replay_connect_headers = { + key.lower(): value for key, value in connect_headers_by_account[alternate_chatgpt_account_id].items() + } + assert replay_connect_headers["x-request-trace"] == "keep-me" + assert ( + not { + "session_id", + "x-codex-session-id", + "x-codex-turn-state", + } + & replay_connect_headers.keys() + ) + replay_payload = json.loads(alternate_upstream.sent_text[0]) + assert "previous_response_id" not in replay_payload + assert replay_payload["input"] == full_resend + owner_conflict = next( + call + for call in selection_calls + if call.get("preferred_account_id") == owner_account.id + and call.get("fallback_on_preferred_account_unavailable") is False + ) + assert owner_conflict["preferred_account_is_continuity_owner"] is True + + @pytest.mark.asyncio async def test_backend_responses_http_bridge_real_selector_recovers_full_resend_without_degrading_pool( async_client, monkeypatch @@ -8178,6 +8453,108 @@ async def fake_connect_responses_websocket( assert connect_count == 2 +@pytest.mark.asyncio +async def test_v1_responses_http_bridge_retries_when_upstream_never_acknowledges_response_create( + async_client, + monkeypatch, +): + _install_bridge_settings_with_limits( + monkeypatch, + enabled=True, + ) + proxy_module.get_settings().http_responses_session_bridge_stuck_gate_retire_after_seconds = 0.01 + account_id = await _import_account( + async_client, + "acc_http_bridge_missing_created_retry", + "http-bridge-missing-created-retry@example.com", + ) + account = await _get_account(account_id) + silent_upstreams = [_SilentUpstreamWebSocket() for _ in range(5)] + recovered_upstream = _FakeBridgeUpstreamWebSocket() + upstreams = [*silent_upstreams, recovered_upstream] + 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, + ): + del preferred_account_id + 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, + ) + 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) + + response = await asyncio.wait_for( + async_client.post( + "/v1/responses", + json={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "retry missing response.created", + "prompt_cache_key": "missing-created-retry-key", + }, + ), + timeout=_TEST_SYNC_TIMEOUT_SECONDS, + ) + + assert response.status_code == 200 + assert connect_count == len(silent_upstreams) + 1 + assert all(upstream.closed for upstream in silent_upstreams) + assert [len(upstream.sent_text) for upstream in silent_upstreams] == [1] * len(silent_upstreams) + assert len(recovered_upstream.sent_text) == 1 + + @pytest.mark.asyncio async def test_backend_responses_http_bridge_retries_precreated_server_overload(async_client, monkeypatch): _install_bridge_settings(monkeypatch, enabled=True) diff --git a/tests/integration/test_proxy_api_extended.py b/tests/integration/test_proxy_api_extended.py index 2ed0938fd1..45cf44e34b 100644 --- a/tests/integration/test_proxy_api_extended.py +++ b/tests/integration/test_proxy_api_extended.py @@ -514,6 +514,32 @@ async def fake_select(*_args, **_kwargs): assert response.json()["error"]["code"] == "no_accounts" +@pytest.mark.asyncio +async def test_thread_goal_get_maps_pool_usage_exhaustion_for_codex(async_client, monkeypatch): + async def fake_select(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="Usage limit reached", + error_code="usage_limit_reached", + ) + + monkeypatch.setattr(proxy_module.ProxyService, "_select_account_with_budget", fake_select) + + response = await async_client.post( + "/backend-api/codex/thread/goal/get", + json={"threadId": "019debd9-2372-7f23-92b9-9f34002a6355"}, + ) + + assert response.status_code == 429 + assert response.json() == { + "error": { + "message": "Usage limit reached", + "type": "usage_limit_reached", + "code": "usage_limit_reached", + } + } + + @pytest.mark.asyncio async def test_thread_goal_set_propagates_upstream_errors(async_client, monkeypatch): await _import_account(async_client, "acc_goal_set_error", "goal-set-error@example.com") diff --git a/tests/integration/test_proxy_responses.py b/tests/integration/test_proxy_responses.py index 001c876133..2fcda344be 100644 --- a/tests/integration/test_proxy_responses.py +++ b/tests/integration/test_proxy_responses.py @@ -175,6 +175,104 @@ async def test_proxy_responses_no_accounts(async_client): assert event["response"]["error"]["code"] == "no_accounts" +def _install_usage_limited_selection(monkeypatch, *, resets_at: int | None = 1_700_003_600) -> None: + async def fake_select_account(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 300s", + error_code="usage_limit_reached", + resets_at=resets_at, + ) + + monkeypatch.setattr( + "app.modules.proxy.load_balancer.LoadBalancer.select_account", + fake_select_account, + ) + + +@pytest.mark.asyncio +async def test_v1_responses_pool_usage_exhaustion_returns_429(async_client, monkeypatch): + _install_usage_limited_selection(monkeypatch) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + + response = await async_client.post("/v1/responses", json=payload) + + assert response.status_code == 429 + error = response.json()["error"] + assert error["type"] == "usage_limit_reached" + assert error["code"] == "usage_limit_reached" + assert error["resets_at"] == 1_700_003_600 + + +@pytest.mark.asyncio +async def test_v1_responses_pool_usage_exhaustion_omits_unknown_reset(async_client, monkeypatch): + _install_usage_limited_selection(monkeypatch, resets_at=None) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + + response = await async_client.post("/v1/responses", json=payload) + + assert response.status_code == 429 + error = response.json()["error"] + assert error["type"] == "usage_limit_reached" + assert error["code"] == "usage_limit_reached" + assert "resets_at" not in error + + +@pytest.mark.asyncio +async def test_backend_responses_pool_usage_exhaustion_returns_429(async_client, monkeypatch): + _install_usage_limited_selection(monkeypatch) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + request_id = "req_stream_usage_limited" + + response = await async_client.post( + "/backend-api/codex/responses", + json=payload, + headers={"x-request-id": request_id}, + ) + + # Codex only classifies a terminal response as usage-limited when it sees + # both HTTP 429 and error.type == "usage_limit_reached" (#1246). + assert response.status_code == 429 + error = response.json()["error"] + assert error["type"] == "usage_limit_reached" + assert error["code"] == "usage_limit_reached" + assert error["resets_at"] == 1_700_003_600 + + +@pytest.mark.asyncio +async def test_v1_responses_mixed_unusable_pool_keeps_no_accounts_semantics(async_client, monkeypatch): + # Paused/deactivated/reauth-only pools must keep the pre-existing + # no_accounts semantics; only usage/quota exhaustion of the whole + # eligible pool may surface the new usage_limit_reached contract. + async def fake_select_account(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="All accounts are paused, deactivated, or require re-authentication", + error_code=None, + ) + + monkeypatch.setattr( + "app.modules.proxy.load_balancer.LoadBalancer.select_account", + fake_select_account, + ) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + + async with async_client.stream("POST", "/v1/responses", json=payload) as resp: + assert resp.status_code == 200 + lines = [line async for line in resp.aiter_lines() if line] + + # The synthetic failure keeps the #1479 SDK stream contract: a sequenced + # synthetic response.created precedes the sequenced response.failed. + created = _extract_first_raw_event(lines) + assert created["type"] == "response.created" + assert created["sequence_number"] == 0 + failed = _extract_first_event(lines) + assert failed["type"] == "response.failed" + assert failed["sequence_number"] == 1 + assert failed["response"]["error"]["code"] == "no_accounts" + assert failed["response"]["error"]["type"] == "server_error" + + @pytest.mark.asyncio async def test_backend_responses_prohibits_fast_model_alias_priority_tier(async_client, monkeypatch): raw_account_id = "acc_prohibit_fast_mode" diff --git a/tests/integration/test_proxy_websocket_responses.py b/tests/integration/test_proxy_websocket_responses.py index 28c795643c..8f3d1a2b99 100644 --- a/tests/integration/test_proxy_websocket_responses.py +++ b/tests/integration/test_proxy_websocket_responses.py @@ -2411,7 +2411,7 @@ async def fake_write_request_log(self, **kwargs): assert cast(dict[str, object], visible_reference_payload["client_metadata"])[marker] == "true" -def test_backend_responses_websocket_keeps_same_response_distinct_tool_call_ids( +def test_backend_responses_websocket_suppresses_same_response_duplicate_side_effect_call_ids( app_instance, monkeypatch, ): @@ -2552,18 +2552,17 @@ async def fake_write_request_log(self, **kwargs): websocket.send_text(json.dumps(request_payload)) created_event = json.loads(websocket.receive_text()) tool_event = json.loads(websocket.receive_text()) - replay_tool_event = json.loads(websocket.receive_text()) - terminal_event = json.loads(websocket.receive_text()) + failed_event = json.loads(websocket.receive_text()) assert created_event["type"] == "response.created" assert tool_event["type"] == "response.output_item.done" assert tool_event["item"]["call_id"] == "call_first" - assert replay_tool_event["type"] == "response.output_item.done" - assert replay_tool_event["item"]["call_id"] == "call_replay" - assert terminal_event["type"] == "response.completed" - assert terminal_event["response"]["id"] == "resp_ws_duplicate_tool" + assert failed_event["type"] == "response.failed" + assert failed_event["response"]["id"] == "resp_ws_duplicate_tool" + assert failed_event["response"]["error"]["code"] == "stream_incomplete" assert len(log_calls) == 1 - assert log_calls[0]["status"] == "success" + assert log_calls[0]["status"] == "error" + assert log_calls[0]["error_code"] == "stream_incomplete" def test_backend_responses_websocket_preserves_image_generation_tool_advertisement(app_instance, monkeypatch): diff --git a/tests/unit/test_db_migrate.py b/tests/unit/test_db_migrate.py index f33537f08b..e95b7de2bc 100644 --- a/tests/unit/test_db_migrate.py +++ b/tests/unit/test_db_migrate.py @@ -1377,6 +1377,18 @@ def test_check_schema_drift_ignores_legacy_live_extra_request_log_column(tmp_pat assert check_schema_drift(url) == () +def test_check_schema_drift_ignores_legacy_live_security_lineage_columns(tmp_path: Path) -> None: + db_path = tmp_path / "legacy-security-lineage-columns.db" + url = _db_url(db_path) + + run_upgrade(url, "head", bootstrap_legacy=False) + + # The live database may already include schema from an older aggregate that + # kept security-lineage persistence. Current code tolerates those columns so + # newer deploys can move past that applied Alembic revision safely. + assert check_schema_drift(url) == () + + def test_check_schema_drift_ignores_sqlite_real_float_reflection_for_sticky_thresholds( monkeypatch, tmp_path: Path, @@ -2026,14 +2038,16 @@ def test_capability_lineage_migration_is_additive_reversible_and_single_head(tmp url = _db_url(db_path) parent_revision = "20260725_000000_add_http_bridge_pending_tool_calls" target_revision = "20260731_000000_add_capability_lineage_markers" + merge_revision = "20260802_000000_merge_bridge_and_capability_lineage_heads" run_upgrade(url, parent_revision, bootstrap_legacy=False) config = _build_alembic_config(url) script_directory = ScriptDirectory.from_config(config) - heads = script_directory.get_heads() - assert len(heads) == 1 - ancestry = {script.revision for script in script_directory.walk_revisions()} - assert target_revision in ancestry + assert script_directory.get_heads() == [merge_revision] + assert script_directory.get_revision(merge_revision).down_revision == ( + "20260729_000000_drop_legacy_bridge_pending_tool_columns", + target_revision, + ) 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 7eb8ec51c2..4ed48abf8e 100644 --- a/tests/unit/test_durable_bridge_sessions.py +++ b/tests/unit/test_durable_bridge_sessions.py @@ -1760,6 +1760,54 @@ async def test_durable_bridge_same_account_closed_takeover_preserves_restart_anc assert reclaimed.latest_response_id == "resp_old" +@pytest.mark.asyncio +async def test_durable_bridge_terminal_release_can_clear_restart_anchor( + coordinator: DurableBridgeSessionCoordinator, +) -> None: + claimed = await coordinator.claim_live_session( + session_key_kind="session_header", + session_key_value="sid-terminal-clear", + api_key_id=None, + instance_id="instance-a", + lease_ttl_seconds=60.0, + account_id="acc-1", + model="gpt-5.4", + service_tier=None, + latest_turn_state="http_turn_old", + latest_response_id="resp_old", + allow_takeover=True, + ) + released = await coordinator.release_live_session( + session_id=claimed.session_id, + instance_id="instance-a", + owner_epoch=claimed.owner_epoch, + draining=False, + clear_continuity=True, + ) + + assert released is not None + assert released.latest_turn_state is None + assert released.latest_response_id is None + + reclaimed = await coordinator.claim_live_session( + session_key_kind="session_header", + session_key_value="sid-terminal-clear", + api_key_id=None, + instance_id="instance-b", + lease_ttl_seconds=60.0, + account_id="acc-1", + model="gpt-5.4", + service_tier=None, + latest_turn_state=None, + latest_response_id=None, + allow_takeover=True, + ) + + assert reclaimed.owner_instance_id == "instance-b" + assert reclaimed.latest_turn_state is None + assert reclaimed.latest_response_id is None + + @pytest.mark.asyncio async def test_durable_bridge_takeover_preserves_existing_anchor_when_replacement_has_none( coordinator: DurableBridgeSessionCoordinator, diff --git a/tests/unit/test_http_bridge_cancel_drain.py b/tests/unit/test_http_bridge_cancel_drain.py index 46f8b818a9..0a0794e95a 100644 --- a/tests/unit/test_http_bridge_cancel_drain.py +++ b/tests/unit/test_http_bridge_cancel_drain.py @@ -93,14 +93,17 @@ async def test_cancelled_stream_settlement_task_releases_reservation( service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) scheduled: list[tuple[str, str]] = [] cleanup_tasks: list[asyncio.Task[None]] = [] + release_retry_flags: list[bool] = [] async def release_unsettled( *, api_key: ApiKeyData, api_key_reservation: ApiKeyUsageReservationData, request_id: str, + retry_persistence_failures: bool = False, ) -> None: scheduled.append((api_key.id, api_key_reservation.reservation_id)) + release_retry_flags.append(retry_persistence_failures) def schedule_cleanup( coro: Any, @@ -133,6 +136,7 @@ def schedule_cleanup( assert ("release_stream_api_key_reservation_after_cancelled_settlement", "req-cancel-settle") in scheduled assert ("key-cancel-settle", "res-cancel-settle") in scheduled + assert release_retry_flags == [True] @pytest.mark.asyncio @@ -209,10 +213,8 @@ async def __aexit__(self, *args: object) -> None: @pytest.mark.asyncio async def test_http_bridge_stream_waits_only_while_completed_delivery_is_active( monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, terminal_outcome: str, ) -> None: - caplog.set_level(logging.INFO, logger="app.modules.proxy.service") service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) queue_waiting = asyncio.Event() terminal_claimed = asyncio.Event() @@ -323,208 +325,6 @@ async def consume_stream() -> list[str]: assert request_state.completed_delivery_scope.active is False assert parse_sse_data_json.call_count == 1 assert parse_sse_event_payload.call_count == 1 - suppression_messages = [ - record.getMessage() - for record in caplog.records - if "HTTP bridge stream idle timeout suppressed during completed delivery" in record.getMessage() - ] - assert len(suppression_messages) == 1 - assert "request_id=req-terminal-race" in suppression_messages[0] - assert "response_id=resp-terminal-race" in suppression_messages[0] - assert "elapsed_seconds=" in suppression_messages[0] - - -@pytest.mark.asyncio -async def test_http_bridge_stream_idle_timeout_revokes_queue_before_completed_claim( - monkeypatch: pytest.MonkeyPatch, -) -> None: - service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) - queue_waiting = asyncio.Event() - - class ObservedQueue(asyncio.Queue[str | None]): - async def get(self) -> str | None: - queue_waiting.set() - return await super().get() - - event_queue = ObservedQueue() - request_state = _make_request_state( - "req-timeout-first", - response_id="resp-timeout-first", - awaiting_response_created=False, - event_queue=event_queue, - ) - session = _make_http_bridge_session(deque(), queued_request_count=0) - backing_lock = anyio.Lock() - queue_revoked_before_lock_release = False - - class ObservedPendingLock: - async def __aenter__(self) -> ObservedPendingLock: - await backing_lock.acquire() - return self - - async def __aexit__(self, *args: object) -> None: - nonlocal queue_revoked_before_lock_release - if request_state.event_queue is None: - queue_revoked_before_lock_release = True - backing_lock.release() - - session.pending_lock = cast(Any, ObservedPendingLock()) - - 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 - async with target_session.pending_lock: - target_session.pending_requests.append(request_state) - target_session.queued_request_count += 1 - - async def fake_detach_http_bridge_request( - target_session: proxy_service._HTTPBridgeSession, - *, - request_state: proxy_service._WebSocketRequestState, - ) -> bool: - del target_session - assert request_state.event_queue is None - return False - - finalize_request = AsyncMock() - monkeypatch.setattr(service, "_submit_http_bridge_request", fake_submit_http_bridge_request) - monkeypatch.setattr(service, "_detach_http_bridge_request", fake_detach_http_bridge_request) - monkeypatch.setattr(service, "_register_http_bridge_previous_response_id", AsyncMock(return_value=True)) - monkeypatch.setattr(service, "_finalize_websocket_request_state", finalize_request) - 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) - monkeypatch.setattr(proxy_service, "_STREAM_KEEPALIVE_MAX_COUNT", 1) - - async def consume_stream() -> list[str]: - return [ - event_block - async for event_block in service._stream_http_bridge_session_events( - session, - request_state=request_state, - text_data="{}", - queue_limit=8, - propagate_http_errors=False, - downstream_turn_state=None, - ) - ] - - stream_task = asyncio.create_task(consume_stream()) - await asyncio.wait_for(queue_waiting.wait(), timeout=1.0) - event_blocks = await asyncio.wait_for(stream_task, timeout=1.0) - - assert queue_revoked_before_lock_release is True - assert request_state.event_queue is None - assert request_state in session.pending_requests - assert "stream_idle_timeout" in "".join(event_blocks) - - terminal_text = '{"type":"response.completed","response":{"id":"resp-timeout-first","status":"completed"}}' - await service._process_http_bridge_upstream_text(session, terminal_text) - - assert request_state.completed_delivery_scope is None - assert event_queue.empty() - finalize_request.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_http_bridge_completed_delivery_stays_dominant_after_recovery_wait( - monkeypatch: pytest.MonkeyPatch, -) -> None: - service = proxy_service.ProxyService(cast(Any, SimpleNamespace())) - recovery_started = asyncio.Event() - release_recovery = asyncio.Event() - event_queue: asyncio.Queue[str | None] = asyncio.Queue() - request_state = _make_request_state( - "req-completed-during-recovery", - response_id=None, - awaiting_response_created=True, - event_queue=event_queue, - ) - session = _make_http_bridge_session(deque(), queued_request_count=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 - async with target_session.pending_lock: - target_session.pending_requests.append(request_state) - target_session.queued_request_count += 1 - - async def block_idle_recovery(*args: Any, **kwargs: Any) -> bool: - del args, kwargs - recovery_started.set() - await release_recovery.wait() - return False - - finalize_request = AsyncMock() - monkeypatch.setattr(service, "_submit_http_bridge_request", fake_submit_http_bridge_request) - monkeypatch.setattr(service, "_retry_http_bridge_precreated_request", block_idle_recovery) - monkeypatch.setattr(service, "_http_bridge_precreated_retry_cooldown_seconds", AsyncMock(return_value=0.0)) - monkeypatch.setattr(service, "_register_http_bridge_previous_response_id", AsyncMock(return_value=True)) - monkeypatch.setattr(service, "_record_http_bridge_retry_circuit_failure", AsyncMock()) - monkeypatch.setattr(service, "_finalize_websocket_request_state", finalize_request) - 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) - monkeypatch.setattr(proxy_service, "_STREAM_KEEPALIVE_MAX_COUNT", 1) - - async def consume_stream() -> list[str]: - return [ - event_block - async for event_block in service._stream_http_bridge_session_events( - session, - request_state=request_state, - text_data="{}", - queue_limit=8, - propagate_http_errors=False, - downstream_turn_state=None, - ) - ] - - stream_task = asyncio.create_task(consume_stream()) - await asyncio.wait_for(recovery_started.wait(), timeout=1.0) - - terminal_text = ( - '{"type":"response.completed","response":{"id":"resp-completed-during-recovery","status":"completed"}}' - ) - await service._process_http_bridge_upstream_text(session, terminal_text) - - assert request_state.completed_delivery_scope is not None - assert request_state.completed_delivery_scope.active is False - assert request_state.completed_delivery_scope.terminal_enqueued is True - release_recovery.set() - event_blocks = await asyncio.wait_for(stream_task, timeout=1.0) - - event_types = [ - payload["type"] - for event_block in event_blocks - if isinstance(payload := proxy_service.parse_sse_data_json(event_block), dict) - ] - assert event_types[-1] == "response.completed" - assert event_types.count("response.completed") == 1 - assert "response.failed" not in event_types - finalize_request.assert_awaited_once() def test_retiring_http_bridge_session_is_not_reusable() -> None: diff --git a/tests/unit/test_load_balancer.py b/tests/unit/test_load_balancer.py index f2d67e1d9e..f6b8c12662 100644 --- a/tests/unit/test_load_balancer.py +++ b/tests/unit/test_load_balancer.py @@ -676,6 +676,343 @@ def test_select_account_skips_rate_limited_until_reset(): assert result.account.account_id == "b" +def test_select_account_reports_pool_wide_usage_exhaustion_structurally(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + reset_at=int(now + 600), + primary_reset_at=int(now + 60), + ), + AccountState( + "b", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 7200), + primary_reset_at=int(now + 3600), + ), + AccountState("paused", AccountStatus.PAUSED), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Rate limit exceeded. Try again in 60s" + assert result.resets_at == int(now + 60) + + +def test_select_account_fails_over_when_one_account_remains_usable(): + now = 1_700_000_000.0 + states = [ + AccountState("exhausted", AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=int(now + 3600)), + AccountState("usable", AccountStatus.ACTIVE, used_percent=40.0), + ] + + result = select_account(states, now=now) + + assert result.account is not None + assert result.account.account_id == "usable" + assert result.error_code is None + assert result.error_message is None + + +def test_select_account_reports_secondary_usage_exhaustion_reset(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=10.0, + secondary_used_percent=100.0, + reset_at=int(now + 60), + secondary_reset_at=int(now + 3600), + ) + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Rate limit exceeded. Try again in 300s" + assert result.resets_at == int(now + 3600) + + +def test_select_account_omits_synthesized_primary_usage_reset(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + reset_at=int(now + 60), + primary_reset_at=None, + ) + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Usage limit reached" + assert result.resets_at is None + + +def test_select_account_waits_for_latest_exhausted_window_per_account(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + secondary_used_percent=100.0, + reset_at=int(now + 600), + primary_reset_at=int(now + 60), + secondary_reset_at=int(now + 3600), + ), + AccountState( + "b", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 7200), + primary_reset_at=int(now + 7200), + ), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Rate limit exceeded. Try again in 300s" + assert result.resets_at == int(now + 3600) + + +def test_select_account_requires_usage_window_evidence_for_quota_exhaustion(): + now = 1_700_000_000.0 + states = [ + AccountState("a", AccountStatus.QUOTA_EXCEEDED, reset_at=int(now + 3600)), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_select_account_can_disable_pool_usage_exhaustion_for_owner_scope(): + now = 1_700_000_000.0 + states = [ + AccountState( + "owner", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 3600), + ) + ] + + result = select_account(states, now=now, allow_usage_exhaustion_error=False) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_budget_safe_selection_uses_full_scope_for_usage_exhaustion() -> None: + now = time.time() + cap_filtered_states = [ + AccountState( + "exhausted", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 3600), + ) + ] + full_scope_states = [ + AccountState( + "capped-but-usable", + AccountStatus.ACTIVE, + used_percent=50.0, + reset_at=int(now + 3600), + ), + *cap_filtered_states, + ] + + result = _select_account_preferring_budget_safe( + cap_filtered_states, + prefer_earlier_reset=False, + routing_strategy="usage_weighted", + budget_threshold_pct=95.0, + usage_exhaustion_states=full_scope_states, + ) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_budget_safe_capacity_selection_forwards_usage_exhaustion_controls() -> None: + now = time.time() + owner_scope = [ + AccountState( + "owner", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 600), + primary_reset_at=int(now + 3600), + ) + ] + full_scope = [ + *owner_scope, + AccountState( + "pool-usable", + AccountStatus.ACTIVE, + used_percent=10.0, + ), + ] + + result = _select_account_preferring_budget_safe( + owner_scope, + prefer_earlier_reset=False, + routing_strategy="capacity_weighted", + budget_threshold_pct=95.0, + allow_usage_exhaustion_error=False, + usage_exhaustion_states=full_scope, + ) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_opportunistic_budget_safe_selection_uses_full_scope_for_usage_exhaustion() -> None: + now = time.time() + cap_filtered_states = [ + AccountState( + "exhausted", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 3600), + ) + ] + full_scope_states = [ + AccountState( + "capped-but-usable", + AccountStatus.ACTIVE, + used_percent=50.0, + reset_at=int(now + 3600), + ), + *cap_filtered_states, + ] + + result = _select_account_preferring_budget_safe( + cap_filtered_states, + prefer_earlier_reset=False, + routing_strategy="usage_weighted", + budget_threshold_pct=95.0, + traffic_class="opportunistic", + usage_exhaustion_states=full_scope_states, + ) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_select_account_uses_raw_priority_usage_for_exhaustion_evidence() -> None: + now = 1_700_000_000.0 + states = [ + AccountState( + "pressure-adjusted", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + priority_used_percent=98.0, + reset_at=int(now + 60), + ) + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "No available accounts" + + +def test_select_account_does_not_treat_generic_rate_limit_as_usage_exhaustion(): + now = 1_700_000_000.0 + states = [ + AccountState("a", AccountStatus.RATE_LIMITED, used_percent=5.0, reset_at=int(now + 60)), + AccountState("b", AccountStatus.RATE_LIMITED, secondary_used_percent=10.0, reset_at=int(now + 120)), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "No available accounts" + + +def test_select_account_does_not_misclassify_transient_backoff_as_usage_exhaustion(): + now = 1_700_000_000.0 + states = [ + AccountState("quota", AccountStatus.QUOTA_EXCEEDED, reset_at=int(now + 3600)), + AccountState( + "transient", + AccountStatus.ACTIVE, + error_count=3, + last_error_at=now, + ), + ] + + result = select_account(states, now=now, allow_backoff_fallback=False) + + assert result.account is None + assert result.error_code is None + + +def test_select_account_does_not_report_ignored_standard_quota_as_pool_exhaustion(): + now = 1_700_000_000.0 + state = AccountState("quota", AccountStatus.QUOTA_EXCEEDED, reset_at=int(now + 3600)) + + result = select_account([state], now=now, ignore_standard_quota=True) + + assert result.account is not None + assert result.account.account_id == "quota" + + +def test_select_account_excludes_per_account_standard_quota_bypass_from_pool_exhaustion(): + now = 1_700_000_000.0 + state = AccountState( + "quota", + AccountStatus.QUOTA_EXCEEDED, + reset_at=int(now + 3600), + cooldown_until=now + 30, + ignore_standard_quota=True, + ) + + result = select_account([state], now=now) + + assert result.account is None + assert result.error_code is None + + +def test_select_account_excludes_scoped_standard_quota_bypass_from_pool_exhaustion(): + now = 1_700_000_000.0 + state = AccountState( + "quota", + AccountStatus.QUOTA_EXCEEDED, + reset_at=int(now + 3600), + cooldown_until=now + 30, + ) + + result = select_account([state], now=now, bypass_quota_exceeded_account_ids={"quota"}) + + assert result.account is None + assert result.error_code is None + + def test_select_account_reports_paused_and_deactivated_without_reauth_reason(): states = [ AccountState("paused", AccountStatus.PAUSED, used_percent=5.0), @@ -1291,12 +1628,14 @@ def test_select_account_caps_quota_exceeded_retry_hint(): AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=far_future_reset, + primary_reset_at=far_future_reset, ), AccountState( "b", AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=int(now + 271_819), + primary_reset_at=int(now + 271_819), ), ] result = select_account(states, now=now) @@ -1315,6 +1654,7 @@ def test_select_account_preserves_short_quota_exceeded_retry_hint(): AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=int(now + 60), + primary_reset_at=int(now + 60), ), ] result = select_account(states, now=now) @@ -1878,6 +2218,52 @@ def test_state_from_account_keeps_active_account_selectable_when_primary_usage_s assert selection.account.account_id == state.account_id +def test_state_from_account_keeps_raw_usage_evidence_separate_from_pressure(monkeypatch): + now = 1_700_000_000.0 + future_reset = int(now + 300) + monkeypatch.setattr("app.modules.proxy.load_balancer.time.time", lambda: now) + monkeypatch.setattr("app.core.usage.quota.time.time", lambda: now) + + state = _state_from_account( + account=_make_test_account(status=AccountStatus.RATE_LIMITED, reset_at=future_reset, blocked_at=int(now)), + primary_entry=_make_test_usage( + window="primary", + used_percent=98.0, + reset_at=future_reset, + recorded_at=_epoch_to_naive_utc(now - 30), + ), + secondary_entry=None, + runtime=RuntimeState(inflight_streams=1), + ) + + assert state.used_percent == 100.0 + assert state.priority_used_percent == 98.0 + + +def test_state_from_account_preserves_pressure_for_active_routing(monkeypatch): + now = 1_700_000_000.0 + future_reset = int(now + 300) + monkeypatch.setattr("app.modules.proxy.load_balancer.time.time", lambda: now) + monkeypatch.setattr("app.core.usage.quota.time.time", lambda: now) + + state = _state_from_account( + account=_make_test_account(status=AccountStatus.ACTIVE), + primary_entry=_make_test_usage( + window="primary", + used_percent=94.0, + reset_at=future_reset, + recorded_at=_epoch_to_naive_utc(now - 30), + ), + secondary_entry=None, + runtime=RuntimeState(inflight_streams=1), + ) + + assert state.status == AccountStatus.ACTIVE + assert state.used_percent == 96.5 + assert state.priority_used_percent is None + assert _state_above_sticky_budget_threshold(state, 95.0) is True + + def test_state_from_account_clears_stale_advisory_account_reset_for_active_account(monkeypatch): now = 1_700_000_000.0 future_reset = int(now + 300) diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index f179ccf2ac..be94c2f7fd 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -120,6 +120,28 @@ async def test_account_lease_uses_explicit_dashboard_cap_snapshot_not_startup_en assert third is None +@pytest.mark.asyncio +async def test_opportunistic_selection_preserves_usage_limit_exhaustion_error() -> None: + account = _make_account("acc-opportunistic-usage-exhausted") + account.status = AccountStatus.QUOTA_EXCEEDED + reset_at = int(time.time() + 300) + account.reset_at = reset_at + usage_repo = _StubUsageRepository( + {account.id: _usage_row_with_percent(1, account.id, used_percent=100.0, reset_at=reset_at)}, + {}, + ) + balancer = LoadBalancer(lambda: _repo_factory(_StubAccountsRepository([account]), usage_repo)) + + result = await balancer.select_account( + routing_strategy="usage_weighted", + traffic_class=load_balancer_module.TRAFFIC_CLASS_OPPORTUNISTIC, + ) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.resets_at == reset_at + + class _StubAccountsRepository: def __init__(self, accounts: list[Account]) -> None: self._accounts = accounts @@ -801,6 +823,57 @@ async def test_account_stream_cap_returns_stable_local_reason_until_released() - assert recovered.lease is not None +@pytest.mark.asyncio +async def test_stream_cap_takes_precedence_over_remaining_quota_exhausted_account() -> None: + now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) + capped = _make_account("acc-stream-cap-mixed-capped") + exhausted = _make_account("acc-stream-cap-mixed-exhausted") + exhausted.status = AccountStatus.QUOTA_EXCEEDED + exhausted.reset_at = now_epoch + 3600 + accounts_repo = _StubAccountsRepository([capped, exhausted]) + usage_repo = _StubUsageRepository( + primary={ + capped.id: _usage_row_with_percent( + 203, + capped.id, + used_percent=50.0, + reset_at=now_epoch + 300, + ), + exhausted.id: _usage_row_with_percent( + 204, + exhausted.id, + used_percent=100.0, + reset_at=now_epoch + 3600, + ), + }, + secondary={}, + ) + balancer = LoadBalancer(lambda: _repo_factory(accounts_repo, usage_repo)) + leases = [ + ( + await balancer.select_account( + routing_strategy="usage_weighted", + lease_kind="stream", + ) + ).lease + for _ in range(8) + ] + + selected = await balancer.select_account( + routing_strategy="usage_weighted", + lease_kind="stream", + ) + + assert selected.account is None + assert selected.error_code == "account_stream_cap" + assert selected.error_message is not None + assert "Account stream capacity is exhausted" in selected.error_message + assert selected.resets_at is None + + for lease in leases: + await balancer.release_account_lease(lease) + + @pytest.mark.asyncio async def test_account_stream_recovery_reserve_keeps_last_slot_for_reattach() -> None: now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) diff --git a/tests/unit/test_openai_errors.py b/tests/unit/test_openai_errors.py index 2534ead7c7..70ec2d136f 100644 --- a/tests/unit/test_openai_errors.py +++ b/tests/unit/test_openai_errors.py @@ -29,6 +29,18 @@ def test_response_failed_event_accepts_incomplete_details(): assert response.get("incomplete_details") == {"reason": "max_output_tokens"} +def test_response_failed_event_preserves_reset_hint(): + event = response_failed_event( + "usage_limit_reached", + "Rate limit exceeded. Try again in 1h", + error_type="usage_limit_reached", + response_id="resp_1", + resets_at=1_700_003_600, + ) + + assert event["response"]["error"]["resets_at"] == 1_700_003_600 + + def test_previous_response_not_found_classifier_covers_openai_shapes(): assert is_previous_response_not_found_error( code="previous_response_not_found", diff --git a/tests/unit/test_openai_requests.py b/tests/unit/test_openai_requests.py index 8074bc0c7d..6a6b32eccf 100644 --- a/tests/unit/test_openai_requests.py +++ b/tests/unit/test_openai_requests.py @@ -2293,6 +2293,42 @@ def test_compact_trimming_keeps_selected_tool_calls_with_matching_outputs(): assert tool_output in dumped_input +def test_compact_trimming_keeps_tool_search_outputs_with_matching_calls(): + tool_call = { + "type": "tool_search_call", + "call_id": "call_search_tail", + "status": "completed", + "execution": "client", + "arguments": {"query": "spawn_agent multi-agent schema", "limit": 8}, + } + tool_output = { + "type": "tool_search_output", + "call_id": "call_search_tail", + "output": "Found matching tools", + } + input_items = [ + {"role": "user", "content": "initial instructions"}, + {"role": "assistant", "content": "x" * 500_000}, + tool_call, + {"role": "assistant", "content": "y" * 500_000}, + tool_output, + {"role": "user", "content": "latest request"}, + ] + payload = { + "model": "gpt-5.1", + "instructions": "hi", + "input": input_items, + } + + request = ResponsesCompactRequest.model_validate(payload) + dumped = request.to_payload() + dumped_input = dumped["input"] + + assert isinstance(dumped_input, list) + assert tool_call in dumped_input + assert tool_output in dumped_input + + def test_compact_trimming_reconciles_duplicate_tool_call_ids_by_occurrence(): first_tool_call = { "type": "function_call", diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index bf9055aa21..0fa9617a59 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -32,7 +32,7 @@ from app.core.config.settings import Settings from app.core.errors import openai_error from app.core.utils.request_id import get_request_id, reset_request_scope_id, set_request_scope_id -from app.db.models import AccountStatus, HttpBridgeSessionState +from app.db.models import Account, AccountStatus, HttpBridgeSessionState from app.modules.proxy import http_bridge_forwarding as http_bridge_forwarding_module from app.modules.proxy import service as proxy_service from app.modules.proxy._service import support as proxy_support_module @@ -149,14 +149,14 @@ def test_http_bridge_eventless_precreated_deadline_uses_current_send_and_client_ request_state, stuck_gate_retire_after_seconds=300.0, ) - == 160.0 + == 115.0 ) assert ( http_bridge_helpers_module._http_bridge_eventless_precreated_deadline( request_state, - stuck_gate_retire_after_seconds=30.0, + stuck_gate_retire_after_seconds=10.0, ) - == 130.0 + == 110.0 ) request_state.latency_first_upstream_event_ms = 25 @@ -166,7 +166,7 @@ def test_http_bridge_eventless_precreated_deadline_uses_current_send_and_client_ request_state, stuck_gate_retire_after_seconds=300.0, ) - == 160.0 + == 115.0 ) @@ -175,7 +175,6 @@ def test_http_bridge_eventless_precreated_deadline_uses_current_send_and_client_ [ ("response_id", "resp-created"), ("latency_response_created_ms", 12), - ("response_event_count", 1), ("downstream_visible", True), ("last_downstream_sequence_number", 0), ("awaiting_response_created", False), @@ -200,6 +199,24 @@ def test_http_bridge_eventless_precreated_deadline_requires_narrow_owner_evidenc ) +def test_http_bridge_eventless_precreated_deadline_survives_reasoning_prelude_without_created() -> None: + request_state = _make_eventless_http_bridge_owner() + request_state.response_event_count = 3 + request_state.upstream_model_output_seen = True + request_state.deferred_reasoning_downstream_texts.append( + 'data: {"type":"response.output_item.added","item":{"type":"reasoning"}}\n\n' + ) + client_safe_cap_seconds = http_bridge_helpers_module._HTTP_BRIDGE_EVENTLESS_RESPONSE_CREATED_MAX_SECONDS + + assert ( + http_bridge_helpers_module._http_bridge_eventless_precreated_deadline( + request_state, + stuck_gate_retire_after_seconds=300.0, + ) + == 100.0 + min(300.0, client_safe_cap_seconds) + ) + + @pytest.mark.asyncio async def test_http_bridge_send_replaces_timestamp_and_wakes_existing_reader( monkeypatch: pytest.MonkeyPatch, @@ -282,6 +299,123 @@ async def send_text(_text: str) -> None: ) +@pytest.mark.asyncio +async def test_http_bridge_missing_response_created_retries_once_before_terminal( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_text = ( + '{"type":"response.create","model":"gpt-5.6-sol","previous_response_id":"resp-owner-anchor","input":"hello"}' + ) + request_state = _make_eventless_http_bridge_owner(request_id="req-missing-created-once") + request_state.previous_response_id = "resp-owner-anchor" + request_state.request_text = request_text + request_state.bridge_request_deadline = time.monotonic() - 1.0 + session = _make_bridge_session( + key_value="missing-created-once", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + send_text = AsyncMock() + session.upstream = cast( + UpstreamWebSocket, + SimpleNamespace(send_text=send_text, close=AsyncMock()), + ) + reconnect = AsyncMock() + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + monkeypatch.setattr(service, "_acquire_account_response_create_lease_or_overload", AsyncMock(return_value=None)) + + first_retry = await service._retry_http_bridge_precreated_request( + session, + allow_expired_deadline=True, + ) + second_retry = await service._retry_http_bridge_precreated_request( + session, + allow_expired_deadline=True, + ) + + assert first_retry is True + assert second_retry is False + reconnect.assert_awaited_once_with( + session, + request_state=request_state, + require_same_account=True, + ) + send_text.assert_awaited_once() + send_text_call = send_text.await_args + assert send_text_call is not None + assert json.loads(send_text_call.args[0])["previous_response_id"] == "resp-owner-anchor" + assert request_state.missing_response_created_retry_count == 1 + assert request_state.replay_count == 1 + assert request_state.awaiting_response_created is True + + +@pytest.mark.asyncio +async def test_http_bridge_missing_response_created_rebinds_hard_owner_when_full_resend_is_safe( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + anchored_text = ( + '{"type":"response.create","model":"gpt-5.6-sol","previous_response_id":"resp-owner-anchor","input":"trimmed"}' + ) + fresh_text = ( + '{"type":"response.create","model":"gpt-5.6-sol",' + '"input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}' + ) + request_state = _make_eventless_http_bridge_owner(request_id="req-missing-created-rebind") + request_state.previous_response_id = "resp-owner-anchor" + request_state.proxy_injected_previous_response_id = True + request_state.fresh_upstream_request_is_retry_safe = True + request_state.fresh_upstream_request_text = fresh_text + request_state.request_text = anchored_text + request_state.bridge_request_deadline = time.monotonic() - 1.0 + session = _make_bridge_session( + key=proxy_service._HTTPBridgeSessionKey("session_header", "missing-created-rebind", None), + key_value="missing-created-rebind", + pending_requests=deque([request_state]), + queued_request_count=1, + ) + original_affinity = session.affinity + send_text = AsyncMock() + session.upstream = cast( + UpstreamWebSocket, + SimpleNamespace(send_text=send_text, close=AsyncMock()), + ) + reconnect = AsyncMock() + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + monkeypatch.setattr(service, "_acquire_account_response_create_lease_or_overload", AsyncMock(return_value=None)) + + retried = await service._retry_http_bridge_precreated_request( + session, + allow_expired_deadline=True, + ) + + assert retried is True + reconnect.assert_awaited_once() + reconnect_call = reconnect.await_args + assert reconnect_call is not None + reconnect_kwargs = reconnect_call.kwargs + assert reconnect_kwargs["request_state"] is request_state + assert reconnect_kwargs["require_same_account"] is False + assert reconnect_kwargs["owner_rebind_affinity"] is original_affinity + selection_affinity = reconnect_kwargs["selection_affinity"] + assert selection_affinity.key is None + assert selection_affinity.kind is None + assert selection_affinity.reallocate_sticky is True + send_text.assert_awaited_once() + send_text_call = send_text.await_args + assert send_text_call is not None + assert json.loads(send_text_call.args[0]).get("previous_response_id") is None + assert request_state.request_text == fresh_text + assert request_state.previous_response_id is None + assert request_state.preferred_account_id is None + assert request_state.excluded_account_ids == {"acc-bridge"} + assert request_state.affinity_policy.reallocate_sticky is True + assert request_state.missing_response_created_retry_count == 1 + assert request_state.replay_count == 1 + assert request_state.awaiting_response_created is True + + def _make_account_neutral_replay_session_key( nonce: str, api_key_id: str | None = None, @@ -1542,6 +1676,88 @@ async def fake_retire( assert session.closed is True +@pytest.mark.asyncio +async def test_http_bridge_stream_gate_wait_retires_stale_pre_submit_holder( + 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() + service._http_bridge_sessions[session.key] = session + await session.response_create_gate.acquire() + old_pending = proxy_service._WebSocketRequestState( + request_id="req-old-pre-submit-holder", + model="gpt-5.4-mini", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic() - 301.0, + transport="http", + response_create_gate=session.response_create_gate, + response_create_gate_acquired=True, + awaiting_response_created=True, + downstream_visible=False, + ) + waiter = proxy_service._WebSocketRequestState( + request_id="req-gate-waiter", + model="gpt-5.4-mini", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=time.monotonic(), + event_queue=asyncio.Queue(), + transport="http", + request_text='{"type":"response.create","model":"gpt-5.4-mini","input":"retry"}', + ) + 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 + + async def no_wait_capacity_sse(**_kwargs: object): + if False: + yield "" + + monkeypatch.setattr(service, "_retire_stale_pending_http_bridge_session", fake_retire) + monkeypatch.setattr(http_bridge_streaming_module, "_iter_account_capacity_wait_sse", no_wait_capacity_sse) + waiter_text = waiter.request_text + assert waiter_text is not None + + events = service._stream_http_bridge_session_events( + session, + request_state=waiter, + text_data=waiter_text, + queue_limit=8, + propagate_http_errors=False, + downstream_turn_state=None, + ) + try: + with pytest.raises(ProxyResponseError) as exc_info: + async for _event in events: + pass + finally: + await events.aclose() + 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[0] == "response_create_gate_timeout_stuck_pending" + assert session.closed is True + + @pytest.mark.asyncio @pytest.mark.parametrize( ( @@ -3757,7 +3973,7 @@ async def test_recovery_completed_alias_persistence_failure_fails_response_and_r assert finalize_call.kwargs["event_type"] == "response.failed" assert await service._retire_http_bridge_after_drain_if_ready(session) is True - close_session.assert_awaited_once_with(session) + close_session.assert_awaited_once_with(session, clear_continuity=False) @pytest.mark.asyncio @@ -6344,6 +6560,112 @@ async def sleep_for_recovery(*_args: object, **kwargs: object) -> bool: assert sleep_calls[0]["max_sleep_seconds"] == pytest.approx(119.5) +@pytest.mark.asyncio +async def test_reconnect_http_bridge_session_skips_capacity_wait_for_usage_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + + async def select_account(_deadline: float, **_kwargs: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + + request_state = proxy_service._WebSocketRequestState( + request_id="req-reconnect-usage-limit-now", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=100.0, + ) + monkeypatch.setattr(proxy_service.time, "monotonic", lambda: 100.5) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=settings)), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr( + http_bridge_mixin_module, + "_sleep_for_account_selection_recovery", + lambda *_args, **_kwargs: pytest.fail("usage_limit_reached must not enter recovery wait"), + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._reconnect_http_bridge_session(session, request_state=request_state) + + assert exc_info.value.status_code == 429 + assert exc_info.value.payload["error"]["code"] == "usage_limit_reached" + assert exc_info.value.payload["error"]["type"] == "usage_limit_reached" + assert exc_info.value.payload["error"]["resets_at"] == 1_700_003_600 + + +@pytest.mark.asyncio +async def test_reconnect_http_bridge_session_preserves_owner_error_for_owner_usage_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + + async def select_account(_deadline: float, **_kwargs: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + + request_state = proxy_service._WebSocketRequestState( + request_id="req-reconnect-owner-usage-limit", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=100.0, + preferred_account_id=session.account.id, + ) + monkeypatch.setattr(proxy_service.time, "monotonic", lambda: 100.5) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=settings)), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr( + http_bridge_mixin_module, + "_sleep_for_account_selection_recovery", + lambda *_args, **_kwargs: pytest.fail("owner-only usage_limit_reached must not enter recovery wait"), + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._reconnect_http_bridge_session( + session, + request_state=request_state, + require_preferred_account=True, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.payload["error"]["code"] == "previous_response_owner_unavailable" + assert exc_info.value.payload["error"]["type"] == "server_error" + + @pytest.mark.asyncio async def test_reconnect_http_bridge_session_preserves_exclusions_after_capacity_wait( monkeypatch: pytest.MonkeyPatch, @@ -8042,6 +8364,33 @@ def fake_prepare( False, id="retained-assistant-output", ), + pytest.param( + [ + { + "type": "message", + "role": "assistant", + "phase": "final_answer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-previous"}, + "content": [{"type": "output_text", "text": "hello back"}], + }, + { + "type": "message", + "role": "user", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "follow up"}], + }, + { + "type": "message", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "new control message"}], + }, + ], + None, + True, + False, + id="retained-assistant-output-with-fresh-developer-followup", + ), pytest.param( [ { @@ -8064,21 +8413,48 @@ def fake_prepare( pytest.param( [ { - "type": "function_call", + "type": "custom_tool_call", "call_id": "call-1", - "name": "lookup", - "arguments": "{}", + "name": "shell", + "input": "pwd", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, }, { - "type": "function_call_output", + "type": "message", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "new control message"}], + }, + { + "type": "custom_tool_call_output", "call_id": "call-1", - "output": "result", + "output": "/workspace", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, }, ], - None, - False, + {"call-1": "custom_tool_call"}, + True, False, - id="tool-loop-with-unknown-manifest", + id="self-contained-tool-loop-with-fresh-developer-interleave", + ), + pytest.param( + [ + { + "type": "function_call", + "call_id": "call-1", + "name": "lookup", + "arguments": "{}", + }, + { + "type": "function_call_output", + "call_id": "call-1", + "output": "result", + }, + ], + None, + False, + False, + id="tool-loop-with-unknown-manifest", ), pytest.param( [{"role": "user", "content": "revise that answer"}], @@ -8846,9 +9222,14 @@ async def test_close_http_bridge_session_bounded_timeout_keeps_close_task_runnin close_finished = asyncio.Event() close_cancelled = False - async def close_http_bridge_session(target: proxy_service._HTTPBridgeSession) -> None: + async def close_http_bridge_session( + target: proxy_service._HTTPBridgeSession, + *, + clear_continuity: bool = False, + ) -> None: nonlocal close_cancelled assert target is session + assert clear_continuity is False close_started.set() try: await release_close.wait() @@ -8970,9 +9351,14 @@ async def test_close_http_bridge_session_bounded_cancellation_keeps_close_task_t close_finished = asyncio.Event() close_cancelled = False - async def close_http_bridge_session(target: proxy_service._HTTPBridgeSession) -> None: + async def close_http_bridge_session( + target: proxy_service._HTTPBridgeSession, + *, + clear_continuity: bool = False, + ) -> None: nonlocal close_cancelled assert target is session + assert clear_continuity is False close_started.set() try: await release_close.wait() @@ -16162,6 +16548,7 @@ async def test_http_bridge_retire_after_drain_waits_for_queued_submission( instance_id="instance-retire-drain", owner_epoch=3, draining=False, + clear_continuity=False, ) release_account_lease.assert_awaited_once_with(lease) assert session.account_lease is None @@ -16227,6 +16614,7 @@ async def test_http_bridge_retire_after_drain_does_not_cancel_current_upstream_r instance_id="instance-reader-retire", owner_epoch=7, draining=False, + clear_continuity=False, ) release_account_lease.assert_awaited_once_with(lease) assert session.account_lease is None @@ -18985,6 +19373,124 @@ async def fail_first_session_before_output( assert all(call["preferred_account_has_continuity_provenance"] is True for call in creation_calls) +@pytest.mark.asyncio +async def test_stream_via_http_bridge_forks_model_transition_after_owner_conflict( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + payload = proxy_service.ResponsesRequest.model_validate( + { + "model": "gpt-5.6-terra", + "instructions": "hi", + "input": [{"role": "user", "content": "continue on the new model"}], + } + ) + durable_lookup = proxy_service.DurableBridgeLookup( + session_id="durable-model-conflict-parent", + canonical_kind="session_header", + canonical_key="shared-root", + api_key_scope="__anonymous__", + account_id="acc-model-owner", + owner_instance_id=None, + owner_epoch=1, + lease_expires_at=datetime.now(timezone.utc) + timedelta(seconds=60), + state=HttpBridgeSessionState.ACTIVE, + latest_turn_state="http_turn_model_parent", + latest_response_id="resp_model_parent", + model="gpt-5.6-sol", + ) + owner_conflict = ProxyResponseError( + 502, + openai_error( + "continuity_owner_conflict", + "Durable continuity aliases resolve to conflicting upstream owners.", + ), + ) + creation_keys: list[proxy_service._HTTPBridgeSessionKey] = [] + creation_calls: list[dict[str, Any]] = [] + + async def fake_get_or_create( + key: proxy_service._HTTPBridgeSessionKey, + **kwargs: Any, + ) -> proxy_service._HTTPBridgeSession: + creation_keys.append(key) + creation_calls.append(kwargs) + if len(creation_calls) == 1: + raise owner_conflict + session = _make_bridge_session(key=key) + session.account = cast( + Any, + SimpleNamespace(id="acc-model-alternate", status=AccountStatus.ACTIVE), + ) + session.request_model = payload.model + return session + + async def fake_stream_events( + _session: proxy_service._HTTPBridgeSession, + **_kwargs: Any, + ): + yield 'data: {"type":"response.completed"}\n\n' + + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: cast( + Any, + SimpleNamespace( + get=AsyncMock( + return_value=SimpleNamespace( + sticky_threads_enabled=False, + openai_cache_affinity_max_age_seconds=1800, + http_responses_session_bridge_prompt_cache_idle_ttl_seconds=3600, + http_responses_session_bridge_gateway_safe_mode=False, + ) + ) + ), + ), + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr(service._durable_bridge, "lookup_request_targets", AsyncMock(return_value=durable_lookup)) + monkeypatch.setattr(service, "_resolve_file_account_for_responses", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_get_or_create_http_bridge_session", fake_get_or_create) + monkeypatch.setattr(service, "_stream_http_bridge_session_events", fake_stream_events) + + chunks = [ + chunk + async for chunk in service._stream_via_http_bridge( + payload, + headers={ + "x-codex-turn-state": "http_turn_model_parent", + "x-codex-session-id": "shared-root", + }, + codex_session_affinity=True, + propagate_http_errors=True, + openai_cache_affinity=True, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=1800.0, + max_sessions=8, + queue_limit=4, + downstream_turn_state="http_turn_model_child", + ) + ] + + assert chunks == ['data: {"type":"response.completed"}\n\n'] + assert len(creation_calls) == 2 + assert creation_keys[0].affinity_kind in {"session_header", "turn_state_header"} + assert is_http_bridge_account_neutral_replay( + kind=creation_keys[1].affinity_kind, + key=creation_keys[1].affinity_key, + ) + assert creation_calls[0]["preferred_account_id"] == "acc-model-owner" + assert creation_calls[0]["preferred_account_has_continuity_provenance"] is True + assert creation_calls[1]["preferred_account_id"] is None + assert creation_calls[1]["preferred_account_has_continuity_provenance"] is False + assert creation_calls[1]["exclude_account_ids"] == {"acc-model-owner"} + assert creation_calls[1]["allow_forward_to_owner"] is False + + @pytest.mark.asyncio async def test_stream_via_http_bridge_preserves_verified_replay_kind_for_durable_model_transition( monkeypatch: pytest.MonkeyPatch, @@ -19344,6 +19850,16 @@ async def close(self) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) upstream = _TrackingUpstream() session = _make_bridge_session(key_value=f"eventless-{leading_telemetry}") + session.account = Account( + id="acc-bridge", + email="bridge@example.com", + plan_type="plus", + access_token_encrypted=b"", + refresh_token_encrypted=b"", + id_token_encrypted=b"", + last_refresh=datetime.now(timezone.utc), + status=AccountStatus.ACTIVE, + ) session.upstream = cast(UpstreamWebSocket, upstream) service._http_bridge_sessions[session.key] = session settings = _make_app_settings( @@ -19355,7 +19871,8 @@ async def close(self) -> None: monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) retry_precreated = AsyncMock(return_value=False) monkeypatch.setattr(service, "_retry_http_bridge_precreated_request", retry_precreated) - monkeypatch.setattr(service, "_handle_stream_error", AsyncMock()) + handle_stream_error = AsyncMock() + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) write_request_log = AsyncMock() monkeypatch.setattr(service, "_write_request_log", write_request_log) record_stuck_retire = Mock() @@ -19378,6 +19895,9 @@ async def close(self) -> None: owner.request_text = '{"type":"response.create","model":"gpt-5.6-sol","input":"hello"}' owner.preferred_account_id = "acc-bridge" owner.excluded_account_ids.add("acc-excluded") + if leading_telemetry: + owner.response_event_count = 1 + owner.upstream_model_output_seen = True sibling_queue: asyncio.Queue[str | None] = asyncio.Queue() sibling = proxy_service._WebSocketRequestState( request_id="req-created-sibling", @@ -19429,15 +19949,20 @@ async def close(self) -> None: assert owner.preferred_account_id == "acc-bridge" assert owner.excluded_account_ids == {"acc-excluded"} assert owner.replay_count == 0 - assert owner.response_event_count == 0 + assert owner.response_event_count == (1 if leading_telemetry else 0) if leading_telemetry: assert owner.latency_first_upstream_event_ms is not None - retry_precreated.assert_awaited_once_with(session) + retry_precreated.assert_awaited_once_with(session, allow_expired_deadline=True) 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() - assert fail_reader.await_args.kwargs["penalize_account"] is False + assert fail_reader.await_args.kwargs["penalize_account"] is True assert fail_reader.await_args.kwargs["force_retire"] is True + handle_stream_error.assert_awaited_once() + handle_stream_error_args = handle_stream_error.await_args + assert handle_stream_error_args is not None + assert handle_stream_error_args.args[0] is session.account + assert handle_stream_error_args.args[2] == "upstream_request_timeout" record_stuck_retire.assert_called_once_with( reason="missing_response_created_timeout", session=session, @@ -20148,6 +20673,7 @@ async def test_retire_stale_pending_http_bridge_session_unregisters_aliases_and_ instance_id="instance-cleanup", owner_epoch=7, draining=False, + clear_continuity=False, ) release_account_lease.assert_awaited_once_with(lease) assert session.account_lease is None @@ -20155,588 +20681,66 @@ async def test_retire_stale_pending_http_bridge_session_unregisters_aliases_and_ @pytest.mark.asyncio -async def test_http_bridge_retirement_does_not_record_midstream_retry_circuit_failure( +async def test_retire_missing_created_http_bridge_session_clears_durable_continuity( 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", + release_live_session = AsyncMock() + service._durable_bridge = cast( + Any, + SimpleNamespace(release_live_session=release_live_session), ) - - 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(), + session = _make_bridge_session(key_value="bridge-missing-created-cleanup") + session.durable_session_id = "durable-missing-created" + session.durable_owner_epoch = 9 + monkeypatch.setattr( + proxy_service, + "get_settings", + lambda: _make_app_settings(http_responses_session_bridge_instance_id="instance-missing-created"), ) - 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() + await service._retire_stale_pending_http_bridge_session(session, detail="missing_response_created_timeout") - -@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"}', + release_live_session.assert_awaited_once_with( + session_id="durable-missing-created", + instance_id="instance-missing-created", + owner_epoch=9, + draining=False, + clear_continuity=True, ) - 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( +async def test_retire_missing_created_http_bridge_session_clears_durable_continuity_after_close_attempt( monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) + release_live_session = AsyncMock() + service._durable_bridge = cast( + Any, + SimpleNamespace(release_live_session=release_live_session), + ) + session = _make_bridge_session(key_value="bridge-missing-created-after-close") + session.durable_session_id = "durable-missing-created-after-close" + session.durable_owner_epoch = 11 + session.upstream_close_attempted = True + close = cast(Any, session.upstream).close monkeypatch.setattr( - http_bridge_retry_circuit_module, - "_HTTP_BRIDGE_RETRY_CIRCUIT_FAILURE_THRESHOLD", - 1, + proxy_service, + "get_settings", + lambda: _make_app_settings(http_responses_session_bridge_instance_id="instance-missing-created-after-close"), ) - 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 + await service._retire_stale_pending_http_bridge_session(session, detail="missing_response_created_timeout") - 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 + close.assert_not_awaited() + release_live_session.assert_awaited_once_with( + session_id="durable-missing-created-after-close", + instance_id="instance-missing-created-after-close", + owner_epoch=11, + draining=False, + clear_continuity=True, + ) @pytest.mark.asyncio @@ -20988,7 +20992,6 @@ async def test_http_bridge_eventless_timeout_force_retires_with_admission_waiter session, error_code="upstream_request_timeout", error_message="missing response.created", - penalize_account=False, retire_detail="missing_response_created_timeout", force_retire=True, ) @@ -20998,7 +21001,7 @@ async def test_http_bridge_eventless_timeout_force_retires_with_admission_waiter 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 + assert fail_pending_await_args.kwargs["penalize_account"] is True @pytest.mark.asyncio @@ -21377,6 +21380,7 @@ async def fail_replay(target_session: proxy_service._HTTPBridgeSession) -> bool: instance_id="instance-log-fails", owner_epoch=3, draining=False, + clear_continuity=False, ) cast(Any, upstream).close.assert_awaited_once() @@ -21458,6 +21462,7 @@ async def test_http_bridge_reader_unexpected_processing_error_fails_pending_requ instance_id="instance-reader-crash", owner_epoch=9, draining=False, + clear_continuity=False, ) cast(Any, upstream).close.assert_awaited_once() write_request_log.assert_awaited_once() diff --git a/tests/unit/test_proxy_tool_call_dedupe.py b/tests/unit/test_proxy_tool_call_dedupe.py index e872dbdeb9..a8c78edd6a 100644 --- a/tests/unit/test_proxy_tool_call_dedupe.py +++ b/tests/unit/test_proxy_tool_call_dedupe.py @@ -20,7 +20,7 @@ def _loads_item_arguments(item: Mapping[str, JsonValue]) -> Any: return json.loads(arguments) -def test_mark_duplicate_tool_call_downstream_event_keeps_distinct_call_ids_with_same_arguments(): +def test_mark_duplicate_tool_call_downstream_event_suppresses_same_response_side_effect_with_new_call_id(): upstream_control = proxy_service._WebSocketUpstreamControl() first_payload: dict[str, JsonValue] = { "type": "response.output_item.done", @@ -67,7 +67,7 @@ def test_mark_duplicate_tool_call_downstream_event_keeps_distinct_call_ids_with_ seen_tool_call_keys=upstream_control.seen_tool_call_keys, response_id="resp_dupe", ) - is False + is True ) assert ( tool_call_dedupe.mark_duplicate_tool_call_downstream_event( @@ -120,6 +120,51 @@ def test_mark_duplicate_tool_call_downstream_event_suppresses_exec_command_with_ ) +def test_mark_duplicate_tool_call_downstream_event_suppresses_exec_command_same_response_new_call_id(): + upstream_control = proxy_service._WebSocketUpstreamControl() + command = ( + "psql -X -d kom -Atc \"select 'global', total_chunks, " + 'chunks_with_any_embedding from rag_embedding_global_stats;"' + ) + + def payload(call_id: str, *, max_output_tokens: int) -> dict[str, JsonValue]: + return { + "type": "response.output_item.done", + "response_id": "resp_live_duplicate", + "item": { + "type": "function_call", + "name": "exec_command", + "arguments": json.dumps( + { + "cmd": command, + "workdir": "/home/kom/Dropbox/Reemxy/tasks-loop", + "yield_time_ms": 10000, + "max_output_tokens": max_output_tokens, + }, + separators=(",", ":"), + ), + "call_id": call_id, + }, + } + + assert ( + tool_call_dedupe.mark_duplicate_tool_call_downstream_event( + payload("call_first", max_output_tokens=12000), + seen_tool_call_keys=upstream_control.seen_tool_call_keys, + response_id="resp_live_duplicate", + ) + is False + ) + assert ( + tool_call_dedupe.mark_duplicate_tool_call_downstream_event( + payload("call_second", max_output_tokens=48000), + seen_tool_call_keys=upstream_control.seen_tool_call_keys, + response_id="resp_live_duplicate", + ) + is True + ) + + def test_mark_duplicate_tool_call_downstream_event_suppresses_code_mode_exec_replay(): upstream_control = proxy_service._WebSocketUpstreamControl() first_payload: dict[str, JsonValue] = { diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 849c9fb3bb..224776ee70 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -2828,6 +2828,49 @@ async def test_opportunistic_admission_uses_api_key_enforced_model(): ) +@pytest.mark.asyncio +async def test_opportunistic_admission_preserves_usage_limit_denial(): + api_key = ApiKeyData( + id="key_opportunistic_usage_limit", + name="opportunistic usage limit", + key_prefix="sk-opportunistic", + allowed_models=None, + enforced_model=None, + enforced_reasoning_effort=None, + enforced_service_tier=None, + traffic_class=proxy_api.TRAFFIC_CLASS_OPPORTUNISTIC, + expires_at=None, + is_active=True, + created_at=utcnow(), + last_used_at=None, + ) + selection = AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + service = SimpleNamespace(check_opportunistic_admission=AsyncMock(return_value=selection)) + context = SimpleNamespace(service=service) + request = Request({"type": "http", "method": "GET", "path": "/v1/opportunistic/admission", "headers": []}) + + response = await proxy_api._opportunistic_admission_denial( + request, + cast(proxy_api.ProxyContext, context), + api_key, + model="gpt-5.1", + ) + + assert response is not None + assert response.status_code == 429 + body = json.loads(bytes(response.body)) + assert body["error"]["code"] == "usage_limit_reached" + assert body["error"]["type"] == "usage_limit_reached" + assert body["error"]["message"] == "Rate limit exceeded. Try again in 1h" + assert body["error"]["resets_at"] == 1_700_003_600 + assert "Retry-After" not in response.headers + + @pytest.mark.asyncio async def test_opportunistic_admission_scopes_single_account_to_selected_account(monkeypatch): settings = _make_proxy_settings() @@ -10365,6 +10408,11 @@ async def test_service_compact_passes_chatgpt_account_id_to_core(monkeypatch): monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + streaming_retry_module, + "_account_selection_recovery_sleep_seconds", + lambda _selection: pytest.fail("usage_limit_reached must not enter recovery wait"), + ) monkeypatch.setattr( service._load_balancer, "select_account", @@ -11607,6 +11655,46 @@ async def test_stream_responses_propagates_selection_error_code(monkeypatch): assert request_logs.calls[0]["error_code"] == "additional_quota_data_unavailable" +@pytest.mark.asyncio +async def test_stream_responses_preserves_usage_limit_reset_hint(monkeypatch): + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + service._load_balancer, + "select_account", + AsyncMock( + return_value=AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ), + ) + + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "stream": True, + } + ) + + chunks = [chunk async for chunk in service.stream_responses(payload, {"session_id": "sid-usage-limit"})] + + event = json.loads(chunks[0].split("data: ", 1)[1]) + assert event["response"]["error"]["code"] == "usage_limit_reached" + assert event["response"]["error"]["type"] == "usage_limit_reached" + assert event["response"]["error"]["resets_at"] == 1_700_003_600 + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert request_logs.calls[0]["error_code"] == "usage_limit_reached" + + @pytest.mark.asyncio async def test_stream_with_retry_keeps_sse_alive_while_account_capacity_recovers(monkeypatch): settings = _make_proxy_settings() @@ -14472,7 +14560,7 @@ async def fake_stream(*_, **__): @pytest.mark.asyncio -async def test_stream_responses_keeps_same_response_http_tool_calls_with_distinct_call_ids(monkeypatch): +async def test_stream_responses_suppresses_same_response_http_tool_calls_with_distinct_call_ids(monkeypatch): settings = _make_proxy_settings() request_logs = _RequestLogsRecorder() service = proxy_service.ProxyService(_repo_factory(request_logs)) @@ -14523,12 +14611,16 @@ async def fake_stream(*_, **__): if isinstance(chunk_payload, dict) and chunk_payload.get("type") == "response.output_item.done": tool_chunks.append(chunk_payload) - assert tool_chunks == [tool_payload, replayed_tool_payload] + assert tool_chunks == [tool_payload] terminal_payload = parse_sse_data_json(chunks[-1]) assert isinstance(terminal_payload, dict) - assert terminal_payload["type"] == "response.completed" + assert terminal_payload["type"] == "response.failed" + terminal_response = cast(dict[str, JsonValue], terminal_payload["response"]) + terminal_error = cast(dict[str, JsonValue], terminal_response["error"]) + assert terminal_error["code"] == "stream_incomplete" assert await service.drain_persistence_tasks(timeout_seconds=1) - assert request_logs.calls[0]["status"] == "success" + assert request_logs.calls[0]["status"] == "error" + assert request_logs.calls[0]["error_code"] == "stream_incomplete" @pytest.mark.asyncio @@ -17921,6 +18013,62 @@ async def test_select_websocket_connect_account_requires_preferred_account_for_p assert select_account.await_args.kwargs["request_stage"] == "reattach" +@pytest.mark.asyncio +async def test_select_websocket_connect_account_preserves_continuity_for_owner_usage_limit(monkeypatch): + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_prev_owner_usage_limit", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + request_stage="reattach", + ) + emit_connect_failure = AsyncMock() + select_account = AsyncMock( + return_value=AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ) + + monkeypatch.setattr(service, "_select_account_with_budget", select_account) + monkeypatch.setattr(service, "_emit_websocket_connect_failure", emit_connect_failure) + + result = await service._select_websocket_connect_account( + time.monotonic() + 10_000.0, + sticky_key=None, + sticky_kind=None, + prefer_earlier_reset=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + model="gpt-5.1", + request_state=request_state, + api_key=None, + client_send_lock=anyio.Lock(), + websocket=cast(WebSocket, SimpleNamespace()), + reallocate_sticky=False, + sticky_max_age_seconds=None, + exclude_account_ids=set(), + preferred_account_id="acc_owner", + require_preferred_account=True, + ) + + assert result is None + emit_connect_failure.assert_awaited_once() + call = emit_connect_failure.await_args + assert call is not None + assert call.kwargs["status_code"] == 502 + assert call.kwargs["error_code"] == "previous_response_owner_unavailable" + assert call.kwargs["account_id"] == "acc_owner" + assert call.kwargs["payload"]["error"]["code"] == "previous_response_owner_unavailable" + assert call.kwargs["payload"]["error"]["type"] == "server_error" + + @pytest.mark.asyncio async def test_select_websocket_connect_account_records_fail_closed_for_preferred_account_mismatch( monkeypatch, @@ -18199,6 +18347,65 @@ async def fake_sleep_for_account_selection_recovery(*_args: object, **kwargs: ob assert sent_payload["request_id"] == "ws_req_capacity_wait" +@pytest.mark.asyncio +async def test_select_websocket_connect_account_skips_capacity_wait_for_usage_limit(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_usage_limit_now", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + ) + websocket_send = AsyncMock() + + monkeypatch.setattr( + service, + "_select_account_with_budget", + AsyncMock( + return_value=AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ), + ) + monkeypatch.setattr( + websocket_mixin_module, + "_sleep_for_account_selection_recovery", + lambda *_args, **_kwargs: pytest.fail("usage_limit_reached must not enter recovery wait"), + ) + + result = await service._select_websocket_connect_account( + time.monotonic() + 10_000.0, + sticky_key=None, + sticky_kind=None, + prefer_earlier_reset=False, + routing_strategy="usage_weighted", + model="gpt-5.1", + request_state=request_state, + api_key=None, + client_send_lock=anyio.Lock(), + websocket=cast(WebSocket, SimpleNamespace(send_text=websocket_send)), + reallocate_sticky=False, + sticky_max_age_seconds=None, + exclude_account_ids=set(), + preferred_account_id=None, + require_preferred_account=False, + ) + + assert result is None + await_args = websocket_send.await_args + assert await_args is not None + sent_payload = json.loads(await_args.args[0]) + assert sent_payload["status"] == 429 + assert sent_payload["error"]["code"] == "usage_limit_reached" + assert sent_payload["error"]["type"] == "usage_limit_reached" + assert sent_payload["error"]["resets_at"] == 1_700_003_600 + + @pytest.mark.asyncio @pytest.mark.parametrize( ("preferred_account_id", "require_preferred_account", "file_required", "defer_no_account_error"), @@ -25363,6 +25570,87 @@ async def release_usage_reservation(self, reservation_id: str) -> None: assert released == ["resv_stream_failed_background"] +@pytest.mark.asyncio +async def test_stream_api_key_release_retries_bound_concurrent_repository_attempts(monkeypatch): + retry_concurrency = proxy_service._STREAM_API_KEY_RELEASE_RETRY_MAX_CONCURRENCY + task_count = retry_concurrency + 1 + active_repository_attempts = 0 + max_active_repository_attempts = 0 + repository_entries = 0 + retry_limit_reached = asyncio.Event() + allow_repository_attempts = asyncio.Event() + released: list[str] = [] + repo = SimpleNamespace(api_keys=object()) + + @asynccontextmanager + async def repo_factory() -> AsyncIterator[SimpleNamespace]: + nonlocal active_repository_attempts, max_active_repository_attempts, repository_entries + active_repository_attempts += 1 + repository_entries += 1 + max_active_repository_attempts = max( + max_active_repository_attempts, + active_repository_attempts, + ) + if active_repository_attempts == retry_concurrency: + retry_limit_reached.set() + try: + await allow_repository_attempts.wait() + yield repo + finally: + active_repository_attempts -= 1 + + class FakeApiKeysService: + def __init__(self, api_keys_repository: object) -> None: + assert api_keys_repository is repo.api_keys + + async def release_usage_reservation(self, reservation_id: str) -> None: + released.append(reservation_id) + + monkeypatch.setattr(proxy_service, "ApiKeysService", FakeApiKeysService) + + service = proxy_service.ProxyService(cast(proxy_service.ProxyRepoFactory, repo_factory)) + api_key = _make_api_key_data("key_stream_release_retry_bound") + reservations = [ + proxy_service.ApiKeyUsageReservationData( + reservation_id=f"resv_stream_release_retry_bound_{index}", + key_id=api_key.id, + model="gpt-5.5", + ) + for index in range(task_count) + ] + for index, reservation in enumerate(reservations): + service._schedule_cancel_safe_cleanup( + service._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=reservation, + request_id=f"req_stream_release_retry_bound_{index}", + retry_persistence_failures=True, + ), + action="release_stream_api_key_reservation_after_failed_settlement", + request_id=f"req_stream_release_retry_bound_{index}", + ) + + drain_task: asyncio.Task[bool] | None = None + try: + await asyncio.wait_for(retry_limit_reached.wait(), timeout=1) + await asyncio.sleep(0) + assert repository_entries == retry_concurrency + assert active_repository_attempts == retry_concurrency + assert len(service._background_cleanup_tasks) == task_count + drain_task = asyncio.create_task(service.drain_persistence_tasks(timeout_seconds=2)) + await asyncio.sleep(0) + assert not drain_task.done() + finally: + allow_repository_attempts.set() + if drain_task is None: + drain_task = asyncio.create_task(service.drain_persistence_tasks(timeout_seconds=2)) + assert await drain_task + + assert max_active_repository_attempts == retry_concurrency + assert sorted(released) == sorted(reservation.reservation_id for reservation in reservations) + assert service._background_cleanup_tasks == set() + + @pytest.mark.asyncio async def test_stream_with_retry_skips_release_after_settlement_transfers_on_cancel(monkeypatch): settings = _make_proxy_settings() @@ -27711,10 +27999,20 @@ async def fake_relay(*args, **kwargs): upstream.send_text.assert_awaited_once() +@pytest.mark.parametrize("release_read_fails", [False, True]) @pytest.mark.asyncio -async def test_stream_with_retry_releases_api_key_reservation_when_owner_lookup_fails(monkeypatch): +async def test_stream_with_retry_releases_api_key_reservation_when_owner_lookup_fails( + monkeypatch, + release_read_fails: bool, +): request_logs = _RequestLogsRecorder() - get_usage_reservation_mock = AsyncMock(return_value=SimpleNamespace(status="reserved", items=[])) + reservation_record = SimpleNamespace(status="reserved", items=[]) + get_usage_reservation_mock = AsyncMock( + side_effect=( + [RuntimeError("transient reservation read failure"), reservation_record] if release_read_fails else None + ), + return_value=reservation_record, + ) transition_usage_reservation_status_mock = AsyncMock(return_value=True) settle_usage_reservation_mock = AsyncMock() commit_mock = AsyncMock() @@ -27746,6 +28044,9 @@ async def __aexit__(self, exc_type, exc, tb) -> bool: return False service = proxy_service.ProxyService(lambda: _RepoContextWithApiKeys()) + # The synchronous stream-finally backstop must not queue behind detached + # retries, even when every retry slot is occupied. + service._stream_api_key_release_retry_semaphore = asyncio.Semaphore(0) settings = _make_proxy_settings() monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) @@ -27792,30 +28093,36 @@ async def __aexit__(self, exc_type, exc, tb) -> bool: monkeypatch.setattr(service, "_select_account_with_budget", select_account) with pytest.raises(proxy_module.ProxyResponseError) as exc_info: - async for _ in service._stream_with_retry( - payload, - {}, - codex_session_affinity=False, - propagate_http_errors=False, - openai_cache_affinity=False, - api_key=api_key, - api_key_reservation=reservation, - suppress_text_done_events=False, - request_transport="http", - ): - pass + async with asyncio.timeout(1): + async for _ in service._stream_with_retry( + payload, + {}, + codex_session_affinity=False, + propagate_http_errors=False, + openai_cache_affinity=False, + api_key=api_key, + api_key_reservation=reservation, + suppress_text_done_events=False, + request_transport="http", + ): + pass assert _proxy_error_code(exc_info.value) == "upstream_unavailable" owner_lookup.assert_awaited_once() select_account.assert_not_called() get_usage_reservation_mock.assert_awaited_once_with(reservation.reservation_id) - transition_usage_reservation_status_mock.assert_awaited_once_with( - reservation.reservation_id, - expected_status="reserved", - new_status="released", - ) - settle_usage_reservation_mock.assert_awaited_once() - commit_mock.assert_awaited_once() + if release_read_fails: + transition_usage_reservation_status_mock.assert_not_awaited() + settle_usage_reservation_mock.assert_not_awaited() + commit_mock.assert_not_awaited() + else: + transition_usage_reservation_status_mock.assert_awaited_once_with( + reservation.reservation_id, + expected_status="reserved", + new_status="released", + ) + settle_usage_reservation_mock.assert_awaited_once() + commit_mock.assert_awaited_once() @pytest.mark.asyncio @@ -28651,7 +28958,7 @@ async def test_process_upstream_websocket_text_masks_previous_response_not_found @pytest.mark.asyncio -async def test_process_upstream_websocket_text_keeps_same_response_distinct_tool_call_ids(monkeypatch): +async def test_process_upstream_websocket_text_suppresses_same_response_distinct_tool_call_ids(monkeypatch): request_logs = _RequestLogsRecorder() service = proxy_service.ProxyService(_repo_factory(request_logs)) finalize_request_state = AsyncMock() @@ -28713,8 +29020,8 @@ async def test_process_upstream_websocket_text_keeps_same_response_distinct_tool assert '"call_id":"call_first"' in first_text assert '"call_id":"call_replayed"' in replay_text - assert replay_control.suppress_downstream_event is False - assert pending_request.suppressed_duplicate_tool_call is False + assert replay_control.suppress_downstream_event is True + assert pending_request.suppressed_duplicate_tool_call is True finalize_request_state.assert_not_awaited() assert list(pending_requests) == [pending_request] @@ -30442,6 +30749,67 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, record_success.assert_not_awaited() +@pytest.mark.asyncio +async def test_stream_previsible_core_eof_with_previous_response_id_retries(monkeypatch): + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + account = _make_account("acc_previsible_core_eof") + request_logs.response_owner_by_id[("resp_parent", None, "sid-stream")] = account.id + handle_stream_error = AsyncMock() + record_success = AsyncMock() + stream_calls = 0 + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_MAX_TRANSIENT_SAME_ACCOUNT_RETRIES", 3) + monkeypatch.setattr(streaming_retry_module.ProcessNetworkRecovery, "wait", AsyncMock(return_value=None)) + monkeypatch.setattr(streaming_retry_module.asyncio, "sleep", AsyncMock()) + monkeypatch.setattr( + service._load_balancer, + "select_account", + AsyncMock(return_value=AccountSelection(account=account, error_message=None)), + ) + monkeypatch.setattr(service._load_balancer, "record_success", record_success) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=account)) + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(service, "_settle_stream_api_key_usage", AsyncMock(return_value=True)) + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False, **kwargs): + nonlocal stream_calls + del payload, headers, access_token, account_id, base_url, raise_for_status, kwargs + stream_calls += 1 + if stream_calls == 1: + return + yield 'data: {"type":"response.completed","response":{"id":"resp_child_retry_ok"}}\n\n' + + monkeypatch.setattr(proxy_service, "core_stream_responses", fake_stream) + + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "stream": True, + "previous_response_id": "resp_parent", + } + ) + + chunks = [chunk async for chunk in service.stream_responses(payload, {"session_id": "sid-stream"})] + + completed = json.loads(chunks[-1].split("data: ", 1)[1]) + assert completed["type"] == "response.completed" + assert completed["response"]["id"] == "resp_child_retry_ok" + assert stream_calls == 2 + assert request_logs.lookup_calls == [("resp_parent", None, "sid-stream")] + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert [call["status"] for call in request_logs.calls] == ["error", "success"] + assert request_logs.calls[0]["error_code"] == "stream_incomplete" + assert request_logs.calls[-1]["request_id"] == "resp_child_retry_ok" + handle_stream_error.assert_not_awaited() + record_success.assert_awaited_once_with(account) + + @pytest.mark.asyncio async def test_stream_missing_tool_output_proxy_error_is_masked_to_stream_incomplete(monkeypatch, caplog): settings = _make_proxy_settings() @@ -38005,7 +38373,7 @@ async def test_retry_http_bridge_precreated_request_does_not_send_after_admissio request_id="req_bridge_retry_deadline", ) acquire_admission = AsyncMock(return_value=replacement_admission) - retry_times = iter((9.0, 11.0)) + retry_times = iter((9.0, 9.5, 11.0)) request_state = proxy_service._WebSocketRequestState( request_id="req_bridge_retry_deadline", model="gpt-5.6-sol", diff --git a/tests/unit/test_replay_safety.py b/tests/unit/test_replay_safety.py index a281ad1835..d7ecaeda84 100644 --- a/tests/unit/test_replay_safety.py +++ b/tests/unit/test_replay_safety.py @@ -676,6 +676,21 @@ def test_full_resend_suffix_rejects_missing_or_misordered_context( False, id="omitted-parallel-call", ), + pytest.param( + [ + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": "{}", + }, + {"role": "developer", "content": "new control message"}, + {"type": "function_call_output", "call_id": "call_1", "output": "result"}, + ], + {"call_1": "function_call"}, + False, + id="inline-developer-message-in-fresh-suffix", + ), ], ) def test_full_resend_suffix_accepts_only_self_contained_tool_loops( @@ -700,6 +715,776 @@ def test_full_resend_suffix_accepts_only_self_contained_tool_loops( ) +def test_full_resend_tool_loop_manifest_tolerates_fresh_developer_interleave_after_historical_one() -> None: + stored_input: list[JsonValue] = [ + {"role": "user", "content": "first question"}, + { + "type": "custom_tool_call", + "id": "ctc_old", + "call_id": "call_old", + "name": "shell", + "input": "pwd", + "status": "completed", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_old"}, + }, + { + "type": "message", + "id": "msg_old_control", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_old"}, + "content": [{"type": "input_text", "text": "historical control message"}], + }, + { + "type": "custom_tool_call_output", + "id": "ctco_old", + "call_id": "call_old", + "output": "/workspace", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_old"}, + }, + ] + suffix: list[JsonValue] = [ + { + "type": "custom_tool_call", + "id": "ctc_current", + "call_id": "call_current", + "name": "shell", + "input": "pwd", + "status": "completed", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + }, + { + "type": "message", + "id": "msg_control", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + "content": [{"type": "input_text", "text": "new control message"}], + }, + { + "type": "custom_tool_call_output", + "id": "ctco_current", + "call_id": "call_current", + "output": "/workspace", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + }, + ] + + projection = project_responses_input_for_account_neutral_fresh_replay( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + assert projection is not None + assert ( + responses_input_suffix_matches_pending_tool_calls( + projection.input_items, + stored_count=projection.stored_prefix_count, + pending_tool_calls={"call_current": "custom_tool_call"}, + ) + is True + ) + + +def test_full_resend_tool_loop_manifest_rejects_fresh_developer_inside_function_call_pair() -> None: + stored_input: list[JsonValue] = [{"role": "user", "content": "first question"}] + suffix: list[JsonValue] = [ + { + "type": "function_call", + "call_id": "call_current", + "name": "lookup", + "arguments": "{}", + }, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "unobserved function control"}], + }, + { + "type": "function_call_output", + "call_id": "call_current", + "output": "ok", + }, + ] + projection = project_responses_input_for_account_neutral_fresh_replay( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + + assert projection is not None + assert ( + responses_input_suffix_matches_pending_tool_calls( + projection.input_items, + stored_count=projection.stored_prefix_count, + pending_tool_calls={"call_current": "function_call"}, + ) + is False + ) + + +def test_full_resend_tool_loop_manifest_rejects_adjacent_pair_nested_in_parallel_batch() -> None: + stored_input: list[JsonValue] = [{"role": "user", "content": "first question"}] + suffix: list[JsonValue] = [ + {"type": "custom_tool_call", "call_id": "call_1", "name": "shell", "input": "pwd"}, + {"type": "custom_tool_call", "call_id": "call_2", "name": "shell", "input": "whoami"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control inside one matching parallel pair"}], + }, + {"type": "custom_tool_call_output", "call_id": "call_2", "output": "worker"}, + {"type": "custom_tool_call_output", "call_id": "call_1", "output": "/workspace"}, + ] + projection = project_responses_input_for_account_neutral_fresh_replay( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + + assert projection is not None + assert ( + responses_input_suffix_matches_pending_tool_calls( + projection.input_items, + stored_count=projection.stored_prefix_count, + pending_tool_calls={"call_1": "custom_tool_call", "call_2": "custom_tool_call"}, + ) + is False + ) + + +@pytest.mark.parametrize( + "suffix", + [ + pytest.param( + [ + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control before call"}], + }, + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + ], + id="developer-before-call", + ), + pytest.param( + [ + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control after output"}], + }, + ], + id="developer-after-output", + ), + pytest.param( + [ + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "first control"}], + }, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "second control"}], + }, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + ], + id="multiple-developer-messages", + ), + pytest.param( + [ + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + { + "role": "developer", + "content": [{"type": "input_text", "text": "control without turn id"}], + }, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + ], + id="developer-without-turn-id", + ), + pytest.param( + [ + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + { + "type": None, + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control with null type"}], + }, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + ], + id="developer-with-null-type", + ), + pytest.param( + [ + { + "type": "message", + "role": "assistant", + "phase": "commentary", + "content": [{"type": "output_text", "text": "working"}], + }, + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control"}], + }, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + ], + id="developer-with-leading-assistant-exception", + ), + pytest.param( + [ + {"type": "custom_tool_call", "call_id": "call_current", "name": "shell", "input": "pwd"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control"}], + }, + {"type": "custom_tool_call_output", "call_id": "call_current", "output": "/workspace"}, + {"role": "user", "content": "follow up"}, + ], + id="developer-with-trailing-user-exception", + ), + ], +) +def test_full_resend_tool_loop_manifest_rejects_unproven_fresh_developer_positions( + suffix: list[JsonValue], +) -> None: + stored_input: list[JsonValue] = [{"role": "user", "content": "first question"}] + projection = project_responses_input_for_account_neutral_fresh_replay( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + + assert projection is not None + assert ( + responses_input_suffix_matches_pending_tool_calls( + projection.input_items, + stored_count=projection.stored_prefix_count, + pending_tool_calls={"call_current": "custom_tool_call"}, + ) + is False + ) + + +@pytest.mark.parametrize( + ("interleaved_item", "expected"), + [ + pytest.param( + {"role": "developer", "content": "historical control"}, + True, + id="implicit-developer-message", + ), + pytest.param( + { + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "historical control"}], + }, + True, + id="explicit-developer-message", + ), + pytest.param( + { + "role": "developer", + "status": "completed", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_historical"}, + "content": "historical control", + }, + True, + id="completed-developer-message-with-neutral-metadata", + ), + pytest.param( + {"role": "developer", "id": "msg_owned", "content": "historical control"}, + False, + id="response-owned-developer-message", + ), + pytest.param( + {"role": "developer", "id": "", "content": "historical control"}, + False, + id="empty-response-owned-developer-message", + ), + pytest.param( + {"role": "developer", "content": " "}, + False, + id="malformed-developer-message", + ), + pytest.param( + {"role": "developer", "phase": "commentary", "content": "historical control"}, + False, + id="phased-developer-message", + ), + pytest.param( + {"role": "developer", "status": "failed", "content": "historical control"}, + False, + id="failed-developer-message", + ), + pytest.param( + {"role": "developer", "content": "historical control", "account_id": "account-scoped"}, + False, + id="unknown-developer-field", + ), + pytest.param( + { + "role": "developer", + "content": "historical control", + "internal_chat_message_metadata_passthrough": {"turn_id": ""}, + }, + False, + id="invalid-developer-metadata", + ), + pytest.param( + {"role": "user", "content": "new user input"}, + False, + id="user-message", + ), + pytest.param( + { + "role": "assistant", + "content": [{"type": "output_text", "text": "assistant output"}], + }, + False, + id="assistant-message", + ), + pytest.param( + {"role": "system", "content": "system control"}, + False, + id="system-message", + ), + ], +) +def test_full_resend_exact_manifest_only_allows_historical_developer_interleaving( + interleaved_item: JsonValue, + expected: bool, +) -> None: + stored_input: list[JsonValue] = [ + {"role": "user", "content": "first question"}, + { + "type": "custom_tool_call", + "call_id": "call_old", + "name": "shell", + "input": "pwd", + }, + interleaved_item, + { + "type": "custom_tool_call_output", + "call_id": "call_old", + "output": "/workspace", + }, + ] + suffix: list[JsonValue] = [ + { + "type": "custom_tool_call", + "call_id": "call_current", + "name": "shell", + "input": "git status --short", + }, + { + "type": "custom_tool_call_output", + "call_id": "call_current", + "output": "", + }, + ] + + input_items = [*stored_input, *suffix] + projection = project_responses_input_for_account_neutral_fresh_replay( + input_items, + stored_count=len(stored_input), + preserve_developer_message_ids=True, + ) + + assert projection is not None + assert ( + responses_input_suffix_matches_pending_tool_calls( + projection.input_items, + stored_count=projection.stored_prefix_count, + pending_tool_calls={"call_current": "custom_tool_call"}, + ) + is expected + ) + + +@pytest.mark.parametrize( + "historical_output", + [ + pytest.param(None, id="missing-output"), + pytest.param( + { + "type": "custom_tool_call_output", + "call_id": "call_other", + "output": "/workspace", + }, + id="mismatched-output", + ), + ], +) +def test_full_resend_exact_manifest_requires_historical_interleaved_call_output( + historical_output: JsonValue | None, +) -> None: + stored_input: list[JsonValue] = [ + {"role": "user", "content": "first question"}, + { + "type": "custom_tool_call", + "call_id": "call_old", + "name": "shell", + "input": "pwd", + }, + {"role": "developer", "content": "historical control"}, + ] + if historical_output is not None: + stored_input.append(historical_output) + suffix: list[JsonValue] = [ + { + "type": "custom_tool_call", + "call_id": "call_current", + "name": "shell", + "input": "git status --short", + }, + { + "type": "custom_tool_call_output", + "call_id": "call_current", + "output": "", + }, + ] + + assert not responses_input_suffix_matches_pending_tool_calls( + [*stored_input, *suffix], + stored_count=len(stored_input), + pending_tool_calls={"call_current": "custom_tool_call"}, + ) + + +def test_full_resend_retained_output_tolerates_fresh_developer_after_user() -> None: + stored_input: list[JsonValue] = [{"role": "user", "content": "first question"}] + suffix: list[JsonValue] = [ + { + "type": "message", + "id": "msg_answer", + "role": "assistant", + "phase": "final_answer", + "status": "completed", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_previous"}, + "content": [{"type": "output_text", "text": "prior answer"}], + }, + { + "type": "message", + "id": "msg_user", + "role": "user", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + "content": [{"type": "input_text", "text": "next question"}], + }, + { + "type": "message", + "id": "msg_control", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_current"}, + "content": [{"type": "input_text", "text": "new control message"}], + }, + ] + + projection = project_responses_input_for_account_neutral_fresh_replay( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + assert projection is not None + assert ( + responses_input_suffix_retains_prior_output( + projection.input_items, + stored_count=projection.stored_prefix_count, + ) + is True + ) + + +@pytest.mark.parametrize( + "suffix", + [ + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control without user"}], + }, + ], + id="developer-without-user", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control"}], + }, + {"role": "user", "content": "later question"}, + ], + id="developer-not-terminal", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + { + "role": "developer", + "content": [{"type": "input_text", "text": "control without turn id"}], + }, + ], + id="developer-without-turn-id", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": { + "turn_id": "turn-current", + "account_id": "acc-1", + }, + "content": [{"type": "input_text", "text": "account-bound control"}], + }, + ], + id="developer-with-account-metadata", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + { + "type": 7, + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control with malformed type"}], + }, + ], + id="developer-with-nonstring-type", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + { + "type": None, + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control with null type"}], + }, + ], + id="developer-with-null-type", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"type": "input_text", "text": "raw next input"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control after raw input"}], + }, + ], + id="developer-after-raw-input-part", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"type": "input_text", "text": "raw next input"}, + {"role": "user", "content": "next question"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control after raw and user"}], + }, + ], + id="developer-after-raw-input-and-user", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "first next question"}, + {"role": "user", "content": "second next question"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control after repeated users"}], + }, + ], + id="developer-after-repeated-user-messages", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "final_answer", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_file", "file_id": "file-account-bound"}], + }, + ], + id="developer-with-account-bound-file-id", + ), + pytest.param( + [ + { + "role": "assistant", + "phase": "commentary", + "content": [{"type": "output_text", "text": "still working"}], + }, + {"role": "user", "content": "next question"}, + { + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn-current"}, + "content": [{"type": "input_text", "text": "control after commentary"}], + }, + ], + id="developer-after-assistant-commentary", + ), + ], +) +def test_full_resend_retained_output_rejects_unproven_fresh_developer_followup( + suffix: list[JsonValue], +) -> None: + stored_input: list[JsonValue] = [{"role": "user", "content": "first question"}] + + assert ( + responses_input_suffix_retains_prior_output( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + is False + ) + + +def test_full_resend_retained_output_rejects_complete_tool_pair_between_output_and_user() -> None: + stored_prefix: list[JsonValue] = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "old"}], + } + ] + full_input: list[JsonValue] = [ + *stored_prefix, + { + "type": "message", + "role": "assistant", + "phase": "final_answer", + "status": "completed", + "content": [{"type": "output_text", "text": "done"}], + }, + { + "type": "custom_tool_call", + "call_id": "call_extra", + "name": "shell", + "input": "pwd", + }, + { + "type": "custom_tool_call_output", + "call_id": "call_extra", + "output": "ok", + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "new"}], + }, + { + "type": "message", + "role": "developer", + "internal_chat_message_metadata_passthrough": {"turn_id": "turn_new"}, + "content": [{"type": "input_text", "text": "control"}], + }, + ] + projected_full_input = project_responses_input_for_account_neutral_fresh_replay( + full_input, + stored_count=len(stored_prefix), + ) + + assert projected_full_input is not None + assert ( + responses_input_suffix_retains_prior_output( + projected_full_input.input_items, + stored_count=projected_full_input.stored_prefix_count, + ) + is False + ) + + +def test_full_resend_retained_output_rejects_historical_developer_interleaving() -> None: + stored_input: list[JsonValue] = [ + {"role": "user", "content": "first question"}, + { + "type": "custom_tool_call", + "call_id": "call_old", + "name": "shell", + "input": "pwd", + }, + {"role": "developer", "content": "historical control"}, + { + "type": "custom_tool_call_output", + "call_id": "call_old", + "output": "/workspace", + }, + ] + suffix: list[JsonValue] = [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "prior answer"}], + }, + {"role": "user", "content": "next question"}, + ] + + assert not responses_input_suffix_retains_prior_output( + [*stored_input, *suffix], + stored_count=len(stored_input), + ) + + def test_full_resend_tool_loop_manifest_rejects_call_id_reused_from_stored_prefix() -> None: stored_input: list[JsonValue] = [ { diff --git a/tests/unit/test_selection_errors.py b/tests/unit/test_selection_errors.py new file mode 100644 index 0000000000..8a670dad33 --- /dev/null +++ b/tests/unit/test_selection_errors.py @@ -0,0 +1,72 @@ +import pytest + +from app.core.resilience.overload import LOCAL_OVERLOAD_CODES +from app.modules.proxy.load_balancer import AccountSelection +from app.modules.proxy.selection_errors import selection_failure_response + + +def test_pool_usage_exhaustion_is_codex_compatible_429(): + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="Usage limit reached", + error_code="usage_limit_reached", + ) + ) + + assert status == 429 + assert payload == { + "error": { + "message": "Usage limit reached", + "type": "usage_limit_reached", + "code": "usage_limit_reached", + } + } + + +def test_unusable_pool_remains_no_accounts_503(): + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="All accounts require re-authentication", + error_code=None, + ) + ) + + assert status == 503 + assert payload["error"]["type"] == "server_error" + assert payload["error"]["code"] == "no_accounts" + + +def test_pool_usage_exhaustion_preserves_authoritative_reset(): + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 300s", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ) + + assert status == 429 + assert payload["error"]["resets_at"] == 1_700_003_600 + + +@pytest.mark.parametrize("local_code", sorted(LOCAL_OVERLOAD_CODES)) +def test_local_overload_codes_keep_rate_limit_contract(local_code: str): + # Covers every canonical local capacity code, including codes added later + # (e.g. api_key_stream_fair_share): local overload must stay a 429 + # rate_limit_error and never be reclassified as upstream usage exhaustion + # or a 503. + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="Local capacity is exhausted", + error_code=local_code, + ) + ) + + assert status == 429 + assert payload["error"]["type"] == "rate_limit_error" + assert payload["error"]["code"] == local_code + assert "resets_at" not in payload["error"]