from __future__ import annotations import hashlib import json import os import stat import tempfile from dataclasses import dataclass, replace from datetime import UTC, datetime from pathlib import Path from collections.abc import Callable from typing import Protocol from account_registry import Account, AccountRegistry, AccountRef, AuthorityBinding, AuthorityMode from health_client import AccountSnapshot, HealthStatus from provider_services import AccountHealthProvider, AccountLifecycle class MigrationError(RuntimeError): pass class MigrationCleanupError(MigrationError): pass @dataclass(frozen=True) class AuthorityProvision: authority_name: str route_id: str proxy_grant_ref: Path generation: int class AuthorityMigrationAdmin(Protocol): def login_and_provision(self, account: Account, authority_name: str) -> AuthorityProvision: ... def status_generation(self, provider: str, account_id: str) -> int: ... def disable_route(self, route_id: str) -> None: ... def remove_route_grants(self, route_id: str) -> None: ... def remove_route(self, route_id: str) -> None: ... def remove_account(self, provider: str, account_id: str) -> None: ... @dataclass(frozen=True) class MigrationAcceptance: cli: bool = False resume: bool = False streaming: bool = False quota_health: bool = False repair: bool = False restart: bool = False exact_route: bool = False no_provider_credentials: bool = False responses: bool = False realtime: bool = False app_server: bool = False http_sse: bool = False def required(self, provider: str) -> tuple[str, ...]: common = ( "cli", "resume", "streaming", "quota_health", "repair", "restart", "exact_route", "no_provider_credentials", ) if provider == "codex": return (*common, "responses", "realtime", "app_server") if provider == "claude": return (*common, "http_sse") raise MigrationError("unsupported authority migration provider") def missing(self, provider: str) -> tuple[str, ...]: return tuple(name for name in self.required(provider) if not getattr(self, name)) def as_dict(self) -> dict[str, bool]: return { name: bool(getattr(self, name)) for name in ( "cli", "resume", "streaming", "quota_health", "repair", "restart", "exact_route", "no_provider_credentials", "responses", "realtime", "app_server", "http_sse", ) } @dataclass(frozen=True) class AccountMigrationReceipt: provider: str authority_name: str account_fingerprint: str route_fingerprint: str state: str generation: int started_at: str updated_at: str promoted_at: str | None = None rollback_started_at: str | None = None rollback_completed_at: str | None = None flow_checks: dict[str, bool] | None = None provider_credential_copied: bool = False native_fallback_used: bool = False cleanup_complete: bool = False def as_dict(self) -> dict[str, object]: return { "schema": "AccountMigrationReceipt/v1", "provider": self.provider, "authority_name": self.authority_name, "account_fingerprint": self.account_fingerprint, "route_fingerprint": self.route_fingerprint, "state": self.state, "generation": self.generation, "started_at": self.started_at, "updated_at": self.updated_at, "promoted_at": self.promoted_at, "rollback_started_at": self.rollback_started_at, "rollback_completed_at": self.rollback_completed_at, "flow_checks": dict(self.flow_checks or {}), "provider_credential_copied": self.provider_credential_copied, "native_fallback_used": self.native_fallback_used, "cleanup_complete": self.cleanup_complete, } class AuthorityMigrationCoordinator: def __init__( self, registry: AccountRegistry, admin: AuthorityMigrationAdmin, *, clock: Callable[[], datetime] | None = None, ) -> None: self._registry = registry self._admin = admin self._clock = clock or (lambda: datetime.now(UTC)) self._receipt_root = registry.base_dir / "authority-migration-receipts" @staticmethod def _fingerprint(value: str) -> str: return hashlib.sha256(value.encode("utf-8")).hexdigest()[:12] def _timestamp(self) -> str: value = self._clock() if not isinstance(value, datetime): raise MigrationError("migration clock returned an invalid value") if value.tzinfo is None: value = value.replace(tzinfo=UTC) return value.astimezone(UTC).isoformat().replace("+00:00", "Z") def _receipt_path(self, account: Account) -> Path: return self._receipt_root / account.tool / f"{self._fingerprint(account.tray_key)}.json" @staticmethod def _native_view(account: Account) -> Account: return Account( ref=AccountRef(account.tool, account.slug), alias=account.alias, account_home=account.account_home, email=account.email, plan=account.plan, account_id=account.account_id, authority_binding=None, ) def _write_receipt(self, account: Account, receipt: AccountMigrationReceipt) -> Path: path = self._receipt_path(account) path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) payload = json.dumps(receipt.as_dict(), indent=2, sort_keys=True) + "\n" fd, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) temporary = Path(name) try: os.fchmod(fd, 0o600) with os.fdopen(fd, "w", encoding="utf-8") as handle: handle.write(payload) handle.flush() os.fsync(handle.fileno()) os.replace(temporary, path) finally: temporary.unlink(missing_ok=True) return path def _read_receipt(self, account: Account) -> AccountMigrationReceipt: path = self._receipt_path(account) try: payload = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: raise MigrationError("migration receipt is unavailable") from exc expected = { "schema", "provider", "authority_name", "account_fingerprint", "route_fingerprint", "state", "generation", "started_at", "updated_at", "promoted_at", "rollback_started_at", "rollback_completed_at", "flow_checks", "provider_credential_copied", "native_fallback_used", "cleanup_complete", } if not isinstance(payload, dict) or set(payload) != expected: raise MigrationError("migration receipt is invalid") if payload.get("schema") != "AccountMigrationReceipt/v1": raise MigrationError("migration receipt is invalid") if payload.get("provider") != account.tool: raise MigrationError("migration receipt provider mismatch") if payload.get("account_fingerprint") != self._fingerprint(account.tray_key): raise MigrationError("migration receipt account mismatch") if payload.get("provider_credential_copied") is not False or payload.get("native_fallback_used") is not False: raise MigrationError("migration receipt violates credential boundary") generation = payload.get("generation") if isinstance(generation, bool) or not isinstance(generation, int) or generation <= 0: raise MigrationError("migration receipt generation is invalid") flow_checks = payload.get("flow_checks") if not isinstance(flow_checks, dict) or any(not isinstance(k, str) or not isinstance(v, bool) for k, v in flow_checks.items()): raise MigrationError("migration receipt flow checks are invalid") return AccountMigrationReceipt( provider=payload["provider"], authority_name=payload["authority_name"], account_fingerprint=payload["account_fingerprint"], route_fingerprint=payload["route_fingerprint"], state=payload["state"], generation=generation, started_at=payload["started_at"], updated_at=payload["updated_at"], promoted_at=payload["promoted_at"], rollback_started_at=payload["rollback_started_at"], rollback_completed_at=payload["rollback_completed_at"], flow_checks=flow_checks, cleanup_complete=payload["cleanup_complete"], ) def stage_dark(self, account: Account, *, authority_name: str) -> AccountMigrationReceipt: if account.tool not in {"codex", "claude"}: raise MigrationError("unsupported authority migration provider") if account.authority_binding is not None: raise MigrationError("account already has an authority binding") provision = self._admin.login_and_provision(account, authority_name) if provision.authority_name != authority_name or provision.generation <= 0: raise MigrationError("authority provisioning result is invalid") if not provision.proxy_grant_ref.is_absolute(): raise MigrationError("authority grant reference must be absolute") binding = AuthorityBinding( mode=AuthorityMode.SUBROUTER_DARK, authority_name=provision.authority_name, route_id=provision.route_id, provider=account.tool, proxy_grant_ref=provision.proxy_grant_ref, ) try: self._registry.set_authority_binding(account.slug, binding) except Exception: self._cleanup_authority(account.tool, account.slug, binding, remove_grant_ref=True) raise now = self._timestamp() receipt = AccountMigrationReceipt( provider=account.tool, authority_name=binding.authority_name, account_fingerprint=self._fingerprint(account.tray_key), route_fingerprint=self._fingerprint(binding.route_id), state="dark", generation=provision.generation, started_at=now, updated_at=now, flow_checks=MigrationAcceptance().as_dict(), ) self._write_receipt(account, receipt) return receipt def promote(self, account: Account, acceptance: MigrationAcceptance) -> AccountMigrationReceipt: current = self._current(account.slug) binding = current.authority_binding if binding is None or binding.mode != AuthorityMode.SUBROUTER_DARK or binding.quiesced: raise MigrationError("account is not in a promotable dark state") missing = acceptance.missing(account.tool) if missing: raise MigrationError("migration acceptance incomplete: " + ", ".join(missing)) generation = self._admin.status_generation(account.tool, account.slug) if generation <= 0: raise MigrationError("authority generation is unavailable") prior = self._read_receipt(current) if prior.route_fingerprint != self._fingerprint(binding.route_id): raise MigrationError("migration route changed during acceptance") promoted_binding = replace(binding, mode=AuthorityMode.SUBROUTER) self._registry.set_authority_binding(account.slug, promoted_binding) now = self._timestamp() receipt = replace( prior, state="authority", generation=generation, updated_at=now, promoted_at=now, flow_checks=acceptance.as_dict(), ) self._write_receipt(current, receipt) return receipt def rollback( self, account: Account, *, native_lifecycle: AccountLifecycle, native_health: AccountHealthProvider, ) -> AccountMigrationReceipt: current = self._current(account.slug) binding = current.authority_binding if binding is None: raise MigrationError("account is not authority-bound") prior = self._read_receipt(current) now = self._timestamp() quiesced = replace(binding, quiesced=True) self._registry.set_authority_binding(current.slug, quiesced) receipt = replace(prior, state="rollback-quiesced", updated_at=now, rollback_started_at=now) self._write_receipt(current, receipt) try: self._admin.disable_route(binding.route_id) except Exception as exc: raise MigrationError("failed to disable authority route") from exc native = self._native_view(current) events = list(native_lifecycle.reauthenticate(native)) if not any(event.kind == "success" for event in events) or any(event.kind == "failure" for event in events): raise MigrationError("fresh native reauthentication did not complete") snapshot: AccountSnapshot = native_health.fetch(native) if snapshot.status != HealthStatus.OK: raise MigrationError("fresh native readiness verification failed") # The mode switch is the transaction commit point. No credential bytes are # copied or restored: native lifecycle just created its own fresh chain. self._registry.set_authority_binding(current.slug, None) completed = self._timestamp() cleanup_complete = False try: self._cleanup_authority(current.tool, current.slug, binding, remove_grant_ref=True) cleanup_complete = True except Exception as exc: receipt = replace( receipt, state="native-cleanup-required", updated_at=completed, rollback_completed_at=completed, cleanup_complete=False, ) self._write_receipt(native, receipt) raise MigrationCleanupError("native routing restored but authority cleanup failed") from exc receipt = replace( receipt, state="native", updated_at=completed, rollback_completed_at=completed, cleanup_complete=cleanup_complete, ) self._write_receipt(native, receipt) return receipt def _current(self, slug: str) -> Account: for account in self._registry.list(): if account.slug == slug: return account raise MigrationError("account disappeared during migration") def _cleanup_authority( self, provider: str, account_id: str, binding: AuthorityBinding, *, remove_grant_ref: bool, ) -> None: errors: list[Exception] = [] for operation in ( lambda: self._admin.remove_route_grants(binding.route_id), lambda: self._admin.remove_route(binding.route_id), lambda: self._admin.remove_account(provider, account_id), ): try: operation() except Exception as exc: errors.append(exc) if remove_grant_ref: try: self._remove_grant_ref(binding.proxy_grant_ref) except Exception as exc: errors.append(exc) if errors: raise MigrationCleanupError("authority cleanup was incomplete") from errors[0] @staticmethod def _remove_grant_ref(path: Path) -> None: try: info = path.lstat() except FileNotFoundError: return if ( not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or info.st_uid != os.geteuid() or stat.S_IMODE(info.st_mode) & 0o077 ): raise MigrationCleanupError("refusing unsafe authority grant cleanup") path.unlink()