diff --git a/app/core/clients/proxy.py b/app/core/clients/proxy.py index 5fb1b15d67..fe50a4085e 100644 --- a/app/core/clients/proxy.py +++ b/app/core/clients/proxy.py @@ -2690,7 +2690,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 @@ -2861,7 +2861,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: @@ -2955,7 +2955,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/account_identity_lock.py b/app/db/account_identity_lock.py new file mode 100644 index 0000000000..a8fd57aedc --- /dev/null +++ b/app/db/account_identity_lock.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from collections.abc import Collection +from hashlib import sha256 + +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +_POSTGRES_ACCOUNT_IDENTITY_LOCK_TIMEOUT_MS = 30_000 + + +def advisory_lock_key(scope: str, value: str) -> int: + digest = sha256(f"{scope}:{value}".encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big", signed=True) + + +def account_identity_lock_key(chatgpt_account_id: str) -> int: + """Return the existing PostgreSQL lock namespace for one upstream identity.""" + return advisory_lock_key("account-id", f"chatgpt:{chatgpt_account_id}") + + +async def lock_postgresql_account_identities( + session: AsyncSession, + chatgpt_account_ids: Collection[str | None], +) -> tuple[int, ...]: + """Lock upstream identity membership in canonical transaction-scoped order.""" + bind = session.get_bind() + if bind is None or bind.dialect.name != "postgresql": + return () + lock_keys = tuple( + sorted( + { + account_identity_lock_key(chatgpt_account_id) + for chatgpt_account_id in chatgpt_account_ids + if chatgpt_account_id + } + ) + ) + try: + if lock_keys: + # Match the database layer's existing 30-second contention budget. + # Transaction-local scope bounds every subsequent lock in this unit + # of work and PostgreSQL restores it at transaction end. + await session.execute( + text("SELECT set_config('lock_timeout', :timeout, true)"), + {"timeout": f"{_POSTGRES_ACCOUNT_IDENTITY_LOCK_TIMEOUT_MS}ms"}, + ) + for lock_key in lock_keys: + await session.execute( + text("SELECT pg_advisory_xact_lock(:lock_key)"), + {"lock_key": lock_key}, + ) + except BaseException: + await session.rollback() + raise + return lock_keys diff --git a/app/modules/accounts/repository.py b/app/modules/accounts/repository.py index c56039e80a..3d1b62fd24 100644 --- a/app/modules/accounts/repository.py +++ b/app/modules/accounts/repository.py @@ -1,6 +1,5 @@ from __future__ import annotations -import hashlib import json import uuid from dataclasses import dataclass @@ -15,6 +14,7 @@ from app.core.upstream_proxy.cache import get_upstream_route_cache from app.core.utils.time import utcnow +from app.db.account_identity_lock import advisory_lock_key, lock_postgresql_account_identities from app.db.models import ( Account, AccountLimitWarmup, @@ -74,6 +74,10 @@ def __init__(self, email: str) -> None: ) +class AccountIdentityRelockError(RuntimeError): + """Raised after identity membership changes across both bounded lock attempts.""" + + class AccountsRepository: def __init__(self, session: AsyncSession) -> None: self._session = session @@ -197,6 +201,7 @@ async def _upsert_unlocked( *, merge_by_email: bool | None = None, merge_by_chatgpt_identity: bool = False, + _identity_lock_attempt: int = 0, ) -> Account: dialect_name = self._dialect_name() sqlite_lock_acquired = False @@ -212,27 +217,34 @@ async def _upsert_unlocked( # exclusive on this dialect. await self._acquire_sqlite_merge_lock() elif dialect_name == "postgresql": - # Identity-keyed advisory lock must always be acquired when - # identity reconciliation is in play, regardless of - # merge_by_email. Two concurrent reauths for the same - # upstream chatgpt_account_id but different email claims - # (e.g. user changed email upstream) would otherwise take - # different email-scoped locks, both miss the canonical-row - # lookup below, and both INSERT a duplicate row for the - # same identity. - # - # Ordering identity-first, then email, gives a stable - # acquisition order across all callers so two concurrent - # reauths that overlap on either key serialize without - # deadlock. - identity_locked = False - if merge_by_chatgpt_identity and account.chatgpt_account_id: - await self._acquire_postgresql_identity_lock(f"chatgpt:{account.chatgpt_account_id}") - identity_locked = True + # Upstream identity membership always serializes before email + # locks and account row locks. This applies to ordinary imports as + # well as explicit identity reconciliation because either can add, + # replace, or remove a membership used by live-usage fallback. + locked_identities = await self._lock_postgresql_upsert_identity_candidates( + account, + include_email=bool(merge_by_email), + ) if merge_by_email: await self._acquire_postgresql_merge_lock(account.email) - elif not identity_locked: + elif not locked_identities: await self._acquire_postgresql_identity_lock(account.id) + if not await self._postgresql_upsert_identity_candidates_are_locked( + account, + include_email=bool(merge_by_email), + locked_identities=locked_identities, + ): + await self._session.rollback() + if _identity_lock_attempt >= 1: + raise AccountIdentityRelockError( + "Account identity candidates changed during PostgreSQL upsert locking" + ) + return await self._upsert_unlocked( + account, + merge_by_email=merge_by_email, + merge_by_chatgpt_identity=merge_by_chatgpt_identity, + _identity_lock_attempt=_identity_lock_attempt + 1, + ) # Identity-aware reconciliation runs before the deterministic-id # check so that a deactivated row whose refresh token was revoked @@ -302,7 +314,13 @@ async def upsert_reauthorized(self, account: Account) -> Account: async def replace_reauthorized(self, account_id: str, account: Account) -> Account | None: """Replace credentials on the exact local row selected for reauthentication.""" async with sqlite_writer_section(): - existing = await self._session.get(Account, account_id) + if self._dialect_name() == "postgresql": + existing = await self._lock_postgresql_account_identity_membership( + account_id, + account.chatgpt_account_id, + ) + else: + existing = await self._session.get(Account, account_id) if existing is None: return None await self._apply_account_replacement(existing, account) @@ -344,6 +362,7 @@ async def _upsert_account_slot_unlocked( *, preserve_unknown_workspace_duplicates: bool | None = None, preserve_identity_slots: bool = False, + _identity_lock_attempt: int = 0, ) -> Account: if preserve_unknown_workspace_duplicates is None: preserve_unknown_workspace_duplicates = not await self._merge_by_email_enabled() @@ -351,6 +370,10 @@ async def _upsert_account_slot_unlocked( if dialect_name == "sqlite": await self._acquire_sqlite_merge_lock() elif dialect_name == "postgresql": + locked_identities = await self._lock_postgresql_upsert_identity_candidates( + account, + include_email=True, + ) for lock_key in sorted( _slot_lock_keys( account, @@ -358,6 +381,22 @@ async def _upsert_account_slot_unlocked( ) ): await self._acquire_postgresql_identity_lock(lock_key) + if not await self._postgresql_upsert_identity_candidates_are_locked( + account, + include_email=True, + locked_identities=locked_identities, + ): + await self._session.rollback() + if _identity_lock_attempt >= 1: + raise AccountIdentityRelockError( + "Account identity candidates changed during PostgreSQL slot locking" + ) + return await self._upsert_account_slot_unlocked( + account, + preserve_unknown_workspace_duplicates=preserve_unknown_workspace_duplicates, + preserve_identity_slots=preserve_identity_slots, + _identity_lock_attempt=_identity_lock_attempt + 1, + ) existing = await self._account_by_slot_identity(account) if existing: @@ -779,6 +818,10 @@ async def update_routing_policy(self, account_id: str, routing_policy: str) -> b async def delete(self, account_id: str, *, delete_history: bool = False) -> bool: async with sqlite_writer_section(): + if self._dialect_name() == "postgresql": + # Identity membership precedes the fold-state lock so live + # settlement and deletion cannot form an identity/fold cycle. + await self._lock_postgresql_account_identity_membership(account_id, None) # Serialize against fold passes before touching the account's # request logs: without the fold-state lock an in-flight hourly # slice could aggregate the pre-delete attribution but commit @@ -843,6 +886,8 @@ async def rotate_tokens( material at all). """ async with sqlite_writer_section(): + if self._dialect_name() == "postgresql": + await self._lock_postgresql_account_identity_membership(account_id, chatgpt_account_id) values: dict[str, bytes | datetime | str] = { "access_token_encrypted": access_token_encrypted, "refresh_token_encrypted": refresh_token_encrypted, @@ -901,6 +946,8 @@ async def update_account_metadata( no-op existence check. """ async with sqlite_writer_section(): + if self._dialect_name() == "postgresql": + await self._lock_postgresql_account_identity_membership(account_id, chatgpt_account_id) values: dict[str, str | datetime] = {} if plan_type is not None: values["plan_type"] = plan_type @@ -1047,6 +1094,75 @@ async def _account_by_slot_identity(self, account: Account) -> Account | None: return matched return None + async def _lock_postgresql_account_identity_membership( + self, + account_id: str, + incoming_chatgpt_account_id: str | None, + *, + second_attempt: bool = False, + ) -> Account | None: + """Lock one row's old/new upstream memberships before mutating it.""" + observed_identity = await self._session.scalar( + select(Account.chatgpt_account_id).where(Account.id == account_id) + ) + await lock_postgresql_account_identities( + self._session, + (observed_identity, incoming_chatgpt_account_id), + ) + locked_account = await self._session.scalar( + select(Account) + .where(Account.id == account_id) + # PostgreSQL FOR NO KEY UPDATE stabilizes identity membership but + # remains compatible with the KEY SHARE lock taken by concurrent + # rollup FK inserts. Deletion upgrades only after the fold lock. + .with_for_update(key_share=True) + .execution_options(populate_existing=True) + ) + locked_identity = locked_account.chatgpt_account_id if locked_account is not None else None + if locked_identity == observed_identity: + return locked_account + await self._session.rollback() + if second_attempt: + raise AccountIdentityRelockError("Account identity changed during PostgreSQL membership lock acquisition") + return await self._lock_postgresql_account_identity_membership( + account_id, + incoming_chatgpt_account_id, + second_attempt=True, + ) + + async def _lock_postgresql_upsert_identity_candidates( + self, + account: Account, + *, + include_email: bool, + ) -> frozenset[str]: + predicates = _upsert_identity_candidate_predicates(account, include_email=include_email) + observed = ( + (await self._session.execute(select(Account.chatgpt_account_id).where(or_(*predicates)))).scalars().all() + ) + identities = frozenset(identity for identity in (*observed, account.chatgpt_account_id) if identity) + await lock_postgresql_account_identities(self._session, identities) + return identities + + async def _postgresql_upsert_identity_candidates_are_locked( + self, + account: Account, + *, + include_email: bool, + locked_identities: frozenset[str], + ) -> bool: + predicates = _upsert_identity_candidate_predicates(account, include_email=include_email) + current = ( + ( + await self._session.execute( + select(Account.chatgpt_account_id).where(or_(*predicates)).with_for_update(key_share=True) + ) + ) + .scalars() + .all() + ) + return all(identity is None or identity in locked_identities for identity in current) + def _dialect_name(self) -> str: return self._session.get_bind().dialect.name @@ -1062,14 +1178,14 @@ async def _acquire_sqlite_merge_lock(self) -> None: await self._session.execute(text("UPDATE accounts SET id = id WHERE 1 = 0")) async def _acquire_postgresql_merge_lock(self, email: str) -> None: - lock_key = _advisory_lock_key("merge-email", email) + lock_key = advisory_lock_key("merge-email", email) await self._session.execute( text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": lock_key}, ) async def _acquire_postgresql_identity_lock(self, account_id: str) -> None: - lock_key = _advisory_lock_key("account-id", account_id) + lock_key = advisory_lock_key("account-id", account_id) await self._session.execute( text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": lock_key}, @@ -1125,6 +1241,15 @@ def _slot_lock_keys(account: Account, *, preserve_unknown_workspace_duplicates: return (f"slot-local:{account.id}",) +def _upsert_identity_candidate_predicates(account: Account, *, include_email: bool) -> list[Any]: + predicates = [Account.id == account.id] + if account.chatgpt_account_id: + predicates.append(Account.chatgpt_account_id == account.chatgpt_account_id) + if include_email and account.email: + predicates.append(Account.email == account.email) + return predicates + + def _same_unknown_workspace_identity(existing: Account, incoming: Account) -> bool: return ( _workspace_slot_key(existing) is None @@ -1176,8 +1301,3 @@ def _can_reuse_email_fallback(existing: Account, incoming: Account) -> bool: or not existing.chatgpt_account_id or existing.chatgpt_account_id == incoming.chatgpt_account_id ) - - -def _advisory_lock_key(scope: str, value: str) -> int: - digest = hashlib.sha256(f"{scope}:{value}".encode("utf-8")).digest() - return int.from_bytes(digest[:8], byteorder="big", signed=True) diff --git a/app/modules/proxy/_service/http_bridge/upstream_events.py b/app/modules/proxy/_service/http_bridge/upstream_events.py index 431aaea01a..b43c2bb3f8 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=session.account.chatgpt_account_id, ) 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..b36a5ea525 100644 --- a/app/modules/usage/live_ingest.py +++ b/app/modules/usage/live_ingest.py @@ -5,22 +5,17 @@ import time 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 -from app.modules.usage.repository import UsageRepository +from app.modules.usage.repository import UsageRepository, UsageWindowWrite 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. @@ -68,7 +63,6 @@ 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._consumer: asyncio.Task[None] | None = None self._dropped = 0 self._last_cache_invalidation = 0.0 @@ -141,14 +135,6 @@ async def _run(self) -> None: ) async def _ingest(self, item: _QueuedSnapshot) -> None: - account_id = item.account_id - if account_id is None: - account_id = await self._resolve_account_id(item.chatgpt_account_id) - if account_id is None: - return - if self._should_skip(account_id, item.snapshot): - return - snapshot = item.snapshot primary = snapshot.primary secondary = snapshot.secondary @@ -162,51 +148,56 @@ async def _ingest(self, item: _QueuedSnapshot) -> None: and primary.window_minutes == usage_core.DEFAULT_WINDOW_MINUTES_MONTHLY ): monthly, primary = primary, None - async with get_background_session() as session: - repo = UsageRepository(session) - if primary is not None: - await repo.add_entry( - account_id=account_id, - used_percent=float(primary.used_percent), - input_tokens=None, - output_tokens=None, + windows: list[UsageWindowWrite] = [] + if primary is not None: + windows.append( + UsageWindowWrite( window="primary", + used_percent=float(primary.used_percent), reset_at=primary.reset_at, window_minutes=primary.window_minutes, credits_has=snapshot.credits_has, credits_unlimited=snapshot.credits_unlimited, credits_balance=snapshot.credits_balance, ) - if secondary is not None: - # Mirror the poller: credits normally ride the primary row. - # A secondary-only snapshot (e.g. the short window is not - # being reported) must still carry the fresh credit state. - secondary_carries_credits = primary is None - await repo.add_entry( - account_id=account_id, - used_percent=float(secondary.used_percent), - input_tokens=None, - output_tokens=None, + ) + if secondary is not None: + # Mirror the poller: credits normally ride the primary row. A + # secondary-only snapshot must still carry fresh credit state. + secondary_carries_credits = primary is None + windows.append( + UsageWindowWrite( window="secondary", + used_percent=float(secondary.used_percent), reset_at=secondary.reset_at, window_minutes=secondary.window_minutes, credits_has=snapshot.credits_has if secondary_carries_credits else None, credits_unlimited=snapshot.credits_unlimited if secondary_carries_credits else None, credits_balance=snapshot.credits_balance if secondary_carries_credits else None, ) - if monthly is not None: - await repo.add_entry( - account_id=account_id, - used_percent=float(monthly.used_percent), - input_tokens=None, - output_tokens=None, + ) + if monthly is not None: + windows.append( + UsageWindowWrite( window="monthly", + used_percent=float(monthly.used_percent), reset_at=monthly.reset_at, window_minutes=monthly.window_minutes, credits_has=snapshot.credits_has, credits_unlimited=snapshot.credits_unlimited, credits_balance=snapshot.credits_balance, ) + ) + + async with get_background_session() as session: + account_id = await UsageRepository(session).settle_live_account_snapshot( + account_id=item.account_id, + chatgpt_account_id=item.chatgpt_account_id, + windows=windows, + should_skip=lambda resolved: self._should_skip(resolved, snapshot), + ) + if account_id is None: + return self._last_write[account_id] = (_fingerprint(snapshot), time.monotonic()) await self._invalidate_caches_throttled() @@ -236,25 +227,6 @@ 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: - 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))) - .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 - _ingestor: LiveUsageIngestor | None = None diff --git a/app/modules/usage/repository.py b/app/modules/usage/repository.py index 3626bd34f5..2965c3efdd 100644 --- a/app/modules/usage/repository.py +++ b/app/modules/usage/repository.py @@ -7,16 +7,17 @@ from datetime import datetime from hashlib import sha256 from threading import RLock -from typing import Any, cast +from typing import Any, Callable, cast from anyio import to_thread -from sqlalchemy import Integer, and_, delete, func, literal_column, or_, select, true, tuple_ +from sqlalchemy import Integer, and_, delete, func, literal_column, or_, select, text, true, tuple_ from sqlalchemy import cast as sqlalchemy_cast from sqlalchemy.ext.asyncio import AsyncSession from app.core.config.settings import get_settings from app.core.usage.types import UsageAggregateRow, UsageTrendBucket from app.core.utils.time import utcnow +from app.db.account_identity_lock import lock_postgresql_account_identities from app.db.models import Account, AdditionalUsageHistory, UsageHistory from app.db.session import relax_commit_durability, sqlite_writer_section from app.db.sqlite_utils import sqlite_db_path_from_url @@ -50,6 +51,35 @@ class UsageWindowWrite: credits_balance: float | None = None +class LiveSnapshotOwnerIdentityRelockError(RuntimeError): + """The selected live-snapshot owner's identity changed twice.""" + + +def _account_snapshot_entries( + account_id: str, + windows: Collection[UsageWindowWrite], + *, + recorded_at: datetime | None = None, +) -> list[UsageHistory]: + captured_at = recorded_at or utcnow() + return [ + UsageHistory( + account_id=account_id, + used_percent=window.used_percent, + input_tokens=None, + output_tokens=None, + window=window.window, + reset_at=window.reset_at, + window_minutes=window.window_minutes, + credits_has=window.credits_has, + credits_unlimited=window.credits_unlimited, + credits_balance=window.credits_balance, + recorded_at=captured_at, + ) + for window in windows + ] + + @dataclass(frozen=True, slots=True) class _BulkHistoryCacheMetadata: row_count: int @@ -635,23 +665,7 @@ async def add_account_snapshot( """Persist one account's standard usage windows atomically.""" if not windows: return [] - captured_at = recorded_at or utcnow() - entries = [ - UsageHistory( - account_id=account_id, - used_percent=window.used_percent, - input_tokens=None, - output_tokens=None, - window=window.window, - reset_at=window.reset_at, - window_minutes=window.window_minutes, - credits_has=window.credits_has, - credits_unlimited=window.credits_unlimited, - credits_balance=window.credits_balance, - recorded_at=captured_at, - ) - for window in windows - ] + entries = _account_snapshot_entries(account_id, windows, recorded_at=recorded_at) try: async with sqlite_writer_section(): # Telemetry write: this transaction only appends usage-history @@ -664,6 +678,133 @@ async def add_account_snapshot( raise return entries + async def _resolve_postgresql_live_snapshot_owner( + self, + account_id: str | None, + chatgpt_account_id: str | None, + ) -> str | None: + locked_identities = (chatgpt_account_id,) + fallback_identity = chatgpt_account_id + relocked = False + + while True: + await lock_postgresql_account_identities(self._session, locked_identities) + locked_identity_values = frozenset(identity for identity in locked_identities if identity) + identity_to_relock: str | None = None + + if account_id is not None: + # Read before taking the row lock so MVCC preserves the + # current recovery identity even when its writer has already + # deleted the local row but not committed yet. + observed = ( + await self._session.execute( + select(Account.id, Account.chatgpt_account_id).where(Account.id == account_id) + ) + ).one_or_none() + if observed is not None: + observed_identity = observed.chatgpt_account_id + if observed_identity and observed_identity not in locked_identity_values: + identity_to_relock = observed_identity + else: + locked = ( + await self._session.execute( + select(Account.id, Account.chatgpt_account_id) + .where(Account.id == account_id) + .with_for_update(key_share=True) + ) + ).one_or_none() + if locked is not None: + if locked.chatgpt_account_id and locked.chatgpt_account_id not in locked_identity_values: + identity_to_relock = locked.chatgpt_account_id + else: + return locked.id + + if identity_to_relock is not None: + if relocked: + raise LiveSnapshotOwnerIdentityRelockError( + "Live snapshot owner identity changed during PostgreSQL relock" + ) + # Release the first lock before adding another identity; the + # shared helper can then reacquire the full set in canonical + # order without inverting an account writer's lock order. + await self._session.rollback() + fallback_identity = identity_to_relock + locked_identities = (chatgpt_account_id, identity_to_relock) + relocked = True + continue + + if fallback_identity: + upstream_stmt = ( + select(Account.id) + .where(Account.chatgpt_account_id == fallback_identity) + .with_for_update(key_share=True) + ) + matches = list((await self._session.execute(upstream_stmt)).scalars().all()) + if len(matches) == 1: + return matches[0] + return None + + async def settle_live_account_snapshot( + self, + *, + account_id: str | None, + chatgpt_account_id: str | None, + windows: Collection[UsageWindowWrite], + should_skip: Callable[[str], bool], + ) -> str | None: + """Resolve a live snapshot owner and atomically persist its windows.""" + if not windows: + return None + + try: + async with sqlite_writer_section(): + bind = self._session.get_bind() + dialect_name = bind.dialect.name if bind is not None else "sqlite" + if dialect_name == "sqlite": + # Acquire SQLite's database-wide writer slot before owner + # lookup. Consolidation then commits before this lookup or + # waits until the snapshot commit, so the chosen FK owner + # cannot disappear between SELECT and INSERT. + await self._session.execute(text("BEGIN IMMEDIATE")) + resolved_account_id = None + if account_id is not None: + resolved_account_id = await self._session.scalar( + select(Account.id).where(Account.id == account_id) + ) + if resolved_account_id is None and chatgpt_account_id: + matches = list( + ( + await self._session.execute( + select(Account.id).where(Account.chatgpt_account_id == chatgpt_account_id) + ) + ) + .scalars() + .all() + ) + if len(matches) == 1: + resolved_account_id = matches[0] + else: + resolved_account_id = await self._resolve_postgresql_live_snapshot_owner( + account_id, + chatgpt_account_id, + ) + + if resolved_account_id is None or should_skip(resolved_account_id): + await self._session.rollback() + return None + + entries = _account_snapshot_entries(resolved_account_id, windows) + # Telemetry write: this transaction only locks the owner and + # appends usage-history rows, so it may skip synchronous WAL + # flush just like add_account_snapshot(). + await relax_commit_durability(self._session) + self._session.add_all(entries) + await self._session.commit() + except BaseException: + await self._session.rollback() + raise + return resolved_account_id + async def aggregate_since( self, since: datetime, diff --git a/openspec/changes/settle-live-usage-after-account-consolidation/design.md b/openspec/changes/settle-live-usage-after-account-consolidation/design.md new file mode 100644 index 0000000000..ad42f7de6c --- /dev/null +++ b/openspec/changes/settle-live-usage-after-account-consolidation/design.md @@ -0,0 +1,177 @@ +## Context + +Live usage publication and account reconciliation run in different ownership +domains. The proxy captures a snapshot and enqueues it without waiting; the +single background consumer later opens its own database session. Meanwhile, +identity-aware account upsert can select canonical account `C`, reparent the +persisted children of duplicate `D`, and delete `D` in one transaction. + +The loss sequence is therefore deterministic: publication records only `D`; +consolidation commits `D -> C`; `_ingest` attempts to append with stale `D`; +the account foreign key rejects the write; and the serving-safe consumer logs +and drops it. Existing history reparenting cannot cover a row that did not +exist when consolidation ran. + +The relevant identity constraint is equally important: an upstream ChatGPT +account id can be shared by distinct real-email slots. Upstream identity is a +safe fallback only when it resolves to exactly one surviving local row. + +## Goals / Non-Goals + +**Goals:** + +- Preserve every already-captured live snapshot across same-slot duplicate + consolidation when one canonical owner survives. +- Keep publication non-blocking and persistence in the background consumer. +- Prefer a valid captured local owner; use upstream identity only to recover a + stale or absent local owner and only when the result is unique. +- Persist each accepted snapshot once under one owner, with all represented + windows committed atomically. +- Prove the stale-local, valid-local, and upstream-only paths without sleeps or + timing-dependent scheduling. + +**Non-Goals:** + +- Changing duplicate-account selection, shared-workspace slot preservation, or + canonical-account choice. +- Guessing between multiple local rows that share an upstream identity. +- Retrying arbitrary ingestion failures or changing the queue's drop-oldest, + throttling, or serving-path isolation behavior. +- Adding a schema migration, configuration flag, or API response field. + +## Decisions + +### D1: Queue an ownership envelope containing local and upstream identities + +Every proxy tap point that knows a local serving account and its upstream +ChatGPT account id will publish both. The queued item remains an in-memory typed +value containing `account_id`, `chatgpt_account_id`, and the snapshot; no +database or wire schema is introduced. Upstream-only callers continue to leave +the local id absent. + +Capturing the upstream id at publication time is necessary because `D` cannot +be queried after consolidation deletes it. Looking up the upstream id only +after detecting stale `D` would already have lost the recovery key. + +### D2: Select and protect the persistence owner at consume time + +Ingestion will settle ownership in this order: + +1. If the captured local id still identifies an account, select it even when + the upstream identity is absent, shared, or points at another candidate. +2. If the local id is absent or no longer exists, resolve the captured upstream + id against current account rows and select it only when exactly one row + survives. +3. If neither rule selects an owner, do not guess; retain the current logged, + serving-safe drop behavior. + +Owner selection and the atomic append of all represented usage windows belong +to one serialized write operation. SQLite acquires `BEGIN IMMEDIATE` before +lookup and keeps its database-wide writer serialization through commit. +PostgreSQL first acquires the existing transaction-scoped advisory-lock +namespace keyed by the captured upstream identity before owner lookup. It then +reads the local owner's current identity without a row lock. When that current +non-null identity is not covered, settlement rolls back to release the initial +lock, reacquires the captured/current identities through the shared canonical +sort, and reselects the owner. This rollback is required: acquiring the current +identity while retaining the captured lock could invert the account-writer lock +order. If reconciliation wins the current-identity lock and deletes the local +row, settlement uses the last observed current identity as the unique fallback. +The reselected owner is held `FOR NO KEY UPDATE` through the append. That row +lock blocks deletion and key-changing writes without blocking the `KEY SHARE` +lock taken by concurrent foreign-key inserts. One relock is allowed; a second +identity change raises a typed terminal error, and null identities add no lock +key. + +Every PostgreSQL writer that can add, replace, move, consolidate, or delete an +`Account.chatgpt_account_id` membership acquires that same upstream lock and +holds it through commit. Old and incoming non-null identities are converted to +the stable advisory keys and acquired in canonical sorted order before any +email/slot advisory locks, account row locks, fold-state lock, or writes. A +local-id writer first reads the current identity without a row lock, acquires +the sorted old/new identity locks, and then row-locks and re-reads the account; +a changed observation rolls back and repeats that lock acquisition at most +once. Upsert candidate changes use the same bounded rollback/restart before any +mutation. Membership re-reads use PostgreSQL `FOR NO KEY UPDATE`, which +stabilizes identity changes while remaining compatible with the `KEY SHARE` +locks taken by concurrent fold rollup foreign-key inserts; deletion upgrades +its lock only after acquiring the fold-state lock. The shared helper applies a +transaction-local 30-second PostgreSQL lock timeout before advisory acquisition, +so request and background transactions propagate lock contention instead of +waiting indefinitely; it performs no polling or retry. + +This ordering gives both legal interleavings the same outcome: a snapshot +committed before consolidation is included when history is reparented, while a +snapshot whose current-identity reconciliation wins first relocks and writes +directly to `C` after the local duplicate disappears. +The per-account fingerprint is evaluated against the selected current owner, +and the successful-write marker is updated only after the atomic append. One +queued item therefore cannot write once to stale `D` and again to `C`. + +### D3: Preserve account-slot ambiguity and consolidation policy + +The fallback reuses the existing unique-upstream resolution rule. Distinct +real-email slots sharing one ChatGPT workspace remain distinct and ambiguous; +the change does not merge them or choose one. Duplicate reconciliation keeps +its current email/workspace candidate filters and canonical selection. It only +runs when the incoming upstream identity is non-null, and its duplicate query +requires `Account.chatgpt_account_id == incoming_identity`; an identity-less +local row therefore cannot be selected or deleted as an identity-reconciliation +duplicate. It only needs to leave the canonical row's existing upstream +identity intact, which it already does. + +This choice rejects two alternatives: always preferring upstream identity +could cross account slots even while the serving local row is valid, and +changing consolidation to force uniqueness would violate the established +shared-workspace account-slot contract. + +### D4: Deterministic regression and authenticated surface QA + +The deterministic transaction regression captures a queued item for `D` with +the shared upstream identity and coordinates independent PostgreSQL sessions at +exact lock and commit events, with no sleep, polling delay, or retry. Database +assertions prove one row per represented window under `C`, no row under `D`, +and no duplicate snapshot in both transaction orderings. A composition test +also drives the real proxied SSE publication tap through the live hub and +background consumer after consolidation, awaiting the exact settlement event +with a bounded timeout. Separate controls prove that an existing local id wins +and that an upstream-only item still resolves uniquely. + +Manual QA will use an isolated database and authenticated backend, execute a +literal `curl -i` request to `GET /api/accounts`, and verify HTTP 200, one +canonical `C`, no `D`, and the injected primary and secondary usage values. +The database diff will independently show one canonical snapshot and no +duplicate-owned history. All QA processes, credentials, database files, ports, +and temporary artifacts will be removed after capture. + +## Risks / Trade-offs + +- **Shared upstream id remains ambiguous.** A stale item can still be dropped + when multiple real-email slots survive. This is deliberate: preserving slot + ownership is safer than attributing usage to the wrong account. +- **Captured upstream identity can be absent.** Publication preserves the valid + local id together with the nullable upstream field, so valid-local settlement + still succeeds. Identity reconciliation cannot delete that identity-less row: + reconciliation requires a non-null incoming identity and selects duplicates + by equality to it. If the local row is already stale, no upstream fallback can + be recovered; genuinely upstream-less callers retain that serving-safe drop. +- **Settlement races consolidation.** A selected-row lock protects a snapshot + when settlement wins the row, but a current-identity consolidator can win + first and delete the local owner while settlement holds only the stale + captured-identity lock. SQLite writer serialization and PostgreSQL's bounded + rollback/relock close both transaction orderings without acquiring locks out + of canonical order. +- **Atomic append changes failure granularity.** If one represented window + cannot be stored, none of that snapshot's windows commit. This is preferable + to a partial snapshot and supports exactly-once settlement. + +## Migration Plan + +Ship publication and ingestion changes atomically. There is no schema or data +migration and no backfill: only snapshots captured after deployment carry both +identities. Rollback reverts the code; existing in-memory queued items disappear +with process shutdown exactly as they do today. + +## Open Questions + +None. diff --git a/openspec/changes/settle-live-usage-after-account-consolidation/proposal.md b/openspec/changes/settle-live-usage-after-account-consolidation/proposal.md new file mode 100644 index 0000000000..b53edb3ec6 --- /dev/null +++ b/openspec/changes/settle-live-usage-after-account-consolidation/proposal.md @@ -0,0 +1,47 @@ +## Why + +A live usage snapshot can be captured for duplicate local account `D` and wait +in the fire-and-forget queue while account reconciliation consolidates `D` into +canonical account `C`. Reconciliation reparents existing history and deletes +`D`, but the delayed ingestor still trusts the captured local id. Its +usage-history insert then violates the account foreign key and the serving-safe +consumer drops the already-captured snapshot. The invariant for this change is: +**an already-captured live snapshot survives duplicate-account consolidation.** + +## What Changes + +- Queue both the serving local account id and its upstream ChatGPT account id + when both identities are available at proxy publication time. +- Settle ownership at ingestion time: prefer a still-valid local account; + otherwise resolve the captured upstream identity only when it identifies one + surviving canonical account. +- Preserve the upstream-only publication path and the existing ambiguity rule + for shared-workspace identities. +- Persist one accepted snapshot atomically under the selected owner so a stale + `D` produces exactly one primary/secondary snapshot under `C` and no history + under `D`. +- Add deterministic, no-sleep regression coverage and authenticated + `/api/accounts` QA for the externally visible canonical result. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `live-usage-ingestion`: retain both ownership identities and settle queued + snapshots against the current account rows before persistence. +- `account-identity`: keep duplicate consolidation's canonical identity usable + for delayed ownership settlement without changing which accounts consolidate. + +## Impact + +- Affected code: live-usage publication call sites and hub contract, + `app/modules/usage/live_ingest.py`, and the existing atomic usage-snapshot + persistence path. +- Affected tests: focused live-ingestion integration coverage for stale-local, + valid-local, and upstream-only ownership paths. +- No database schema migration, new setting, API schema change, or account + consolidation policy change. diff --git a/openspec/changes/settle-live-usage-after-account-consolidation/specs/account-identity/spec.md b/openspec/changes/settle-live-usage-after-account-consolidation/specs/account-identity/spec.md new file mode 100644 index 0000000000..4cbaa1fd9e --- /dev/null +++ b/openspec/changes/settle-live-usage-after-account-consolidation/specs/account-identity/spec.md @@ -0,0 +1,23 @@ +## ADDED Requirements + +### Requirement: Duplicate consolidation preserves a recoverable canonical identity + +Identity reconciliation MUST preserve the upstream ChatGPT account id on the canonical row, reparent existing account-owned usage history to that row, and remove selected duplicate rows when it consolidates duplicate local accounts under the existing email and workspace-slot policy. Reconciliation MUST NOT consolidate distinct real-email account slots solely to make an upstream identity unique. On PostgreSQL, every account insertion, replacement, token/metadata identity update, duplicate consolidation, and deletion that changes upstream-identity membership MUST acquire the same transaction-scoped upstream-identity advisory lock as live-usage settlement before row or fold-state locks and hold it through commit. An old-to-new membership move MUST acquire both stable identity lock keys in canonical sorted order. + +#### Scenario: Same-slot duplicate leaves one upstream-resolvable canonical row + +- **GIVEN** canonical account `C` and duplicate account `D` are selected for consolidation by the existing identity policy +- **AND** both rows carry the same upstream ChatGPT account id +- **WHEN** reconciliation consolidates `D` into `C` +- **THEN** `C` remains with that upstream ChatGPT account id +- **AND** existing usage history formerly owned by `D` is owned by `C` +- **AND** `D` no longer exists +- **AND** the upstream ChatGPT account id resolves uniquely to `C` + +#### Scenario: Shared-workspace sibling slots remain distinct + +- **GIVEN** two current accounts have different real email addresses +- **AND** they share the same upstream ChatGPT account id +- **WHEN** identity reconciliation evaluates the accounts +- **THEN** it preserves both local account slots +- **AND** it does not consolidate either account solely to make upstream resolution unique diff --git a/openspec/changes/settle-live-usage-after-account-consolidation/specs/live-usage-ingestion/spec.md b/openspec/changes/settle-live-usage-after-account-consolidation/specs/live-usage-ingestion/spec.md new file mode 100644 index 0000000000..0b1d75ccb9 --- /dev/null +++ b/openspec/changes/settle-live-usage-after-account-consolidation/specs/live-usage-ingestion/spec.md @@ -0,0 +1,77 @@ +## ADDED Requirements + +### Requirement: Captured live snapshots survive account consolidation + +The proxy MUST enqueue both the serving local account id and the upstream +ChatGPT account id when both are available. At consumption, live usage +ingestion MUST prefer the captured local id when it still identifies an +account. If that local id is absent or no longer exists, ingestion MUST use the +captured upstream id only when it resolves to exactly one current local account. +For a selected owner, ingestion MUST atomically persist no more than one history +row for each window represented by the queued snapshot. On PostgreSQL, +ingestion MUST acquire a transaction-scoped advisory lock keyed by the captured +upstream identity before either owner lookup and hold it through snapshot +commit. If the selected local owner's current non-null upstream identity is not +already locked, ingestion MUST roll back the initial transaction, reacquire the +captured and current identity locks in canonical sorted order, and reselect and +revalidate the owner before persistence. If that local owner was consolidated +while the current-identity lock was acquired, ingestion MUST use the last +observed current identity only when it resolves to exactly one surviving local +account. Ingestion MUST perform at most one such relock and MUST raise a typed +error if the selected owner's identity changes again; a null current identity +MUST NOT create an advisory-lock key. Every account writer that can change +membership in an upstream identity MUST acquire the same lock before row locks +or mutation and hold it through commit. Writers moving membership between two +non-null upstream identities MUST acquire both stable lock keys in canonical +sorted order. + +#### Scenario: Stale duplicate settles under the unique canonical account + +- **GIVEN** a primary/secondary live snapshot was queued for duplicate account `D` +- **AND** the queued item contains `D` and the upstream identity shared with canonical account `C` +- **AND** duplicate reconciliation reparents existing history to `C` and deletes `D` +- **WHEN** the queued snapshot is consumed +- **THEN** exactly one primary row and one secondary row are persisted under `C` +- **AND** no usage-history row is persisted under `D` +- **AND** the persisted values equal the captured snapshot + +#### Scenario: A valid local owner takes precedence + +- **GIVEN** a queued snapshot contains a local account id that still exists +- **AND** it also contains an upstream identity usable for fallback +- **WHEN** the queued snapshot is consumed +- **THEN** the snapshot is persisted under the captured local account +- **AND** ingestion does not substitute another account selected by the upstream identity + +#### Scenario: A selected owner's current identity is revalidated + +- **GIVEN** a queued snapshot contains local account `A` and captured identity `X` +- **AND** `A` currently belongs to identity `Y` +- **WHEN** settlement overlaps reconciliation of `A` into a canonical `Y` owner +- **THEN** settlement releases its initial `X` lock before acquiring the canonical sorted lock set for `X` and `Y` +- **AND** settlement reselects and revalidates the owner under that full lock set +- **AND** exactly one row per represented window survives under the canonical `Y` owner +- **AND** a second selected-owner identity change raises a typed terminal error without persisting the snapshot + +#### Scenario: Upstream-only publication still resolves + +- **GIVEN** a queued snapshot has no local account id +- **AND** its upstream identity resolves to exactly one current local account +- **WHEN** the queued snapshot is consumed +- **THEN** the snapshot is persisted once under that local account + +#### Scenario: Consolidation cannot delete a snapshot inserted after reparenting + +- **GIVEN** PostgreSQL settlement has selected duplicate `D` for a captured upstream identity +- **AND** reconciliation would reparent `D` history to `C` and then delete `D` +- **WHEN** settlement and reconciliation overlap across independent sessions +- **THEN** their shared transaction-scoped upstream-identity lock serializes the complete membership change +- **AND** the snapshot is either committed under `D` before reparenting or directly under `C` after reconciliation +- **AND** exactly one row per represented window survives under `C` + +#### Scenario: Ambiguous fallback does not guess an owner + +- **GIVEN** the captured local account id is absent or no longer exists +- **AND** the captured upstream identity matches multiple current local accounts +- **WHEN** the queued snapshot is consumed +- **THEN** no usage-history row is persisted for that snapshot diff --git a/openspec/changes/settle-live-usage-after-account-consolidation/tasks.md b/openspec/changes/settle-live-usage-after-account-consolidation/tasks.md new file mode 100644 index 0000000000..aa25b1fca3 --- /dev/null +++ b/openspec/changes/settle-live-usage-after-account-consolidation/tasks.md @@ -0,0 +1,72 @@ +## 1. Deterministic regression coverage + +- [x] 1.1 Add a no-sleep integration test that queues a primary/secondary live + snapshot for duplicate `D` with local and upstream identities, completes + same-slot reconciliation into canonical `C`, then directly consumes the + captured item. +- [x] 1.2 Assert exactly one persisted row for each represented window under + `C`, no usage row under `D`, no duplicate snapshot, and preservation of the + injected usage/reset/credits values. +- [x] 1.3 Add controls proving a still-valid local id is preferred even when an + upstream fallback exists, and an upstream-only queued item still resolves to + its unique local account. +- [x] 1.4 Capture the focused failing-first command and RED output before any + production edit; do not use sleeps, polling delays, retries, or a background + consumer timing race. +- [x] 1.5 Add a deterministic two-session PostgreSQL regression that pauses at + exact transaction events and proves consolidation cannot reparent before a + snapshot append and then cascade-delete that append. +- [x] 1.6 Add both legal PostgreSQL interleavings for queued identity `X` when + the selected local owner currently belongs to `Y`, including the causal RED + where `Y` reconciliation wins the owner row and deletes it before lookup. + +## 2. Publication ownership envelope + +- [x] 2.1 Retain both local account id and upstream ChatGPT account id in the + typed live-usage hub/queue contract. +- [x] 2.2 Update every local-account HTTP/SSE and WebSocket publication tap + point to supply the upstream identity when available, preserving the existing + upstream-only path and no-op hub behavior. + +## 3. Consume-time settlement + +- [x] 3.1 Resolve the persistence owner in the background ingestion session: + prefer an existing local row; if it is stale or absent, accept only one + current row matching the captured upstream identity. +- [x] 3.2 Protect owner resolution through persistence and write all represented + windows atomically so the item settles once under one account on SQLite and + PostgreSQL; use one shared transaction-scoped upstream-identity lock across + settlement, ordinary/slot upserts, replacement, rotation, metadata update, + consolidation, and deletion before row/fold locks. +- [x] 3.3 Keep ambiguous/missing ownership serving-safe and logged; do not alter + account consolidation policy, queue overflow, throttling, or retry behavior. +- [x] 3.4 Add no Alembic revision, model column, setting, or API schema change. +- [x] 3.5 Roll back before bounded relock of the canonical captured/current + identity set, reselect and revalidate ownership, and raise a typed terminal + error on a second identity change without fabricating a null lock key. + +## 4. Automated verification + +- [x] 4.1 Run the focused live-ingestion integration selection once to GREEN, + proving stale-local consolidation, valid-local preference, and upstream-only + resolution. +- [x] 4.2 Run diagnostics on every changed Python file and the affected backend + lint/type/test gates on both supported database paths where registered. +- [x] 4.3 Run `openspec validate settle-live-usage-after-account-consolidation --strict`. +- [x] 4.4 Run the deterministic PostgreSQL race repeatedly plus lock-routing, + identity, live-ingest, snapshot, and HTTP publication regressions after the + shared lock implementation is complete. +- [x] 4.5 Run the selected-owner identity race in both transaction orders and + the focused no-relock, one-relock, terminal-change, rollback, sorted-lock, + and null-identity unit coverage. + +## 5. Authenticated QA and cleanup + +- [x] 5.1 Start an isolated QA database/backend, reproduce `D -> C` settlement, + and execute authenticated `curl -i GET /api/accounts` with the QA bearer key. +- [x] 5.2 Capture HTTP 200 evidence showing exactly one canonical `C`, no `D`, + and the injected primary and secondary usage values; capture an independent + database diff showing one canonical row per represented window and no + duplicate-owned row. +- [x] 5.3 Stop and remove every QA process, listener, credential, database file, + and temporary artifact; record the cleanup receipt. diff --git a/tests/integration/test_live_usage_ingest.py b/tests/integration/test_live_usage_ingest.py index db079ee890..dd4d4d66c7 100644 --- a/tests/integration/test_live_usage_ingest.py +++ b/tests/integration/test_live_usage_ingest.py @@ -1,17 +1,27 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any, cast import pytest +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.sql.dml import Delete +from app.core.clients import proxy as core_proxy from app.core.crypto import TokenEncryptor +from app.core.openai.requests import ResponsesRequest from app.core.usage import live_hub from app.core.usage.live_snapshots import LiveRateLimitSnapshot, LiveUsageWindow from app.core.utils.time import utcnow from app.db.models import Account, AccountStatus, UsageHistory from app.db.session import SessionLocal +from app.modules.accounts import repository as accounts_repository_module from app.modules.accounts.repository import AccountsRepository from app.modules.usage import live_ingest +from app.modules.usage import repository as usage_repository_module from app.modules.usage.repository import UsageRepository pytestmark = pytest.mark.integration @@ -43,6 +53,21 @@ def _snapshot() -> LiveRateLimitSnapshot: ) +async def _usage_rows_for(*account_ids: str) -> list[UsageHistory]: + async with SessionLocal() as session: + return list( + ( + await session.execute( + select(UsageHistory) + .where(UsageHistory.account_id.in_(account_ids)) + .order_by(UsageHistory.account_id, UsageHistory.window, UsageHistory.id) + ) + ) + .scalars() + .all() + ) + + async def _wait_for_rows(account_id: str, *, timeout: float = 5.0) -> tuple[UsageHistory | None, UsageHistory | None]: deadline = asyncio.get_event_loop().time() + timeout while True: @@ -234,22 +259,585 @@ async def test_live_ingestor_normalizes_monthly_only_snapshots(db_setup) -> None @pytest.mark.asyncio -async def test_live_ingestor_resolves_chatgpt_account_id(db_setup) -> None: +async def test_live_ingestor_settles_snapshot_after_duplicate_account_consolidation(db_setup) -> None: del db_setup + canonical_id = "acc_live_consolidated" + duplicate_id = "acc_live_consolidated__copy" + upstream_id = "workspace-live-consolidated" + email = "live-consolidated@example.com" + async with SessionLocal() as session: - await AccountsRepository(session).upsert( - _make_account("acc_live_resolved", "live-resolved@example.com", chatgpt_account_id="workspace-live-1") + repo = AccountsRepository(session) + await repo.upsert( + _make_account(canonical_id, email, chatgpt_account_id=upstream_id), + merge_by_email=False, + ) + await repo.upsert( + _make_account(duplicate_id, email, chatgpt_account_id=upstream_id), + merge_by_email=False, ) + snapshot = _snapshot() ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) - ingestor.start() + ingestor.publish( + snapshot, + account_id=duplicate_id, + chatgpt_account_id=upstream_id, + ) + queued = ingestor._queue.get_nowait() + + async with SessionLocal() as session: + saved = await AccountsRepository(session).upsert( + _make_account("acc_live_consolidated_reauth", email, chatgpt_account_id=upstream_id), + merge_by_email=False, + merge_by_chatgpt_identity=True, + ) + assert saved.id == canonical_id + assert await session.get(Account, duplicate_id) is None + + await ingestor._ingest(queued) + + rows = await _usage_rows_for(canonical_id, duplicate_id) + assert [row.account_id for row in rows] == [canonical_id, canonical_id] + primary_rows = [row for row in rows if row.window == "primary"] + secondary_rows = [row for row in rows if row.window == "secondary"] + assert len(primary_rows) == 1 + assert len(secondary_rows) == 1 + + primary = primary_rows[0] + secondary = secondary_rows[0] + assert snapshot.primary is not None + assert snapshot.secondary is not None + assert primary.used_percent == pytest.approx(snapshot.primary.used_percent) + assert primary.window_minutes == snapshot.primary.window_minutes + assert primary.reset_at == snapshot.primary.reset_at + assert primary.credits_has == snapshot.credits_has + assert primary.credits_unlimited == snapshot.credits_unlimited + assert primary.credits_balance == pytest.approx(snapshot.credits_balance) + assert secondary.used_percent == pytest.approx(snapshot.secondary.used_percent) + assert secondary.window_minutes == snapshot.secondary.window_minutes + assert secondary.reset_at == snapshot.secondary.reset_at + + +@pytest.mark.asyncio +async def test_sse_publication_tap_settles_queued_duplicate_snapshot_under_canonical_account( + monkeypatch: pytest.MonkeyPatch, + db_setup, +) -> None: + del db_setup + canonical_id = "acc_live_sse_canonical" + duplicate_id = "acc_live_sse_duplicate" + upstream_id = "workspace-live-sse" + email = "live-sse@example.com" + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert(_make_account(canonical_id, email, chatgpt_account_id=upstream_id), merge_by_email=False) + await repo.upsert(_make_account(duplicate_id, email, chatgpt_account_id=upstream_id), merge_by_email=False) + + rate_limit_event = ( + 'data: {"type":"codex.rate_limits","rate_limits":' + '{"primary":{"used_percent":33,"window_minutes":300,"reset_at":1700000300},' + '"secondary":{"used_percent":44,"window_minutes":10080,"reset_at":1700604800}}}\n\n' + ) + + @asynccontextmanager + async def _fake_http_session(_session): + yield cast(Any, object()) + + async def _fake_upstream_stream(**kwargs): + assert kwargs["account_id"] == upstream_id + assert kwargs["codex_lb_account_id"] == duplicate_id + yield rate_limit_event + + monkeypatch.setattr(core_proxy, "lease_http_session", _fake_http_session) + monkeypatch.setattr(core_proxy, "_stream_responses_with_session", _fake_upstream_stream) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingest_completed = asyncio.Event() + ingest_snapshot = ingestor._ingest + + async def _observed_ingest(item: live_ingest._QueuedSnapshot) -> None: + await ingest_snapshot(item) + ingest_completed.set() + + monkeypatch.setattr(ingestor, "_ingest", _observed_ingest) + live_hub.register_live_usage_publisher(ingestor.publish) try: - ingestor.publish(_snapshot(), chatgpt_account_id="workspace-live-1") - primary, secondary = await _wait_for_rows("acc_live_resolved") + events = [ + event + async for event in core_proxy.stream_responses( + ResponsesRequest(model="gpt-5.1", instructions="", input="hello", stream=True), + {}, + "access-token", + upstream_id, + session=cast(Any, object()), + codex_lb_account_id=duplicate_id, + ) + ] + assert events == [rate_limit_event] + + async with SessionLocal() as session: + saved = await AccountsRepository(session).upsert( + _make_account("acc_live_sse_reauth", email, chatgpt_account_id=upstream_id), + merge_by_email=False, + merge_by_chatgpt_identity=True, + ) + assert saved.id == canonical_id + + ingestor.start() + await asyncio.wait_for(ingest_completed.wait(), timeout=5.0) finally: await ingestor.stop() + live_hub.register_live_usage_publisher(None) - assert primary is not None and secondary is not None + rows = await _usage_rows_for(canonical_id, duplicate_id) + assert [row.account_id for row in rows] == [canonical_id, canonical_id] + assert {row.window for row in rows} == {"primary", "secondary"} + assert {row.used_percent for row in rows} == {33.0, 44.0} + + +@pytest.mark.asyncio +async def test_postgresql_live_ingest_serializes_identity_membership_through_snapshot_commit( + monkeypatch: pytest.MonkeyPatch, + db_setup, +) -> None: + del db_setup + bind = SessionLocal.kw["bind"] + if bind.dialect.name != "postgresql": + pytest.skip("PostgreSQL transaction-lock regression") + + canonical_id = "acc_live_pg_canonical" + duplicate_id = "acc_live_pg_duplicate" + upstream_id = "workspace-live-pg-consolidated" + email = "live-pg-consolidated@example.com" + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account(canonical_id, email, chatgpt_account_id=upstream_id), + merge_by_email=False, + ) + await repo.upsert( + _make_account(duplicate_id, email, chatgpt_account_id=upstream_id), + merge_by_email=False, + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.publish( + _snapshot(), + account_id=duplicate_id, + chatgpt_account_id=upstream_id, + ) + queued = ingestor._queue.get_nowait() + + settlement_commit_started = asyncio.Event() + release_settlement_commit = asyncio.Event() + writer_lock_attempted = asyncio.Event() + release_writer_delete = asyncio.Event() + settlement_lock_keys: list[int] = [] + writer_lock_keys: list[int] = [] + settlement_session = SessionLocal() + writer_session = SessionLocal() + settlement_task: asyncio.Task[None] | None = None + writer_task: asyncio.Task[Account] | None = None + + def _lock_key(args: tuple[Any, ...], kwargs: dict[str, Any]) -> int: + parameters = args[0] if args else kwargs.get("params") + assert isinstance(parameters, dict) + lock_key = parameters["lock_key"] + assert isinstance(lock_key, int) + return lock_key + + settlement_execute = settlement_session.execute + + async def _settlement_execute(statement: Any, *args: Any, **kwargs: Any): + if "pg_advisory_xact_lock" in str(statement): + settlement_lock_keys.append(_lock_key(args, kwargs)) + return await settlement_execute(statement, *args, **kwargs) + + settlement_commit = settlement_session.commit + + async def _settlement_commit() -> None: + settlement_commit_started.set() + await asyncio.wait_for(release_settlement_commit.wait(), timeout=5.0) + await settlement_commit() + + writer_execute = writer_session.execute + + async def _writer_execute(statement: Any, *args: Any, **kwargs: Any): + if "pg_advisory_xact_lock" in str(statement): + writer_lock_keys.append(_lock_key(args, kwargs)) + writer_lock_attempted.set() + if isinstance(statement, Delete) and statement.table.name == Account.__tablename__: + await asyncio.wait_for(release_writer_delete.wait(), timeout=5.0) + return await writer_execute(statement, *args, **kwargs) + + monkeypatch.setattr(settlement_session, "execute", _settlement_execute) + monkeypatch.setattr(settlement_session, "commit", _settlement_commit) + monkeypatch.setattr(writer_session, "execute", _writer_execute) + + @asynccontextmanager + async def _settlement_session() -> AsyncIterator[AsyncSession]: + yield settlement_session + + monkeypatch.setattr(live_ingest, "get_background_session", _settlement_session) + + try: + settlement_task = asyncio.create_task(ingestor._ingest(queued)) + await asyncio.wait_for(settlement_commit_started.wait(), timeout=5.0) + assert settlement_lock_keys, "settlement must take the upstream identity lock" + + writer_task = asyncio.create_task( + AccountsRepository(writer_session).upsert( + _make_account("acc_live_pg_reauth", email, chatgpt_account_id=upstream_id), + merge_by_email=False, + merge_by_chatgpt_identity=True, + ) + ) + await asyncio.wait_for(writer_lock_attempted.wait(), timeout=5.0) + + assert writer_lock_keys[0] == settlement_lock_keys[0] + release_settlement_commit.set() + await asyncio.wait_for(settlement_task, timeout=5.0) + release_writer_delete.set() + saved = await asyncio.wait_for(writer_task, timeout=5.0) + assert saved.id == canonical_id + finally: + release_settlement_commit.set() + release_writer_delete.set() + for task in (settlement_task, writer_task): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather( + *(task for task in (settlement_task, writer_task) if task is not None), + return_exceptions=True, + ) + await settlement_session.rollback() + await writer_session.rollback() + await settlement_session.close() + await writer_session.close() + + async with SessionLocal() as session: + accounts = list((await session.execute(select(Account).order_by(Account.id))).scalars().all()) + rows = list((await session.execute(select(UsageHistory).order_by(UsageHistory.id))).scalars().all()) + assert [account.id for account in accounts] == [canonical_id] + assert [row.account_id for row in rows] == [canonical_id, canonical_id] + assert {row.window for row in rows} == {"primary", "secondary"} + + +@pytest.mark.asyncio +async def test_postgresql_live_ingest_waits_for_identity_consolidation_commit( + monkeypatch: pytest.MonkeyPatch, + db_setup, +) -> None: + del db_setup + bind = SessionLocal.kw["bind"] + if bind.dialect.name != "postgresql": + pytest.skip("PostgreSQL transaction-lock regression") + + canonical_id = "acc_live_pg_writer_first_canonical" + duplicate_id = "acc_live_pg_writer_first_duplicate" + upstream_id = "workspace-live-pg-writer-first" + email = "live-pg-writer-first@example.com" + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert(_make_account(canonical_id, email, chatgpt_account_id=upstream_id), merge_by_email=False) + await repo.upsert(_make_account(duplicate_id, email, chatgpt_account_id=upstream_id), merge_by_email=False) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.publish(_snapshot(), account_id=duplicate_id, chatgpt_account_id=upstream_id) + queued = ingestor._queue.get_nowait() + + writer_commit_started = asyncio.Event() + release_writer_commit = asyncio.Event() + settlement_lock_attempted = asyncio.Event() + writer_session = SessionLocal() + writer_commit = writer_session.commit + real_settlement_lock = usage_repository_module.lock_postgresql_account_identities + writer_task: asyncio.Task[Account] | None = None + settlement_task: asyncio.Task[None] | None = None + + async def _writer_commit() -> None: + writer_commit_started.set() + await asyncio.wait_for(release_writer_commit.wait(), timeout=5.0) + await writer_commit() + + async def _observed_settlement_lock(session: AsyncSession, identities): + settlement_lock_attempted.set() + return await real_settlement_lock(session, identities) + + monkeypatch.setattr(writer_session, "commit", _writer_commit) + monkeypatch.setattr(usage_repository_module, "lock_postgresql_account_identities", _observed_settlement_lock) + + try: + writer_task = asyncio.create_task( + AccountsRepository(writer_session).upsert( + _make_account("acc_live_pg_writer_first_reauth", email, chatgpt_account_id=upstream_id), + merge_by_email=False, + merge_by_chatgpt_identity=True, + ) + ) + await asyncio.wait_for(writer_commit_started.wait(), timeout=5.0) + + settlement_task = asyncio.create_task(ingestor._ingest(queued)) + await asyncio.wait_for(settlement_lock_attempted.wait(), timeout=5.0) + assert not settlement_task.done() + + release_writer_commit.set() + saved = await asyncio.wait_for(writer_task, timeout=5.0) + await asyncio.wait_for(settlement_task, timeout=5.0) + assert saved.id == canonical_id + finally: + release_writer_commit.set() + for task in (writer_task, settlement_task): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather( + *(task for task in (writer_task, settlement_task) if task is not None), + return_exceptions=True, + ) + await writer_session.rollback() + await writer_session.close() + + async with SessionLocal() as session: + accounts = list((await session.execute(select(Account).order_by(Account.id))).scalars().all()) + rows = list((await session.execute(select(UsageHistory).order_by(UsageHistory.id))).scalars().all()) + assert [account.id for account in accounts] == [canonical_id] + assert [row.account_id for row in rows] == [canonical_id, canonical_id] + assert {row.window for row in rows} == {"primary", "secondary"} + + +@pytest.mark.asyncio +async def test_postgresql_live_ingest_recovers_when_current_identity_reconciliation_wins_owner_lock( + monkeypatch: pytest.MonkeyPatch, + db_setup, +) -> None: + del db_setup + bind = SessionLocal.kw["bind"] + if bind.dialect.name != "postgresql": + pytest.skip("PostgreSQL selected-owner lock regression") + + canonical_id = "acc_live_pg_current_identity_canonical" + selected_id = "acc_live_pg_current_identity_selected" + queued_identity = "workspace-live-pg-current-before" + current_identity = "workspace-live-pg-current-after" + email = "live-pg-current-identity@example.com" + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account(canonical_id, email, chatgpt_account_id=current_identity), + merge_by_email=False, + ) + selected = await repo.upsert( + _make_account(selected_id, email, chatgpt_account_id=queued_identity), + merge_by_email=False, + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.publish( + _snapshot(), + account_id=selected_id, + chatgpt_account_id=queued_identity, + ) + queued = ingestor._queue.get_nowait() + + async with SessionLocal() as session: + moved = await AccountsRepository(session).rotate_tokens( + selected.id, + selected.access_token_encrypted, + selected.refresh_token_encrypted, + selected.id_token_encrypted, + utcnow(), + expected_refresh_token_encrypted=selected.refresh_token_encrypted, + chatgpt_account_id=current_identity, + ) + assert moved is True + + reconciliation_commit_started = asyncio.Event() + release_reconciliation_commit = asyncio.Event() + settlement_local_lookup_started = asyncio.Event() + settlement_session = SessionLocal() + reconciliation_session = SessionLocal() + settlement_task: asyncio.Task[None] | None = None + reconciliation_task: asyncio.Task[Account] | None = None + settlement_execute = settlement_session.execute + reconciliation_commit = reconciliation_session.commit + + async def _settlement_execute(statement: Any, *args: Any, **kwargs: Any): + sql = str(statement) + if sql.startswith("SELECT accounts.id, accounts.chatgpt_account_id") and "WHERE accounts.id =" in sql: + settlement_local_lookup_started.set() + return await settlement_execute(statement, *args, **kwargs) + + async def _reconciliation_commit() -> None: + reconciliation_commit_started.set() + await asyncio.wait_for(release_reconciliation_commit.wait(), timeout=5.0) + await reconciliation_commit() + + monkeypatch.setattr(settlement_session, "execute", _settlement_execute) + monkeypatch.setattr(reconciliation_session, "commit", _reconciliation_commit) + + @asynccontextmanager + async def _settlement_session() -> AsyncIterator[AsyncSession]: + yield settlement_session + + monkeypatch.setattr(live_ingest, "get_background_session", _settlement_session) + + try: + reconciliation_task = asyncio.create_task( + AccountsRepository(reconciliation_session).upsert( + _make_account("acc_live_pg_current_identity_reauth", email, chatgpt_account_id=current_identity), + merge_by_email=False, + merge_by_chatgpt_identity=True, + ) + ) + await asyncio.wait_for(reconciliation_commit_started.wait(), timeout=5.0) + + settlement_task = asyncio.create_task(ingestor._ingest(queued)) + await asyncio.wait_for(settlement_local_lookup_started.wait(), timeout=5.0) + assert not settlement_task.done() + + release_reconciliation_commit.set() + saved = await asyncio.wait_for(reconciliation_task, timeout=5.0) + await asyncio.wait_for(settlement_task, timeout=5.0) + assert saved.id == canonical_id + finally: + release_reconciliation_commit.set() + for task in (settlement_task, reconciliation_task): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather( + *(task for task in (settlement_task, reconciliation_task) if task is not None), + return_exceptions=True, + ) + await settlement_session.rollback() + await reconciliation_session.rollback() + await settlement_session.close() + await reconciliation_session.close() + + async with SessionLocal() as session: + accounts = list((await session.execute(select(Account).order_by(Account.id))).scalars().all()) + rows = list((await session.execute(select(UsageHistory).order_by(UsageHistory.id))).scalars().all()) + assert [account.id for account in accounts] == [canonical_id] + assert [account.chatgpt_account_id for account in accounts] == [current_identity] + assert [row.account_id for row in rows] == [canonical_id, canonical_id] + assert {row.window for row in rows} == {"primary", "secondary"} + assert {row.used_percent for row in rows} == {33.0, 44.0} + + +@pytest.mark.asyncio +async def test_postgresql_opposite_identity_moves_use_one_sorted_lock_order( + monkeypatch: pytest.MonkeyPatch, + db_setup, +) -> None: + del db_setup + bind = SessionLocal.kw["bind"] + if bind.dialect.name != "postgresql": + pytest.skip("PostgreSQL transaction-lock regression") + + first = _make_account("acc_identity_move_a", "identity-move-a@example.com", chatgpt_account_id="workspace-a") + second = _make_account("acc_identity_move_b", "identity-move-b@example.com", chatgpt_account_id="workspace-b") + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert(first, merge_by_email=False) + await repo.upsert(second, merge_by_email=False) + + real_identity_lock = accounts_repository_module.lock_postgresql_account_identities + arrival_guard = asyncio.Lock() + both_arrived = asyncio.Event() + arrival_count = 0 + + async def _synchronized_identity_lock(session: AsyncSession, identities): + nonlocal arrival_count + async with arrival_guard: + arrival_count += 1 + if arrival_count == 2: + both_arrived.set() + await asyncio.wait_for(both_arrived.wait(), timeout=5.0) + return await real_identity_lock(session, identities) + + monkeypatch.setattr(accounts_repository_module, "lock_postgresql_account_identities", _synchronized_identity_lock) + + async def _move(account: Account, incoming_identity: str) -> bool: + async with SessionLocal() as session: + return await AccountsRepository(session).rotate_tokens( + account.id, + account.access_token_encrypted, + account.refresh_token_encrypted, + account.id_token_encrypted, + utcnow(), + expected_refresh_token_encrypted=account.refresh_token_encrypted, + chatgpt_account_id=incoming_identity, + ) + + moved_first, moved_second = await asyncio.wait_for( + asyncio.gather(_move(first, "workspace-b"), _move(second, "workspace-a")), + timeout=5.0, + ) + assert moved_first is True + assert moved_second is True + + async with SessionLocal() as session: + identities = { + account_id: chatgpt_account_id + for account_id, chatgpt_account_id in ( + await session.execute(select(Account.id, Account.chatgpt_account_id)) + ).all() + } + assert identities == { + first.id: "workspace-b", + second.id: "workspace-a", + } + + +@pytest.mark.asyncio +async def test_live_ingestor_prefers_valid_local_owner_over_upstream_fallback(db_setup) -> None: + del db_setup + local_id = "acc_live_valid_local" + sibling_id = "acc_live_valid_local_sibling" + upstream_id = "workspace-live-shared" + async with SessionLocal() as session: + repo = AccountsRepository(session) + await repo.upsert( + _make_account(local_id, "live-valid-local@example.com", chatgpt_account_id=upstream_id), + merge_by_email=False, + ) + await repo.upsert( + _make_account(sibling_id, "live-valid-sibling@example.com", chatgpt_account_id=upstream_id), + merge_by_email=False, + ) + + snapshot = _snapshot() + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.publish(snapshot, account_id=local_id, chatgpt_account_id=upstream_id) + queued = ingestor._queue.get_nowait() + + await ingestor._ingest(queued) + + rows = await _usage_rows_for(local_id, sibling_id) + assert len(rows) == 2 + assert {row.account_id for row in rows} == {local_id} + assert {row.window for row in rows} == {"primary", "secondary"} + + +@pytest.mark.asyncio +async def test_live_ingestor_resolves_chatgpt_account_id(db_setup) -> None: + del db_setup + account_id = "acc_live_resolved" + async with SessionLocal() as session: + await AccountsRepository(session).upsert( + _make_account(account_id, "live-resolved@example.com", chatgpt_account_id="workspace-live-1") + ) + + ingestor = live_ingest.LiveUsageIngestor(queue_size=8, write_min_interval_seconds=0.0) + ingestor.publish(_snapshot(), chatgpt_account_id="workspace-live-1") + queued = ingestor._queue.get_nowait() + + await ingestor._ingest(queued) + + rows = await _usage_rows_for(account_id) + assert len(rows) == 2 + assert {row.account_id for row in rows} == {account_id} + assert {row.window for row in rows} == {"primary", "secondary"} @pytest.mark.asyncio diff --git a/tests/integration/test_repositories.py b/tests/integration/test_repositories.py index 9d5e639015..08983ba27d 100644 --- a/tests/integration/test_repositories.py +++ b/tests/integration/test_repositories.py @@ -1240,6 +1240,40 @@ async def test_accounts_upsert_merge_by_chatgpt_identity_skips_without_upstream_ assert saved.id.startswith("acc_no_id__copy") +@pytest.mark.asyncio +async def test_identity_reconciliation_does_not_select_identityless_local_row_as_duplicate(db_setup): + async with SessionLocal() as session: + repo = AccountsRepository(session) + canonical = _make_account_with_chatgpt_id( + "acc_identity_canonical", + "identity-invariant@example.com", + "chatgpt_identity_invariant", + ) + identityless = _make_account("acc_identityless_local", "identity-invariant@example.com") + await repo.upsert(canonical, merge_by_email=False) + await repo.upsert(identityless, merge_by_email=False) + + saved = await repo.upsert( + _make_account_with_chatgpt_id( + "acc_identity_reauth", + "identity-invariant@example.com", + "chatgpt_identity_invariant", + ), + merge_by_email=False, + merge_by_chatgpt_identity=True, + ) + + assert saved.id == canonical.id + remaining = { + account.id: account.chatgpt_account_id + for account in (await session.execute(select(Account).order_by(Account.id))).scalars().all() + } + assert remaining == { + canonical.id: "chatgpt_identity_invariant", + identityless.id: None, + } + + @pytest.mark.asyncio async def test_usage_repository_aggregate(db_setup): async with SessionLocal() as session: diff --git a/tests/unit/test_accounts_repository_locks.py b/tests/unit/test_accounts_repository_locks.py index e4c6238abc..ce61f9ba7a 100644 --- a/tests/unit/test_accounts_repository_locks.py +++ b/tests/unit/test_accounts_repository_locks.py @@ -1,6 +1,7 @@ from __future__ import annotations from contextlib import asynccontextmanager +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -8,8 +9,9 @@ import app.modules.accounts.repository as repository_module from app.core.crypto import TokenEncryptor from app.core.utils.time import utcnow +from app.db.account_identity_lock import account_identity_lock_key, lock_postgresql_account_identities from app.db.models import Account, AccountStatus -from app.modules.accounts.repository import AccountsRepository +from app.modules.accounts.repository import AccountIdentityRelockError, AccountsRepository def _stub_account(account_id: str, email: str, chatgpt_id: str | None = None) -> Account: @@ -41,18 +43,40 @@ def _make_postgres_repo(monkeypatch: pytest.MonkeyPatch) -> tuple[AccountsReposi session.execute = AsyncMock() session.commit = AsyncMock() session.refresh = AsyncMock() + session.rollback = AsyncMock() session.add = MagicMock() session.get = AsyncMock(return_value=None) repo = AccountsRepository(session) - recorded: dict[str, list[str]] = {"identity": [], "email": []} + recorded: dict[str, list[str]] = {"upstream": [], "identity": [], "email": [], "order": []} async def fake_identity_lock(key: str) -> None: recorded["identity"].append(key) + recorded["order"].append(f"identity:{key}") async def fake_email_lock(email: str) -> None: recorded["email"].append(email) + recorded["order"].append(f"email:{email}") + + async def fake_upstream_identity_locks(account: Account, *, include_email: bool) -> frozenset[str]: + del include_email + if account.chatgpt_account_id: + recorded["upstream"].append(account.chatgpt_account_id) + recorded["order"].append(f"upstream:{account.chatgpt_account_id}") + return frozenset((account.chatgpt_account_id,)) + return frozenset() + + async def fake_candidates_are_locked( + account: Account, + *, + include_email: bool, + locked_identities: frozenset[str], + ) -> bool: + del account + del include_email + del locked_identities + return True async def fake_merge_by_email_enabled() -> bool: # only used when merge_by_email is None return True @@ -76,9 +100,13 @@ async def fake_next_available_account_id(account_id: str) -> str: monkeypatch.setattr(repo, "_dialect_name", lambda: "postgresql") monkeypatch.setattr(repo, "_acquire_postgresql_identity_lock", fake_identity_lock) monkeypatch.setattr(repo, "_acquire_postgresql_merge_lock", fake_email_lock) + monkeypatch.setattr(repo, "_lock_postgresql_upsert_identity_candidates", fake_upstream_identity_locks) + monkeypatch.setattr(repo, "_postgresql_upsert_identity_candidates_are_locked", fake_candidates_are_locked) monkeypatch.setattr(repo, "_merge_by_email_enabled", fake_merge_by_email_enabled) monkeypatch.setattr(repo, "_account_by_chatgpt_identity", fake_account_by_chatgpt_identity) + monkeypatch.setattr(repo, "_account_by_slot_identity", AsyncMock(return_value=None)) monkeypatch.setattr(repo, "_single_account_by_email", fake_single_account_by_email) + monkeypatch.setattr(repo, "_single_unknown_workspace_account_by_email", fake_single_account_by_email) monkeypatch.setattr(repo, "_next_available_account_id", fake_next_available_account_id) return repo, recorded @@ -90,6 +118,35 @@ def _make_result(value: str | None = "acc") -> MagicMock: return result +@pytest.mark.asyncio +async def test_postgresql_upstream_identity_locks_use_existing_namespace_in_sorted_order() -> None: + session = MagicMock() + session.get_bind.return_value.dialect.name = "postgresql" + session.execute = AsyncMock() + + lock_keys = await lock_postgresql_account_identities(session, ("workspace-z", None, "workspace-a", "workspace-z")) + + expected = tuple(sorted((account_identity_lock_key("workspace-a"), account_identity_lock_key("workspace-z")))) + assert lock_keys == expected + assert session.execute.await_args_list[0].args[1] == {"timeout": "30000ms"} + assert [call.args[1]["lock_key"] for call in session.execute.await_args_list[1:]] == list(expected) + + +@pytest.mark.asyncio +async def test_postgresql_upstream_identity_lock_failure_rolls_back_and_propagates() -> None: + session = MagicMock() + session.get_bind.return_value.dialect.name = "postgresql" + lock_error = RuntimeError("injected lock timeout") + session.execute = AsyncMock(side_effect=[MagicMock(), lock_error]) + session.rollback = AsyncMock() + + with pytest.raises(RuntimeError) as exc_info: + await lock_postgresql_account_identities(session, ("workspace-timeout",)) + + assert exc_info.value is lock_error + session.rollback.assert_awaited_once() + + @pytest.mark.asyncio async def test_account_update_status_uses_sqlite_writer_section(monkeypatch): session = MagicMock() @@ -183,9 +240,10 @@ async def test_upsert_takes_identity_lock_even_when_merge_by_email_enabled(monke await repo.upsert(account, merge_by_email=True, merge_by_chatgpt_identity=True) - assert recorded["identity"] == ["chatgpt:chatgpt_xyz"], ( - "identity lock must be acquired even when merge_by_email is True" + assert recorded["upstream"] == ["chatgpt_xyz"], ( + "upstream identity lock must be acquired even when merge_by_email is True" ) + assert recorded["identity"] == [] assert recorded["email"] == ["a@example.com"], "email lock must still be acquired when merge_by_email is True" @@ -200,7 +258,8 @@ async def test_upsert_takes_identity_lock_when_merge_by_email_disabled(monkeypat await repo.upsert(account, merge_by_email=False, merge_by_chatgpt_identity=True) - assert recorded["identity"] == ["chatgpt:chatgpt_zzz"] + assert recorded["upstream"] == ["chatgpt_zzz"] + assert recorded["identity"] == [] assert recorded["email"] == [] @@ -216,6 +275,7 @@ async def test_upsert_falls_back_to_id_lock_without_identity(monkeypatch): await repo.upsert(account, merge_by_email=False, merge_by_chatgpt_identity=False) + assert recorded["upstream"] == [] assert recorded["identity"] == ["acc_c"] assert recorded["email"] == [] @@ -231,5 +291,148 @@ async def test_upsert_email_only_when_identity_not_in_play(monkeypatch): await repo.upsert(account, merge_by_email=True, merge_by_chatgpt_identity=False) - assert recorded["identity"] == [], "no identity lock when merge_by_chatgpt_identity is False" + assert recorded["upstream"] == ["chatgpt_qqq"] + assert recorded["identity"] == [] assert recorded["email"] == ["d@example.com"] + assert recorded["order"] == ["upstream:chatgpt_qqq", "email:d@example.com"] + + +@pytest.mark.asyncio +async def test_ordinary_identity_upsert_uses_upstream_membership_lock(monkeypatch): + repo, recorded = _make_postgres_repo(monkeypatch) + account = _stub_account("acc_e", "e@example.com", chatgpt_id="chatgpt_ordinary") + + await repo.upsert(account, merge_by_email=False, merge_by_chatgpt_identity=False) + + assert recorded["upstream"] == ["chatgpt_ordinary"] + assert recorded["order"] == ["upstream:chatgpt_ordinary"] + + +@pytest.mark.asyncio +async def test_account_slot_upsert_locks_upstream_before_slot_keys(monkeypatch): + repo, recorded = _make_postgres_repo(monkeypatch) + account = _stub_account("acc_slot", "slot@example.com", chatgpt_id="chatgpt_slot") + account.workspace_id = "workspace-slot" + + await repo.upsert_account_slot(account, preserve_unknown_workspace_duplicates=False) + + assert recorded["upstream"] == ["chatgpt_slot"] + assert recorded["order"][0] == "upstream:chatgpt_slot" + assert all(item.startswith("identity:") for item in recorded["order"][1:]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slot_upsert", [False, True]) +async def test_identity_candidate_revalidation_restarts_once_then_succeeds(monkeypatch, slot_upsert: bool): + repo, _recorded = _make_postgres_repo(monkeypatch) + account = _stub_account("acc_retry", "retry@example.com", chatgpt_id="chatgpt_retry") + candidates_are_locked = AsyncMock(side_effect=[False, True]) + monkeypatch.setattr(repo, "_postgresql_upsert_identity_candidates_are_locked", candidates_are_locked) + + if slot_upsert: + saved = await repo.upsert_account_slot(account, preserve_unknown_workspace_duplicates=False) + else: + saved = await repo.upsert(account, merge_by_email=False) + + assert saved is account + assert candidates_are_locked.await_count == 2 + assert cast(Any, repo.session.rollback).await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slot_upsert", [False, True]) +async def test_identity_candidate_revalidation_raises_typed_error_after_second_change( + monkeypatch, + slot_upsert: bool, +): + repo, _recorded = _make_postgres_repo(monkeypatch) + account = _stub_account("acc_terminal", "terminal@example.com", chatgpt_id="chatgpt_terminal") + monkeypatch.setattr( + repo, + "_postgresql_upsert_identity_candidates_are_locked", + AsyncMock(side_effect=[False, False]), + ) + + with pytest.raises(AccountIdentityRelockError): + if slot_upsert: + await repo.upsert_account_slot(account, preserve_unknown_workspace_duplicates=False) + else: + await repo.upsert(account, merge_by_email=False) + + assert cast(Any, repo.session.rollback).await_count == 2 + + +@pytest.mark.asyncio +async def test_local_identity_membership_relocks_after_observed_identity_changes(monkeypatch): + session = MagicMock() + changed = _stub_account("acc_relock", "relock@example.com", chatgpt_id="chatgpt_changed") + session.scalar = AsyncMock(side_effect=["chatgpt_old", changed, "chatgpt_changed", changed]) + session.rollback = AsyncMock() + repo = AccountsRepository(session) + identity_locks = AsyncMock() + monkeypatch.setattr(repository_module, "lock_postgresql_account_identities", identity_locks) + + locked = await repo._lock_postgresql_account_identity_membership("acc_relock", "chatgpt_incoming") + + assert locked is changed + assert [call.args[1] for call in identity_locks.await_args_list] == [ + ("chatgpt_old", "chatgpt_incoming"), + ("chatgpt_changed", "chatgpt_incoming"), + ] + session.rollback.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_local_identity_membership_raises_typed_error_after_second_change(monkeypatch): + session = MagicMock() + changed_once = _stub_account("acc_relock", "relock@example.com", chatgpt_id="chatgpt_changed") + changed_twice = _stub_account("acc_relock", "relock@example.com", chatgpt_id="chatgpt_changed_again") + session.scalar = AsyncMock(side_effect=["chatgpt_old", changed_once, "chatgpt_changed", changed_twice]) + session.rollback = AsyncMock() + repo = AccountsRepository(session) + monkeypatch.setattr(repository_module, "lock_postgresql_account_identities", AsyncMock()) + + with pytest.raises(AccountIdentityRelockError): + await repo._lock_postgresql_account_identity_membership("acc_relock", "chatgpt_incoming") + + assert session.rollback.await_count == 2 + + +@pytest.mark.asyncio +async def test_local_identity_writers_lock_old_and_incoming_membership(monkeypatch): + repo, _recorded = _make_postgres_repo(monkeypatch) + existing = _stub_account("acc_writer", "writer@example.com", chatgpt_id="chatgpt_old") + membership_locks: list[tuple[str, str | None]] = [] + cast(Any, repo.session.execute).return_value = _make_result("acc_writer") + + async def fake_membership_lock(account_id: str, incoming: str | None) -> Account: + membership_locks.append((account_id, incoming)) + return existing + + monkeypatch.setattr(repo, "_lock_postgresql_account_identity_membership", fake_membership_lock) + monkeypatch.setattr(repo, "_apply_account_replacement", AsyncMock()) + monkeypatch.setattr(repository_module, "lock_fold_state", AsyncMock()) + monkeypatch.setattr(repository_module, "mirror_account_soft_delete_into_time_rollups", AsyncMock()) + + await repo.replace_reauthorized( + existing.id, + _stub_account("incoming", existing.email, chatgpt_id="chatgpt_new"), + ) + assert await repo.rotate_tokens( + existing.id, + b"access", + b"refresh", + b"id", + utcnow(), + expected_refresh_token_encrypted=b"expected", + chatgpt_account_id="chatgpt_new", + ) + assert await repo.update_account_metadata(existing.id, chatgpt_account_id="chatgpt_new") + assert await repo.delete(existing.id) + + assert membership_locks == [ + (existing.id, "chatgpt_new"), + (existing.id, "chatgpt_new"), + (existing.id, "chatgpt_new"), + (existing.id, None), + ] diff --git a/tests/unit/test_live_snapshot_owner_relock.py b/tests/unit/test_live_snapshot_owner_relock.py new file mode 100644 index 0000000000..11d59f28ab --- /dev/null +++ b/tests/unit/test_live_snapshot_owner_relock.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, call + +import pytest + +from app.db.account_identity_lock import lock_postgresql_account_identities +from app.modules.usage import repository as usage_repository_module +from app.modules.usage.repository import ( + LiveSnapshotOwnerIdentityRelockError, + UsageRepository, + UsageWindowWrite, +) + +pytestmark = pytest.mark.unit + + +def _identity_result(account_id: str | None, chatgpt_account_id: str | None = None) -> MagicMock: + result = MagicMock() + if account_id is None: + result.one_or_none.return_value = None + else: + result.one_or_none.return_value = MagicMock( + id=account_id, + chatgpt_account_id=chatgpt_account_id, + ) + return result + + +def _postgresql_session(results: list[MagicMock]) -> MagicMock: + session = MagicMock() + session.get_bind.return_value.dialect.name = "postgresql" + session.execute = AsyncMock(side_effect=results) + session.add_all = MagicMock() + session.commit = AsyncMock() + session.rollback = AsyncMock() + return session + + +async def _settle(session: MagicMock) -> str | None: + return await UsageRepository(session).settle_live_account_snapshot( + account_id="acc-selected", + chatgpt_account_id="workspace-x", + windows=[UsageWindowWrite(window="primary", used_percent=25.0)], + should_skip=lambda _account_id: False, + ) + + +@pytest.mark.asyncio +async def test_postgresql_live_snapshot_same_identity_does_not_relock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = _postgresql_session( + [ + _identity_result("acc-selected", "workspace-x"), + _identity_result("acc-selected", "workspace-x"), + ] + ) + identity_lock = AsyncMock() + monkeypatch.setattr(usage_repository_module, "lock_postgresql_account_identities", identity_lock) + monkeypatch.setattr(usage_repository_module, "relax_commit_durability", AsyncMock()) + + resolved = await _settle(session) + + assert resolved == "acc-selected" + identity_lock.assert_awaited_once_with(session, ("workspace-x",)) + session.rollback.assert_not_awaited() + session.commit.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_postgresql_live_snapshot_relocks_once_for_current_owner_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = _postgresql_session( + [ + _identity_result("acc-selected", "workspace-y"), + _identity_result("acc-selected", "workspace-y"), + _identity_result("acc-selected", "workspace-y"), + ] + ) + identity_lock = AsyncMock() + monkeypatch.setattr(usage_repository_module, "lock_postgresql_account_identities", identity_lock) + monkeypatch.setattr(usage_repository_module, "relax_commit_durability", AsyncMock()) + + resolved = await _settle(session) + + assert resolved == "acc-selected" + assert identity_lock.await_args_list == [ + call(session, ("workspace-x",)), + call(session, ("workspace-x", "workspace-y")), + ] + session.rollback.assert_awaited_once() + session.add_all.assert_called_once() + session.commit.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_postgresql_live_snapshot_second_owner_identity_change_is_terminal( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = _postgresql_session( + [ + _identity_result("acc-selected", "workspace-y"), + _identity_result("acc-selected", "workspace-z"), + ] + ) + identity_lock = AsyncMock() + monkeypatch.setattr(usage_repository_module, "lock_postgresql_account_identities", identity_lock) + + with pytest.raises(LiveSnapshotOwnerIdentityRelockError): + await _settle(session) + + assert identity_lock.await_args_list == [ + call(session, ("workspace-x",)), + call(session, ("workspace-x", "workspace-y")), + ] + assert session.rollback.await_count == 2 + session.add_all.assert_not_called() + session.commit.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_postgresql_identity_lock_does_not_fabricate_none_key() -> None: + session = MagicMock() + session.get_bind.return_value.dialect.name = "postgresql" + session.execute = AsyncMock() + + lock_keys = await lock_postgresql_account_identities(session, (None,)) + + assert lock_keys == () + session.execute.assert_not_awaited() diff --git a/tests/unit/test_live_usage_ingest.py b/tests/unit/test_live_usage_ingest.py index c1108e56f3..c37335d451 100644 --- a/tests/unit/test_live_usage_ingest.py +++ b/tests/unit/test_live_usage_ingest.py @@ -281,7 +281,6 @@ async def fake_lease(session: Any = None): assert (account_id, chatgpt_account_id) == (None, "workspace-live") 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. + # Local attribution stays preferred while retaining its recovery identity. _, 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_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index bc361b6f56..076082e0fb 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -779,7 +779,15 @@ def _make_bridge_session( kind=proxy_service.StickySessionKind.CODEX_SESSION, ), request_model="gpt-5.2", - account=cast(Any, SimpleNamespace(id="acc-bridge", status=AccountStatus.ACTIVE, plan_type="plus")), + account=cast( + Any, + SimpleNamespace( + id="acc-bridge", + chatgpt_account_id="workspace-bridge", + status=AccountStatus.ACTIVE, + plan_type="plus", + ), + ), upstream=cast(UpstreamWebSocket, SimpleNamespace(close=AsyncMock())), upstream_control=proxy_service._WebSocketUpstreamControl(), pending_requests=pending_requests or deque(), @@ -5319,6 +5327,7 @@ async def test_http_bridge_relay_publishes_live_rate_limit_events( service = proxy_service.ProxyService(cast(Any, nullcontext())) session = _make_bridge_session(key_value="bridge-live-rate-limits") + session.account.chatgpt_account_id = "workspace-bridge-live-rate-limits" rate_limit_text = ( '{"type":"codex.rate_limits","rate_limits":{"primary":' '{"used_percent":72,"window_minutes":300,"reset_at":1700000300}}}' @@ -5342,9 +5351,11 @@ async def test_http_bridge_relay_publishes_live_rate_limit_events( monkeypatch.setattr(service, "_fail_http_bridge_reader_and_maybe_retire", AsyncMock()) monkeypatch.setattr(service, "_fail_pending_websocket_requests", AsyncMock()) - captured: list[tuple[Any, str | None]] = [] + captured: list[tuple[Any, str | None, str | None]] = [] live_hub.register_live_usage_publisher( - lambda snapshot, *, account_id=None, chatgpt_account_id=None: captured.append((snapshot, account_id)) + lambda snapshot, *, account_id=None, chatgpt_account_id=None: captured.append( + (snapshot, account_id, chatgpt_account_id) + ) ) try: await service._relay_http_bridge_upstream_messages(session) @@ -5352,8 +5363,8 @@ async def test_http_bridge_relay_publishes_live_rate_limit_events( live_hub.register_live_usage_publisher(None) assert len(captured) == 1 - snapshot, account_id = captured[0] - assert account_id == session.account.id + snapshot, account_id, chatgpt_account_id = captured[0] + assert (account_id, chatgpt_account_id) == (session.account.id, session.account.chatgpt_account_id) assert snapshot.primary is not None assert snapshot.primary.used_percent == pytest.approx(72.0) diff --git a/tests/unit/test_usage_snapshot_repository.py b/tests/unit/test_usage_snapshot_repository.py index 117ab1142d..10b696fa32 100644 --- a/tests/unit/test_usage_snapshot_repository.py +++ b/tests/unit/test_usage_snapshot_repository.py @@ -3,13 +3,16 @@ from collections.abc import AsyncIterator, Collection from contextlib import asynccontextmanager from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock import pytest from sqlalchemy import event, func, select +from sqlalchemy.dialects import postgresql from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from app.db.models import Account, AccountStatus, Base, UsageHistory from app.modules.usage import background_repository as background_repository_module +from app.modules.usage import repository as usage_repository_module from app.modules.usage.background_repository import BackgroundUsageRepository from app.modules.usage.repository import UsageRepository, UsageWindowWrite @@ -26,6 +29,43 @@ async def session_factory() -> AsyncIterator[async_sessionmaker[AsyncSession]]: await engine.dispose() +@pytest.mark.asyncio +async def test_postgresql_settlement_lookups_compile_for_no_key_update( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = MagicMock() + session.get_bind.return_value.dialect.name = "postgresql" + observed_result = MagicMock() + observed_result.one_or_none.return_value = MagicMock( + id="acc_current", + chatgpt_account_id="workspace-current", + ) + locked_result = MagicMock() + locked_result.one_or_none.return_value = MagicMock( + id="acc_current", + chatgpt_account_id="workspace-current", + ) + session.execute = AsyncMock(side_effect=[observed_result, locked_result]) + session.add_all = MagicMock() + session.commit = AsyncMock() + session.rollback = AsyncMock() + monkeypatch.setattr(usage_repository_module, "lock_postgresql_account_identities", AsyncMock()) + monkeypatch.setattr(usage_repository_module, "relax_commit_durability", AsyncMock()) + + resolved = await UsageRepository(session).settle_live_account_snapshot( + account_id="acc_stale", + chatgpt_account_id="workspace-current", + windows=[UsageWindowWrite(window="primary", used_percent=25.0)], + should_skip=lambda _account_id: False, + ) + + assert resolved == "acc_current" + observed_stmt = session.execute.await_args_list[0].args[0] + locked_stmt = session.execute.await_args_list[1].args[0] + assert "FOR NO KEY UPDATE" not in str(observed_stmt.compile(dialect=postgresql.dialect())) + assert "FOR NO KEY UPDATE" in str(locked_stmt.compile(dialect=postgresql.dialect())) + + def _account(account_id: str) -> Account: return Account( id=account_id,