From 028476bbae23599d06456cde29e390245e56a432 Mon Sep 17 00:00:00 2001 From: mastertyko <11311479+mastertyko@users.noreply.github.com> Date: Sat, 15 Aug 2026 13:42:23 +0200 Subject: [PATCH 1/5] fix(proxy): abandon unavailable owner on thread-scoped goal restart Current Codex sends thread-id with the process session, so affinity classified the restart as thread_header and never retired the raw legacy owner. Grant the one-shot abandonment flag when a process session is present, allow retirement CAS from thread_header requests, and consult the raw row as session_header interpretation so later thread turns stay on the replacement. --- .../proxy/_load_balancer/sticky_selection.py | 2 +- app/modules/proxy/affinity.py | 13 +++-- app/modules/proxy/load_balancer.py | 5 +- .../.openspec.yaml | 2 + .../context.md | 24 ++++++++ .../design.md | 56 +++++++++++++++++++ .../proposal.md | 43 ++++++++++++++ .../specs/sticky-session-operations/spec.md | 43 ++++++++++++++ .../tasks.md | 23 ++++++++ .../sticky-session-operations/context.md | 2 +- .../specs/sticky-session-operations/spec.md | 28 ++++++++++ tests/unit/test_load_balancer_concurrency.py | 48 ++++++++++++++++ tests/unit/test_proxy_utils.py | 44 +++++++++++++++ 13 files changed, 326 insertions(+), 7 deletions(-) create mode 100644 openspec/changes/goal-restart-thread-header-abandonment/.openspec.yaml create mode 100644 openspec/changes/goal-restart-thread-header-abandonment/context.md create mode 100644 openspec/changes/goal-restart-thread-header-abandonment/design.md create mode 100644 openspec/changes/goal-restart-thread-header-abandonment/proposal.md create mode 100644 openspec/changes/goal-restart-thread-header-abandonment/specs/sticky-session-operations/spec.md create mode 100644 openspec/changes/goal-restart-thread-header-abandonment/tasks.md diff --git a/app/modules/proxy/_load_balancer/sticky_selection.py b/app/modules/proxy/_load_balancer/sticky_selection.py index 981d65403f..1dc29cbf71 100644 --- a/app/modules/proxy/_load_balancer/sticky_selection.py +++ b/app/modules/proxy/_load_balancer/sticky_selection.py @@ -511,7 +511,7 @@ def _direct_error( abandon_unavailable_legacy_owner and hard_sticky and sticky_existing_is_legacy - and sticky_source == "session_header" + and sticky_source in {"session_header", "thread_header"} and legacy_sticky_key is not None and isinstance(sticky_existing_account_id, str) and legacy_owner_in_effective_policy_scope diff --git a/app/modules/proxy/affinity.py b/app/modules/proxy/affinity.py index 4dd41a114d..a8e5d25a36 100644 --- a/app/modules/proxy/affinity.py +++ b/app/modules/proxy/affinity.py @@ -735,10 +735,15 @@ def _sticky_key_for_responses_request( else: policy = _AffinityPolicy() if ( - # Only typed process-session provenance can represent the legacy row - # this escape hatch targets. An explicit turn-state header stays hard - # even when a client includes the same goal marker. - policy.codex_session_source == "session_header" + # The raw row this escape hatch retires is the process-session key. + # Current Codex also sends thread-id, so locality source is often + # thread_header; that must not hide the process-session exception. + # An explicit turn-state header stays hard even with the same marker. + policy.codex_session_source in {"session_header", "thread_header"} + and ( + policy.codex_session_source == "session_header" + or _codex_backend_identity(headers).process_session is not None + ) and _request_allows_unavailable_legacy_owner_abandonment(payload) ): policy = replace(policy, abandon_unavailable_legacy_owner=True) diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 0bc5717dd5..20c081adac 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -731,7 +731,10 @@ async def load_selection_inputs() -> _SelectionInputs: # Raw rows may be historical turn-state ownership. The # bounded thread TTL must never age out that hard evidence. max_age_seconds=None, - continuity_source=sticky_source, + # This key is the process-session compatibility row. + # Thread-header requests still consult it as that row, + # so session_header-scoped abandonment must hide it. + continuity_source="session_header", ) legacy_existing_account_id = legacy_owner_lookup.account_id abandoned_account_id = legacy_owner_lookup.abandoned_account_id diff --git a/openspec/changes/goal-restart-thread-header-abandonment/.openspec.yaml b/openspec/changes/goal-restart-thread-header-abandonment/.openspec.yaml new file mode 100644 index 0000000000..0c73c8f54e --- /dev/null +++ b/openspec/changes/goal-restart-thread-header-abandonment/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-08-15 diff --git a/openspec/changes/goal-restart-thread-header-abandonment/context.md b/openspec/changes/goal-restart-thread-header-abandonment/context.md new file mode 100644 index 0000000000..ba6c52edab --- /dev/null +++ b/openspec/changes/goal-restart-thread-header-abandonment/context.md @@ -0,0 +1,24 @@ +## Purpose + +Close the `#1703` × `#1680` composition hole: current Codex always +sends `thread-id`, so the merged goal-restart recovery never fires. + +## Decision + +Abandonment stays a `session_header` *interpretation* of the raw +process-session key. Request locality may be `thread_header`. Explicit +`turn_state` is unchanged. + +## Failure modes + +- Incremental or file-pinned restarts must still fail closed on the + required owner. +- After retirement, a later thread-id turn must not revive the raw + row as hard ownership. + +## Example + +Process session `sid` maps to quota-exceeded account A. Codex resends +an account-neutral goal body with `session-id: sid` and +`thread-id: t1`. Selection retires `sid` for `session_header`, routes +to B, and later `t1` turns stay on B. diff --git a/openspec/changes/goal-restart-thread-header-abandonment/design.md b/openspec/changes/goal-restart-thread-header-abandonment/design.md new file mode 100644 index 0000000000..3f2988efc6 --- /dev/null +++ b/openspec/changes/goal-restart-thread-header-abandonment/design.md @@ -0,0 +1,56 @@ +## Context + +`#1679` / `#1680` added a proof-gated exception that retires an +unavailable raw `codex_session` owner for `session_header` +interpretation. `#1703` then made `thread-id` the winning locality +source for current Codex. The two compose incorrectly: the flag and +CAS both require `sticky_source == "session_header"`, which current +Codex never is. + +The raw compatibility row is the process-session key. Looking it up +with `continuity_source=thread_header` treats a `session_header` +tombstone as a live hard owner, so even a successful session-only +restart is undone by the next thread-id turn. + +## Goals / Non-Goals + +**Goals:** + +- Account-neutral goal restart with `session-id` + `thread-id` retires + the unavailable raw owner for process-session interpretation and + routes to a replacement. +- Later same-thread turns without a new hard owner stay on that + replacement. +- Explicit `turn_state` of the same text stays hard-bound. + +**Non-Goals:** + +- Changing file-pin, previous-response, conversation, or tool-state + fail-closed ownership. +- Making `thread_header` an abandonment scope on the raw row. +- Dashboard, settings, or schema changes. + +## Decisions + +- Grant `abandon_unavailable_legacy_owner` for `thread_header` only + when a process session is also present. Thread-only clients have no + process-session raw row to retire. +- Allow retirement CAS when request source is `thread_header`. The + write remains `abandonment_scope=session_header`. +- Load the raw `legacy_sticky_key` with `continuity_source=session_header`. + That lookup is process-session interpretation, not thread identity. + +**Alternative considered:** keep CAS gated on request source and only +set the flag. Rejected because the CAS would still not run. + +**Alternative considered:** abandon the raw row for every source. +Rejected because colliding explicit `turn_state` must stay hard. + +## Risks / Trade-offs + +- [Risk] A thread-header request could retire a raw row that was + written as turn-state with equal text. → Mitigation: CAS still + writes `session_header` scope only; turn-state lookup of that text + keeps the stored owner. +- [Risk] Existing tests only exercise `session_id` without `thread-id`. + → Mitigation: add the missing header combination next to those tests. diff --git a/openspec/changes/goal-restart-thread-header-abandonment/proposal.md b/openspec/changes/goal-restart-thread-header-abandonment/proposal.md new file mode 100644 index 0000000000..c4a953aae3 --- /dev/null +++ b/openspec/changes/goal-restart-thread-header-abandonment/proposal.md @@ -0,0 +1,43 @@ +## Why + +Current Codex sends both a shared process `session-id` and a distinct +`thread-id` on a self-contained goal restart. Affinity classifies that +request as `thread_header`, so the one-shot +`abandon_unavailable_legacy_owner` flag never sets and retirement CAS +never runs. The restart stays fail-closed on the unavailable legacy +owner even though the payload is account-neutral. + +## What Changes + +- Grant goal-restart abandonment when a thread-scoped request still + carries a process session, not only when locality source is + `session_header`. +- Let retirement CAS retire the raw process-session row for + `session_header` interpretation from that thread-scoped request. +- Consult the raw process-session row as `session_header` + interpretation so a scoped tombstone hides it from later thread-id + turns. Explicit `turn_state` of the same text stays hard. +- Keep incremental, file-pinned, conversation-bound, and unresolved + tool-state requests fail-closed. + +## Capabilities + +### New Capabilities + +- None. + +### Modified Capabilities + +- `sticky-session-operations`: Current Codex `thread-id` on a + self-contained goal restart MUST still abandon the unavailable raw + process-session owner for `session_header` interpretation and keep + later same-thread continuity on the replacement. + +## Impact + +- `app/modules/proxy/affinity.py` restart-capability gate. +- `app/modules/proxy/_load_balancer/sticky_selection.py` retirement CAS + source check. +- `app/modules/proxy/load_balancer.py` raw-row lookup source. +- Focused affinity and sticky-selection tests. +- No API, schema, setting, dashboard, or wire-format change. diff --git a/openspec/changes/goal-restart-thread-header-abandonment/specs/sticky-session-operations/spec.md b/openspec/changes/goal-restart-thread-header-abandonment/specs/sticky-session-operations/spec.md new file mode 100644 index 0000000000..655f1d7553 --- /dev/null +++ b/openspec/changes/goal-restart-thread-header-abandonment/specs/sticky-session-operations/spec.md @@ -0,0 +1,43 @@ +## ADDED Requirements + +### Requirement: Thread-scoped current Codex restarts still abandon a raw process-session owner + +A self-contained Codex goal-continuation restart that also carries a distinct `thread-id` MUST still be eligible for the existing process-session abandonment exception. The request's thread-scoped locality source MUST NOT prevent the one-shot abandonment capability or the compare-and-set retirement of the raw process-session row. + +The retirement write MUST remain scoped to `session_header` +interpretation of that raw key. An explicit `turn_state` lookup of the +same text MUST stay hard-bound to the stored account. After a +successful retirement, later same-thread turns that have no new hard +owner MUST keep continuity on the replacement account and MUST NOT +treat the `session_header`-abandoned raw row as live hard ownership. + +Ordinary incremental, file-pinned, conversation-bound, and unresolved +tool-state requests MUST remain fail-closed on their required owner. + +#### Scenario: Goal restart with process session and thread-id abandons the unavailable raw owner + +- **GIVEN** a process-session identifier has a raw legacy `codex_session` mapping to account A +- **AND** account A is paused, rate-limited, or quota-exceeded +- **AND** account B is eligible +- **AND** the request also carries a distinct `thread-id` +- **WHEN** Codex sends the recognized goal-continuation marker with an account-neutral self-contained full resend and no other continuity dependency +- **THEN** the proxy marks the still-current raw mapping to account A abandoned only for process-session interpretation +- **AND** it routes the restarted turn to account B +- **AND** subsequent same-thread continuity remains on account B + +#### Scenario: Thread-id on a goal restart cannot erase colliding explicit turn-state ownership + +- **GIVEN** a raw legacy `codex_session` row was written as explicit turn-state ownership for account A +- **AND** a later request carries the same text as a process-session header plus a distinct `thread-id` +- **WHEN** a marked self-contained goal restart abandons that text for process-session interpretation +- **THEN** the restart may select account B +- **AND** an explicit turn-state lookup of the same text remains hard-bound to account A + +#### Scenario: Account-dependent thread-scoped restart stays fail-closed + +- **GIVEN** a process-session identifier has a raw legacy mapping to unavailable account A +- **AND** the request carries a distinct `thread-id` +- **AND** the body has a previous response, conversation, file pin, or unresolved tool state +- **WHEN** the request is selected +- **THEN** the request fails closed on account A +- **AND** the raw mapping is neither deleted nor rebound diff --git a/openspec/changes/goal-restart-thread-header-abandonment/tasks.md b/openspec/changes/goal-restart-thread-header-abandonment/tasks.md new file mode 100644 index 0000000000..788859600e --- /dev/null +++ b/openspec/changes/goal-restart-thread-header-abandonment/tasks.md @@ -0,0 +1,23 @@ +## 1. Implementation + +- [x] 1.1 Grant `abandon_unavailable_legacy_owner` for `thread_header` + when a process session is present and the payload is + account-neutral. +- [x] 1.2 Allow retirement CAS when request source is `thread_header`. + Keep the write scoped to `session_header`. +- [x] 1.3 Load the raw `legacy_sticky_key` as `session_header` + interpretation so a scoped tombstone hides it from later + thread-id turns. + +## 2. Regression coverage + +- [x] 2.1 Assert session-id + thread-id goal restart sets the + abandonment flag; turn-state and account-dependent payloads do + not. +- [x] 2.2 Assert sticky selection retires the raw owner and selects a + replacement when source is `thread_header`. + +## 3. Validation + +- [x] 3.1 Run the focused affinity and sticky-selection tests. +- [x] 3.2 Run strict OpenSpec validation for this change. diff --git a/openspec/specs/sticky-session-operations/context.md b/openspec/specs/sticky-session-operations/context.md index a85e9c0220..3dcc5d12ad 100644 --- a/openspec/specs/sticky-session-operations/context.md +++ b/openspec/specs/sticky-session-operations/context.md @@ -12,7 +12,7 @@ See `openspec/specs/sticky-session-operations/spec.md` for normative requirement - Bare process-session headers use a header-inaccessible, source-separated storage key and are soft only for self-contained pre-visible work. - Account-cap spillover is request-local: it selects an alternate without deleting or rebinding the process-session row. - Raw and legacy Codex rows remain hard during rolling upgrades because they may represent explicit turn-state ownership. -- A raw legacy Codex owner can be abandoned only for an explicit goal-continuation restart whose canonical upstream payload passes the account-neutral fresh-replay proof, and only while that owner has a persisted unavailable status. Canonicalization keeps accepted compatibility fields and transport envelopes from changing classification. The compare-and-set marker is scoped to `session_header`, so an explicit turn-state lookup with colliding raw text retains the stored owner; a concurrent rebind or owner recovery still wins. The scoped marker deliberately leaves the historical global-tombstone timestamp empty, so replicas that do not understand scope continue to fail closed on the retained owner. +- A raw legacy Codex owner can be abandoned only for an explicit goal-continuation restart whose canonical upstream payload passes the account-neutral fresh-replay proof, and only while that owner has a persisted unavailable status. Canonicalization keeps accepted compatibility fields and transport envelopes from changing classification. The compare-and-set marker is scoped to `session_header`, so an explicit turn-state lookup with colliding raw text retains the stored owner; a concurrent rebind or owner recovery still wins. The scoped marker deliberately leaves the historical global-tombstone timestamp empty, so replicas that do not understand scope continue to fail closed on the retained owner. Current Codex also sends `thread-id`; that locality source does not block the process-session exception. The raw compatibility lookup stays a `session_header` interpretation so a scoped tombstone cannot revive the retired owner on later thread-id turns. - Restart mutation authority is the authenticated account-assignment and security-policy scope before model and service-tier eligibility. Model filtering constrains only replacement selection. - Goal-restart retirement is an account-selection capability. An existing HTTP bridge cannot consume the request first through local reuse, durable-owner promotion, or forwarding. The retired owner is excluded from stale account snapshots for the remainder of the request, including when another selector wrote the scoped marker and this selector discovers it after losing the compare-and-set. - Canonical bridge replacement preserves request-owned pre-submit admission on the detached predecessor, but that predecessor cannot publish new continuity aliases under the replacement's key. Every detached generation remains lifecycle-owned and capacity-counted until resource closure ends, including an idle predecessor already marked closed for admission. diff --git a/openspec/specs/sticky-session-operations/spec.md b/openspec/specs/sticky-session-operations/spec.md index babf08b41a..3472080f29 100644 --- a/openspec/specs/sticky-session-operations/spec.md +++ b/openspec/specs/sticky-session-operations/spec.md @@ -66,6 +66,34 @@ A later security-authorized bridge replacement that revalidates a raw legacy row - **THEN** the process-session restart may select account B - **AND** an explicit turn-state lookup of the same text remains hard-bound to account A +#### Scenario: Goal restart with process session and thread-id abandons the unavailable raw owner + +- **GIVEN** a process-session identifier has a raw legacy `codex_session` mapping to account A +- **AND** account A is paused, rate-limited, or quota-exceeded +- **AND** account B is eligible +- **AND** the request also carries a distinct `thread-id` +- **WHEN** Codex sends the recognized goal-continuation marker with an account-neutral self-contained full resend and no other continuity dependency +- **THEN** the proxy marks the still-current raw mapping to account A abandoned only for process-session interpretation +- **AND** it routes the restarted turn to account B +- **AND** subsequent same-thread continuity remains on account B + +#### Scenario: Thread-id on a goal restart cannot erase colliding explicit turn-state ownership + +- **GIVEN** a raw legacy `codex_session` row was written as explicit turn-state ownership for account A +- **AND** a later request carries the same text as a process-session header plus a distinct `thread-id` +- **WHEN** a marked self-contained goal restart abandons that text for process-session interpretation +- **THEN** the restart may select account B +- **AND** an explicit turn-state lookup of the same text remains hard-bound to account A + +#### Scenario: Account-dependent thread-scoped restart stays fail-closed + +- **GIVEN** a process-session identifier has a raw legacy mapping to unavailable account A +- **AND** the request carries a distinct `thread-id` +- **AND** the body has a previous response, conversation, file pin, or unresolved tool state +- **WHEN** the request is selected +- **THEN** the request fails closed on account A +- **AND** the raw mapping is neither deleted nor rebound + #### Scenario: Source-qualified retirement fails closed on an older replica - **GIVEN** a current replica marks a raw account A mapping abandoned only for `session_header` interpretation diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index dcb1bfab87..9699e5a605 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -3126,6 +3126,54 @@ async def test_goal_restart_does_not_repin_retired_owner_from_stale_selection_sn await balancer.release_account_lease(selected.lease) +@pytest.mark.asyncio +async def test_goal_restart_with_thread_header_retires_unavailable_legacy_owner() -> None: + now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) + stale_owner = _make_account("goal-restart-thread-header-owner") + replacement = _make_account("goal-restart-thread-header-replacement") + raw_session = "goal-restart-thread-header-session" + thread_key = _codex_backend_identity( + {"session-id": raw_session, "thread-id": "goal-restart-thread"} + ).thread_selection_key + assert thread_key is not None + sticky_repo = _RetiringStaleOwnerStickySessionsRepository( + raw_key=raw_session, + owner_account_id=stale_owner.id, + ) + balancer = LoadBalancer( + lambda: _repo_factory( + _StubAccountsRepository([stale_owner, replacement]), + _StubUsageRepository( + { + stale_owner.id: _usage_row(311, stale_owner.id, window="primary", reset_at=now_epoch + 300), + replacement.id: _usage_row(312, replacement.id, window="primary", reset_at=now_epoch + 300), + }, + {}, + ), + sticky_repo, + ) + ) + + selected = await balancer.select_account( + sticky_key=thread_key, + sticky_kind=StickySessionKind.PROMPT_CACHE, + sticky_source="thread_header", + sticky_max_age_seconds=300, + legacy_sticky_key=raw_session, + abandon_unavailable_legacy_owner=True, + routing_strategy="single_account", + lease_kind="stream", + ) + + assert selected.account is not None + assert selected.account.id == replacement.id + assert sticky_repo.tombstones == [(raw_session, stale_owner.id)] + assert sticky_repo.account_ids_by_key[raw_session] == stale_owner.id + assert sticky_repo.account_ids_by_key[thread_key] == replacement.id + assert all(account_id != stale_owner.id for _, account_id, _ in sticky_repo.upserts) + await balancer.release_account_lease(selected.lease) + + @pytest.mark.asyncio async def test_goal_restart_cas_loser_does_not_repin_concurrently_retired_owner() -> None: now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 43a1f8c887..c54f0c09d8 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -10506,8 +10506,33 @@ def test_goal_restart_affinity_can_abandon_only_legacy_session_owner(): sticky_threads_enabled=False, ) + thread_policy = proxy_service._sticky_key_for_responses_request( + payload, + headers={ + "session_id": "goal-restart-session", + "thread-id": "goal-restart-thread", + }, + codex_session_affinity=True, + openai_cache_affinity=False, + openai_cache_affinity_max_age_seconds=300, + sticky_threads_enabled=False, + ) + thread_only_policy = proxy_service._sticky_key_for_responses_request( + payload, + headers={"thread-id": "goal-restart-thread"}, + codex_session_affinity=True, + openai_cache_affinity=False, + openai_cache_affinity_max_age_seconds=300, + sticky_threads_enabled=False, + ) + assert policy.codex_session_source == "session_header" assert policy.abandon_unavailable_legacy_owner is True + assert thread_policy.codex_session_source == "thread_header" + assert thread_policy.abandon_unavailable_legacy_owner is True + assert thread_policy.legacy_selection_key == "goal-restart-session" + assert thread_only_policy.codex_session_source == "thread_header" + assert thread_only_policy.abandon_unavailable_legacy_owner is False assert turn_state_policy.codex_session_source == "turn_state" assert turn_state_policy.abandon_unavailable_legacy_owner is False @@ -10571,6 +10596,25 @@ def test_goal_restart_affinity_preserves_owner_for_account_dependent_payloads( assert policy.abandon_unavailable_legacy_owner is False +def test_goal_restart_affinity_preserves_owner_for_account_dependent_thread_payloads(): + payload = _goal_restart_payload(previous_response_id="resp_owner") + + policy = proxy_service._sticky_key_for_responses_request( + payload, + headers={ + "session_id": "goal-restart-session", + "thread-id": "goal-restart-thread", + }, + codex_session_affinity=True, + openai_cache_affinity=False, + openai_cache_affinity_max_age_seconds=300, + sticky_threads_enabled=False, + ) + + assert policy.codex_session_source == "thread_header" + assert policy.abandon_unavailable_legacy_owner is False + + def test_full_resend_without_goal_marker_cannot_abandon_legacy_owner(): payload = ResponsesRequest.model_validate( { From 076f0c77a5721c05c3ffe82cee94eafcdb8992e5 Mon Sep 17 00:00:00 2001 From: mastertyko <11311479+mastertyko@users.noreply.github.com> Date: Sat, 15 Aug 2026 13:50:11 +0200 Subject: [PATCH 2/5] test(proxy): avoid optional sticky-map subscript in thread restart test ty rejects subscripting account_ids_by_key because the stub field is optional. Compare the whole mapping like the existing session-header test. --- tests/unit/test_load_balancer_concurrency.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index 9699e5a605..685f22fae1 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -3168,8 +3168,10 @@ async def test_goal_restart_with_thread_header_retires_unavailable_legacy_owner( assert selected.account is not None assert selected.account.id == replacement.id assert sticky_repo.tombstones == [(raw_session, stale_owner.id)] - assert sticky_repo.account_ids_by_key[raw_session] == stale_owner.id - assert sticky_repo.account_ids_by_key[thread_key] == replacement.id + assert sticky_repo.account_ids_by_key == { + raw_session: stale_owner.id, + thread_key: replacement.id, + } assert all(account_id != stale_owner.id for _, account_id, _ in sticky_repo.upserts) await balancer.release_account_lease(selected.lease) From 13f520162606e1d58fd13ba6e920d6b7e19b263c Mon Sep 17 00:00:00 2001 From: mastertyko <11311479+mastertyko@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:02:36 +0200 Subject: [PATCH 3/5] fix(proxy): keep thread-only raw owners and cover route restart Codex review: a hardcoded session_header lookup hid thread-only raw rows after a process-session tombstone, and the restart path lacked /backend-api/codex/responses coverage for session-id plus thread-id. Look up the raw key with the source that actually wrote it, and add the route-level restart plus follow-up continuity test. --- app/modules/proxy/_service/codex_control.py | 3 + app/modules/proxy/_service/websocket/mixin.py | 1 + app/modules/proxy/affinity.py | 9 + app/modules/proxy/load_balancer.py | 10 +- app/modules/proxy/service.py | 4 + .../integration/test_proxy_sticky_sessions.py | 160 +++++++++++++++++- tests/unit/test_proxy_utils.py | 2 + 7 files changed, 184 insertions(+), 5 deletions(-) diff --git a/app/modules/proxy/_service/codex_control.py b/app/modules/proxy/_service/codex_control.py index 242f2bbd30..5c3fad98a2 100644 --- a/app/modules/proxy/_service/codex_control.py +++ b/app/modules/proxy/_service/codex_control.py @@ -229,6 +229,7 @@ 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, + legacy_continuity_source=affinity.legacy_continuity_source, sticky_seed_key=affinity.seed_selection_key, sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, @@ -398,6 +399,7 @@ 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, + legacy_continuity_source=affinity.legacy_continuity_source, sticky_seed_key=affinity.seed_selection_key, sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, @@ -492,6 +494,7 @@ 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, + legacy_continuity_source=affinity.legacy_continuity_source, sticky_seed_key=affinity.seed_selection_key, sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index ea866d4bf8..0e1894e559 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -3476,6 +3476,7 @@ 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, + legacy_continuity_source=request_state.affinity_policy.legacy_continuity_source, 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, diff --git a/app/modules/proxy/affinity.py b/app/modules/proxy/affinity.py index a8e5d25a36..a8e3e35761 100644 --- a/app/modules/proxy/affinity.py +++ b/app/modules/proxy/affinity.py @@ -43,6 +43,7 @@ class _AffinitySelectionKwargs(TypedDict): reallocate_sticky: bool sticky_source: _CodexSessionSource | None legacy_sticky_key: str | None + legacy_continuity_source: _CodexSessionSource | None sticky_seed_key: str | None sticky_seed_kind: StickySessionKind | None spill_bare_session_on_account_cap: bool @@ -69,6 +70,10 @@ class _AffinityPolicy: # compatibility lookup explicit instead of trying to reconstruct it from # the new opaque thread key. legacy_codex_session_key: str | None = None + # Interpretation used when consulting that raw key. Process-session text + # is session_header even on a thread-scoped request; a thread-only raw + # key stays thread_header so a session_header tombstone cannot hide it. + legacy_continuity_source: _CodexSessionSource | 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 @@ -109,6 +114,9 @@ def selection_kwargs(self) -> _AffinitySelectionKwargs: "reallocate_sticky": self.reallocate_sticky, "sticky_source": self.codex_session_source, "legacy_sticky_key": self.legacy_selection_key, + "legacy_continuity_source": ( + None if self.legacy_selection_key is None else (self.legacy_continuity_source or "session_header") + ), "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, @@ -433,6 +441,7 @@ def _thread_codex_session_affinity( max_age_seconds=max_age_seconds, codex_session_source="thread_header", legacy_codex_session_key=legacy_key, + legacy_continuity_source=("session_header" if identity.process_session is not None else "thread_header"), seed_selection_key=( _codex_session_selection_key(identity.process_session) if identity.process_session is not None else None ), diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 20c081adac..3a6dd88e2f 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -532,6 +532,7 @@ async def select_account( reallocate_sticky: bool = False, sticky_source: _CodexSessionSource | None = None, legacy_sticky_key: str | None = None, + legacy_continuity_source: _CodexSessionSource | None = None, sticky_seed_key: str | None = None, sticky_seed_kind: StickySessionKind | None = None, spill_bare_session_on_account_cap: bool = False, @@ -731,10 +732,11 @@ async def load_selection_inputs() -> _SelectionInputs: # Raw rows may be historical turn-state ownership. The # bounded thread TTL must never age out that hard evidence. max_age_seconds=None, - # This key is the process-session compatibility row. - # Thread-header requests still consult it as that row, - # so session_header-scoped abandonment must hide it. - continuity_source="session_header", + # Process-session raw text is session_header even when + # request locality is thread_header. Thread-only raw keys + # keep thread_header so a session_header tombstone cannot + # hide a distinct thread owner. + continuity_source=legacy_continuity_source or "session_header", ) legacy_existing_account_id = legacy_owner_lookup.account_id abandoned_account_id = legacy_owner_lookup.abandoned_account_id diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index 18db15a76f..e7c62e2da7 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -1084,6 +1084,7 @@ 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, + legacy_continuity_source=affinity.legacy_continuity_source, sticky_seed_key=affinity.seed_selection_key, sticky_seed_kind=affinity.seed_selection_kind, sticky_max_age_seconds=affinity.max_age_seconds, @@ -1702,6 +1703,7 @@ async def _select_account_with_budget( reallocate_sticky: bool = False, sticky_source: _CodexSessionSource | None = None, legacy_sticky_key: str | None = None, + legacy_continuity_source: _CodexSessionSource | None = None, sticky_seed_key: str | None = None, sticky_seed_kind: StickySessionKind | None = None, spill_bare_session_on_account_cap: bool = False, @@ -1862,6 +1864,7 @@ 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], + legacy_continuity_source=legacy_continuity_source, # Exact ownership chooses the account; a first-ever thread # still seeds atomically without overwriting a process default. sticky_seed_key=sticky_seed_key, @@ -1924,6 +1927,7 @@ def log_account_id(account_id: str | None) -> str | None: reallocate_sticky=reallocate_sticky, sticky_source=sticky_source, legacy_sticky_key=legacy_sticky_key, + legacy_continuity_source=legacy_continuity_source, sticky_seed_key=sticky_seed_key, sticky_seed_kind=sticky_seed_kind, spill_bare_session_on_account_cap=_AffinityPolicy.cap_spillover_allowed( diff --git a/tests/integration/test_proxy_sticky_sessions.py b/tests/integration/test_proxy_sticky_sessions.py index 5ad05217e2..0302976c05 100644 --- a/tests/integration/test_proxy_sticky_sessions.py +++ b/tests/integration/test_proxy_sticky_sessions.py @@ -23,7 +23,7 @@ _REALTIME_CALL_AFFINITY_MAX_AGE_SECONDS, realtime_call_affinity_key, ) -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.usage.repository import UsageRepository pytestmark = pytest.mark.integration @@ -401,6 +401,164 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, assert legacy_replica_owner == owner_id +@pytest.mark.asyncio +async def test_codex_goal_restart_with_thread_id_retires_unavailable_legacy_owner_and_stays_on_replacement( + async_client, + monkeypatch, +): + from sqlalchemy import select + + from app.db.models import StickySession + from app.modules.proxy.sticky_repository import StickySessionsRepository + + _install_proxy_settings_cache(monkeypatch, sticky_threads_enabled=False) + owner_id = await _import_account( + async_client, + "acc_goal_restart_thread_owner", + "goal-restart-thread-owner@example.com", + ) + replacement_id = await _import_account( + async_client, + "acc_goal_restart_thread_replacement", + "goal-restart-thread-replacement@example.com", + ) + raw_session = "goal-restart-thread-session" + thread_id = "goal-restart-thread" + headers = {"session_id": raw_session, "thread-id": thread_id} + thread_key = _codex_backend_identity(headers).thread_selection_key + assert thread_key is not None + + now_epoch = int(utcnow().replace(tzinfo=timezone.utc).timestamp()) + async with SessionLocal() as session: + usage_repo = UsageRepository(session) + await usage_repo.add_entry( + account_id=owner_id, + used_percent=10.0, + window="primary", + reset_at=now_epoch + 3600, + window_minutes=300, + ) + await usage_repo.add_entry( + account_id=replacement_id, + used_percent=20.0, + window="primary", + reset_at=now_epoch + 3600, + window_minutes=300, + ) + await StickySessionsRepository(session).upsert( + raw_session, + owner_id, + kind=StickySessionKind.CODEX_SESSION, + ) + + seen: list[str] = [] + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False, **kwargs): + del payload, headers, access_token, base_url, raise_for_status, kwargs + seen.append(account_id) + yield f'data: {{"type":"response.completed","response":{{"id":"resp_goal_thread_{len(seen)}"}}}}\n\n' + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + restart_payload = { + "model": "gpt-5.1", + "instructions": "Continue the existing task.", + "input": [ + { + "role": "developer", + "content": ('\nContinue working toward the active thread goal.'), + }, + {"role": "user", "content": [{"type": "input_text", "text": "continue"}]}, + ], + "stream": True, + } + + healthy_response = await async_client.post( + "/backend-api/codex/responses", + json=restart_payload, + headers=headers, + ) + assert healthy_response.status_code == 200 + assert seen == ["acc_goal_restart_thread_owner"] + + async with SessionLocal() as session: + await session.execute(update(Account).where(Account.id == owner_id).values(status=AccountStatus.QUOTA_EXCEEDED)) + await session.commit() + + restart_response = await async_client.post( + "/backend-api/codex/responses", + json=restart_payload, + headers=headers, + ) + assert restart_response.status_code == 200 + assert seen == ["acc_goal_restart_thread_owner", "acc_goal_restart_thread_replacement"] + + follow_up_response = await async_client.post( + "/backend-api/codex/responses", + json={"model": "gpt-5.1", "instructions": "continue", "input": [], "stream": True}, + headers=headers, + ) + assert follow_up_response.status_code == 200 + assert seen == [ + "acc_goal_restart_thread_owner", + "acc_goal_restart_thread_replacement", + "acc_goal_restart_thread_replacement", + ] + + turn_state_response = await async_client.post( + "/backend-api/codex/responses", + json={"model": "gpt-5.1", "instructions": "continue", "input": [], "stream": True}, + headers={"x-codex-turn-state": raw_session}, + ) + assert turn_state_response.status_code == 502 + assert turn_state_response.json()["error"]["code"] == "turn_state_owner_unavailable" + assert seen == [ + "acc_goal_restart_thread_owner", + "acc_goal_restart_thread_replacement", + "acc_goal_restart_thread_replacement", + ] + + async with SessionLocal() as session: + repo = StickySessionsRepository(session) + raw_row = ( + await session.execute( + select(StickySession).where( + StickySession.key == raw_session, + StickySession.kind == StickySessionKind.CODEX_SESSION, + ) + ) + ).scalar_one() + thread_row = ( + await session.execute( + select(StickySession).where( + StickySession.key == thread_key, + StickySession.kind == StickySessionKind.PROMPT_CACHE, + ) + ) + ).scalar_one() + session_header_lookup = await repo.get_account_id_and_abandonment( + raw_session, + kind=StickySessionKind.CODEX_SESSION, + continuity_source="session_header", + ) + thread_legacy_lookup = await repo.get_account_id_and_abandonment( + raw_session, + kind=StickySessionKind.CODEX_SESSION, + continuity_source="thread_header", + ) + turn_state_owner = await repo.get_account_id( + raw_session, + kind=StickySessionKind.CODEX_SESSION, + continuity_source="turn_state", + ) + assert raw_row.account_id == owner_id + assert raw_row.continuity_abandonment_scope == "session_header" + assert thread_row.account_id == replacement_id + assert session_header_lookup.account_id is None + assert session_header_lookup.continuity_abandoned is True + assert thread_legacy_lookup.account_id == owner_id + assert turn_state_owner == owner_id + + @pytest.mark.asyncio async def test_codex_goal_restart_cas_miss_reloads_concurrently_rebound_raw_owner( async_client, diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index c54f0c09d8..3737e4ba44 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -10531,8 +10531,10 @@ def test_goal_restart_affinity_can_abandon_only_legacy_session_owner(): assert thread_policy.codex_session_source == "thread_header" assert thread_policy.abandon_unavailable_legacy_owner is True assert thread_policy.legacy_selection_key == "goal-restart-session" + assert thread_policy.legacy_continuity_source == "session_header" assert thread_only_policy.codex_session_source == "thread_header" assert thread_only_policy.abandon_unavailable_legacy_owner is False + assert thread_only_policy.legacy_continuity_source == "thread_header" assert turn_state_policy.codex_session_source == "turn_state" assert turn_state_policy.abandon_unavailable_legacy_owner is False From d7082250a9a76296841257703ef16ddc60d525e9 Mon Sep 17 00:00:00 2001 From: mastertyko <11311479+mastertyko@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:13:29 +0200 Subject: [PATCH 4/5] test(proxy): accept legacy continuity source in control selection mock --- tests/unit/test_proxy_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 3737e4ba44..d917385d86 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -5751,6 +5751,7 @@ async def test_select_codex_control_account_without_budget_uses_balancer(monkeyp reallocate_sticky=False, sticky_source=None, legacy_sticky_key=None, + legacy_continuity_source=None, sticky_seed_key=None, sticky_seed_kind=None, sticky_max_age_seconds=123, From f0e7a83d97b7346116d946fa248eeea925e2ac1c Mon Sep 17 00:00:00 2001 From: mastertyko <11311479+mastertyko@users.noreply.github.com> Date: Sat, 15 Aug 2026 14:19:49 +0200 Subject: [PATCH 5/5] fix(proxy): use legacy continuity source on HTTP-bridge rebind Security-authorized replacement still looked up the raw process-session row with thread_header, so a session_header tombstone resurrected the retired owner as a continuity conflict. Lookup the raw key with legacy_continuity_source instead. --- .../_service/http_bridge/request_submit.py | 2 +- tests/unit/test_proxy_utils.py | 85 +++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/app/modules/proxy/_service/http_bridge/request_submit.py b/app/modules/proxy/_service/http_bridge/request_submit.py index 7662fee5d8..e885d7f28f 100644 --- a/app/modules/proxy/_service/http_bridge/request_submit.py +++ b/app/modules/proxy/_service/http_bridge/request_submit.py @@ -3599,7 +3599,7 @@ async def _claim_http_bridge_replacement_before_swap( # remains durable hard ownership. kind=StickySessionKind.CODEX_SESSION, max_age_seconds=None, - continuity_source=owner_rebind_affinity.codex_session_source, + continuity_source=(owner_rebind_affinity.legacy_continuity_source or "session_header"), ) if legacy_owner_id is not None and legacy_owner_id != account_id: raise ProxyResponseError( diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index d917385d86..3e34ab342b 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -18112,6 +18112,91 @@ async def __aexit__(self, exc_type, exc, tb) -> bool: assert session.closed is False +@pytest.mark.asyncio +async def test_http_bridge_replacement_uses_legacy_continuity_source_for_raw_row() -> None: + rejected_account = _make_account("acc_bridge_thread_restart_owner") + authorized_account = _make_account("acc_bridge_thread_restart_replacement") + sticky_sessions = AsyncMock() + seen_sources: list[str | None] = [] + + async def legacy_owner_for_source( + _key: str, + *, + kind: StickySessionKind, + max_age_seconds: int | None = None, + continuity_source: str | None = None, + ) -> str | None: + del kind, max_age_seconds + seen_sources.append(continuity_source) + return rejected_account.id if continuity_source == "thread_header" else None + + sticky_sessions.get_account_id.side_effect = legacy_owner_for_source + + class _TrackingRepoContext: + def __init__(self) -> None: + self._repos = ProxyRepositories( + accounts=cast(AccountsRepository, AsyncMock()), + usage=cast(UsageRepository, AsyncMock()), + request_logs=cast(RequestLogsRepository, _RequestLogsRecorder()), + sticky_sessions=cast(StickySessionsRepository, sticky_sessions), + api_keys=cast(ApiKeysRepository, AsyncMock()), + additional_usage=cast(AdditionalUsageRepository, AsyncMock()), + ) + + async def __aenter__(self) -> ProxyRepositories: + return self._repos + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + service = proxy_service.ProxyService(_TrackingRepoContext) + replacement_upstream = AsyncMock() + affinity = proxy_service._AffinityPolicy( + key="thread-restart-rebind", + kind=StickySessionKind.PROMPT_CACHE, + codex_session_source="thread_header", + legacy_codex_session_key="process-restart-rebind", + legacy_continuity_source="session_header", + ) + session = proxy_service._HTTPBridgeSession( + key=proxy_service._HTTPBridgeSessionKey("thread_header", "thread-restart-rebind", None), + headers={"session_id": "process-restart-rebind", "thread-id": "thread-restart-rebind"}, + affinity=affinity, + request_model="gpt-5.1", + account=rejected_account, + upstream=AsyncMock(), + upstream_control=proxy_service._WebSocketUpstreamControl(), + pending_requests=deque(), + pending_lock=anyio.Lock(), + response_create_gate=asyncio.Semaphore(1), + queued_request_count=0, + last_used_at=1.0, + idle_ttl_seconds=300.0, + durable_session_id="durable-thread-restart-rebind", + durable_owner_epoch=2, + ) + durable_claim = AsyncMock() + service._claim_durable_http_bridge_session = durable_claim + + await service._claim_http_bridge_replacement_before_swap( + session, + account_id=authorized_account.id, + upstream=replacement_upstream, + release_selected_account_lease=AsyncMock(), + owner_rebind_affinity=affinity, + ) + + assert seen_sources == ["session_header"] + sticky_sessions.get_account_id.assert_awaited_once_with( + "process-restart-rebind", + kind=StickySessionKind.CODEX_SESSION, + max_age_seconds=None, + continuity_source="session_header", + ) + durable_claim.assert_awaited_once() + replacement_upstream.close.assert_not_awaited() + + @pytest.mark.asyncio async def test_http_bridge_security_retry_restores_codex_affinity_and_turn_aliases_on_failure( monkeypatch: pytest.MonkeyPatch,