@oh-my-pi/pi-ai 18.12.2 → 18.12.4

package/src/auth-broker/discover.ts

@@ -27,7 +27,7 @@ import {
 } from "../auth-storage";
 import * as AIError from "../error";
 import { AuthBrokerClient, AuthBrokerError } from "./client";
-import { type AuthBrokerAccountPool, RemoteAuthCredentialStore } from "./remote-store";
+import { type AuthBrokerAccountPool, type AuthBrokerProviderPool, RemoteAuthCredentialStore } from "./remote-store";
 import { readAuthBrokerSnapshotCache, scheduleAuthBrokerSnapshotCacheWrite } from "./snapshot-cache";
 import { DEFAULT_SNAPSHOT_CACHE_TTL_MS, type SnapshotResponse } from "./types";
 
@@ -122,7 +122,12 @@ function parseAuthAccountPolicies(value: unknown): AuthAccountPolicies {
 		}
 		const policy = entry as Record<string, unknown>;
 		const unknownPolicyFields = Object.keys(policy).filter(
-			key => key !== "provider" && key !== "account" && key !== "priority" && key !== "reservePct",
+			key =>
+				key !== "provider" &&
+				key !== "account" &&
+				key !== "priority" &&
+				key !== "reservePct" &&
+				key !== "reclaimAbovePct",
 		);
 		if (unknownPolicyFields.length > 0) {
 			throw new AIError.ConfigurationError(`${path} has unknown fields: ${unknownPolicyFields.join(", ")}`);
@@ -171,6 +176,12 @@ function parseAuthAccountPolicies(value: unknown): AuthAccountPolicies {
 		) {
 			throw new AIError.ConfigurationError(`${path}.reservePct must be between 0 and 100`);
 		}
+		if (
+			policy.reclaimAbovePct !== undefined &&
+			(typeof policy.reclaimAbovePct !== "number" || !Number.isFinite(policy.reclaimAbovePct))
+		) {
+			throw new AIError.ConfigurationError(`${path}.reclaimAbovePct must be a finite number`);
+		}
 
 		return {
 			provider,
@@ -182,6 +193,7 @@ function parseAuthAccountPolicies(value: unknown): AuthAccountPolicies {
 			},
 			...(typeof policy.priority === "number" ? { priority: policy.priority } : {}),
 			...(typeof policy.reservePct === "number" ? { reservePct: policy.reservePct } : {}),
+			...(typeof policy.reclaimAbovePct === "number" ? { reclaimAbovePct: policy.reclaimAbovePct } : {}),
 		};
 	});
 }
@@ -272,7 +284,7 @@ export async function loadAuthBrokerAccountPool(): Promise<AuthBrokerAccountPool
 		throw new AIError.ConfigurationError("OMP_AUTH_BROKER_ACCOUNT_POOL_FILE must contain a JSON object");
 	}
 
-	const accountPool = new Map<string, ReadonlySet<string>>();
+	const accountPool = new Map<string, ReadonlySet<string> | AuthBrokerProviderPool>();
 	for (const [provider, value] of Object.entries(parsed)) {
 		const normalizedProvider = provider.trim();
 		if (normalizedProvider.length === 0) {
@@ -283,28 +295,66 @@ export async function loadAuthBrokerAccountPool(): Promise<AuthBrokerAccountPool
 				"OMP_AUTH_BROKER_ACCOUNT_POOL_FILE contains a provider id with surrounding whitespace",
 			);
 		}
-		if (!Array.isArray(value)) {
-			throw new AIError.ConfigurationError(
-				`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} must be an array of identity keys`,
-			);
-		}
-		const identities = new Set<string>();
-		for (const identity of value) {
-			if (typeof identity !== "string" || identity.length === 0) {
+		accountPool.set(provider, parseAccountPoolEntry(provider, value));
+	}
+	return accountPool;
+}
+
+/**
+ * One pool-file entry: an identity array (visible OAuth accounts, API keys
+ * kept), `true` (everything visible), `false` (nothing visible), or a rule
+ * object `{ "accounts"?: [...], "apiKeys"?: boolean }`.
+ */
+function parseAccountPoolEntry(provider: string, value: unknown): ReadonlySet<string> | AuthBrokerProviderPool {
+	if (Array.isArray(value)) return parseAccountPoolIdentities(provider, value);
+	if (value === true) return {};
+	if (value === false) return { accounts: new Set(), apiKeys: false };
+	if (value === null || typeof value !== "object") {
+		throw new AIError.ConfigurationError(
+			`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} must be an array of identity keys, a boolean, or an object with "accounts"/"apiKeys"`,
+		);
+	}
+	const rule: { accounts?: ReadonlySet<string>; apiKeys?: boolean } = {};
+	for (const [key, field] of Object.entries(value)) {
+		if (key === "accounts") {
+			if (!Array.isArray(field)) {
 				throw new AIError.ConfigurationError(
-					`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an invalid identity key`,
+					`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} has "accounts" that is not an array of identity keys`,
 				);
 			}
-			if (identity !== identity.trim()) {
+			rule.accounts = parseAccountPoolIdentities(provider, field);
+		} else if (key === "apiKeys") {
+			if (typeof field !== "boolean") {
 				throw new AIError.ConfigurationError(
-					`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an identity key with surrounding whitespace`,
+					`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} has "apiKeys" that is not a boolean`,
 				);
 			}
-			identities.add(identity);
+			rule.apiKeys = field;
+		} else {
+			throw new AIError.ConfigurationError(
+				`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} has unknown field ${JSON.stringify(key)}`,
+			);
 		}
-		accountPool.set(provider, identities);
 	}
-	return accountPool;
+	return rule;
+}
+
+function parseAccountPoolIdentities(provider: string, value: unknown[]): Set<string> {
+	const identities = new Set<string>();
+	for (const identity of value) {
+		if (typeof identity !== "string" || identity.length === 0) {
+			throw new AIError.ConfigurationError(
+				`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an invalid identity key`,
+			);
+		}
+		if (identity !== identity.trim()) {
+			throw new AIError.ConfigurationError(
+				`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an identity key with surrounding whitespace`,
+			);
+		}
+		identities.add(identity);
+	}
+	return identities;
 }
 
 function resolveSnapshotTtlMs(): number {

package/src/auth-broker/remote-store.ts

@@ -41,21 +41,61 @@ import type {
 const REMOTE_CACHE_MAX_ENTRIES = 4096;
 const REMOTE_CACHE_MAX_CHARACTERS = 4 * 1024 * 1024;
 
+/** Visibility rule for one provider's broker credentials. */
+export interface AuthBrokerProviderPool {
+	/** OAuth `identityKey`s visible to this client; omitted = every OAuth account. */
+	readonly accounts?: ReadonlySet<string>;
+	/** Whether API-key credentials are visible. Default `true`. */
+	readonly apiKeys?: boolean;
+}
+
+/** Pool key whose rule applies to every provider without its own entry. */
+const AUTH_BROKER_ACCOUNT_POOL_DEFAULT_KEY = "*";
+
 /**
- * Per-provider OAuth identities visible to this trusted broker client.
- * Missing providers are unrestricted; an empty set excludes that provider's
- * OAuth credentials. API keys are never filtered.
+ * Per-provider credential visibility for this trusted broker client. A set
+ * value lists the provider's visible OAuth identities and keeps its API keys
+ * (an empty set excludes every OAuth credential); a rule object can also hide
+ * API keys. Providers without an entry follow the `"*"` rule, which defaults
+ * to unrestricted.
  */
-export type AuthBrokerAccountPool = ReadonlyMap<string, ReadonlySet<string>>;
+export type AuthBrokerAccountPool = ReadonlyMap<string, ReadonlySet<string> | AuthBrokerProviderPool>;
+
+interface ResolvedProviderPool {
+	accounts: ReadonlySet<string> | undefined;
+	apiKeys: boolean;
+}
+
+/** Copies `pool` so later caller mutation cannot widen it, normalizing both value forms. */
+function resolveAccountPool(pool: AuthBrokerAccountPool): ReadonlyMap<string, ResolvedProviderPool> {
+	const resolved = new Map<string, ResolvedProviderPool>();
+	for (const [provider, value] of pool) {
+		resolved.set(
+			provider,
+			"has" in value
+				? { accounts: new Set(value), apiKeys: true }
+				: { accounts: value.accounts && new Set(value.accounts), apiKeys: value.apiKeys ?? true },
+		);
+	}
+	return resolved;
+}
+
+function providerPoolRule(
+	accountPool: ReadonlyMap<string, ResolvedProviderPool> | undefined,
+	provider: string,
+): ResolvedProviderPool | undefined {
+	return accountPool?.get(provider) ?? accountPool?.get(AUTH_BROKER_ACCOUNT_POOL_DEFAULT_KEY);
+}
 
 function isCredentialInAccountPool(
 	entry: Pick<SnapshotEntry, "provider" | "credential" | "identityKey">,
-	accountPool: AuthBrokerAccountPool | undefined,
+	accountPool: ReadonlyMap<string, ResolvedProviderPool> | undefined,
 ): boolean {
-	if (entry.credential.type !== "oauth") return true;
-	const identities = accountPool?.get(entry.provider);
-	if (identities === undefined) return true;
-	return entry.identityKey !== null && identities.has(entry.identityKey);
+	const rule = providerPoolRule(accountPool, entry.provider);
+	if (!rule) return true;
+	if (entry.credential.type !== "oauth") return rule.apiKeys;
+	if (!rule.accounts) return true;
+	return entry.identityKey !== null && rule.accounts.has(entry.identityKey);
 }
 
 /**
@@ -245,8 +285,9 @@ export interface RemoteAuthCredentialStoreOptions {
 	 */
 	onSnapshot?: (snapshot: SnapshotResponse, generation: number) => void;
 	/**
-	 * OAuth identities visible through this store. This is a trusted-client
-	 * routing policy, not broker authorization.
+	 * Broker credentials visible through this store (OAuth identities and API
+	 * keys per provider). This is a trusted-client routing policy, not broker
+	 * authorization.
 	 */
 	accountPool?: AuthBrokerAccountPool;
 	/** Flush cadence for batched observed-usage reports. Default 10s. */
@@ -265,7 +306,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 	readonly #client: AuthBrokerClient;
 	readonly #streamSnapshots: boolean;
 	readonly #onSnapshot?: (snapshot: SnapshotResponse, generation: number) => void;
-	readonly #accountPool?: AuthBrokerAccountPool;
+	readonly #accountPool?: ReadonlyMap<string, ResolvedProviderPool>;
 	#snapshot: SnapshotResponse = emptySnapshot();
 	#snapshotReceivedAt = Date.now();
 	#generation = 0;
@@ -336,9 +377,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		this.#streamSnapshots = opts.streamSnapshots ?? true;
 		this.#observedUsageFlushMs = opts.observedUsageFlushMs ?? 10_000;
 		this.#backgroundIdleMs = opts.backgroundIdleMs ?? BACKGROUND_IDLE_MS;
-		this.#accountPool = opts.accountPool
-			? new Map([...opts.accountPool].map(([provider, identities]) => [provider, new Set(identities)]))
-			: undefined;
+		this.#accountPool = opts.accountPool ? resolveAccountPool(opts.accountPool) : undefined;
 		this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0);
 		this.#acknowledgedRevision = this.#credentialRevision;
 		this.#onSnapshot = opts.onSnapshot;
@@ -1372,7 +1411,11 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		if (memo && memo.input === reports && memo.byProvider === lookup.byProvider) return memo.output;
 		const byProvider = lookup.byProvider;
 		const output = reports.filter(report => {
-			if (!accountPool.has(report.provider)) return true;
+			// Unrestricted for this provider: every report stays. Otherwise a
+			// report must belong to a visible OAuth credential, so reports from
+			// hidden accounts or API keys fail closed.
+			const rule = providerPoolRule(accountPool, report.provider);
+			if (!rule || (rule.accounts === undefined && rule.apiKeys)) return true;
 			const credentials = byProvider.get(report.provider);
 			if (!credentials) return false;
 			return credentials.some(credential => usageReportMatchesCredential(report, credential));

package/src/error/flags.ts

@@ -167,7 +167,7 @@ function matchesPayloadRejectionText(text: string): boolean {
 
 const TIMEOUT_PATTERN = /\b(?:operation\s+)?timed?\s*out\b|\btimeout\b|\bstream stall\b/i;
 const TRANSIENT_ENVELOPE_PATTERN = /anthropic stream envelope error:/i;
-const TRANSIENT_ENVELOPE_TRUNCATION_PATTERN = /before message_(?:start|stop)/i;
+const TRANSIENT_ENVELOPE_TRUNCATION_PATTERN = /before message_(?:start|delta|stop)/i;
 export const STREAM_READ_ERROR_PATTERN = /stream[_ -]?read[_ -]?error/i;
 /** Python h2/httpx diagnostics forwarded through provider or proxy error events. */
 export const PYTHON_HTTP2_STREAM_RESET_PATTERN = /<StreamReset stream_id:\d+, error_code:(?:2|7), remote_reset:True>/;
@@ -1103,7 +1103,7 @@ export function isTransientStreamParseError(error: unknown): boolean {
  * non-terminal status treat it as transient (#11805).
  */
 const STREAM_DROP_PATTERN =
-	/stream disconnected before completion|(?:stream|websocket) closed before response\.completed|stream was interrupted|stream ended before terminal (?:chunk|completion event)|socket disconnected before secure tls connection/i;
+	/stream disconnected before completion|(?:stream|websocket) closed before response\.completed|stream was interrupted|stream ended before (?:terminal (?:chunk|completion event)|message_(?:start|delta|stop))|socket disconnected before secure tls connection/i;
 
 /**
  * Transport stream-drop diagnostic (see {@link STREAM_DROP_PATTERN}). Unlike

package/src/images/google-antigravity.ts

@@ -99,7 +99,6 @@ function buildRequest(request: ImageGenerationRequest, model: string, projectId:
 				{ category: "HARM_CATEGORY_CIVIC_INTEGRITY", threshold: "BLOCK_ONLY_HIGH" },
 			],
 		},
-		requestType: "agent",
 		requestId: `agent-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`,
 		userAgent: "antigravity",
 	};