fix(proxy): recover goal restarts from unavailable owners

app/core/openai/requests.py

@@ -1452,6 +1452,36 @@ def _compact_item_texts(item: Mapping[str, JsonValue]) -> list[str]:
     return texts
 
 
+def responses_input_contains_goal_continuation_context(input_value: JsonValue) -> bool:
+    """Return whether Responses input carries Codex's goal-continuation marker."""
+
+    if not is_json_list(input_value):
+        return False
+    for item in input_value:
+        if not is_json_mapping(item):
+            continue
+        for text in _compact_item_texts(item):
+            if text.lstrip().startswith(_GOAL_CONTINUATION_CONTEXT_PREFIX):
+                return True
+    return False
+
+
+def responses_request_contains_goal_continuation_context(payload: ResponsesRequest) -> bool:
+    """Return whether a normalized request carries Codex's goal restart marker."""
+
+    # ResponsesRequest normalization lifts developer/system input messages into
+    # ``instructions``. The marker can therefore disappear from ``input`` and
+    # follow pre-existing instruction text by the time affinity is classified.
+    # Keep both locations in this check or a harmless parser refactor can
+    # silently break restart recovery while marker-preservation tests still pass.
+    instructions = payload.instructions
+    if isinstance(instructions, str) and any(
+        line.lstrip().startswith(_GOAL_CONTINUATION_CONTEXT_PREFIX) for line in instructions.splitlines()
+    ):
+        return True
+    return responses_input_contains_goal_continuation_context(payload.input)
+
+
 def _compact_trimmed_input_with_markers(
     input_value: list[JsonValue], token_counts: list[int], selected_indices: set[int]
 ) -> list[JsonValue]:

app/db/alembic/versions/20260811_000000_add_sticky_abandonment_scope.py

@@ -0,0 +1,43 @@
+"""add source scope to sticky continuity abandonment
+
+Revision ID: 20260811_000000_add_sticky_abandonment_scope
+Revises: 20260806_120000_add_http_bridge_owner_process_epoch
+Create Date: 2026-08-11 00:00:00.000000
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.engine import Connection
+
+revision = "20260811_000000_add_sticky_abandonment_scope"
+down_revision = "20260806_120000_add_http_bridge_owner_process_epoch"
+branch_labels = None
+depends_on = None
+
+_TABLE = "sticky_sessions"
+_COLUMN = "continuity_abandonment_scope"
+
+
+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()
+    if _COLUMN in _columns(bind):
+        return
+    with op.batch_alter_table(_TABLE) as batch_op:
+        batch_op.add_column(sa.Column(_COLUMN, sa.String(length=32), nullable=True))
+
+
+def downgrade() -> None:
+    bind = op.get_bind()
+    if _COLUMN not in _columns(bind):
+        return
+    with op.batch_alter_table(_TABLE) as batch_op:
+        batch_op.drop_column(_COLUMN)

app/db/models.py

@@ -755,17 +755,15 @@ class StickySession(Base):
         onupdate=func.now(),
         nullable=False,
     )
-    # Set only by purge_stale_hard_codex_session_mappings's first pass. A hard
-    # codex_session row normally proves ownership for `conversation`-continuity
-    # requests (see affinity.py's require_unambiguous_account), which have no
-    # other owner index. Once the durably-unavailable owner's proof is this
-    # stale, we stop treating the row as a live pin (so a fresh account can be
-    # selected) but keep it around with this marker set instead of deleting it
-    # outright, so selection can tell "this key was deliberately abandoned,
-    # picking a new owner is authorized" apart from "this key was never seen,
-    # ambiguity must fail closed." The row is only ever hard-deleted once it
-    # has sat abandoned past a further grace window with nobody claiming it.
+    # A non-null timestamp marks continuity abandoned either globally (scope
+    # is NULL, as written by stale-hard cleanup) or only for one typed source.
+    # Source-scoped abandonment is what lets a process-session restart stop
+    # consulting an ambiguous raw compatibility row without erasing that
+    # row's retained account_id for an explicit turn-state lookup that happens
+    # to use the same client-controlled text. The row is only hard-deleted by
+    # global stale-hard cleanup after its additional grace window.
     continuity_abandoned_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
+    continuity_abandonment_scope: Mapped[str | None] = mapped_column(String(32), nullable=True, default=None)
 
 
 class CapabilityLineageMarker(Base):

app/modules/proxy/_load_balancer/sticky_selection.py

@@ -77,6 +77,9 @@ class SelectionInputsProtocol(Protocol):
     @property
     def effective_continuity_owner_candidates(self) -> list[Account]: ...
 
+    @property
+    def effective_sticky_mutation_authority_account_ids(self) -> frozenset[str]: ...
+
 
 SelectionInputsT = TypeVar("SelectionInputsT", bound=SelectionInputsProtocol)
 
@@ -204,6 +207,7 @@ class StickySelectionRequest(Generic[SelectionInputsT]):
     legacy_sticky_key: str | None
     legacy_existing_account_id: str | None
     spill_bare_session_on_account_cap: bool
+    abandon_unavailable_legacy_owner: bool
     require_unambiguous_account: bool
     sticky_max_age_seconds: int | None
     prefer_earlier_reset_accounts: bool
@@ -266,6 +270,7 @@ async def run_sticky_selection_path(
     legacy_sticky_key = request.legacy_sticky_key
     legacy_existing_account_id = request.legacy_existing_account_id
     spill_bare_session_on_account_cap = request.spill_bare_session_on_account_cap
+    abandon_unavailable_legacy_owner = request.abandon_unavailable_legacy_owner
     require_unambiguous_account = request.require_unambiguous_account
     sticky_max_age_seconds = request.sticky_max_age_seconds
     prefer_earlier_reset_accounts = request.prefer_earlier_reset_accounts
@@ -313,6 +318,7 @@ def _direct_error(
 
     sticky_existing_account_id: str | None | object = _STICKY_EXISTING_UNSET
     sticky_continuity_abandoned = False
+    retired_legacy_owner_account_ids: set[str] = set()
     attempt = 0
     suppress_recovery_probe_candidates = False
     while True:
@@ -326,6 +332,7 @@ def _direct_error(
                     sticky_key,
                     kind=sticky_kind,
                     max_age_seconds=sticky_max_age_seconds,
+                    continuity_source=sticky_source,
                 )
                 sticky_existing_account_id = sticky_owner_lookup.account_id
                 # `is True` (not a truthy check): an un-configured test double
@@ -348,6 +355,17 @@ def _direct_error(
                 required_account_id=required_account_id,
                 redact_sensitive_details=redact_sensitive_details,
             )
+            if retired_legacy_owner_account_ids:
+                # Retirement is authoritative even when this selector loaded a
+                # pre-retirement account snapshot (or another replica still has
+                # one cached). Never let that stale snapshot immediately repin
+                # the account this request just proved durably unavailable.
+                states = [state for state in states if state.account_id not in retired_legacy_owner_account_ids]
+                account_map = {
+                    account_id: account
+                    for account_id, account in account_map.items()
+                    if account_id not in retired_legacy_owner_account_ids
+                }
             effective_routing_costs = (
                 routing_costs_by_account_id
                 if routing_costs_by_account_id is not None
@@ -462,6 +480,63 @@ def _direct_error(
                     traffic_class=traffic_class,
                 )
             probe_reservation: ProbeReservation | None = None
+        # Raw sticky rows are global, while account-assigned API keys and
+        # other authenticated policies narrow a request's mutation authority.
+        # Keep this check on the pre-health continuity pool: quota exhaustion
+        # may authorize retirement, but being outside policy scope never does.
+        legacy_owner_in_effective_policy_scope = (
+            isinstance(sticky_existing_account_id, str)
+            and sticky_existing_account_id in selection_inputs.effective_sticky_mutation_authority_account_ids
+        )
+        if (
+            abandon_unavailable_legacy_owner
+            and hard_sticky
+            and sticky_existing_is_legacy
+            and sticky_source == "session_header"
+            and legacy_sticky_key is not None
+            and isinstance(sticky_existing_account_id, str)
+            and legacy_owner_in_effective_policy_scope
+        ):
+            async with owner._repo_factory() as repos:
+                owner_retired = await repos.sticky_sessions.abandon_legacy_session_header_owner_if_unavailable(
+                    legacy_sticky_key,
+                    kind=StickySessionKind.CODEX_SESSION,
+                    expected_account_id=sticky_existing_account_id,
+                )
+                authoritative_legacy_owner = None
+                if not owner_retired:
+                    authoritative_legacy_owner = await repos.sticky_sessions.get_account_id_and_abandonment(
+                        legacy_sticky_key,
+                        kind=StickySessionKind.CODEX_SESSION,
+                        continuity_source="session_header",
+                    )
+            # One guarded write is authoritative for this selection. Repeating
+            # it in capacity-wait retries would add write pressure and could
+            # reinterpret a later status transition as restart authorization.
+            abandon_unavailable_legacy_owner = False
+            if owner_retired:
+                # The raw compatibility row is now a tombstone. Drop only the
+                # selection loop's cached legacy owner and run the normal path
+                # again so namespaced affinity, leases, and admission checks
+                # are established through the existing selection path.
+                logger.info(
+                    "Legacy Codex session-header owner abandoned for self-contained goal restart account_id=%s",
+                    "<redacted>" if redact_sensitive_details else sticky_existing_account_id,
+                )
+                retired_legacy_owner_account_ids.add(sticky_existing_account_id)
+                legacy_existing_account_id = None
+                continue
+            # A failed compare-and-set means the cached owner is no longer
+            # authoritative: it may have recovered, another request may have
+            # rebound the raw row, or another worker may already have
+            # tombstoned it. Re-read under a fresh transaction and restart the
+            # loop so each outcome is handled by normal selection. Retaining
+            # the stale owner here would defeat the CAS and can fail a restart
+            # even though a concurrent operation already established a valid
+            # replacement.
+            assert authoritative_legacy_owner is not None
+            legacy_existing_account_id = authoritative_legacy_owner.account_id
+            continue
         sticky_outcome = _StickySelectionOutcome(selection=SelectionResult(None, None))
         if fair_share_denial is not None:
             # Denial parks in the transport capacity-wait loop like a cap

app/modules/proxy/_service/websocket/mixin.py

@@ -431,6 +431,7 @@
     _is_synthesized_turn_state,
     _owner_lookup_session_id_from_headers,
     _prompt_cache_key_from_request_model,
+    _request_allows_unavailable_legacy_owner_abandonment,
     _sticky_key_for_responses_request,
     _sticky_key_from_session_header,  # noqa: F401
     _sticky_key_from_turn_state_header,
@@ -532,23 +533,24 @@ async def _reject_websocket_owner_switch_blocked(
     api_key: ApiKeyData | None,
     response_create_gate: asyncio.Semaphore,
     downstream_activity: _DownstreamWebSocketActivity,
-) -> None:
-    error_message = (
+    error_code: str = "previous_response_owner_unavailable",
+    error_message: str = (
         "Previous response owner differs while another response is still streaming; retry after the terminal frame."
-    )
+    ),
+) -> None:
     await proxy._release_websocket_request_state_reservation(request_state)
     await proxy._write_websocket_connect_failure(
         account_id=account.id,
         api_key=api_key,
         request_state=request_state,
-        error_code="previous_response_owner_unavailable",
+        error_code=error_code,
         error_message=error_message,
     )
     await proxy._emit_websocket_terminal_error(
         websocket,
         client_send_lock=client_send_lock,
         request_state=request_state,
-        error_code="previous_response_owner_unavailable",
+        error_code=error_code,
         error_message=error_message,
         downstream_activity=downstream_activity,
     )
@@ -1238,7 +1240,14 @@ async def proxy_responses_websocket(
         account_lease: AccountLease | None = None
         upstream_requires_security_work_authorized: bool | None = None
         upstream_turn_state: str | None = _sticky_key_from_turn_state_header(headers)
-        client_turn_state_header: str | None = _sticky_key_from_turn_state_header(filtered_headers)
+        # The API inserts its generated downstream turn state into ``headers``
+        # before entering this service. Preserve a turn-state header as
+        # client-owned only when no synthesized value accompanied it; otherwise
+        # account-switch cleanup must remain able to remove the old account's
+        # generated token from ``filtered_headers``.
+        client_turn_state_header: str | None = (
+            _sticky_key_from_turn_state_header(filtered_headers) if synthesized_turn_state is None else None
+        )
         upstream_account_id: str | None = None
         downstream_activity = _DownstreamWebSocketActivity()
         replay_request_state: _WebSocketRequestState | None = None
@@ -1792,6 +1801,53 @@ def take_reader_replay_request_state() -> _WebSocketRequestState | None:
                     payload = None
                     continue
 
+                if (
+                    request_state is not None
+                    and upstream is not None
+                    and account is not None
+                    and request_state.affinity_policy.abandon_unavailable_legacy_owner
+                ):
+                    # Reusing the existing socket would bypass sticky
+                    # selection, so the unavailable raw owner would never be
+                    # compared, tombstoned, or replaced. A restart is movable
+                    # only before dispatch and cannot retire a socket that
+                    # still owns another response.
+                    async with pending_lock:
+                        restart_switch_blocked = _websocket_owner_switch_has_other_pending_requests(
+                            request_state,
+                            pending_requests,
+                        )
+                    if restart_switch_blocked:
+                        await _reject_websocket_owner_switch_blocked(
+                            proxy,
+                            websocket,
+                            client_send_lock=client_send_lock,
+                            request_state=request_state,
+                            account=account,
+                            api_key=api_key,
+                            response_create_gate=response_create_gate,
+                            downstream_activity=downstream_activity,
+                            error_code="stream_incomplete",
+                            error_message=(
+                                "Goal restart cannot switch accounts while another response is still streaming; "
+                                "retry after the terminal frame."
+                            ),
+                        )
+                        request_state = None
+                        text_data = None
+                        payload = None
+                        continue
+                    await retire_current_upstream()
+                    upstream_turn_state = None
+                    if client_turn_state_header is None:
+                        # Provenance was captured before this switch: absence
+                        # here means the API synthesized the forwarded token.
+                        # Such account-local state must die with its upstream;
+                        # an actual client anchor remains fail-closed instead.
+                        filtered_headers = {
+                            key: value for key, value in filtered_headers.items() if key.lower() != "x-codex-turn-state"
+                        }
+
                 if (
                     request_state is not None
                     and upstream is not None
@@ -2724,11 +2780,21 @@ async def _prepare_websocket_response_create_request(
         original_full_resend_payload: ResponsesRequest | None = None
         original_input_item_count: int | None = None
         original_input_fingerprint: str | None = None
-        session_anchor = _websocket_continuity_anchor_for_payload(
-            continuity_state,
-            responses_payload=responses_payload,
-            codex_session_affinity=codex_session_affinity,
-        )
+        # Classify restart authority from the complete normalized client body,
+        # before ordinary direct-WebSocket continuity injects a
+        # ``previous_response_id`` and trims historical input. That injected
+        # anchor is account-owned and would both erase the restart capability
+        # and make the payload unsafe for the replacement account. A proven
+        # goal restart must retain the complete resend through selection.
+        goal_restart_full_resend = _request_allows_unavailable_legacy_owner_abandonment(responses_payload)
+        restart_affinity_payload = responses_payload
+        session_anchor = None
+        if not goal_restart_full_resend:
+            session_anchor = _websocket_continuity_anchor_for_payload(
+                continuity_state,
+                responses_payload=responses_payload,
+                codex_session_affinity=codex_session_affinity,
+            )
         if session_anchor is not None:
             original_input_items = cast(list[JsonValue], responses_payload.input)
             original_input_item_count = len(original_input_items)
@@ -2892,7 +2958,10 @@ async def _prepare_websocket_response_create_request(
                     request_state.input_item_count,
                 )
         affinity_policy = _sticky_key_for_responses_request(
-            responses_payload,
+            # Only the proven restart uses the pre-injection body. Ordinary
+            # full resends must be classified after anchor injection so they
+            # cannot accidentally gain soft-session mobility.
+            restart_affinity_payload if goal_restart_full_resend else responses_payload,
             headers,
             codex_session_affinity=codex_session_affinity,
             openai_cache_affinity=openai_cache_affinity,
@@ -3293,6 +3362,7 @@ async def _select_websocket_connect_account(
                     sticky_source=request_state.affinity_policy.codex_session_source,
                     legacy_sticky_key=request_state.affinity_policy.legacy_selection_key,
                     spill_bare_session_on_account_cap=request_state.affinity_policy.spill_on_account_cap,
+                    abandon_unavailable_legacy_owner=(request_state.affinity_policy.abandon_unavailable_legacy_owner),
                     require_unambiguous_account=request_state.affinity_policy.require_unambiguous_account,
                     sticky_max_age_seconds=sticky_max_age_seconds,
                     prefer_earlier_reset_accounts=prefer_earlier_reset,