@f5-sales-demo/pi-ai 20.4.2 → 20.4.3

package/package.json

@@ -1,7 +1,7 @@
 {
 	"type": "module",
 	"name": "@f5-sales-demo/pi-ai",
-	"version": "20.4.2",
+	"version": "20.4.3",
 	"description": "Unified LLM API with automatic model discovery and provider configuration",
 	"homepage": "https://github.com/f5-sales-demo/xcsh",
 	"author": "Can Boluk",
@@ -46,7 +46,7 @@
 		"@anthropic-ai/sdk": "^0.115",
 		"@aws-sdk/client-bedrock-runtime": "^3",
 		"@bufbuild/protobuf": "^2.11",
-		"@f5-sales-demo/pi-utils": "20.4.2",
+		"@f5-sales-demo/pi-utils": "20.4.3",
 		"@google/genai": "^2.15",
 		"@sinclair/typebox": "^0.34",
 		"@smithy/node-http-handler": "^4.4",

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

@@ -1,7 +1,7 @@
 /**
  * Google Gemini CLI / Antigravity provider.
  * Shared implementation for both google-gemini-cli and google-antigravity providers.
- * Uses the Cloud Code Assist API endpoint to access Gemini and Claude models.
+ * Uses Cloud Code Assist plus explicit enterprise Vertex routes for supported Antigravity models.
  */
 import { createHash, randomBytes, randomUUID } from "node:crypto";
 import { abortableSleep, readSseJson } from "@f5-sales-demo/pi-utils";
@@ -62,6 +62,9 @@ const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com";
 const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com";
 const ANTIGRAVITY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";
 const ANTIGRAVITY_ENDPOINT_FALLBACKS = [ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_SANDBOX_ENDPOINT] as const;
+const ANTIGRAVITY_VERTEX_MODELS: Readonly<Record<string, string>> = Object.freeze({
+	"gemini-3.1-pro-high-vertex": "gemini-3.1-pro-preview",
+});
 
 /**
  * Build a User-Agent string that identifies as Gemini CLI to unlock higher rate limits.
@@ -421,34 +424,36 @@ interface CloudCodeAssistRequest {
 	requestId?: string;
 }
 
-interface CloudCodeAssistResponseChunk {
-	response?: {
-		candidates?: Array<{
-			content?: {
-				role: string;
-				parts?: Array<{
-					text?: string;
-					thought?: boolean;
-					thoughtSignature?: string;
-					functionCall?: {
-						name: string;
-						args: Record<string, unknown>;
-						id?: string;
-					};
-				}>;
-			};
-			finishReason?: string;
-		}>;
-		usageMetadata?: {
-			promptTokenCount?: number;
-			candidatesTokenCount?: number;
-			thoughtsTokenCount?: number;
-			totalTokenCount?: number;
-			cachedContentTokenCount?: number;
+interface GeminiGenerateContentResponse {
+	candidates?: Array<{
+		content?: {
+			role: string;
+			parts?: Array<{
+				text?: string;
+				thought?: boolean;
+				thoughtSignature?: string;
+				functionCall?: {
+					name: string;
+					args: Record<string, unknown>;
+					id?: string;
+				};
+			}>;
 		};
-		modelVersion?: string;
-		responseId?: string;
+		finishReason?: string;
+	}>;
+	usageMetadata?: {
+		promptTokenCount?: number;
+		candidatesTokenCount?: number;
+		thoughtsTokenCount?: number;
+		totalTokenCount?: number;
+		cachedContentTokenCount?: number;
 	};
+	modelVersion?: string;
+	responseId?: string;
+}
+
+interface CloudCodeAssistResponseChunk extends GeminiGenerateContentResponse {
+	response?: GeminiGenerateContentResponse;
 	traceId?: string;
 }
 
@@ -492,11 +497,16 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 			const parsedCredentials = parseGeminiCliCredentials(apiKeyRaw);
 			const activeCredentials = await refreshGeminiCliCredentialsIfNeeded(parsedCredentials, isAntigravity);
 			const { accessToken, projectId } = activeCredentials;
+			const vertexModelId = isAntigravity ? ANTIGRAVITY_VERTEX_MODELS[model.id] : undefined;
+			const serviceName = vertexModelId ? "Vertex AI" : "Cloud Code Assist";
 
 			const baseUrl = model.baseUrl?.trim();
 			const endpoints = baseUrl ? [baseUrl] : isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT];
 
-			let requestBody = buildRequest(model, context, projectId, options, isAntigravity);
+			const cloudCodeRequest = buildRequest(model, context, projectId, options, isAntigravity);
+			let requestBody: CloudCodeAssistRequest | CloudCodeAssistRequest["request"] = vertexModelId
+				? cloudCodeRequest.request
+				: cloudCodeRequest;
 			const replacementPayload = await options?.onPayload?.(requestBody, model);
 			if (replacementPayload !== undefined) {
 				requestBody = replacementPayload as typeof requestBody;
@@ -534,7 +544,9 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 
 				try {
 					const endpoint = endpoints[Math.min(attempt, endpoints.length - 1)];
-					requestUrl = `${endpoint}/v1internal:streamGenerateContent?alt=sse`;
+					requestUrl = vertexModelId
+						? `${endpoint}/v1/projects/${encodeURIComponent(projectId)}/locations/global/publishers/google/models/${encodeURIComponent(vertexModelId)}:streamGenerateContent?alt=sse`
+						: `${endpoint}/v1internal:streamGenerateContent?alt=sse`;
 					response = await fetch(requestUrl, {
 						method: "POST",
 						headers: requestHeaders,
@@ -552,7 +564,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					if (response.status === 429) {
 						if (/quota|exhausted/i.test(errorText)) {
 							throw withHttpStatus(
-								new Error(`Cloud Code Assist API error (429): ${extractErrorMessage(errorText)}`),
+								new Error(`${serviceName} API error (429): ${extractErrorMessage(errorText)}`),
 								429,
 							);
 						}
@@ -590,7 +602,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 
 					// Not retryable or budget exceeded
 					throw withHttpStatus(
-						new Error(`Cloud Code Assist API error (${response.status}): ${extractErrorMessage(errorText)}`),
+						new Error(`${serviceName} API error (${response.status}): ${extractErrorMessage(errorText)}`),
 						response.status,
 					);
 				} catch (error) {
@@ -665,8 +677,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					activeResponse.body!,
 					options?.signal,
 				)) {
-					const responseData = chunk.response;
-					if (!responseData) continue;
+					const responseData = chunk.response ?? chunk;
 
 					const candidate = responseData.candidates?.[0];
 					if (candidate?.content?.parts) {
@@ -878,7 +889,7 @@ 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}`),
+							new Error(`${serviceName} API error (${currentResponse.status}): ${retryErrorText}`),
 							currentResponse.status,
 						);
 					}
@@ -896,7 +907,7 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 			}
 
 			if (!receivedContent) {
-				throw new Error("Cloud Code Assist API returned an empty response");
+				throw new Error(`${serviceName} API returned an empty response`);
 			}
 
 			if (options?.signal?.aborted) {

package/src/utils/discovery/antigravity.ts

@@ -8,6 +8,13 @@ const DEFAULT_ANTIGRAVITY_DISCOVERY_ENDPOINTS = [
 	"https://daily-cloudcode-pa.sandbox.googleapis.com",
 ] as const;
 const FETCH_AVAILABLE_MODELS_PATH = "/v1internal:fetchAvailableModels";
+const ANTIGRAVITY_VERTEX_ENDPOINT = "https://aiplatform.googleapis.com";
+const ANTIGRAVITY_VERTEX_VARIANTS: Readonly<Record<string, { id: string; name: string }>> = Object.freeze({
+	"gemini-3.1-pro-high": {
+		id: "gemini-3.1-pro-high-vertex",
+		name: "Gemini 3.1 Pro High (Vertex, Antigravity)",
+	},
+});
 
 const DEFAULT_CONTEXT_WINDOW = 200_000;
 const DEFAULT_MAX_TOKENS = 64_000;
@@ -222,12 +229,13 @@ export async function fetchAntigravityDiscoveryModels(
 			}
 
 			const supportsImages = model.supportsImages === true;
+			const vertexVariant = ANTIGRAVITY_VERTEX_VARIANTS[modelId];
 			models.push({
-				id: modelId,
-				name: model.displayName ? `${model.displayName} (Antigravity)` : modelId,
+				id: vertexVariant?.id ?? modelId,
+				name: vertexVariant?.name ?? (model.displayName ? `${model.displayName} (Antigravity)` : modelId),
 				api: "google-gemini-cli",
 				provider: "google-antigravity",
-				baseUrl: endpoint,
+				baseUrl: vertexVariant ? ANTIGRAVITY_VERTEX_ENDPOINT : endpoint,
 				reasoning: model.supportsThinking === true,
 				input: supportsImages ? ["text", "image"] : ["text"],
 				cost: {

package/src/utils/oauth/google-antigravity.ts

@@ -30,6 +30,7 @@ const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth";
 const TOKEN_URL = "https://oauth2.googleapis.com/token";
 const CLOUD_CODE_ENDPOINT = "https://cloudcode-pa.googleapis.com";
 const TIER_LEGACY = "legacy-tier";
+const TIER_STANDARD = "standard-tier";
 const PROJECT_ONBOARD_MAX_ATTEMPTS = 5;
 const PROJECT_ONBOARD_INTERVAL_MS = 2000;
 const PROJECT_ID_PATTERN = /^[a-z][a-z0-9-]{4,28}[a-z0-9]$/;
@@ -248,6 +249,25 @@ async function onboardProjectWithRetries(
 	);
 }
 
+async function loadCodeAssist(
+	endpoint: string,
+	headers: Record<string, string>,
+	projectRequest: AntigravityProjectRequest,
+): Promise<LoadCodeAssistPayload> {
+	const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
+		method: "POST",
+		headers,
+		body: JSON.stringify(projectRequest),
+	});
+
+	if (!response.ok) {
+		const errorText = await response.text();
+		throw new Error(`loadCodeAssist failed: ${response.status} ${response.statusText}: ${errorText}`);
+	}
+
+	return (await response.json()) as LoadCodeAssistPayload;
+}
+
 async function discoverProject(
 	accessToken: string,
 	configuredProjectId: string | undefined,
@@ -263,31 +283,42 @@ async function discoverProject(
 	onProgress?.("Checking for existing project...");
 	const endpoint = CLOUD_CODE_ENDPOINT;
 	try {
-		const loadResponse = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
-			method: "POST",
-			headers,
-			body: JSON.stringify(projectRequest),
-		});
+		const loadPayload = await loadCodeAssist(endpoint, headers, projectRequest);
+		const existingProject = readProjectId(loadPayload.cloudaicompanionProject);
 
-		if (!loadResponse.ok) {
-			const errorText = await loadResponse.text();
-			throw new Error(`loadCodeAssist failed: ${loadResponse.status} ${loadResponse.statusText}: ${errorText}`);
+		if (!configuredProjectId) {
+			if (existingProject) {
+				return existingProject;
+			}
+
+			const tierId = getDefaultTierId(loadPayload.allowedTiers);
+			onProgress?.("Provisioning project...");
+			const onboardBody = {
+				tierId,
+				...projectRequest,
+			};
+			return await onboardProjectWithRetries(endpoint, headers, onboardBody, onProgress);
 		}
 
-		const loadPayload = (await loadResponse.json()) as LoadCodeAssistPayload;
-		const existingProject = readProjectId(loadPayload.cloudaicompanionProject);
-		if (existingProject) {
-			return configuredProjectId ?? existingProject;
+		if (existingProject === configuredProjectId && loadPayload.currentTier?.id === TIER_STANDARD) {
+			return configuredProjectId;
+		}
+
+		const standardTierAvailable = loadPayload.allowedTiers?.some(tier => tier.id === TIER_STANDARD) ?? false;
+		if (!standardTierAvailable) {
+			throw new Error(`standard-tier is not available for configured project ${configuredProjectId}`);
 		}
 
-		const tierId = getDefaultTierId(loadPayload.allowedTiers);
 		onProgress?.("Provisioning project...");
 		const onboardBody = {
-			tierId,
+			tierId: TIER_STANDARD,
 			...projectRequest,
 		};
 		const provisionedProject = await onboardProjectWithRetries(endpoint, headers, onboardBody, onProgress);
-		return configuredProjectId ?? provisionedProject;
+		if (provisionedProject !== configuredProjectId) {
+			throw new Error(`onboardUser returned project ${provisionedProject} instead of ${configuredProjectId}`);
+		}
+		return configuredProjectId;
 	} catch (error) {
 		throw new Error(
 			`Could not discover or provision an Antigravity project. ${error instanceof Error ? error.message : String(error)}`,