diff --git a/app/modules/proxy/affinity.py b/app/modules/proxy/affinity.py index 4dd41a114d..029270ec82 100644 --- a/app/modules/proxy/affinity.py +++ b/app/modules/proxy/affinity.py @@ -142,7 +142,7 @@ def preferred_owner_sticky_inputs( _CodexSessionSource | None, str | None, ]: - if sticky_source != "session_header": + if sticky_source not in {"session_header", "thread_header"}: return ( sticky_key, sticky_kind, @@ -151,10 +151,11 @@ def preferred_owner_sticky_inputs( sticky_source, legacy_sticky_key, ) - # 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. The caller also + # A resolved response/file/bridge owner bypasses the current-Codex + # soft row (process-session or thread PROMPT_CACHE). 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. 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 diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 0bc5717dd5..464c2ba29b 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -817,6 +817,27 @@ async def load_selection_inputs() -> _SelectionInputs: error_message=error_message, error_code=selection_error_code, ) + if ( + selected_snapshot is not None + and selected_lease is not None + and sticky_seed_key is not None + and sticky_seed_kind is not None + and sticky_seed_account_id is None + ): + # Required-owner selection bypasses the thread row, but a + # first-ever process preference still has to land so later + # unpinned siblings inherit that exact owner. + try: + async with self._repo_factory() as repos: + await repos.sticky_sessions.insert_if_absent( + sticky_seed_key, + selected_snapshot.id, + sticky_seed_kind, + ) + except BaseException: + await self.release_account_lease(selected_lease) + selected_lease = None + raise else: sticky_outcome = await run_sticky_selection_path( self, diff --git a/openspec/changes/file-pin-does-not-rewrite-thread-locality/.openspec.yaml b/openspec/changes/file-pin-does-not-rewrite-thread-locality/.openspec.yaml new file mode 100644 index 0000000000..0c73c8f54e --- /dev/null +++ b/openspec/changes/file-pin-does-not-rewrite-thread-locality/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-08-15 diff --git a/openspec/changes/file-pin-does-not-rewrite-thread-locality/context.md b/openspec/changes/file-pin-does-not-rewrite-thread-locality/context.md new file mode 100644 index 0000000000..65772f6eba --- /dev/null +++ b/openspec/changes/file-pin-does-not-rewrite-thread-locality/context.md @@ -0,0 +1,14 @@ +## Purpose + +Stop a live file pin from rewriting current-Codex thread locality. + +## Decision + +The required owner bypasses the thread PROMPT_CACHE row the same way +it already bypasses the process-session soft row. + +## Example + +Upload pins `file_xyz` to account A. Thread `t1` is already mapped to +account B. A Responses turn that references `file_xyz` goes to A; the +`t1` row remains B. The next unpinned `t1` turn still uses B. diff --git a/openspec/changes/file-pin-does-not-rewrite-thread-locality/design.md b/openspec/changes/file-pin-does-not-rewrite-thread-locality/design.md new file mode 100644 index 0000000000..08dff38f7f --- /dev/null +++ b/openspec/changes/file-pin-does-not-rewrite-thread-locality/design.md @@ -0,0 +1,42 @@ +## Context + +`#1521` made file pins durable hard ownership and bypassed the +process-session soft row. `#1703` then made `thread_header` / +PROMPT_CACHE the current Codex soft mapping. The bypass was not +updated, so a required file owner still enters sticky persist and +upserts thread B→A. + +## Goals / Non-Goals + +**Goals:** + +- File-pinned routing stays on the pin account. +- An existing thread PROMPT_CACHE row is not rewritten. +- Process-session seed remains insert-if-absent. + +**Non-Goals:** + +- Changing 1011 file-pin reconnect. +- Weakening file-pin fail-closed or hard-owner conflict checks. +- Dashboard or settings changes. + +## Decisions + +- Null the writable sticky key for both `session_header` and + `thread_header` in `preferred_owner_sticky_inputs`. Selection then + takes the unbound required-owner path. +- Keep `legacy_sticky_key` so a conflicting raw process-session owner + still fail-closes. +- Leave `sticky_seed_key` to the caller so a missing process + preference can still initialize without writing the thread row. + +**Alternative considered:** persist the thread row onto the file +owner so later unpinned turns stay there. Rejected: the file pin is +hard only for this turn; thread locality is a separate soft mapping. + +## Risks / Trade-offs + +- [Risk] A later unpinned turn on the same thread stays on the + pre-file account and cannot see the upload. → Mitigation: that is + the existing unpinned-file compatibility path; the pin still binds + any turn that references the file. diff --git a/openspec/changes/file-pin-does-not-rewrite-thread-locality/proposal.md b/openspec/changes/file-pin-does-not-rewrite-thread-locality/proposal.md new file mode 100644 index 0000000000..80167449d5 --- /dev/null +++ b/openspec/changes/file-pin-does-not-rewrite-thread-locality/proposal.md @@ -0,0 +1,36 @@ +## Why + +A live `input_file.file_id` pin is hard ownership. After thread-scoped +affinity, current Codex locality is the `thread_header` PROMPT_CACHE +row, but `preferred_owner_sticky_inputs` only bypasses +`session_header`. A file-pinned Responses turn therefore rewrites the +thread mapping to the upload account, so later unpinned turns follow +the file owner. + +## What Changes + +- Treat `thread_header` as the current-Codex soft row that a resolved + file/response/bridge owner must bypass. +- Keep consulting the raw process-session compatibility row for hard + conflicts. +- Keep process-session seed insert-if-absent. Do not write or rebind + the thread row on the required-owner path. +- Keep explicit `turn_state` as hard ownership. + +## Capabilities + +### New Capabilities + +- None. + +### Modified Capabilities + +- `sticky-session-operations`: A resolved file-pin owner MUST be + selected without consulting or rewriting the thread-scoped soft + mapping. + +## Impact + +- `app/modules/proxy/affinity.py` preferred-owner sticky inputs. +- Focused selection tests. +- No API, schema, setting, dashboard, or wire-format change. diff --git a/openspec/changes/file-pin-does-not-rewrite-thread-locality/specs/sticky-session-operations/spec.md b/openspec/changes/file-pin-does-not-rewrite-thread-locality/specs/sticky-session-operations/spec.md new file mode 100644 index 0000000000..3ec60a18b6 --- /dev/null +++ b/openspec/changes/file-pin-does-not-rewrite-thread-locality/specs/sticky-session-operations/spec.md @@ -0,0 +1,21 @@ +## ADDED Requirements + +### Requirement: File-pin required owner does not rewrite thread locality + +A resolved live `input_file.file_id` pin MUST be selected as the required owner without consulting or rewriting the current-Codex thread-scoped soft mapping. The process-session compatibility row MAY still be consulted as independent hard ownership. If that raw row conflicts with the pin account, the request MUST fail closed. A missing process-session preference MAY still initialize insert-if-absent. + +#### Scenario: File-pinned request owner overrides thread locality + +- **GIVEN** a request carries a `thread-id` whose bounded mapping points to account A +- **AND** its `input_file.file_id` is durably pinned to account B +- **WHEN** the request is routed +- **THEN** account B is treated as the required owner +- **AND** the thread mapping is neither consulted as an owner nor rewritten + +#### Scenario: File pin still conflicts with a raw process-session owner + +- **GIVEN** a raw process-session `codex_session` row points to account A +- **AND** a live file pin points to account B +- **WHEN** the request is routed +- **THEN** the service fails with `continuity_owner_conflict` before upstream dispatch +- **AND** neither the raw row nor the thread row is rewritten diff --git a/openspec/changes/file-pin-does-not-rewrite-thread-locality/tasks.md b/openspec/changes/file-pin-does-not-rewrite-thread-locality/tasks.md new file mode 100644 index 0000000000..c016d94535 --- /dev/null +++ b/openspec/changes/file-pin-does-not-rewrite-thread-locality/tasks.md @@ -0,0 +1,19 @@ +## 1. Implementation + +- [x] 1.1 Bypass the writable `thread_header` sticky key in + `preferred_owner_sticky_inputs` the same way as `session_header`. + +## 2. Regression coverage + +- [x] 2.1 Assert preferred-owner selection nulls the thread sticky key + and keeps the process seed / raw legacy key. +- [x] 2.2 Assert an existing thread row is not upserted when a file + pin is the required owner. +- [x] 2.3 Cover the same file-pin plus existing-thread case through + `/backend-api/codex/responses`, including the later unpinned + thread turn and process-seed sibling. + +## 3. Validation + +- [x] 3.1 Run the focused 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..e43665c465 100644 --- a/openspec/specs/sticky-session-operations/context.md +++ b/openspec/specs/sticky-session-operations/context.md @@ -18,7 +18,7 @@ See `openspec/specs/sticky-session-operations/spec.md` for normative requirement - 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. - Drain status counts unsettled pending or queued work after detachment closes a generation for admission. If an idle predecessor alone fills the cap, the verified restart owns its bounded close synchronously and rechecks capacity before opening a replacement. - Resource close is single-flight across reader retirement, account invalidation, and shutdown. Capacity is released only after resource finalization, not after detachment or a bounded-close timeout. Close finalization defers caller cancellation until owned resources are released, while shutdown starts all snapshotted closes before propagating cancellation and retains failed generations for a later close pass. Durable claims are fenced per websocket generation as well as per replica, so a replacement for a row still owned by the same configured replica advances the owner epoch before serving work even when model-transition isolation no longer uses that row for routing. Security-authorized rebind keeps typed continuity provenance so a source-qualified session-header tombstone cannot reappear as an untyped hard owner. -- Durable file pins, responses, conversations, live/durable bridges, replay, and reattach sources are independent hard evidence; conflicting evidence fails closed instead of using source precedence. Opaque file IDs with no live durable pin remain unpinned for compatibility with uploads that occurred outside the current process. +- Durable file pins, responses, conversations, live/durable bridges, replay, and reattach sources are independent hard evidence; conflicting evidence fails closed instead of using source precedence. Opaque file IDs with no live durable pin remain unpinned for compatibility with uploads that occurred outside the current process. A resolved file-pin owner bypasses the current-Codex thread PROMPT_CACHE row the same way it bypasses process-session locality, so an upload does not rebind later unpinned turns on that thread. - Dashboard prompt-cache TTL is persisted in settings so operators can adjust it without restart. - Background cleanup removes stale prompt-cache rows proactively, while manual delete and purge endpoints provide operator override. diff --git a/openspec/specs/sticky-session-operations/spec.md b/openspec/specs/sticky-session-operations/spec.md index babf08b41a..6203a222c2 100644 --- a/openspec/specs/sticky-session-operations/spec.md +++ b/openspec/specs/sticky-session-operations/spec.md @@ -332,6 +332,14 @@ A nonblank `conversation` without a dedicated resolved owner MUST proceed only w - **THEN** account B is treated as the required owner - **AND** the process-session mapping is neither consulted as an owner nor rewritten +#### Scenario: File-pinned request owner overrides thread locality + +- **GIVEN** a request carries a `thread-id` whose bounded mapping points to account A +- **AND** its `input_file.file_id` is durably pinned to account B +- **WHEN** the request is routed +- **THEN** account B is treated as the required owner +- **AND** the thread mapping is neither consulted as an owner nor rewritten + #### Scenario: Conflicting hard owners fail closed - **GIVEN** a turn state, previous response, bridge, or input file resolves to account A diff --git a/tests/integration/test_proxy_files.py b/tests/integration/test_proxy_files.py index fc9cb926eb..45068df607 100644 --- a/tests/integration/test_proxy_files.py +++ b/tests/integration/test_proxy_files.py @@ -24,9 +24,11 @@ from app.core.auth.refresh import RefreshError from app.core.clients.files import FileProxyError from app.core.clients.proxy import ProxyResponseError -from app.db.models import FileAccountPin +from app.db.models import FileAccountPin, StickySessionKind from app.db.session import SessionLocal +from app.modules.proxy.affinity import _codex_backend_identity, _codex_session_selection_key from app.modules.proxy.file_pin_repository import FileAccountPinRepository +from app.modules.proxy.sticky_repository import StickySessionsRepository pytestmark = pytest.mark.integration @@ -53,11 +55,12 @@ def _make_auth_json(account_id: str, email: str) -> dict: } -async def _import_account(async_client, account_id: str, email: str) -> None: +async def _import_account(async_client, account_id: str, email: str) -> str: auth_json = _make_auth_json(account_id, email) files = {"auth_json": ("auth.json", json.dumps(auth_json), "application/json")} response = await async_client.post("/api/accounts/import", files=files) assert response.status_code == 200 + return response.json()["accountId"] @pytest.mark.asyncio @@ -865,6 +868,112 @@ async def fake_stream( assert resolved is not None +@pytest.mark.asyncio +async def test_backend_responses_file_pin_does_not_rewrite_existing_thread_row( + async_client, + monkeypatch, +): + from app.dependencies import get_proxy_service_for_app + + thread_owner_chatgpt_id = "acc_file_pin_thread_owner" + file_owner_chatgpt_id = "acc_file_pin_file_owner" + thread_owner_id = await _import_account( + async_client, + thread_owner_chatgpt_id, + "file-pin-thread-owner@example.com", + ) + file_owner_id = await _import_account( + async_client, + file_owner_chatgpt_id, + "file-pin-file-owner@example.com", + ) + process_session = "file-pin-process" + thread_headers = {"session-id": process_session, "thread-id": "file-pin-thread"} + sibling_headers = {"session-id": process_session, "thread-id": "file-pin-sibling"} + thread_key = _codex_backend_identity(thread_headers).thread_selection_key + sibling_key = _codex_backend_identity(sibling_headers).thread_selection_key + process_key = _codex_session_selection_key(process_session) + assert thread_key is not None + assert sibling_key is not None + + async with SessionLocal() as session: + await StickySessionsRepository(session).upsert( + thread_key, + thread_owner_id, + kind=StickySessionKind.PROMPT_CACHE, + ) + + service = get_proxy_service_for_app(async_client._transport.app) + await service._pin_file_account("file_thread_locality", file_owner_id) + seen: list[str] = [] + + async def fake_stream(payload, headers, access_token, account_id, **kwargs): + del payload, headers, access_token, kwargs + seen.append(account_id) + yield 'data: {"type":"response.completed","response":{"id":"resp_file_pin_thread"}}\n\n' + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + + pinned_response = await async_client.post( + "/backend-api/codex/responses", + headers=thread_headers, + json={ + "model": "gpt-5.2", + "instructions": "You are a helpful assistant.", + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Read the file."}, + {"type": "input_file", "file_id": "file_thread_locality"}, + ], + } + ], + "stream": True, + }, + ) + assert pinned_response.status_code == 200 + assert seen == [file_owner_chatgpt_id] + + async with SessionLocal() as session: + repo = StickySessionsRepository(session) + assert await repo.get_account_id(thread_key, kind=StickySessionKind.PROMPT_CACHE) == thread_owner_id + assert await repo.get_account_id(process_key, kind=StickySessionKind.CODEX_SESSION) == file_owner_id + assert await repo.get_account_id(sibling_key, kind=StickySessionKind.PROMPT_CACHE) is None + + unpinned_response = await async_client.post( + "/backend-api/codex/responses", + headers=thread_headers, + json={ + "model": "gpt-5.2", + "instructions": "You are a helpful assistant.", + "input": "Continue without the file.", + "stream": True, + }, + ) + assert unpinned_response.status_code == 200 + assert seen == [file_owner_chatgpt_id, thread_owner_chatgpt_id] + + sibling_response = await async_client.post( + "/backend-api/codex/responses", + headers=sibling_headers, + json={ + "model": "gpt-5.2", + "instructions": "You are a helpful assistant.", + "input": "Sibling thread without a file.", + "stream": True, + }, + ) + assert sibling_response.status_code == 200 + assert seen == [file_owner_chatgpt_id, thread_owner_chatgpt_id, file_owner_chatgpt_id] + + async with SessionLocal() as session: + repo = StickySessionsRepository(session) + assert await repo.get_account_id(thread_key, kind=StickySessionKind.PROMPT_CACHE) == thread_owner_id + assert await repo.get_account_id(process_key, kind=StickySessionKind.CODEX_SESSION) == file_owner_id + assert await repo.get_account_id(sibling_key, kind=StickySessionKind.PROMPT_CACHE) == file_owner_id + + @pytest.mark.asyncio async def test_derived_prompt_cache_key_does_not_block_file_id_pin(async_client): """Regression: a ``prompt_cache_key`` that the proxy itself derived diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index dcb1bfab87..f62cc99f54 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -29,7 +29,11 @@ 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_backend_identity, _codex_session_selection_key +from app.modules.proxy.affinity import ( + _AffinityPolicy, + _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 @@ -3023,6 +3027,106 @@ async def test_first_codex_thread_initializes_process_preference_once_for_later_ ] +@pytest.mark.asyncio +async def test_required_file_owner_does_not_rewrite_existing_thread_row() -> None: + balancer, thread_owner, file_owner, sticky_repo = _make_cap_spillover_balancer("file-pin-thread") + assert file_owner is not None + process_session = "file-pin-process" + thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "file-pin-thread"} + ).thread_selection_key + assert thread_key is not None + sticky_repo.account_ids_by_key = {thread_key: thread_owner.id} + preferred = _AffinityPolicy.preferred_owner_sticky_inputs( + thread_key, + StickySessionKind.PROMPT_CACHE, + False, + 300, + "thread_header", + process_session, + ) + + selected = await balancer.select_account( + sticky_key=preferred[0], + sticky_kind=preferred[1], + reallocate_sticky=preferred[2], + sticky_max_age_seconds=preferred[3], + sticky_source=preferred[4], + legacy_sticky_key=preferred[5], + required_account_id=file_owner.id, + required_account_is_ownership_constraint=True, + routing_strategy="usage_weighted", + lease_kind="stream", + ) + + assert selected.account is not None + assert selected.account.id == file_owner.id + assert sticky_repo.account_ids_by_key == {thread_key: thread_owner.id} + assert sticky_repo.upserts == [] + await balancer.release_account_lease(selected.lease) + + +@pytest.mark.asyncio +async def test_required_file_owner_seeds_process_preference_for_later_sibling() -> None: + balancer, thread_owner, file_owner, sticky_repo = _make_cap_spillover_balancer("file-pin-seed") + assert file_owner is not None + process_session = "file-pin-seed-process" + process_key = _codex_session_selection_key(process_session) + first_thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "file-pin-first"} + ).thread_selection_key + sibling_thread_key = _codex_backend_identity( + {"session-id": process_session, "thread-id": "file-pin-sibling"} + ).thread_selection_key + assert first_thread_key is not None + assert sibling_thread_key is not None + sticky_repo.account_ids_by_key = {} + preferred = _AffinityPolicy.preferred_owner_sticky_inputs( + first_thread_key, + StickySessionKind.PROMPT_CACHE, + False, + 300, + "thread_header", + process_session, + ) + + first = await balancer.select_account( + sticky_key=preferred[0], + sticky_kind=preferred[1], + reallocate_sticky=preferred[2], + sticky_max_age_seconds=preferred[3], + sticky_source=preferred[4], + legacy_sticky_key=preferred[5], + sticky_seed_key=process_key, + sticky_seed_kind=StickySessionKind.CODEX_SESSION, + required_account_id=file_owner.id, + required_account_is_ownership_constraint=True, + routing_strategy="usage_weighted", + lease_kind="stream", + ) + assert first.account is not None + assert first.account.id == file_owner.id + assert first_thread_key not in (sticky_repo.account_ids_by_key or {}) + assert sticky_repo.account_ids_by_key == {process_key: file_owner.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", + lease_kind="stream", + ) + assert sibling.account is not None + assert sibling.account.id == file_owner.id + assert sticky_repo.account_ids_by_key[process_key] == file_owner.id + await balancer.release_account_lease(first.lease) + await balancer.release_account_lease(sibling.lease) + + @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") diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 43a1f8c887..fa5d129127 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -38130,7 +38130,10 @@ async def test_select_account_with_budget_keeps_thread_seed_for_first_exact_owne 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_key"] is None + assert select_account.await_args.kwargs["sticky_kind"] == proxy_service.StickySessionKind.CODEX_SESSION + assert select_account.await_args.kwargs["sticky_source"] == "thread_header" + assert select_account.await_args.kwargs["legacy_sticky_key"] == "process-first-exact-owner" 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