Skip to content

fix(proxy): scope backend Codex affinity by thread identity #1703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Komzpa merged 2 commits into from
Aug 15, 2026
Merged

fix(proxy): scope backend Codex affinity by thread identity #1703

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
26b1134
fix(proxy): scope backend Codex affinity by thread identity
leventov Aug 12, 2026
8d274c8
fix(proxy): extract bridge thread fallback helpers
Komzpa Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 74 additions & 5 deletions app/modules/proxy/_load_balancer/sticky_selection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)


Expand Down
6 changes: 6 additions & 0 deletions app/modules/proxy/_service/codex_control.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down
15 changes: 14 additions & 1 deletion app/modules/proxy/_service/compact.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
77 changes: 67 additions & 10 deletions app/modules/proxy/_service/http_bridge/helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -1617,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,
Expand Down Expand Up @@ -1708,7 +1762,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
)
Expand Down Expand Up @@ -2398,7 +2452,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
Expand Down Expand Up @@ -2604,7 +2661,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
Expand Down
Loading
Loading