From 61dc17a5ae6ed90ad1031157f125ae4577269d3a Mon Sep 17 00:00:00 2001 From: Darafei Praliaskouski Date: Wed, 19 Aug 2026 19:17:24 +0400 Subject: [PATCH] fix(proxy): keep hard continuity owners ahead of affinity Rebased onto current main; squashed the branch's review-iteration commits. The identity-index revision was relinked onto the current Alembic head, the live-usage resolver imports main's header rewrite dropped were restored, and the telemetry migration test keeps main's head lookup instead of pinning one revision id. --- ...000_add_accounts_chatgpt_identity_index.py | 32 ++ app/db/models.py | 1 + app/modules/proxy/_service/compact.py | 3 +- app/modules/proxy/_service/streaming/retry.py | 38 ++- app/modules/usage/live_ingest.py | 91 ++++- .../.openspec.yaml | 2 + .../proposal.md | 35 ++ .../specs/live-usage-ingestion/spec.md | 80 +++++ .../tasks.md | 22 ++ tests/integration/test_live_usage_ingest.py | 322 +++++++++++++++++- tests/integration/test_proxy_compact.py | 105 +++++- tests/integration/test_proxy_responses.py | 138 +++++++- tests/unit/test_db_migrate.py | 6 +- tests/unit/test_live_usage_ingest.py | 27 ++ tests/unit/test_proxy_utils.py | 26 +- 15 files changed, 909 insertions(+), 19 deletions(-) create mode 100644 app/db/alembic/versions/20260814_000000_add_accounts_chatgpt_identity_index.py create mode 100644 openspec/changes/normalize-live-usage-account-identity/.openspec.yaml create mode 100644 openspec/changes/normalize-live-usage-account-identity/proposal.md create mode 100644 openspec/changes/normalize-live-usage-account-identity/specs/live-usage-ingestion/spec.md create mode 100644 openspec/changes/normalize-live-usage-account-identity/tasks.md diff --git a/app/db/alembic/versions/20260814_000000_add_accounts_chatgpt_identity_index.py b/app/db/alembic/versions/20260814_000000_add_accounts_chatgpt_identity_index.py new file mode 100644 index 0000000000..d083490764 --- /dev/null +++ b/app/db/alembic/versions/20260814_000000_add_accounts_chatgpt_identity_index.py @@ -0,0 +1,32 @@ +"""add account ChatGPT identity lookup index + +Revision ID: 20260814_000000_add_accounts_chatgpt_identity_index +Revises: 20260806_030000_add_api_key_allowed_reasoning_efforts +Create Date: 2026-08-14 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "20260814_000000_add_accounts_chatgpt_identity_index" +down_revision = "20260806_030000_add_api_key_allowed_reasoning_efforts" +branch_labels = None +depends_on = None + +_INDEX = "idx_accounts_chatgpt_account_id" +_TABLE = "accounts" + + +def upgrade() -> None: + bind = op.get_bind() + if bind.dialect.name == "postgresql": + with op.get_context().autocommit_block(): + op.execute(sa.text(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX} ON {_TABLE} (chatgpt_account_id)")) + else: + op.execute(sa.text(f"CREATE INDEX IF NOT EXISTS {_INDEX} ON {_TABLE} (chatgpt_account_id)")) + + +def downgrade() -> None: + op.drop_index(_INDEX, table_name=_TABLE, if_exists=True) diff --git a/app/db/models.py b/app/db/models.py index f5634cdb8a..37c3d6dab9 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -2154,6 +2154,7 @@ class HttpBridgeRetryCircuit(Base): postgresql_where=text("delete_requested_at IS NOT NULL"), sqlite_where=text("delete_requested_at IS NOT NULL"), ) +Index("idx_accounts_chatgpt_account_id", Account.chatgpt_account_id) Index("idx_api_keys_name", ApiKey.name) Index("idx_logs_account_time", RequestLog.account_id, RequestLog.requested_at) Index("idx_logs_model_source_time", RequestLog.model_source_id, RequestLog.requested_at) diff --git a/app/modules/proxy/_service/compact.py b/app/modules/proxy/_service/compact.py index 682cfb4bb9..79b495d7ee 100644 --- a/app/modules/proxy/_service/compact.py +++ b/app/modules/proxy/_service/compact.py @@ -731,12 +731,11 @@ async def _resolve_compact_turn_state_owner( if owner_account_id == resolved_owner and session_identity is not None } if len(session_identities) > 1: - sources = ", ".join(source for source, _account_id, _session_id in owner_refs) raise ProxyResponseError( 502, openai_error( "continuity_owner_conflict", - f"Account-owned continuity sources conflict ({sources}); retry the logical turn.", + "Turn-state owner sessions conflict; retry the logical turn.", error_type="server_error", ), ) diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 0b28853a44..4566c751d0 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -1376,16 +1376,27 @@ async def _retry_account_model_rejection( ) return if require_preferred_account and preferred_account_id is not None: + error_code = "previous_response_owner_unavailable" message = "Previous response owner account is unavailable; retry later." + reason = "owner_account_unavailable" + upstream_error_code = "no_accounts" + if selection.error_code == "continuity_owner_conflict": + error_code = "continuity_owner_conflict" + message = ( + selection.error_message + or "Account-owned continuity sources conflict; retry the logical turn" + ) + reason = "owner_conflict" + upstream_error_code = selection.error_code _record_continuity_fail_closed( surface="http_stream", - reason="owner_account_unavailable", + reason=reason, previous_response_id=payload.previous_response_id, session_id=headers.get("x-codex-turn-state") or headers.get("session_id"), - upstream_error_code="no_accounts", + upstream_error_code=upstream_error_code, ) event = response_failed_event( - "previous_response_owner_unavailable", + error_code, message, response_id=request_id, ) @@ -1397,7 +1408,7 @@ async def _retry_account_model_rejection( model=payload.model, latency_ms=int((time.monotonic() - start) * 1000), status="error", - error_code="previous_response_owner_unavailable", + error_code=error_code, error_message=message, reasoning_effort=payload.reasoning.effort if payload.reasoning else None, transport=request_transport, @@ -1516,16 +1527,27 @@ async def _retry_account_model_rejection( request_id, ) else: + error_code = "previous_response_owner_unavailable" message = "Previous response owner account is unavailable; retry later." + reason = "owner_account_unavailable" + upstream_error_code = "upstream_unavailable" + if selection.error_code == "continuity_owner_conflict": + error_code = "continuity_owner_conflict" + message = ( + selection.error_message + or "Account-owned continuity sources conflict; retry the logical turn" + ) + reason = "owner_conflict" + upstream_error_code = selection.error_code _record_continuity_fail_closed( surface="http_stream", - reason="owner_account_unavailable", + reason=reason, previous_response_id=payload.previous_response_id, session_id=headers.get("x-codex-turn-state") or headers.get("session_id"), - upstream_error_code="upstream_unavailable", + upstream_error_code=upstream_error_code, ) event = response_failed_event( - "previous_response_owner_unavailable", + error_code, message, response_id=request_id, ) @@ -1537,7 +1559,7 @@ async def _retry_account_model_rejection( model=payload.model, latency_ms=int((time.monotonic() - start) * 1000), status="error", - error_code="previous_response_owner_unavailable", + error_code=error_code, error_message=message, reasoning_effort=payload.reasoning.effort if payload.reasoning else None, transport=request_transport, diff --git a/app/modules/usage/live_ingest.py b/app/modules/usage/live_ingest.py index af4ecba115..485145a786 100644 --- a/app/modules/usage/live_ingest.py +++ b/app/modules/usage/live_ingest.py @@ -6,10 +6,13 @@ import weakref from dataclasses import dataclass +from sqlalchemy import select + from app.core import usage as usage_core from app.core.config.settings import get_settings from app.core.usage.live_hub import register_live_usage_publisher from app.core.usage.live_snapshots import LiveRateLimitSnapshot, LiveUsageWindow +from app.db.models import Account from app.db.session import get_background_session from app.modules.proxy.account_cache import get_account_selection_cache from app.modules.proxy.rate_limit_cache import get_rate_limit_headers_cache @@ -17,6 +20,8 @@ logger = logging.getLogger(__name__) +_RESOLUTION_TTL_SECONDS = 300.0 + # Write-coalescing tuning (fixed; issue #1340 / PRINCIPLES.md P2). The # ingestor keeps both as constructor fields so tests can exercise queue # overflow and coalescing with small values. @@ -108,6 +113,8 @@ def __init__( self._queue: asyncio.Queue[_QueuedSnapshot] = asyncio.Queue(maxsize=max(1, queue_size)) self._write_min_interval_seconds = write_min_interval_seconds self._last_write: dict[str, tuple[tuple[object, ...], float]] = {} + self._resolution_cache: dict[str, tuple[str | None, float]] = {} + self._resolution_aliases: dict[str, tuple[str, str | None]] = {} self._consumer: asyncio.Task[None] | None = None self._dropped = 0 self._last_cache_invalidation = 0.0 @@ -121,7 +128,13 @@ def publish( chatgpt_account_id: str | None = None, ) -> None: item = _QueuedSnapshot(account_id=account_id, chatgpt_account_id=chatgpt_account_id, snapshot=snapshot) - if account_id is not None and self._should_skip(account_id, snapshot): + coalesce_account_id = account_id + if account_id is not None: + alias = self._resolution_aliases.get(account_id) + if alias is not None and chatgpt_account_id is None: + alias_account_id, _alias_chatgpt_account_id = alias + coalesce_account_id = alias_account_id + if coalesce_account_id is not None and self._should_skip(coalesce_account_id, snapshot): return try: self._queue.put_nowait(item) @@ -185,6 +198,15 @@ async def _run(self) -> None: ) async def _ingest(self, item: _QueuedSnapshot) -> None: + raw_account_id = item.account_id + account_id = await self._resolve_persisted_account_id(raw_account_id, item.chatgpt_account_id) + if account_id is None: + return + if raw_account_id and raw_account_id != account_id: + self._resolution_aliases[raw_account_id] = (account_id, item.chatgpt_account_id) + if self._should_skip(account_id, item.snapshot): + return + snapshot = item.snapshot primary = snapshot.primary secondary = snapshot.secondary @@ -281,6 +303,73 @@ async def _invalidate_caches_now(self) -> None: # values before the TTL expires. await get_rate_limit_headers_cache().invalidate() + async def _resolve_persisted_account_id( + self, + account_id: str | None, + chatgpt_account_id: str | None, + ) -> str | None: + if account_id is not None: + exact = await self._resolve_account_id_by_id(account_id) + if exact is not None: + if chatgpt_account_id: + resolved = await self._resolve_account_id(chatgpt_account_id) + if resolved is not None and exact != resolved: + return None + return exact + resolved = await self._resolve_account_id(account_id) + if resolved is not None: + if chatgpt_account_id and chatgpt_account_id != account_id: + chatgpt_resolved = await self._resolve_account_id(chatgpt_account_id) + if chatgpt_resolved is not None and chatgpt_resolved != resolved: + return None + return resolved + return await self._resolve_account_id(chatgpt_account_id) + + async def _resolve_account_id_by_id(self, account_id: str | None) -> str | None: + if not account_id: + return None + cache_key = f"id:{account_id}" + cached = self._resolution_cache.get(cache_key) + now = time.monotonic() + if cached is not None and now - cached[1] < _RESOLUTION_TTL_SECONDS: + cached_account_id = cached[0] + if cached_account_id is None: + return None + if await self._account_id_exists(cached_account_id): + return cached_account_id + self._resolution_cache.pop(cache_key, None) + async with get_background_session() as session: + resolved = await session.scalar(select(Account.id).where(Account.id == account_id)) + if not isinstance(resolved, str): + resolved = None + self._resolution_cache[cache_key] = (resolved, now) + return resolved + + async def _resolve_account_id_by_account_id(self, account_id: str) -> str | None: + return await self._resolve_account_id_by_id(account_id) + + async def _resolve_account_id(self, chatgpt_account_id: str | None) -> str | None: + if not chatgpt_account_id: + return None + cache_key = f"chatgpt:{chatgpt_account_id}" + now = time.monotonic() + async with get_background_session() as session: + rows = ( + (await session.execute(select(Account.id).where(Account.chatgpt_account_id == chatgpt_account_id))) + .scalars() + .all() + ) + # Ambiguous identities (multiple workspace slots) are dropped rather + # than guessed; the poller stays authoritative for them. + resolved = rows[0] if len(rows) == 1 else None + self._resolution_cache[cache_key] = (resolved, now) + return resolved + + async def _account_id_exists(self, account_id: str) -> bool: + async with get_background_session() as session: + resolved = await session.scalar(select(Account.id).where(Account.id == account_id)) + return isinstance(resolved, str) + _ingestor: LiveUsageIngestor | None = None # Registrations a nested startup displaced, innermost-last. A stack rather diff --git a/openspec/changes/normalize-live-usage-account-identity/.openspec.yaml b/openspec/changes/normalize-live-usage-account-identity/.openspec.yaml new file mode 100644 index 0000000000..4af864176c --- /dev/null +++ b/openspec/changes/normalize-live-usage-account-identity/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-08-14 diff --git a/openspec/changes/normalize-live-usage-account-identity/proposal.md b/openspec/changes/normalize-live-usage-account-identity/proposal.md new file mode 100644 index 0000000000..29dc92746d --- /dev/null +++ b/openspec/changes/normalize-live-usage-account-identity/proposal.md @@ -0,0 +1,35 @@ +## Why + +Live usage snapshots can arrive with either the internal codex-lb account ID +or the upstream ChatGPT account identity. Some live paths include a +workspace-local suffix on the internal ID, which does not necessarily identify +the persisted codex-lb account row. Guessing by string shape can attribute one +workspace slot's quota state to another slot and break routing decisions. + +## What Changes + +- Resolve supplied internal account IDs only when the exact persisted account + row exists. +- Resolve suffixed or otherwise unknown internal IDs only through a unique + upstream ChatGPT account identity when that identity is also supplied. +- Drop unresolved or ambiguous live snapshots rather than guessing. +- Coalesce later snapshots under the resolved account ID after a raw ID has + been resolved. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `live-usage-ingestion`: live snapshot attribution uses explicit identity + evidence and preserves per-account write coalescing after normalization. + +## Impact + +- Code: live usage ingestor identity resolution and coalescing. +- Tests: integration coverage for hub-published snapshots, unproven suffix + drops, and resolved-alias coalescing. +- Configuration, schema, response shapes, and background polling are unchanged. diff --git a/openspec/changes/normalize-live-usage-account-identity/specs/live-usage-ingestion/spec.md b/openspec/changes/normalize-live-usage-account-identity/specs/live-usage-ingestion/spec.md new file mode 100644 index 0000000000..a92f58ba0c --- /dev/null +++ b/openspec/changes/normalize-live-usage-account-identity/specs/live-usage-ingestion/spec.md @@ -0,0 +1,80 @@ +# live-usage-ingestion Delta + +## ADDED Requirements + +### Requirement: Live usage account attribution requires explicit identity proof + +Live usage ingestion MUST attribute a snapshot to a persisted account row only +when the supplied internal account ID exactly matches that row, or when the +snapshot also supplies a ChatGPT account identity that resolves to exactly one +persisted account row. A workspace-local suffix, shard suffix, or other string +shape on an unknown internal account ID MUST NOT by itself justify stripping or +rewriting the ID. + +When an exact internal account ID and a ChatGPT account identity both resolve +but identify different persisted account rows, the snapshot MUST be dropped +without failing the proxied request. The ingestor MUST NOT let cached ChatGPT +identity resolution attribute a snapshot to a row that no longer owns that +ChatGPT identity, and it MUST revalidate cached exact internal account matches +before using them for attribution. + +When neither identity path resolves to exactly one persisted account row, the +snapshot MUST be dropped without failing the proxied request. After a raw +snapshot identity resolves to a persisted account row, duplicate coalescing +MUST use that resolved account ID so repeated live snapshots do not bypass the +per-account write interval. When a later snapshot supplies a ChatGPT account +identity, the ingestor MUST NOT skip that snapshot using a previously cached +raw-ID alias before revalidating the supplied ChatGPT identity. ChatGPT identity +lookups MUST use an indexed lookup path and only need to inspect enough rows to +distinguish a unique match from ambiguity. + +#### Scenario: Hub-published snapshot resolves through ChatGPT identity + +- **GIVEN** a proxy path publishes a live usage snapshot with a workspace-suffixed internal ID +- **AND** it also supplies a ChatGPT account identity that maps to exactly one persisted account +- **WHEN** the live usage ingestor persists the snapshot +- **THEN** the usage rows are written for the uniquely mapped persisted account + +#### Scenario: Unproven suffix is dropped + +- **GIVEN** a proxy path publishes a live usage snapshot with an internal ID that does not exactly match a persisted account row +- **AND** no unique ChatGPT account identity is supplied for that snapshot +- **WHEN** the live usage ingestor processes it +- **THEN** no usage rows are written for a guessed prefix account +- **AND** the proxied request is not failed + +#### Scenario: Normalized aliases are coalesced + +- **GIVEN** a raw snapshot identity has resolved to a persisted account row +- **WHEN** the same raw identity publishes an unchanged snapshot inside the write coalescing interval +- **THEN** the duplicate snapshot is skipped using the resolved account ID + +#### Scenario: Conflicting identities are dropped + +- **GIVEN** a proxy path publishes a live usage snapshot with an exact internal account ID for one persisted account +- **AND** it also supplies a ChatGPT account identity that uniquely maps to a different persisted account +- **WHEN** the live usage ingestor processes it +- **THEN** no usage rows are written for either conflicting account + +#### Scenario: Moved ChatGPT identities are revalidated + +- **GIVEN** a ChatGPT account identity previously resolved to one persisted account +- **AND** that ChatGPT account identity now belongs to a different persisted account +- **WHEN** a later snapshot uses the same ChatGPT account identity with an unknown internal account ID +- **THEN** the usage rows are written only for the current persisted account that owns the ChatGPT identity + +#### Scenario: Supplied ChatGPT identity bypasses stale raw alias coalescing + +- **GIVEN** a raw snapshot identity previously resolved to a persisted account row +- **AND** an unchanged later snapshot supplies the same raw identity with a ChatGPT account identity +- **WHEN** that ChatGPT account identity now resolves to a different persisted account row +- **THEN** the later snapshot is not skipped by the stale raw identity alias +- **AND** the usage rows are written only for the current persisted account that owns the ChatGPT identity + +#### Scenario: Cached exact account matches are revalidated + +- **GIVEN** an exact internal account ID previously resolved to a persisted account row +- **AND** that persisted account row has since been deleted or rewritten +- **WHEN** a later snapshot uses the same internal account ID +- **THEN** the ingestor revalidates the cached match before attribution +- **AND** the stale cached account ID is not used for a usage-history write diff --git a/openspec/changes/normalize-live-usage-account-identity/tasks.md b/openspec/changes/normalize-live-usage-account-identity/tasks.md new file mode 100644 index 0000000000..3a71f77a8d --- /dev/null +++ b/openspec/changes/normalize-live-usage-account-identity/tasks.md @@ -0,0 +1,22 @@ +## 1. Implementation + +- [x] 1.1 Require exact account-row proof for direct internal account IDs. +- [x] 1.2 Resolve unknown or suffixed internal IDs only through a unique supplied ChatGPT account identity. +- [x] 1.3 Drop unresolved or ambiguous snapshots without failing the serving path. +- [x] 1.4 Coalesce duplicate snapshots under the normalized account ID after resolution. +- [x] 1.5 Revalidate cached exact account IDs and bypass stale raw aliases when a ChatGPT identity is supplied. +- [x] 1.6 Add an indexed, bounded ChatGPT identity lookup path. + +## 2. Regression Coverage + +- [x] 2.1 Cover the proxy hub path that publishes both a raw internal ID and a ChatGPT identity. +- [x] 2.2 Cover that an unproven suffixed internal ID is dropped instead of guessed. +- [x] 2.3 Cover duplicate coalescing after a raw ID is resolved to the stored account. +- [x] 2.4 Cover moved ChatGPT identities, conflicting supplied identities, and deleted cached exact IDs. + +## 3. Validation + +- [x] 3.1 Run focused live usage ingest tests. +- [x] 3.2 Run Ruff on changed Python files. +- [x] 3.3 Run strict OpenSpec validation. +- [x] 3.4 Run migration policy and schema drift validation. diff --git a/tests/integration/test_live_usage_ingest.py b/tests/integration/test_live_usage_ingest.py index 9e16704921..77244e5eac 100644 --- a/tests/integration/test_live_usage_ingest.py +++ b/tests/integration/test_live_usage_ingest.py @@ -1,15 +1,17 @@ from __future__ import annotations import asyncio +import json from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Any, cast import pytest -from sqlalchemy import select +from sqlalchemy import delete, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql.dml import Delete +import app.core.clients.proxy as proxy_client_module from app.core.clients import proxy as core_proxy from app.core.crypto import TokenEncryptor from app.core.openai.requests import ResponsesRequest @@ -840,6 +842,324 @@ async def test_live_ingestor_resolves_chatgpt_account_id(db_setup) -> None: assert {row.window for row in rows} == {"primary", "secondary"} +@pytest.mark.asyncio +async def test_live_ingestor_resolves_workspace_suffixed_account_id(db_setup) -> None: + del db_setup + stored_account_id = "acc_live_workspace" + async with SessionLocal() as session: + await AccountsRepository(session).upsert( + _make_account( + stored_account_id, + "live-workspace-suffix@example.com", + chatgpt_account_id="workspace-live-suffix", + ) + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.start() + try: + live_hub.register_live_usage_publisher(ingestor.publish) + live_hub.publish_live_usage( + _snapshot(), + account_id=f"{stored_account_id}_49115a1d", + chatgpt_account_id="workspace-live-suffix", + ) + primary, secondary = await _wait_for_rows(stored_account_id) + finally: + live_hub.register_live_usage_publisher(None) + await ingestor.stop() + + assert primary is not None + assert primary.account_id == stored_account_id + assert secondary is not None + + +@pytest.mark.asyncio +async def test_live_usage_stream_tap_persists_workspace_suffixed_account_id( + db_setup, + monkeypatch: pytest.MonkeyPatch, +) -> None: + del db_setup + stored_account_id = "acc_live_stream_workspace" + chatgpt_account_id = "workspace-live-stream" + raw_account_id = f"{stored_account_id}_49115a1d" + async with SessionLocal() as session: + await AccountsRepository(session).upsert( + _make_account( + stored_account_id, + "live-stream-workspace@example.com", + chatgpt_account_id=chatgpt_account_id, + ) + ) + + blocks = [ + 'data: {"type":"response.created","response":{"id":"resp_live_stream"}}\n\n', + "data: " + + json.dumps( + { + "type": "codex.rate_limits", + "rate_limits": { + "primary": {"used_percent": 55, "window_minutes": 300, "reset_at": 1700000300}, + "secondary": {"used_percent": 12, "window_minutes": 10080, "reset_at": 1700600000}, + }, + }, + separators=(",", ":"), + ) + + "\n\n", + 'data: {"type":"response.completed","response":{"id":"resp_live_stream"}}\n\n', + ] + + async def fake_stream(**kwargs): + del kwargs + for block in blocks: + yield block + + import contextlib + + @contextlib.asynccontextmanager + async def fake_lease(session=None): + yield session + + monkeypatch.setattr(proxy_client_module, "_stream_responses_with_session", lambda **kwargs: fake_stream(**kwargs)) + monkeypatch.setattr(proxy_client_module, "lease_http_session", fake_lease) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.start() + try: + live_hub.register_live_usage_publisher(ingestor.publish) + request = proxy_client_module.ResponsesRequest.model_validate( + {"model": "gpt-5.1", "instructions": "hi", "input": "live", "stream": True} + ) + seen = [ + block + async for block in proxy_client_module.stream_responses( + request, + {}, + "access-token", + chatgpt_account_id, + codex_lb_account_id=raw_account_id, + ) + ] + primary, secondary = await _wait_for_rows(stored_account_id) + finally: + live_hub.register_live_usage_publisher(None) + await ingestor.stop() + + assert seen == blocks + assert primary is not None + assert primary.account_id == stored_account_id + assert secondary is not None + + +@pytest.mark.asyncio +async def test_live_ingestor_drops_unproven_workspace_suffix(db_setup) -> None: + del db_setup + stored_account_id = "acc_live_workspace_a" + async with SessionLocal() as session: + await AccountsRepository(session).upsert( + _make_account( + stored_account_id, + "live-workspace-a@example.com", + chatgpt_account_id="workspace-live-a", + ) + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.start() + try: + ingestor.publish(_snapshot(), account_id=f"{stored_account_id}_49115a1d") + primary, secondary = await _wait_for_rows(stored_account_id, timeout=0.3) + finally: + await ingestor.stop() + + assert primary is None + assert secondary is None + + +@pytest.mark.asyncio +async def test_live_ingestor_coalesces_resolved_account_alias(db_setup) -> None: + del db_setup + stored_account_id = "acc_live_alias" + raw_account_id = f"{stored_account_id}_49115a1d" + async with SessionLocal() as session: + await AccountsRepository(session).upsert( + _make_account(stored_account_id, "live-alias@example.com", chatgpt_account_id="workspace-live-alias") + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=30.0) + ingestor.start() + try: + ingestor.publish(_snapshot(), account_id=raw_account_id, chatgpt_account_id="workspace-live-alias") + primary, secondary = await _wait_for_rows(stored_account_id) + assert primary is not None and secondary is not None + + ingestor.publish(_snapshot(), account_id=raw_account_id) + await asyncio.sleep(0) + assert ingestor._queue.qsize() == 0 + finally: + await ingestor.stop() + + +@pytest.mark.asyncio +async def test_live_ingestor_does_not_alias_coalesce_different_chatgpt_identity(db_setup) -> None: + del db_setup + raw_account_id = "stale-local-slot_49115a1d" + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account("acc_live_alias_a", "live-alias-a@example.com", chatgpt_account_id="workspace-alias-a") + ) + await repo.upsert( + _make_account("acc_live_alias_b", "live-alias-b@example.com", chatgpt_account_id="workspace-alias-b") + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=30.0) + ingestor.start() + try: + ingestor.publish(_snapshot(), account_id=raw_account_id, chatgpt_account_id="workspace-alias-a") + primary_a, secondary_a = await _wait_for_rows("acc_live_alias_a") + assert primary_a is not None and secondary_a is not None + + ingestor.publish(_snapshot(), account_id=raw_account_id, chatgpt_account_id="workspace-alias-b") + primary_b, secondary_b = await _wait_for_rows("acc_live_alias_b") + finally: + await ingestor.stop() + + assert primary_b is not None + assert primary_b.account_id == "acc_live_alias_b" + assert secondary_b is not None + + +@pytest.mark.asyncio +async def test_live_ingestor_drops_conflicting_exact_and_chatgpt_identity(db_setup) -> None: + del db_setup + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account("acc_live_conflict_a", "live-conflict-a@example.com", chatgpt_account_id="workspace-a") + ) + await repo.upsert( + _make_account("acc_live_conflict_b", "live-conflict-b@example.com", chatgpt_account_id="workspace-b") + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.start() + try: + ingestor.publish(_snapshot(), account_id="acc_live_conflict_a", chatgpt_account_id="workspace-b") + primary_a, secondary_a = await _wait_for_rows("acc_live_conflict_a", timeout=0.3) + primary_b, secondary_b = await _wait_for_rows("acc_live_conflict_b", timeout=0.3) + finally: + await ingestor.stop() + + assert primary_a is None + assert secondary_a is None + assert primary_b is None + assert secondary_b is None + + +@pytest.mark.asyncio +async def test_live_ingestor_revalidates_moved_chatgpt_identity(db_setup) -> None: + del db_setup + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account("acc_live_move_a", "live-move-a@example.com", chatgpt_account_id="workspace-moving") + ) + await repo.upsert(_make_account("acc_live_move_b", "live-move-b@example.com")) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.start() + try: + ingestor.publish(_snapshot(), account_id="stale-local-slot_49115a1d", chatgpt_account_id="workspace-moving") + primary_a, secondary_a = await _wait_for_rows("acc_live_move_a") + assert primary_a is not None and secondary_a is not None + + async with SessionLocal() as session: + await session.execute( + update(Account).where(Account.id == "acc_live_move_a").values(chatgpt_account_id=None) + ) + await session.execute( + update(Account).where(Account.id == "acc_live_move_b").values(chatgpt_account_id="workspace-moving") + ) + await session.commit() + + moved = _snapshot() + assert moved.primary is not None + moved = LiveRateLimitSnapshot( + primary=LiveUsageWindow( + used_percent=primary_a.used_percent + 7.0, + window_minutes=moved.primary.window_minutes, + reset_at=moved.primary.reset_at + 60 if moved.primary.reset_at is not None else None, + ), + secondary=moved.secondary, + credits_has=moved.credits_has, + credits_unlimited=moved.credits_unlimited, + credits_balance=moved.credits_balance, + ) + ingestor.publish(moved, account_id="stale-local-slot_49115a1d", chatgpt_account_id="workspace-moving") + primary_b, secondary_b = await _wait_for_rows("acc_live_move_b") + finally: + await ingestor.stop() + + assert primary_b is not None + assert primary_b.account_id == "acc_live_move_b" + assert secondary_b is not None + + +@pytest.mark.asyncio +async def test_live_ingestor_revalidates_moved_chatgpt_identity_before_alias_coalesce(db_setup) -> None: + del db_setup + raw_account_id = "moved-local-slot_49115a1d" + snapshot = _snapshot() + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account("acc_live_move_same_a", "live-move-same-a@example.com", chatgpt_account_id="workspace-same") + ) + await repo.upsert(_make_account("acc_live_move_same_b", "live-move-same-b@example.com")) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=30.0) + ingestor.start() + try: + ingestor.publish(snapshot, account_id=raw_account_id, chatgpt_account_id="workspace-same") + primary_a, secondary_a = await _wait_for_rows("acc_live_move_same_a") + assert primary_a is not None and secondary_a is not None + + async with SessionLocal() as session: + await session.execute( + update(Account).where(Account.id == "acc_live_move_same_a").values(chatgpt_account_id=None) + ) + await session.execute( + update(Account).where(Account.id == "acc_live_move_same_b").values(chatgpt_account_id="workspace-same") + ) + await session.commit() + + ingestor.publish(snapshot, account_id=raw_account_id, chatgpt_account_id="workspace-same") + primary_b, secondary_b = await _wait_for_rows("acc_live_move_same_b") + finally: + await ingestor.stop() + + assert primary_b is not None + assert primary_b.account_id == "acc_live_move_same_b" + assert secondary_b is not None + + +@pytest.mark.asyncio +async def test_live_ingestor_revalidates_cached_exact_account_id(db_setup) -> None: + del db_setup + async with SessionLocal() as session: + await AccountsRepository(session).upsert(_make_account("acc_live_deleted", "live-deleted@example.com")) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + assert await ingestor._resolve_account_id_by_account_id("acc_live_deleted") == "acc_live_deleted" + + async with SessionLocal() as session: + await session.execute(delete(Account).where(Account.id == "acc_live_deleted")) + await session.commit() + + assert await ingestor._resolve_account_id_by_account_id("acc_live_deleted") is None + + @pytest.mark.asyncio async def test_live_ingestion_kill_switch_disables_publishing(monkeypatch, db_setup) -> None: del db_setup diff --git a/tests/integration/test_proxy_compact.py b/tests/integration/test_proxy_compact.py index 6bb285ca26..d1737c3a5b 100644 --- a/tests/integration/test_proxy_compact.py +++ b/tests/integration/test_proxy_compact.py @@ -4,6 +4,7 @@ import json import logging from datetime import timedelta, timezone +from types import SimpleNamespace from typing import cast from unittest.mock import AsyncMock @@ -18,7 +19,7 @@ from app.core.openai.models import CompactResponsePayload, OpenAIResponsePayload from app.core.openai.requests import ResponsesCompactRequest from app.core.utils.time import utcnow -from app.db.models import Account, AccountStatus +from app.db.models import Account, AccountStatus, StickySessionKind from app.db.session import SessionLocal from app.modules.api_keys.repository import ApiKeysRepository from app.modules.api_keys.service import ApiKeyCreateData, ApiKeysService @@ -613,6 +614,108 @@ async def fake_compact(payload, headers, access_token, account_id): assert "file_compact_trimmed_optional" not in json.dumps(upstream_payload) +@pytest.mark.asyncio +async def test_proxy_compact_route_rejects_live_durable_turn_state_session_drift(async_client, monkeypatch): + from app.dependencies import get_proxy_service_for_app + + email = "compact-turn-state-drift-route@example.com" + raw_account_id = "acc_compact_turn_state_drift_route" + auth_json = _make_auth_json(raw_account_id, email) + response = await async_client.post( + "/api/accounts/import", + files={"auth_json": ("auth.json", json.dumps(auth_json), "application/json")}, + ) + assert response.status_code == 200 + account_id = generate_unique_account_id(raw_account_id, email) + + service = get_proxy_service_for_app(async_client._transport.app) + turn_state = "turn-compact-route-session-drift" + bridge_key = proxy_module._HTTPBridgeSessionKey("session_header", "compact-route-live-session", None) + service._http_bridge_turn_state_index[(turn_state, None)] = bridge_key + service._http_bridge_sessions[bridge_key] = SimpleNamespace( + key=bridge_key, + account=SimpleNamespace(id=account_id), + durable_session_id="durable-live-session", + ) + monkeypatch.setattr( + service._durable_bridge, + "lookup_turn_state_target", + AsyncMock(return_value=SimpleNamespace(account_id=account_id, session_id="durable-other-session")), + ) + compact = AsyncMock( + return_value=CompactResponsePayload.model_validate({"object": "response.compaction", "output": []}) + ) + monkeypatch.setattr(proxy_module, "core_compact_responses", compact) + + try: + response = await async_client.post( + "/backend-api/codex/responses/compact", + headers={"x-codex-turn-state": turn_state}, + json={"model": "gpt-5.1", "instructions": "hi", "input": []}, + ) + finally: + service._http_bridge_turn_state_index.pop((turn_state, None), None) + service._http_bridge_sessions.pop(bridge_key, None) + + assert response.status_code == 502 + assert response.json()["error"]["code"] == "continuity_owner_conflict" + compact.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_proxy_compact_route_preserves_legacy_raw_owner_conflict(async_client, monkeypatch): + from app.modules.proxy.sticky_repository import StickySessionsRepository + + raw_owner_id = "acc_compact_raw_owner_route" + raw_owner_email = "compact-raw-owner-route@example.com" + previous_owner_id = "acc_compact_previous_owner_route" + previous_owner_email = "compact-previous-owner-route@example.com" + for account_id, email in ( + (raw_owner_id, raw_owner_email), + (previous_owner_id, previous_owner_email), + ): + response = await async_client.post( + "/api/accounts/import", + files={"auth_json": ("auth.json", json.dumps(_make_auth_json(account_id, email)), "application/json")}, + ) + assert response.status_code == 200 + raw_owner_account_id = generate_unique_account_id(raw_owner_id, raw_owner_email) + previous_owner_account_id = generate_unique_account_id(previous_owner_id, previous_owner_email) + raw_session = "compact-legacy-raw-conflict" + + async with SessionLocal() as session: + await StickySessionsRepository(session).upsert( + raw_session, + raw_owner_account_id, + kind=StickySessionKind.CODEX_SESSION, + ) + + async def fake_previous_owner(self, *, previous_response_id, api_key, session_id=None, surface): + del self, previous_response_id, api_key, session_id, surface + return previous_owner_account_id + + compact = AsyncMock( + return_value=CompactResponsePayload.model_validate({"object": "response.compaction", "output": []}) + ) + monkeypatch.setattr(proxy_module.ProxyService, "_resolve_websocket_previous_response_owner", fake_previous_owner) + monkeypatch.setattr(proxy_module, "core_compact_responses", compact) + + response = await async_client.post( + "/backend-api/codex/responses/compact", + headers={"session_id": raw_session}, + json={ + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "previous_response_id": "resp_compact_previous_owner_route", + }, + ) + + assert response.status_code == 503 + assert response.json()["error"]["code"] == "continuity_owner_conflict" + compact.assert_not_awaited() + + @pytest.mark.asyncio async def test_proxy_compact_normalizes_summary_output_for_codex_remote_v2(async_client, monkeypatch): email = "compact-v2-summary@example.com" diff --git a/tests/integration/test_proxy_responses.py b/tests/integration/test_proxy_responses.py index 6d373aeea2..f4c355ec8e 100644 --- a/tests/integration/test_proxy_responses.py +++ b/tests/integration/test_proxy_responses.py @@ -19,7 +19,7 @@ from app.core.openai.models import CompactResponsePayload from app.core.types import JsonValue from app.core.utils.time import utcnow -from app.db.models import Account, DashboardSettings, RequestLog +from app.db.models import Account, DashboardSettings, RequestLog, StickySessionKind from app.db.session import SessionLocal from app.modules.api_keys.service import ApiKeyUsageReservationData from app.modules.proxy._service.streaming import retry as streaming_retry_module @@ -1371,6 +1371,142 @@ async def fake_resolve_owner(self, *, previous_response_id, api_key, session_id, assert response.json()["error"]["message"] == "Previous response owner account is unavailable; retry later." +@pytest.mark.asyncio +async def test_v1_responses_preserves_legacy_raw_owner_conflict(async_client, monkeypatch): + from app.modules.proxy.sticky_repository import StickySessionsRepository + + raw_owner_id = "acc_prev_http_raw_owner" + raw_owner_email = "prev-http-raw-owner@example.com" + previous_owner_id = "acc_prev_http_previous_owner" + previous_owner_email = "prev-http-previous-owner@example.com" + for account_id, email in ( + (raw_owner_id, raw_owner_email), + (previous_owner_id, previous_owner_email), + ): + 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 + raw_owner_account_id = generate_unique_account_id(raw_owner_id, raw_owner_email) + previous_owner_account_id = generate_unique_account_id(previous_owner_id, previous_owner_email) + raw_session = "stream-legacy-raw-conflict" + + async with SessionLocal() as session: + await StickySessionsRepository(session).upsert( + raw_session, + raw_owner_account_id, + kind=StickySessionKind.CODEX_SESSION, + ) + + async def fake_stream( + payload, + headers, + access_token, + account_id, + base_url=None, + raise_for_status=False, + **kwargs, + ): + del payload, headers, access_token, account_id, base_url, raise_for_status, kwargs + raise AssertionError("legacy raw continuity conflict must fail before upstream stream attempt") + if False: + yield "" + + async def fake_resolve_owner(self, *, previous_response_id, api_key, session_id, surface): + del self, previous_response_id, api_key, session_id, surface + return previous_owner_account_id + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + monkeypatch.setattr(proxy_module.ProxyService, "_resolve_websocket_previous_response_owner", fake_resolve_owner) + + async with async_client.stream( + "POST", + "/backend-api/codex/responses", + json={ + "model": "gpt-5.1", + "input": "continue", + "previous_response_id": "resp_prev_http_raw_conflict", + "stream": True, + }, + headers={"session_id": raw_session}, + ) as response: + assert response.status_code == 200 + lines = [line async for line in response.aiter_lines() if line] + + event = _extract_first_event(lines) + assert event["type"] == "response.failed" + assert event["response"]["error"]["code"] == "continuity_owner_conflict" + + +@pytest.mark.asyncio +async def test_v1_responses_preserves_turn_state_owner_proof(async_client, monkeypatch): + from app.modules.proxy.sticky_repository import StickySessionsRepository + + owner_id = "acc_prev_http_turn_owner" + owner_email = "prev-http-turn-owner@example.com" + other_id = "acc_prev_http_other_owner" + other_email = "prev-http-other-owner@example.com" + for account_id, email in ( + (owner_id, owner_email), + (other_id, other_email), + ): + 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 + turn_state = "http_turn_state_owner_proof" + async with SessionLocal() as session: + owner_account_id = await session.scalar(select(Account.id).where(Account.email == owner_email)) + assert isinstance(owner_account_id, str) + await StickySessionsRepository(session).upsert( + turn_state, + owner_account_id, + kind=StickySessionKind.CODEX_SESSION, + ) + + selected_ids: list[str | None] = [] + + 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 + selected_ids.append(account_id) + yield 'data: {"type":"response.created","response":{"id":"resp_turn_owner"}}\n\n' + yield 'data: {"type":"response.completed","response":{"id":"resp_turn_owner"}}\n\n' + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + + async def fake_turn_state_owner(self, *, turn_state, api_key, fail_on_missing): + del self, turn_state, api_key, fail_on_missing + return owner_account_id + + monkeypatch.setattr(proxy_module.ProxyService, "_resolve_compact_turn_state_owner", fake_turn_state_owner) + + async with async_client.stream( + "POST", + "/backend-api/codex/responses", + json={ + "model": "gpt-5.1", + "input": "continue", + "conversation": "conv-turn-owner-proof", + "stream": True, + }, + headers={"x-codex-turn-state": turn_state}, + ) as response: + assert response.status_code == 200 + lines = [line async for line in response.aiter_lines() if line] + + event = _extract_first_event(lines) + assert event["type"] == "response.created" + assert selected_ids == [owner_id] + + @pytest.mark.asyncio async def test_v1_responses_previous_response_owner_lookup_failure_without_http_bridge_returns_upstream_unavailable( async_client, diff --git a/tests/unit/test_db_migrate.py b/tests/unit/test_db_migrate.py index deb79efac4..0e203577c1 100644 --- a/tests/unit/test_db_migrate.py +++ b/tests/unit/test_db_migrate.py @@ -2231,7 +2231,11 @@ def test_api_key_reasoning_policy_migration_round_trips_from_current_parent(tmp_ config = _build_alembic_config(url) script_directory = ScriptDirectory.from_config(config) assert script_directory.get_revision(target_revision).down_revision == parent_revision - assert target_revision in script_directory.get_heads() + # Assert reachability from the single head rather than "is the head": every + # later migration would otherwise have to edit this test. + heads = script_directory.get_heads() + assert len(heads) == 1 + assert target_revision in {revision.revision for revision in script_directory.iterate_revisions(heads[0], "base")} engine = create_engine(to_sync_database_url(url)) try: diff --git a/tests/unit/test_live_usage_ingest.py b/tests/unit/test_live_usage_ingest.py index c37335d451..96735a4730 100644 --- a/tests/unit/test_live_usage_ingest.py +++ b/tests/unit/test_live_usage_ingest.py @@ -192,6 +192,33 @@ async def test_ingestor_queue_overflow_drops_oldest() -> None: ] +@pytest.mark.asyncio +async def test_ingestor_resolves_upstream_account_id_before_history_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=60.0) + by_id_calls: list[str | None] = [] + by_chatgpt_calls: list[str | None] = [] + + async def fake_by_id(account_id: str | None) -> str | None: + by_id_calls.append(account_id) + return "acc-internal" if account_id == "acc-internal" else None + + async def fake_by_chatgpt(chatgpt_account_id: str | None) -> str | None: + by_chatgpt_calls.append(chatgpt_account_id) + return "acc-internal" if chatgpt_account_id == "chatgpt-raw" else None + + monkeypatch.setattr(ingestor, "_resolve_account_id_by_id", fake_by_id) + monkeypatch.setattr(ingestor, "_resolve_account_id", fake_by_chatgpt) + + assert await ingestor._resolve_persisted_account_id("acc-internal", "unused") == "acc-internal" + assert await ingestor._resolve_persisted_account_id("chatgpt-raw", None) == "acc-internal" + assert await ingestor._resolve_persisted_account_id(None, "chatgpt-raw") == "acc-internal" + assert await ingestor._resolve_persisted_account_id("missing", "missing-too") is None + assert by_id_calls == ["acc-internal", "chatgpt-raw", "missing"] + assert by_chatgpt_calls == ["unused", "chatgpt-raw", "chatgpt-raw", "missing", "missing-too"] + + @pytest.mark.asyncio async def test_stream_responses_tap_publishes_rate_limit_events(monkeypatch: pytest.MonkeyPatch) -> None: import app.core.clients.proxy as proxy_client_module diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index b8d03ed5a2..e04675697f 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -4369,10 +4369,10 @@ async def test_compact_turn_state_owner_lookup_is_api_key_scoped_and_fails_close @pytest.mark.asyncio -async def test_compact_turn_state_owner_fails_closed_when_same_account_sessions_conflict() -> None: +async def test_compact_turn_state_owner_rejects_same_account_session_identity_drift() -> None: service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) owner = SimpleNamespace(id="account-owner") - turn_state = "turn-owner-session-conflict" + turn_state = "turn-owner-session-drift" owner_key = proxy_service._http_bridge_turn_state_alias_key(turn_state, None) service._http_bridge_turn_state_index[owner_key] = "bridge-live" # type: ignore[assignment] service._http_bridge_sessions["bridge-live"] = SimpleNamespace( # type: ignore[index] @@ -4440,6 +4440,10 @@ async def select_account(_deadline: float, **kwargs: object) -> AccountSelection assert seen_selection["preferred_account_id"] == owner.id assert seen_selection["fallback_on_preferred_account_unavailable"] is False + affinity = cast(proxy_service._AffinityPolicy, seen_selection["affinity_policy"]) + assert affinity.key == turn_state + assert affinity.kind == StickySessionKind.CODEX_SESSION + assert affinity.codex_session_source == "turn_state" @pytest.mark.asyncio @@ -4479,7 +4483,7 @@ async def test_compact_previous_response_owner_ignores_legacy_session_header_aff 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"] is None - assert select_account.await_args.kwargs["sticky_kind"] == proxy_service.StickySessionKind.CODEX_SESSION + assert select_account.await_args.kwargs["sticky_kind"] == StickySessionKind.CODEX_SESSION assert select_account.await_args.kwargs["sticky_source"] == "session_header" assert select_account.await_args.kwargs["legacy_sticky_key"] == "sid-root" @@ -4519,6 +4523,10 @@ async def select_account(_deadline: float, **kwargs: object) -> AccountSelection assert seen_selection["preferred_account_id"] == owner.id assert seen_selection["fallback_on_preferred_account_unavailable"] is False + affinity = cast(proxy_service._AffinityPolicy, seen_selection["affinity_policy"]) + assert affinity.key == turn_state + assert affinity.kind == StickySessionKind.CODEX_SESSION + assert affinity.codex_session_source == "turn_state" @pytest.mark.asyncio @@ -4620,6 +4628,8 @@ async def select_account(_deadline: float, **kwargs: object) -> AccountSelection assert seen_selection["preferred_account_id"] == account.id assert seen_selection["fallback_on_preferred_account_unavailable"] is False affinity = cast(proxy_service._AffinityPolicy, seen_selection["affinity_policy"]) + assert affinity.key == "soft-process-session" + assert affinity.kind == StickySessionKind.CODEX_SESSION assert affinity.codex_session_source == "session_header" assert await service.drain_persistence_tasks(timeout_seconds=1) @@ -13427,6 +13437,14 @@ async def test_plain_stream_resolves_http_bridge_turn_state_owner(monkeypatch: p async def select_account(_deadline: float, **kwargs: object) -> AccountSelection: selections.append(kwargs) + affinity = kwargs["affinity_policy"] + assert isinstance(affinity, proxy_service._AffinityPolicy) + # The stale session_id header must not reach selection: the hard + # turn-state owner supplies the key, exactly as on the compact path. + assert affinity.key == turn_state + assert affinity.kind == StickySessionKind.CODEX_SESSION + assert affinity.codex_session_source == "turn_state" + assert affinity.legacy_codex_session_key is None return AccountSelection(account=owner, error_message=None) async def fake_core_stream_responses(*_args: object, **_kwargs: object): @@ -13446,7 +13464,7 @@ async def fake_core_stream_responses(*_args: object, **_kwargs: object): chunk async for chunk in service._stream_with_retry( payload, - {"x-codex-turn-state": turn_state}, + {"session_id": "stale-session-owner", "x-codex-turn-state": turn_state}, codex_session_affinity=True, propagate_http_errors=False, openai_cache_affinity=False,