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

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. */

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 "@oh-my-pi/pi-catalog/compat/auth";
 import { $pickenv, logger } from "@oh-my-pi/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";

package/src/auth-broker/discover.ts

@@ -28,7 +28,7 @@ import {
 import * as AIError from "../error";
 import { AuthBrokerClient, AuthBrokerError } from "./client";
 import { type AuthBrokerAccountPool, RemoteAuthCredentialStore } from "./remote-store";
-import { readAuthBrokerSnapshotCache, writeAuthBrokerSnapshotCache } from "./snapshot-cache";
+import { readAuthBrokerSnapshotCache, scheduleAuthBrokerSnapshotCacheWrite } from "./snapshot-cache";
 import { DEFAULT_SNAPSHOT_CACHE_TTL_MS, type SnapshotResponse } from "./types";
 
 export interface AuthBrokerClientConfig {
@@ -396,16 +396,17 @@ export async function openAuthCredentialStore(
 		const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token });
 		const cachePath = options.cachePath ?? getAuthBrokerSnapshotCachePath();
 		const ttlMs = resolveSnapshotTtlMs();
+		// Coalesced fire-and-forget: the first snapshot lands immediately (so a
+		// fresh boot seeds the cache), bursts collapse into one trailing write,
+		// and anything pending is flushed on shutdown.
 		const persist =
 			ttlMs > 0
 				? (snapshot: SnapshotResponse): void => {
-						void writeAuthBrokerSnapshotCache({
+						scheduleAuthBrokerSnapshotCacheWrite({
 							path: cachePath,
 							token: brokerConfig.token,
 							url: brokerConfig.url,
 							snapshot,
-						}).catch(error => {
-							logger.debug("auth-broker snapshot cache write failed", { error: String(error) });
 						});
 					}
 				: undefined;