@prometheus-ai/ai 0.5.4 → 0.5.8

package/src/providers/google-auth.ts

@@ -17,6 +17,7 @@ import * as os from "node:os";
 import * as path from "node:path";
 import { $envpos, isEnoent, logger } from "@prometheus-ai/utils";
 import type { FetchImpl } from "../types";
+import { raceWithSignal } from "../utils/abort";
 
 const OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token";
 const METADATA_TOKEN_URL = "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token";
@@ -42,7 +43,14 @@ interface AuthorizedUserCredentials {
 	refresh_token: string;
 }
 
-type AdcFileCredentials = ServiceAccountCredentials | AuthorizedUserCredentials;
+interface ImpersonatedServiceAccountCredentials {
+	type: "impersonated_service_account";
+	service_account_impersonation_url: string;
+	source_credentials: AuthorizedUserCredentials | ServiceAccountCredentials;
+	delegates?: string[];
+}
+
+type AdcFileCredentials = ServiceAccountCredentials | AuthorizedUserCredentials | ImpersonatedServiceAccountCredentials;
 
 interface TokenResponse {
 	access_token: string;
@@ -196,10 +204,52 @@ async function resolveAccessTokenUncached(
 ): Promise<{ source: string; token: TokenResponse }> {
 	const adc = await loadAdcCredentials();
 	if (adc) {
-		const token =
-			adc.creds.type === "service_account"
-				? await exchangeJwtForToken(adc.creds, signal, fetchImpl)
-				: await exchangeRefreshToken(adc.creds, signal, fetchImpl);
+		const creds = adc.creds;
+		let token: TokenResponse;
+
+		if (creds.type === "impersonated_service_account") {
+			const targetPrincipalMatch = /(?<target>[^/]+):(generateAccessToken|generateIdToken)$/.exec(
+				creds.service_account_impersonation_url,
+			);
+			const targetPrincipal = targetPrincipalMatch?.groups?.target;
+			if (!targetPrincipal) {
+				throw new RangeError(`Cannot extract target principal from ${creds.service_account_impersonation_url}`);
+			}
+
+			const sourceToken =
+				creds.source_credentials.type === "service_account"
+					? await exchangeJwtForToken(creds.source_credentials, signal, fetchImpl)
+					: await exchangeRefreshToken(creds.source_credentials, signal, fetchImpl);
+
+			const response = await fetchImpl(
+				`https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/${targetPrincipal}:generateAccessToken`,
+				{
+					method: "POST",
+					headers: {
+						"Content-Type": "application/json",
+						Authorization: `Bearer ${sourceToken.access_token}`,
+					},
+					body: JSON.stringify({
+						delegates: creds.delegates ?? [],
+						scope: [CLOUD_PLATFORM_SCOPE],
+						lifetime: "3600s",
+					}),
+					signal,
+				},
+			);
+			if (!response.ok) {
+				const detail = await response.text().catch(() => "");
+				throw new Error(`Google Impersonation token exchange failed (${response.status}): ${detail}`);
+			}
+			const data = (await response.json()) as { accessToken: string; expireTime: string };
+			const expiresIn = Math.max(0, Math.floor((new Date(data.expireTime).getTime() - Date.now()) / 1000));
+			token = { access_token: data.accessToken, expires_in: expiresIn, token_type: "Bearer" };
+		} else {
+			token =
+				creds.type === "service_account"
+					? await exchangeJwtForToken(creds, signal, fetchImpl)
+					: await exchangeRefreshToken(creds, signal, fetchImpl);
+		}
 		return { source: adc.source, token };
 	}
 	const metadata = await fetchMetadataToken(signal, fetchImpl);
@@ -209,6 +259,13 @@ async function resolveAccessTokenUncached(
 	);
 }
 
+/**
+ * Bound for the detached (signal-free) shared token resolution: a hung OAuth
+ * exchange or metadata fetch must not pin the inflight slot forever — every
+ * later call would await the stuck promise until process restart.
+ */
+const SHARED_TOKEN_RESOLVE_TIMEOUT_MS = 30_000;
+
 /**
  * Returns a Bearer access token suitable for the `Authorization` header on Vertex AI calls.
  * The token is cached in module scope and refreshed `GOOGLE_VERTEX_REFRESH_SKEW_MS` ms before it expires.
@@ -228,11 +285,17 @@ export async function getVertexAccessToken(options?: { signal?: AbortSignal; fet
 
 	const cacheKey = "vertex-adc";
 	const existing = inflight.get(cacheKey);
-	if (existing) return existing;
+	if (existing) return raceWithSignal(existing, options?.signal);
 
+	// Deliberately resolve without any caller's signal: the in-flight promise is shared
+	// by every concurrent caller, so aborting one request must not fail the whole batch.
+	// Each caller races its own signal against the shared promise instead.
 	const promise = (async () => {
 		try {
-			const { source, token } = await resolveAccessTokenUncached(options?.signal, fetchImpl);
+			const { source, token } = await resolveAccessTokenUncached(
+				AbortSignal.timeout(SHARED_TOKEN_RESOLVE_TIMEOUT_MS),
+				fetchImpl,
+			);
 			const expiresAtMs = Date.now() + Math.max(0, token.expires_in * 1000);
 			tokenCache.set(source, { token: token.access_token, expiresAtMs });
 			logger.debug("vertex.adc acquired access token", { source, expiresInSec: token.expires_in });
@@ -242,7 +305,7 @@ export async function getVertexAccessToken(options?: { signal?: AbortSignal; fet
 		}
 	})();
 	inflight.set(cacheKey, promise);
-	return promise;
+	return raceWithSignal(promise, options?.signal);
 }
 
 /** Test seam: clears every cached token. */

package/src/providers/google-gemini-cli.ts

@@ -5,8 +5,15 @@
  */
 import { createHash, randomBytes, randomUUID } from "node:crypto";
 import { scheduler } from "node:timers/promises";
+import { calculateCost } from "@prometheus-ai/catalog/models";
+import {
+	ANTIGRAVITY_NO_PREAMBLE_INSTRUCTION,
+	ANTIGRAVITY_SYSTEM_INSTRUCTION,
+	getAntigravityUserAgent,
+	getGeminiCliHeaders,
+} from "@prometheus-ai/catalog/wire/gemini-headers";
 import { extractHttpStatusFromError, fetchWithRetry, readSseJson } from "@prometheus-ai/utils";
-import { calculateCost } from "../models";
+import { ProviderHttpError } from "../errors";
 import type {
 	Api,
 	AssistantMessage,
@@ -20,17 +27,20 @@ import type {
 } from "../types";
 import { normalizeSystemPrompts } from "../utils";
 import { AssistantMessageEventStream } from "../utils/event-stream";
-import { appendRawHttpRequestDumpFor400, type RawHttpRequestDump, withHttpStatus } from "../utils/http-inspector";
+import { appendRawHttpRequestDumpFor400, type RawHttpRequestDump } from "../utils/http-inspector";
+import { getStreamFirstEventTimeoutMs } from "../utils/idle-iterator";
 // Refresh is the sole responsibility of AuthStorage (broker-aware, single-flighted);
 // the stream provider trusts the access token threaded through `options.apiKey`.
 import { normalizeSchemaForCCA } from "../utils/schema";
-import { ANTIGRAVITY_SYSTEM_INSTRUCTION, getAntigravityUserAgent, getGeminiCliHeaders } from "./google-gemini-headers";
 import type { Content, FunctionCallingConfigMode, ThinkingConfig } from "./google-shared";
 import {
 	convertMessages,
 	convertTools,
+	EMPTY_STREAM_BASE_DELAY_MS,
 	type GoogleThinkingLevel,
+	hasMeaningfulGoogleContent,
 	isThinkingPart,
+	MAX_EMPTY_STREAM_RETRIES,
 	mapStopReasonString,
 	mapToolChoice,
 	nextToolCallId,
@@ -46,6 +56,11 @@ import {
  */
 export type { GoogleThinkingLevel };
 
+/** Non-2xx response (or in-stream error chunk) from the Cloud Code Assist API. */
+export class GeminiCliApiError extends ProviderHttpError {
+	override readonly name = "GeminiCliApiError";
+}
+
 export interface GoogleGeminiCliOptions extends StreamOptions {
 	/**
 	 * Tool selection mode. String forms map directly to Gemini
@@ -66,7 +81,19 @@ export interface GoogleGeminiCliOptions extends StreamOptions {
 		budgetTokens?: number;
 		/** Thinking level. Use for Gemini 3 models (LOW/HIGH for Pro, MINIMAL/LOW/MEDIUM/HIGH for Flash). */
 		level?: GoogleThinkingLevel;
+		/**
+		 * Explicit wire suppression when `enabled` is false. Cloud Code Assist
+		 * re-applies the per-id baked server default when thinkingConfig is
+		 * omitted, so models with `thinking.suppressWhenOff` must send
+		 * `includeThoughts: false` plus a MINIMAL level (or zero budget).
+		 */
+		suppress?: { level: GoogleThinkingLevel } | { budget: number };
 	};
+	/**
+	 * Upstream wire model id override for collapsed effort-tier variants.
+	 * Serialized as `requestModelId ?? model.requestModelId ?? model.id`.
+	 */
+	requestModelId?: string;
 	projectId?: string;
 }
 
@@ -76,17 +103,16 @@ const ANTIGRAVITY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googlea
 const ANTIGRAVITY_ENDPOINT_FALLBACKS = [ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_SANDBOX_ENDPOINT] as const;
 
 export {
+	ANTIGRAVITY_NO_PREAMBLE_INSTRUCTION,
 	ANTIGRAVITY_SYSTEM_INSTRUCTION,
 	getAntigravityUserAgent,
 	getGeminiCliHeaders,
 	getGeminiCliUserAgent,
-} from "./google-gemini-headers";
+} from "@prometheus-ai/catalog/wire/gemini-headers";
 
 // Retry configuration
 const MAX_RETRIES = 3;
 const BASE_DELAY_MS = 1000;
-const MAX_EMPTY_STREAM_RETRIES = 2;
-const EMPTY_STREAM_BASE_DELAY_MS = 500;
 const RATE_LIMIT_BUDGET_MS = 5 * 60 * 1000;
 const CLAUDE_THINKING_BETA_HEADER = "interleaved-thinking-2025-05-14";
 const GOOGLE_GEMINI_REFRESH_SKEW_MS = 60_000;
@@ -253,7 +279,10 @@ interface CloudCodeAssistResponseChunk {
 		};
 		modelVersion?: string;
 		responseId?: string;
+		promptFeedback?: { blockReason?: string; blockReasonMessage?: string };
 	};
+	/** In-band stream failure (quota, internal error) delivered as a final JSON event. */
+	error?: { code?: number; message?: string; status?: string };
 	traceId?: string;
 }
 
@@ -339,29 +368,48 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 				headers: requestHeaders,
 			};
 
+			// Direct callers that skip `register-builtins` (which installs the
+			// iterator-level watchdog) need a pre-response timer alongside
+			// `timeout: false`; otherwise a stalled Cloud Code Assist proxy
+			// would hang forever. Floor matches the lazy wrapper's 5min default.
+			const firstEventTimeoutMs =
+				options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(undefined, 300_000);
+			const preResponseWatchdog =
+				firstEventTimeoutMs !== undefined && firstEventTimeoutMs > 0
+					? AbortSignal.timeout(firstEventTimeoutMs)
+					: undefined;
+			const callerSignal = options?.signal;
+			const fetchSignal = preResponseWatchdog
+				? callerSignal
+					? AbortSignal.any([callerSignal, preResponseWatchdog])
+					: preResponseWatchdog
+				: callerSignal;
 			const response = await fetchWithRetry(
 				attempt => `${endpoints[Math.min(attempt, endpoints.length - 1)]}/v1internal:streamGenerateContent?alt=sse`,
 				{
 					method: "POST",
 					headers: requestHeaders,
 					body: requestBodyJson,
-					signal: options?.signal,
+					signal: fetchSignal,
 					maxAttempts: MAX_RETRIES + 1,
 					defaultDelayMs: attempt => BASE_DELAY_MS * 2 ** attempt,
 					maxDelayMs: options?.maxRetryDelayMs ?? RATE_LIMIT_BUDGET_MS,
 					fetch: options?.fetch,
+					timeout: false,
 				},
 			);
 			if (!response.ok) {
 				const errorText = await response.text();
-				throw withHttpStatus(
-					new Error(`Cloud Code Assist API error (${response.status}): ${extractErrorMessage(errorText)}`),
+				throw new GeminiCliApiError(
+					`Cloud Code Assist API error (${response.status}): ${extractErrorMessage(errorText)}`,
 					response.status,
+					{ headers: response.headers },
 				);
 			}
 			const requestUrl = response.url;
 
 			let started = false;
+			let sawFinishReason = false;
 			const ensureStarted = () => {
 				if (!started) {
 					if (!firstTokenTime) firstTokenTime = Date.now();
@@ -383,7 +431,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 				output.stopReason = "stop";
 				output.errorMessage = undefined;
 				output.timestamp = Date.now();
-				started = false;
+				sawFinishReason = false;
 			};
 
 			const streamResponse = async (activeResponse: Response): Promise<boolean> => {
@@ -391,7 +439,6 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					throw new Error("No response body");
 				}
 
-				let hasContent = false;
 				let currentBlock: TextContent | ThinkingContent | null = null;
 				const blocks = output.content;
 				const blockIndex = () => blocks.length - 1;
@@ -401,14 +448,26 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					options?.signal,
 					event => options?.onSseEvent?.({ event: event.event, data: event.data, raw: [...event.raw] }, model),
 				)) {
+					if (chunk.error) {
+						const detail = chunk.error.message || chunk.error.status || "unknown error";
+						const message = `Cloud Code Assist stream error: ${detail}`;
+						throw typeof chunk.error.code === "number" && chunk.error.code >= 400
+							? new GeminiCliApiError(message, chunk.error.code)
+							: new Error(message);
+					}
 					const responseData = chunk.response;
 					if (!responseData) continue;
+					if (!responseData.candidates?.length && responseData.promptFeedback?.blockReason) {
+						const detail = responseData.promptFeedback.blockReasonMessage;
+						throw new Error(
+							`Request blocked by Google (${responseData.promptFeedback.blockReason})${detail ? `: ${detail}` : ""}`,
+						);
+					}
 
 					const candidate = responseData.candidates?.[0];
 					if (candidate?.content?.parts) {
 						for (const part of candidate.content.parts) {
-							if (part.text !== undefined) {
-								hasContent = true;
+							if (part.text !== undefined && part.text !== "") {
 								const isThinking = isThinkingPart(part);
 								if (
 									!currentBlock ||
@@ -445,10 +504,21 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 										partial: output,
 									});
 								}
+							} else if (part.text === "" && part.thoughtSignature && currentBlock && !part.functionCall) {
+								if (currentBlock.type === "thinking") {
+									currentBlock.thinkingSignature = retainThoughtSignature(
+										currentBlock.thinkingSignature,
+										part.thoughtSignature,
+									);
+								} else {
+									currentBlock.textSignature = retainThoughtSignature(
+										currentBlock.textSignature,
+										part.thoughtSignature,
+									);
+								}
 							}
 
 							if (part.functionCall) {
-								hasContent = true;
 								if (currentBlock) {
 									pushBlockEndEvent(currentBlock, blockIndex(), output, stream);
 									currentBlock = null;
@@ -463,7 +533,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 									type: "toolCall",
 									id: toolCallId,
 									name: part.functionCall.name || "",
-									arguments: part.functionCall.args as Record<string, unknown>,
+									arguments: (part.functionCall.args ?? {}) as Record<string, unknown>,
 									...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
 								};
 
@@ -475,9 +545,17 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					}
 
 					if (candidate?.finishReason) {
-						output.stopReason = mapStopReasonString(candidate.finishReason);
-						if (output.content.some(b => b.type === "toolCall")) {
+						sawFinishReason = true;
+						const mapped = mapStopReasonString(candidate.finishReason);
+						// Only let a trailing tool call upgrade benign finishes; error finishes
+						// (SAFETY, MALFORMED_FUNCTION_CALL, ...) must surface even with tool calls present.
+						if ((mapped === "stop" || mapped === "length") && output.content.some(b => b.type === "toolCall")) {
 							output.stopReason = "toolUse";
+						} else {
+							output.stopReason = mapped;
+							if (mapped === "error") {
+								output.errorMessage = `Generation failed with finish reason: ${candidate.finishReason}`;
+							}
 						}
 					}
 
@@ -509,7 +587,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					pushBlockEndEvent(currentBlock, blockIndex(), output, stream);
 				}
 
-				return hasContent;
+				return hasMeaningfulGoogleContent(output);
 			};
 
 			let receivedContent = false;
@@ -542,9 +620,10 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 
 					if (!currentResponse.ok) {
 						const retryErrorText = await currentResponse.text();
-						throw withHttpStatus(
-							new Error(`Cloud Code Assist API error (${currentResponse.status}): ${retryErrorText}`),
+						throw new GeminiCliApiError(
+							`Cloud Code Assist API error (${currentResponse.status}): ${retryErrorText}`,
 							currentResponse.status,
+							{ headers: currentResponse.headers },
 						);
 					}
 				}
@@ -568,6 +647,12 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 				throw new Error("Request was aborted");
 			}
 
+			if (!sawFinishReason) {
+				throw new Error(
+					"Cloud Code Assist stream ended without a finish reason (connection dropped or response truncated)",
+				);
+			}
+
 			if (output.stopReason === "aborted" || output.stopReason === "error") {
 				throw new Error(output.errorMessage ?? "An unknown error occurred");
 			}
@@ -726,6 +811,17 @@ export function buildRequest(
 		} else if (options.thinking.budgetTokens !== undefined) {
 			generationConfig.thinkingConfig.thinkingBudget = options.thinking.budgetTokens;
 		}
+	} else if (options.thinking?.suppress && model.reasoning) {
+		// Explicit off: omitting thinkingConfig re-applies the per-id baked
+		// server default (the model silently thinks and bills the tokens).
+		const suppress = options.thinking.suppress;
+		generationConfig.thinkingConfig = { includeThoughts: false };
+		if ("level" in suppress) {
+			// Cast to any since our GoogleThinkingLevel mirrors Google's ThinkingLevel enum values
+			generationConfig.thinkingConfig.thinkingLevel = suppress.level as any;
+		} else {
+			generationConfig.thinkingConfig.thinkingBudget = suppress.budget;
+		}
 	}
 
 	const request: CloudCodeAssistRequest["request"] = {
@@ -785,10 +881,10 @@ export function buildRequest(
 	if (isAntigravity && shouldInjectAntigravitySystemInstruction(model.id)) {
 		const existingParts = request.systemInstruction?.parts ?? [];
 		request.systemInstruction = {
-			role: "user",
 			parts: [
 				{ text: ANTIGRAVITY_SYSTEM_INSTRUCTION },
 				{ text: `Please ignore following [ignore]${ANTIGRAVITY_SYSTEM_INSTRUCTION}[/ignore]` },
+				{ text: ANTIGRAVITY_NO_PREAMBLE_INSTRUCTION },
 				...existingParts,
 			],
 		};
@@ -796,7 +892,7 @@ export function buildRequest(
 
 	return {
 		project: projectId,
-		model: model.id,
+		model: options.requestModelId ?? model.requestModelId ?? model.id,
 		request,
 		...(isAntigravity
 			? {