diff --git a/app/core/clients/proxy.py b/app/core/clients/proxy.py index d39f068154..10080e6378 100644 --- a/app/core/clients/proxy.py +++ b/app/core/clients/proxy.py @@ -2688,7 +2688,7 @@ async def stream_responses( publish_live_usage( parse_rate_limit_event_text(event_block), account_id=codex_lb_account_id, - chatgpt_account_id=None if codex_lb_account_id else account_id, + chatgpt_account_id=account_id, ) yield event_block @@ -2859,7 +2859,7 @@ async def _stream_via_http_attempt( publish_live_usage( parse_rate_limit_headers(getattr(raw_resp, "headers", None)), account_id=codex_lb_account_id, - chatgpt_account_id=None if codex_lb_account_id else account_id, + chatgpt_account_id=account_id, ) if resp.status >= 400: if raise_for_status: @@ -2953,7 +2953,7 @@ async def _stream_via_http_attempt( publish_live_usage( parse_rate_limit_headers(getattr(resp, "headers", None)), account_id=codex_lb_account_id, - chatgpt_account_id=None if codex_lb_account_id else account_id, + chatgpt_account_id=account_id, ) if resp.status >= 400: if raise_for_status: 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..0a8d9b8185 --- /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_000000_add_anonymous_telemetry +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_000000_add_anonymous_telemetry" +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 609520e2b8..b7acf4053d 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -2113,6 +2113,7 @@ class HttpBridgeRetryCircuit(Base): postgresql_include=["used_percent", "reset_at", "window_minutes", "id"], ) Index("idx_accounts_email", Account.email) +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/http_bridge/upstream_events.py b/app/modules/proxy/_service/http_bridge/upstream_events.py index 431aaea01a..bcc86efac8 100644 --- a/app/modules/proxy/_service/http_bridge/upstream_events.py +++ b/app/modules/proxy/_service/http_bridge/upstream_events.py @@ -1337,6 +1337,7 @@ async def _relay_http_bridge_upstream_messages( publish_live_usage( parse_rate_limit_event_text(message.text), account_id=session.account.id, + chatgpt_account_id=getattr(session.account, "chatgpt_account_id", None), ) await self._process_http_bridge_upstream_text(session, message.text) if await self._retire_http_bridge_after_drain_if_ready(session): diff --git a/app/modules/usage/live_ingest.py b/app/modules/usage/live_ingest.py index d26c2e7e2c..e46a7aa7bb 100644 --- a/app/modules/usage/live_ingest.py +++ b/app/modules/usage/live_ingest.py @@ -68,7 +68,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_cache: dict[str, tuple[str, 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 @@ -82,7 +83,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) @@ -142,10 +149,21 @@ async def _run(self) -> None: async def _ingest(self, item: _QueuedSnapshot) -> None: account_id = item.account_id + raw_account_id = account_id if account_id is None: - account_id = await self._resolve_account_id(item.chatgpt_account_id) + account_id = await self._resolve_account_id_by_chatgpt_account_id(item.chatgpt_account_id) + else: + resolved_account_id = await self._resolve_account_id( + chatgpt_account_id=item.chatgpt_account_id, + account_id=account_id, + ) + if resolved_account_id is None: + return + account_id = resolved_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 @@ -236,25 +254,62 @@ async def _invalidate_caches_now(self) -> None: # values before the TTL expires. await get_rate_limit_headers_cache().invalidate() - async def _resolve_account_id(self, chatgpt_account_id: str | None) -> str | None: + async def _resolve_account_id( + self, + chatgpt_account_id: str | None, + *, + account_id: str | None = None, + ) -> str | None: + if account_id: + if chatgpt_account_id: + resolved = await self._resolve_account_id_by_chatgpt_account_id(chatgpt_account_id) + if resolved is not None: + exact = await self._resolve_account_id_by_account_id(account_id) + if exact is not None and exact != resolved: + return None + return resolved + return await self._resolve_account_id_by_account_id(account_id) + return await self._resolve_account_id_by_chatgpt_account_id(chatgpt_account_id) + + async def _resolve_account_id_by_chatgpt_account_id(self, chatgpt_account_id: str | None) -> str | None: if not chatgpt_account_id: return None - cached = self._resolution_cache.get(chatgpt_account_id) - now = time.monotonic() - if cached is not None and now - cached[1] < _RESOLUTION_TTL_SECONDS: - return cached[0] async with get_background_session() as session: rows = ( - (await session.execute(select(Account.id).where(Account.chatgpt_account_id == chatgpt_account_id))) + ( + await session.execute( + select(Account.id).where(Account.chatgpt_account_id == chatgpt_account_id).limit(2) + ) + ) .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[chatgpt_account_id] = (resolved, now) return resolved + async def _resolve_account_id_by_account_id(self, account_id: str) -> str | None: + cache_key = f"account:{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 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.execute(select(Account.id).where(Account.id == account_id))).scalar_one_or_none() + if resolved is not 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.execute(select(Account.id).where(Account.id == account_id))).scalar_one_or_none() + return resolved is not None + _ingestor: LiveUsageIngestor | None = None 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 db079ee890..809de4904e 100644 --- a/tests/integration/test_live_usage_ingest.py +++ b/tests/integration/test_live_usage_ingest.py @@ -3,6 +3,7 @@ import asyncio import pytest +from sqlalchemy import delete, update from app.core.crypto import TokenEncryptor from app.core.usage import live_hub @@ -252,6 +253,247 @@ async def test_live_ingestor_resolves_chatgpt_account_id(db_setup) -> None: assert primary is not None and secondary is not None +@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_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/unit/test_live_usage_ingest.py b/tests/unit/test_live_usage_ingest.py index c1108e56f3..6f21c553da 100644 --- a/tests/unit/test_live_usage_ingest.py +++ b/tests/unit/test_live_usage_ingest.py @@ -282,6 +282,7 @@ async def fake_lease(session: Any = None): assert snapshot.primary is not None assert snapshot.primary.used_percent == pytest.approx(55.0) # When the caller knows the selected internal account, attribution - # prefers it so multi-seat workspaces are not dropped as ambiguous. + # includes the upstream ChatGPT identity so suffixed local account IDs can + # be resolved without guessing from their string shape. _, account_id_internal, chatgpt_internal = captured[1] - assert (account_id_internal, chatgpt_internal) == ("acc-internal", None) + assert (account_id_internal, chatgpt_internal) == ("acc-internal", "workspace-live") diff --git a/tests/unit/test_telemetry_migration.py b/tests/unit/test_telemetry_migration.py index 3428c5e7ab..56e2ba9897 100644 --- a/tests/unit/test_telemetry_migration.py +++ b/tests/unit/test_telemetry_migration.py @@ -16,6 +16,7 @@ async def test_telemetry_migration_upgrade_defaults_and_downgrade(tmp_path) -> N db_url = f"sqlite+aiosqlite:///{tmp_path / 'telemetry.sqlite'}" parent = "20260803_000000_merge_http_bridge_recovery_and_capability_lineage_heads" revision = "20260806_000000_add_anonymous_telemetry" + head_revision = "20260814_000000_add_accounts_chatgpt_identity_index" telemetry_columns = { "telemetry_consent", "telemetry_instance_id", @@ -54,7 +55,7 @@ async def columns_and_rows(engine): assert not telemetry_columns & columns result = await to_thread.run_sync(lambda: run_upgrade(db_url, "head", bootstrap_legacy=False)) - assert result.current_revision == revision + assert result.current_revision == head_revision columns, _ = await columns_and_rows(engine) assert telemetry_columns <= columns finally: