@linxiraos/pi-ai 1.1.26 → 1.2.0

package/src/auth/sqlite-credential-store.ts

@@ -9,7 +9,7 @@ import * as fs from "node:fs/promises";
 import * as path from "node:path";
 import { authPolicyFor } from "@linxiraos/pi-catalog/compat/auth";
 import { parseAlibabaTokenPlanCredential } from "@linxiraos/pi-catalog/wire/alibaba-token-plan";
-import { parseCloudflareAiGatewayCredential } from "@linxiraos/pi-catalog/wire/cloudflare-ai-gateway";
+import { parseCloudflareCredential } from "@linxiraos/pi-catalog/wire/cloudflare";
 import {
 	getAgentDbPath,
 	getDbBusyTimeoutMs,
@@ -110,6 +110,12 @@ const SQLITE_NOW_EPOCH = "CAST(strftime('%s','now') AS INTEGER)";
 const LEGACY_CODEX_BLOCK_PROVIDER_KEY = "openai-codex:oauth";
 const LEGACY_CODEX_BLOCK_SCOPE = "shared";
 const CODEX_METER_BLOCK_SCOPES = ["chat", "spark"] as const;
+/**
+ * Minimum spacing between read-path sweeps of expired block rows. Reads already
+ * filter `blocked_until_ms > now`; the sweep only bounds table growth, so it
+ * must not take the SQLite write lock on every selection.
+ */
+const CREDENTIAL_BLOCK_SWEEP_INTERVAL_MS = 60_000;
 
 // SQLite error classifiers live in pi-utils so the credential store and the
 // model cache share one implementation; re-exported here to preserve the
@@ -211,6 +217,12 @@ function resolveProviderCredentialIdentityKey(provider: string, identifiers: str
 	return null;
 }
 
+/**
+ * Identity key of a credential: the `identityKey` that broker snapshots carry
+ * and that broker account pools and `sessions.restrict` match, such as
+ * `email:<address>|org:<id>` for org-scoped providers. `null` for an API key or
+ * an OAuth credential without any account, email, project, or org identity.
+ */
 export function resolveCredentialIdentityKey(provider: string, credential: AuthCredential): string | null {
 	if (credential.type === "api_key") return null;
 	return resolveProviderCredentialIdentityKey(provider, extractOAuthCredentialIdentifiers(credential));
@@ -223,7 +235,11 @@ function resolveRowCredentialIdentityKey(provider: string, row: AuthRow): string
 	return credential?.type === "oauth" ? resolveCredentialIdentityKey(provider, credential) : null;
 }
 
-function matchesReplacementCredential(
+/**
+ * Whether upserting `incoming` replaces the active row holding `existing` (stored
+ * under `existingIdentityKey`) instead of inserting a new row beside it.
+ */
+export function matchesReplacementCredential(
 	provider: string,
 	existing: AuthCredential | null,
 	existingIdentityKey: string | null,
@@ -238,9 +254,9 @@ function matchesReplacementCredential(
 			const incomingToken = parseAlibabaTokenPlanCredential(incoming.key)?.token;
 			return existingToken !== undefined && existingToken === incomingToken;
 		}
-		if (provider === "cloudflare-ai-gateway") {
-			const existingToken = parseCloudflareAiGatewayCredential(existing.key)?.token;
-			const incomingToken = parseCloudflareAiGatewayCredential(incoming.key)?.token;
+		if (provider === "cloudflare-ai-gateway" || provider === "cloudflare-workers-ai") {
+			const existingToken = parseCloudflareCredential(existing.key)?.token;
+			const incomingToken = parseCloudflareCredential(incoming.key)?.token;
 			return existingToken !== undefined && existingToken === incomingToken;
 		}
 		return false;
@@ -294,6 +310,23 @@ function matchesReplacementCredential(
 	return false;
 }
 
+/**
+ * Whether two OAuth credentials belong to the same account of the same organization,
+ * whichever identity key each is stored under. The email decides when both carry one;
+ * otherwise a matching account id does, unless it is the organization's own id (a
+ * ChatGPT workspace id shared by every member). Project ids are shared between users
+ * (every consumer Antigravity login has the same one) and never count.
+ */
+export function isSameOAuthAccount(a: OAuthCredential, b: OAuthCredential): boolean {
+	const orgId = normalizeStoredAccountId(a.orgId);
+	if (orgId !== normalizeStoredAccountId(b.orgId)) return false;
+	const emailA = normalizeStoredEmail(a.email);
+	const emailB = normalizeStoredEmail(b.email);
+	if (emailA !== null && emailB !== null) return emailA === emailB;
+	const accountId = normalizeStoredAccountId(a.accountId);
+	return accountId !== null && accountId !== orgId && accountId === normalizeStoredAccountId(b.accountId);
+}
+
 function extractOAuthCredentialIdentifiers(credential: OAuthCredential): string[] {
 	const identifiers = new Set<string>();
 	const accountId = normalizeStoredAccountId(credential.accountId);
@@ -386,11 +419,14 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 	#updateIfMatchesWithLeaseStmt: Statement;
 	#deleteIfMatchesWithLeaseStmt: Statement;
 	#getCredentialBlockStmt: Statement;
-	#listCredentialBlocksByCredentialStmt: Statement;
+	#listCredentialBlockScopesStmt: Statement;
+	#listCredentialBlocksStmt: Statement;
 	#upsertCredentialBlockStmt: Statement;
 	#deleteCredentialBlocksStmt: Statement;
 	#deleteCredentialBlockStmt: Statement;
 	#deleteExpiredCredentialBlocksStmt: Statement;
+	/** Wall-clock time of the last expired-block sweep; throttles read-path sweeps. */
+	#lastCredentialBlockSweepMs = Number.NEGATIVE_INFINITY;
 	#acquireCredentialRefreshLeaseStmt: Statement;
 	#getCredentialRefreshLeaseStmt: Statement;
 	#renewCredentialRefreshLeaseStmt: Statement;
@@ -475,12 +511,18 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 		this.#getCredentialBlockStmt = this.#db.prepare(
 			"SELECT blocked_until_ms, updated_at FROM auth_credential_blocks WHERE credential_id = ? AND provider_key = ? AND block_scope = ? AND blocked_until_ms > ?",
 		);
-		this.#listCredentialBlocksByCredentialStmt = this.#db.prepare(
-			`SELECT credential_id, provider_key, block_scope, blocked_until_ms, updated_at
-			FROM auth_credential_blocks
-			WHERE credential_id = ? AND blocked_until_ms > ?
-				AND NOT (provider_key = ? AND block_scope = ?)
-			ORDER BY provider_key ASC, block_scope ASC`,
+		this.#listCredentialBlockScopesStmt = this.#db.prepare(
+			"SELECT block_scope, blocked_until_ms FROM auth_credential_blocks WHERE credential_id = ? AND provider_key = ? AND blocked_until_ms > ?",
+		);
+		// One statement for any number of ids: `json_each` keeps the caller's id
+		// order (`ids.key`) and the per-credential provider/scope order.
+		this.#listCredentialBlocksStmt = this.#db.prepare(
+			`SELECT b.credential_id, b.provider_key, b.block_scope, b.blocked_until_ms, b.updated_at
+			FROM json_each(?) AS ids
+			JOIN auth_credential_blocks AS b ON b.credential_id = ids.value
+			WHERE b.blocked_until_ms > ?
+				AND NOT (b.provider_key = ? AND b.block_scope = ?)
+			ORDER BY ids.key ASC, b.provider_key ASC, b.block_scope ASC`,
 		);
 		this.#upsertCredentialBlockStmt = this.#db.prepare(
 			`INSERT INTO auth_credential_blocks (credential_id, provider_key, block_scope, blocked_until_ms, updated_at)
@@ -576,7 +618,7 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 	 * Install the per-connection busy handler so lock-taking statements wait for
 	 * a contended writer instead of failing immediately (Bun defaults
 	 * `busy_timeout` to 0). MUST run before the first lock-taking statement on
-	 * the connection: concurrent zeta startups race WAL recovery and the leases
+	 * the connection: concurrent omp startups race WAL recovery and the leases
 	 * DDL. Uses the centralized timeout so headless hosts keep their bounded
 	 * busy wait instead of the interactive 5s value. See issues #2421, #7298.
 	 */
@@ -587,7 +629,7 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 	#initializeSchema(): void {
 		// Install the busy handler BEFORE any lock-taking statement (incl.
 		// `PRAGMA journal_mode=WAL`, which acquires an exclusive lock during WAL
-		// recovery). Without this, concurrent zeta startups can crash here with
+		// recovery). Without this, concurrent omp startups can crash here with
 		// `SQLITE_BUSY` / `SQLITE_BUSY_RECOVERY`. Re-setting when opened via
 		// `open()` (which already installed it) is idempotent. See issue #2421.
 		SqliteAuthCredentialStore.#installBusyTimeout(this.#db);
@@ -1585,27 +1627,38 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 	}
 
 	getCredentialBlock(credentialId: number, providerKey: string, blockScope: string): number | undefined {
-		const nowMs = Date.now();
-		const isCodexBlock = providerKey === LEGACY_CODEX_BLOCK_PROVIDER_KEY;
 		// Current callers use meter scopes. The physical shared row exists only
 		// for direct SQLite readers from pre-meter releases.
-		if (isCodexBlock && blockScope === LEGACY_CODEX_BLOCK_SCOPE) {
+		if (providerKey === LEGACY_CODEX_BLOCK_PROVIDER_KEY && blockScope === LEGACY_CODEX_BLOCK_SCOPE) {
 			return undefined;
 		}
-		if (!isCodexBlock) this.#deleteExpiredCredentialBlocksStmt.run(nowMs);
-		const row = this.#getCredentialBlockStmt.get(credentialId, providerKey, blockScope, nowMs) as
+		this.#sweepExpiredCredentialBlocks();
+		const row = this.#getCredentialBlockStmt.get(credentialId, providerKey, blockScope, Date.now()) as
 			| { blocked_until_ms?: number; updated_at?: number }
 			| undefined;
 		return typeof row?.blocked_until_ms === "number" ? row.blocked_until_ms : undefined;
 	}
 
+	getCredentialBlockScopes(credentialId: number, providerKey: string): Map<string, number> {
+		this.#sweepExpiredCredentialBlocks();
+		const rows = this.#listCredentialBlockScopesStmt.all(credentialId, providerKey, Date.now()) as Array<
+			Pick<CredentialBlockRow, "block_scope" | "blocked_until_ms">
+		>;
+		const hidesLegacyShared = providerKey === LEGACY_CODEX_BLOCK_PROVIDER_KEY;
+		const scopes = new Map<string, number>();
+		for (const row of rows) {
+			if (hidesLegacyShared && row.block_scope === LEGACY_CODEX_BLOCK_SCOPE) continue;
+			scopes.set(row.block_scope, row.blocked_until_ms);
+		}
+		return scopes;
+	}
+
 	getCredentialBlockReconcileAfter(credentialId: number, providerKey: string, blockScope: string): number | undefined {
-		const nowMs = Date.now();
-		const isCodexBlock = providerKey === LEGACY_CODEX_BLOCK_PROVIDER_KEY;
-		if (isCodexBlock && blockScope === LEGACY_CODEX_BLOCK_SCOPE) {
+		if (providerKey === LEGACY_CODEX_BLOCK_PROVIDER_KEY && blockScope === LEGACY_CODEX_BLOCK_SCOPE) {
 			return undefined;
 		}
-		if (!isCodexBlock) this.#deleteExpiredCredentialBlocksStmt.run(nowMs);
+		this.#sweepExpiredCredentialBlocks();
+		const nowMs = Date.now();
 		const row = this.#getCredentialBlockStmt.get(credentialId, providerKey, blockScope, nowMs) as
 			| { blocked_until_ms?: number; updated_at?: number }
 			| undefined;
@@ -1622,7 +1675,10 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 		const isLegacyCodexBlock =
 			block.providerKey === LEGACY_CODEX_BLOCK_PROVIDER_KEY && block.blockScope === LEGACY_CODEX_BLOCK_SCOPE;
 		const blockScopes = isLegacyCodexBlock ? CODEX_METER_BLOCK_SCOPES : [block.blockScope];
+		const nowMs = Date.now();
 		const upsert = this.#db.transaction(() => {
+			// A mark already holds the write lock: sweep expired rows in the same transaction.
+			this.#deleteExpiredCredentialBlocksStmt.run(nowMs);
 			for (const blockScope of blockScopes) {
 				this.#upsertCredentialBlockStmt.run(
 					block.credentialId,
@@ -1633,8 +1689,10 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 			}
 		});
 		upsert.immediate();
+		this.#lastCredentialBlockSweepMs = nowMs;
+		this.#pruneCredentialBlockReconcileAfter(nowMs);
 
-		const reconcileAfterMs = Math.min(block.blockedUntilMs, Date.now() + USAGE_REPORT_TTL_MS);
+		const reconcileAfterMs = Math.min(block.blockedUntilMs, nowMs + USAGE_REPORT_TTL_MS);
 		for (const blockScope of blockScopes) {
 			this.#credentialBlockReconcileAfter.set(
 				`${block.credentialId}\0${block.providerKey}\0${blockScope}`,
@@ -1661,38 +1719,49 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 	}
 
 	cleanExpiredCredentialBlocks(nowMs: number): void {
+		// Stamp wall-clock time first: callers may pass a future `nowMs`, and a
+		// failed sweep must not be retried by every following read.
+		this.#lastCredentialBlockSweepMs = Date.now();
 		this.#deleteExpiredCredentialBlocksStmt.run(nowMs);
+		this.#pruneCredentialBlockReconcileAfter(nowMs);
+	}
+
+	#pruneCredentialBlockReconcileAfter(nowMs: number): void {
 		for (const [key, reconcileAfterMs] of this.#credentialBlockReconcileAfter) {
 			if (reconcileAfterMs <= nowMs) this.#credentialBlockReconcileAfter.delete(key);
 		}
 	}
 
-	listCredentialBlocks(credentialIds: readonly number[]): StoredCredentialBlock[] {
-		if (credentialIds.length === 0) return [];
+	/**
+	 * Read-path expiry sweep, at most once per {@link CREDENTIAL_BLOCK_SWEEP_INTERVAL_MS}.
+	 * A busy write lock only postpones housekeeping; it never fails the read.
+	 */
+	#sweepExpiredCredentialBlocks(): void {
 		const nowMs = Date.now();
-		this.cleanExpiredCredentialBlocks(nowMs);
-		const seenCredentialIds = new Set<number>();
-		const blocks: StoredCredentialBlock[] = [];
-		for (const credentialId of credentialIds) {
-			if (seenCredentialIds.has(credentialId)) continue;
-			seenCredentialIds.add(credentialId);
-			const rows = this.#listCredentialBlocksByCredentialStmt.all(
-				credentialId,
-				nowMs,
-				LEGACY_CODEX_BLOCK_PROVIDER_KEY,
-				LEGACY_CODEX_BLOCK_SCOPE,
-			) as CredentialBlockRow[];
-			for (const row of rows) {
-				blocks.push({
-					credentialId: row.credential_id,
-					providerKey: row.provider_key,
-					blockScope: row.block_scope,
-					blockedUntilMs: row.blocked_until_ms,
-					updatedAtMs: row.updated_at * 1000,
-				});
-			}
+		if (nowMs - this.#lastCredentialBlockSweepMs < CREDENTIAL_BLOCK_SWEEP_INTERVAL_MS) return;
+		try {
+			this.cleanExpiredCredentialBlocks(nowMs);
+		} catch (err) {
+			if (!isSqliteBusyError(err)) throw err;
 		}
-		return blocks;
+	}
+
+	listCredentialBlocks(credentialIds: readonly number[]): StoredCredentialBlock[] {
+		if (credentialIds.length === 0) return [];
+		this.#sweepExpiredCredentialBlocks();
+		const rows = this.#listCredentialBlocksStmt.all(
+			JSON.stringify([...new Set(credentialIds)]),
+			Date.now(),
+			LEGACY_CODEX_BLOCK_PROVIDER_KEY,
+			LEGACY_CODEX_BLOCK_SCOPE,
+		) as CredentialBlockRow[];
+		return rows.map(row => ({
+			credentialId: row.credential_id,
+			providerKey: row.provider_key,
+			blockScope: row.block_scope,
+			blockedUntilMs: row.blocked_until_ms,
+			updatedAtMs: row.updated_at * 1000,
+		}));
 	}
 
 	tryAcquireCredentialRefreshLease(credentialId: number, owner: string, expiresAtMs: number): boolean {
@@ -2050,7 +2119,8 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore {
 		this.#upsertCacheStmt.finalize();
 		this.#deleteExpiredCacheStmt.finalize();
 		this.#getCredentialBlockStmt.finalize();
-		this.#listCredentialBlocksByCredentialStmt.finalize();
+		this.#listCredentialBlockScopesStmt.finalize();
+		this.#listCredentialBlocksStmt.finalize();
 		this.#upsertCredentialBlockStmt.finalize();
 		this.#deleteCredentialBlocksStmt.finalize();
 		this.#deleteCredentialBlockStmt.finalize();

package/src/auth/store.ts

@@ -98,6 +98,12 @@ export interface CredentialCacheStore {
 export interface CredentialBlockStore {
 	/** Non-expired block for one (credential, providerKey, scope) key, or undefined. */
 	getCredentialBlock(credentialId: number, providerKey: string, blockScope: string): number | undefined;
+	/**
+	 * Non-expired blocks for one credential and providerKey, keyed by block
+	 * scope (`""` is unscoped), in one read. Optional: selectors fall back to
+	 * {@link CredentialBlockStore.getCredentialBlock} per scope when absent.
+	 */
+	getCredentialBlockScopes?(credentialId: number, providerKey: string): ReadonlyMap<string, number>;
 	/** Earliest time a shared-store block should be eligible for live-usage reconciliation. */
 	getCredentialBlockReconcileAfter(credentialId: number, providerKey: string, blockScope: string): number | undefined;
 	/** Upsert with MAX semantics: keep the later blockedUntilMs on conflict. */
@@ -146,8 +152,8 @@ export interface UsageLedgerStore {
 
 /** Broker-delegated OAuth refresh and usage-report operations. */
 export interface CredentialUpstream {
-	/** Optional hook to notify the underlying store that usage report cache is stale. */
-	invalidateUsageCache(signal?: AbortSignal): Promise<void>;
+	/** Optional hook to notify the underlying store that usage reports (one provider's, when given) are stale. */
+	invalidateUsageCache(provider?: string, signal?: AbortSignal): Promise<void>;
 	/**
 	 * Optional store-supplied OAuth refresh. When present, `AuthStorage` uses
 	 * it before the per-provider local refresh path. `RemoteAuthCredentialStore`

package/src/auth/types.ts

@@ -257,6 +257,8 @@ export interface CheckCredentialsOptions {
 	completionProbe?: CompletionProbe;
 	/** Per-credential completion probe timeout (ms). Defaults to `timeoutMs`. */
 	completionTimeoutMs?: number;
+	/** Providers whose credentials are skipped entirely (not probed, not reported). */
+	excludeProviders?: ReadonlySet<string>;
 }
 
 // ─────────────────────────────────────────────────────────────────────────────
@@ -371,6 +373,12 @@ export type AuthStorageOptions = {
 		signal?: AbortSignal,
 		reason?: OAuthRefreshReason,
 	) => Promise<OAuthCredentials>;
+	/**
+	 * Set when {@link refreshOAuthCredential} exchanges the refresh token with the
+	 * provider itself instead of delegating (e.g. to a broker), so its tokens count
+	 * as this process's mints for auth-recovery reuse.
+	 */
+	refreshOAuthCredentialMints?: boolean;
 	/**
 	 * Human-readable description of the credential store backing this
 	 * AuthStorage instance. Surfaced through {@link AuthStorage.keys.describe}
@@ -470,7 +478,7 @@ export interface ModelUsageHealthOptions {
 export type AuthApiKeyOptions = {
 	baseUrl?: string;
 	modelId?: string;
-	/** Provider account ids known to serve `modelId` from multi-account discovery; OAuth selection prefers them and tries other accounts only as a last resort. */
+	/** {@link oauthAccountKey}s of accounts known to serve `modelId` from multi-account discovery; OAuth selection prefers them and tries other accounts only as a last resort. */
 	accountIds?: readonly string[];
 	/**
 	 * Caller's cancel signal. Threaded into any broker-bound OAuth refresh so
@@ -487,8 +495,26 @@ export type AuthApiKeyOptions = {
 	forceRefresh?: boolean;
 	/** Explicit provider-401 recovery; generic force refreshes leave this unset. */
 	refreshReason?: OAuthRefreshReason;
+	/** When false, select as `sessionId` would without recording the choice as that session's sticky credential. */
+	recordAffinity?: boolean;
 };
 
+/**
+ * Identity that multi-account discovery records in `Model.accountAccess` and
+ * that {@link AuthApiKeyOptions.accountIds} selection matches: the provider
+ * account id (Codex `chatgpt_account_id`), else the login email for providers
+ * whose credentials carry no account id (Antigravity), else the project id for
+ * a login whose email lookup failed. Same precedence as the credential store's
+ * identity key.
+ */
+export function oauthAccountKey(identity: {
+	accountId?: string;
+	email?: string;
+	projectId?: string;
+}): string | undefined {
+	return identity.accountId ?? identity.email ?? identity.projectId;
+}
+
 /** Non-secret identity bound to the OAuth credential selected for one request attempt. */
 export interface OAuthRequestIdentity {
 	orgId?: string;
@@ -890,9 +916,22 @@ export interface KeysApi {
 	 * 4. API key persisted by a successful `/login`
 	 * 5. Environment variable
 	 * 6. Stored API key (e.g. a broker-migrated copy) — last resort, so an explicit env var wins
+	 * 7. OMP compatibility mirror — below every stored credential, so a key the
+	 *    user saves locally always wins over the upstream fallback
+	 *
+	 * Resolves `undefined` when no permitted credential can serve, including when
+	 * OAuth refresh failed transiently (network, timeout, 5xx), so availability
+	 * probes move on to their next candidate. Request paths use
+	 * {@link KeysApi.getWithCredential}, which surfaces that failure.
 	 */
 	get(provider: string, sessionId?: string, options?: AuthApiKeyOptions): Promise<string | undefined>;
-	/** Resolve a bearer together with its durable stored credential row id, when known. */
+	/**
+	 * Resolve a bearer together with its durable stored credential row id, when known.
+	 * Same precedence as {@link KeysApi.get}, but rejects with
+	 * `OAuthRefreshUnavailableError` (transient, retryable) when OAuth refresh failed
+	 * with a retryable error and no other permitted source could serve the request;
+	 * the stored credential is kept.
+	 */
 	getWithCredential(
 		provider: string,
 		sessionId?: string,
@@ -951,10 +990,14 @@ export interface KeysApi {
 	 * `fallback: true` ranks the value below stored OAuth and `/login`
 	 * credentials instead, so a provider's default key reference cannot shadow
 	 * a key the user logged in with.
+	 *
+	 * `mirror: true` ranks the value below everything else, under stored
+	 * api_key credentials — a read-only compatibility mirror (upstream OMP
+	 * agent.db key) that must never shadow a locally stored key.
 	 */
-	setConfig(provider: string, apiKeyConfig: string, options?: { fallback?: boolean }): void;
+	setConfig(provider: string, apiKeyConfig: string, options?: { fallback?: boolean; mirror?: boolean }): void;
 	/**
-	 * Remove a single config-sourced API key (override or fallback).
+	 * Remove a single config-sourced API key (override, fallback, or mirror).
 	 */
 	removeConfig(provider: string): void;
 	/**
@@ -983,7 +1026,7 @@ export interface KeysApi {
 	resolver(provider: string, options?: { sessionId?: string; baseUrl?: string; modelId?: string }): ApiKeyResolver;
 }
 
-/** Controls whether a row-id refresh may reuse a token minted by this refresher. */
+/** Controls whether a row-id refresh may reuse a recently minted token. */
 export interface OAuthRefreshByIdOptions {
 	/**
 	 * Return the stored credential when it still holds a fresh access token this
@@ -991,6 +1034,12 @@ export interface OAuthRefreshByIdOptions {
 	 * refresh tokens repeatedly when a provider rejects every valid bearer.
 	 */
 	reuseRecentMint?: boolean;
+	/**
+	 * Provider-401 recovery intent. Forwarded to a delegated refresh (an auth
+	 * broker), which may then return a token it minted in the last five minutes.
+	 * Local reuse still requires {@link reuseRecentMint}.
+	 */
+	reason?: OAuthRefreshReason;
 }
 
 /** OAuth login, access, account identity, and refresh operations. */
@@ -1013,9 +1062,11 @@ export interface OAuthApi {
 	 * `enterpriseUrl`). For pure "give me the bytes for `Authorization`"
 	 * scenarios, prefer {@link AuthStorage.keys.get}.
 	 *
-	 * Returns `undefined` when no OAuth credential is available, the
-	 * credential fails to refresh, or runtime/config overrides have replaced
-	 * OAuth with an explicit API key.
+	 * Returns `undefined` when no usable OAuth credential is available (none
+	 * stored, or every one definitively failed to refresh) or runtime/config
+	 * overrides have replaced OAuth with an explicit API key. Rejects with
+	 * `OAuthRefreshUnavailableError` (transient, retryable) when a retryable
+	 * refresh failure left no usable credential.
 	 */
 	access(provider: string, sessionId?: string, options?: AuthApiKeyOptions): Promise<OAuthAccess | undefined>;
 	/**
@@ -1031,9 +1082,10 @@ export interface OAuthApi {
 	 * Resolve one stored OAuth credential by its durable storage row id.
 	 *
 	 * Unlike the normal session resolver, this method never ranks, rotates, or
-	 * falls back to sibling credentials. A forced refresh re-mints only the
-	 * requested row, preserving exact-account affinity for operations whose
-	 * provenance and policy boundary are tied to one workspace.
+	 * falls back to sibling credentials. A generic forced refresh re-mints only the
+	 * requested row; with `refreshReason: "auth-recovery"` it may instead reuse a
+	 * still-usable recent mint of that row. Either way exact-account affinity holds
+	 * for operations whose provenance and policy boundary are tied to one workspace.
 	 *
 	 * Returns `undefined` when the row does not exist for `provider` or an
 	 * explicit runtime/config API-key override suppresses OAuth.
@@ -1071,8 +1123,9 @@ export interface OAuthApi {
 	 * Refresh the OAuth credential with the given id through a per-credential
 	 * single-flight. Concurrent callers for the same row await the same upstream
 	 * refresh attempt, which is required for providers that rotate refresh tokens
-	 * on every successful refresh. Mints unconditionally unless
-	 * {@link OAuthRefreshByIdOptions.reuseRecentMint} is set.
+	 * on every successful refresh. Mints unless {@link OAuthRefreshByIdOptions.reuseRecentMint}
+	 * finds a recent local mint, or a delegated refresh given
+	 * {@link OAuthRefreshByIdOptions.reason} `"auth-recovery"` returns its recent mint.
 	 */
 	refresh(id: number, signal?: AbortSignal, options?: OAuthRefreshByIdOptions): Promise<AuthCredentialSnapshotEntry>;
 	/**
@@ -1084,6 +1137,12 @@ export interface OAuthApi {
 	): Promise<StoredOAuthRefreshResult<T>>;
 }
 
+/**
+ * Handle for one {@link SessionsApi.restrict} call. {@link SessionsApi.unrestrict}
+ * lifts a restriction only with the lease of the call that installed it.
+ */
+export type SessionRestrictionLease = symbol;
+
 /** Session credential affinity operations. */
 export interface SessionsApi {
 	/**
@@ -1101,14 +1160,46 @@ export interface SessionsApi {
 	 */
 	pin(provider: string, sessionId: string, credentialId: number, options?: { restoredAtMs?: number }): boolean;
 	/**
-	 * Copy every stored credential affinity from one live session to another.
+	 * Copy stored credential affinities from one live session to another.
 	 *
 	 * The target receives its own sticky entries, so request resolution, usage
 	 * blocking, credential rotation, metadata, and persisted pins all continue
 	 * through the target session id without retaining a live dependency on the
-	 * source session.
-	 */
-	inherit(sourceSessionId: string, targetSessionId: string): number;
+	 * source session. `include` limits the copy to the providers it accepts,
+	 * given whether the source's affinity is an explicit user pin.
+	 */
+	inherit(
+		sourceSessionId: string,
+		targetSessionId: string,
+		include?: (provider: string, explicit: boolean) => boolean,
+	): number;
+	/**
+	 * Restrict one session's credentials for `provider` to the OAuth accounts
+	 * whose identity key (`email:<address>|org:<id>` for org-scoped providers;
+	 * otherwise `account:`, `email:`, or `project:` — the keys broker account
+	 * pools use) is listed. Selection, session pins, inherited affinity,
+	 * fallback passes, and rotation never leave the list, and runtime, config,
+	 * environment, and stored API keys are not used for that session; when no
+	 * listed account can serve, key resolution fails instead of borrowing
+	 * another account. An empty list allows no account. Replaces any earlier
+	 * restriction for the same provider and session; {@link inherit} does not
+	 * copy restrictions. Restrictions are never evicted, because that would
+	 * widen a live session; the owner passes the returned lease to
+	 * {@link unrestrict} when the session ends.
+	 */
+	restrict(provider: string, sessionId: string, identityKeys: readonly string[]): SessionRestrictionLease;
+	/**
+	 * Lift the restriction that `lease` installed. Once a later {@link restrict}
+	 * call has replaced it for the same provider and session, this does nothing,
+	 * so a stale owner never lifts the restriction a newer owner relies on.
+	 */
+	unrestrict(provider: string, sessionId: string, lease: SessionRestrictionLease): void;
+	/**
+	 * Whether stored credential `credentialId` may serve `sessionId`: always for
+	 * an unrestricted session, otherwise only an OAuth account its
+	 * {@link restrict} list names.
+	 */
+	permits(provider: string, sessionId: string, credentialId: number): boolean;
 	/**
 	 * Release a session's sticky credential so its next {@link getApiKey} call
 	 * re-runs native pool ranking. This never blocks or penalizes the released

package/src/auth/usage-cache.ts

@@ -6,6 +6,7 @@ import type { UsageCredential, UsageProvider, UsageReport } from "../usage";
 
 const USAGE_CACHE_PREFIX = "usage_cache:";
 const USAGE_FORCE_REFRESH_CACHE_PREFIX = "force-refresh:";
+const USAGE_REPORT_KEY_PREFIX = "report:";
 /** Minimum interval between non-exhausted header snapshots; used by UsageService. */
 export const USAGE_HEADER_INGEST_INTERVAL_MS = 60_000;
 const USAGE_LAST_GOOD_RETENTION_MS = 24 * 60 * 60_000;
@@ -45,6 +46,13 @@ function parseUsageCacheEntry<T>(raw: string): UsageCacheEntry<T> | undefined {
 	}
 }
 
+/** Cached reports drop the provider's upstream `raw` payload; only a fresh fetch carries it. */
+function withoutRawPayload(value: unknown): unknown {
+	if (value === null || typeof value !== "object" || !("raw" in value)) return value;
+	const { raw: _raw, ...rest } = value;
+	return rest;
+}
+
 /** Convert a stored credential for usage providers; used by usage and health services. */
 export function buildUsageCredential(credential: AuthCredential): UsageCredential {
 	if (credential.type === "api_key") {
@@ -196,12 +204,13 @@ export class UsageCache {
 	}
 
 	set<T>(key: string, entry: UsageCacheEntry<T>): void {
+		const value = withoutRawPayload(entry.value);
 		const payload = JSON.stringify({
-			value: entry.value,
+			value,
 			expiresAt: entry.expiresAt,
 		});
 		const durableExpiresAt =
-			entry.value === null ? entry.expiresAt : Math.max(entry.expiresAt, Date.now() + USAGE_LAST_GOOD_RETENTION_MS);
+			value === null ? entry.expiresAt : Math.max(entry.expiresAt, Date.now() + USAGE_LAST_GOOD_RETENTION_MS);
 		this.#store.setCache(`${USAGE_CACHE_PREFIX}${key}`, payload, Math.floor(durableExpiresAt / 1000));
 	}
 
@@ -272,7 +281,7 @@ export class UsageCache {
 		const baseUrl = this.#normalizeUsageBaseUrl(request.baseUrl) || "default";
 		const identity = usageCacheIdentity(request.credential);
 		const providerKey = this.#usageCacheProviderKey(request.provider);
-		const cacheKey = `report:${providerKey}:${baseUrl}:${identity}`;
+		const cacheKey = `${USAGE_REPORT_KEY_PREFIX}${providerKey}:${baseUrl}:${identity}`;
 		const cacheKeys = this.#usageReportCacheKeysByProvider.get(request.provider) ?? new Set<string>();
 		cacheKeys.add(cacheKey);
 		this.#usageReportCacheKeysByProvider.set(request.provider, cacheKeys);
@@ -334,7 +343,7 @@ export class UsageCache {
 		this.#epoch += 1;
 		this.#bumpRefreshEpoch(provider);
 		const expired = Date.now() - 1;
-		const prefix = `report:${this.#usageCacheProviderKey(provider)}:`;
+		const prefix = `${USAGE_REPORT_KEY_PREFIX}${this.#usageCacheProviderKey(provider)}:`;
 		if (this.deletePrefix(prefix)) return;
 		const cacheKeys = new Set(this.#usageReportCacheKeysByProvider.get(provider));
 		for (const entry of this.#pool.entries(provider)) {
@@ -361,7 +370,9 @@ export class UsageCache {
 	): Promise<void> {
 		this.#epoch += 1;
 		this.#bumpRefreshEpoch(provider);
-		const prefix = provider ? `report:${this.#usageCacheProviderKey(provider)}:` : "report:";
+		const prefix = provider
+			? `${USAGE_REPORT_KEY_PREFIX}${this.#usageCacheProviderKey(provider)}:`
+			: USAGE_REPORT_KEY_PREFIX;
 		if (!this.deletePrefix(prefix)) {
 			// Third-party stores may not support prefix deletion. Clear every active
 			// request key instead, including API-key and environment credentials.

package/src/auth/usage.ts

@@ -1,7 +1,7 @@
 import { authPolicyFor } from "@linxiraos/pi-catalog/compat/auth";
 import { $pickenv, logger } from "@linxiraos/pi-utils";
+import { getEnvApiKey } from "../env-api-key";
 import * as AIError from "../error";
-import { getEnvApiKey } from "../stream";
 import type { OAuthCredentials } from "../registry/oauth/types";
 import type { Provider } from "../types";
 import { resolveUsedFraction } from "../usage";
@@ -871,7 +871,7 @@ export class UsageService implements UsageApi {
 		await this.#deps.cache.clearReports(provider, () => this.#collectUsageRequests());
 
 		if (this.#deps.store.invalidateUsageCache) {
-			await this.#deps.store.invalidateUsageCache(signal).catch(err => {
+			await this.#deps.store.invalidateUsageCache(provider, signal).catch(err => {
 				logger.debug("Failed to notify store of stale usage", { err });
 			});
 		}

package/src/auth-broker/client.ts

@@ -1,8 +1,8 @@
 /**
- * HTTP client for the omp auth-broker server.
+ * HTTP client for the zetacode auth-broker server.
  *
  * Used by {@link RemoteAuthCredentialStore} (snapshot pulls) and by
- * `omp auth-broker status` (liveness checks). All endpoints except
+ * `zetacode auth-broker status` (liveness checks). All endpoints except
  * `/v1/healthz` require a bearer token.
  */
 
@@ -309,8 +309,11 @@ export class AuthBrokerClient {
 		});
 	}
 
-	notifyUsageStale(signal?: AbortSignal): Promise<UsageStaleResponse> {
-		return this.#request<UsageStaleResponse>("POST", "/v1/usage/stale", {
+	notifyUsageStale(provider?: string, signal?: AbortSignal): Promise<UsageStaleResponse> {
+		const params = new URLSearchParams();
+		if (provider) params.set("provider", provider);
+		const path = `/v1/usage/stale${params.size > 0 ? `?${params.toString()}` : ""}`;
+		return this.#request<UsageStaleResponse>("POST", path, {
 			schema: "usageStaleResponseSchema",
 			signal,
 		});