@oh-my-pi/pi-ai 18.6.3 → 18.8.0

package/src/auth/affinity.ts

@@ -1,5 +1,6 @@
 import { logger } from "@oh-my-pi/pi-utils";
-import { getEnvApiKey } from "../stream";
+import { LRUCache } from "@oh-my-pi/pi-utils/lru";
+import { getEnvApiKey } from "../env-api-key";
 import type { AuthCredential, OAuthCredential, SessionsApi } from "./types";
 import type { AuthCredentialStore } from "./store";
 import type { CredentialPool } from "./pool";
@@ -14,6 +15,11 @@ const SESSION_STICKY_TTL_SEC = 30 * 24 * 60 * 60;
  * in-memory sticky stays exact; only the copy other processes resume from lags.
  */
 const SESSION_STICKY_PERSIST_INTERVAL_MS = 60_000;
+/**
+ * In-memory pins kept per provider. Long-lived gateways mint a session id per
+ * conversation; evicted sessions resume from their persisted sticky row.
+ */
+const SESSION_AFFINITY_MAX_SESSIONS_PER_PROVIDER = 256;
 
 /** What this process last wrote (or read) for one persisted sticky row. */
 type PersistedSticky = {
@@ -36,9 +42,9 @@ export type SessionCredential = {
 /** Session → credential affinity (pins), persisted in the store cache. */
 export class SessionAffinity implements SessionsApi {
 	/** Tracks the last used credential per provider for a session (used for rate-limit switching). */
-	#sessionLastCredential: Map<string, Map<string, SessionCredential>> = new Map();
-	/** Persisted sticky rows keyed by cache key, so unchanged re-records skip the write. */
-	#persistedSticky: Map<string, PersistedSticky> = new Map();
+	#sessionLastCredential: Map<string, LRUCache<string, SessionCredential>> = new Map();
+	/** Persisted sticky rows per provider, keyed by session id, so unchanged re-records skip the write. */
+	#persistedSticky: Map<string, LRUCache<string, PersistedSticky>> = new Map();
 	#store: AuthCredentialStore;
 	#pool: CredentialPool;
 	#overrides: KeyOverrides;
@@ -49,13 +55,21 @@ export class SessionAffinity implements SessionsApi {
 		this.#overrides = overrides;
 	}
 
+	/** Bounded per-provider session map, created on first use. */
+	static #sessionsFor<V>(maps: Map<string, LRUCache<string, V>>, provider: string): LRUCache<string, V> {
+		let sessions = maps.get(provider);
+		if (!sessions) {
+			sessions = new LRUCache<string, V>({ max: SESSION_AFFINITY_MAX_SESSIONS_PER_PROVIDER });
+			maps.set(provider, sessions);
+		}
+		return sessions;
+	}
+
 	/** Drop every pin for a provider, both in memory and in the persisted cache. */
 	clearProvider(provider: string): void {
 		this.#sessionLastCredential.delete(provider);
+		this.#persistedSticky.delete(provider);
 		const prefix = `${SESSION_STICKY_CACHE_PREFIX}${provider}:`;
-		for (const cacheKey of this.#persistedSticky.keys()) {
-			if (cacheKey.startsWith(prefix)) this.#persistedSticky.delete(cacheKey);
-		}
 		try {
 			this.#store.deleteCachePrefix?.(prefix);
 		} catch (err) {
@@ -90,7 +104,7 @@ export class SessionAffinity implements SessionsApi {
 		if (!sessionId) return;
 		const nowMs = lastUsedAtMs ?? Date.now();
 		const credentialId = this.#pool.entries(provider)[index]?.id;
-		const sessionMap = this.#sessionLastCredential.get(provider) ?? new Map();
+		const sessionMap = SessionAffinity.#sessionsFor(this.#sessionLastCredential, provider);
 		const previous = sessionMap.get(sessionId);
 		const sameCredential =
 			previous?.type === type &&
@@ -104,12 +118,12 @@ export class SessionAffinity implements SessionsApi {
 			...(isExplicit ? { explicit: true as const } : {}),
 		};
 		sessionMap.set(sessionId, sessionCredential);
-		this.#sessionLastCredential.set(provider, sessionMap);
 
 		if (credentialId === undefined) return;
 		const cacheKey = `${SESSION_STICKY_CACHE_PREFIX}${provider}:${sessionId}`;
 		const expiresAtSec = Math.floor(nowMs / 1000) + SESSION_STICKY_TTL_SEC;
-		const persisted = this.#persistedSticky.get(cacheKey);
+		const persistedSessions = SessionAffinity.#sessionsFor(this.#persistedSticky, provider);
+		const persisted = persistedSessions.get(sessionId);
 		if (
 			persisted &&
 			persisted.type === type &&
@@ -121,14 +135,14 @@ export class SessionAffinity implements SessionsApi {
 		}
 		try {
 			this.#store.setCache(cacheKey, JSON.stringify(sessionCredential), expiresAtSec);
-			this.#persistedSticky.set(cacheKey, {
+			persistedSessions.set(sessionId, {
 				type,
 				credentialId,
 				explicit: isExplicit,
 				lastUsedAtMs: nowMs,
 			});
 		} catch (err) {
-			this.#persistedSticky.delete(cacheKey);
+			persistedSessions.delete(sessionId);
 			logger.debug("Failed to write session sticky credential to persistent store cache", { err });
 		}
 	}
@@ -136,8 +150,7 @@ export class SessionAffinity implements SessionsApi {
 	/** Retrieves the last credential used by a session. */
 	get(provider: string, sessionId: string | undefined): SessionCredential | undefined {
 		if (!sessionId) return undefined;
-		let sessionMap = this.#sessionLastCredential.get(provider);
-		const live = sessionMap?.get(sessionId);
+		const live = this.#sessionLastCredential.get(provider)?.get(sessionId);
 		if (live) {
 			// Another process can add or drop rows mid-session and the pool is an
 			// index-ordered snapshot, so re-resolve the pin through its durable row
@@ -147,7 +160,7 @@ export class SessionAffinity implements SessionsApi {
 			const stored = this.#pool.entries(provider);
 			const actualIndex = stored.findIndex(entry => entry.id === live.credentialId);
 			if (actualIndex === -1 || stored[actualIndex]?.credential.type !== live.type) {
-				sessionMap?.delete(sessionId);
+				this.#sessionLastCredential.get(provider)?.delete(sessionId);
 				return undefined;
 			}
 			live.index = actualIndex;
@@ -163,22 +176,18 @@ export class SessionAffinity implements SessionsApi {
 					const stored = this.#pool.entries(provider);
 					const actualIndex = stored.findIndex(entry => entry.id === val.credentialId);
 					if (actualIndex === -1 || stored[actualIndex]?.credential.type !== val.type) {
-						this.#persistedSticky.delete(cacheKey);
+						this.#persistedSticky.get(provider)?.delete(sessionId);
 						this.#store.setCache(cacheKey, "", 0);
 						return undefined;
 					}
 					val.index = actualIndex;
 				} else {
 					// Fallback: drop unsafe index-only cache rows to prevent wrong-account routing
-					this.#persistedSticky.delete(cacheKey);
+					this.#persistedSticky.get(provider)?.delete(sessionId);
 					this.#store.setCache(cacheKey, "", 0);
 					return undefined;
 				}
 
-				if (!sessionMap) {
-					sessionMap = new Map();
-					this.#sessionLastCredential.set(provider, sessionMap);
-				}
 				const sessionVal: SessionCredential = {
 					type: val.type,
 					index: val.index,
@@ -187,14 +196,14 @@ export class SessionAffinity implements SessionsApi {
 					...(val.explicit === true ? { explicit: true } : {}),
 				};
 				if (typeof val.lastUsedAtMs === "number") {
-					this.#persistedSticky.set(cacheKey, {
+					SessionAffinity.#sessionsFor(this.#persistedSticky, provider).set(sessionId, {
 						type: val.type,
 						credentialId: val.credentialId,
 						explicit: val.explicit === true,
 						lastUsedAtMs: val.lastUsedAtMs,
 					});
 				}
-				sessionMap.set(sessionId, sessionVal);
+				SessionAffinity.#sessionsFor(this.#sessionLastCredential, provider).set(sessionId, sessionVal);
 				return sessionVal;
 			}
 		} catch (err) {
@@ -214,7 +223,7 @@ export class SessionAffinity implements SessionsApi {
 			}
 		}
 		const cacheKey = `${SESSION_STICKY_CACHE_PREFIX}${provider}:${sessionId}`;
-		this.#persistedSticky.delete(cacheKey);
+		this.#persistedSticky.get(provider)?.delete(sessionId);
 		try {
 			this.#store.setCache(cacheKey, "", 0);
 		} catch (err) {

package/src/auth/blocks.ts

@@ -64,6 +64,14 @@ export function credentialBlockScopesForRequest(
  */
 type CredentialBackoff = { until: number; timed: boolean; probeAfter: number; persisted?: boolean };
 
+/**
+ * One read of a credential's persisted blocks under a providerKey, keyed by
+ * scope (`""` is unscoped). `undefined`: the store has no bulk read, so use
+ * per-scope point reads. `null`: unreadable (damaged store or failed read), so
+ * it neither contributes deadlines nor retires mirrored in-memory blocks.
+ */
+type PersistedBlockScopes = ReadonlyMap<string, number> | null | undefined;
+
 /** Latch for unrecoverably corrupt persisted block stores, shared with the credential pool. */
 export class BlockStoreHealth {
 	readonly #sourceLabel: string | undefined;
@@ -134,8 +142,16 @@ export class CredentialBlocks implements BlocksApi {
 		}
 	}
 
-	/** Returns in-memory block expiry timestamp for a credential/key pair, cleaning up expired entries. */
-	#getCredentialBlockedUntilForKey(backoffKey: string, credentialId: number, nowMs: number): number | undefined {
+	/**
+	 * Returns in-memory block expiry timestamp for a credential/key pair, cleaning up expired entries.
+	 * `persistedScopes` is a bulk read for the key's providerKey; omitted, the store is point-read.
+	 */
+	#getCredentialBlockedUntilForKey(
+		backoffKey: string,
+		credentialId: number,
+		nowMs: number,
+		persistedScopes: PersistedBlockScopes = undefined,
+	): number | undefined {
 		const block = this.#credentialBackoff.get(backoffKey)?.get(credentialId);
 		if (!block) return undefined;
 		// Once mirrored successfully, the store is authoritative for cross-process
@@ -143,12 +159,19 @@ export class CredentialBlocks implements BlocksApi {
 		// Upserts are longest-wins, so a persisted row shorter than the deadline this
 		// process mirrored means the row was deleted and replaced: the old deadline
 		// was superseded and must not win the next longest-wins merge.
-		if (block.persisted && !this.#deps.health.damaged && this.#deps.store.getCredentialBlock) {
+		if (
+			block.persisted &&
+			!this.#deps.health.damaged &&
+			persistedScopes !== null &&
+			(persistedScopes !== undefined || this.#deps.store.getCredentialBlock)
+		) {
 			const separator = backoffKey.indexOf("\0");
 			const providerKey = separator < 0 ? backoffKey : backoffKey.slice(0, separator);
 			const scope = separator < 0 ? "" : backoffKey.slice(separator + 1);
 			try {
-				const persisted = this.#deps.store.getCredentialBlock(credentialId, providerKey, scope);
+				const persisted = persistedScopes
+					? persistedScopes.get(scope)
+					: this.#deps.store.getCredentialBlock?.(credentialId, providerKey, scope);
 				if (persisted === undefined || persisted < block.until) {
 					this.#deleteCredentialBackoff(backoffKey, credentialId);
 					return undefined;
@@ -175,7 +198,9 @@ export class CredentialBlocks implements BlocksApi {
 		credentialId: number,
 		providerKey: string,
 		blockScope: string | undefined,
+		persistedScopes: PersistedBlockScopes = undefined,
 	): number | undefined {
+		if (persistedScopes !== undefined) return persistedScopes?.get(blockScope ?? "");
 		if (this.#deps.health.damaged) return undefined;
 		const getCredentialBlock = this.#deps.store.getCredentialBlock?.bind(this.#deps.store);
 		if (!getCredentialBlock) return undefined;
@@ -193,6 +218,20 @@ export class CredentialBlocks implements BlocksApi {
 		}
 	}
 
+	/** Every persisted scope of one credential/providerKey in a single store read, when supported. */
+	#readPersistedCredentialBlockScopes(credentialId: number, providerKey: string): PersistedBlockScopes {
+		const store = this.#deps.store;
+		if (!store.getCredentialBlockScopes) return undefined;
+		if (this.#deps.health.damaged) return null;
+		try {
+			return store.getCredentialBlockScopes(credentialId, providerKey);
+		} catch (err) {
+			if (this.#deps.health.handle(err)) return null;
+			logger.debug("Failed to read credential blocks from persistent store", { err, credentialId, providerKey });
+			return null;
+		}
+	}
+
 	#readPersistedCredentialBlockReconcileAfter(credentialId: number, providerKey: string, blockScope: string): number {
 		if (this.#deps.health.damaged) return 0;
 		const getCredentialBlockReconcileAfter = this.#deps.store.getCredentialBlockReconcileAfter?.bind(
@@ -215,13 +254,19 @@ export class CredentialBlocks implements BlocksApi {
 		}
 	}
 
-	#exactBlockedUntil(credentialId: number, providerKey: string, scope: string | undefined): number | undefined {
+	#exactBlockedUntil(
+		credentialId: number,
+		providerKey: string,
+		scope: string | undefined,
+		persistedScopes: PersistedBlockScopes = this.#readPersistedCredentialBlockScopes(credentialId, providerKey),
+	): number | undefined {
 		const memory = this.#getCredentialBlockedUntilForKey(
 			scopedBackoffKey(providerKey, scope),
 			credentialId,
 			Date.now(),
+			persistedScopes,
 		);
-		const persisted = this.#readPersistedCredentialBlock(credentialId, providerKey, scope);
+		const persisted = this.#readPersistedCredentialBlock(credentialId, providerKey, scope, persistedScopes);
 		return memory === undefined ? persisted : persisted === undefined ? memory : Math.max(memory, persisted);
 	}
 
@@ -231,6 +276,22 @@ export class CredentialBlocks implements BlocksApi {
 		providerKey: string,
 		credentialIndex: number,
 		blockScopeOrScopes: string | readonly string[] | undefined = undefined,
+	): number | undefined {
+		const credentialId = this.#deps.pool.entries(provider)[credentialIndex]?.id;
+		if (credentialId === undefined) return undefined;
+		return this.#blockedUntilById(
+			credentialId,
+			providerKey,
+			blockScopeOrScopes,
+			this.#readPersistedCredentialBlockScopes(credentialId, providerKey),
+		);
+	}
+
+	#blockedUntilById(
+		credentialId: number,
+		providerKey: string,
+		blockScopeOrScopes: string | readonly string[] | undefined,
+		persistedScopes: PersistedBlockScopes,
 	): number | undefined {
 		const nowMs = Date.now();
 		// A request honours its own scope plus any legacy catch-all scope, so a
@@ -239,21 +300,25 @@ export class CredentialBlocks implements BlocksApi {
 			...PROTECTED_BLOCK_SCOPES,
 			...(typeof blockScopeOrScopes === "string" ? [blockScopeOrScopes] : (blockScopeOrScopes ?? [])),
 		].filter(scope => scope.length > 0);
-		const credentialId = this.#deps.pool.entries(provider)[credentialIndex]?.id;
-		if (credentialId === undefined) return undefined;
-		let blockedUntil = this.#getCredentialBlockedUntilForKey(providerKey, credentialId, nowMs);
+		let blockedUntil = this.#getCredentialBlockedUntilForKey(providerKey, credentialId, nowMs, persistedScopes);
 		for (const blockScope of scopes) {
 			const scopedBlockedUntil = this.#getCredentialBlockedUntilForKey(
 				scopedBackoffKey(providerKey, blockScope),
 				credentialId,
 				nowMs,
+				persistedScopes,
 			);
 			if (scopedBlockedUntil !== undefined && (blockedUntil === undefined || scopedBlockedUntil > blockedUntil)) {
 				blockedUntil = scopedBlockedUntil;
 			}
 		}
 
-		const persistedGlobalBlockedUntil = this.#readPersistedCredentialBlock(credentialId, providerKey, "");
+		const persistedGlobalBlockedUntil = this.#readPersistedCredentialBlock(
+			credentialId,
+			providerKey,
+			"",
+			persistedScopes,
+		);
 		if (
 			persistedGlobalBlockedUntil !== undefined &&
 			(blockedUntil === undefined || persistedGlobalBlockedUntil > blockedUntil)
@@ -261,7 +326,12 @@ export class CredentialBlocks implements BlocksApi {
 			blockedUntil = persistedGlobalBlockedUntil;
 		}
 		for (const blockScope of scopes) {
-			const persistedScopedBlockedUntil = this.#readPersistedCredentialBlock(credentialId, providerKey, blockScope);
+			const persistedScopedBlockedUntil = this.#readPersistedCredentialBlock(
+				credentialId,
+				providerKey,
+				blockScope,
+				persistedScopes,
+			);
 			if (
 				persistedScopedBlockedUntil !== undefined &&
 				(blockedUntil === undefined || persistedScopedBlockedUntil > blockedUntil)
@@ -440,14 +510,20 @@ export class CredentialBlocks implements BlocksApi {
 		if (!this.supportsHealing(provider)) return false;
 		const credentialId = this.#deps.pool.entries(provider)[credentialIndex]?.id;
 		if (credentialId === undefined) return false;
-		if (PROTECTED_BLOCK_SCOPES.some(scope => this.#exactBlockedUntil(credentialId, providerKey, scope) !== undefined))
+		// One persisted read serves the protected-scope, global, and requested checks.
+		const persistedScopes = this.#readPersistedCredentialBlockScopes(credentialId, providerKey);
+		if (
+			PROTECTED_BLOCK_SCOPES.some(
+				scope => this.#exactBlockedUntil(credentialId, providerKey, scope, persistedScopes) !== undefined,
+			)
+		)
 			return false;
 		if (
 			!this.#deps.strategies(provider)?.healsGlobalBlocks &&
-			this.blockedUntil(provider, providerKey, credentialIndex) !== undefined
+			this.#blockedUntilById(credentialId, providerKey, undefined, persistedScopes) !== undefined
 		)
 			return false;
-		return this.blockedUntil(provider, providerKey, credentialIndex, blockScopeOrScopes) !== undefined;
+		return this.#blockedUntilById(credentialId, providerKey, blockScopeOrScopes, persistedScopes) !== undefined;
 	}
 
 	/**

package/src/auth/cascade.ts

@@ -1,10 +1,10 @@
 import { authPolicyFor } from "@oh-my-pi/pi-catalog/compat/auth";
 import { $env, $envExact } from "@oh-my-pi/pi-utils";
 import { type ApiKeyResolver, markAfterSiblingWait, type ResolvedApiKey } from "../auth-retry";
+import { getEnvApiKey, getEnvApiKeyName } from "../env-api-key";
 import * as AIError from "../error";
 import { isUsageLimitOutcome } from "../error/rate-limit";
 import { AUTHENTICATED_SENTINEL } from "../registry/types";
-import { getEnvApiKey, getEnvApiKeyName } from "../stream";
 import type { SessionAffinity } from "./affinity";
 import type { CredentialPool } from "./pool";
 import type { CredentialSelector } from "./select";

package/src/auth/select.ts

@@ -28,7 +28,13 @@ import {
 } from "./rank";
 import { mergeRefreshedCredential, OAUTH_REFRESH_SKEW_MS, type OAuthRefresher } from "./refresh";
 import type { AuthCredentialStore } from "./store";
-import type { ApiKeyCredential, AuthApiKeyOptions, AuthCredential, OAuthCredential } from "./types";
+import type {
+	ApiKeyCredential,
+	AuthApiKeyOptions,
+	AuthCredential,
+	OAuthCredential,
+	StoredAuthCredential,
+} from "./types";
 import type { UsageService } from "./usage";
 import {
 	isUsageLimitReached,
@@ -690,6 +696,10 @@ export class CredentialSelector {
 		const forceRefreshIndex = options?.forceRefresh
 			? (sessionPreferredIndex ?? candidates[0]?.selection.index)
 			: undefined;
+		// Each candidate's synchronous prefix below runs back to back inside `map`,
+		// before any await, so one provider re-list serves every initial rebind.
+		// Resyncs after an await (refresh, disable) re-read the store.
+		let preflightRows: StoredAuthCredential[] | undefined;
 		await Promise.all(
 			candidates.map(async candidate => {
 				const force = forceRefreshIndex !== undefined && candidate.selection.index === forceRefreshIndex;
@@ -697,7 +707,8 @@ export class CredentialSelector {
 				let syncedPeerCredential = false;
 				if (initialCredentialId !== undefined) {
 					const beforeSync = candidate.selection.credential;
-					if (!this.#syncOAuthSelectionFromStore(provider, candidate.selection, initialCredentialId)) return;
+					preflightRows ??= this.#reloadProviderRows(provider);
+					if (!this.#rebindOAuthSelection(preflightRows, candidate.selection, initialCredentialId)) return;
 					syncedPeerCredential = !authCredentialEquals(beforeSync, candidate.selection.credential);
 				}
 				const hasFreshAccess = Date.now() + OAUTH_REFRESH_SKEW_MS < candidate.selection.credential.expires;
@@ -881,11 +892,25 @@ export class CredentialSelector {
 		selection: { credential: OAuthCredential; index: number },
 		credentialId: number,
 	): boolean {
+		return this.#rebindOAuthSelection(this.#reloadProviderRows(provider), selection, credentialId);
+	}
+
+	/** Re-list one provider's active rows and adopt them in the pool. */
+	#reloadProviderRows(provider: string): StoredAuthCredential[] {
 		const latestRows = this.#deps.store.listAuthCredentials(provider);
 		this.#deps.pool.replace(
 			provider,
 			latestRows.map(row => ({ id: row.id, credential: row.credential })),
 		);
+		return latestRows;
+	}
+
+	/** Point `selection` at row `credentialId` in `latestRows`; false when it is gone or no longer OAuth. */
+	#rebindOAuthSelection(
+		latestRows: readonly StoredAuthCredential[],
+		selection: { credential: OAuthCredential; index: number },
+		credentialId: number,
+	): boolean {
 		const latestIndex = latestRows.findIndex(row => row.id === credentialId);
 		if (latestIndex === -1) return false;
 		const latest = latestRows[latestIndex];

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

@@ -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
@@ -386,11 +392,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 +484,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)
@@ -1585,27 +1600,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 +1648,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 +1662,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 +1692,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 +2092,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. */