From 26b1134d408ac367dbaa13584e056c0b11cac8a7 Mon Sep 17 00:00:00 2001 From: Roman Leventov Date: Wed, 12 Aug 2026 14:41:38 +0800 Subject: [PATCH 1/2] fix(proxy): scope backend Codex affinity by thread identity --- .../proxy/_load_balancer/sticky_selection.py | 79 +++- app/modules/proxy/_service/codex_control.py | 6 + app/modules/proxy/_service/compact.py | 15 +- .../proxy/_service/http_bridge/helpers.py | 44 ++- .../proxy/_service/http_bridge/mixin.py | 29 +- .../_service/http_bridge/owner_forwarding.py | 3 +- .../_service/http_bridge/request_submit.py | 12 +- .../proxy/_service/http_bridge/streaming.py | 28 +- app/modules/proxy/_service/streaming/retry.py | 7 +- app/modules/proxy/_service/support.py | 2 + app/modules/proxy/_service/websocket/mixin.py | 78 +++- .../proxy/_service/websocket/protocol.py | 1 + app/modules/proxy/affinity.py | 250 ++++++++++++- app/modules/proxy/continuity.py | 9 +- app/modules/proxy/load_balancer.py | 20 +- app/modules/proxy/service.py | 28 +- app/modules/proxy/sticky_repository.py | 36 ++ .../.openspec.yaml | 2 + .../scope-codex-affinity-by-thread/design.md | 77 ++++ .../proposal.md | 38 ++ .../specs/responses-api-compat/spec.md | 221 +++++++++++ .../specs/sticky-session-operations/spec.md | 301 +++++++++++++++ .../scope-codex-affinity-by-thread/tasks.md | 18 + .../integration/test_proxy_sticky_sessions.py | 138 +++++++ tests/unit/test_load_balancer_concurrency.py | 238 +++++++++++- tests/unit/test_proxy_http_bridge.py | 289 ++++++++++++++ tests/unit/test_proxy_utils.py | 352 ++++++++++++++++++ 27 files changed, 2238 insertions(+), 83 deletions(-) create mode 100644 openspec/changes/scope-codex-affinity-by-thread/.openspec.yaml create mode 100644 openspec/changes/scope-codex-affinity-by-thread/design.md create mode 100644 openspec/changes/scope-codex-affinity-by-thread/proposal.md create mode 100644 openspec/changes/scope-codex-affinity-by-thread/specs/responses-api-compat/spec.md create mode 100644 openspec/changes/scope-codex-affinity-by-thread/specs/sticky-session-operations/spec.md create mode 100644 openspec/changes/scope-codex-affinity-by-thread/tasks.md diff --git a/app/modules/proxy/_load_balancer/sticky_selection.py b/app/modules/proxy/_load_balancer/sticky_selection.py index a1b61c2639..d05e5cc450 100644 --- a/app/modules/proxy/_load_balancer/sticky_selection.py +++ b/app/modules/proxy/_load_balancer/sticky_selection.py @@ -185,6 +185,7 @@ async def _select_with_stickiness( sticky_repo: StickySessionsRepository | None, routing_costs_by_account_id: RoutingCostsByAccount | None, sticky_existing_account_id: str | None | object, + initial_preferred_account_id: str | None, preserve_existing_mapping_on_fallback: bool, traffic_class: TrafficClass, ignore_standard_quota: bool, @@ -203,6 +204,9 @@ class StickySelectionRequest(Generic[SelectionInputsT]): sticky_source: _CodexSessionSource | None legacy_sticky_key: str | None legacy_existing_account_id: str | None + sticky_seed_key: str | None + sticky_seed_kind: StickySessionKind | None + sticky_seed_account_id: str | None spill_bare_session_on_account_cap: bool require_unambiguous_account: bool sticky_max_age_seconds: int | None @@ -265,6 +269,9 @@ async def run_sticky_selection_path( sticky_source = request.sticky_source legacy_sticky_key = request.legacy_sticky_key legacy_existing_account_id = request.legacy_existing_account_id + sticky_seed_key = request.sticky_seed_key + sticky_seed_kind = request.sticky_seed_kind + sticky_seed_account_id = request.sticky_seed_account_id spill_bare_session_on_account_cap = request.spill_bare_session_on_account_cap require_unambiguous_account = request.require_unambiguous_account sticky_max_age_seconds = request.sticky_max_age_seconds @@ -335,7 +342,7 @@ def _direct_error( # always has, rather than silently bypassing the ambiguous # owner check below. sticky_continuity_abandoned = sticky_owner_lookup.continuity_abandoned is True - if sticky_kind == StickySessionKind.CODEX_SESSION and sticky_existing_is_legacy: + if sticky_existing_is_legacy: # Mixed-version replicas can create both rows on # different accounts. The raw row was loaded before # branch selection and always wins as possible hard @@ -368,10 +375,8 @@ def _direct_error( and not sticky_existing_is_legacy ) cap_spillover_allowed = spill_bare_session_on_account_cap and lease_kind is not None and bare_session_key - hard_sticky = ( - sticky_kind == StickySessionKind.CODEX_SESSION - and isinstance(sticky_existing_account_id, str) - and not bare_session_key + hard_sticky = isinstance(sticky_existing_account_id, str) and ( + sticky_existing_is_legacy or (sticky_kind == StickySessionKind.CODEX_SESSION and not bare_session_key) ) if hard_sticky and required_account_id is not None and sticky_existing_account_id != required_account_id: return _direct_error( @@ -534,6 +539,11 @@ def _direct_error( relative_availability_top_k=relative_availability_top_k, sticky_repo=repos.sticky_sessions, sticky_existing_account_id=sticky_existing_account_id, + initial_preferred_account_id=( + sticky_seed_account_id + if not isinstance(sticky_existing_account_id, str) and not sticky_continuity_abandoned + else None + ), preserve_existing_mapping_on_fallback=preserve_existing_mapping, traffic_class=traffic_class, ignore_standard_quota=False, @@ -808,6 +818,11 @@ def _direct_error( assert sticky_mutation is not None try: async with owner._repo_factory() as repos: + # A recovery-probe reservation is still reversible until + # the runtime CAS below succeeds. Persist its thread row so + # existing rollback machinery can restore it, but do not + # publish an immutable process seed that cannot be safely + # deleted after a concurrent sibling observes it. await _persist_sticky_mutation( sticky_repo=repos.sticky_sessions, sticky_key=sticky_key, @@ -953,6 +968,12 @@ def _direct_error( sticky_key=sticky_key, sticky_kind=sticky_kind, mutation=sticky_mutation, + initialize_seed_key=( + sticky_seed_key + if sticky_source == "thread_header" and sticky_seed_account_id is None + else None + ), + initialize_seed_kind=sticky_seed_kind, ) except BaseException: # Runtime admission may already be committed. Preserve @@ -991,6 +1012,7 @@ async def _select_with_stickiness( sticky_repo: StickySessionsRepository | None, routing_costs_by_account_id: RoutingCostsByAccount | None = None, sticky_existing_account_id: str | None | object = _STICKY_EXISTING_UNSET, + initial_preferred_account_id: str | None = None, preserve_existing_mapping_on_fallback: bool = False, traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, @@ -1046,6 +1068,32 @@ def finish_selection( persist_fallback = not preserve_existing_mapping_on_fallback apply_sticky_secondary_budget_threshold = False + if not existing and initial_preferred_account_id is not None: + initial_preferred = next( + (state for state in states if state.account_id == initial_preferred_account_id), + None, + ) + if initial_preferred is not None: + initial_result = select_account( + [initial_preferred], + prefer_earlier_reset=prefer_earlier_reset_accounts, + prefer_earlier_reset_window=prefer_earlier_reset_window, + routing_strategy=routing_strategy, + allow_backoff_fallback=False, + relative_availability_power=relative_availability_power, + relative_availability_top_k=relative_availability_top_k, + traffic_class=traffic_class, + ignore_standard_quota=ignore_standard_quota, + routing_costs=routing_costs_by_account_id, + ) + if initial_result.account is not None: + # Persist only the new thread row. The process mapping supplied + # the preference but is deliberately outside this mutation. + return finish_selection( + initial_result, + persist_account_id=initial_preferred.account_id, + ) + if existing: pinned = next((state for state in states if state.account_id == existing), None) if pinned is not None: @@ -1249,10 +1297,31 @@ async def _persist_sticky_mutation( sticky_key: str, sticky_kind: StickySessionKind, mutation: _StickyMutation, + initialize_seed_key: str | None = None, + initialize_seed_kind: StickySessionKind | None = None, ) -> None: if mutation.account_id is None: await sticky_repo.delete(sticky_key, kind=sticky_kind) return + if initialize_seed_key is not None: + if initialize_seed_kind is None: + raise ValueError("initialize_seed_kind is required when initialize_seed_key is provided") + # Current Codex sends thread-id on the first root request, so a fresh + # process has no older bare-session request available to create its + # default. Initialize it exactly once from the first admitted thread. + # insert-if-absent is essential: failover or a later child may move its + # own bounded row but can never rewrite the process/sibling default. + # The repository operation is intentionally atomic; splitting it into + # the public insert/upsert methods would commit a process default even + # when persistence of the initiating thread fails. + await sticky_repo.upsert_with_seed_if_absent( + sticky_key, + mutation.account_id, + kind=sticky_kind, + seed_key=initialize_seed_key, + seed_kind=initialize_seed_kind, + ) + return await sticky_repo.upsert(sticky_key, mutation.account_id, kind=sticky_kind) diff --git a/app/modules/proxy/_service/codex_control.py b/app/modules/proxy/_service/codex_control.py index 6879c7c3c4..242f2bbd30 100644 --- a/app/modules/proxy/_service/codex_control.py +++ b/app/modules/proxy/_service/codex_control.py @@ -229,6 +229,8 @@ async def _select_codex_control_account_without_budget( reallocate_sticky=affinity.reallocate_sticky, sticky_source=affinity.codex_session_source, legacy_sticky_key=affinity.legacy_selection_key, + sticky_seed_key=affinity.seed_selection_key, + sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, account_ids=scoped_account_ids, prefer_earlier_reset_window=prefer_earlier_reset_window, @@ -396,6 +398,8 @@ async def _select_control_failover(excluded_account_ids: set[str]) -> AccountSel reallocate_sticky=affinity.reallocate_sticky, sticky_source=affinity.codex_session_source, legacy_sticky_key=affinity.legacy_selection_key, + sticky_seed_key=affinity.seed_selection_key, + sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, prefer_earlier_reset_accounts=settings.prefer_earlier_reset_accounts, routing_strategy=routing_strategy, @@ -488,6 +492,8 @@ async def _select_control_failover(excluded_account_ids: set[str]) -> AccountSel reallocate_sticky=affinity.reallocate_sticky, sticky_source=affinity.codex_session_source, legacy_sticky_key=affinity.legacy_selection_key, + sticky_seed_key=affinity.seed_selection_key, + sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, prefer_earlier_reset_accounts=settings.prefer_earlier_reset_accounts, prefer_earlier_reset_window=_prefer_earlier_reset_window(settings), diff --git a/app/modules/proxy/_service/compact.py b/app/modules/proxy/_service/compact.py index cdcc8b9d1d..08b032cf93 100644 --- a/app/modules/proxy/_service/compact.py +++ b/app/modules/proxy/_service/compact.py @@ -48,6 +48,7 @@ _resolve_prompt_cache_key, _sticky_key_from_session_header, _sticky_key_from_turn_state_header, + _thread_codex_session_affinity, ) from app.modules.proxy.api_key_usage import estimate_api_key_request_usage from app.modules.proxy.continuity import resolve_required_account_id @@ -432,6 +433,14 @@ def _sticky_key_for_compact_request( kind=StickySessionKind.CODEX_SESSION, codex_session_source="turn_state", ) + elif ( + thread_affinity := _thread_codex_session_affinity( + headers, + enabled=codex_session_affinity, + max_age_seconds=openai_cache_affinity_max_age_seconds, + ) + ) is not None: + policy = thread_affinity elif ( session_affinity := _bare_codex_session_affinity( headers, @@ -618,7 +627,11 @@ async def compact_responses( api_key=api_key, ) sticky_key_source = "none" - if affinity.kind == StickySessionKind.CODEX_SESSION: + if affinity.codex_session_source == "thread_header": + # The payload cache hint remains unchanged; diagnostics must not + # imply that it supplied the internal thread-local routing key. + sticky_key_source = "thread_header" + elif affinity.kind == StickySessionKind.CODEX_SESSION: if _sticky_key_from_turn_state_header(headers) is not None: sticky_key_source = "turn_state_header" elif _sticky_key_from_session_header(headers) is not None: diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index 27072ceaf1..3571746338 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -165,6 +165,7 @@ from app.modules.proxy.account_cache import is_account_routing_unavailable from app.modules.proxy.affinity import ( _AffinityPolicy, + _codex_backend_identity, _extract_model_class, _sticky_key_from_session_header, _sticky_key_from_turn_state_header, @@ -1053,6 +1054,7 @@ def _http_bridge_incompatible_model_fork_key( ) -> "_HTTPBridgeSessionKey | None": if key.affinity_kind not in { "session_header", + "thread_header", "turn_state_header", "internal_unanchored_parallel", "internal_model_parallel", @@ -1121,7 +1123,11 @@ def _http_bridge_parallel_fork_key( """Give incompatible or concurrent requests an independent websocket lane.""" reason: str | None = None - if key.affinity_kind == "session_header" and incoming_turn_state is None and previous_response_id is None: + if ( + key.affinity_kind in {"session_header", "thread_header"} + and incoming_turn_state is None + and previous_response_id is None + ): if inflight_creation: reason = "session_creation_inflight" elif session is not None and not session.closed: @@ -1199,7 +1205,11 @@ def _http_bridge_request_needs_unanchored_handoff( ) -> bool: if forwarded_request: return forwarded_original_request_unanchored - return key.affinity_kind == "session_header" and incoming_turn_state is None and previous_response_id is None + return ( + key.affinity_kind in {"session_header", "thread_header"} + and incoming_turn_state is None + and previous_response_id is None + ) def _reserve_http_bridge_unanchored_handoff( @@ -1567,14 +1577,19 @@ def _make_http_bridge_session_key( affinity_key = turn_state_key affinity_kind = "turn_state_header" strength: Literal["hard", "soft"] = "hard" + elif (thread_key := _codex_backend_identity(headers).thread_selection_key) is not None: + # prompt_cache_key is intentionally shared by current Codex root trees. + # The thread key is canonical identity; once a bridge exists it is hard + # transport continuity even though pre-bridge account locality is soft. + affinity_key = thread_key + affinity_kind = "thread_header" + strength = "hard" else: session_key = _sticky_key_from_session_header(headers) if session_key is not None: - # One Codex process session can host several independent agent - # threads. Codex keeps the process-level session header shared but - # gives every thread a stable explicit prompt_cache_key. Keying - # only by the header makes a later, non-overlapping child reuse the - # parent's upstream conversation and receive the wrong history. + # Compatibility path for clients that do not expose thread-id. + # Current Codex reaches the thread_header branch above; do not + # reintroduce prompt_cache_key as thread identity here. session_header_key = _make_http_bridge_session_header_fallback_key( headers=headers, api_key=api_key, @@ -1602,6 +1617,12 @@ def _make_http_bridge_session_header_fallback_key( api_key: ApiKeyData | None, explicit_prompt_cache_key: str | None, ) -> _HTTPBridgeSessionKey | None: + if _codex_backend_identity(headers).thread_id is not None: + # Never let a current thread attach to the legacy + # (session-id, prompt_cache_key) lane: both values are shared across + # siblings. Exact turn-state/previous-response aliases are handled by + # durable lookup independently and remain the only safe migration path. + return None session_key = _sticky_key_from_session_header(headers) if session_key is None: return None @@ -1708,7 +1729,7 @@ def _http_bridge_can_local_recover_without_ring( ): return True return ( - key.affinity_kind == "session_header" + key.affinity_kind in {"session_header", "thread_header"} and previous_response_id is None and _sticky_key_from_turn_state_header(headers) is None ) @@ -2398,7 +2419,10 @@ def _effective_http_bridge_idle_ttl_seconds( codex_idle_ttl_seconds: float, prompt_cache_idle_ttl_seconds: float | None = None, ) -> float: - if affinity.kind == StickySessionKind.CODEX_SESSION: + if affinity.kind == StickySessionKind.CODEX_SESSION or affinity.codex_session_source == "thread_header": + # The DB row is bounded soft locality, but a live thread bridge owns + # upstream socket history and therefore receives the Codex continuity + # lifetime once created. return max(idle_ttl_seconds, codex_idle_ttl_seconds) if affinity.kind == StickySessionKind.PROMPT_CACHE and prompt_cache_idle_ttl_seconds is not None: return prompt_cache_idle_ttl_seconds @@ -2604,7 +2628,7 @@ def _http_bridge_should_attempt_local_bootstrap_rebind( headers: Mapping[str, str], previous_response_id: str | None, ) -> bool: - if key.affinity_kind != "session_header": + if key.affinity_kind not in {"session_header", "thread_header"}: return False if previous_response_id is not None: return False diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index c5b8b4d126..6dd9f39552 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -204,6 +204,7 @@ ) from app.modules.proxy.affinity import ( _AffinityPolicy, + _codex_backend_identity, _extract_model_class, _sticky_key_from_session_header, _sticky_key_from_turn_state_header, @@ -429,8 +430,19 @@ async def _get_or_create_http_bridge_session( request_scope_id = ensure_request_scope_id() api_key_id = api_key.id if api_key is not None else None incoming_turn_state = _sticky_key_from_turn_state_header(headers) - incoming_session_key = _sticky_key_from_session_header(headers) - initial_session_key = session_header_fallback_key or (key if key.affinity_kind == "session_header" else None) + thread_selection_key = _codex_backend_identity(headers).thread_selection_key + thread_fallback_key = None + if thread_selection_key is not None: + thread_fallback_key = _HTTPBridgeSessionKey("thread_header", thread_selection_key, api_key_id) + # Exact aliases may fall back to this thread, never the raw process lane + # where mixed-version replicas may have stored a sibling. V2 forwards use + # their signed key; this derived fallback protects legacy raw headers. + incoming_session_key = None if thread_fallback_key is not None else _sticky_key_from_session_header(headers) + initial_session_key = ( + session_header_fallback_key + or thread_fallback_key + or (key if key.affinity_kind == "session_header" else None) + ) original_request_unanchored = _http_bridge_request_needs_unanchored_handoff( key, incoming_turn_state, previous_response_id, forwarded_request, forwarded_original_request_unanchored ) @@ -662,10 +674,11 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: key=key.affinity_key, ): key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id) + elif initial_session_key is not None: + key = initial_session_key + used_session_header_fallback = True elif incoming_session_key is not None: - key = initial_session_key or _HTTPBridgeSessionKey( - "session_header", incoming_session_key, api_key_id - ) + key = _HTTPBridgeSessionKey("session_header", incoming_session_key, api_key_id) used_session_header_fallback = True else: key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id) @@ -1975,7 +1988,7 @@ async def _create_http_bridge_session( lifecycle_lock=anyio.Lock(), last_used_at=_service_time().monotonic(), idle_ttl_seconds=idle_ttl_seconds, - codex_session=affinity.kind == StickySessionKind.CODEX_SESSION, + codex_session=(affinity.kind == StickySessionKind.CODEX_SESSION or key.affinity_kind == "thread_header"), prewarm_lock=anyio.Lock(), upstream_turn_state=_upstream_turn_state_from_socket(upstream), downstream_turn_state=None, @@ -2395,7 +2408,9 @@ async def abort_selected_handoff() -> None: session.last_completed_input_prefix_fingerprint = None session.last_pending_tool_calls.clear() session.affinity = selection_affinity or session.affinity - session.codex_session = False + # Clearing stale response/turn aliases makes an account move + # safe; it does not make the canonical thread lane soft. + session.codex_session = session.key.affinity_kind == "thread_header" session.upstream_turn_state = None session.downstream_turn_state = None session.headers = { diff --git a/app/modules/proxy/_service/http_bridge/owner_forwarding.py b/app/modules/proxy/_service/http_bridge/owner_forwarding.py index 2155d8ebed..6739a45ea8 100644 --- a/app/modules/proxy/_service/http_bridge/owner_forwarding.py +++ b/app/modules/proxy/_service/http_bridge/owner_forwarding.py @@ -381,7 +381,8 @@ async def _forward_http_bridge_request_to_owner( original_request_unanchored=( recovery_forward or ( - owner_forward.key.affinity_kind in {"session_header", "internal_unanchored_parallel"} + owner_forward.key.affinity_kind + in {"session_header", "thread_header", "internal_unanchored_parallel"} and incoming_turn_state is None and payload.previous_response_id is None ) diff --git a/app/modules/proxy/_service/http_bridge/request_submit.py b/app/modules/proxy/_service/http_bridge/request_submit.py index 23216a8c36..f1fe4eaa09 100644 --- a/app/modules/proxy/_service/http_bridge/request_submit.py +++ b/app/modules/proxy/_service/http_bridge/request_submit.py @@ -61,6 +61,7 @@ set_request_id, ) from app.core.utils.sse import format_sse_event, parse_sse_data_json +from app.db.models import StickySessionKind from app.modules.api_keys.service import ( ApiKeyData, ApiKeyUsageReservationData, @@ -3409,7 +3410,7 @@ def mark_security_retry_send_started() -> None: request_state.account_response_create_release = self._load_balancer.release_account_lease if session.account.id != owner_account_id: if ( - previous_session_affinity.codex_session_source == "session_header" + previous_session_affinity.codex_session_source in {"session_header", "thread_header"} and previous_session_affinity.selection_key is not None and previous_session_affinity.kind is not None ): @@ -3554,12 +3555,15 @@ async def _claim_http_bridge_replacement_before_swap( if account_id == session.account.id: return try: - if owner_rebind_affinity.legacy_selection_key is not None and owner_rebind_affinity.kind is not None: + if owner_rebind_affinity.legacy_selection_key is not None: async with self._repo_factory() as repos: legacy_owner_id = await repos.sticky_sessions.get_account_id( owner_rebind_affinity.legacy_selection_key, - kind=owner_rebind_affinity.kind, - max_age_seconds=owner_rebind_affinity.max_age_seconds, + # The new thread row may be PROMPT_CACHE, but the raw + # compatibility row has always been CODEX_SESSION and + # remains durable hard ownership. + kind=StickySessionKind.CODEX_SESSION, + max_age_seconds=None, ) if legacy_owner_id is not None and legacy_owner_id != account_id: raise ProxyResponseError( diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index d7f2d8f599..066d37cef5 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -204,6 +204,7 @@ ) from app.modules.proxy.affinity import ( _AffinityPolicy, + _codex_backend_identity, _extract_model_class, _prompt_cache_key_from_request_model, _request_allows_bare_session_cap_spillover, @@ -1183,7 +1184,9 @@ async def release_unowned_bridge_lifecycle( api_key=api_key, ) sticky_key_source = "none" - if affinity.kind == StickySessionKind.CODEX_SESSION: + if affinity.codex_session_source == "thread_header": + sticky_key_source = "thread_header" + elif affinity.kind == StickySessionKind.CODEX_SESSION: sticky_key_source = ( "turn_state_header" if _sticky_key_from_turn_state_header(headers) is not None else "session_header" ) @@ -1227,6 +1230,13 @@ async def release_unowned_bridge_lifecycle( if not forwarded_request else None ) + durable_session_header_alias = ( + None + if _codex_backend_identity(headers).thread_id is not None + else session_header_fallback_key.affinity_key + if explicit_prompt_cache_key is not None and session_header_fallback_key is not None + else incoming_session_header + ) legacy_anchor_lookup = await _legacy_forward_anchor_lookup( durable_bridge=self._durable_bridge, bridge_session_key=bridge_session_key, @@ -1254,11 +1264,9 @@ async def release_unowned_bridge_lifecycle( session_key_value=bridge_session_key.affinity_key, api_key_id=bridge_session_key.api_key_id, turn_state=durable_lookup_turn_state, - session_header=( - session_header_fallback_key.affinity_key - if explicit_prompt_cache_key is not None and session_header_fallback_key is not None - else incoming_session_header - ), + # A raw process alias is ambiguous when thread-id exists. + # Exact turn/response aliases remain independent inputs. + session_header=durable_session_header_alias, previous_response_id=payload.previous_response_id, ) except ProxyResponseError: @@ -1641,6 +1649,7 @@ def classify_durable_full_resend( affinity = _AffinityPolicy() incoming_turn_state_header = None session_header_fallback_key = None + durable_session_header_alias = None owner_bound_full_resend_ignores_broad_session = ( not forwarded_request and durable_full_resend_fresh_bridge_proof is not None @@ -1655,6 +1664,7 @@ def classify_durable_full_resend( affinity = _AffinityPolicy(kind=StickySessionKind.CODEX_SESSION) incoming_session_header = None session_header_fallback_key = None + durable_session_header_alias = None _log_http_bridge_event( "fresh_reattach_broad_session_owner_ignored", bridge_session_key, @@ -2209,11 +2219,7 @@ def switch_to_account_neutral_replay() -> None: session_key_value=bridge_session_key.affinity_key, api_key_id=bridge_session_key.api_key_id, turn_state=takeover_turn_state, - session_header=( - session_header_fallback_key.affinity_key - if explicit_prompt_cache_key is not None and session_header_fallback_key is not None - else incoming_session_header - ), + session_header=durable_session_header_alias, previous_response_id=effective_payload.previous_response_id, ) except Exception: diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index f09346a162..5f55d6cb02 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -63,6 +63,7 @@ _sticky_key_for_responses_request, _sticky_key_from_session_header, _sticky_key_from_turn_state_header, + _websocket_continuity_key_from_headers, ) from app.modules.proxy.api_key_usage import estimate_api_key_request_usage from app.modules.proxy.continuity import resolve_required_account_id @@ -130,7 +131,7 @@ def _verified_cross_transport_fresh_replay( input_items = cast(list[Any], input_value) if not _websocket_input_items_are_self_contained_fresh_replay(input_items): return None - session_id = _owner_lookup_session_id_from_headers(headers) + session_id = _websocket_continuity_key_from_headers(headers) if session_id is None: return None api_key_id = api_key.id if api_key is not None else None @@ -336,7 +337,9 @@ async def _stream_with_retry( fail_on_missing=not _is_synthesized_turn_state(turn_state), ) sticky_key_source = "none" - if affinity.kind == StickySessionKind.CODEX_SESSION: + if affinity.codex_session_source == "thread_header": + sticky_key_source = "thread_header" + elif affinity.kind == StickySessionKind.CODEX_SESSION: sticky_key_source = "session_header" elif affinity.key: sticky_key_source = "payload" if had_prompt_cache_key else "derived" diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index c9de31d37b..975078d89b 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -72,6 +72,7 @@ { "turn_state_header", "session_header", + "thread_header", "internal_unanchored_parallel", "internal_model_parallel", "internal_request_parallel", @@ -936,6 +937,7 @@ class _WebSocketRequestState: account_response_create_release: Callable[[AccountLease | None], Coroutine[Any, Any, None]] | None = None websocket_stream_lease: AccountLease | None = None affinity_policy: _AffinityPolicy = field(default_factory=_AffinityPolicy) + thread_affinity_last_touch_at: float = field(default_factory=time.monotonic) suppressed_downstream_tool_call: bool = False suppressed_duplicate_tool_call: bool = False pending_function_call_ids: list[str] = field(default_factory=list) diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index 3b7a1b1e2d..f50dd0e9e9 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -436,6 +436,7 @@ _sticky_key_for_responses_request, _sticky_key_from_session_header, # noqa: F401 _sticky_key_from_turn_state_header, + _websocket_continuity_aliases_from_headers, ) from app.modules.proxy.api_key_usage import estimate_api_key_request_usage from app.modules.proxy.capability_routing import ( @@ -1213,6 +1214,45 @@ async def _process_upstream_websocket_transport_end( class _WebSocketMixin: + async def _touch_active_websocket_thread_affinity( + self, + request_state: _WebSocketRequestState, + account: Account, + ) -> None: + """Refresh bounded thread locality without turning it into ownership.""" + + proxy = cast(_WebSocketServiceProtocol, self) + policy = request_state.affinity_policy + if ( + policy.codex_session_source != "thread_header" + or policy.selection_key is None + or policy.kind != StickySessionKind.PROMPT_CACHE + or policy.max_age_seconds is None + ): + return + now = time.monotonic() + touch_interval = max(1.0, min(float(policy.max_age_seconds) / 2.0, 60.0)) + if now - request_state.thread_affinity_last_touch_at < touch_interval: + return + try: + # A response can outlive the selection TTL. Throttled event-time + # touches keep reconnect locality current, while exact response or + # bridge ownership remains the hard authority for this turn. + async with proxy._repo_factory() as repos: + await repos.sticky_sessions.upsert( + policy.selection_key, + account.id, + kind=policy.kind, + ) + except Exception: + _facade().logger.warning( + "Failed to refresh active Codex thread affinity account_id=%s", + account.id, + exc_info=True, + ) + return + request_state.thread_affinity_last_touch_at = now + def _websocket_continuity_state_for_request( self, headers: Mapping[str, str], @@ -1225,28 +1265,37 @@ def _websocket_continuity_state_for_request( _ = proxy if not codex_session_affinity: return _WebSocketContinuityState() - session_id = _owner_lookup_session_id_from_headers(headers, synthesized_turn_state=synthesized_turn_state) api_key_id = api_key.id if api_key is not None else None - cache_keys: list[tuple[str, str | None]] = [] - if session_id is not None: - cache_keys.append((session_id, api_key_id)) - if synthesized_turn_state is not None: - generated_key = (synthesized_turn_state, api_key_id) - if generated_key not in cache_keys: - cache_keys.append(generated_key) + cache_keys = [ + (continuity_key, api_key_id) + for continuity_key in _websocket_continuity_aliases_from_headers( + headers, + synthesized_turn_state=synthesized_turn_state, + ) + ] if not cache_keys: return _WebSocketContinuityState() + explicit_turn_state = _sticky_key_from_turn_state_header(headers) + exact_client_turn = explicit_turn_state is not None and explicit_turn_state != synthesized_turn_state + # An exact client turn state is hard continuity. If its alias is + # unknown, do not borrow retained response/tool state from the broader + # thread key; the turn may have a different owner. Once the exact alias + # resolves, publishing that same state under the thread key is safe and + # keeps a later unanchored reconnect thread-local. + lookup_keys = cache_keys[:1] if exact_client_turn else cache_keys continuity_state = next( ( existing_state - for key in cache_keys + for key in lookup_keys if (existing_state := proxy._websocket_continuity_index.get(key)) is not None ), None, ) + exact_alias_resolved = continuity_state is not None if continuity_state is None: continuity_state = _WebSocketContinuityState() - for key in cache_keys: + publish_keys = cache_keys if not exact_client_turn or exact_alias_resolved else lookup_keys + for key in publish_keys: proxy._websocket_continuity_index.pop(key, None) proxy._websocket_continuity_index[key] = continuity_state while len(proxy._websocket_continuity_index) > _facade()._WEBSOCKET_CONTINUITY_CACHE_LIMIT: @@ -2966,7 +3015,9 @@ async def _prepare_websocket_response_create_request( synthesized_turn_state=synthesized_turn_state, ) sticky_key_source = "none" - if affinity_policy.kind == StickySessionKind.CODEX_SESSION: + if affinity_policy.codex_session_source == "thread_header": + sticky_key_source = "thread_header" + elif affinity_policy.kind == StickySessionKind.CODEX_SESSION: turn_state_key = _sticky_key_from_turn_state_header(headers) if turn_state_key is not None and turn_state_key == synthesized_turn_state: sticky_key_source = "generated_turn_state" @@ -3356,6 +3407,8 @@ async def _select_websocket_connect_account( reallocate_sticky=reallocate_sticky, sticky_source=request_state.affinity_policy.codex_session_source, legacy_sticky_key=request_state.affinity_policy.legacy_selection_key, + sticky_seed_key=request_state.affinity_policy.seed_selection_key, + sticky_seed_kind=request_state.affinity_policy.seed_selection_kind, spill_bare_session_on_account_cap=request_state.affinity_policy.spill_on_account_cap, require_unambiguous_account=request_state.affinity_policy.require_unambiguous_account, sticky_max_age_seconds=sticky_max_age_seconds, @@ -4971,6 +5024,9 @@ async def _process_upstream_websocket_text( if event_type == "response.created" and release_create_gate and created_request_state is not None: await _release_websocket_response_create_gate(created_request_state, response_create_gate) + if request_state is not None: + await proxy._touch_active_websocket_thread_affinity(request_state, account) + if len(grouped_previous_response_request_states) > 1: upstream_control.reconnect_requested = True downstream_texts: list[str] = [] diff --git a/app/modules/proxy/_service/websocket/protocol.py b/app/modules/proxy/_service/websocket/protocol.py index 519cd7c868..f0dcd2c343 100644 --- a/app/modules/proxy/_service/websocket/protocol.py +++ b/app/modules/proxy/_service/websocket/protocol.py @@ -59,6 +59,7 @@ class _WebSocketServiceProtocol(Protocol): _settle_stream_api_key_usage: Any _start_request_state_api_key_reservation_heartbeat: Any _try_open_websocket_connect_attempt: Any + _touch_active_websocket_thread_affinity: Any _websocket_continuity_index: Any _websocket_continuity_state_for_request: Any _websocket_previous_response_account_index: Any diff --git a/app/modules/proxy/affinity.py b/app/modules/proxy/affinity.py index 52e941c45f..7cdcee0529 100644 --- a/app/modules/proxy/affinity.py +++ b/app/modules/proxy/affinity.py @@ -13,7 +13,7 @@ from collections.abc import Mapping from dataclasses import dataclass, replace from hashlib import sha256 -from typing import Literal, cast +from typing import Literal, TypedDict, cast from uuid import uuid4 from app.core.config.settings import get_settings @@ -23,7 +23,7 @@ # This typed provenance is a routing capability: callers must never recover it # from key text, because a client-controlled turn state can mimic any prefix. -_CodexSessionSource = Literal["session_header", "turn_state"] +_CodexSessionSource = Literal["session_header", "thread_header", "turn_state"] # Request headers are stripped and HTTP forbids CR/LF, while PostgreSQL/SQLite # text keys can safely retain LF. This sentinel makes the internal namespace # structurally unreachable by every legacy raw header, even if its digest is @@ -31,6 +31,19 @@ _CODEX_SELECTION_KEY_PREFIX = "\ncodex-lb-affinity-v1" +class _AffinitySelectionKwargs(TypedDict): + sticky_key: str | None + sticky_kind: StickySessionKind | None + reallocate_sticky: bool + sticky_source: _CodexSessionSource | None + legacy_sticky_key: str | None + sticky_seed_key: str | None + sticky_seed_kind: StickySessionKind | None + spill_bare_session_on_account_cap: bool + require_unambiguous_account: bool + sticky_max_age_seconds: int | None + + @dataclass(frozen=True, slots=True) class _AffinityPolicy: key: str | None = None @@ -41,6 +54,17 @@ class _AffinityPolicy: spill_on_account_cap: bool = False max_age_seconds: int | None = None codex_session_source: _CodexSessionSource | None = None + # A thread row is soft locality, but old replicas may have persisted the + # raw process/session value as hard CODEX_SESSION ownership. Keep that + # compatibility lookup explicit instead of trying to reconstruct it from + # the new opaque thread key. + legacy_codex_session_key: str | None = None + # A previously unseen thread should inherit the healthy process preference + # once, then persist its own bounded row. This is never ownership: a + # missing process default may be initialized once by insert-if-absent, but + # no thread request may update or delete an established process row. + seed_selection_key: str | None = None + seed_selection_kind: StickySessionKind | None = None # ``conversation`` has no dedicated owner index. Preserve that provenance # until selection can prove one hard owner or a one-account pool. require_unambiguous_account: bool = False @@ -59,8 +83,29 @@ def legacy_selection_key(self) -> str | None: # Old replicas persisted bare session headers as raw CODEX_SESSION # keys. Always consult this alongside the soft row: any raw hit may be # hard turn-state ownership and therefore takes precedence. + if self.legacy_codex_session_key is not None: + return self.legacy_codex_session_key return self.key if self.codex_session_source == "session_header" else None + def selection_kwargs(self) -> _AffinitySelectionKwargs: + """Expand routing policy once at the account-selection boundary.""" + + # Keep the compatibility edge from silently omitting new policy + # fields. In particular, thread locality is incomplete if callers pass + # its row but forget the process seed or legacy hard-owner lookup. + return { + "sticky_key": self.selection_key, + "sticky_kind": self.kind, + "reallocate_sticky": self.reallocate_sticky, + "sticky_source": self.codex_session_source, + "legacy_sticky_key": self.legacy_selection_key, + "sticky_seed_key": self.seed_selection_key, + "sticky_seed_kind": self.seed_selection_kind, + "spill_bare_session_on_account_cap": self.spill_on_account_cap, + "require_unambiguous_account": self.require_unambiguous_account, + "sticky_max_age_seconds": self.max_age_seconds, + } + @staticmethod def cap_spillover_allowed( capability: bool, @@ -98,7 +143,8 @@ def preferred_owner_sticky_inputs( # A resolved response/file/bridge owner bypasses the new soft row, but # the raw compatibility row still has to be checked for conflicting # legacy hard ownership. Selection receives no writable sticky key, so - # a raw miss cannot manufacture or rebind a mapping. + # a raw miss cannot manufacture or rebind a mapping. The caller also + # deliberately omits any broader process seed in this exact-owner path. return None, StickySessionKind.CODEX_SESSION, False, sticky_max_age_seconds, sticky_source, legacy_sticky_key @@ -109,6 +155,73 @@ def _codex_session_selection_key(key: str) -> str: return f"{_CODEX_SELECTION_KEY_PREFIX}:session_header:{digest}" +@dataclass(frozen=True, slots=True) +class _CodexBackendIdentity: + """Independently parsed process-tree and logical-thread identities.""" + + process_session: str | None + thread_id: str | None + + @property + def thread_selection_key(self) -> str | None: + if self.thread_id is None: + return None + # The explicit scope tag prevents the thread-only compatibility form + # from colliding with (process, thread). Length framing keeps distinct + # client tuples distinct even if a future non-HTTP caller admits NULs + # or other delimiters. The LF namespace remains unreachable by headers. + if self.process_session is None: + parts = ("thread-only", self.thread_id) + scope = "thread_only" + else: + parts = ("process-thread", self.process_session, self.thread_id) + scope = "process_thread" + encoded_parts = (part.encode() for part in parts) + framed = b"".join(len(part).to_bytes(8, "big") + part for part in encoded_parts) + digest = sha256(framed).hexdigest() + return f"{_CODEX_SELECTION_KEY_PREFIX}:thread_header:{scope}:{digest}" + + +_CODEX_PROCESS_SESSION_HEADERS = ( + "session_id", + "session-id", + "x-codex-session-id", + "x-codex-conversation-id", +) + + +def _normalized_header_value(headers: Mapping[str, str], names: tuple[str, ...]) -> str | None: + normalized = {key.lower(): value for key, value in headers.items()} + for name in names: + value = normalized.get(name) + if not isinstance(value, str): + continue + stripped = value.strip() + if stripped: + return stripped + return None + + +def _process_session_key_from_headers(headers: Mapping[str, str]) -> str | None: + return _normalized_header_value(headers, _CODEX_PROCESS_SESSION_HEADERS) + + +def _thread_id_from_headers(headers: Mapping[str, str]) -> str | None: + return _normalized_header_value(headers, ("thread-id",)) + + +def _codex_backend_identity( + headers: Mapping[str, str], + *, + thread_id: str | None = None, +) -> _CodexBackendIdentity: + normalized_thread_id = thread_id.strip() if isinstance(thread_id, str) and thread_id.strip() else None + return _CodexBackendIdentity( + process_session=_process_session_key_from_headers(headers), + thread_id=normalized_thread_id if thread_id is not None else _thread_id_from_headers(headers), + ) + + def _prompt_cache_key_from_request_model(payload: ResponsesRequest | ResponsesCompactRequest) -> str | None: typed_value = getattr(payload, "prompt_cache_key", None) if isinstance(typed_value, str) and typed_value: @@ -252,15 +365,10 @@ def _sticky_key_from_payload(payload: ResponsesRequest) -> str | None: def _sticky_key_from_session_header(headers: Mapping[str, str]) -> str | None: - normalized = {key.lower(): value for key, value in headers.items()} - for key in ("session_id", "session-id", "x-codex-session-id", "x-codex-conversation-id", "thread-id"): - value = normalized.get(key) - if not isinstance(value, str): - continue - stripped = value.strip() - if stripped: - return stripped - return None + # Legacy owner/request-log callers still need the historical alias order. + # New account, bridge, and replay locality MUST use the typed process/thread + # helpers above; otherwise a shared process id silently hides thread-id. + return _process_session_key_from_headers(headers) or _thread_id_from_headers(headers) def _sticky_key_from_turn_state_header(headers: Mapping[str, str]) -> str | None: @@ -291,6 +399,36 @@ def _bare_codex_session_affinity( ) +def _thread_codex_session_affinity( + headers: Mapping[str, str], + *, + enabled: bool, + max_age_seconds: int, + thread_id: str | None = None, +) -> _AffinityPolicy | None: + if not enabled: + return None + identity = _codex_backend_identity(headers, thread_id=thread_id) + thread_key = identity.thread_selection_key + if thread_key is None: + return None + # Current Codex shares process session and prompt_cache_key across a root + # tree. Thread locality therefore reuses the bounded PROMPT_CACHE lifecycle + # but does not rewrite the upstream cache hint or create durable child rows. + legacy_key = identity.process_session or identity.thread_id + return _AffinityPolicy( + key=thread_key, + kind=StickySessionKind.PROMPT_CACHE, + max_age_seconds=max_age_seconds, + codex_session_source="thread_header", + legacy_codex_session_key=legacy_key, + seed_selection_key=( + _codex_session_selection_key(identity.process_session) if identity.process_session is not None else None + ), + seed_selection_kind=(StickySessionKind.CODEX_SESSION if identity.process_session is not None else None), + ) + + def _request_allows_bare_session_cap_spillover( payload: ResponsesRequest | ResponsesCompactRequest, ) -> bool: @@ -347,6 +485,37 @@ def _sticky_key_for_codex_control_request( return _AffinityPolicy() +def _sticky_key_for_thread_goal_request( + payload: Mapping[str, object], + headers: Mapping[str, str], + codex_session_affinity: bool, + max_age_seconds: int, +) -> _AffinityPolicy: + turn_state_key = _sticky_key_from_turn_state_header(headers) + if turn_state_key is not None: + return _AffinityPolicy( + key=turn_state_key, + kind=StickySessionKind.CODEX_SESSION, + codex_session_source="turn_state", + ) + payload_thread_id = payload.get("threadId") + if isinstance(payload_thread_id, str) and payload_thread_id.strip(): + thread_affinity = _thread_codex_session_affinity( + headers, + enabled=codex_session_affinity, + max_age_seconds=max_age_seconds, + thread_id=payload_thread_id, + ) + if thread_affinity is not None: + return thread_affinity + # Routing only consumes a valid nonblank identity. The upstream thread-goal + # protocol remains authoritative for payload validation and error shape. + return _sticky_key_for_codex_control_request( + headers, + codex_session_affinity=codex_session_affinity, + ) + + def _owner_lookup_session_id_from_headers( headers: Mapping[str, str], *, @@ -363,6 +532,55 @@ def _owner_lookup_session_id_from_headers( return _sticky_key_from_session_header(headers) +def _websocket_continuity_key_from_headers( + headers: Mapping[str, str], + *, + synthesized_turn_state: str | None = None, +) -> str | None: + """Return the primary count-bounded direct-WebSocket continuity key.""" + + explicit_turn_state = _sticky_key_from_turn_state_header(headers) + if explicit_turn_state is not None and explicit_turn_state != synthesized_turn_state: + # Exact client continuation must outrank broader thread locality. A + # synthesized handshake placeholder is only an alias for the current + # connection and therefore does not gain this hard precedence. + return explicit_turn_state + identity = _codex_backend_identity(headers) + if identity.thread_selection_key is not None: + return identity.thread_selection_key + return _owner_lookup_session_id_from_headers( + headers, + synthesized_turn_state=synthesized_turn_state, + ) + + +def _websocket_continuity_aliases_from_headers( + headers: Mapping[str, str], + *, + synthesized_turn_state: str | None = None, +) -> tuple[str, ...]: + """Keep exact turn aliases without restoring process-wide thread state.""" + + aliases: list[str] = [] + primary = _websocket_continuity_key_from_headers( + headers, + synthesized_turn_state=synthesized_turn_state, + ) + if primary is not None: + aliases.append(primary) + thread_key = _codex_backend_identity(headers).thread_selection_key + if thread_key is not None: + # When an exact turn resolved first, refresh the thread alias to that + # same state so a later unanchored reconnect remains thread-local. + aliases.append(thread_key) + explicit_turn_state = _sticky_key_from_turn_state_header(headers) + if explicit_turn_state is not None and explicit_turn_state != synthesized_turn_state: + aliases.append(explicit_turn_state) + if synthesized_turn_state is not None: + aliases.append(synthesized_turn_state) + return tuple(dict.fromkeys(aliases)) + + # Pattern matching turn-state values synthesized by the helpers below. # A 32-char lowercase hex (uuid4().hex) suffix follows the prefix. _SYNTHESIZED_TURN_STATE_PATTERN = re.compile(r"^(?:http_)?turn_[0-9a-f]{32}$") @@ -450,6 +668,14 @@ def _sticky_key_for_responses_request( kind=StickySessionKind.CODEX_SESSION, codex_session_source="turn_state", ) + elif ( + thread_affinity := _thread_codex_session_affinity( + headers, + enabled=codex_session_affinity, + max_age_seconds=openai_cache_affinity_max_age_seconds, + ) + ) is not None: + policy = thread_affinity elif ( session_affinity := _bare_codex_session_affinity( headers, diff --git a/app/modules/proxy/continuity.py b/app/modules/proxy/continuity.py index c26ba9a93b..ca6cb46d0c 100644 --- a/app/modules/proxy/continuity.py +++ b/app/modules/proxy/continuity.py @@ -11,7 +11,14 @@ HTTP_BRIDGE_ACCOUNT_NEUTRAL_REPLAY_KIND = "internal_unanchored_parallel" HTTP_BRIDGE_ACCOUNT_NEUTRAL_REPLAY_KEY_PREFIX = "account-neutral-replay:v1:" -HTTP_BRIDGE_ACCOUNT_NEUTRAL_REPLAY_REBINDABLE_KINDS = frozenset({"prompt_cache", "session_header", "turn_state_header"}) +# These are canonical lanes whose exact aliases may move only after the +# existing full-resend validator has proved the request account-neutral. A +# thread lane is hard during ordinary use, just like a session-header lane; +# omitting it here would accidentally remove safe owner-unavailable recovery +# merely because Codex now supplies a more precise canonical identity. +HTTP_BRIDGE_ACCOUNT_NEUTRAL_REPLAY_REBINDABLE_KINDS = frozenset( + {"prompt_cache", "session_header", "thread_header", "turn_state_header"} +) _HTTP_BRIDGE_SESSION_AFFINITY_HEADERS = frozenset( { "session_id", diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index ec6ca10917..47b1b53bbf 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -521,6 +521,8 @@ async def select_account( reallocate_sticky: bool = False, sticky_source: _CodexSessionSource | None = None, legacy_sticky_key: str | None = None, + sticky_seed_key: str | None = None, + sticky_seed_kind: StickySessionKind | None = None, spill_bare_session_on_account_cap: bool = False, require_unambiguous_account: bool = False, sticky_max_age_seconds: int | None = None, @@ -686,12 +688,14 @@ async def load_selection_inputs() -> _SelectionInputs: 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: + if legacy_sticky_key is not None: async with self._repo_factory() as repos: legacy_existing_account_id = await repos.sticky_sessions.get_account_id( legacy_sticky_key, kind=StickySessionKind.CODEX_SESSION, - max_age_seconds=sticky_max_age_seconds, + # Raw rows may be historical turn-state ownership. The + # bounded thread TTL must never age out that hard evidence. + max_age_seconds=None, ) if required_account_id is not None and ( legacy_existing_account_id is not None and legacy_existing_account_id != required_account_id @@ -704,6 +708,13 @@ async def load_selection_inputs() -> _SelectionInputs: error_message="Account-owned continuity sources conflict; retry the logical turn", error_code="continuity_owner_conflict", ) + sticky_seed_account_id: str | None = None + if sticky_seed_key is not None and sticky_seed_kind is not None: + async with self._repo_factory() as repos: + sticky_seed_account_id = await repos.sticky_sessions.get_account_id( + sticky_seed_key, + kind=sticky_seed_kind, + ) # Resolve uniqueness from the model/API-key/security-scoped pool before # runtime health, budget, or cap filtering. Transient pressure cannot # prove that another candidate does not own an upstream conversation. @@ -776,6 +787,9 @@ async def load_selection_inputs() -> _SelectionInputs: sticky_source=sticky_source, legacy_sticky_key=legacy_sticky_key, legacy_existing_account_id=legacy_existing_account_id, + sticky_seed_key=sticky_seed_key, + sticky_seed_kind=sticky_seed_kind, + sticky_seed_account_id=sticky_seed_account_id, spill_bare_session_on_account_cap=spill_bare_session_on_account_cap, require_unambiguous_account=require_unambiguous_account, sticky_max_age_seconds=sticky_max_age_seconds, @@ -1539,6 +1553,7 @@ async def _select_with_stickiness( sticky_repo: StickySessionsRepository | None, routing_costs_by_account_id: RoutingCostsByAccount | None = None, sticky_existing_account_id: str | None | object = _STICKY_EXISTING_UNSET, + initial_preferred_account_id: str | None = None, preserve_existing_mapping_on_fallback: bool = False, traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, @@ -1562,6 +1577,7 @@ async def _select_with_stickiness( sticky_repo=sticky_repo, routing_costs_by_account_id=routing_costs_by_account_id, sticky_existing_account_id=sticky_existing_account_id, + initial_preferred_account_id=initial_preferred_account_id, preserve_existing_mapping_on_fallback=preserve_existing_mapping_on_fallback, traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index ed3ed91013..6b52113671 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -709,7 +709,7 @@ from app.modules.proxy.affinity import ( _AffinityPolicy, _CodexSessionSource, - _sticky_key_for_codex_control_request, + _sticky_key_for_thread_goal_request, _sticky_key_from_session_header, # noqa: F401 ) from app.modules.proxy.affinity import ( @@ -995,9 +995,8 @@ async def thread_goal_request( base_settings = get_settings() deadline = start + base_settings.proxy_request_budget_seconds settings = await get_settings_cache().get() - affinity = _sticky_key_for_codex_control_request( - headers, - codex_session_affinity=codex_session_affinity, + affinity = _sticky_key_for_thread_goal_request( + payload, headers, codex_session_affinity, settings.openai_cache_affinity_max_age_seconds ) selection_model = api_key.enforced_model if api_key is not None else None routing_strategy = _routing_strategy(settings) @@ -1093,6 +1092,8 @@ async def _select_goal_failover(excluded_account_ids: set[str]) -> AccountSelect reallocate_sticky=affinity.reallocate_sticky, sticky_source=affinity.codex_session_source, legacy_sticky_key=affinity.legacy_selection_key, + sticky_seed_key=affinity.seed_selection_key, + sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, prefer_earlier_reset_accounts=settings.prefer_earlier_reset_accounts, routing_strategy=routing_strategy, @@ -1427,16 +1428,7 @@ async def _select_account_with_budget_compatible( affinity_policy = kwargs.pop("affinity_policy", None) if isinstance(affinity_policy, _AffinityPolicy): # Expand once at the compatibility edge so transport callers cannot drift. - kwargs.update( - sticky_key=affinity_policy.selection_key, - sticky_kind=affinity_policy.kind, - reallocate_sticky=affinity_policy.reallocate_sticky, - sticky_source=affinity_policy.codex_session_source, - legacy_sticky_key=affinity_policy.legacy_selection_key, - spill_bare_session_on_account_cap=affinity_policy.spill_on_account_cap, - require_unambiguous_account=affinity_policy.require_unambiguous_account, - sticky_max_age_seconds=affinity_policy.max_age_seconds, - ) + kwargs.update(affinity_policy.selection_kwargs()) required_capability_kwargs = {} if kwargs.get("require_security_work_authorized") is True: required_capability_kwargs["require_security_work_authorized"] = kwargs.pop( @@ -1722,6 +1714,8 @@ async def _select_account_with_budget( reallocate_sticky: bool = False, sticky_source: _CodexSessionSource | None = None, legacy_sticky_key: str | None = None, + sticky_seed_key: str | None = None, + sticky_seed_kind: StickySessionKind | None = None, spill_bare_session_on_account_cap: bool = False, require_unambiguous_account: bool = False, sticky_max_age_seconds: int | None = None, @@ -1879,6 +1873,10 @@ def log_account_id(account_id: str | None) -> str | None: sticky_max_age_seconds=preferred_sticky_inputs[3], sticky_source=preferred_sticky_inputs[4], legacy_sticky_key=preferred_sticky_inputs[5], + # Exact ownership chooses the account; a first-ever thread + # still seeds atomically without overwriting a process default. + sticky_seed_key=sticky_seed_key, + sticky_seed_kind=sticky_seed_kind, prefer_earlier_reset_accounts=prefer_earlier_reset_accounts, prefer_earlier_reset_window=prefer_earlier_reset_window, routing_strategy=routing_strategy, @@ -1936,6 +1934,8 @@ def log_account_id(account_id: str | None) -> str | None: reallocate_sticky=reallocate_sticky, sticky_source=sticky_source, legacy_sticky_key=legacy_sticky_key, + sticky_seed_key=sticky_seed_key, + sticky_seed_kind=sticky_seed_kind, spill_bare_session_on_account_cap=_AffinityPolicy.cap_spillover_allowed( spill_bare_session_on_account_cap, preferred_account_id, diff --git a/app/modules/proxy/sticky_repository.py b/app/modules/proxy/sticky_repository.py index 0948da94c5..bd0794e96d 100644 --- a/app/modules/proxy/sticky_repository.py +++ b/app/modules/proxy/sticky_repository.py @@ -191,6 +191,42 @@ async def insert_if_absent(self, key: str, account_id: str, kind: StickySessionK raise RuntimeError("StickySession immutable insert did not resolve an owner") return owner_id + async def upsert_with_seed_if_absent( + self, + key: str, + account_id: str, + *, + kind: StickySessionKind, + seed_key: str, + seed_kind: StickySessionKind, + ) -> StickySession: + """Upsert one mapping and initialize its immutable seed atomically.""" + + # Keep these writes in one transaction. A process seed without the + # initiating thread row is false placement evidence, while a thread + # row without its seed makes the first admitted thread invisible to + # later siblings. Do not replace the seed's DO NOTHING with an upsert: + # another thread may have won first-writer initialization already. + seed_statement = self._build_insert_do_nothing_statement(seed_key, account_id, seed_kind) + mapping_statement = self._build_upsert_statement(key, account_id, kind).returning(StickySession) + async with sqlite_writer_section(): + try: + await self._session.execute(seed_statement) + result = await self._session.execute( + mapping_statement, + execution_options={"populate_existing": True}, + ) + row = result.scalar_one_or_none() + if row is None: + raise RuntimeError(f"StickySession seeded upsert failed for key={key!r} kind={kind.value!r}") + await self._session.commit() + except BaseException: + # This method owns both writes as one unit even when a caller + # catches the error and keeps using the same session. + await self._session.rollback() + raise + return row + async def delete(self, key: str, *, kind: StickySessionKind) -> bool: if not key: return False diff --git a/openspec/changes/scope-codex-affinity-by-thread/.openspec.yaml b/openspec/changes/scope-codex-affinity-by-thread/.openspec.yaml new file mode 100644 index 0000000000..5081c98763 --- /dev/null +++ b/openspec/changes/scope-codex-affinity-by-thread/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-08-12 diff --git a/openspec/changes/scope-codex-affinity-by-thread/design.md b/openspec/changes/scope-codex-affinity-by-thread/design.md new file mode 100644 index 0000000000..a3125eb2d5 --- /dev/null +++ b/openspec/changes/scope-codex-affinity-by-thread/design.md @@ -0,0 +1,77 @@ +## Context + +Upstream Codex deliberately separates three roles: `session-id` identifies a +root process tree, `thread-id` identifies one root/child/fork/resumed thread, +and `prompt_cache_key` defaults to the shared process session to colocate +cache entries. The old codex-lb bridge assumption that explicit cache keys +distinguish children stopped holding when upstream adopted the shared cache +key. + +## Decisions + +1. **Use a typed thread identity.** Parse process session and `thread-id` + independently. Derive a versioned, header-inaccessible opaque key from both, + with a separately namespaced thread-only fallback when no process session + exists. Never infer identity from subagent or parent markers. +2. **Reuse bounded prompt-cache rows for soft locality.** Thread locality uses + the existing `prompt_cache` kind and configured freshness window. A missing + thread row first prefers the eligible source-separated process-session row; + successful admission persists the selected account under the thread key + without rewriting the process row. Because current Codex includes + `thread-id` on its first root request, the first admitted thread also + initializes a missing process preference with an atomic insert-if-absent. + Later thread movement can neither overwrite that first-writer default nor + mutate a sibling row. A provisional recovery-probe reservation persists + only its reversible thread row; it cannot publish the immutable process + default until a normal admission does so, because a failed probe CAS cannot + safely delete a seed that a concurrent sibling may already have observed. +3. **Keep ownership separate from locality.** Raw legacy `codex_session` rows + are looked up independently and remain hard. Exact turn state, response, + file, conversation, bridge, replay, and reattach evidence keeps its existing + precedence and conflict behavior. +4. **Use the same logical identity at transport boundaries.** Direct + WebSocket retained replay/tool state is count-bounded in memory by thread. + An HTTP bridge canonical lane is hard while live/durable, but its pre-bridge + account hint remains soft and bounded. +5. **Migrate bridge lanes only through exact aliases.** With current Codex, + legacy `(session-id, prompt_cache_key)` is shared by siblings. A request + carrying `thread-id` must not fall back to that canonical key. Existing + lanes remain recoverable through exact turn-state or previous-response + aliases and otherwise expire naturally. Authenticated forwarded affinity + keys are accepted verbatim and never derived again. +6. **Preserve upstream cache intent.** `prompt_cache_key` is forwarded + unchanged. Request-log conversation grouping continues using raw + `thread-id`. + +## Rejected Alternatives + +- Adapting the marker/TTL/schema design in #1309: thread identity is already + explicit, so marker inference, a migration, settings, and dashboard controls + add lifecycle without improving identity. +- Using subagent or parent-thread markers: they describe role/provenance and + can group siblings rather than identify the current thread. +- Using `prompt_cache_key` or `(session-id, prompt_cache_key)`: current Codex + intentionally sends the same value for root and children. +- Rewriting `prompt_cache_key` to `thread-id`: this defeats upstream's intended + tree-wide cache colocation. +- Balancing every unseen thread independently: a thread boundary contains + justified divergence; it is not a reason to discard the healthy process + preference on first placement. +- Durable per-thread Codex rows or a new sticky kind/setting: existing bounded + rows already express soft locality, while bridge/object ownership remains + durable separately. +- Falling back to the old bridge canonical key when `thread-id` exists: that + can attach a sibling's history. Only exact hard aliases are safe migration + evidence. + +## Risks / Trade-offs + +- Thread rows expire. Selection and active response completion refresh the row + so a long-lived active thread does not lose reconnect locality solely to the + freshness window. +- A health, quota, explicit restart, or proven hard-owner transition may still + move a thread. The fix removes cross-thread collisions; it does not promise + an account never changes. +- Mixed-version replicas may retain old shared bridge rows. New requests with + a thread identity ignore those rows unless an exact alias proves ownership, + allowing safe coexistence until cleanup. diff --git a/openspec/changes/scope-codex-affinity-by-thread/proposal.md b/openspec/changes/scope-codex-affinity-by-thread/proposal.md new file mode 100644 index 0000000000..0014440f6d --- /dev/null +++ b/openspec/changes/scope-codex-affinity-by-thread/proposal.md @@ -0,0 +1,38 @@ +## Why + +Current Codex sends one process `session-id` and one `prompt_cache_key` across +an entire root/subagent tree while giving each logical conversation its own +stable `thread-id`. codex-lb still keys backend account locality, direct +WebSocket replay state, and HTTP bridge lanes primarily from the shared +process/cache identities. Sibling threads therefore overwrite one another's +locality and can reuse replay or bridge history that belongs to another +thread. + +## What Changes + +- Derive one source-separated bounded locality key from the process session + and `thread-id` for backend Responses and compact requests. +- Seed a new thread from an eligible process-session preference, then persist + only the thread-local bounded row so later failover does not move siblings. +- Key direct WebSocket retained state and HTTP bridge canonical lanes by the + same logical thread identity. +- Route thread-goal operations from their payload `threadId`. +- Preserve explicit turn state, previous response, file, conversation, bridge, + replay, and legacy raw Codex rows as hard ownership; keep `prompt_cache_key` + unchanged as the upstream cache hint. + +## Capabilities + +### Modified Capabilities + +- `responses-api-compat`: Scope backend Responses, compact, direct WebSocket, + HTTP bridge, and thread-goal locality by Codex thread identity. +- `sticky-session-operations`: Keep process preference and legacy hard-owner + semantics while introducing bounded per-thread locality. + +## Impact + +The change touches proxy affinity parsing, account selection, direct +WebSocket continuity, HTTP bridge identity, thread-goal routing, and focused +tests. It adds no setting, schema, migration, dashboard surface, or upstream +payload rewrite. diff --git a/openspec/changes/scope-codex-affinity-by-thread/specs/responses-api-compat/spec.md b/openspec/changes/scope-codex-affinity-by-thread/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..e9d3df3d81 --- /dev/null +++ b/openspec/changes/scope-codex-affinity-by-thread/specs/responses-api-compat/spec.md @@ -0,0 +1,221 @@ +## MODIFIED Requirements + +### Requirement: Codex backend session_id preserves account affinity + +When a backend Codex Responses or compact request includes a nonblank +`thread-id`, the service MUST use a source-separated bounded key derived from +the independently parsed process session and thread identity for soft account +locality. If the thread has no mapping, selection MUST first prefer an eligible +source-separated process-session mapping and then persist the admitted thread +mapping. If no process-session mapping exists, the first admitted thread MUST +initialize that soft process preference atomically without overwriting a +concurrent or later first writer, unless its account is admitted only through +a recovery-probe reservation. A recovery-probe admission MUST NOT initialize +the immutable process preference; its reversible thread row MAY be persisted +independently until a normal admission establishes the process default. + +When `thread-id` is absent, a non-empty accepted process-session header MUST +retain its established account-affinity behavior. Accepted process-session +headers are `session_id`, `session-id`, `x-codex-session-id`, and +`x-codex-conversation-id`, in that priority order. A client-supplied nonblank +`x-codex-turn-state` remains a more specific hard continuity key. If the +request lacks a client-supplied `prompt_cache_key`, the service MUST derive and +attach a stable `prompt_cache_key` before upstream forwarding so account +affinity and upstream prompt-cache routing can coexist. A client-supplied +`prompt_cache_key` MUST be forwarded unchanged and MUST NOT be used as thread +identity. + +A turn state synthesized by the proxy for the current downstream WebSocket +handshake MUST NOT override client-supplied process/thread identity or a +prompt-cache key for routing or WebSocket continuity selection. The proxy MUST +seed WebSocket continuity storage under that synthesized turn state so a later +client echo can reuse the completed-turn owner. The proxy MUST continue to +forward that synthesized turn state upstream. A turn state sent by the client, +including one that the proxy generated and the client later echoed, remains a +client-supplied turn-state affinity key. + +When a WebSocket handshake has neither a client-supplied turn state nor an +accepted process/thread identity, the proxy MUST store its generated turn state +as the WebSocket continuity key. A later connection that echoes that accepted +value MUST recover the same continuity state. Direct WebSocket retained +response, input-prefix, Responses Lite, and unresolved-tool state MUST use the +derived thread identity plus API-key scope, with count-bounded storage. +Request-log conversation grouping MUST continue to use raw `thread-id`. + +#### Scenario: Backend Codex request derives prompt_cache_key before codex-session routing + +- **WHEN** `/backend-api/codex/responses` is called with `session_id` and without `thread-id` or `prompt_cache_key` +- **THEN** the routing decision retains process-session `codex_session` affinity +- **AND** the forwarded upstream payload includes a derived stable `prompt_cache_key` + +#### Scenario: backend WebSocket reconnect retains session affinity despite a generated turn state + +- **WHEN** two backend Codex Responses WebSocket connections include the same process session and `thread-id` and omit `x-codex-turn-state` +- **AND** the proxy generates a distinct turn state for each handshake +- **THEN** both account selections use the same bounded thread-local affinity key +- **AND** each generated turn state is still forwarded to the upstream + +#### Scenario: echoed generated turn state remains a client continuation key + +- **WHEN** a client reconnects with a non-empty `x-codex-turn-state` value it received from an earlier proxy handshake +- **THEN** that turn state remains the routing and WebSocket continuity key ahead of broader process/thread locality +- **AND** full-resend continuity for that echoed turn state can reuse the earlier completed response anchor + +#### Scenario: generated turn state seeds continuity without a session header + +- **WHEN** a backend Codex Responses WebSocket handshake omits process/thread identity and `x-codex-turn-state` +- **AND** the proxy generates and returns a turn state for that handshake +- **THEN** the proxy stores its WebSocket continuity state under that generated value +- **AND WHEN** a later connection sends that value in `x-codex-turn-state` +- **THEN** it recovers the stored continuity state + +#### Scenario: Root and child keep separate locality with one cache hint + +- **GIVEN** root and child requests share a process session and explicit `prompt_cache_key` +- **AND** they carry different stable `thread-id` values +- **WHEN** backend Responses or compact routes them +- **THEN** they use different bounded internal thread keys +- **AND** both upstream payloads retain the original `prompt_cache_key` + +#### Scenario: New thread inherits process preference without coupling siblings + +- **GIVEN** a process-session soft row points to eligible account A +- **AND** a previously unseen thread in that process arrives +- **WHEN** selection admits the request +- **THEN** it prefers account A and persists a bounded row for that thread +- **AND** later movement of that thread does not rewrite the process row or a sibling row + +#### Scenario: First thread initializes the process preference + +- **GIVEN** a fresh process has no process-session or thread mapping +- **WHEN** its first thread is admitted on account A +- **THEN** it initializes the process preference to A with insert-if-absent +- **AND** a later sibling prefers A without gaining authority to rewrite that process preference + +#### Scenario: Exact owner admission still initializes first-thread locality + +- **GIVEN** a fresh process has no process-session or thread mapping +- **AND** an exact response, file, or bridge owner requires account A +- **WHEN** the first thread is admitted on account A through that hard owner +- **THEN** the thread row and absent process preference are persisted atomically +- **AND** the process preference remains insert-only if another thread already initialized it + +#### Scenario: Recovery probe does not seed the process + +- **GIVEN** a fresh process has no process-session mapping +- **WHEN** a thread is selected on probing account A through a recovery reservation +- **THEN** account A is not published as the immutable process preference +- **AND** a failed reservation commit can restore the reversible thread placement + +#### Scenario: Direct WebSocket siblings do not share replay state + +- **GIVEN** sibling threads share one process session and cache key +- **WHEN** each uses direct WebSocket Responses and one reconnects +- **THEN** retained response, prefix, Lite, and pending-tool state is read only from that thread +- **AND** the reconnect cannot inject or replay its sibling's state + +#### Scenario: Unknown exact turn does not borrow broader thread replay + +- **GIVEN** a direct WebSocket thread has retained replay or tool state +- **WHEN** a request supplies a nonblank client turn state with no exact in-memory alias +- **THEN** it does not reuse or replace the broader thread state +- **AND** only a previously resolved exact alias may refresh the thread alias + +### Requirement: HTTP Responses routes preserve upstream websocket session continuity + +When serving HTTP `/v1/responses` or HTTP `/backend-api/codex/responses`, the +service MUST preserve upstream Responses websocket session continuity on a +stable per-session bridge key instead of opening a brand new upstream session +for every eligible request. For backend Codex requests carrying `thread-id`, +the canonical bridge key MUST use the same derived logical thread identity used +for account locality and compact routing. Otherwise the bridge key MUST use an +explicit session/conversation header when present, then normalized +`prompt_cache_key`, deriving a stable key from the existing cache-affinity +inputs when the client omits one. While bridged, the service MUST preserve the +external HTTP/SSE contract, continue request logging with `transport = "http"`, +and keep requests from different bridge keys isolated. + +An established live or durable thread bridge is hard continuity. The bridge +MUST retain request-scoped fork lanes for concurrent unanchored requests and +MUST preserve exact turn-state and previous-response aliases. A request with +`thread-id` MUST NOT fall back to the legacy canonical key derived from process +session and `prompt_cache_key`, because current Codex may share both across +siblings. It MAY recover an old bridge only through an exact hard alias. +Authenticated forwarded affinity kind/key values MUST remain verbatim and MUST +NOT be namespaced or hashed again. + +#### Scenario: bridge forwards hard continuity keys to the owner replica + +- **WHEN** operators configure multiple eligible bridge instance ids +- **AND** a request uses a bridge key derived from `x-codex-turn-state`, an explicit legacy session header, or `thread-id` +- **AND** that request lands on a non-owner instance +- **THEN** the service MUST forward the request internally to the owner replica +- **AND** it MUST NOT return a topology-bearing `bridge_instance_mismatch` error to the client for that owner mismatch alone + +#### Scenario: gateway-style prompt-cache bridge requests tolerate wrong-replica arrival + +- **WHEN** a request uses a bridge key derived only from `prompt_cache_key` or a derived prompt-cache key +- **AND** that request lands on a non-owner instance +- **THEN** the service MAY create or reuse a local bridge session on that instance +- **AND** it MUST treat the owner mismatch as a locality miss instead of a continuity failure + +#### Scenario: forwarded bridge requests fail closed when owner forwarding loops + +- **WHEN** a forwarded hard-continuity bridge request reaches another non-owner replica +- **THEN** the service MUST fail the request with a generic 5xx bridge-forward error +- **AND** it MUST NOT attempt another owner handoff + +#### Scenario: local restart orphan is recovered by the replacement instance + +- **WHEN** a single local bridge instance is replaced while durable hard-continuity ownership still references the old instance id +- **AND** the old owner has no distinct active forwarding endpoint from the current replacement instance +- **THEN** the replacement instance MUST treat the row as restart-orphaned and may claim durable ownership locally +- **AND** same-account takeover MUST preserve the latest persisted response anchor until a replacement response id is recorded +- **AND** normal client retries MUST NOT be stranded waiting for the old instance lease to expire + +When request aliases resolve to different durable rows for the same account, +an explicitly requested previous-response alias MUST select its row even if +that row has since advanced to a newer response id. Without an explicitly +resolved previous-response alias, recovery MUST select the freshest row that +contains a persisted response anchor rather than using alias enumeration order. + +#### Scenario: requested durable response alias survives same-account row divergence + +- **GIVEN** turn-state and previous-response aliases resolve to different durable rows for the same account +- **AND** the request names the previous-response alias whose row has since advanced to a newer response id +- **WHEN** the service resolves durable continuity +- **THEN** it selects the row resolved by the requested previous-response alias +- **AND** it preserves that row's latest persisted response anchor + +#### Scenario: Sequential siblings use distinct canonical bridges + +- **GIVEN** root and child requests share process session and `prompt_cache_key` +- **AND** they carry different `thread-id` values +- **WHEN** the child starts after the root request completes +- **THEN** the child uses a different canonical bridge identity +- **AND** another child request reuses only the child's lane + +#### Scenario: Old shared canonical lane is not a thread fallback + +- **GIVEN** an old bridge exists under `(session-id, prompt_cache_key)` +- **WHEN** a request carries a new thread identity but no exact turn-state or previous-response alias +- **THEN** it does not attach to the old shared lane +- **AND** it creates or reuses its thread-canonical lane + +## ADDED Requirements + +### Requirement: Thread-goal routing uses payload thread identity + +Thread-goal get, set, and clear operations carrying a nonblank payload +`threadId` MUST select account locality from that exact thread identity, +combined with the process session when available. An explicit client turn +state remains hard continuity and MUST retain its existing precedence or +conflict behavior. Other generic request headers MUST NOT cause one sibling +thread's goal operation to follow another sibling's locality. Existing +protocol forwarding and error behavior MUST remain unchanged. + +#### Scenario: Sibling goal operations follow their own threads + +- **GIVEN** sibling threads share a process session but have distinct `threadId` values +- **WHEN** each invokes thread-goal get, set, or clear +- **THEN** each operation uses its own bounded thread locality diff --git a/openspec/changes/scope-codex-affinity-by-thread/specs/sticky-session-operations/spec.md b/openspec/changes/scope-codex-affinity-by-thread/specs/sticky-session-operations/spec.md new file mode 100644 index 0000000000..7bb53c9577 --- /dev/null +++ b/openspec/changes/scope-codex-affinity-by-thread/specs/sticky-session-operations/spec.md @@ -0,0 +1,301 @@ +## MODIFIED Requirements + +### Requirement: Bare process-session cap spillover is non-mutating + +The system MUST parse process-session and thread headers independently. A bare +process-session mapping and a bounded thread-local mapping MUST use distinct, +header-inaccessible storage identities, so a client-supplied hard turn-state +value cannot alias either derived soft row. A current replica MUST consult a +legacy raw Codex-session key independently even when a namespaced process or +thread row exists. Any raw hit MUST take precedence as hard ownership. If a +resolved file, response, bridge, or other exact owner conflicts with that raw +legacy owner, the request MUST fail closed without creating or rewriting any +of those rows. + +A missing thread row MAY use an eligible process-session soft row as its +initial placement preference. If that process row is missing, the first +admitted thread MUST initialize it with insert-if-absent and MUST persist its +own bounded thread row. A concurrent or later thread MUST NOT overwrite that +first-writer process preference. Account-cap spillover or later thread movement +MUST NOT rewrite or delete the process-session mapping or a sibling's thread +mapping. A provisional recovery-probe reservation MUST NOT initialize a +missing process preference, because its thread mapping may still require +rollback and a probing account is not a stable process default. A later normal +admission MAY initialize the missing process preference. + +When the mapped account for a bare process-session key is locally capped and +another eligible account is selected, the spillover MUST apply only to that +request. Selection MUST NOT update or delete the stored process-session mapping +because of account-cap spillover. If the mapped account is below cap, normal +sticky selection MUST retain it. + +#### Scenario: Capped bare-session owner spills without rebinding + +- **GIVEN** a bare process-session mapping points to account A +- **AND** account A is locally capped +- **AND** account B is eligible and below cap +- **WHEN** a self-contained pre-visible request is selected +- **THEN** the request uses account B +- **AND** the stored process-session mapping still points to account A + +#### Scenario: Unsaturated bare-session owner retains locality + +- **GIVEN** a bare process-session mapping points to eligible account A below its local caps +- **WHEN** a self-contained request is selected +- **THEN** the request uses account A +- **AND** the mapping remains unchanged + +#### Scenario: Equal session and turn-state values remain isolated + +- **GIVEN** a process-session header and an explicit turn-state header have equal text values +- **WHEN** their affinity mappings are resolved +- **THEN** the process-session mapping uses a source-separated opaque key +- **AND** the explicit turn-state mapping continues to use the legacy raw key as hard ownership + +#### Scenario: Derived soft key cannot be reused as raw hard turn state + +- **GIVEN** a process-session or thread value has a derived internal storage key +- **WHEN** a client submits the visible representation of that key as a turn-state header +- **THEN** header normalization cannot reproduce the internal storage identity +- **AND** hard turn-state selection cannot read or rewrite the soft row + +#### Scenario: Legacy raw mapping remains hard + +- **GIVEN** a legacy replica persisted a raw Codex-session mapping +- **WHEN** a current replica receives the matching process or legacy thread header +- **THEN** it does not reinterpret or mutate the legacy raw row as spillable affinity +- **AND** mixed-version operation remains fail-closed for that row + +#### Scenario: Coexisting legacy and namespaced rows prefer hard ownership + +- **GIVEN** mixed-version replicas created a raw row and a namespaced process or thread row for the same request identity +- **AND** the rows point to different accounts +- **WHEN** a current replica selects the request +- **THEN** the raw row's account is treated as the hard owner +- **AND** neither row is deleted or rewritten by account-cap spillover + +#### Scenario: Legacy hard owner conflicts with resolved owner + +- **GIVEN** a raw legacy session row points to account A +- **AND** a file, previous response, or bridge resolves to account B +- **WHEN** the request is routed +- **THEN** the service fails with `continuity_owner_conflict` +- **AND** it neither bypasses nor rewrites the raw row + +#### Scenario: Process preference seeds only the new thread + +- **GIVEN** a process-session soft row points to account A +- **AND** no bounded row exists for thread T +- **WHEN** T is admitted on account A or a safely selected alternate +- **THEN** the admitted account is persisted under T's bounded key +- **AND** the process-session row remains unchanged + +#### Scenario: Missing process preference is initialized once + +- **GIVEN** no process-session mapping exists +- **WHEN** the first thread is admitted on account A and a concurrent or later thread is admitted on account B +- **THEN** insert-if-absent preserves the first persisted process owner +- **AND** each thread persists only its own bounded locality after that initialization + +#### Scenario: Provisional probe placement does not escape rollback + +- **GIVEN** neither process nor thread has a bounded mapping +- **WHEN** a probing-account placement persists provisionally and then loses its runtime commit +- **THEN** its thread mutation is restored +- **AND** no immutable process preference is left behind + +#### Scenario: Legacy raw owner wins over thread locality + +- **GIVEN** a raw legacy Codex row points to account A +- **AND** a bounded thread row points to account B +- **WHEN** the request is routed +- **THEN** the raw row remains hard ownership evidence and account A wins +- **AND** neither mapping is rewritten to reconcile the disagreement + +### Requirement: Unanchored process-session concurrency uses independent bridge lanes + +When multiple Responses requests share a process-level session header but +carry neither `previous_response_id` nor nonblank turn-state continuity, the +service MUST NOT queue an independent request behind an active response-create +gate. If the canonical bridge is still being created, reserved by another +request before submit, already has a visible request, or belongs to a different +model class, the service MUST create a server request-scoped bridge lane. The +lane identity MUST NOT depend on a client-controlled request ID. The fork MUST +leave the canonical bridge and its model metadata unchanged. + +When such requests carry nonblank `thread-id`, each thread MUST have a stable +canonical bridge identity derived from process and thread identity regardless +of `prompt_cache_key`; distinct threads MUST remain isolated even when they +execute sequentially, and repeated requests from one thread MUST retain one +identity. Requests without `thread-id` MUST retain the legacy session-header +identity, including the established explicit-prompt-cache composition. + +A pre-submit handoff reservation MUST protect its bridge from idle pruning and +capacity eviction, and any cancellation or error between lookup and visible +submission MUST release it. Owner forwarding MUST preserve whether a +session-header, thread-header, or internal-fork request was unanchored instead +of treating a proxy-generated downstream turn-state as an explicit client +anchor, but MUST NOT attach that v2-only state to prompt-cache or unrelated +affinity families. It MUST fail closed when a mixed-version hop cannot +authenticate required unanchored state. The v2 primary signature MUST bind +whether client-IP metadata was present, while the companion signature MUST bind +its value. When the canonical owner itself creates a fork for a forwarded +request, it MUST own that fork locally instead of re-hashing it into another +forwarding hop. Explicitly anchored owner forwards MUST retain the +legacy-compatible primary signature during rolling upgrades, and a receiving +instance MUST reject ambiguous delimiter-bearing legacy fields. Durable aliases +derived from the forked lane MUST retain hard owner and account continuity. If +durable ownership fencing rejects a stale owner's new alias, the stale owner +MUST remove the matching local alias without removing a newer local +generation's mapping. + +#### Scenario: sequential child agent does not reuse parent bridge history + +- **GIVEN** a parent and child Codex agent share one process session and `prompt_cache_key` +- **AND** each agent supplies its own stable `thread-id` +- **WHEN** the child starts after the parent's visible request has completed +- **THEN** the child uses a different bridge identity from the parent +- **AND** another request from that same child keeps the child's bridge identity + +#### Scenario: Background requests do not block behind a foreground turn + +- **GIVEN** a foreground request is active on a session-header or thread-header bridge +- **WHEN** two unanchored background requests arrive with the same canonical identity +- **THEN** each background request uses an independent response-create gate +- **AND** neither request waits for the foreground response to complete +- **AND** the foreground bridge's model metadata remains unchanged + +#### Scenario: Lookup-to-submit requests remain isolated + +- **GIVEN** an unanchored request has reserved an idle canonical bridge but has not yet made queued activity visible +- **WHEN** another unanchored request arrives with the same canonical identity and client request ID +- **THEN** the second request uses a distinct server-scoped bridge lane +- **AND** it does not reuse the reserved canonical bridge + +#### Scenario: Durable refresh publishes the handoff reservation + +- **GIVEN** an unanchored request reuses an idle durable canonical bridge +- **WHEN** refreshing the durable lease yields before lookup returns +- **THEN** the canonical bridge is already reserved for that request +- **AND** a concurrent unanchored request uses a distinct server-scoped lane + +#### Scenario: Cancelled pre-submit handoff does not strand a reservation + +- **GIVEN** an unanchored request is reusing an idle canonical bridge +- **WHEN** the request is cancelled after claiming the bridge but before queued activity becomes visible +- **THEN** the canonical bridge remains unreserved +- **AND** later requests are not forced onto fork lanes by the cancelled lookup + +#### Scenario: Payload preparation failure does not strand a reservation + +- **GIVEN** an unanchored request has reserved an idle canonical bridge +- **WHEN** anchor injection, trimming, or payload validation fails before submission +- **THEN** request-scope cleanup releases the reservation +- **AND** later requests may reuse the canonical bridge + +#### Scenario: Remote owner preserves unanchored concurrency + +- **GIVEN** an unanchored request is forwarded to the canonical bridge owner +- **AND** the proxy generated a downstream turn-state for response aliasing +- **WHEN** the owner receives the forwarded request while the canonical lane is active +- **THEN** the owner still treats the request as unanchored +- **AND** the request uses an independent bridge lane +- **AND** the pre-submit handoff remains reserved until submission becomes visible + +#### Scenario: Owner-side fork does not start a second forwarding hop + +- **GIVEN** an unanchored request has reached its canonical owner +- **AND** that owner creates an independent fork because the canonical lane is active +- **WHEN** rendezvous hashing the generated fork key would select another instance +- **THEN** the canonical owner creates and durably claims the fork locally +- **AND** the request is not rejected as a forwarding loop + +#### Scenario: Blank turn-state is not an anchor + +- **GIVEN** a request has process/thread identity and an empty or whitespace-only turn-state header +- **WHEN** the request is forwarded to its owner +- **THEN** the signed forwarding context marks the original request as unanchored +- **AND** the generated downstream turn-state does not collapse it onto the canonical gate + +#### Scenario: Forwarding downgrade fails closed + +- **GIVEN** an owner-forward request requires unanchored concurrency semantics +- **WHEN** the signed unanchored boolean is changed, removed, or repacked into affinity fields, or either instance only supports the legacy signature +- **THEN** the owner-forward hop fails closed +- **AND** the request is not attached to the shared canonical response-create gate + +#### Scenario: Anchored forwarding remains rolling-upgrade compatible + +- **GIVEN** an owner-forward request carries explicit previous-response or turn-state continuity +- **WHEN** the origin and owner run different bridge protocol versions +- **THEN** the primary signature remains valid under the legacy contract +- **AND** the anchored request can continue without weakening unanchored fail-closed behavior + +#### Scenario: Prompt-cache forwarding remains rolling-upgrade compatible + +- **GIVEN** an unanchored first-turn request uses a prompt-cache affinity lane +- **WHEN** that request is forwarded to its canonical owner +- **THEN** the origin does not attach session/thread-header unanchored v2 state +- **AND** an older owner may accept the legacy-compatible forwarding contract + +#### Scenario: Legacy session-header canonical lane proves its turn-state anchor + +- **GIVEN** a legacy-signed owner forward has no previous-response ID and its durable canonical key is still `session_header` +- **WHEN** its forwarded turn state is a registered durable alias for that exact canonical lane +- **THEN** the current owner accepts it as anchored continuity +- **AND** an unknown turn state or an alias for another canonical lane fails closed with `bridge_forward_upgrade_required` + +#### Scenario: Legacy proof precedes compact and bridge fallback branches + +- **GIVEN** a legacy-signed owner forward requires turn-state anchor proof +- **WHEN** the request contains a terminal compaction trigger or bypasses the websocket bridge +- **THEN** exact alias proof runs before compact, HTTP fallback, admission, or upstream work + +#### Scenario: Current origin proves a turn-state alias before legacy owner forwarding + +- **GIVEN** a current origin resolves a nonblank turn state only through a shared `session_header` durable lane +- **WHEN** that request would be forwarded to another owner with the legacy signature contract +- **THEN** the origin proves an exact turn-state alias row for that canonical lane before sending the owner request +- **AND** an unknown alias fails closed with `bridge_forward_upgrade_required` + +#### Scenario: Latest-state metadata is not proof of alias registration + +- **GIVEN** a durable session records a latest turn state but has no matching turn-state alias row +- **WHEN** that value is presented by a legacy-signed owner forward +- **THEN** the owner rejects it with `bridge_forward_upgrade_required` + +#### Scenario: Stale owners cannot register continuity aliases after takeover + +- **GIVEN** durable ownership advanced to a new owner epoch +- **WHEN** the stale owner attempts to register a turn-state or previous-response alias with its old epoch +- **THEN** alias registration writes nothing +- **AND** the stale owner removes the rejected value from its local alias index +- **AND** a newer local generation's mapping for the same value remains intact +- **AND** the stale value cannot satisfy legacy anchor proof + +#### Scenario: Ambiguous legacy signature fields fail closed + +- **GIVEN** a legacy owner-forward signature contains a delimiter in any signed header field +- **WHEN** field boundaries are repacked without changing the legacy joined byte string +- **THEN** a current owner rejects the forwarding context as invalid +- **AND** the repacked affinity kind cannot weaken hard continuity + +#### Scenario: V2 client-IP metadata cannot be removed or blanked + +- **GIVEN** an unanchored v2 owner-forward request carries signed client-IP metadata +- **WHEN** both client-IP headers are removed, the value is blanked, or the value is changed +- **THEN** the owner rejects the forwarding context as invalid +- **AND** a genuinely no-IP v2 request remains valid + +#### Scenario: Durable fork continuation remains owner-bound + +- **GIVEN** a forked lane has produced a durable turn-state or previous-response alias +- **WHEN** a later request resolves that alias on another instance +- **THEN** the request follows the hard owner-bound continuity path +- **AND** the original account binding is preserved + +#### Scenario: Explicit continuation is not split + +- **WHEN** a request carries `previous_response_id` or a turn-state header +- **THEN** the service keeps the request on the hard owner-bound continuity path +- **AND** it does not apply unanchored parallel-session isolation diff --git a/openspec/changes/scope-codex-affinity-by-thread/tasks.md b/openspec/changes/scope-codex-affinity-by-thread/tasks.md new file mode 100644 index 0000000000..0df723dd2d --- /dev/null +++ b/openspec/changes/scope-codex-affinity-by-thread/tasks.md @@ -0,0 +1,18 @@ +## 1. Identity and account locality + +- [x] 1.1 Parse process session and thread identity independently and derive source-separated opaque thread keys +- [x] 1.2 Route backend Responses and compact through bounded thread locality with process-preference seeding +- [x] 1.3 Preserve raw legacy Codex rows and all exact hard-owner precedence/conflict behavior + +## 2. Transport continuity + +- [x] 2.1 Scope direct WebSocket replay/tool continuity by thread and refresh active thread locality +- [x] 2.2 Scope HTTP bridge canonical lanes by thread while preserving exact-alias migration and forwarded-key behavior +- [x] 2.3 Route thread-goal operations from payload `threadId` + +## 3. Regression evidence + +- [x] 3.1 Cover identity parsing, source separation, process seeding, Responses/compact parity, and unchanged cache hints +- [x] 3.2 Cover sibling direct-WebSocket replay isolation and thread-goal account selection +- [x] 3.3 Cover sibling bridge isolation, exact legacy alias recovery, and no old-canonical fallback +- [x] 3.4 Run focused tests, Ruff, type checks, OpenSpec checks available in the checkout, and review the final diff diff --git a/tests/integration/test_proxy_sticky_sessions.py b/tests/integration/test_proxy_sticky_sessions.py index a3726b3b89..0639f4342c 100644 --- a/tests/integration/test_proxy_sticky_sessions.py +++ b/tests/integration/test_proxy_sticky_sessions.py @@ -494,6 +494,70 @@ async def fake_compact(payload, headers, access_token, account_id): assert stream_seen == ["acc_sid_a", "acc_sid_a"] +@pytest.mark.asyncio +async def test_backend_thread_rows_route_sibling_responses_and_compact_independently( + async_client, + monkeypatch, +): + from app.modules.proxy.affinity import _codex_backend_identity + from app.modules.proxy.sticky_repository import StickySessionsRepository + + await _set_routing_settings(async_client, sticky_threads_enabled=False) + account_a_id = await _import_account(async_client, "acc_thread_route_a", "thread-route-a@example.com") + account_b_id = await _import_account(async_client, "acc_thread_route_b", "thread-route-b@example.com") + process_session = "process-thread-route-shared" + root_headers = {"session-id": process_session, "thread-id": "thread-route-root"} + child_headers = {"session-id": process_session, "thread-id": "thread-route-child"} + root_key = _codex_backend_identity(root_headers).thread_selection_key + child_key = _codex_backend_identity(child_headers).thread_selection_key + assert root_key is not None + assert child_key is not None + + async with SessionLocal() as session: + repo = StickySessionsRepository(session) + await repo.upsert(root_key, account_a_id, kind=StickySessionKind.PROMPT_CACHE) + await repo.upsert(child_key, account_b_id, kind=StickySessionKind.PROMPT_CACHE) + + observed: list[tuple[str, str, str | None]] = [] + + async def fake_stream(payload, headers, access_token, account_id, **kwargs): + del headers, access_token, kwargs + observed.append(("responses", account_id, payload.prompt_cache_key)) + yield 'data: {"type":"response.completed","response":{"id":"resp_thread_route"}}\n\n' + + async def fake_compact(payload, headers, access_token, account_id): + del headers, access_token + observed.append(("compact", account_id, payload.prompt_cache_key)) + return OpenAIResponsePayload.model_validate({"output": []}) + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + monkeypatch.setattr(proxy_module, "core_compact_responses", fake_compact) + payload = { + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "prompt_cache_key": process_session, + } + + responses_response = await async_client.post( + "/backend-api/codex/responses", + json={**payload, "stream": True}, + headers=root_headers, + ) + compact_response = await async_client.post( + "/backend-api/codex/responses/compact", + json=payload, + headers=child_headers, + ) + + assert responses_response.status_code == 200 + assert compact_response.status_code == 200 + assert observed == [ + ("responses", "acc_thread_route_a", process_session), + ("compact", "acc_thread_route_b", process_session), + ] + + @pytest.mark.asyncio async def test_proxy_unregistered_turn_state_fails_closed_for_stream_and_compact( async_client, @@ -1665,6 +1729,80 @@ async def test_sticky_insert_if_absent_never_rebinds_existing_owner(db_setup): assert persisted_owner == "acc_live_immutable_a" +@pytest.mark.asyncio +async def test_seeded_sticky_upsert_is_atomic_and_preserves_first_seed_owner(db_setup, monkeypatch): + from sqlalchemy.exc import IntegrityError + from sqlalchemy.sql import Insert + + from app.modules.proxy.sticky_repository import StickySessionsRepository + + encryptor = TokenEncryptor() + async with SessionLocal() as session: + accounts = AccountsRepository(session) + for account_id in ("acc_seeded_a", "acc_seeded_b"): + await accounts.upsert( + Account( + id=account_id, + email=f"{account_id}@example.com", + plan_type="plus", + access_token_encrypted=encryptor.encrypt("access"), + refresh_token_encrypted=encryptor.encrypt("refresh"), + id_token_encrypted=encryptor.encrypt("id"), + last_refresh=utcnow(), + status=AccountStatus.ACTIVE, + deactivation_reason=None, + ) + ) + + seed_key = "seeded-process" + async with SessionLocal() as session: + repo = StickySessionsRepository(session) + first_thread = await repo.upsert_with_seed_if_absent( + "seeded-thread-a", + "acc_seeded_a", + kind=StickySessionKind.PROMPT_CACHE, + seed_key=seed_key, + seed_kind=StickySessionKind.CODEX_SESSION, + ) + second_thread = await repo.upsert_with_seed_if_absent( + "seeded-thread-b", + "acc_seeded_b", + kind=StickySessionKind.PROMPT_CACHE, + seed_key=seed_key, + seed_kind=StickySessionKind.CODEX_SESSION, + ) + + assert first_thread.account_id == "acc_seeded_a" + assert second_thread.account_id == "acc_seeded_b" + assert await repo.get_account_id(seed_key, kind=StickySessionKind.CODEX_SESSION) == "acc_seeded_a" + + original_build_upsert = repo._build_upsert_statement + + def _build_failing_upsert(key: str, account_id: str, kind: StickySessionKind) -> Insert: + del account_id + return original_build_upsert(key, "missing-account", kind) + + monkeypatch.setattr(repo, "_build_upsert_statement", _build_failing_upsert) + with pytest.raises(IntegrityError): + await repo.upsert_with_seed_if_absent( + "seeded-thread-failing", + "acc_seeded_a", + kind=StickySessionKind.PROMPT_CACHE, + seed_key="seeded-process-failing", + seed_kind=StickySessionKind.CODEX_SESSION, + ) + + # The repository rolls back both statements itself, so even a caller + # that catches the failure cannot accidentally commit the seed later. + assert ( + await repo.get_account_id( + "seeded-process-failing", + kind=StickySessionKind.CODEX_SESSION, + ) + is None + ) + + @pytest.mark.asyncio async def test_stale_expiry_cleanup_cannot_delete_fresh_rebound_owner(async_client) -> None: from sqlalchemy import select, update diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index d07334c14b..b9ae1055b9 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -29,7 +29,7 @@ from app.core.crypto import TokenEncryptor from app.db.models import Account, AccountStatus, StickySessionKind, UsageHistory from app.modules.api_keys.repository import ApiKeysRepository -from app.modules.proxy.affinity import _codex_session_selection_key +from app.modules.proxy.affinity import _codex_backend_identity, _codex_session_selection_key from app.modules.proxy.cap_partitioning import CapPartition from app.modules.proxy.load_balancer import LoadBalancer, RuntimeState, effective_account_concurrency_caps from app.modules.proxy.repo_bundle import ProxyRepositories @@ -224,6 +224,8 @@ def __init__(self) -> None: self.abandoned_keys: set[str] = set() self.deleted: list[tuple[str, StickySessionKind | None]] = [] self.upserts: list[tuple[str, str, StickySessionKind | None]] = [] + self.insert_if_absent_calls: list[tuple[str, str, StickySessionKind]] = [] + self.seeded_upserts: list[tuple[str, str, StickySessionKind, str, StickySessionKind]] = [] async def get_account_id(self, *args: Any, **kwargs: Any) -> str | None: lookup = await self.get_account_id_and_abandonment(*args, **kwargs) @@ -242,12 +244,44 @@ async def upsert(self, *args: Any, **kwargs: Any) -> Any: sticky_key = cast(str, args[0]) account_id = cast(str, args[1]) self.account_id = account_id + if self.account_ids_by_key is not None: + self.account_ids_by_key[sticky_key] = account_id self.upserts.append((sticky_key, account_id, kwargs.get("kind"))) return None + async def insert_if_absent( + self, + key: str, + account_id: str, + kind: StickySessionKind, + ) -> str: + self.insert_if_absent_calls.append((key, account_id, kind)) + if self.account_ids_by_key is None: + self.account_ids_by_key = {} + return self.account_ids_by_key.setdefault(key, account_id) + + async def upsert_with_seed_if_absent( + self, + key: str, + account_id: str, + *, + kind: StickySessionKind, + seed_key: str, + seed_kind: StickySessionKind, + ) -> None: + self.seeded_upserts.append((key, account_id, kind, seed_key, seed_kind)) + if self.account_ids_by_key is None: + self.account_ids_by_key = {} + self.account_ids_by_key.setdefault(seed_key, account_id) + self.account_ids_by_key[key] = account_id + self.account_id = account_id + self.upserts.append((key, account_id, kind)) + async def delete(self, *args: Any, **kwargs: Any) -> bool: sticky_key = cast(str, args[0]) self.deleted.append((sticky_key, kwargs.get("kind"))) + if self.account_ids_by_key is not None: + self.account_ids_by_key.pop(sticky_key, None) self.account_id = None return True @@ -259,13 +293,20 @@ async def restore_if_current( expected_account_id: str | None, restore_account_id: str | None, ) -> bool: - if self.account_id != expected_account_id: + current_account_id = ( + self.account_ids_by_key.get(key) if self.account_ids_by_key is not None else self.account_id + ) + if current_account_id != expected_account_id: return False if restore_account_id is None: self.deleted.append((key, kind)) + if self.account_ids_by_key is not None: + self.account_ids_by_key.pop(key, None) self.account_id = None return True self.upserts.append((key, restore_account_id, kind)) + if self.account_ids_by_key is not None: + self.account_ids_by_key[key] = restore_account_id self.account_id = restore_account_id return True @@ -2111,6 +2152,74 @@ async def test_sticky_probe_reservation_restores_affinity_after_repeated_commit_ await balancer.release_account_lease(selected.lease) +@pytest.mark.asyncio +async def test_provisional_recovery_probe_does_not_publish_process_seed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) + healthy = _make_account("acc-probe-seed-healthy") + probing = _make_account("acc-probe-seed-probing") + accounts_repo = _StubAccountsRepository([healthy, probing]) + usage_repo = _StubUsageRepository( + primary={ + healthy.id: _usage_row_with_percent( + 211, + healthy.id, + used_percent=30.0, + reset_at=now_epoch + 300, + ), + probing.id: _usage_row_with_percent( + 212, + probing.id, + used_percent=10.0, + reset_at=now_epoch + 300, + ), + }, + secondary={}, + ) + sticky_repo = _StubStickySessionsRepository() + sticky_repo.account_ids_by_key = {} + balancer = LoadBalancer(lambda: _repo_factory(accounts_repo, usage_repo, sticky_repo)) + balancer._runtime[probing.id] = RuntimeState( + health_tier=HEALTH_TIER_PROBING, + last_selected_at=0.0, + version=37, + health_version=13, + ) + monkeypatch.setattr(balancer, "_commit_due_probe_reservation_locked", lambda *args, **kwargs: False) + + selected = await balancer.select_account( + sticky_key="thread-after-probe-loss", + sticky_kind=StickySessionKind.PROMPT_CACHE, + sticky_source="thread_header", + legacy_sticky_key="process-after-probe-loss", + sticky_seed_key="process-seed-after-probe-loss", + sticky_seed_kind=StickySessionKind.CODEX_SESSION, + sticky_max_age_seconds=300, + routing_strategy="usage_weighted", + lease_kind="stream", + ) + + assert selected.account is not None + assert selected.account.id == healthy.id + assert selected.lease is not None + assert sticky_repo.account_ids_by_key == { + "process-seed-after-probe-loss": healthy.id, + "thread-after-probe-loss": healthy.id, + } + assert sticky_repo.seeded_upserts == [ + ( + "thread-after-probe-loss", + healthy.id, + StickySessionKind.PROMPT_CACHE, + "process-seed-after-probe-loss", + StickySessionKind.CODEX_SESSION, + ) + ] + + await balancer.release_account_lease(selected.lease) + + @pytest.mark.asyncio async def test_sticky_probe_reservation_restore_does_not_clobber_newer_owner( monkeypatch: pytest.MonkeyPatch, @@ -2764,6 +2873,131 @@ async def test_legacy_raw_session_mapping_wins_when_namespaced_row_also_exists() await balancer.release_account_lease(lease) +@pytest.mark.asyncio +async def test_new_codex_thread_is_seeded_from_process_preference_without_rewriting_process_row() -> None: + balancer, owner, alternate, sticky_repo = _make_cap_spillover_balancer("thread-process-seed") + assert alternate is not None + process_session = "process-seed" + process_key = _codex_session_selection_key(process_session) + thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "thread-new"} + ).thread_selection_key + assert thread_key is not None + sticky_repo.account_ids_by_key = {process_key: owner.id} + + selected = await balancer.select_account( + sticky_key=thread_key, + sticky_kind=StickySessionKind.PROMPT_CACHE, + sticky_source="thread_header", + legacy_sticky_key=process_session, + sticky_seed_key=process_key, + sticky_seed_kind=StickySessionKind.CODEX_SESSION, + sticky_max_age_seconds=300, + routing_strategy="usage_weighted", + ) + + assert selected.account is not None + assert selected.account.id == owner.id + assert sticky_repo.account_ids_by_key == { + process_key: owner.id, + thread_key: owner.id, + } + assert sticky_repo.deleted == [] + assert sticky_repo.upserts == [(thread_key, owner.id, StickySessionKind.PROMPT_CACHE)] + + +@pytest.mark.asyncio +async def test_first_codex_thread_initializes_process_preference_once_for_later_siblings() -> None: + balancer, owner, alternate, sticky_repo = _make_cap_spillover_balancer("thread-first-process") + assert alternate is not None + process_session = "process-first-thread" + process_key = _codex_session_selection_key(process_session) + first_thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "thread-first"} + ).thread_selection_key + sibling_thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "thread-sibling"} + ).thread_selection_key + assert first_thread_key is not None + assert sibling_thread_key is not None + sticky_repo.account_ids_by_key = {} + + first = await balancer.select_account( + sticky_key=first_thread_key, + sticky_kind=StickySessionKind.PROMPT_CACHE, + sticky_source="thread_header", + legacy_sticky_key=process_session, + sticky_seed_key=process_key, + sticky_seed_kind=StickySessionKind.CODEX_SESSION, + sticky_max_age_seconds=300, + routing_strategy="usage_weighted", + ) + assert first.account is not None + first_account_id = first.account.id + + sibling = await balancer.select_account( + sticky_key=sibling_thread_key, + sticky_kind=StickySessionKind.PROMPT_CACHE, + sticky_source="thread_header", + legacy_sticky_key=process_session, + sticky_seed_key=process_key, + sticky_seed_kind=StickySessionKind.CODEX_SESSION, + sticky_max_age_seconds=300, + routing_strategy="usage_weighted", + ) + + assert sibling.account is not None + assert sibling.account.id == first_account_id + assert sticky_repo.account_ids_by_key[process_key] == first_account_id + assert sticky_repo.insert_if_absent_calls == [] + assert sticky_repo.seeded_upserts == [ + ( + first_thread_key, + first_account_id, + StickySessionKind.PROMPT_CACHE, + process_key, + StickySessionKind.CODEX_SESSION, + ) + ] + assert sticky_repo.upserts == [ + (first_thread_key, first_account_id, StickySessionKind.PROMPT_CACHE), + (sibling_thread_key, first_account_id, StickySessionKind.PROMPT_CACHE), + ] + + +@pytest.mark.asyncio +async def test_legacy_raw_process_owner_wins_over_thread_locality() -> None: + balancer, owner, alternate, sticky_repo = _make_cap_spillover_balancer("thread-legacy-owner") + assert alternate is not None + process_session = "legacy-process-owner" + thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "thread-existing"} + ).thread_selection_key + assert thread_key is not None + sticky_repo.account_ids_by_key = { + process_session: owner.id, + thread_key: alternate.id, + } + + selected = await balancer.select_account( + sticky_key=thread_key, + sticky_kind=StickySessionKind.PROMPT_CACHE, + sticky_source="thread_header", + legacy_sticky_key=process_session, + sticky_max_age_seconds=300, + routing_strategy="usage_weighted", + ) + + assert selected.account is not None + assert selected.account.id == owner.id + assert sticky_repo.account_ids_by_key == { + process_session: owner.id, + thread_key: alternate.id, + } + assert sticky_repo.deleted == [] + assert sticky_repo.upserts == [] + + @pytest.mark.asyncio async def test_legacy_raw_owner_conflict_blocks_resolved_preferred_owner() -> None: balancer, owner, alternate, sticky_repo = _make_cap_spillover_balancer("legacy-preferred-conflict") diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 1ede9d331b..09cd10a510 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -39,6 +39,7 @@ 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, Base, HttpBridgeSessionState +from app.modules.proxy import affinity as proxy_affinity 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 @@ -6915,6 +6916,7 @@ async def fake_sleep(seconds: float) -> None: def test_http_bridge_session_key_infers_strength_from_affinity_kind() -> None: assert proxy_service._HTTPBridgeSessionKey("turn_state_header", "turn", None).strength == "hard" assert proxy_service._HTTPBridgeSessionKey("session_header", "session", None).strength == "hard" + assert proxy_service._HTTPBridgeSessionKey("thread_header", "thread", None).strength == "hard" assert proxy_service._HTTPBridgeSessionKey("prompt_cache", "cache", None).strength == "soft" assert proxy_service._HTTPBridgeSessionKey("request", "request", None).strength == "soft" @@ -6975,6 +6977,146 @@ def test_http_bridge_session_header_key_without_prompt_cache_key_stays_legacy_co assert key == proxy_service._HTTPBridgeSessionKey("session_header", "legacy-session", None) +def test_http_bridge_thread_keys_isolate_siblings_with_shared_process_cache_identity() -> None: + payload = proxy_service.ResponsesRequest.model_validate( + { + "model": "gpt-5.6-sol", + "instructions": "", + "input": [], + "prompt_cache_key": "process-shared", + } + ) + root_headers = {"session-id": "process-shared", "thread-id": "thread-root"} + child_headers = {"session-id": "process-shared", "thread-id": "thread-child"} + + root_key = proxy_service._make_http_bridge_session_key( + payload, + headers=root_headers, + affinity=proxy_service._AffinityPolicy(), + api_key=None, + request_id="request-root", + explicit_prompt_cache_key="process-shared", + ) + root_retry_key = proxy_service._make_http_bridge_session_key( + payload, + headers=root_headers, + affinity=proxy_service._AffinityPolicy(), + api_key=None, + request_id="request-root-retry", + explicit_prompt_cache_key="process-shared", + ) + child_key = proxy_service._make_http_bridge_session_key( + payload, + headers=child_headers, + affinity=proxy_service._AffinityPolicy(), + api_key=None, + request_id="request-child", + explicit_prompt_cache_key="process-shared", + ) + + assert root_key.affinity_kind == "thread_header" + assert root_key.strength == "hard" + assert root_key == root_retry_key + assert root_key != child_key + assert ( + http_bridge_helpers_module._make_http_bridge_session_header_fallback_key( + headers=root_headers, + api_key=None, + explicit_prompt_cache_key="process-shared", + ) + is None + ) + + +def test_http_bridge_same_thread_unanchored_concurrency_keeps_request_scoped_fork() -> None: + canonical = proxy_service._HTTPBridgeSessionKey("thread_header", "thread-canonical", None) + + fork = http_bridge_helpers_module._http_bridge_parallel_fork_key( + key=canonical, + session=None, + inflight_creation=True, + incoming_turn_state=None, + previous_response_id=None, + request_model="gpt-5.6-sol", + request_service_tier=None, + request_scope_id="second-request", + ) + + assert fork is not None + assert fork.affinity_kind == "internal_unanchored_parallel" + assert fork.affinity_key != canonical.affinity_key + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turn_state", [None, "turn-exact-thread-alias"]) +async def test_thread_bridge_durable_lookup_preserves_exact_response_alias_without_legacy_session_fallback( + monkeypatch: pytest.MonkeyPatch, + turn_state: str | None, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + payload = proxy_service.ResponsesRequest.model_validate( + { + "model": "gpt-5.6-sol", + "instructions": "", + "input": [], + "prompt_cache_key": "process-shared", + "previous_response_id": "resp-exact-legacy-alias", + } + ) + lookup = AsyncMock( + side_effect=ProxyResponseError( + 409, + openai_error("stop_after_lookup", "stop after durable lookup"), + ) + ) + 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(service._durable_bridge, "lookup_request_targets", lookup) + + headers = {"session-id": "process-shared", "thread-id": "thread-child"} + if turn_state is not None: + headers["x-codex-turn-state"] = turn_state + stream = service._stream_via_http_bridge( + payload, + headers=headers, + codex_session_affinity=True, + propagate_http_errors=False, + 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, + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await anext(stream) + + assert exc_info.value.payload["error"]["code"] == "stop_after_lookup" + assert lookup.await_args is not None + lookup_kwargs = lookup.await_args.kwargs + assert lookup_kwargs["session_key_kind"] == ("thread_header" if turn_state is None else "turn_state_header") + assert lookup_kwargs["previous_response_id"] == "resp-exact-legacy-alias" + assert lookup_kwargs["session_header"] is None + + def test_http_bridge_owner_check_required_keeps_prompt_cache_soft() -> None: key = proxy_service._HTTPBridgeSessionKey("prompt_cache", "cache", None) @@ -9178,6 +9320,30 @@ def test_make_http_bridge_session_key_prefers_signed_forwarded_affinity_over_gen assert key.strength == "hard" +def test_make_http_bridge_session_key_keeps_forwarded_thread_affinity_verbatim() -> None: + payload = proxy_service.ResponsesRequest.model_validate({"model": "gpt-5.4", "instructions": "hi", "input": "hi"}) + + key = proxy_service._make_http_bridge_session_key( + payload, + headers={ + "session-id": "process-local", + "thread-id": "thread-local", + "x-codex-bridge-affinity-kind": "thread_header", + "x-codex-bridge-affinity-key": "opaque-key-from-owner", + }, + affinity=proxy_service._AffinityPolicy(), + api_key=None, + request_id="req-forwarded-thread", + allow_forwarded_affinity_headers=True, + ) + + assert key == proxy_service._HTTPBridgeSessionKey( + "thread_header", + "opaque-key-from-owner", + None, + ) + + def test_make_http_bridge_session_key_keeps_forwarded_parallel_lane_hard() -> None: payload = proxy_service.ResponsesRequest.model_validate({"model": "gpt-5.4", "instructions": "hi", "input": "hi"}) @@ -15364,6 +15530,129 @@ async def fake_create_http_bridge_session( assert captured["key"] == fallback_key +@pytest.mark.asyncio +async def test_missing_turn_alias_with_thread_uses_thread_canonical_not_legacy_process_lane( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + headers = { + "session-id": "process-shared", + "thread-id": "thread-current", + "x-codex-turn-state": "http_turn_missing_thread_alias", + } + requested_key = proxy_service._HTTPBridgeSessionKey( + "turn_state_header", + "http_turn_missing_thread_alias", + None, + ) + thread_key_value = proxy_affinity._codex_backend_identity(headers).thread_selection_key + assert thread_key_value is not None + thread_key = proxy_service._HTTPBridgeSessionKey("thread_header", thread_key_value, None) + legacy_process_key = proxy_service._HTTPBridgeSessionKey("session_header", "process-shared", None) + service._http_bridge_sessions[legacy_process_key] = _make_bridge_session( + key=legacy_process_key, + key_value="legacy-sibling", + ) + captured: dict[str, object] = {} + + async def create_session(key: proxy_service._HTTPBridgeSessionKey, **kwargs: object): + del kwargs + captured["key"] = key + return _make_bridge_session(key=key, key_value="current-thread") + + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_create_http_bridge_session", create_session) + monkeypatch.setattr(service, "_claim_durable_http_bridge_session", AsyncMock()) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr(proxy_service, "_http_bridge_owner_instance", AsyncMock(return_value="instance-a")) + monkeypatch.setattr( + proxy_service, + "_active_http_bridge_instance_ring", + AsyncMock(return_value=("instance-a", ["instance-a"])), + ) + + resolved = await service._get_or_create_http_bridge_session( + requested_key, + headers=headers, + affinity=proxy_service._AffinityPolicy( + key=thread_key_value, + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + codex_session_source="thread_header", + ), + api_key=None, + request_model="gpt-5.6-sol", + idle_ttl_seconds=120.0, + max_sessions=8, + previous_response_id="resp-missing-thread-alias", + ) + + assert resolved.key == thread_key + assert captured["key"] == thread_key + assert resolved is not service._http_bridge_sessions[legacy_process_key] + + +@pytest.mark.asyncio +async def test_legacy_forward_missing_turn_alias_with_thread_avoids_process_lane( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + headers = { + "session-id": "process-shared", + "thread-id": "thread-forwarded", + "x-codex-turn-state": "http_turn_missing_forwarded_alias", + } + requested_key = proxy_service._HTTPBridgeSessionKey( + "turn_state_header", + "http_turn_missing_forwarded_alias", + None, + ) + thread_key_value = proxy_affinity._codex_backend_identity(headers).thread_selection_key + assert thread_key_value is not None + thread_key = proxy_service._HTTPBridgeSessionKey("thread_header", thread_key_value, None) + legacy_process_key = proxy_service._HTTPBridgeSessionKey("session_header", "process-shared", None) + service._http_bridge_sessions[legacy_process_key] = _make_bridge_session( + key=legacy_process_key, + key_value="legacy-sibling", + ) + captured: dict[str, object] = {} + + async def create_session(key: proxy_service._HTTPBridgeSessionKey, **kwargs: object): + del kwargs + captured["key"] = key + return _make_bridge_session(key=key, key_value="forwarded-thread") + + monkeypatch.setattr(service, "_prune_http_bridge_sessions_locked", Mock(return_value=[])) + monkeypatch.setattr(service, "_create_http_bridge_session", create_session) + monkeypatch.setattr(service, "_claim_durable_http_bridge_session", AsyncMock()) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr(proxy_service, "_http_bridge_owner_instance", AsyncMock(return_value="instance-a")) + monkeypatch.setattr( + proxy_service, + "_active_http_bridge_instance_ring", + AsyncMock(return_value=("instance-a", ["instance-a"])), + ) + + resolved = await service._get_or_create_http_bridge_session( + requested_key, + headers=headers, + affinity=proxy_service._AffinityPolicy( + key=thread_key_value, + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + codex_session_source="thread_header", + ), + api_key=None, + request_model="gpt-5.6-sol", + idle_ttl_seconds=120.0, + max_sessions=8, + previous_response_id="resp-missing-forwarded-alias", + forwarded_request=True, + ) + + assert resolved.key == thread_key + assert captured["key"] == thread_key + assert resolved is not service._http_bridge_sessions[legacy_process_key] + + @pytest.mark.asyncio async def test_get_or_create_http_bridge_session_preserves_durable_canonical_prompt_cache_key( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 9ed5b4d98f..85f1118fe6 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -5703,6 +5703,8 @@ async def test_select_codex_control_account_without_budget_uses_balancer(monkeyp reallocate_sticky=False, sticky_source=None, legacy_sticky_key=None, + sticky_seed_key=None, + sticky_seed_kind=None, sticky_max_age_seconds=123, prefer_earlier_reset_window="primary", routing_strategy="usage_weighted", @@ -5835,6 +5837,47 @@ async def thread_goal_request(*_args: object, **_kwargs: object) -> dict[str, Js assert request_logs.calls[0]["conversation_id"] == "conv-thread-goal" +@pytest.mark.asyncio +async def test_thread_goal_request_routes_from_payload_thread_identity(monkeypatch: pytest.MonkeyPatch) -> None: + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + account = _make_account("acc-thread-goal-payload") + selection_kwargs: list[dict[str, object]] = [] + + async def select_account(_deadline: float, **kwargs: object) -> AccountSelection: + selection_kwargs.append(kwargs) + return AccountSelection(account=account, error_message=None) + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(return_value=account)) + monkeypatch.setattr(proxy_service, "core_thread_goal_request", AsyncMock(return_value={"ok": True})) + + payload = {"threadId": "thread-from-payload", "goal": "finish"} + headers = { + "session-id": "process-shared", + # The payload is the endpoint's exact subject and must not be replaced + # by a broader or stale generic request header. + "thread-id": "thread-from-header", + "user-agent": "codex/1.2", + } + response = await service.thread_goal_request("set", payload, headers) + + assert response == {"ok": True} + affinity = cast(proxy_service._AffinityPolicy, selection_kwargs[0]["affinity_policy"]) + expected_key = proxy_affinity._codex_backend_identity( + headers, + thread_id="thread-from-payload", + ).thread_selection_key + assert affinity.selection_key == expected_key + assert affinity.codex_session_source == "thread_header" + assert affinity.legacy_selection_key == "process-shared" + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert request_logs.calls[0]["conversation_id"] == "thread-from-header" + + @pytest.mark.asyncio async def test_thread_goal_401_failover_preserves_dashboard_reset_window(monkeypatch): settings = _make_proxy_settings() @@ -10401,6 +10444,133 @@ def test_sticky_key_for_compact_request_prefers_codex_session_affinity(): assert policy.max_age_seconds is None +def test_backend_codex_thread_affinity_is_shared_by_responses_and_compact_without_rewriting_cache_hint() -> None: + headers_by_thread = { + "root": {"session-id": "process-shared", "thread-id": "thread-root"}, + "child": {"session-id": "process-shared", "thread-id": "thread-child"}, + } + keys_by_surface: dict[str, dict[str, str]] = {"responses": {}, "compact": {}} + + for surface, payload in ( + ( + "responses", + ResponsesRequest.model_validate( + { + "model": "gpt-5.6-sol", + "instructions": "hi", + "input": [], + "prompt_cache_key": "process-shared", + } + ), + ), + ( + "compact", + ResponsesCompactRequest.model_validate( + { + "model": "gpt-5.6-sol", + "instructions": "hi", + "input": [], + "prompt_cache_key": "process-shared", + } + ), + ), + ): + for thread_name, headers in headers_by_thread.items(): + request_payload = payload.model_copy() + if isinstance(request_payload, ResponsesRequest): + policy = proxy_service._sticky_key_for_responses_request( + request_payload, + headers=headers, + codex_session_affinity=True, + openai_cache_affinity=True, + openai_cache_affinity_max_age_seconds=300, + sticky_threads_enabled=False, + ) + else: + policy = proxy_service._sticky_key_for_compact_request( + request_payload, + headers=headers, + codex_session_affinity=True, + openai_cache_affinity=True, + openai_cache_affinity_max_age_seconds=300, + sticky_threads_enabled=False, + ) + + assert policy.key is not None + keys_by_surface[surface][thread_name] = policy.key + assert policy.kind == StickySessionKind.PROMPT_CACHE + assert policy.codex_session_source == "thread_header" + assert policy.legacy_selection_key == "process-shared" + assert policy.seed_selection_key == proxy_affinity._codex_session_selection_key("process-shared") + assert policy.seed_selection_kind == StickySessionKind.CODEX_SESSION + assert policy.max_age_seconds == 300 + assert request_payload.prompt_cache_key == "process-shared" + + assert keys_by_surface["responses"] == keys_by_surface["compact"] + assert keys_by_surface["responses"]["root"] != keys_by_surface["responses"]["child"] + + +def test_codex_thread_identity_is_source_separated_and_supports_thread_only_clients() -> None: + process_thread = proxy_affinity._codex_backend_identity( + {"session-id": "same-value", "thread-id": "same-value"} + ).thread_selection_key + thread_only = proxy_affinity._codex_backend_identity({"thread-id": "same-value"}).thread_selection_key + + assert process_thread is not None + assert thread_only is not None + assert process_thread != thread_only + assert process_thread.startswith("\n") + assert thread_only.startswith("\n") + assert "same-value" not in process_thread + assert "same-value" not in thread_only + + +def test_codex_thread_identity_length_frames_client_values() -> None: + left = proxy_affinity._codex_backend_identity( + { + "session-id": "process-a\0thread\0thread-b", + "thread-id": "thread-c", + } + ).thread_selection_key + right = proxy_affinity._codex_backend_identity( + { + "session-id": "process-a", + "thread-id": "thread-b\0thread\0thread-c", + } + ).thread_selection_key + + assert left is not None + assert right is not None + assert left != right + + +def test_codex_thread_identity_normalizes_header_names_and_blank_values() -> None: + identity = proxy_affinity._codex_backend_identity( + { + "SESSION-ID": " process-mixed-case ", + "THREAD-ID": " thread-mixed-case ", + } + ) + expected = proxy_affinity._codex_backend_identity( + { + "session-id": "process-mixed-case", + "thread-id": "thread-mixed-case", + } + ) + blank = proxy_affinity._codex_backend_identity( + { + "session-id": " ", + "thread-id": "\t", + } + ) + + assert identity == expected + assert identity.thread_selection_key is not None + assert blank.process_session is None + assert blank.thread_id is None + assert blank.thread_selection_key is None + + @pytest.mark.parametrize("codex_session_affinity", [False, True]) def test_sticky_key_for_compact_request_prefers_turn_state_over_session_and_cache( codex_session_affinity: bool, @@ -21528,6 +21698,147 @@ def test_websocket_continuity_state_reuses_codex_session_scope(): assert unscoped is not first +def test_websocket_continuity_state_isolates_sibling_codex_threads() -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + root_headers = {"session-id": "process-shared", "thread-id": "thread-root"} + child_headers = {"session-id": "process-shared", "thread-id": "thread-child"} + + root = service._websocket_continuity_state_for_request( + root_headers, + api_key=None, + codex_session_affinity=True, + ) + root.last_completed_response_id = "resp-root" + root.last_pending_tool_call_types["call-root"] = "function_call" + + child = service._websocket_continuity_state_for_request( + child_headers, + api_key=None, + codex_session_affinity=True, + ) + root_reconnect = service._websocket_continuity_state_for_request( + root_headers, + api_key=None, + codex_session_affinity=True, + ) + + assert child is not root + assert child.last_completed_response_id is None + assert child.last_pending_tool_call_types == {} + assert root_reconnect is root + assert root_reconnect.last_completed_response_id == "resp-root" + assert root_reconnect.last_pending_tool_call_types == {"call-root": "function_call"} + + +def test_websocket_explicit_turn_state_precedes_broader_thread_state() -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + thread_headers = {"session-id": "process-shared", "thread-id": "thread-exact"} + explicit_turn_state = "turn_client_exact" + + broader_thread_state = service._websocket_continuity_state_for_request( + thread_headers, + api_key=None, + codex_session_affinity=True, + ) + broader_thread_state.last_completed_response_id = "resp-thread-broad" + exact_turn_state = service._websocket_continuity_state_for_request( + {"x-codex-turn-state": explicit_turn_state}, + api_key=None, + codex_session_affinity=True, + ) + exact_turn_state.last_completed_response_id = "resp-turn-exact" + + resolved = service._websocket_continuity_state_for_request( + {**thread_headers, "x-codex-turn-state": explicit_turn_state}, + api_key=None, + codex_session_affinity=True, + ) + thread_reconnect = service._websocket_continuity_state_for_request( + thread_headers, + api_key=None, + codex_session_affinity=True, + ) + + assert resolved is exact_turn_state + assert resolved is not broader_thread_state + assert resolved.last_completed_response_id == "resp-turn-exact" + assert thread_reconnect is exact_turn_state + + +def test_websocket_unknown_explicit_turn_state_does_not_reuse_broader_thread_state() -> None: + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + thread_headers = {"session-id": "process-shared", "thread-id": "thread-hard-turn"} + + broader_thread_state = service._websocket_continuity_state_for_request( + thread_headers, + api_key=None, + codex_session_affinity=True, + ) + broader_thread_state.last_completed_response_id = "resp-thread-broad" + + unknown_turn_state = service._websocket_continuity_state_for_request( + {**thread_headers, "x-codex-turn-state": "turn_client_unknown"}, + api_key=None, + codex_session_affinity=True, + ) + thread_reconnect = service._websocket_continuity_state_for_request( + thread_headers, + api_key=None, + codex_session_affinity=True, + ) + + assert unknown_turn_state is not broader_thread_state + assert unknown_turn_state.last_completed_response_id is None + assert thread_reconnect is broader_thread_state + assert thread_reconnect.last_completed_response_id == "resp-thread-broad" + + +@pytest.mark.asyncio +async def test_active_websocket_refreshes_only_its_bounded_thread_affinity() -> None: + sticky_sessions = AsyncMock(spec=StickySessionsRepository) + + class _ThreadAffinityRepoContext: + async def __aenter__(self) -> SimpleNamespace: + return SimpleNamespace(sticky_sessions=sticky_sessions) + + async def __aexit__(self, exc_type: object, exc: object, tb: object) -> bool: + del exc_type, exc, tb + return False + + service = proxy_service.ProxyService(cast(Any, lambda: _ThreadAffinityRepoContext())) + thread_key = proxy_affinity._codex_backend_identity( + {"session-id": "process-active", "thread-id": "thread-active"} + ).thread_selection_key + assert thread_key is not None + request_state = proxy_service._WebSocketRequestState( + request_id="req-active-thread-touch", + model="gpt-5.6-sol", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + affinity_policy=proxy_service._AffinityPolicy( + key=thread_key, + kind=StickySessionKind.PROMPT_CACHE, + max_age_seconds=300, + codex_session_source="thread_header", + legacy_codex_session_key="process-active", + seed_selection_key=proxy_affinity._codex_session_selection_key("process-active"), + seed_selection_kind=StickySessionKind.CODEX_SESSION, + ), + thread_affinity_last_touch_at=0.0, + ) + account = _make_account("acc-active-thread-touch") + + await service._touch_active_websocket_thread_affinity(request_state, account) + + sticky_sessions.upsert.assert_awaited_once_with( + thread_key, + account.id, + kind=StickySessionKind.PROMPT_CACHE, + ) + + def test_websocket_continuity_state_seeds_generated_turn_state_alias(): service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) generated_turn_state = "turn_0123456789abcdef0123456789abcdef" @@ -37436,6 +37747,47 @@ async def test_select_account_with_budget_reconciles_sticky_mapping_for_preferre assert select_account.await_args.kwargs["legacy_sticky_key"] is None +@pytest.mark.asyncio +async def test_select_account_with_budget_keeps_thread_seed_for_first_exact_owner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + owner = _make_account("acc-thread-first-exact-owner") + select_account = AsyncMock(return_value=AccountSelection(account=owner, error_message=None)) + process_key = proxy_affinity._codex_session_selection_key("process-first-exact-owner") + thread_policy = proxy_service._AffinityPolicy( + key="thread-first-exact-owner", + kind=proxy_service.StickySessionKind.PROMPT_CACHE, + codex_session_source="thread_header", + legacy_codex_session_key="process-first-exact-owner", + seed_selection_key=process_key, + seed_selection_kind=proxy_service.StickySessionKind.CODEX_SESSION, + ) + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(proxy_service, "_remaining_budget_seconds", lambda _deadline: 10.0) + + await service._select_account_with_budget( + deadline=123.0, + request_id="req-thread-first-exact-owner", + kind="stream", + request_stage="first_turn", + **thread_policy.selection_kwargs(), + preferred_account_id=owner.id, + preferred_account_is_continuity_owner=True, + lease_kind="stream", + ) + + select_account.assert_awaited_once() + assert select_account.await_args is not None + assert select_account.await_args.kwargs["required_account_id"] == owner.id + assert select_account.await_args.kwargs["sticky_key"] == thread_policy.selection_key + assert select_account.await_args.kwargs["sticky_seed_key"] == process_key + assert select_account.await_args.kwargs["sticky_seed_kind"] == proxy_service.StickySessionKind.CODEX_SESSION + + @pytest.mark.asyncio async def test_select_account_with_budget_preserves_conversation_check_for_preferred_owner( monkeypatch: pytest.MonkeyPatch, From 8d274c8e1f373e631768994a1de134398b8a7949 Mon Sep 17 00:00:00 2001 From: Darafei Praliaskouski Date: Fri, 14 Aug 2026 23:56:48 +0400 Subject: [PATCH 2/2] fix(proxy): extract bridge thread fallback helpers --- .../proxy/_service/http_bridge/helpers.py | 33 +++++++++++++++++++ .../proxy/_service/http_bridge/mixin.py | 29 ++++------------ 2 files changed, 40 insertions(+), 22 deletions(-) diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index 3571746338..7e4967e6e7 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -1638,6 +1638,39 @@ def _make_http_bridge_session_header_fallback_key( ) +def _turn_keys( + headers: Mapping[str, str], + api_key: ApiKeyData | None, + requested_key: _HTTPBridgeSessionKey, + fallback_key: _HTTPBridgeSessionKey | None, +) -> tuple[str | None, _HTTPBridgeSessionKey | None]: + thread_key = _codex_backend_identity(headers).thread_selection_key + thread_fallback_key = ( + _HTTPBridgeSessionKey("thread_header", thread_key, api_key.id if api_key is not None else None) + if thread_key is not None + else None + ) + incoming_session_key = None if thread_fallback_key is not None else _sticky_key_from_session_header(headers) + initial_session_key = ( + fallback_key + or thread_fallback_key + or (requested_key if requested_key.affinity_kind == "session_header" else None) + ) + return incoming_session_key, initial_session_key + + +def _alias_fallback_key( + incoming_session_key: str | None, + initial_session_key: _HTTPBridgeSessionKey | None, + api_key_id: str | None, +) -> _HTTPBridgeSessionKey | None: + if initial_session_key is not None: + return initial_session_key + if incoming_session_key is None: + return None + return _HTTPBridgeSessionKey("session_header", incoming_session_key, api_key_id) + + async def _http_bridge_should_wait_for_registration( self, key: _HTTPBridgeSessionKey, diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index 6dd9f39552..3106270813 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -70,6 +70,7 @@ _HTTP_BRIDGE_BACKGROUND_CLOSE_TIMEOUT_SECONDS, _HTTP_BRIDGE_INFLIGHT_STARTED_AT_ATTR, _active_http_bridge_instance_ring, + _alias_fallback_key, _durable_bridge_lookup_active_owner, _durable_bridge_lookup_allows_local_reuse, _forwarded_http_bridge_session_key, @@ -121,6 +122,7 @@ _require_http_bridge_bound_account_not_excluded, _reserve_http_bridge_unanchored_handoff, _settle_failed_http_bridge_creation, + _turn_keys, ) from app.modules.proxy._service.http_bridge.helpers import ( _close_http_bridge_session as _helpers_close_http_bridge_session, @@ -204,9 +206,7 @@ ) from app.modules.proxy.affinity import ( _AffinityPolicy, - _codex_backend_identity, _extract_model_class, - _sticky_key_from_session_header, _sticky_key_from_turn_state_header, ) from app.modules.proxy.continuity import ( @@ -430,19 +430,7 @@ async def _get_or_create_http_bridge_session( request_scope_id = ensure_request_scope_id() api_key_id = api_key.id if api_key is not None else None incoming_turn_state = _sticky_key_from_turn_state_header(headers) - thread_selection_key = _codex_backend_identity(headers).thread_selection_key - thread_fallback_key = None - if thread_selection_key is not None: - thread_fallback_key = _HTTPBridgeSessionKey("thread_header", thread_selection_key, api_key_id) - # Exact aliases may fall back to this thread, never the raw process lane - # where mixed-version replicas may have stored a sibling. V2 forwards use - # their signed key; this derived fallback protects legacy raw headers. - incoming_session_key = None if thread_fallback_key is not None else _sticky_key_from_session_header(headers) - initial_session_key = ( - session_header_fallback_key - or thread_fallback_key - or (key if key.affinity_kind == "session_header" else None) - ) + incoming_session_key, initial_session_key = _turn_keys(headers, api_key, key, session_header_fallback_key) original_request_unanchored = _http_bridge_request_needs_unanchored_handoff( key, incoming_turn_state, previous_response_id, forwarded_request, forwarded_original_request_unanchored ) @@ -674,11 +662,10 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: key=key.affinity_key, ): key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id) - elif initial_session_key is not None: - key = initial_session_key - used_session_header_fallback = True - elif incoming_session_key is not None: - key = _HTTPBridgeSessionKey("session_header", incoming_session_key, api_key_id) + elif ( + fallback_key := _alias_fallback_key(incoming_session_key, initial_session_key, api_key_id) + ) is not None: + key = fallback_key used_session_header_fallback = True else: key = _HTTPBridgeSessionKey("turn_state_header", incoming_turn_state, api_key_id) @@ -2408,8 +2395,6 @@ async def abort_selected_handoff() -> None: session.last_completed_input_prefix_fingerprint = None session.last_pending_tool_calls.clear() session.affinity = selection_affinity or session.affinity - # Clearing stale response/turn aliases makes an account move - # safe; it does not make the canonical thread lane soft. session.codex_session = session.key.affinity_kind == "thread_header" session.upstream_turn_state = None session.downstream_turn_state = None