diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4cd98e03..0117f86b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,6 +59,7 @@ jobs: --no-extensions \ --extension "$package/extensions/file-search/index.ts" \ --extension "$package/extensions/git-info/index.ts" \ + --extension "$package/extensions/ai-providers/index.ts" \ = { + chat_20706: true, + chat_23310: true, + "gemini-2.5-pro": true, +}; + +const DEFAULT_CONTEXT_WINDOW = 200_000; +const DEFAULT_MAX_TOKENS = 64_000; + +const DiscoveryModelSchema = Type.Object({ + displayName: Type.Optional(Type.String()), + supportsImages: Type.Optional(Type.Boolean()), + supportsThinking: Type.Optional(Type.Boolean()), + maxTokens: Type.Optional(Type.Number()), + maxOutputTokens: Type.Optional(Type.Number()), + isInternal: Type.Optional(Type.Boolean()), +}); + +const DiscoveryResponseSchema = Type.Object({ + models: Type.Optional(Type.Record(Type.String(), DiscoveryModelSchema)), +}); + +type DiscoveryResponse = Static; + +type DiscoveredModel = AntigravityProviderModel & AntigravityModelDefinition; + +function positiveNumber(value: number | undefined, fallback: number): number { + return value !== undefined && Number.isFinite(value) && value > 0 + ? value + : fallback; +} + +function toModelDefinition( + id: string, + meta: Static, +): DiscoveredModel { + const advertisedMaxTokens = positiveNumber( + meta.maxOutputTokens, + DEFAULT_MAX_TOKENS, + ); + return { + id, + name: meta.displayName ?? id, + api: "antigravity-cloudcode", + provider: "google-antigravity", + baseUrl: ANTIGRAVITY_API_URL, + reasoning: meta.supportsThinking === true, + input: meta.supportsImages === true ? ["text", "image"] : ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: positiveNumber(meta.maxTokens, DEFAULT_CONTEXT_WINDOW), + maxTokens: id.toLowerCase().includes("claude") + ? Math.min(advertisedMaxTokens, DEFAULT_MAX_TOKENS) + : advertisedMaxTokens, + }; +} + +export async function fetchAntigravityModels( + context: RefreshModelsContext, +): Promise { + if (!context.allowNetwork) return []; + context.signal.throwIfAborted(); + const credential = context.credential; + if (!credential || credential.type !== "oauth") return []; + await ensureAntigravityVersion(context.signal); + + for (const endpoint of DISCOVERY_ENDPOINTS) { + if (context.signal.aborted) break; + try { + const response = await fetch( + `${endpoint}${FETCH_AVAILABLE_MODELS_PATH}`, + { + method: "POST", + headers: { + Authorization: `Bearer ${credential.access}`, + "Content-Type": "application/json", + "User-Agent": getAntigravityUserAgent(), + }, + body: "{}", + signal: context.signal, + }, + ); + if (!response.ok) continue; + const parsed = Value.Parse( + DiscoveryResponseSchema, + await response.json(), + ) as DiscoveryResponse; + if (!parsed.models) continue; + const discovered = Object.entries(parsed.models) + .filter( + ([id, meta]) => + !Object.hasOwn(DISCOVERY_DENYLIST, id) && meta.isInternal !== true, + ) + .map(([id, meta]) => toModelDefinition(id, meta)); + return collapseAntigravityModels(discovered); + } catch { + // Try the next endpoint; total failure must preserve the stored catalog. + } + } + context.signal.throwIfAborted(); + throw new Error("Antigravity model discovery failed on all endpoints"); +} diff --git a/extensions/ai-providers/antigravity/google-conversion.ts b/extensions/ai-providers/antigravity/google-conversion.ts new file mode 100644 index 00000000..1c1165e3 --- /dev/null +++ b/extensions/ai-providers/antigravity/google-conversion.ts @@ -0,0 +1,455 @@ +/** + * Google/Cloud Code Assist message conversion used by the Antigravity adapter. + * + * Kept local because Pi only exposes its documented runtime modules to loaded + * extensions; `@earendil-works/pi-ai/api/google-shared` is not one of them. + * Behavior is adapted from pi-ai 0.84.1's MIT-licensed google-shared converter. + */ + +import type { + Api, + AssistantMessage, + Context, + ImageContent, + Message, + Model, + StopReason, + Tool, + ToolCall, +} from "@earendil-works/pi-ai/compat"; + +const NON_VISION_USER_IMAGE_PLACEHOLDER = + "(image omitted: model does not support images)"; +const NON_VISION_TOOL_IMAGE_PLACEHOLDER = + "(tool image omitted: model does not support images)"; +const BASE64_SIGNATURE_PATTERN = /^[A-Za-z0-9+/]+={0,2}$/; +const JSON_SCHEMA_META_DECLARATIONS = new Set([ + "$schema", + "$id", + "$anchor", + "$dynamicAnchor", + "$vocabulary", + "$comment", + "$defs", + "definitions", +]); + +interface GoogleFunctionCall { + id?: string; + name: string; + args: Record; +} + +interface GoogleFunctionResponse { + id?: string; + name: string; + response: { output: string } | { error: string }; + parts?: GooglePart[]; +} + +export interface GooglePart { + text?: string; + thought?: boolean; + thoughtSignature?: string; + inlineData?: { mimeType: string; data: string }; + functionCall?: GoogleFunctionCall; + functionResponse?: GoogleFunctionResponse; +} + +export interface GoogleContent { + role: "user" | "model"; + parts: GooglePart[]; +} + +function sanitizeSurrogates(text: string): string { + return text.replace( + /[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?, +): Message[] { + if (model.input.includes("image")) return messages; + return messages.map((message) => { + if (message.role === "user" && Array.isArray(message.content)) { + return { + ...message, + content: replaceContentImages( + message.content, + NON_VISION_USER_IMAGE_PLACEHOLDER, + ), + }; + } + if (message.role === "toolResult") { + return { + ...message, + content: replaceContentImages( + message.content, + NON_VISION_TOOL_IMAGE_PLACEHOLDER, + ), + }; + } + return message; + }); +} + +function normalizeMessageContent(message: Message): Message { + if (message.content !== null && message.content !== undefined) return message; + return { ...message, content: [] }; +} + +function transformMessages( + messages: Message[], + model: Model, + normalizeToolCallId: (id: string) => string, +): Message[] { + const toolCallIdMap = new Map(); + const normalizedMessages = messages.map(normalizeMessageContent); + const imageAwareMessages = downgradeUnsupportedImages( + normalizedMessages, + model, + ); + const transformed = imageAwareMessages.map((message): Message => { + if (message.role === "user") return message; + if (message.role === "toolResult") { + const normalizedId = toolCallIdMap.get(message.toolCallId); + return normalizedId && normalizedId !== message.toolCallId + ? { ...message, toolCallId: normalizedId } + : message; + } + + const isSameModel = + message.provider === model.provider && + message.api === model.api && + message.model === model.id; + const content: AssistantMessage["content"] = []; + for (const block of message.content) { + if (block.type === "thinking") { + if (block.redacted) { + if (isSameModel) content.push(block); + continue; + } + if (isSameModel && block.thinkingSignature) { + content.push(block); + continue; + } + if (!block.thinking.trim()) continue; + content.push( + isSameModel ? block : { type: "text", text: block.thinking }, + ); + continue; + } + if (block.type === "text") { + content.push(isSameModel ? block : { type: "text", text: block.text }); + continue; + } + + let normalizedToolCall: ToolCall = block; + if (!isSameModel && block.thoughtSignature) { + const { thoughtSignature: _thoughtSignature, ...unsignedToolCall } = + normalizedToolCall; + normalizedToolCall = unsignedToolCall; + } + if (!isSameModel) { + const normalizedId = normalizeToolCallId(block.id); + if (normalizedId !== block.id) { + toolCallIdMap.set(block.id, normalizedId); + normalizedToolCall = { ...normalizedToolCall, id: normalizedId }; + } + } + content.push(normalizedToolCall); + } + return { ...message, content }; + }); + + const result: Message[] = []; + let pendingToolCalls: ToolCall[] = []; + let existingToolResultIds = new Set(); + const insertSyntheticToolResults = () => { + for (const toolCall of pendingToolCalls) { + if (existingToolResultIds.has(toolCall.id)) continue; + result.push({ + role: "toolResult", + toolCallId: toolCall.id, + toolName: toolCall.name, + content: [{ type: "text", text: "No result provided" }], + isError: true, + timestamp: Date.now(), + }); + } + pendingToolCalls = []; + existingToolResultIds = new Set(); + }; + + for (const message of transformed) { + if (message.role === "assistant") { + insertSyntheticToolResults(); + if (message.stopReason === "error" || message.stopReason === "aborted") { + continue; + } + pendingToolCalls = message.content.filter( + (block): block is ToolCall => block.type === "toolCall", + ); + result.push(message); + continue; + } + if (message.role === "toolResult") { + existingToolResultIds.add(message.toolCallId); + result.push(message); + continue; + } + insertSyntheticToolResults(); + result.push(message); + } + insertSyntheticToolResults(); + return result; +} + +function getGeminiMajorVersion(modelId: string): number | undefined { + const match = modelId.toLowerCase().match(/^gemini(?:-live)?-(\d+)/); + return match ? Number.parseInt(match[1], 10) : undefined; +} + +function requiresToolCallId(modelId: string): boolean { + const geminiMajorVersion = getGeminiMajorVersion(modelId); + return ( + modelId.startsWith("claude-") || + modelId.startsWith("gpt-oss-") || + (geminiMajorVersion !== undefined && geminiMajorVersion >= 3) + ); +} + +function supportsMultimodalFunctionResponse(modelId: string): boolean { + const geminiMajorVersion = getGeminiMajorVersion(modelId); + return geminiMajorVersion === undefined || geminiMajorVersion >= 3; +} + +function resolveThoughtSignature( + isSameProviderAndModel: boolean, + signature: string | undefined, +): string | undefined { + if (!isSameProviderAndModel || !signature || signature.length % 4 !== 0) { + return undefined; + } + return BASE64_SIGNATURE_PATTERN.test(signature) ? signature : undefined; +} + +export function convertMessages( + model: Model, + context: Context, +): GoogleContent[] { + const contents: GoogleContent[] = []; + const normalizeToolCallId = (id: string) => + requiresToolCallId(model.id) + ? id.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 64) + : id; + const messages = transformMessages( + context.messages, + model, + normalizeToolCallId, + ); + + for (const message of messages) { + if (message.role === "user") { + const parts = + typeof message.content === "string" + ? [{ text: sanitizeSurrogates(message.content) }] + : message.content.map( + (item): GooglePart => + item.type === "text" + ? { text: sanitizeSurrogates(item.text) } + : { + inlineData: { + mimeType: item.mimeType, + data: item.data, + }, + }, + ); + if (parts.length > 0) contents.push({ role: "user", parts }); + continue; + } + + if (message.role === "assistant") { + const parts: GooglePart[] = []; + const isSameProviderAndModel = + message.provider === model.provider && message.model === model.id; + for (const block of message.content) { + if (block.type === "text") { + const thoughtSignature = resolveThoughtSignature( + isSameProviderAndModel, + block.textSignature, + ); + if (!block.text.trim() && !thoughtSignature) continue; + parts.push({ + text: sanitizeSurrogates(block.text), + ...(thoughtSignature ? { thoughtSignature } : {}), + }); + continue; + } + if (block.type === "thinking") { + if (!isSameProviderAndModel) { + if (block.thinking.trim()) { + parts.push({ text: sanitizeSurrogates(block.thinking) }); + } + continue; + } + const thoughtSignature = resolveThoughtSignature( + true, + block.thinkingSignature, + ); + if (!block.thinking.trim() && !thoughtSignature) continue; + parts.push({ + thought: true, + text: sanitizeSurrogates(block.thinking), + ...(thoughtSignature ? { thoughtSignature } : {}), + }); + continue; + } + + const thoughtSignature = resolveThoughtSignature( + isSameProviderAndModel, + block.thoughtSignature, + ); + parts.push({ + functionCall: { + name: block.name, + args: block.arguments ?? {}, + ...(requiresToolCallId(model.id) ? { id: block.id } : {}), + }, + ...(thoughtSignature ? { thoughtSignature } : {}), + }); + } + if (parts.length > 0) contents.push({ role: "model", parts }); + continue; + } + + const textResult = message.content + .filter((part) => part.type === "text") + .map((part) => part.text) + .join("\n"); + const imageContent = model.input.includes("image") + ? message.content.filter( + (part): part is ImageContent => part.type === "image", + ) + : []; + const hasImages = imageContent.length > 0; + const responseValue = textResult + ? sanitizeSurrogates(textResult) + : hasImages + ? "(see attached image)" + : ""; + const imageParts: GooglePart[] = imageContent.map((image) => ({ + inlineData: { mimeType: image.mimeType, data: image.data }, + })); + const supportsMultimodal = supportsMultimodalFunctionResponse(model.id); + const functionResponsePart: GooglePart = { + functionResponse: { + name: message.toolName, + response: message.isError + ? { error: responseValue } + : { output: responseValue }, + ...(hasImages && supportsMultimodal ? { parts: imageParts } : {}), + ...(requiresToolCallId(model.id) ? { id: message.toolCallId } : {}), + }, + }; + const lastContent = contents.at(-1); + if ( + lastContent?.role === "user" && + lastContent.parts.some((part) => part.functionResponse) + ) { + lastContent.parts.push(functionResponsePart); + } else { + contents.push({ role: "user", parts: [functionResponsePart] }); + } + if (hasImages && !supportsMultimodal) { + contents.push({ + role: "user", + parts: [{ text: "Tool result image:" }, ...imageParts], + }); + } + } + return contents; +} + +function sanitizeForOpenApi( + schema: unknown, + insidePropertiesMap = false, +): unknown { + if (typeof schema !== "object" || schema === null || Array.isArray(schema)) { + return schema; + } + const result: Record = {}; + for (const [key, value] of Object.entries(schema)) { + if (insidePropertiesMap) { + result[key] = sanitizeForOpenApi(value); + continue; + } + if (!JSON_SCHEMA_META_DECLARATIONS.has(key)) { + result[key] = sanitizeForOpenApi(value, key === "properties"); + } + } + return result; +} + +export function convertTools( + tools: Tool[], + useParameters = false, +): { functionDeclarations: Record[] }[] | undefined { + if (tools.length === 0) return undefined; + return [ + { + functionDeclarations: tools.map((tool) => ({ + name: tool.name, + description: tool.description, + ...(useParameters + ? { parameters: sanitizeForOpenApi(tool.parameters) } + : { parametersJsonSchema: tool.parameters }), + })), + }, + ]; +} + +export function isThinkingPart(part: { + thought?: boolean; + thoughtSignature?: string; +}): boolean { + return part.thought === true; +} + +export function retainThoughtSignature( + existing: string | undefined, + incoming: string | undefined, +): string | undefined { + return typeof incoming === "string" && incoming.length > 0 + ? incoming + : existing; +} + +export function mapStopReasonString(reason: string): StopReason { + if (reason === "STOP") return "stop"; + if (reason === "MAX_TOKENS") return "length"; + return "error"; +} diff --git a/extensions/ai-providers/antigravity/models.ts b/extensions/ai-providers/antigravity/models.ts new file mode 100644 index 00000000..a17db229 --- /dev/null +++ b/extensions/ai-providers/antigravity/models.ts @@ -0,0 +1,84 @@ +/** + * Static Antigravity model table. + * + * This is the offline fallback so the provider is usable before (or without) + * discovery. The live source of truth is the Cloud Code Assist + * `fetchAvailableModels` endpoint, wired via `createProvider.fetchModels` (see + * oh-my-pi packages/catalog/src/discovery/antigravity.ts for the reference + * implementation and its denylist). Wire ids follow the real client: + * `gemini-*` for Gemini routes, `claude-*[-thinking]` for Claude routes. + */ + +import type { Model } from "@earendil-works/pi-ai"; + +export const ANTIGRAVITY_API_URL = "https://daily-cloudcode-pa.googleapis.com"; + +export type AntigravityProviderModel = Model<"antigravity-cloudcode"> & { + requestModelId?: string; +}; + +const ZERO_COST: AntigravityProviderModel["cost"] = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, +}; + +// Defaults mirror omp's discovery normalization (200k context, 64k output). +const DEFAULT_CONTEXT_WINDOW = 200_000; +const DEFAULT_MAX_TOKENS = 64_000; + +function define( + id: string, + name: string, + reasoning: boolean, + requestModelId?: string, + capabilities?: Pick< + AntigravityProviderModel, + "input" | "contextWindow" | "maxTokens" + >, +): AntigravityProviderModel { + return { + id, + name, + api: "antigravity-cloudcode", + provider: "google-antigravity", + baseUrl: ANTIGRAVITY_API_URL, + reasoning, + input: capabilities?.input ?? ["text", "image"], + cost: ZERO_COST, + contextWindow: capabilities?.contextWindow ?? DEFAULT_CONTEXT_WINDOW, + maxTokens: capabilities?.maxTokens ?? DEFAULT_MAX_TOKENS, + ...(requestModelId ? { requestModelId } : {}), + }; +} + +export const ANTIGRAVITY_MODELS: AntigravityProviderModel[] = [ + define( + "gemini-3.7-flash", + "Gemini 3.7 Flash (Antigravity)", + true, + "gemini-3.7-flash-low", + ), + define( + "gemini-3.5-flash", + "Gemini 3.5 Flash (Antigravity)", + true, + "gemini-3.5-flash-extra-low", + ), + define( + "gemini-3.1-pro", + "Gemini 3.1 Pro (Antigravity)", + true, + "gemini-3.1-pro-low", + ), + define("claude-sonnet-4-6", "Claude Sonnet 4.6 (Antigravity)", true), + define("claude-opus-4-6", "Claude Opus 4.6 (Antigravity)", true), + define( + "gpt-oss-120b", + "GPT OSS 120B (Antigravity)", + true, + "gpt-oss-120b-medium", + { input: ["text"], contextWindow: 131_072, maxTokens: 32_768 }, + ), +]; diff --git a/extensions/ai-providers/antigravity/oauth.ts b/extensions/ai-providers/antigravity/oauth.ts new file mode 100644 index 00000000..ce74ae1f --- /dev/null +++ b/extensions/ai-providers/antigravity/oauth.ts @@ -0,0 +1,700 @@ +/** + * Google OAuth login + Cloud Code Assist provisioning for Antigravity. + * + * Flow (mirrors the real Antigravity client, reference: + * oh-my-pi packages/ai/src/registry/oauth/google-antigravity.ts): + * + * 1. Serve a localhost callback (127.0.0.1:51121/oauth-callback, ephemeral + * fallback) and hand the Google consent URL to pi via `callbacks.onAuth`. + * 2. Exchange the authorization code (no PKCE — the Antigravity client is a + * confidential client with a public client_secret). + * 3. Provision: loadCodeAssist → (onboardUser free-tier poll when the + * account has no tier yet) → loadCodeAssist for the project id. + * 4. Return pi's OAuthCredentials plus { projectId, email } extras; pi + * persists the whole object in auth.json and hands it back verbatim. + */ + +import * as http from "node:http"; +import type { AddressInfo } from "node:net"; +import type { OAuthCredentials } from "@earendil-works/pi-ai/compat"; +import type { CancellableOAuthLoginCallbacks } from "../oauth-adapter.ts"; +import type { AntigravityCredentials } from "./credentials.ts"; + +// Public OAuth client identity of the Antigravity IDE (also used by omp and +// the pi-agy project; it is designed to be embedded in shipped binaries). +// Stored as base64 so GitHub secret scanning does not treat the public +// confidential-client pair as a leaked secret. +const CLIENT_ID = atob( + "MTA3MTAwNjA2MDU5MS10bWhzc2luMmgyMWxjcmUyMzV2dG9sb2poNGc0MDNlcC5hcHBzLmdvb2dsZXVzZXJjb250ZW50LmNvbQ==", +); +const CLIENT_SECRET = atob("R09DU1BYLUs1OEZXUjQ4NkxkTEoxbUxCOHNYQzR6NnFEQWY="); + +const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"; +const TOKEN_URL = "https://oauth2.googleapis.com/token"; +const USERINFO_URL = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json"; + +const SCOPES = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", + "https://www.googleapis.com/auth/experimentsandconfigs", +]; + +const CALLBACK_PORT = 51121; +const CALLBACK_PATH = "/oauth-callback"; +const LOGIN_TIMEOUT_MS = 5 * 60 * 1000; +const REQUEST_TIMEOUT_MS = 30_000; + +const CLOUD_CODE_ASSIST_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com"; +const LOAD_CODE_ASSIST_URL = `${CLOUD_CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist`; +const ONBOARD_USER_URL = `${CLOUD_CODE_ASSIST_ENDPOINT}/v1internal:onboardUser`; +const OPERATIONS_URL = `${CLOUD_CODE_ASSIST_ENDPOINT}/v1internal`; +const FREE_TIER_ID = "free-tier"; +const ONBOARD_TIMEOUT_MS = 30_000; +const ONBOARD_POLL_INTERVAL_MS = 1_000; + +/** Refresh tokens are considered stale this long before their actual expiry. */ +const EXPIRY_MARGIN_MS = 5 * 60 * 1000; + +const LOAD_CODE_ASSIST_METADATA = { ideType: "ANTIGRAVITY" }; + +// --------------------------------------------------------------------------- +// User-Agent (version-gated upstream; see omp catalog/wire/gemini-headers.ts) +// --------------------------------------------------------------------------- + +const DEFAULT_ANTIGRAVITY_VERSION = "2.8.0"; +const VERSION_MANIFEST_URL = + "https://antigravity-hub-auto-updater-974169037036.us-central1.run.app/manifest/latest-arm64-mac.yml"; +const VERSION_FETCH_TIMEOUT_MS = 5_000; + +let discoveredVersion: string | undefined; +let versionFetch: Promise | undefined; + +export function getAntigravityUserAgent(): string { + const version = + process.env.OPENPI_ANTIGRAVITY_VERSION || + discoveredVersion || + DEFAULT_ANTIGRAVITY_VERSION; + const cl = process.env.OPENPI_ANTIGRAVITY_CL || "963137146"; + return `antigravity/hub/${version} (aidev_client; os_type=darwin; arch=arm64; cl=${cl})`; +} + +/** Best-effort version discovery from the official update manifest. */ +export function ensureAntigravityVersion( + signal?: AbortSignal, + fetcher: typeof fetch = fetch, +): Promise { + if (process.env.OPENPI_ANTIGRAVITY_VERSION || discoveredVersion) { + return Promise.resolve(); + } + if (versionFetch) return versionFetch; + versionFetch = (async () => { + try { + const timeout = AbortSignal.timeout(VERSION_FETCH_TIMEOUT_MS); + const response = await fetcher(VERSION_MANIFEST_URL, { + headers: { + "Cache-Control": "no-cache", + "User-Agent": "electron-builder", + }, + signal: signal ? AbortSignal.any([signal, timeout]) : timeout, + }); + if (response.ok) { + discoveredVersion = parseManifestVersion(await response.text()); + } + } catch { + // Silent: the pinned fallback stays valid when discovery fails. + } finally { + if (!discoveredVersion) versionFetch = undefined; + } + })(); + return versionFetch; +} + +function parseManifestVersion(yamlText: string): string | undefined { + for (const line of yamlText.split(/\r?\n/)) { + const match = + /^\s*version\s*:\s*(?:"([^"]*)"|'([^']*)'|([^\s#]+))\s*(?:#.*)?$/.exec( + line, + ); + if (!match) continue; + const version = (match[1] ?? match[2] ?? match[3] ?? "").trim(); + return /^\d+\.\d+\.\d+$/.test(version) ? version : undefined; + } + return undefined; +} + +// --------------------------------------------------------------------------- +// HTTP helpers +// --------------------------------------------------------------------------- + +function throwIfCancelled(signal: AbortSignal | undefined): void { + if (signal?.aborted) { + throw new Error(`Login cancelled: ${String(signal.reason ?? "aborted")}`); + } +} + +function withCancellation( + promise: Promise, + signal: AbortSignal | undefined, +): Promise { + if (!signal) return promise; + throwIfCancelled(signal); + return new Promise((resolve, reject) => { + const onAbort = () => { + reject(new Error("Login cancelled")); + }; + signal.addEventListener("abort", onAbort, { once: true }); + promise.then( + (value) => { + signal.removeEventListener("abort", onAbort); + resolve(value); + }, + (error) => { + signal.removeEventListener("abort", onAbort); + reject(error); + }, + ); + }); +} + +/** fetch with login cancellation + per-request timeout composed. */ +async function oauthFetch( + url: string, + init: RequestInit, + signal: AbortSignal | undefined, + timeoutMs = REQUEST_TIMEOUT_MS, +): Promise { + const timeout = AbortSignal.timeout(timeoutMs); + const combined = signal ? AbortSignal.any([signal, timeout]) : timeout; + try { + return await fetch(url, { ...init, signal: combined }); + } catch (error) { + throwIfCancelled(signal); + if (timeout.aborted) { + throw new Error(`Timed out after ${timeoutMs}ms waiting for ${url}`); + } + throw error; + } +} + +function sleep(ms: number, signal: AbortSignal | undefined): Promise { + const { promise, resolve, reject } = Promise.withResolvers(); + const onAbort = () => { + clearTimeout(timer); + reject(new Error("Login cancelled")); + }; + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + signal?.addEventListener("abort", onAbort, { once: true }); + return promise; +} + +// --------------------------------------------------------------------------- +// Localhost callback server +// --------------------------------------------------------------------------- + +interface CallbackResult { + code: string; + state: string; +} + +interface CallbackServer { + port: number; + waitForCallback( + expectedState: string, + signal: AbortSignal | undefined, + ): Promise; + close(): void; +} + +function parseCallbackInput(input: string): CallbackResult | undefined { + const value = input.trim(); + if (!value) return undefined; + try { + const url = new URL(value); + const code = url.searchParams.get("code"); + if (!code) return undefined; + return { code, state: url.searchParams.get("state") ?? "" }; + } catch { + return { code: value, state: "" }; + } +} + +const LOGIN_SUCCESS_HTML = + "

Antigravity login complete

" + + "

You can return to your terminal.

"; + +const LOGIN_ERROR_HTML = + "

Antigravity login failed

" + + "

Missing or mismatched authorization response.

"; + +async function startCallbackServer(): Promise { + const server = http.createServer(); + const pending: ((result: CallbackResult) => void)[] = []; + const failures: { expectedState: string; reject(error: Error): void }[] = []; + + server.on("request", (request, response) => { + const url = new URL(request.url ?? "/", "http://127.0.0.1"); + if (url.pathname !== CALLBACK_PATH) { + response.writeHead(404).end(); + return; + } + const code = url.searchParams.get("code"); + const state = url.searchParams.get("state"); + const error = url.searchParams.get("error"); + const matching = state + ? failures.findIndex((entry) => entry.expectedState === state) + : -1; + if (error || !code || matching < 0) { + response.writeHead(400, { "Content-Type": "text/html" }); + response.end(LOGIN_ERROR_HTML); + // Only a denial carrying our nonce can terminate the flow. A random + // local request must not be able to cancel an in-flight login. + if (error && matching >= 0) { + const [failure] = failures.splice(matching, 1); + failure?.reject(new Error(`Authorization failed: ${error}`)); + } + return; + } + response.writeHead(200, { "Content-Type": "text/html" }); + response.end(LOGIN_SUCCESS_HTML); + for (const resolve of pending.splice(0)) resolve({ code, state: state! }); + failures.length = 0; + }); + + async function listen(port: number): Promise { + const { promise, resolve, reject } = Promise.withResolvers(); + server.once("error", reject); + server.listen(port, "127.0.0.1", () => { + server.removeListener("error", reject); + resolve((server.address() as AddressInfo).port); + }); + return promise; + } + + let port: number; + try { + port = await listen(CALLBACK_PORT); + } catch { + // Preferred port busy: fall back to an ephemeral port; the redirect_uri + // is built from whatever we actually bound. + port = await listen(0); + } + + return { + port, + waitForCallback(expectedState, signal) { + const { promise, resolve, reject } = + Promise.withResolvers(); + let settled = false; + let onAbort: (() => void) | undefined; + const pendingCallback = (result: CallbackResult) => { + if (result.state !== expectedState) { + return; + } + settle(() => resolve(result)); + }; + const failure = { + expectedState, + reject(error: Error) { + settle(() => reject(error)); + }, + }; + const cleanup = () => { + const pendingIndex = pending.indexOf(pendingCallback); + if (pendingIndex >= 0) pending.splice(pendingIndex, 1); + const failureIndex = failures.indexOf(failure); + if (failureIndex >= 0) failures.splice(failureIndex, 1); + if (onAbort) signal?.removeEventListener("abort", onAbort); + }; + const settle = (complete: () => void) => { + if (settled) return; + settled = true; + cleanup(); + complete(); + }; + pending.push(pendingCallback); + failures.push(failure); + if (signal) { + onAbort = () => settle(() => reject(new Error("Login cancelled"))); + if (signal.aborted) onAbort(); + else signal.addEventListener("abort", onAbort, { once: true }); + } + return promise; + }, + close() { + for (const failure of [...failures]) { + failure.reject(new Error("Login callback server closed")); + } + pending.length = 0; + failures.length = 0; + server.close(); + }, + }; +} + +// --------------------------------------------------------------------------- +// Cloud Code Assist provisioning +// --------------------------------------------------------------------------- + +interface LoadCodeAssistResponse { + currentTier?: { id?: string } | null; + paidTier?: { id?: string } | null; + allowedTiers?: { id?: string }[]; + ineligibleTiers?: { + tierId?: string; + reasonMessage?: string; + validationUrl?: string; + }[]; + cloudaicompanionProject?: string; +} + +interface OnboardOperation { + name?: string; + done?: boolean; + error?: { code?: number; message?: string } | null; + response?: { cloudaicompanionProject?: string } | null; +} + +async function postCloudCodeAssist( + label: string, + url: string, + accessToken: string, + body: Record, + signal: AbortSignal | undefined, + timeoutMs = REQUEST_TIMEOUT_MS, +): Promise { + throwIfCancelled(signal); + const response = await oauthFetch( + url, + { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + "User-Agent": getAntigravityUserAgent(), + }, + body: JSON.stringify(body), + }, + signal, + timeoutMs, + ); + if (response.status !== 200) { + const errorText = await response.text(); + throw new Error( + `${label} failed: ${response.status} ${response.statusText}: ${errorText}`, + ); + } + return response.json(); +} + +function asLoadCodeAssistResponse(payload: unknown): LoadCodeAssistResponse { + if (payload === null || typeof payload !== "object") { + throw new Error("loadCodeAssist returned a non-object response"); + } + return payload as LoadCodeAssistResponse; +} + +async function loadCodeAssist( + accessToken: string, + signal: AbortSignal | undefined, +): Promise { + let payload = asLoadCodeAssistResponse( + await postCloudCodeAssist( + "loadCodeAssist", + LOAD_CODE_ASSIST_URL, + accessToken, + { metadata: LOAD_CODE_ASSIST_METADATA }, + signal, + ), + ); + const projectId = payload.cloudaicompanionProject; + if (payload.paidTier === undefined && projectId) { + payload = asLoadCodeAssistResponse( + await postCloudCodeAssist( + "loadCodeAssist", + LOAD_CODE_ASSIST_URL, + accessToken, + { + cloudaicompanionProject: projectId, + metadata: LOAD_CODE_ASSIST_METADATA, + }, + signal, + ), + ); + } + return payload; +} + +function assertFreeTierEligible(payload: LoadCodeAssistResponse): void { + if (payload.allowedTiers?.some((tier) => tier.id === FREE_TIER_ID)) return; + const ineligible = payload.ineligibleTiers?.find( + (tier) => tier.tierId === FREE_TIER_ID, + ); + if (!ineligible?.reasonMessage) return; + const validation = ineligible.validationUrl + ? `\n${ineligible.validationUrl}` + : ""; + throw new Error(`${ineligible.reasonMessage}${validation}`); +} + +async function onboardUser( + accessToken: string, + signal: AbortSignal | undefined, +): Promise { + const deadline = Date.now() + ONBOARD_TIMEOUT_MS; + const remaining = () => { + const left = deadline - Date.now(); + if (left <= 0) { + throw new Error(`onboardUser timed out after ${ONBOARD_TIMEOUT_MS}ms`); + } + return left; + }; + + let operation = (await postCloudCodeAssist( + "onboardUser", + ONBOARD_USER_URL, + accessToken, + { tierId: FREE_TIER_ID, metadata: LOAD_CODE_ASSIST_METADATA }, + signal, + remaining(), + )) as OnboardOperation; + + while (true) { + if (operation.done === true) { + if (operation.error) { + const detail = + operation.error.message ?? JSON.stringify(operation.error); + throw new Error(`OnboardUser operation failed: ${detail}`); + } + return; + } + await sleep(Math.min(ONBOARD_POLL_INTERVAL_MS, remaining()), signal); + throwIfCancelled(signal); + if (!operation.name) { + throw new Error("onboardUser returned an operation without a name"); + } + const polled = await oauthFetch( + `${OPERATIONS_URL}/${operation.name}`, + { + headers: { + Authorization: `Bearer ${accessToken}`, + "User-Agent": getAntigravityUserAgent(), + }, + }, + signal, + remaining(), + ); + if (polled.status !== 200) { + throw new Error( + `onboardUser poll failed: ${polled.status} ${await polled.text()}`, + ); + } + operation = (await polled.json()) as OnboardOperation; + } +} + +async function discoverProject( + accessToken: string, + onProgress: ((message: string) => void) | undefined, + signal: AbortSignal | undefined, +): Promise { + onProgress?.("Checking Cloud Code Assist account status..."); + const initial = await loadCodeAssist(accessToken, signal); + assertFreeTierEligible(initial); + if (initial.currentTier === undefined || initial.currentTier === null) { + onProgress?.("Provisioning the Antigravity free tier..."); + await onboardUser(accessToken, signal); + } + onProgress?.("Refreshing Cloud Code Assist project..."); + const refreshed = await loadCodeAssist(accessToken, signal); + const projectId = refreshed.cloudaicompanionProject; + if (!projectId) { + throw new Error("loadCodeAssist did not return a cloudaicompanionProject"); + } + return projectId; +} + +async function fetchUserEmail( + accessToken: string, + signal: AbortSignal | undefined, +): Promise { + try { + const response = await oauthFetch( + USERINFO_URL, + { headers: { Authorization: `Bearer ${accessToken}` } }, + signal, + ); + if (response.ok) { + const data: unknown = await response.json(); + if (data !== null && typeof data === "object" && "email" in data) { + const email = data.email; + if (typeof email === "string") return email; + } + } + } catch { + // Email is display metadata only; never fail the login over it. + } + return undefined; +} + +// --------------------------------------------------------------------------- +// Public flow +// --------------------------------------------------------------------------- + +export async function loginAntigravity( + callbacks: CancellableOAuthLoginCallbacks, +): Promise { + const signal = callbacks.signal; + throwIfCancelled(signal); + await ensureAntigravityVersion(signal); + throwIfCancelled(signal); + + const server = await startCallbackServer(); + try { + throwIfCancelled(signal); + const state = crypto.randomUUID(); + const redirectUri = `http://127.0.0.1:${server.port}${CALLBACK_PATH}`; + const authParams = new URLSearchParams({ + client_id: CLIENT_ID, + response_type: "code", + redirect_uri: redirectUri, + scope: SCOPES.join(" "), + state, + access_type: "offline", + prompt: "consent", + }); + callbacks.onAuth({ + url: `${AUTH_URL}?${authParams.toString()}`, + instructions: "Complete the Google sign-in in your browser.", + }); + + const timeout = AbortSignal.timeout(LOGIN_TIMEOUT_MS); + const loginSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; + const raceController = new AbortController(); + const raceSignal = AbortSignal.any([loginSignal, raceController.signal]); + let callback: CallbackResult; + try { + const callbackPromise = server.waitForCallback(state, raceSignal); + const onManualCodeInput = callbacks.onManualCodeInput; + if (onManualCodeInput) { + const manualPromise = (async () => { + while (true) { + const parsed = parseCallbackInput( + await withCancellation(onManualCodeInput(raceSignal), raceSignal), + ); + if (parsed && (!parsed.state || parsed.state === state)) + return parsed; + } + })(); + callback = await Promise.race([callbackPromise, manualPromise]); + } else { + callback = await callbackPromise; + } + } catch (error) { + if (timeout.aborted) { + throw new Error("Login timed out waiting for the browser callback"); + } + throw error; + } finally { + // The callback and manual prompt are alternatives. Cancel the loser as + // soon as either one supplies a valid authorization code. + raceController.abort(); + } + + callbacks.onProgress?.("Exchanging authorization code for tokens..."); + throwIfCancelled(signal); + const tokenResponse = await oauthFetch( + TOKEN_URL, + { + method: "POST", + headers: { "Content-Type": "application/x-www-form-urlencoded" }, + body: new URLSearchParams({ + client_id: CLIENT_ID, + client_secret: CLIENT_SECRET, + code: callback.code, + grant_type: "authorization_code", + redirect_uri: redirectUri, + }), + }, + signal, + ); + if (!tokenResponse.ok) { + throw new Error( + `Token exchange failed: ${tokenResponse.status} ${await tokenResponse.text()}`, + ); + } + const tokens = (await tokenResponse.json()) as { + access_token: string; + expires_in: number; + refresh_token?: string; + }; + if (!tokens.refresh_token) { + throw new Error( + "Google did not return a refresh token — revoke the app's access " + + "at myaccount.google.com/permissions and try again.", + ); + } + + const projectId = await discoverProject( + tokens.access_token, + callbacks.onProgress, + signal, + ); + const email = await fetchUserEmail(tokens.access_token, signal); + throwIfCancelled(signal); + + const credentials: AntigravityCredentials = { + refresh: tokens.refresh_token, + access: tokens.access_token, + expires: Date.now() + tokens.expires_in * 1000 - EXPIRY_MARGIN_MS, + projectId, + email, + }; + return credentials; + } finally { + server.close(); + } +} + +export async function refreshAntigravityToken( + credentials: OAuthCredentials, + signal: AbortSignal, +): Promise { + const existing = credentials as AntigravityCredentials; + const response = await oauthFetch( + TOKEN_URL, + { + method: "POST", + headers: { "Content-Type": "application/x-www-form-urlencoded" }, + body: new URLSearchParams({ + client_id: CLIENT_ID, + client_secret: CLIENT_SECRET, + refresh_token: credentials.refresh, + grant_type: "refresh_token", + }), + }, + signal, + ); + if (!response.ok) { + throw new Error( + `Antigravity token refresh failed: ${response.status} ${await response.text()}`, + ); + } + const data = (await response.json()) as { + access_token: string; + expires_in: number; + refresh_token?: string; + }; + const refreshed: AntigravityCredentials = { + refresh: data.refresh_token || credentials.refresh, + access: data.access_token, + expires: Date.now() + data.expires_in * 1000 - EXPIRY_MARGIN_MS, + projectId: existing.projectId, + email: existing.email, + }; + return refreshed; +} diff --git a/extensions/ai-providers/antigravity/provider.ts b/extensions/ai-providers/antigravity/provider.ts new file mode 100644 index 00000000..484c5ff5 --- /dev/null +++ b/extensions/ai-providers/antigravity/provider.ts @@ -0,0 +1,1126 @@ +/** + * Cloud Code Assist (Antigravity) streamSimple implementation. + * + * Wire shape (reference: oh-my-pi packages/ai/src/providers/google-gemini-cli.ts, + * shared google-gemini-cli/google-antigravity implementation): + * + * POST {endpoint}/v1internal:streamGenerateContent?alt=sse + * { project, requestId, model, userAgent: "antigravity", requestType: "agent", + * request: { contents, systemInstruction, tools, toolConfig, + * generationConfig, labels, sessionId } } + * + * SSE frames carry `{ response: GenerateContentResponse }` envelopes. Endpoint + * failover: daily-cloudcode-pa → daily-cloudcode-pa.sandbox. Message/tool + * conversion is kept local because Pi extensions cannot resolve pi-ai's + * internal google-shared module at installed runtime. + */ + +import { createHash, randomUUID } from "node:crypto"; +import type { + Api, + AssistantMessage, + AssistantMessageEventStream, + Context, + Model, + SimpleStreamOptions, + ToolCall, +} from "@earendil-works/pi-ai/compat"; +import { createAssistantMessageEventStream } from "@earendil-works/pi-ai/compat"; +import { decodeApiKey } from "./credentials.ts"; +import { + convertMessages, + convertTools, + isThinkingPart, + mapStopReasonString, + retainThoughtSignature, +} from "./google-conversion.ts"; +import { ensureAntigravityVersion, getAntigravityUserAgent } from "./oauth.ts"; +import { routeAntigravityModel } from "./routing.ts"; + +const ENDPOINTS = [ + "https://daily-cloudcode-pa.googleapis.com", + "https://daily-cloudcode-pa.sandbox.googleapis.com", +] as const; + +const CLAUDE_THINKING_BETA_HEADER = "interleaved-thinking-2025-05-14"; + +const FLASH_FIRST_EVENT_TIMEOUT_MS = 60_000; +const DEFAULT_FIRST_EVENT_TIMEOUT_MS = 300_000; +const MAX_ERROR_BODY_BYTES = 64 * 1024; + +type AntigravityStreamOptions = Omit & { + /** Keep Antigravity's richer modes compatible across Pi 0.84.1 and 0.84.3+. */ + toolChoice?: + | "auto" + | "none" + | "any" + | { mode: "ANY"; allowedFunctionNames: [string, ...string[]] }; +}; + +const FORCED_TOOL_DIRECTIVE = + "TOOL-ONLY TURN. This turn accepts a tool call and nothing else; " + + "a text reply here is discarded unread and you will be re-prompted. " + + "Emit the tool call now."; + +interface AntigravitySessionState { + agentId: string; + trajectoryId: string; + sessionId: string; + stepIndex: number; + lastExecutionId?: string; +} + +const sessionStates = new Map(); +const MAX_SESSION_STATES = 64; + +// Cloud Code Assist 400s on these JSON Schema keywords (reference: omp +// packages/ai/src/utils/schema/fields.ts — union of UNSUPPORTED_SCHEMA_FIELDS +// and LIFTABLE_TO_DESCRIPTION_FIELDS). Constraints that remain useful to the +// model are serialized into the sibling description before being removed. +const CCA_UNSUPPORTED_SCHEMA_FIELDS: Record = { + $schema: true, + $ref: true, + $defs: true, + $dynamicRef: true, + $dynamicAnchor: true, + $comment: true, + examples: true, + prefixItems: true, + unevaluatedProperties: true, + unevaluatedItems: true, + patternProperties: true, + additionalProperties: true, + propertyNames: true, + minItems: true, + maxItems: true, + minLength: true, + maxLength: true, + minProperties: true, + maxProperties: true, + minimum: true, + maximum: true, + exclusiveMinimum: true, + exclusiveMaximum: true, + multipleOf: true, + uniqueItems: true, + pattern: true, + format: true, + default: true, + deprecated: true, + readOnly: true, + writeOnly: true, + dependencies: true, + dependentSchemas: true, + dependentRequired: true, + "x-mcp-header": true, +}; +// Stripped keywords whose constraint stays model-visible by spilling into the +// node's description (omp LIFTABLE_TO_DESCRIPTION_FIELDS, "spill" format). +const CCA_LIFTABLE_TO_DESCRIPTION: Record = { + pattern: true, + format: true, + minLength: true, + maxLength: true, + minimum: true, + maximum: true, + exclusiveMinimum: true, + exclusiveMaximum: true, + multipleOf: true, + minItems: true, + maxItems: true, + uniqueItems: true, + minProperties: true, + maxProperties: true, + default: true, + examples: true, +}; + +function sanitizeSchemaForCcaValue( + value: unknown, + insidePropertiesMap: boolean, +): unknown { + if (Array.isArray(value)) { + return value.map((entry) => sanitizeSchemaForCcaValue(entry, false)); + } + if (value === null || typeof value !== "object") return value; + const out: Record = {}; + const spill: Array<[string, unknown]> = []; + for (const [key, entry] of Object.entries(value)) { + // Keys below `properties` are user-defined parameter names, not JSON + // Schema keywords. A tool parameter named `pattern`, for example, must be + // preserved while its schema value is sanitized normally. + if (insidePropertiesMap) { + out[key] = sanitizeSchemaForCcaValue(entry, false); + continue; + } + if (Object.hasOwn(CCA_UNSUPPORTED_SCHEMA_FIELDS, key)) { + if ( + entry !== undefined && + Object.hasOwn(CCA_LIFTABLE_TO_DESCRIPTION, key) + ) { + spill.push([key, entry]); + } + continue; + } + out[key] = sanitizeSchemaForCcaValue(entry, key === "properties"); + } + if (spill.length > 0) { + const formatted = `{${spill.map(([key, entry]) => `${key}: ${JSON.stringify(entry)}`).join(", ")}}`; + const existing = typeof out.description === "string" ? out.description : ""; + out.description = existing ? `${existing}\n\n${formatted}` : formatted; + } + return out; +} + +/** Recursively drop schema keywords Cloud Code Assist rejects. */ +export function sanitizeSchemaForCca(value: unknown): unknown { + return sanitizeSchemaForCcaValue(value, false); +} + +// --------------------------------------------------------------------------- +// Request construction +// --------------------------------------------------------------------------- + +function isClaudeRoute(modelId: string): boolean { + return modelId.toLowerCase().includes("claude"); +} + +function normalizeSystemPrompts( + systemPrompt: Context["systemPrompt"], +): string[] { + if (!systemPrompt) return []; + return Array.isArray(systemPrompt) ? systemPrompt : [systemPrompt]; +} + +/** Deterministic conversation id: hash of the first user text, like the client. */ +function deriveSessionId(context: Context): string { + for (const message of context.messages) { + if (message.role !== "user") continue; + const content = message.content; + const text = + typeof content === "string" + ? content + : content.find((part) => part.type === "text")?.text; + if (text && text.trim().length > 0) { + const digest = createHash("sha256").update(text).digest(); + let value = 0n; + for (let i = 0; i < 8; i++) { + value = (value << 8n) | BigInt(digest[i]); + } + // The real client formats its bounded int63 identifier as a negative + // decimal string rather than using a UUID on the wire. + return `-${String(value & 0x7fffffffffffffffn)}`; + } + break; + } + const random = BigInt(`0x${randomUUID().replaceAll("-", "").slice(0, 16)}`); + return `-${String(random & 0x7fffffffffffffffn)}`; +} + +function getSessionState( + options: AntigravityStreamOptions | undefined, + context: Context, +): AntigravitySessionState | undefined { + const key = options?.sessionId; + if (!key) return undefined; + const existing = sessionStates.get(key); + if (existing) return existing; + if (sessionStates.size >= MAX_SESSION_STATES) { + const oldest = sessionStates.keys().next().value; + if (oldest) sessionStates.delete(oldest); + } + const created: AntigravitySessionState = { + agentId: randomUUID(), + trajectoryId: randomUUID(), + sessionId: deriveSessionId(context), + stepIndex: 1, + }; + sessionStates.set(key, created); + return created; +} + +function buildToolConfig( + model: Model, + hasTools: boolean, + toolChoice: AntigravityStreamOptions["toolChoice"], +): Record | undefined { + if (!hasTools) { + return isClaudeRoute(model.id) && toolChoice !== "none" + ? { functionCallingConfig: { mode: "VALIDATED" } } + : undefined; + } + if (toolChoice === "none") { + return { functionCallingConfig: { mode: "NONE" } }; + } + if (toolChoice === "any") { + return { functionCallingConfig: { mode: "ANY" } }; + } + if (typeof toolChoice === "object") { + return { + functionCallingConfig: { + mode: "ANY", + allowedFunctionNames: [...toolChoice.allowedFunctionNames], + }, + }; + } + // Antigravity's default tool mode is VALIDATED (verified upstream for both + // Gemini and Claude routes). + return { functionCallingConfig: { mode: "VALIDATED" } }; +} + +/** Convert pi tools to CCA functionDeclarations with sanitized schemas. */ +function buildTools( + context: Context, + toolChoice: AntigravityStreamOptions["toolChoice"], +): Record[] | undefined { + if (toolChoice === "none") return undefined; + const tools = context.tools; + if (!tools || tools.length === 0) return undefined; + const converted = convertTools([...tools], true) as + | { functionDeclarations: Record[] }[] + | undefined; + if (!converted) return undefined; + return converted.map((group) => ({ + ...group, + functionDeclarations: group.functionDeclarations.map((declaration) => ({ + ...declaration, + parameters: sanitizeSchemaForCca(declaration.parameters), + })), + })); +} + +export function buildRequestBody( + model: Model, + context: Context, + options: AntigravityStreamOptions | undefined, + projectId: string, + state?: AntigravitySessionState, +): Record { + const contents = convertMessages(model, context); + + const request: Record = { contents }; + const systemPrompts = normalizeSystemPrompts(context.systemPrompt); + if (systemPrompts.length > 0) { + request.systemInstruction = { + role: "user", + parts: systemPrompts.map((text) => ({ text })), + }; + } + + const tools = buildTools(context, options?.toolChoice); + if (tools) request.tools = tools; + const toolConfig = buildToolConfig( + model, + Boolean(tools), + options?.toolChoice, + ); + if (toolConfig) request.toolConfig = toolConfig; + if ( + tools && + !isClaudeRoute(model.id) && + (options?.toolChoice === "any" || typeof options?.toolChoice === "object") + ) { + contents.push({ role: "user", parts: [{ text: FORCED_TOOL_DIRECTIVE }] }); + } + + const route = routeAntigravityModel( + model.id, + model.reasoning ? options?.reasoning : undefined, + options?.thinkingBudgets, + model as Model & { + requestModelId?: string; + antigravityEffortRouting?: Partial< + Record, string> + >; + }, + ); + const requestedMaxTokens = options?.maxTokens ?? model.maxTokens; + const generationConfig: Record = { + maxOutputTokens: isClaudeRoute(route.wireModelId) + ? Math.min(requestedMaxTokens, 64_000) + : requestedMaxTokens, + }; + if (options?.temperature !== undefined) { + generationConfig.temperature = options.temperature; + } + const thinkingConfig = model.reasoning ? route.thinkingConfig : undefined; + if (thinkingConfig) generationConfig.thinkingConfig = thinkingConfig; + request.generationConfig = generationConfig; + + const agentId = state?.agentId ?? randomUUID(); + const trajectoryId = state?.trajectoryId ?? randomUUID(); + const sessionId = state?.sessionId ?? deriveSessionId(context); + const stepIndex = state ? ++state.stepIndex : 2; + request.sessionId = sessionId; + request.labels = { + ...(state?.lastExecutionId + ? { last_execution_id: state.lastExecutionId } + : {}), + last_step_index: String(stepIndex - 1), + trajectory_id: trajectoryId, + used_claude: String(isClaudeRoute(model.id)), + used_claude_conservative: String(isClaudeRoute(model.id)), + }; + + return { + project: projectId, + requestId: `agent/${agentId}/${Date.now()}/${trajectoryId}/${stepIndex}`, + model: route.wireModelId, + userAgent: "antigravity", + requestType: "agent", + request, + }; +} + +// --------------------------------------------------------------------------- +// SSE decoding +// --------------------------------------------------------------------------- + +interface CcaPart { + text?: string; + thought?: boolean; + thoughtSignature?: string; + functionCall?: { + id?: string; + name?: string; + args?: Record; + }; +} + +interface CcaChunk { + response?: { + responseId?: string; + candidates?: { + content?: { parts?: CcaPart[] }; + finishReason?: string; + }[]; + promptFeedback?: { + blockReason?: string; + blockReasonMessage?: string; + }; + usageMetadata?: { + promptTokenCount?: number; + cachedContentTokenCount?: number; + candidatesTokenCount?: number; + thoughtsTokenCount?: number; + totalTokenCount?: number; + }; + }; + error?: { code?: number; message?: string; status?: string }; +} + +async function readErrorResponseBody( + body: ReadableStream | null, + signal: AbortSignal | undefined, + deadline: number, +): Promise { + if (!body) return ""; + const reader = body.getReader(); + const decoder = new TextDecoder(); + const chunks: string[] = []; + let bytesRead = 0; + let reachedEnd = false; + let truncated = false; + const timeoutError = () => + new Error("Timed out reading Cloud Code Assist error response body"); + try { + while (bytesRead < MAX_ERROR_BODY_BYTES) { + if (signal?.aborted) throw new Error("Request aborted"); + const remaining = deadline - Date.now(); + if (remaining <= 0) throw timeoutError(); + let timer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + const gates: Promise>[] = [ + reader.read(), + new Promise((_, reject) => { + timer = setTimeout( + () => reject(timeoutError()), + Math.max(1, remaining), + ); + }), + ]; + if (signal) { + gates.push( + new Promise((_, reject) => { + onAbort = () => reject(new Error("Request aborted")); + signal.addEventListener("abort", onAbort, { once: true }); + }), + ); + } + const result = await Promise.race(gates).finally(() => { + if (timer) clearTimeout(timer); + if (onAbort) signal?.removeEventListener("abort", onAbort); + }); + if (Date.now() >= deadline) throw timeoutError(); + if (result.done) { + reachedEnd = true; + break; + } + const remainingBytes = MAX_ERROR_BODY_BYTES - bytesRead; + const value = result.value.subarray(0, remainingBytes); + bytesRead += value.byteLength; + chunks.push(decoder.decode(value, { stream: true })); + if (value.byteLength < result.value.byteLength) { + truncated = true; + break; + } + if (bytesRead === MAX_ERROR_BODY_BYTES) truncated = true; + } + chunks.push(decoder.decode()); + return chunks.join("") + (truncated ? "… [truncated]" : ""); + } finally { + if (!reachedEnd) await reader.cancel().catch(() => {}); + reader.releaseLock(); + } +} + +async function* readSseChunks( + body: ReadableStream, + signal: AbortSignal | undefined, + firstEventDeadline: number, + requestDeadline: number | undefined, +): AsyncGenerator { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + let dataLines: string[] = []; + const flush = (): CcaChunk | undefined => { + if (dataLines.length === 0) return undefined; + const payload = dataLines.join("\n"); + dataLines = []; + if (payload === "[DONE]") return undefined; + return JSON.parse(payload) as CcaChunk; + }; + let sawEvent = false; + let reachedEnd = false; + const timeoutError = () => + new Error( + sawEvent + ? "Timed out waiting for the next SSE event" + : "Timed out waiting for the first SSE event", + ); + try { + while (true) { + if (signal?.aborted) throw new Error("Request aborted"); + const deadline = sawEvent ? requestDeadline : firstEventDeadline; + const remaining = + deadline === undefined ? undefined : deadline - Date.now(); + if (remaining !== undefined && remaining <= 0) throw timeoutError(); + let timer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + const read = reader.read(); + const gates: Promise>[] = [read]; + if (remaining !== undefined) { + gates.push( + new Promise((_, reject) => { + timer = setTimeout( + () => reject(timeoutError()), + Math.max(1, remaining), + ); + }), + ); + } + if (signal) { + gates.push( + new Promise((_, reject) => { + onAbort = () => reject(new Error("Request aborted")); + signal.addEventListener("abort", onAbort, { once: true }); + }), + ); + } + const result = await Promise.race(gates).finally(() => { + if (timer) clearTimeout(timer); + if (onAbort) signal?.removeEventListener("abort", onAbort); + }); + if (deadline !== undefined && Date.now() >= deadline) + throw timeoutError(); + const { done, value } = result; + if (done) { + reachedEnd = true; + break; + } + buffer += decoder.decode(value, { stream: true }); + let newline = buffer.indexOf("\n"); + while (newline !== -1) { + const line = buffer.slice(0, newline).replace(/\r$/, ""); + buffer = buffer.slice(newline + 1); + newline = buffer.indexOf("\n"); + if (line === "") { + const event = flush(); + if (event) { + sawEvent = true; + yield event; + } + } else if (line.startsWith("data:")) { + dataLines.push(line.slice(5).trimStart()); + } + // event:/id:/retry: lines carry no payload for this API. + } + } + buffer += decoder.decode(); + const tail = buffer.trim(); + if (tail.startsWith("data:")) dataLines.push(tail.slice(5).trimStart()); + const event = flush(); + if (event) yield event; + } finally { + if (!reachedEnd) await reader.cancel().catch(() => {}); + reader.releaseLock(); + } +} + +async function* prependChunks( + initial: readonly CcaChunk[], + rest: AsyncGenerator, +): AsyncGenerator { + yield* initial; + yield* rest; +} + +function transientStatus(status: number | undefined): boolean { + return ( + status === 408 || status === 429 || (status !== undefined && status >= 500) + ); +} + +class EndpointAttemptError extends Error { + readonly retryable: boolean; + + constructor(message: string, retryable: boolean) { + super(message); + this.name = "EndpointAttemptError"; + this.retryable = retryable; + } +} + +async function preflightChunks( + chunks: AsyncGenerator, +): Promise { + const initial: CcaChunk[] = []; + while (true) { + const next = await chunks.next(); + if (next.done) { + throw new EndpointAttemptError( + "Cloud Code Assist stream ended before returning content", + true, + ); + } + const chunk = next.value; + initial.push(chunk); + if (chunk.error) { + const code = chunk.error.code; + const detail = + chunk.error.message || chunk.error.status || "unknown error"; + throw new EndpointAttemptError( + `Cloud Code Assist stream error: ${detail}`, + transientStatus(code), + ); + } + const data = chunk.response; + if (!data) continue; + if (!data.candidates?.length && data.promptFeedback?.blockReason) { + const detail = data.promptFeedback.blockReasonMessage; + throw new EndpointAttemptError( + `Request blocked by Google (${data.promptFeedback.blockReason})` + + (detail ? `: ${detail}` : ""), + false, + ); + } + const candidate = data.candidates?.[0]; + if ( + (candidate?.content?.parts ?? []).some( + (part) => + Boolean(part.functionCall) || + (Boolean(part.text?.trim()) && part.thought !== true), + ) + ) { + return initial; + } + if (candidate?.finishReason) { + throw new EndpointAttemptError( + "Cloud Code Assist returned an empty response", + true, + ); + } + } +} + +function mergeRequestHeaders( + defaults: Record, + overrides: AntigravityStreamOptions["headers"], +): Record { + const headers = new Headers(defaults); + for (const [name, value] of Object.entries(overrides ?? {})) { + if (value === null) headers.delete(name); + else headers.set(name, value); + } + return Object.fromEntries(headers.entries()); +} + +// --------------------------------------------------------------------------- +// Stream entry point +// --------------------------------------------------------------------------- + +function emptyUsage(): AssistantMessage["usage"] { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +let toolCallCounter = 0; + +export function streamAntigravity( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): AssistantMessageEventStream { + const requestOptions = options as AntigravityStreamOptions | undefined; + const stream = createAssistantMessageEventStream(); + + (async () => { + const output: AssistantMessage = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: emptyUsage(), + stopReason: "pending", + timestamp: Date.now(), + }; + + const fail = (error: unknown) => { + output.stopReason = options?.signal?.aborted ? "aborted" : "error"; + output.errorMessage = + error instanceof Error ? error.message : String(error); + stream.push({ + type: "error", + reason: output.stopReason, + error: output, + }); + stream.end(); + }; + + try { + const { token, projectId } = decodeApiKey(requestOptions?.apiKey ?? ""); + if (!token) { + throw new Error( + "No Antigravity access token — run /login google-antigravity", + ); + } + if (!projectId) { + throw new Error( + "No Cloud Code Assist project id in credentials — " + + "re-run /login google-antigravity", + ); + } + + const fetcher = requestOptions?.fetch ?? fetch; + await ensureAntigravityVersion(requestOptions?.signal, fetcher); + const providerState = getSessionState(requestOptions, context); + let payload: unknown = buildRequestBody( + model, + context, + requestOptions, + projectId, + providerState, + ); + const replacement = await requestOptions?.onPayload?.(payload, model); + if (replacement !== undefined) payload = replacement; + const body = JSON.stringify(payload); + const headers = mergeRequestHeaders( + { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + Accept: "text/event-stream", + "User-Agent": getAntigravityUserAgent(), + ...(isClaudeRoute(model.id) && model.reasoning + ? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER } + : {}), + }, + requestOptions?.headers, + ); + const defaultFirstEventTimeout = model.id.includes("flash") + ? FLASH_FIRST_EVENT_TIMEOUT_MS + : DEFAULT_FIRST_EVENT_TIMEOUT_MS; + const requestTimeout = + requestOptions?.timeoutMs && requestOptions.timeoutMs > 0 + ? requestOptions.timeoutMs + : undefined; + const firstEventTimeout = + requestTimeout !== undefined + ? Math.min(requestTimeout, defaultFirstEventTimeout) + : defaultFirstEventTimeout; + + let chunks: AsyncGenerator | undefined; + let lastError: Error | undefined; + for (const endpoint of ENDPOINTS) { + const isLast = endpoint === ENDPOINTS[ENDPOINTS.length - 1]; + const attemptStartedAt = Date.now(); + const firstEventDeadline = attemptStartedAt + firstEventTimeout; + const requestDeadline = + requestTimeout === undefined + ? undefined + : attemptStartedAt + requestTimeout; + const attemptAbort = new AbortController(); + const attemptSignal = requestOptions?.signal + ? AbortSignal.any([requestOptions.signal, attemptAbort.signal]) + : attemptAbort.signal; + let attempt: Response; + let headersTimer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + try { + const pending = fetcher( + `${endpoint}/v1internal:streamGenerateContent?alt=sse`, + { method: "POST", headers, body, signal: attemptSignal }, + ); + const gates: Promise[] = [pending]; + gates.push( + new Promise((_, reject) => { + headersTimer = setTimeout( + () => { + const error = new Error( + "Timed out waiting for Cloud Code Assist response headers", + ); + reject(error); + attemptAbort.abort(error); + }, + Math.max(1, firstEventDeadline - Date.now()), + ); + }), + ); + if (requestOptions?.signal) { + gates.push( + new Promise((_, reject) => { + onAbort = () => { + const error = new Error("Request aborted"); + reject(error); + attemptAbort.abort(error); + }; + requestOptions.signal?.addEventListener("abort", onAbort, { + once: true, + }); + }), + ); + } + attempt = await Promise.race(gates); + } catch (error) { + // Network/transport failure: fail over before any bytes stream. + if (requestOptions?.signal?.aborted) throw error; + lastError = error instanceof Error ? error : new Error(String(error)); + if (isLast) throw lastError; + continue; + } finally { + if (headersTimer) clearTimeout(headersTimer); + if (onAbort) { + requestOptions?.signal?.removeEventListener("abort", onAbort); + } + } + await requestOptions?.onResponse?.( + { + status: attempt.status, + headers: Object.fromEntries(attempt.headers.entries()), + }, + model, + ); + if (!attempt.ok) { + let errorText: string; + try { + errorText = await readErrorResponseBody( + attempt.body, + attemptSignal, + requestDeadline ?? firstEventDeadline, + ); + } catch (error) { + attemptAbort.abort(error); + if (requestOptions?.signal?.aborted) throw error; + errorText = error instanceof Error ? error.message : String(error); + } + lastError = new Error( + `Cloud Code Assist API error (${attempt.status}): ${errorText}`, + ); + if (!transientStatus(attempt.status) || isLast) throw lastError; + continue; + } + if (!attempt.body) { + lastError = new Error( + "Cloud Code Assist returned an empty response body", + ); + if (isLast) throw lastError; + continue; + } + const candidateChunks = readSseChunks( + attempt.body, + attemptSignal, + firstEventDeadline, + requestDeadline, + ); + try { + const initial = await preflightChunks(candidateChunks); + chunks = prependChunks(initial, candidateChunks); + break; + } catch (error) { + await candidateChunks.return(undefined).catch(() => {}); + if (requestOptions?.signal?.aborted) throw error; + lastError = error instanceof Error ? error : new Error(String(error)); + if ( + isLast || + (error instanceof EndpointAttemptError && !error.retryable) + ) { + throw lastError; + } + } + } + if (!chunks) throw lastError ?? new Error("No endpoint reachable"); + + stream.push({ type: "start", partial: output }); + + let sawFinishReason = false; + let sawMeaningfulContent = false; + const contentIndex = () => output.content.length - 1; + const closeOpenBlock = () => { + const block = output.content[contentIndex()]; + if (!block) return; + if (block.type === "text") { + stream.push({ + type: "text_end", + contentIndex: contentIndex(), + content: block.text, + partial: output, + }); + } else if (block.type === "thinking") { + stream.push({ + type: "thinking_end", + contentIndex: contentIndex(), + content: block.thinking, + partial: output, + }); + } + }; + + let lastResponseId: string | undefined; + for await (const chunk of chunks) { + if (chunk.error) { + const detail = + chunk.error.message || chunk.error.status || "unknown error"; + throw new Error(`Cloud Code Assist stream error: ${detail}`); + } + const data = chunk.response; + if (!data) continue; + if (data.responseId) lastResponseId = data.responseId; + if (!data.candidates?.length && data.promptFeedback?.blockReason) { + const detail = data.promptFeedback.blockReasonMessage; + throw new Error( + `Request blocked by Google (${data.promptFeedback.blockReason})` + + (detail ? `: ${detail}` : ""), + ); + } + + const candidate = data.candidates?.[0]; + for (const part of candidate?.content?.parts ?? []) { + if (part.text !== undefined && part.text !== "") { + if (isThinkingPart(part)) { + const open = output.content[contentIndex()]; + if (open?.type !== "thinking") { + closeOpenBlock(); + output.content.push({ + type: "thinking", + thinking: "", + }); + stream.push({ + type: "thinking_start", + contentIndex: contentIndex(), + partial: output, + }); + } + const block = output.content[contentIndex()]; + if (block.type === "thinking") { + block.thinking += part.text; + block.thinkingSignature = retainThoughtSignature( + block.thinkingSignature, + part.thoughtSignature, + ); + stream.push({ + type: "thinking_delta", + contentIndex: contentIndex(), + delta: part.text, + partial: output, + }); + } + } else { + if (part.text.trim().length > 0) sawMeaningfulContent = true; + const open = output.content[contentIndex()]; + if (open?.type !== "text") { + closeOpenBlock(); + output.content.push({ type: "text", text: "" }); + stream.push({ + type: "text_start", + contentIndex: contentIndex(), + partial: output, + }); + } + const block = output.content[contentIndex()]; + if (block.type === "text") { + block.text += part.text; + block.textSignature = retainThoughtSignature( + block.textSignature, + part.thoughtSignature, + ); + stream.push({ + type: "text_delta", + contentIndex: contentIndex(), + delta: part.text, + partial: output, + }); + } + } + } else if ( + part.text === "" && + part.thoughtSignature && + !part.functionCall + ) { + const open = output.content[contentIndex()]; + if (open?.type === "thinking") { + open.thinkingSignature = retainThoughtSignature( + open.thinkingSignature, + part.thoughtSignature, + ); + } else if (open?.type === "text") { + open.textSignature = retainThoughtSignature( + open.textSignature, + part.thoughtSignature, + ); + } + } + + if (part.functionCall) { + sawMeaningfulContent = true; + closeOpenBlock(); + const call = part.functionCall; + const providedId = call.id; + const duplicated = + providedId !== undefined && + output.content.some( + (b) => b.type === "toolCall" && b.id === providedId, + ); + const toolCall: ToolCall = { + type: "toolCall", + id: + providedId && !duplicated + ? providedId + : `call_${call.name ?? "tool"}_${++toolCallCounter}`, + name: call.name ?? "", + arguments: call.args ?? {}, + ...(part.thoughtSignature + ? { thoughtSignature: part.thoughtSignature } + : {}), + }; + output.content.push(toolCall); + const index = contentIndex(); + stream.push({ + type: "toolcall_start", + contentIndex: index, + partial: output, + }); + stream.push({ + type: "toolcall_delta", + contentIndex: index, + delta: JSON.stringify(toolCall.arguments), + partial: output, + }); + stream.push({ + type: "toolcall_end", + contentIndex: index, + toolCall, + partial: output, + }); + } + } + + if (candidate?.finishReason) { + sawFinishReason = true; + const mapped = mapStopReasonString(candidate.finishReason); + const hasToolCalls = output.content.some( + (b) => b.type === "toolCall", + ); + if ((mapped === "stop" || mapped === "length") && hasToolCalls) { + output.stopReason = "toolUse"; + } else { + output.stopReason = mapped; + if (mapped === "error") { + output.errorMessage = `Generation failed with finish reason: ${candidate.finishReason}`; + } + } + } + + if (data.usageMetadata) { + const usage = data.usageMetadata; + const promptTokens = usage.promptTokenCount ?? 0; + const cacheReadTokens = usage.cachedContentTokenCount ?? 0; + const thinkingTokens = usage.thoughtsTokenCount ?? 0; + output.usage = { + input: Math.max(0, promptTokens - cacheReadTokens), + output: (usage.candidatesTokenCount ?? 0) + thinkingTokens, + cacheRead: cacheReadTokens, + cacheWrite: 0, + totalTokens: usage.totalTokenCount ?? 0, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }; + } + } + + closeOpenBlock(); + + if (requestOptions?.signal?.aborted) { + output.stopReason = "aborted"; + } else if (output.stopReason === "pending") { + if (!sawFinishReason) { + throw new Error( + "Cloud Code Assist stream ended without a finish reason " + + "(connection dropped or response truncated)", + ); + } + output.stopReason = "stop"; + } + if (!sawMeaningfulContent && output.stopReason === "stop") { + throw new Error("Cloud Code Assist API returned an empty response"); + } + if ( + providerState && + output.stopReason !== "error" && + output.stopReason !== "aborted" + ) { + providerState.lastExecutionId = lastResponseId; + } + + if (output.stopReason === "error" || output.stopReason === "aborted") { + stream.push({ + type: "error", + reason: output.stopReason, + error: output, + }); + } else { + stream.push({ + type: "done", + reason: output.stopReason, + message: output, + }); + } + stream.end(); + } catch (error) { + fail(error); + } + })(); + + return stream; +} diff --git a/extensions/ai-providers/antigravity/routing.ts b/extensions/ai-providers/antigravity/routing.ts new file mode 100644 index 00000000..bed0cf69 --- /dev/null +++ b/extensions/ai-providers/antigravity/routing.ts @@ -0,0 +1,340 @@ +import type { Model, ThinkingLevel } from "@earendil-works/pi-ai/compat"; + +export type AntigravityModelDefinition = Omit< + Model, + "api" | "provider" | "baseUrl" +> & { + /** Provider-private fields preserved by pi's extension model composer. */ + requestModelId?: string; + antigravityEffortRouting?: Partial>; +}; + +type Family = { + id: string; + name: string; + members: readonly string[]; + defaultWireId: string; + routes?: Partial>; + mode: "budget" | "google-level"; + budgets?: Partial>; + mandatory?: boolean; + retiredMembers?: readonly string[]; + preserveAbsentEffortRoutes?: boolean; +}; + +function thinkingPairs( + pairs: readonly (readonly [id: string, name: string])[], +): Family[] { + return pairs.map(([id, name]) => ({ + id, + name, + members: [id, `${id}-thinking`], + defaultWireId: id, + routes: { + minimal: `${id}-thinking`, + low: `${id}-thinking`, + medium: `${id}-thinking`, + high: `${id}-thinking`, + xhigh: `${id}-thinking`, + max: `${id}-thinking`, + }, + mode: "budget", + preserveAbsentEffortRoutes: true, + })); +} + +const FAMILIES: readonly Family[] = [ + { + id: "gemini-3.7-flash", + name: "Gemini 3.7 Flash", + members: [ + "gemini-3.7-flash-low", + "gemini-3.7-flash-medium", + "gemini-3.7-flash-high", + ], + defaultWireId: "gemini-3.7-flash-low", + routes: { + minimal: "gemini-3.7-flash-low", + low: "gemini-3.7-flash-low", + medium: "gemini-3.7-flash-medium", + high: "gemini-3.7-flash-high", + xhigh: "gemini-3.7-flash-high", + max: "gemini-3.7-flash-high", + }, + mode: "google-level", + mandatory: true, + }, + { + id: "gemini-3.6-flash", + name: "Gemini 3.6 Flash", + members: [ + "gemini-3.6-flash-low", + "gemini-3.6-flash-medium", + "gemini-3.6-flash-high", + "gemini-3.6-flash-tiered", + ], + defaultWireId: "gemini-3.6-flash-low", + routes: { + minimal: "gemini-3.6-flash-low", + low: "gemini-3.6-flash-low", + medium: "gemini-3.6-flash-medium", + high: "gemini-3.6-flash-high", + xhigh: "gemini-3.6-flash-high", + max: "gemini-3.6-flash-high", + }, + mode: "google-level", + mandatory: true, + }, + { + id: "gemini-3.5-flash", + name: "Gemini 3.5 Flash", + members: [ + "gemini-3.5-flash-extra-low", + "gemini-3.5-flash-low", + "gemini-3-flash-agent", + ], + defaultWireId: "gemini-3.5-flash-extra-low", + routes: { + minimal: "gemini-3.5-flash-extra-low", + low: "gemini-3.5-flash-extra-low", + medium: "gemini-3.5-flash-low", + high: "gemini-3-flash-agent", + xhigh: "gemini-3-flash-agent", + max: "gemini-3-flash-agent", + }, + mode: "budget", + budgets: { + minimal: 1_000, + low: 1_000, + medium: 4_000, + high: 10_000, + xhigh: 10_000, + max: 10_000, + }, + }, + { + id: "gemini-3.1-pro", + name: "Gemini 3.1 Pro", + members: [ + "gemini-3.1-pro-low", + "gemini-pro-agent", + // Discovery still publishes this deployment, but requests always fail. + "gemini-3.1-pro-high", + ], + defaultWireId: "gemini-3.1-pro-low", + routes: { + minimal: "gemini-3.1-pro-low", + low: "gemini-3.1-pro-low", + medium: "gemini-3.1-pro-low", + high: "gemini-pro-agent", + xhigh: "gemini-pro-agent", + max: "gemini-pro-agent", + }, + mode: "budget", + budgets: { + minimal: 1_001, + low: 1_001, + medium: 1_001, + high: 10_001, + xhigh: 10_001, + max: 10_001, + }, + retiredMembers: ["gemini-3.1-pro-high"], + }, + { + id: "gemini-3-pro", + name: "Gemini 3 Pro", + members: ["gemini-3-pro-low", "gemini-3-pro-high"], + defaultWireId: "gemini-3-pro-low", + routes: { + minimal: "gemini-3-pro-low", + low: "gemini-3-pro-low", + medium: "gemini-3-pro-low", + high: "gemini-3-pro-high", + xhigh: "gemini-3-pro-high", + max: "gemini-3-pro-high", + }, + mode: "google-level", + }, + { + id: "gpt-oss-120b", + name: "GPT-OSS 120B", + members: ["gpt-oss-120b-medium"], + defaultWireId: "gpt-oss-120b-medium", + mode: "budget", + }, + { + id: "claude-sonnet-4-6", + name: "Claude Sonnet 4.6", + members: ["claude-sonnet-4-6", "claude-sonnet-4-6-thinking"], + defaultWireId: "claude-sonnet-4-6", + mode: "budget", + retiredMembers: ["claude-sonnet-4-6-thinking"], + }, + { + id: "claude-opus-4-6", + name: "Claude Opus 4.6", + members: ["claude-opus-4-6-thinking", "claude-opus-4-6"], + defaultWireId: "claude-opus-4-6-thinking", + mode: "budget", + retiredMembers: ["claude-opus-4-6"], + }, + ...thinkingPairs([ + ["claude-sonnet-4-5", "Claude Sonnet 4.5"], + ["claude-opus-4-5", "Claude Opus 4.5"], + ["gemini-2.5-flash", "Gemini 2.5 Flash"], + ]), +]; + +const familyById = new Map(); +for (const family of FAMILIES) { + familyById.set(family.id, family); + for (const member of family.members) familyById.set(member, family); +} + +function clampLevel( + level: ThinkingLevel, +): "minimal" | "low" | "medium" | "high" { + if (level === "xhigh" || level === "max") return "high"; + return level; +} + +export function collapseAntigravityModels( + models: readonly T[], +): T[] { + const byId = new Map(models.map((model) => [model.id, model])); + const consumed = new Set(); + const collapsed: T[] = []; + + for (const model of models) { + if (consumed.has(model.id)) continue; + const family = familyById.get(model.id); + if (!family) { + collapsed.push(model); + continue; + } + if (collapsed.some((entry) => entry.id === family.id)) continue; + for (const member of family.members) consumed.add(member); + consumed.add(family.id); + const logical = byId.get(family.id); + const source = + logical ?? family.members.map((id) => byId.get(id)).find(Boolean); + if (!source) continue; + const retired = new Set(family.retiredMembers ?? []); + const liveWireIds = new Set( + family.members.filter((id) => byId.has(id) && !retired.has(id)), + ); + if (logical?.requestModelId && !retired.has(logical.requestModelId)) { + liveWireIds.add(logical.requestModelId); + } else if (logical && family.members.includes(family.id)) { + liveWireIds.add(family.id); + } + // A discovery response containing only a retired deployment is unusable; + // consume it without publishing a logical model that cannot be requested. + if (liveWireIds.size === 0) continue; + const requestModelId = liveWireIds.has(family.defaultWireId) + ? family.defaultWireId + : liveWireIds.values().next().value; + if (!requestModelId) continue; + const effortRouting: Partial> = {}; + for (const [effort, target] of Object.entries(family.routes ?? {}) as [ + ThinkingLevel, + string, + ][]) { + if ( + !retired.has(target) && + (liveWireIds.has(target) || family.preserveAbsentEffortRoutes) + ) { + effortRouting[effort] = target; + } + } + collapsed.push({ + ...source, + id: family.id, + name: family.name, + reasoning: true, + ...(requestModelId !== family.id ? { requestModelId } : {}), + ...(Object.keys(effortRouting).length > 0 + ? { antigravityEffortRouting: effortRouting } + : {}), + }); + } + + return collapsed; +} + +export function routeAntigravityModel( + modelId: string, + reasoning: ThinkingLevel | undefined, + thinkingBudgets: Partial> | undefined, + overrides?: Pick< + AntigravityModelDefinition, + "requestModelId" | "antigravityEffortRouting" + >, +): { wireModelId: string; thinkingConfig?: Record } { + const family = familyById.get(modelId); + if (!family) { + if (!reasoning) return { wireModelId: modelId }; + if (modelId.toLowerCase().includes("claude")) { + const level = clampLevel(reasoning); + const budget = + thinkingBudgets?.[level] ?? + { + minimal: 1_024, + low: 8_192, + medium: 16_384, + high: 32_768, + }[level]; + return { + wireModelId: modelId, + thinkingConfig: { includeThoughts: true, thinkingBudget: budget }, + }; + } + return { + wireModelId: modelId, + thinkingConfig: { + includeThoughts: true, + thinkingLevel: clampLevel(reasoning).toUpperCase(), + }, + }; + } + + const effective = reasoning ?? (family.mandatory ? "minimal" : undefined); + const discoveredRouting = overrides?.antigravityEffortRouting; + const wireModelId = effective + ? ((discoveredRouting + ? (discoveredRouting[effective] ?? overrides?.requestModelId) + : (family.routes?.[effective] ?? overrides?.requestModelId)) ?? + family.defaultWireId) + : (overrides?.requestModelId ?? family.defaultWireId); + if (!effective) { + return { + wireModelId, + thinkingConfig: { includeThoughts: false, thinkingBudget: 0 }, + }; + } + if (family.mode === "google-level") { + const level = clampLevel(effective); + return { + wireModelId, + thinkingConfig: { + includeThoughts: true, + thinkingLevel: level === "minimal" ? "LOW" : level.toUpperCase(), + }, + }; + } + const level = clampLevel(effective); + const budget = + thinkingBudgets?.[level] ?? + family.budgets?.[effective] ?? + { + minimal: 1_024, + low: 8_192, + medium: 16_384, + high: 32_768, + }[level]; + return { + wireModelId, + thinkingConfig: { includeThoughts: true, thinkingBudget: budget }, + }; +} diff --git a/extensions/ai-providers/antigravity/with-resolvers.d.ts b/extensions/ai-providers/antigravity/with-resolvers.d.ts new file mode 100644 index 00000000..377e1a17 --- /dev/null +++ b/extensions/ai-providers/antigravity/with-resolvers.d.ts @@ -0,0 +1,19 @@ +/** + * Promise.withResolvers type shim. + * + * The repo tsconfig targets ES2022 libs, but openpi runs on Node >=22.19 + * (package.json engines), where Promise.withResolvers is available since + * Node 22.0. This declaration matches lib.es2024.promise.withresolvers.d.ts. + */ + +declare global { + interface PromiseConstructor { + withResolvers(): { + promise: Promise; + resolve: (value: T | PromiseLike) => void; + reject: (reason?: unknown) => void; + }; + } +} + +export {}; diff --git a/extensions/ai-providers/cursor/constants.ts b/extensions/ai-providers/cursor/constants.ts new file mode 100644 index 00000000..d9e4a325 --- /dev/null +++ b/extensions/ai-providers/cursor/constants.ts @@ -0,0 +1,5 @@ +/** Wire constants shared by Cursor OAuth, discovery, and AgentService. */ +export const CURSOR_API_URL = "https://api2.cursor.sh"; +export const CURSOR_CLIENT_VERSION = "cli-2026.07.23-e383d2b"; +export const CURSOR_RUN_PATH = "/agent.v1.AgentService/Run"; +export const CURSOR_MODELS_PATH = "/agent.v1.AgentService/GetUsableModels"; diff --git a/extensions/ai-providers/cursor/credentials.ts b/extensions/ai-providers/cursor/credentials.ts new file mode 100644 index 00000000..5b1b7b98 --- /dev/null +++ b/extensions/ai-providers/cursor/credentials.ts @@ -0,0 +1,14 @@ +import type { OAuthCredentials } from "@earendil-works/pi-ai/compat"; + +/** Cursor stores the short-lived access JWT and its refresh token together. */ +export type CursorCredentials = OAuthCredentials; + +/** pi passes the return value of this function to `streamSimple.apiKey`. */ +export function getCursorApiKey(credentials: CursorCredentials): string { + return credentials.access; +} + +/** Accept a bare access token for hand-written or older auth.json entries. */ +export function decodeCursorApiKey(value: string | undefined): string { + return value?.trim() ?? ""; +} diff --git a/extensions/ai-providers/cursor/discovery.ts b/extensions/ai-providers/cursor/discovery.ts new file mode 100644 index 00000000..e8d1a3b9 --- /dev/null +++ b/extensions/ai-providers/cursor/discovery.ts @@ -0,0 +1,291 @@ +import { randomUUID } from "node:crypto"; +import * as http2 from "node:http2"; +import type { Api, Model, RefreshModelsContext } from "@earendil-works/pi-ai"; +import { + CURSOR_API_URL, + CURSOR_CLIENT_VERSION, + CURSOR_MODELS_PATH, +} from "./constants.ts"; +import type { CursorModelDefinition } from "./models.ts"; +import { + GetUsableModelsRequestSchema, + GetUsableModelsResponseSchema, + type ModelDetails, +} from "./proto.ts"; +import { create, fromBinary, toBinary } from "./protobuf.ts"; +import { connectCursorHttp2 } from "./proxy.ts"; + +export interface CursorModelDiscoveryOptions { + apiKey: string; + baseUrl?: string; + clientVersion?: string; + timeoutMs?: number; + signal?: AbortSignal; + customModelIds?: string[]; +} + +const FALLBACK_CONTEXT_WINDOW = 200_000; +const DEFAULT_MAX_TOKENS = 64_000; +const MAX_DISCOVERY_RESPONSE_BYTES = 16 * 1024 * 1024; +const ONE_MILLION_CONTEXT_WINDOW = 1_000_000; + +/** + * GetUsableModels has no numeric context-window field. Recover 1M only from + * signals Cursor does send; use 200k as the conservative unknown-model fallback. + */ +function resolveContextWindow(details: ModelDetails, id: string): number { + const labels = [ + id, + details.displayName, + details.displayNameShort, + details.displayModelId, + ...details.aliases, + ].join(" "); + if (/\b1m\b/i.test(labels)) return ONE_MILLION_CONTEXT_WINDOW; + if (details.maxMode && /claude|gemini|gpt-5\.6-sol/i.test(id)) { + return ONE_MILLION_CONTEXT_WINDOW; + } + if (isNativeOneMillionModel(id)) return ONE_MILLION_CONTEXT_WINDOW; + return FALLBACK_CONTEXT_WINDOW; +} + +/** Cursor serves these coding families with a native, unlabeled 1M window. */ +function isNativeOneMillionModel(id: string): boolean { + const bareId = id.split("/").at(-1)?.toLowerCase() ?? id.toLowerCase(); + if (/^(?:kimi-)?k3$/.test(bareId)) return true; + + const glm = + /^glm-(\d{1,2})(?:\.(\d+))?(v)?(?:-(air|turbo|flashx|flash|preview))?$/.exec( + bareId, + ); + if (!glm || glm[3]) return false; + const variant = glm[4]; + if (variant && variant !== "air" && variant !== "turbo") return false; + const major = Number(glm[1]); + const minor = Number(glm[2] ?? 0); + return major > 5 || (major === 5 && minor >= 2); +} + +/** + * Fetch account-specific models over Cursor's HTTP/2 Connect endpoint. + * `null` means transport/protocol failure; an empty successful response is + * deliberately represented as `[]` so callers can choose their fallback policy. + */ +export async function fetchCursorUsableModels( + options: CursorModelDiscoveryOptions, +): Promise { + const token = options.apiKey.trim(); + if (!token || options.signal?.aborted) return null; + const baseUrl = (options.baseUrl ?? CURSOR_API_URL).replace(/\/+$/, ""); + const request = create(GetUsableModelsRequestSchema, { + customModelIds: normalizeModelIds(options.customModelIds), + }); + const requestBytes = toBinary(GetUsableModelsRequestSchema, request); + const responseBytes = await requestHttp2( + baseUrl, + requestBytes, + token, + options, + ); + if (!responseBytes) return null; + + const payload = decodeUnaryPayload(responseBytes); + if (!payload) return null; + let response; + try { + response = fromBinary(GetUsableModelsResponseSchema, payload); + } catch { + return null; + } + + const models: CursorModelDefinition[] = []; + const seen = new Set(); + for (const details of response.models) { + const normalized = normalizeCursorModel(details, baseUrl); + if (!normalized || seen.has(normalized.id)) continue; + seen.add(normalized.id); + models.push(normalized); + } + models.sort((left, right) => left.id.localeCompare(right.id)); + return models; +} + +/** Provider extension hook: discovery always uses the credential passed by pi. */ +export async function fetchCursorModels( + context: RefreshModelsContext, +): Promise { + if (!context.allowNetwork) return []; + context.signal.throwIfAborted(); + const credential = context.credential; + const apiKey = + credential?.type === "oauth" + ? credential.access + : credential?.type === "api_key" + ? credential.key + : undefined; + if (!apiKey) return []; + const discovered = await fetchCursorUsableModels({ + apiKey, + signal: context.signal, + }); + context.signal.throwIfAborted(); + if (discovered === null) { + throw new Error("Cursor model discovery failed"); + } + return discovered; +} + +function normalizeModelIds(ids: readonly string[] | undefined): string[] { + if (!ids) return []; + const result = new Set(); + for (const id of ids) { + if (typeof id !== "string") continue; + const value = id.trim(); + if (value) result.add(value); + } + return [...result]; +} + +function normalizeCursorModel( + details: ModelDetails, + baseUrl: string, +): CursorModelDefinition | null { + const id = details.modelId.trim(); + if (!id) return null; + const name = + [ + details.displayName, + details.displayNameShort, + details.displayModelId, + ...details.aliases, + ] + .map((value) => value.trim()) + .find(Boolean) ?? id; + const multimodal = supportsCursorImages(id); + return { + id, + name, + api: "cursor-agent", + provider: "cursor", + baseUrl, + reasoning: Boolean(details.thinkingDetails) || /^cursor-grok-\d/i.test(id), + input: multimodal ? ["text", "image"] : ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: resolveContextWindow(details, id), + maxTokens: DEFAULT_MAX_TOKENS, + ...(details.maxMode ? { cursorMaxMode: true } : {}), + }; +} + +/** GetUsableModels omits modality metadata for Cursor-native image families. */ +function supportsCursorImages(id: string): boolean { + const lower = id.toLowerCase(); + if (/claude|gemini|gpt-|codex/.test(lower)) return true; + const bareId = lower.split("/").at(-1) ?? lower; + return ( + /^(?:kimi-)?k3(?:$|[._:-])/.test(bareId) || + /^cursor-grok-4(?:$|[._:-])/.test(bareId) || + /^(?:cursor-)?composer-2\.5(?:$|[._:-])/.test(bareId) + ); +} + +function buildHeaders( + apiKey: string, + clientVersion: string, +): Record { + return { + "content-type": "application/proto", + te: "trailers", + authorization: `Bearer ${apiKey}`, + "x-ghost-mode": "true", + "x-cursor-client-version": clientVersion, + "x-cursor-client-type": "cli", + "x-request-id": randomUUID(), + }; +} + +async function requestHttp2( + baseUrl: string, + body: Uint8Array, + apiKey: string, + options: CursorModelDiscoveryOptions, +): Promise { + const timeoutMs = options.timeoutMs ?? 5_000; + let client: http2.ClientHttp2Session; + try { + client = await connectCursorHttp2(baseUrl, { + signal: options.signal, + timeoutMs, + }); + } catch { + return null; + } + const { promise, resolve } = Promise.withResolvers(); + let settled = false; + let timer: ReturnType | undefined; + let removeAbortListener: (() => void) | undefined; + const finish = (result: Uint8Array | null, destroy = false) => { + if (settled) return; + settled = true; + if (timer) clearTimeout(timer); + removeAbortListener?.(); + if (destroy) client.destroy(); + else client.close(); + resolve(result); + }; + timer = setTimeout(() => finish(null, true), timeoutMs); + client.once("error", () => finish(null, true)); + const req = client.request({ + ":method": "POST", + ":path": CURSOR_MODELS_PATH, + ...buildHeaders(apiKey, options.clientVersion ?? CURSOR_CLIENT_VERSION), + }); + const chunks: Buffer[] = []; + let responseBytes = 0; + req.on("response", (headers) => { + const status = Number(headers[":status"] ?? 0); + if (status < 200 || status >= 300) finish(null, true); + }); + req.on("data", (chunk: Buffer) => { + responseBytes += chunk.length; + if (responseBytes > MAX_DISCOVERY_RESPONSE_BYTES) { + req.close(http2.constants.NGHTTP2_CANCEL); + finish(null, true); + return; + } + chunks.push(chunk); + }); + req.on("end", () => finish(new Uint8Array(Buffer.concat(chunks)))); + req.once("error", () => finish(null, true)); + const onAbort = () => { + options.signal?.removeEventListener("abort", onAbort); + req.close(http2.constants.NGHTTP2_CANCEL); + finish(null, true); + }; + if (options.signal) { + if (options.signal.aborted) onAbort(); + else { + options.signal.addEventListener("abort", onAbort, { once: true }); + removeAbortListener = () => + options.signal?.removeEventListener("abort", onAbort); + } + } + req.end(Buffer.from(body)); + return promise; +} + +/** Decode the first uncompressed Connect data frame, or accept raw unary proto. */ +export function decodeUnaryPayload(body: Uint8Array): Uint8Array | null { + if (body.length === 0) return body; + if (body.length < 5) return body; + const flags = body[0]!; + const size = new DataView(body.buffer, body.byteOffset, 5).getUint32(1); + const end = 5 + size; + // Unary Connect responses are normally one data frame followed by an + // optional end-stream frame. If the prefix is not a valid uncompressed frame, + // accept the raw protobuf response used by older Cursor deployments. + if (flags > 3 || end > body.length || (flags & 1) !== 0) return body; + const data = body.subarray(5, end); + if ((flags & 2) !== 0) return null; + return data; +} diff --git a/extensions/ai-providers/cursor/input-images.ts b/extensions/ai-providers/cursor/input-images.ts new file mode 100644 index 00000000..d67d7dbc --- /dev/null +++ b/extensions/ai-providers/cursor/input-images.ts @@ -0,0 +1,106 @@ +import { readFile, stat } from "node:fs/promises"; +import { basename, isAbsolute } from "node:path"; +import type { + ExtensionContext, + InputEvent, + InputEventResult, +} from "@earendil-works/pi-coding-agent"; +import type { ImageContent } from "@earendil-works/pi-ai/compat"; + +const MAX_IMAGE_BYTES = 10 * 1024 * 1024; + +function leadingPath(text: string): { path: string; end: number } | undefined { + const match = /^\s*(?:"([^"\n]+)"|'([^'\n]+)'|(\S+))/.exec(text); + if (!match) return undefined; + const value = match[1] ?? match[2] ?? match[3]; + if (!value || !isAbsolute(value)) return undefined; + if (!/\.(?:png|jpe?g|gif|webp)$/i.test(value)) return undefined; + return { path: value, end: match[0].length }; +} + +function detectImageMimeType( + bytes: Uint8Array, +): ImageContent["mimeType"] | undefined { + if ( + bytes.length >= 8 && + bytes[0] === 0x89 && + bytes[1] === 0x50 && + bytes[2] === 0x4e && + bytes[3] === 0x47 && + bytes[4] === 0x0d && + bytes[5] === 0x0a && + bytes[6] === 0x1a && + bytes[7] === 0x0a + ) { + return "image/png"; + } + if ( + bytes.length >= 3 && + bytes[0] === 0xff && + bytes[1] === 0xd8 && + bytes[2] === 0xff + ) { + return "image/jpeg"; + } + if (bytes.length >= 6) { + const signature = Buffer.from(bytes.subarray(0, 6)).toString("ascii"); + if (signature === "GIF87a" || signature === "GIF89a") return "image/gif"; + } + if ( + bytes.length >= 12 && + Buffer.from(bytes.subarray(0, 4)).toString("ascii") === "RIFF" && + Buffer.from(bytes.subarray(8, 12)).toString("ascii") === "WEBP" + ) { + return "image/webp"; + } + return undefined; +} + +/** + * Pi's TUI represents a clipboard image as a leading local path. Cursor's + * chat-only provider cannot ask a native read-file tool to resolve that path, + * so convert an explicit leading image path into the same ImageContent shape + * used by CLI/RPC attachments before the agent turn starts. + */ +export async function transformCursorImageInput( + event: InputEvent, + ctx: ExtensionContext, +): Promise { + if ( + event.source !== "interactive" || + ctx.model?.provider !== "cursor" || + (event.images?.length ?? 0) > 0 + ) { + return { action: "continue" }; + } + + const candidate = leadingPath(event.text); + if (!candidate) return { action: "continue" }; + + try { + const file = await stat(candidate.path); + if (!file.isFile() || file.size === 0 || file.size > MAX_IMAGE_BYTES) { + return { action: "continue" }; + } + const bytes = await readFile(candidate.path); + if (bytes.length > MAX_IMAGE_BYTES) return { action: "continue" }; + const mimeType = detectImageMimeType(bytes); + if (!mimeType) return { action: "continue" }; + + const question = event.text.slice(candidate.end).trimStart(); + const attachment = `Attached image: ${JSON.stringify(basename(candidate.path))}`; + return { + action: "transform", + text: question ? `${attachment}\n${question}` : attachment, + images: [ + { + type: "image", + data: bytes.toString("base64"), + mimeType, + }, + ], + }; + } catch { + return { action: "continue" }; + } +} diff --git a/extensions/ai-providers/cursor/models.ts b/extensions/ai-providers/cursor/models.ts new file mode 100644 index 00000000..aebd9800 --- /dev/null +++ b/extensions/ai-providers/cursor/models.ts @@ -0,0 +1,45 @@ +import type { Model } from "@earendil-works/pi-ai/compat"; +import { CURSOR_API_URL } from "./constants.ts"; + +export type CursorModelDefinition = { + id: string; + name: string; + api: "cursor-agent"; + provider: "cursor"; + baseUrl: string; + reasoning: boolean; + input: ("text" | "image")[]; + cost: Model["cost"]; + contextWindow: number; + maxTokens: number; + cursorMaxMode?: boolean; +}; + +const ZERO_COST: CursorModelDefinition["cost"] = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, +}; + +/** + * Minimal offline catalog. Cursor's usable-model list is account-specific; + * `default` is the server-side Auto route and remains valid when discovery is + * unavailable or the account has no models response. + */ +export const CURSOR_MODELS: CursorModelDefinition[] = [ + { + id: "default", + name: "Auto", + api: "cursor-agent", + provider: "cursor", + baseUrl: CURSOR_API_URL, + reasoning: false, + input: ["text", "image"], + cost: ZERO_COST, + contextWindow: 200_000, + maxTokens: 64_000, + }, +]; + +export const CURSOR_STATIC_MODELS = CURSOR_MODELS; diff --git a/extensions/ai-providers/cursor/oauth.ts b/extensions/ai-providers/cursor/oauth.ts new file mode 100644 index 00000000..47b9b9a5 --- /dev/null +++ b/extensions/ai-providers/cursor/oauth.ts @@ -0,0 +1,263 @@ +import { createHash, randomBytes, randomUUID } from "node:crypto"; +import type { + OAuthCredentials, + OAuthLoginCallbacks, +} from "@earendil-works/pi-ai/compat"; +import type { CursorCredentials } from "./credentials.ts"; + +const CURSOR_LOGIN_URL = "https://cursor.com/loginDeepControl"; +const CURSOR_POLL_URL = "https://api2.cursor.sh/auth/poll"; +const CURSOR_REFRESH_URL = "https://api2.cursor.sh/auth/exchange_user_api_key"; + +const POLL_MAX_ATTEMPTS = 150; +const POLL_BASE_DELAY_MS = 1_000; +const POLL_MAX_DELAY_MS = 10_000; +const POLL_BACKOFF_MULTIPLIER = 1.2; +const LOGIN_TIMEOUT_MS = 5 * 60 * 1_000; +const REQUEST_TIMEOUT_MS = 30 * 1_000; +const EXPIRY_MARGIN_MS = 5 * 60 * 1_000; + +export interface CursorAuthParams { + verifier: string; + challenge: string; + uuid: string; + loginUrl: string; +} + +export interface CursorPollOptions { + /** Test/bridge override; production uses Cursor's auth endpoint. */ + pollUrl?: string; + maxAttempts?: number; + baseDelayMs?: number; + maxDelayMs?: number; + backoffMultiplier?: number; +} + +function throwIfAborted(signal: AbortSignal | undefined): void { + if (signal?.aborted) { + throw new Error("Cursor authentication cancelled"); + } +} + +function wait(ms: number, signal: AbortSignal | undefined): Promise { + throwIfAborted(signal); + const { promise, resolve, reject } = Promise.withResolvers(); + const timer = setTimeout( + () => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, + Math.max(0, ms), + ); + const onAbort = () => { + clearTimeout(timer); + reject(new Error("Cursor authentication cancelled")); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + return promise; +} + +async function fetchWithTimeout( + url: string, + init: RequestInit, + signal: AbortSignal | undefined, +): Promise { + throwIfAborted(signal); + const timeout = AbortSignal.timeout(REQUEST_TIMEOUT_MS); + const combined = signal ? AbortSignal.any([signal, timeout]) : timeout; + try { + return await fetch(url, { ...init, signal: combined }); + } catch (error) { + throwIfAborted(signal); + if (timeout.aborted) { + throw new Error(`Cursor request timed out after ${REQUEST_TIMEOUT_MS}ms`); + } + throw error; + } +} + +export async function generateCursorAuthParams(): Promise { + const verifierBytes = randomBytes(96); + const verifier = verifierBytes.toString("base64url"); + const challenge = createHash("sha256").update(verifier).digest("base64url"); + const uuid = randomUUID(); + const query = new URLSearchParams({ + challenge, + uuid, + mode: "login", + redirectTarget: "cli", + }); + return { + verifier, + challenge, + uuid, + loginUrl: `${CURSOR_LOGIN_URL}?${query.toString()}`, + }; +} + +/** Poll Cursor's loginDeepControl handoff until the browser finishes. */ +export async function pollCursorAuth( + uuid: string, + verifier: string, + signal?: AbortSignal, + options?: CursorPollOptions, +): Promise<{ accessToken: string; refreshToken: string }> { + const pollUrl = options?.pollUrl ?? CURSOR_POLL_URL; + const maxAttempts = options?.maxAttempts ?? POLL_MAX_ATTEMPTS; + const maxDelay = options?.maxDelayMs ?? POLL_MAX_DELAY_MS; + const multiplier = options?.backoffMultiplier ?? POLL_BACKOFF_MULTIPLIER; + let delay = options?.baseDelayMs ?? POLL_BASE_DELAY_MS; + let consecutiveErrors = 0; + + for (let attempt = 0; attempt < maxAttempts; attempt++) { + await wait(delay, signal); + const url = new URL(pollUrl); + url.searchParams.set("uuid", uuid); + url.searchParams.set("verifier", verifier); + try { + const response = await fetchWithTimeout(url.toString(), {}, signal); + if (response.status === 404) { + consecutiveErrors = 0; + delay = Math.min(delay * multiplier, maxDelay); + continue; + } + if (!response.ok) { + throw new Error(`Cursor auth poll failed: HTTP ${response.status}`); + } + const payload: unknown = await response.json(); + if (!isTokenPayload(payload) || !payload.refreshToken) { + throw new Error("Cursor auth poll returned an invalid token payload"); + } + return { + accessToken: payload.accessToken, + refreshToken: payload.refreshToken, + }; + } catch (error) { + if (signal?.aborted) throw error; + consecutiveErrors++; + delay = Math.min(delay * multiplier, maxDelay); + if (consecutiveErrors >= 3) { + throw new Error( + "Too many consecutive errors during Cursor authentication polling", + ); + } + } + } + throw new Error("Cursor authentication polling timed out"); +} + +function isTokenPayload( + value: unknown, +): value is { accessToken: string; refreshToken?: string } { + if (value === null || typeof value !== "object") return false; + const record = value as Record; + return ( + typeof record.accessToken === "string" && + record.accessToken.length > 0 && + (record.refreshToken === undefined || + typeof record.refreshToken === "string") + ); +} + +export async function loginCursor( + callbacks: OAuthLoginCallbacks, +): Promise { + const auth = await generateCursorAuthParams(); + callbacks.onAuth({ + url: auth.loginUrl, + instructions: "Complete the Cursor sign-in in your browser.", + }); + callbacks.onProgress?.("Waiting for browser authentication..."); + + const timeout = AbortSignal.timeout(LOGIN_TIMEOUT_MS); + const signal = callbacks.signal + ? AbortSignal.any([callbacks.signal, timeout]) + : timeout; + const tokens = await pollCursorAuth(auth.uuid, auth.verifier, signal); + return { + access: tokens.accessToken, + refresh: tokens.refreshToken, + expires: getCursorTokenExpiry(tokens.accessToken), + }; +} + +export async function refreshCursorToken( + credentials: OAuthCredentials, + signal: AbortSignal, +): Promise { + const response = await fetchWithTimeout( + CURSOR_REFRESH_URL, + { + method: "POST", + headers: { + Authorization: `Bearer ${credentials.refresh}`, + "Content-Type": "application/json", + }, + body: "{}", + }, + signal, + ); + if (!response.ok) { + throw new Error( + `Cursor token refresh failed: ${response.status} ${await response.text()}`, + ); + } + const payload: unknown = await response.json(); + if (!isTokenPayload(payload)) { + throw new Error("Cursor token refresh returned an invalid token payload"); + } + return { + access: payload.accessToken, + refresh: payload.refreshToken || credentials.refresh, + expires: getCursorTokenExpiry(payload.accessToken), + }; +} + +function decodeCursorJwtPayload(token: string): unknown | undefined { + const parts = token.split("."); + if (parts.length !== 3 || !parts[1]) return undefined; + const normalized = parts[1].replace(/-/g, "+").replace(/_/g, "/"); + const padded = normalized.padEnd(Math.ceil(normalized.length / 4) * 4, "="); + return JSON.parse(Buffer.from(padded, "base64").toString("utf8")); +} + +/** Returns an expiry with the same five-minute safety margin as the Cursor client. */ +export function getCursorTokenExpiry(token: string): number { + try { + const payload = decodeCursorJwtPayload(token); + if ( + payload !== null && + typeof payload === "object" && + "exp" in payload && + typeof payload.exp === "number" && + Number.isFinite(payload.exp) + ) { + return payload.exp * 1_000 - EXPIRY_MARGIN_MS; + } + } catch { + // Cursor occasionally returns opaque access tokens; use a conservative hour. + } + return Date.now() + 60 * 60 * 1_000; +} + +export const getTokenExpiry = getCursorTokenExpiry; + +export function isCursorTokenExpiringSoon( + token: string, + thresholdSeconds = 300, +): boolean { + try { + const payload = decodeCursorJwtPayload(token); + if ( + payload === null || + typeof payload !== "object" || + !("exp" in payload) || + typeof payload.exp !== "number" + ) { + return true; + } + return payload.exp - Math.floor(Date.now() / 1_000) < thresholdSeconds; + } catch { + return true; + } +} diff --git a/extensions/ai-providers/cursor/proto.ts b/extensions/ai-providers/cursor/proto.ts new file mode 100644 index 00000000..a4e71b74 --- /dev/null +++ b/extensions/ai-providers/cursor/proto.ts @@ -0,0 +1,1064 @@ +/** + * Small Cursor AgentService protobuf surface. + * + * The field numbers and message names are vendored from + * oh-my-pi@eab72e88e4, packages/catalog/src/discovery/cursor-proto.ts + * (MIT). Only the chat, image, model-discovery, and exec-rejection messages + * used by this chat-only provider are retained. Unknown fields are skipped by + * the local protobuf codec so newer Cursor messages remain forward-compatible. + */ + +import { type MessageCodec, type ProtoMessage, pb } from "./protobuf.ts"; + +export interface AgentClientMessage extends ProtoMessage { + message: + | { case: undefined; value?: undefined } + | { case: "runRequest"; value: AgentRunRequest } + | { case: "execClientMessage"; value: ExecClientMessage } + | { case: "execClientControlMessage"; value: ExecClientControlMessage } + | { case: "kvClientMessage"; value: KvClientMessage } + | { case: "clientHeartbeat"; value: ClientHeartbeat }; +} + +export const AgentClientMessageSchema: MessageCodec = + pb("agent.v1.AgentClientMessage", [ + { + kind: "oneof", + name: "message", + variants: [ + { + no: 1, + name: "runRequest", + kind: "message", + T: () => AgentRunRequestSchema, + }, + { + no: 2, + name: "execClientMessage", + kind: "message", + T: () => ExecClientMessageSchema, + }, + { + no: 5, + name: "execClientControlMessage", + kind: "message", + T: () => ExecClientControlMessageSchema, + }, + { + no: 3, + name: "kvClientMessage", + kind: "message", + T: () => KvClientMessageSchema, + }, + { + no: 7, + name: "clientHeartbeat", + kind: "message", + T: () => ClientHeartbeatSchema, + }, + ], + }, + ]); + +export interface AgentRunRequest extends ProtoMessage { + conversationState?: ConversationStateStructure; + action?: ConversationAction; + modelDetails?: ModelDetails; + requestedModel?: RequestedModel; + conversationId?: string; + customSystemPrompt?: string; +} + +export const AgentRunRequestSchema: MessageCodec = + pb("agent.v1.AgentRunRequest", [ + { + no: 1, + name: "conversationState", + kind: "message", + T: () => ConversationStateStructureSchema, + }, + { + no: 2, + name: "action", + kind: "message", + T: () => ConversationActionSchema, + }, + { + no: 3, + name: "modelDetails", + kind: "message", + T: () => ModelDetailsSchema, + }, + { + no: 9, + name: "requestedModel", + kind: "message", + T: () => RequestedModelSchema, + }, + { no: 5, name: "conversationId", kind: "string", optional: true }, + { no: 8, name: "customSystemPrompt", kind: "string", optional: true }, + ]); + +export interface ConversationStateStructure extends ProtoMessage { + rootPromptMessagesJson: Uint8Array[]; + turns: Uint8Array[]; + pendingToolCalls: string[]; +} + +export const ConversationStateStructureSchema: MessageCodec = + pb("agent.v1.ConversationStateStructure", [ + { no: 1, name: "rootPromptMessagesJson", kind: "bytes", repeat: true }, + { no: 4, name: "pendingToolCalls", kind: "string", repeat: true }, + { no: 8, name: "turns", kind: "bytes", repeat: true }, + ]); + +export interface ConversationAction extends ProtoMessage { + action: + | { case: undefined; value?: undefined } + | { case: "userMessageAction"; value: UserMessageAction } + | { case: "resumeAction"; value: ResumeAction }; +} + +export const ConversationActionSchema: MessageCodec = + pb("agent.v1.ConversationAction", [ + { + kind: "oneof", + name: "action", + variants: [ + { + no: 1, + name: "userMessageAction", + kind: "message", + T: () => UserMessageActionSchema, + }, + { + no: 2, + name: "resumeAction", + kind: "message", + T: () => ResumeActionSchema, + }, + ], + }, + ]); + +export interface UserMessageAction extends ProtoMessage { + userMessage?: UserMessage; +} + +export const UserMessageActionSchema: MessageCodec = + pb("agent.v1.UserMessageAction", [ + { no: 1, name: "userMessage", kind: "message", T: () => UserMessageSchema }, + ]); + +export interface ResumeAction extends ProtoMessage {} + +export const ResumeActionSchema: MessageCodec = pb( + "agent.v1.ResumeAction", + [], +); + +export interface UserMessage extends ProtoMessage { + text: string; + messageId: string; + selectedContext?: SelectedContext; + mode: number; +} + +export const UserMessageSchema: MessageCodec = pb( + "agent.v1.UserMessage", + [ + { no: 1, name: "text", kind: "string" }, + { no: 2, name: "messageId", kind: "string" }, + { + no: 3, + name: "selectedContext", + kind: "message", + T: () => SelectedContextSchema, + }, + { no: 4, name: "mode", kind: "int32" }, + ], +); + +export interface SelectedContext extends ProtoMessage { + selectedImages: SelectedImage[]; +} + +export const SelectedContextSchema: MessageCodec = + pb("agent.v1.SelectedContext", [ + { + no: 1, + name: "selectedImages", + kind: "message", + T: () => SelectedImageSchema, + repeat: true, + }, + ]); + +export interface SelectedImage extends ProtoMessage { + uuid: string; + path: string; + dimension?: SelectedImage_Dimension; + mimeType: string; + dataOrBlobId: + | { case: undefined; value?: undefined } + | { case: "blobId"; value: Uint8Array } + | { case: "data"; value: Uint8Array } + | { case: "blobIdWithData"; value: SelectedImage_BlobIdWithData }; +} + +export const SelectedImageSchema: MessageCodec = + pb("agent.v1.SelectedImage", [ + { no: 2, name: "uuid", kind: "string" }, + { no: 3, name: "path", kind: "string" }, + { + no: 4, + name: "dimension", + kind: "message", + T: () => SelectedImage_DimensionSchema, + }, + { no: 7, name: "mimeType", kind: "string" }, + { + kind: "oneof", + name: "dataOrBlobId", + variants: [ + { no: 1, name: "blobId", kind: "bytes" }, + { no: 8, name: "data", kind: "bytes" }, + { + no: 9, + name: "blobIdWithData", + kind: "message", + T: () => SelectedImage_BlobIdWithDataSchema, + }, + ], + }, + ]); + +export interface SelectedImage_BlobIdWithData extends ProtoMessage { + blobId: Uint8Array; + data: Uint8Array; +} + +export const SelectedImage_BlobIdWithDataSchema: MessageCodec = + pb("agent.v1.SelectedImage_BlobIdWithData", [ + { no: 1, name: "blobId", kind: "bytes" }, + { no: 2, name: "data", kind: "bytes" }, + ]); + +export interface SelectedImage_Dimension extends ProtoMessage { + width: number; + height: number; +} + +export const SelectedImage_DimensionSchema: MessageCodec = + pb("agent.v1.SelectedImage_Dimension", [ + { no: 1, name: "width", kind: "int32" }, + { no: 2, name: "height", kind: "int32" }, + ]); + +export interface ConversationTurnStructure extends ProtoMessage { + turn: + | { case: undefined; value?: undefined } + | { case: "agentConversationTurn"; value: AgentConversationTurnStructure }; +} + +export const ConversationTurnStructureSchema: MessageCodec = + pb("agent.v1.ConversationTurnStructure", [ + { + kind: "oneof", + name: "turn", + variants: [ + { + no: 1, + name: "agentConversationTurn", + kind: "message", + T: () => AgentConversationTurnStructureSchema, + }, + ], + }, + ]); + +export interface AgentConversationTurnStructure extends ProtoMessage { + userMessage: Uint8Array; + steps: Uint8Array[]; + requestId?: string; +} + +export const AgentConversationTurnStructureSchema: MessageCodec = + pb( + "agent.v1.AgentConversationTurnStructure", + [ + { no: 1, name: "userMessage", kind: "bytes" }, + { no: 2, name: "steps", kind: "bytes", repeat: true }, + { no: 3, name: "requestId", kind: "string", optional: true }, + ], + ); + +export interface ConversationStep extends ProtoMessage { + message: + | { case: undefined; value?: undefined } + | { case: "assistantMessage"; value: AssistantMessage } + | { case: "thinkingMessage"; value: ThinkingMessage }; +} + +export const ConversationStepSchema: MessageCodec = + pb("agent.v1.ConversationStep", [ + { + kind: "oneof", + name: "message", + variants: [ + { + no: 1, + name: "assistantMessage", + kind: "message", + T: () => AssistantMessageSchema, + }, + { + no: 3, + name: "thinkingMessage", + kind: "message", + T: () => ThinkingMessageSchema, + }, + ], + }, + ]); + +export interface AssistantMessage extends ProtoMessage { + text: string; +} + +export const AssistantMessageSchema: MessageCodec = + pb("agent.v1.AssistantMessage", [ + { no: 1, name: "text", kind: "string" }, + ]); + +export interface ThinkingMessage extends ProtoMessage { + text: string; + durationMs: number; +} + +export const ThinkingMessageSchema: MessageCodec = + pb("agent.v1.ThinkingMessage", [ + { no: 1, name: "text", kind: "string" }, + { no: 2, name: "durationMs", kind: "uint32" }, + ]); + +export interface ModelDetails extends ProtoMessage { + modelId: string; + displayModelId: string; + displayName: string; + displayNameShort: string; + aliases: string[]; + thinkingDetails?: ThinkingDetails; + maxMode?: boolean; +} + +export const ModelDetailsSchema: MessageCodec = pb( + "agent.v1.ModelDetails", + [ + { no: 1, name: "modelId", kind: "string" }, + { no: 3, name: "displayModelId", kind: "string" }, + { no: 4, name: "displayName", kind: "string" }, + { no: 5, name: "displayNameShort", kind: "string" }, + { no: 6, name: "aliases", kind: "string", repeat: true }, + { + no: 2, + name: "thinkingDetails", + kind: "message", + T: () => ThinkingDetailsSchema, + }, + { no: 7, name: "maxMode", kind: "bool", optional: true }, + ], +); + +export interface ThinkingDetails extends ProtoMessage {} + +export const ThinkingDetailsSchema: MessageCodec = + pb("agent.v1.ThinkingDetails", []); + +export interface RequestedModel extends ProtoMessage { + modelId: string; + maxMode: boolean; + parameters: RequestedModel_ModelParameterbytes[]; +} + +export const RequestedModelSchema: MessageCodec = + pb("agent.v1.RequestedModel", [ + { no: 1, name: "modelId", kind: "string" }, + { no: 2, name: "maxMode", kind: "bool" }, + { + no: 3, + name: "parameters", + kind: "message", + T: () => RequestedModel_ModelParameterbytesSchema, + repeat: true, + }, + ]); + +export interface RequestedModel_ModelParameterbytes extends ProtoMessage { + id: string; + value: string; +} + +export const RequestedModel_ModelParameterbytesSchema: MessageCodec = + pb( + "agent.v1.RequestedModel_ModelParameterbytes", + [ + { no: 1, name: "id", kind: "string" }, + { no: 2, name: "value", kind: "string" }, + ], + ); + +export interface ClientHeartbeat extends ProtoMessage {} + +export const ClientHeartbeatSchema: MessageCodec = + pb("agent.v1.ClientHeartbeat", []); + +export interface GetBlobArgs extends ProtoMessage { + blobId: Uint8Array; +} + +export const GetBlobArgsSchema: MessageCodec = pb( + "agent.v1.GetBlobArgs", + [{ no: 1, name: "blobId", kind: "bytes" }], +); + +export interface GetBlobResult extends ProtoMessage { + blobData?: Uint8Array; +} + +export const GetBlobResultSchema: MessageCodec = + pb("agent.v1.GetBlobResult", [ + { no: 1, name: "blobData", kind: "bytes", optional: true }, + ]); + +export interface SetBlobArgs extends ProtoMessage { + blobId: Uint8Array; + blobData: Uint8Array; +} + +export const SetBlobArgsSchema: MessageCodec = pb( + "agent.v1.SetBlobArgs", + [ + { no: 1, name: "blobId", kind: "bytes" }, + { no: 2, name: "blobData", kind: "bytes" }, + ], +); + +export interface SetBlobResult extends ProtoMessage {} + +export const SetBlobResultSchema: MessageCodec = + pb("agent.v1.SetBlobResult", []); + +export interface KvClientMessage extends ProtoMessage { + id: number; + message: + | { case: undefined; value?: undefined } + | { case: "getBlobResult"; value: GetBlobResult } + | { case: "setBlobResult"; value: SetBlobResult }; +} + +export const KvClientMessageSchema: MessageCodec = + pb("agent.v1.KvClientMessage", [ + { no: 1, name: "id", kind: "uint32" }, + { + kind: "oneof", + name: "message", + variants: [ + { + no: 2, + name: "getBlobResult", + kind: "message", + T: () => GetBlobResultSchema, + }, + { + no: 3, + name: "setBlobResult", + kind: "message", + T: () => SetBlobResultSchema, + }, + ], + }, + ]); + +export interface KvServerMessage extends ProtoMessage { + id: number; + message: + | { case: undefined; value?: undefined } + | { case: "getBlobArgs"; value: GetBlobArgs } + | { case: "setBlobArgs"; value: SetBlobArgs }; +} + +export const KvServerMessageSchema: MessageCodec = + pb("agent.v1.KvServerMessage", [ + { no: 1, name: "id", kind: "uint32" }, + { + kind: "oneof", + name: "message", + variants: [ + { + no: 2, + name: "getBlobArgs", + kind: "message", + T: () => GetBlobArgsSchema, + }, + { + no: 3, + name: "setBlobArgs", + kind: "message", + T: () => SetBlobArgsSchema, + }, + ], + }, + ]); + +/** Exec context is the only server-side interaction answered successfully. */ +export interface CursorRuleTypeGlobal extends ProtoMessage {} + +export const CursorRuleTypeGlobalSchema: MessageCodec = + pb("agent.v1.CursorRuleTypeGlobal", []); + +export interface CursorRuleType extends ProtoMessage { + type: + | { case: undefined; value?: undefined } + | { case: "global"; value: CursorRuleTypeGlobal }; +} + +export const CursorRuleTypeSchema: MessageCodec = + pb("agent.v1.CursorRuleType", [ + { + kind: "oneof", + name: "type", + variants: [ + { + no: 1, + name: "global", + kind: "message", + T: () => CursorRuleTypeGlobalSchema, + }, + ], + }, + ]); + +export interface CursorRule extends ProtoMessage { + fullPath: string; + content: string; + type?: CursorRuleType; + source: number; +} + +export const CursorRuleSchema: MessageCodec = pb( + "agent.v1.CursorRule", + [ + { no: 1, name: "fullPath", kind: "string" }, + { no: 2, name: "content", kind: "string" }, + { no: 3, name: "type", kind: "message", T: () => CursorRuleTypeSchema }, + { no: 4, name: "source", kind: "int32" }, + ], +); + +/** Empty definitions deliberately make the request-context tool list empty. */ +export interface McpToolDefinition extends ProtoMessage {} + +export const McpToolDefinitionSchema: MessageCodec = + pb("agent.v1.McpToolDefinition", []); + +export interface RequestContext extends ProtoMessage { + rules: CursorRule[]; + tools: McpToolDefinition[]; +} + +export const RequestContextSchema: MessageCodec = + pb("agent.v1.RequestContext", [ + { + no: 2, + name: "rules", + kind: "message", + T: () => CursorRuleSchema, + repeat: true, + }, + { + no: 7, + name: "tools", + kind: "message", + T: () => McpToolDefinitionSchema, + repeat: true, + }, + ]); + +export interface RequestContextArgs extends ProtoMessage { + notesSessionId?: string; + workspaceId?: string; + readOnlyPinnedTreeSha?: string; + readOnlyPluginCacheRoot?: string; + useCached?: boolean; +} + +export const RequestContextArgsSchema: MessageCodec = + pb("agent.v1.RequestContextArgs", [ + { no: 2, name: "notesSessionId", kind: "string", optional: true }, + { no: 3, name: "workspaceId", kind: "string", optional: true }, + { no: 4, name: "readOnlyPinnedTreeSha", kind: "string", optional: true }, + { no: 5, name: "readOnlyPluginCacheRoot", kind: "string", optional: true }, + { no: 7, name: "useCached", kind: "bool", optional: true }, + ]); + +export interface RequestContextSuccess extends ProtoMessage { + requestContext?: RequestContext; + servedFromDiskCache?: boolean; +} + +export const RequestContextSuccessSchema: MessageCodec = + pb("agent.v1.RequestContextSuccess", [ + { + no: 1, + name: "requestContext", + kind: "message", + T: () => RequestContextSchema, + }, + { no: 2, name: "servedFromDiskCache", kind: "bool", optional: true }, + ]); + +export interface RequestContextError extends ProtoMessage { + error: string; +} + +export const RequestContextErrorSchema: MessageCodec = + pb("agent.v1.RequestContextError", [ + { no: 1, name: "error", kind: "string" }, + ]); + +export interface RequestContextRejected extends ProtoMessage { + reason: string; +} + +export const RequestContextRejectedSchema: MessageCodec = + pb("agent.v1.RequestContextRejected", [ + { no: 1, name: "reason", kind: "string" }, + ]); + +export interface RequestContextResult extends ProtoMessage { + result: + | { case: undefined; value?: undefined } + | { case: "success"; value: RequestContextSuccess } + | { case: "error"; value: RequestContextError } + | { case: "rejected"; value: RequestContextRejected }; +} + +export const RequestContextResultSchema: MessageCodec = + pb("agent.v1.RequestContextResult", [ + { + kind: "oneof", + name: "result", + variants: [ + { + no: 1, + name: "success", + kind: "message", + T: () => RequestContextSuccessSchema, + }, + { + no: 2, + name: "error", + kind: "message", + T: () => RequestContextErrorSchema, + }, + { + no: 3, + name: "rejected", + kind: "message", + T: () => RequestContextRejectedSchema, + }, + ], + }, + ]); + +export interface ExecClientMessage extends ProtoMessage { + id: number; + execId: string; + message: + | { case: undefined; value?: undefined } + | { case: "requestContextResult"; value: RequestContextResult }; +} + +export const ExecClientMessageSchema: MessageCodec = + pb("agent.v1.ExecClientMessage", [ + { no: 1, name: "id", kind: "uint32" }, + { no: 15, name: "execId", kind: "string" }, + { + kind: "oneof", + name: "message", + variants: [ + { + no: 10, + name: "requestContextResult", + kind: "message", + T: () => RequestContextResultSchema, + }, + ], + }, + ]); + +export interface ExecClientStreamClose extends ProtoMessage { + id: number; +} + +export const ExecClientStreamCloseSchema: MessageCodec = + pb("agent.v1.ExecClientStreamClose", [ + { no: 1, name: "id", kind: "uint32" }, + ]); + +export interface ExecClientThrow extends ProtoMessage { + id: number; + error: string; + stackTrace?: string; + errorCode?: string; +} + +export const ExecClientThrowSchema: MessageCodec = + pb("agent.v1.ExecClientThrow", [ + { no: 1, name: "id", kind: "uint32" }, + { no: 2, name: "error", kind: "string" }, + { no: 3, name: "stackTrace", kind: "string", optional: true }, + { no: 4, name: "errorCode", kind: "string", optional: true }, + ]); + +export interface ExecClientControlMessage extends ProtoMessage { + message: + | { case: undefined; value?: undefined } + | { case: "streamClose"; value: ExecClientStreamClose } + | { case: "throw"; value: ExecClientThrow }; +} + +export const ExecClientControlMessageSchema: MessageCodec = + pb("agent.v1.ExecClientControlMessage", [ + { + kind: "oneof", + name: "message", + variants: [ + { + no: 1, + name: "streamClose", + kind: "message", + T: () => ExecClientStreamCloseSchema, + }, + { + no: 2, + name: "throw", + kind: "message", + T: () => ExecClientThrowSchema, + }, + ], + }, + ]); + +export interface ExecServerMessage extends ProtoMessage { + id: number; + execId: string; + message: + | { case: undefined; value?: undefined } + | { case: "requestContextArgs"; value: RequestContextArgs }; +} + +export const ExecServerMessageSchema: MessageCodec = + pb("agent.v1.ExecServerMessage", [ + { no: 1, name: "id", kind: "uint32" }, + { no: 15, name: "execId", kind: "string" }, + { + kind: "oneof", + name: "message", + variants: [ + { + no: 10, + name: "requestContextArgs", + kind: "message", + T: () => RequestContextArgsSchema, + }, + ], + }, + ]); + +export interface AgentServerMessage extends ProtoMessage { + message: + | { case: undefined; value?: undefined } + | { case: "interactionUpdate"; value: InteractionUpdate } + | { case: "execServerMessage"; value: ExecServerMessage } + | { case: "kvServerMessage"; value: KvServerMessage } + | { case: "interactionQuery"; value: InteractionQuery }; +} + +export const AgentServerMessageSchema: MessageCodec = + pb("agent.v1.AgentServerMessage", [ + { + kind: "oneof", + name: "message", + variants: [ + { + no: 1, + name: "interactionUpdate", + kind: "message", + T: () => InteractionUpdateSchema, + }, + { + no: 2, + name: "execServerMessage", + kind: "message", + T: () => ExecServerMessageSchema, + }, + { + no: 4, + name: "kvServerMessage", + kind: "message", + T: () => KvServerMessageSchema, + }, + { + no: 7, + name: "interactionQuery", + kind: "message", + T: () => InteractionQuerySchema, + }, + ], + }, + ]); + +/** + * Queries require an interactive client answer. This provider has no UI or + * tool execution channel, so it recognizes the envelope and fails the turn + * explicitly instead of silently dropping a server request. + */ +export interface InteractionQuery extends ProtoMessage { + id: number; + query: + | { case: undefined; value?: undefined } + | { case: "webSearchRequestQuery"; value: InteractionQueryPayload } + | { case: "askQuestionInteractionQuery"; value: InteractionQueryPayload } + | { case: "switchModeRequestQuery"; value: InteractionQueryPayload } + | { case: "exaSearchRequestQuery"; value: InteractionQueryPayload } + | { case: "exaFetchRequestQuery"; value: InteractionQueryPayload } + | { case: "createPlanRequestQuery"; value: InteractionQueryPayload } + | { case: "setupVmEnvironmentArgs"; value: InteractionQueryPayload } + | { case: "webFetchRequestQuery"; value: InteractionQueryPayload }; +} + +export interface InteractionQueryPayload extends ProtoMessage {} + +export const InteractionQueryPayloadSchema: MessageCodec = + pb("agent.v1.InteractionQueryPayload", []); + +export const InteractionQuerySchema: MessageCodec = + pb("agent.v1.InteractionQuery", [ + { no: 1, name: "id", kind: "uint32" }, + { + kind: "oneof", + name: "query", + variants: [ + { + no: 2, + name: "webSearchRequestQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 3, + name: "askQuestionInteractionQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 4, + name: "switchModeRequestQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 5, + name: "exaSearchRequestQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 6, + name: "exaFetchRequestQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 7, + name: "createPlanRequestQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 8, + name: "setupVmEnvironmentArgs", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + { + no: 9, + name: "webFetchRequestQuery", + kind: "message", + T: () => InteractionQueryPayloadSchema, + }, + ], + }, + ]); + +export interface InteractionUpdate extends ProtoMessage { + message: + | { case: undefined; value?: undefined } + | { case: "textDelta"; value: TextDeltaUpdate } + | { case: "partialToolCall"; value: ToolInteractionUpdate } + | { case: "toolCallDelta"; value: ToolInteractionUpdate } + | { case: "toolCallStarted"; value: ToolInteractionUpdate } + | { case: "toolCallCompleted"; value: ToolInteractionUpdate } + | { case: "thinkingDelta"; value: ThinkingDeltaUpdate } + | { case: "thinkingCompleted"; value: ThinkingCompletedUpdate } + | { case: "tokenDelta"; value: TokenDeltaUpdate } + | { case: "heartbeat"; value: HeartbeatUpdate } + | { case: "turnEnded"; value: TurnEndedUpdate }; +} + +export const InteractionUpdateSchema: MessageCodec = + pb("agent.v1.InteractionUpdate", [ + { + kind: "oneof", + name: "message", + variants: [ + { + no: 1, + name: "textDelta", + kind: "message", + T: () => TextDeltaUpdateSchema, + }, + { + no: 7, + name: "partialToolCall", + kind: "message", + T: () => ToolInteractionUpdateSchema, + }, + { + no: 15, + name: "toolCallDelta", + kind: "message", + T: () => ToolInteractionUpdateSchema, + }, + { + no: 2, + name: "toolCallStarted", + kind: "message", + T: () => ToolInteractionUpdateSchema, + }, + { + no: 3, + name: "toolCallCompleted", + kind: "message", + T: () => ToolInteractionUpdateSchema, + }, + { + no: 4, + name: "thinkingDelta", + kind: "message", + T: () => ThinkingDeltaUpdateSchema, + }, + { + no: 5, + name: "thinkingCompleted", + kind: "message", + T: () => ThinkingCompletedUpdateSchema, + }, + { + no: 8, + name: "tokenDelta", + kind: "message", + T: () => TokenDeltaUpdateSchema, + }, + { + no: 13, + name: "heartbeat", + kind: "message", + T: () => HeartbeatUpdateSchema, + }, + { + no: 14, + name: "turnEnded", + kind: "message", + T: () => TurnEndedUpdateSchema, + }, + ], + }, + ]); + +export interface ToolInteractionUpdate extends ProtoMessage {} + +export const ToolInteractionUpdateSchema: MessageCodec = + pb("agent.v1.ToolInteractionUpdate", []); + +export interface TextDeltaUpdate extends ProtoMessage { + text: string; +} + +export const TextDeltaUpdateSchema: MessageCodec = + pb("agent.v1.TextDeltaUpdate", [ + { no: 1, name: "text", kind: "string" }, + ]); + +export interface ThinkingDeltaUpdate extends ProtoMessage { + text: string; +} + +export const ThinkingDeltaUpdateSchema: MessageCodec = + pb("agent.v1.ThinkingDeltaUpdate", [ + { no: 1, name: "text", kind: "string" }, + ]); + +export interface ThinkingCompletedUpdate extends ProtoMessage { + thinkingDurationMs: number; +} + +export const ThinkingCompletedUpdateSchema: MessageCodec = + pb("agent.v1.ThinkingCompletedUpdate", [ + { no: 1, name: "thinkingDurationMs", kind: "int32" }, + ]); + +export interface TokenDeltaUpdate extends ProtoMessage { + tokens: number; +} + +export const TokenDeltaUpdateSchema: MessageCodec = + pb("agent.v1.TokenDeltaUpdate", [ + { no: 1, name: "tokens", kind: "int32" }, + ]); + +export interface HeartbeatUpdate extends ProtoMessage {} + +export const HeartbeatUpdateSchema: MessageCodec = + pb("agent.v1.HeartbeatUpdate", []); + +export interface TurnEndedUpdate extends ProtoMessage {} + +export const TurnEndedUpdateSchema: MessageCodec = + pb("agent.v1.TurnEndedUpdate", []); + +export interface GetUsableModelsRequest extends ProtoMessage { + customModelIds: string[]; +} + +export const GetUsableModelsRequestSchema: MessageCodec = + pb("agent.v1.GetUsableModelsRequest", [ + { no: 1, name: "customModelIds", kind: "string", repeat: true }, + ]); + +export interface GetUsableModelsResponse extends ProtoMessage { + models: ModelDetails[]; +} + +export const GetUsableModelsResponseSchema: MessageCodec = + pb("agent.v1.GetUsableModelsResponse", [ + { + no: 1, + name: "models", + kind: "message", + T: () => ModelDetailsSchema, + repeat: true, + }, + ]); diff --git a/extensions/ai-providers/cursor/protobuf.ts b/extensions/ai-providers/cursor/protobuf.ts new file mode 100644 index 00000000..49b8330c --- /dev/null +++ b/extensions/ai-providers/cursor/protobuf.ts @@ -0,0 +1,1171 @@ +/** + * Vendored from oh-my-pi@eab72e88e4, + * packages/catalog/src/discovery/protobuf.ts (MIT). Kept local so this + * extension has no @oh-my-pi package or subprocess dependency. + */ +import { Buffer } from "node:buffer"; + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === "object" && !Array.isArray(value); +} + +/** + * High-performance, zero-builder protobuf wire codecs for @oh-my-pi/pi-catalog. + * + * Schemas are declared as static IR descriptors with near-zero module load overhead + * and lazy compilation on first encode/decode/create invocation. + */ + +/** JSON values carried by `google.protobuf.Value` fields. */ +export type JsonValue = + | null + | boolean + | number + | string + | JsonValue[] + | { [key: string]: JsonValue }; + +const textEncoder = new TextEncoder(); +const textDecoder = new TextDecoder("utf-8", { fatal: true }); + +/** An unrecognised wire field retained for forward-compatible round-trips. */ +export interface ProtoUnknownField { + no: number; + wireType: number; + data: Uint8Array; +} + +/** Shared internal metadata present on every decoded protocol message. */ +export interface ProtoMessage { + $typeName?: string; + $unknown?: ProtoUnknownField[]; +} + +/** A bidirectional codec for one protobuf message type. */ +export interface MessageCodec { + (value: T): Uint8Array; + (value: Uint8Array): T; + /** Creates a message with protobuf defaults for omitted fields. */ + create(value?: Partial): T; + /** Encodes one message into protobuf wire bytes. */ + encode(value: T): Uint8Array; + /** Decodes one protobuf message from wire bytes. */ + decode(value: Uint8Array): T; + /** Converts a message to its protobuf JSON representation. */ + toJson(value: T): JsonValue; +} + +/** Infers a message shape from a codec result. */ +export type InferMessage = + TCodec extends MessageCodec ? TMessage : never; + +/** Erases a referenced message's concrete shape for static field descriptors. */ +export interface MessageReference { + encode(value: unknown): Uint8Array; + decode(value: Uint8Array): ProtoMessage; + toJson(value: unknown): JsonValue; +} + +export type ScalarKind = + | "bool" + | "bytes" + | "double" + | "enum" + | "float" + | "int32" + | "int64" + | "string" + | "uint32" + | "uint64"; + +export type WireType = 0 | 1 | 2 | 5; + +export interface ScalarFieldDesc { + readonly no: number; + readonly name: string; + readonly kind: ScalarKind; + readonly optional?: boolean; + readonly repeat?: boolean; +} + +export interface MessageFieldDesc { + readonly no: number; + readonly name: string; + readonly kind: "message"; + readonly T: () => MessageReference; + readonly repeat?: boolean; +} + +export interface EnumFieldDesc { + readonly no: number; + readonly name: string; + readonly kind: "enum"; + readonly optional?: boolean; + readonly repeat?: boolean; +} + +export interface MapFieldDesc { + readonly no: number; + readonly name: string; + readonly kind: "map"; + readonly K: "string"; + readonly V: ScalarKind | (() => MessageReference); +} + +export type VariantDesc = + | { readonly no: number; readonly name: string; readonly kind: ScalarKind } + | { + readonly no: number; + readonly name: string; + readonly kind: "message"; + readonly T: () => MessageReference; + }; + +export interface OneofFieldDesc { + readonly kind: "oneof"; + readonly name: string; + readonly variants: readonly VariantDesc[]; +} + +export type FieldDesc = + | ScalarFieldDesc + | MessageFieldDesc + | EnumFieldDesc + | MapFieldDesc + | OneofFieldDesc; + +/** Runtime representation shared by fields and variants. */ +interface ValueCodec { + readonly wireType: WireType; + readonly defaultValue: TValue | undefined; + encode(value: unknown, writer: Writer): void; + decode(reader: Reader): TValue; + toJson(value: unknown): JsonValue; + isDefault(value: unknown): boolean; +} + +interface CompiledField { + readonly number: number; + encode(message: object, writer: Writer): void; + decode( + message: object, + reader: Reader, + wireType: WireType, + fieldNumber: number, + ): void; + toJson(message: object, output: { [key: string]: JsonValue }): void; + initDefault(message: object): void; +} + +/** Creates a high-performance, lazy protobuf message codec from an IR field descriptor list. */ +export function pb( + typeName: string, + fields: readonly FieldDesc[] = [], +): MessageCodec { + let compiled: MessageCodec | undefined; + + function getCodec(): MessageCodec { + if (!compiled) compiled = compileCodec(typeName, fields); + return compiled; + } + + const codec = ((arg: T | Uint8Array) => { + if (arg instanceof Uint8Array) return getCodec().decode(arg); + return getCodec().encode(arg); + }) as MessageCodec; + + codec.create = (value?: Partial): T => getCodec().create(value); + codec.encode = (value: T): Uint8Array => getCodec().encode(value); + codec.decode = (value: Uint8Array): T => getCodec().decode(value); + codec.toJson = (value: T): JsonValue => getCodec().toJson(value); + + return codec; +} + +/** Creates a message using its codec's protobuf defaults. */ +export function create( + codec: MessageCodec, + value?: Partial, +): TMessage { + return codec.create(value); +} + +/** Encodes a message using its codec. */ +export function toBinary( + codec: MessageCodec, + value: TMessage, +): Uint8Array { + return codec.encode(value); +} + +/** Decodes wire bytes using a message codec. */ +export function fromBinary( + codec: MessageCodec, + value: Uint8Array, +): TMessage { + return codec.decode(value); +} + +/** Converts a message to protobuf JSON using its codec. */ +export function toJson( + codec: MessageCodec, + value: TMessage, +): JsonValue { + return codec.toJson(value); +} + +/** Encodes a JSON value as `google.protobuf.Value`. */ +export function encodeJsonValue(value: JsonValue): Uint8Array { + const writer = new Writer(); + writeJsonValue(writer, value); + return writer.finish(); +} + +/** Decodes `google.protobuf.Value` wire bytes into a JSON value. */ +export function decodeJsonValue(value: Uint8Array): JsonValue { + return readJsonValue(new Reader(value)); +} + +function compileCodec( + typeName: string, + fieldDescs: readonly FieldDesc[], +): MessageCodec { + const compiledFields: CompiledField[] = []; + const byNumber = new Map(); + + for (const desc of fieldDescs) { + if (desc.kind === "oneof") { + const oneofHandler = compileOneofField(desc); + compiledFields.push(oneofHandler); + for (const v of desc.variants) { + byNumber.set(v.no, oneofHandler); + } + } else if (desc.kind === "map") { + const mapHandler = compileMapField(desc); + compiledFields.push(mapHandler); + byNumber.set(desc.no, mapHandler); + } else if (desc.kind === "message") { + const msgHandler = desc.repeat + ? compileRepeatedField(desc.name, desc.no, messageValue(desc.T)) + : compileSingularField(desc.name, desc.no, messageValue(desc.T)); + compiledFields.push(msgHandler); + byNumber.set(desc.no, msgHandler); + } else { + const valCodec = scalarValue(desc.kind); + const scalarHandler = desc.repeat + ? compileRepeatedField(desc.name, desc.no, valCodec) + : compileSingularField(desc.name, desc.no, valCodec, desc.optional); + compiledFields.push(scalarHandler); + byNumber.set(desc.no, scalarHandler); + } + } + + const codec: MessageCodec = ((arg: T | Uint8Array) => { + if (arg instanceof Uint8Array) return codec.decode(arg); + return codec.encode(arg); + }) as MessageCodec; + + codec.create = (value?: Partial): T => { + const message = (typeName ? { $typeName: typeName } : {}) as T; + for (const f of compiledFields) { + f.initDefault(message); + } + if (value) { + for (const key in value) { + const v = Reflect.get(value, key); + if (v !== undefined) { + Reflect.set(message, key, v); + } + } + } + return message; + }; + + codec.encode = (value: T): Uint8Array => { + const writer = new Writer(); + for (const f of compiledFields) { + f.encode(value, writer); + } + writeUnknownFields(value, writer); + return writer.finish(); + }; + + codec.decode = (value: Uint8Array): T => { + const reader = new Reader(value); + const message = (typeName ? { $typeName: typeName } : {}) as T; + for (const f of compiledFields) { + f.initDefault(message); + } + + while (reader.pos < reader.len) { + const tag = reader.uint32(); + const fieldNumber = tag >>> 3; + const wireType = tag & 7; + if (!isWireType(wireType)) { + throw new Error( + `Unsupported protobuf wire type ${wireType} at byte ${reader.pos}`, + ); + } + const field = byNumber.get(fieldNumber); + if (field) { + field.decode(message, reader, wireType, fieldNumber); + } else { + const start = reader.pos; + reader.skip(wireType); + appendUnknownField(message, { + no: fieldNumber, + wireType, + data: reader.slice(start, reader.pos), + }); + } + } + return message; + }; + + codec.toJson = (value: T): JsonValue => { + const output: { [key: string]: JsonValue } = {}; + for (const f of compiledFields) { + f.toJson(value, output); + } + return output; + }; + + return codec; +} + +function compileSingularField( + name: string, + number: number, + value: ValueCodec, + optional = false, +): CompiledField { + return { + number, + initDefault(message) { + if (!optional && value.defaultValue !== undefined) { + Reflect.set(message, name, value.defaultValue); + } + }, + encode(message, writer) { + const input = Reflect.get(message, name); + if (input === undefined || (!optional && value.isDefault(input))) return; + writer.tag(number, value.wireType); + value.encode(input, writer); + }, + decode(message, reader, wireType, _fieldNumber) { + assertWireType(wireType, value.wireType); + Reflect.set(message, name, value.decode(reader)); + }, + toJson(message, output) { + const input = Reflect.get(message, name); + if (input === undefined || (!optional && value.isDefault(input))) return; + output[name] = value.toJson(input); + }, + }; +} + +function compileRepeatedField( + name: string, + number: number, + value: ValueCodec, +): CompiledField { + return { + number, + initDefault(message) { + Reflect.set(message, name, []); + }, + encode(message, writer) { + const items = Reflect.get(message, name); + if (!Array.isArray(items) || items.length === 0) return; + + if (value.wireType !== 2 && isPackableScalar(value)) { + const packed = new Writer(); + for (const item of items) { + value.encode(item, packed); + } + writer.tag(number, 2); + writer.lengthDelimited(packed.finish()); + return; + } + + for (const item of items) { + writer.tag(number, value.wireType); + value.encode(item, writer); + } + }, + decode(message, reader, wireType, _fieldNumber) { + const target = arrayField(message, name); + if (wireType === 2 && value.wireType !== 2 && isPackableScalar(value)) { + const limit = reader.uint32(); + const end = reader.pos + limit; + while (reader.pos < end) { + target.push(value.decode(reader)); + } + return; + } + assertWireType(wireType, value.wireType); + target.push(value.decode(reader)); + }, + toJson(message, output) { + const items = Reflect.get(message, name); + if (!Array.isArray(items) || items.length === 0) return; + output[name] = items.map((item) => value.toJson(item)); + }, + }; +} + +function compileMapField(desc: MapFieldDesc): CompiledField { + const name = desc.name; + const number = desc.no; + const key = scalarValue("string"); + const valCodec = + typeof desc.V === "function" ? messageValue(desc.V) : scalarValue(desc.V); + + return { + number, + initDefault(message) { + Reflect.set(message, name, Object.create(null)); + }, + encode(message, writer) { + const input = Reflect.get(message, name); + if (!isMessageObject(input)) return; + for (const entryKey in input) { + const entry = new Writer(); + if (!key.isDefault(entryKey)) { + entry.tag(1, key.wireType); + key.encode(entryKey, entry); + } + const entryValue = input[entryKey]; + if (!valCodec.isDefault(entryValue)) { + entry.tag(2, valCodec.wireType); + valCodec.encode(entryValue, entry); + } + writer.tag(number, 2); + writer.lengthDelimited(entry.finish()); + } + }, + decode(message, reader, wireType, _fieldNumber) { + assertWireType(wireType, 2); + const target = mapField(message, name); + const limit = reader.uint32(); + const end = reader.pos + limit; + let entryKey = ""; + let entryValue: unknown = valCodec.defaultValue; + + while (reader.pos < end) { + const tag = reader.uint32(); + const entryNumber = tag >>> 3; + const entryWireType = tag & 7; + if (!isWireType(entryWireType)) { + throw new Error( + `Unsupported wire type ${entryWireType} in map entry`, + ); + } + if (entryNumber === 1) { + assertWireType(entryWireType, key.wireType); + entryKey = requireString(key.decode(reader)); + } else if (entryNumber === 2) { + assertWireType(entryWireType, valCodec.wireType); + entryValue = valCodec.decode(reader); + } else { + reader.skip(entryWireType); + } + } + + target[entryKey] = entryValue; + }, + toJson(message, output) { + const input = Reflect.get(message, name); + if (!isMessageObject(input)) return; + const mapOutput: { [key: string]: JsonValue } = Object.create(null); + for (const entryKey in input) { + mapOutput[entryKey] = valCodec.toJson(input[entryKey]); + } + output[name] = mapOutput; + }, + }; +} + +function compileOneofField(desc: OneofFieldDesc): CompiledField { + const name = desc.name; + const variantsByName = new Map< + string, + { no: number; codec: ValueCodec } + >(); + const variantsByNumber = new Map< + number, + { name: string; codec: ValueCodec } + >(); + + for (const variant of desc.variants) { + const codec = + variant.kind === "message" + ? messageValue(variant.T) + : scalarValue(variant.kind); + variantsByName.set(variant.name, { no: variant.no, codec }); + variantsByNumber.set(variant.no, { name: variant.name, codec }); + } + + return { + number: 0, + initDefault(message) { + Reflect.set(message, name, { case: undefined }); + }, + encode(message, writer) { + const oneof = Reflect.get(message, name); + if ( + !oneof || + typeof oneof !== "object" || + !("case" in oneof) || + typeof oneof.case !== "string" + ) + return; + const variant = variantsByName.get(oneof.case); + if (!variant) return; + const value = Reflect.get(oneof, "value"); + if (value === undefined) return; + writer.tag(variant.no, variant.codec.wireType); + variant.codec.encode(value, writer); + }, + decode(message, reader, wireType, fieldNumber) { + const variant = variantsByNumber.get(fieldNumber); + if (!variant) throw new Error(`Unknown oneof field ${fieldNumber}`); + assertWireType(wireType, variant.codec.wireType); + Reflect.set(message, name, { + case: variant.name, + value: variant.codec.decode(reader), + }); + }, + toJson(message, output) { + const oneof = Reflect.get(message, name); + if ( + !oneof || + typeof oneof !== "object" || + !("case" in oneof) || + typeof oneof.case !== "string" + ) + return; + const variant = variantsByName.get(oneof.case); + if (!variant) return; + const value = Reflect.get(oneof, "value"); + if (value === undefined) return; + output[oneof.case] = variant.codec.toJson(value); + }, + }; +} + +function isPackableScalar(value: ValueCodec): boolean { + return value.wireType === 0 || value.wireType === 1 || value.wireType === 5; +} + +function scalarValue(kind: ScalarKind): ValueCodec { + switch (kind) { + case "bool": + return scalar( + 0, + false, + requireBoolean, + (value, writer) => writer.bool(value), + (reader) => reader.bool(), + (value) => value, + ); + case "bytes": + return scalar( + 2, + new Uint8Array(0), + requireBytes, + (value, writer) => writer.bytes(value), + (reader) => reader.bytes(), + (value) => Buffer.from(value).toString("base64"), + (value) => value.byteLength === 0, + ); + case "double": + return scalar( + 1, + 0, + requireNumber, + (value, writer) => writer.double(value), + (reader) => reader.double(), + (value) => value, + ); + case "enum": + return scalar( + 0, + 0, + requireInt32, + (value, writer) => writer.int32(value), + (reader) => reader.int32(), + (value) => value, + ); + case "float": + return scalar( + 5, + 0, + requireNumber, + (value, writer) => writer.float(value), + (reader) => reader.float(), + (value) => value, + ); + case "int32": + return scalar( + 0, + 0, + requireInt32, + (value, writer) => writer.int32(value), + (reader) => reader.int32(), + (value) => value, + ); + case "int64": + return scalar( + 0, + 0n, + requireBigInt, + (value, writer) => writer.int64(value), + (reader) => reader.int64(), + (value) => value.toString(), + ); + case "string": + return scalar( + 2, + "", + requireString, + (value, writer) => writer.string(value), + (reader) => reader.string(), + (value) => value, + ); + case "uint32": + return scalar( + 0, + 0, + requireUint32, + (value, writer) => writer.uint32(value), + (reader) => reader.uint32(), + (value) => value, + ); + case "uint64": + return scalar( + 0, + 0n, + requireUnsignedBigInt, + (value, writer) => writer.uint64(value), + (reader) => reader.uint64(), + (value) => value.toString(), + ); + } +} + +function scalar( + wireType: WireType, + defaultValue: TValue, + validate: (value: unknown) => TValue, + write: (value: TValue, writer: Writer) => void, + read: (reader: Reader) => TValue, + toJson: (value: TValue) => JsonValue, + isDefault: (value: TValue) => boolean = (value) => value === defaultValue, +): ValueCodec { + return { + wireType, + defaultValue, + encode(value, writer) { + write(validate(value), writer); + }, + decode(reader) { + return read(reader); + }, + toJson(value) { + return toJson(validate(value)); + }, + isDefault(value) { + return isDefault(validate(value)); + }, + }; +} + +function messageValue(factory: () => MessageReference): ValueCodec { + let cached: MessageReference | undefined; + + function getCodec(): MessageReference { + if (!cached) cached = factory(); + return cached; + } + + return { + wireType: 2, + defaultValue: undefined, + encode(value, writer) { + writer.lengthDelimited(getCodec().encode(value)); + }, + decode(reader) { + return getCodec().decode(reader.bytes()); + }, + toJson(value) { + return getCodec().toJson(value); + }, + isDefault(value) { + return value === undefined; + }, + }; +} + +function requireBoolean(value: unknown): boolean { + if (typeof value === "boolean") return value; + throw new Error(`Expected boolean, got ${typeof value}`); +} + +function requireBytes(value: unknown): Uint8Array { + if (value instanceof Uint8Array) return value; + throw new Error("Expected Uint8Array"); +} + +function requireNumber(value: unknown): number { + if (typeof value === "number" && Number.isFinite(value)) return value; + throw new Error(`Expected number, got ${typeof value}`); +} + +function requireInt32(value: unknown): number { + if (typeof value === "number" && Number.isInteger(value)) return value | 0; + throw new Error(`Expected int32, got ${typeof value}`); +} +function requireString(value: unknown): string { + if (typeof value === "string") return value; + throw new Error(`Expected string, got ${typeof value}`); +} + +function requireBigInt(value: unknown): bigint { + if (typeof value === "bigint") return value; + if (typeof value === "number" && Number.isInteger(value)) + return BigInt(value); + if (typeof value === "string") return BigInt(value); + throw new Error(`Expected bigint, got ${typeof value}`); +} + +function requireUint32(value: unknown): number { + if (typeof value === "number" && Number.isInteger(value) && value >= 0) + return value >>> 0; + throw new Error(`Expected uint32, got ${typeof value}`); +} + +function requireUnsignedBigInt(value: unknown): bigint { + const b = requireBigInt(value); + if (b < 0n) throw new Error("Expected unsigned bigint"); + return b; +} + +function arrayField(message: object, name: string): unknown[] { + let arr = Reflect.get(message, name); + if (!Array.isArray(arr)) { + arr = []; + Reflect.set(message, name, arr); + } + return arr; +} + +function mapField(message: object, name: string): Record { + const value = Reflect.get(message, name); + if (isRecord(value)) return value; + const map: Record = Object.create(null); + Reflect.set(message, name, map); + return map; +} + +function appendUnknownField(message: object, field: ProtoUnknownField): void { + const existing = Reflect.get(message, "$unknown"); + if (isUnknownFields(existing)) { + existing.push(field); + return; + } + Reflect.set(message, "$unknown", [field]); +} + +function isUnknownFields(value: unknown): value is ProtoUnknownField[] { + return Array.isArray(value) && value.every(isUnknownField); +} + +function writeUnknownFields(message: object, writer: Writer): void { + const bag = Reflect.get(message, "$unknown"); + if (!Array.isArray(bag)) return; + for (const field of bag) { + if (isUnknownField(field)) { + writer.tag(field.no, field.wireType); + writer.raw(field.data); + } + } +} + +function isUnknownField( + value: unknown, +): value is ProtoUnknownField & { wireType: WireType } { + return ( + isRecord(value) && + typeof value.no === "number" && + typeof value.wireType === "number" && + isWireType(value.wireType) && + value.data instanceof Uint8Array + ); +} + +function isMessageObject(value: unknown): value is { [key: string]: unknown } { + return isRecord(value) && !(value instanceof Uint8Array); +} + +function isWireType(value: number): value is WireType { + return value === 0 || value === 1 || value === 2 || value === 5; +} + +function assertWireType(actual: WireType, expected: WireType): void { + if (actual !== expected) + throw new Error( + `Unexpected protobuf wire type ${actual}; expected ${expected}`, + ); +} + +class Writer { + #chunks: Uint8Array[] = []; + #length = 0; + + tag(number: number, wireType: WireType): void { + this.uint32((number << 3) | wireType); + } + + uint32(value: number): void { + let v = value >>> 0; + const buffer = new Uint8Array(5); + let pos = 0; + while (v > 0x7f) { + buffer[pos++] = (v & 0x7f) | 0x80; + v >>>= 7; + } + buffer[pos++] = v; + this.raw(buffer.subarray(0, pos)); + } + + int32(value: number): void { + if (value >= 0) { + this.uint32(value); + return; + } + this.int64(BigInt(value)); + } + + int64(value: bigint): void { + let v = BigInt.asUintN(64, value); + const buffer = new Uint8Array(10); + let pos = 0; + while (v > 0x7fn) { + buffer[pos++] = Number(v & 0x7fn) | 0x80; + v >>= 7n; + } + buffer[pos++] = Number(v); + this.raw(buffer.subarray(0, pos)); + } + + uint64(value: bigint): void { + this.int64(value); + } + + bool(value: boolean): void { + this.raw(new Uint8Array([value ? 1 : 0])); + } + + float(value: number): void { + const buffer = new Uint8Array(4); + new DataView(buffer.buffer).setFloat32(0, value, true); + this.raw(buffer); + } + + double(value: number): void { + const buffer = new Uint8Array(8); + new DataView(buffer.buffer).setFloat64(0, value, true); + this.raw(buffer); + } + + string(value: string): void { + this.lengthDelimited(textEncoder.encode(value)); + } + + bytes(value: Uint8Array): void { + this.lengthDelimited(value); + } + + lengthDelimited(value: Uint8Array): void { + this.uint32(value.byteLength); + this.raw(value); + } + + raw(chunk: Uint8Array): void { + this.#chunks.push(chunk); + this.#length += chunk.byteLength; + } + + finish(): Uint8Array { + if (this.#chunks.length === 1) return this.#chunks[0]; + const result = new Uint8Array(this.#length); + let offset = 0; + for (const chunk of this.#chunks) { + result.set(chunk, offset); + offset += chunk.byteLength; + } + return result; + } +} + +class Reader { + readonly buf: Uint8Array; + readonly len: number; + pos = 0; + + constructor(buf: Uint8Array) { + this.buf = buf; + this.len = buf.byteLength; + } + + uint32(): number { + let result = 0; + let shift = 0; + while (this.pos < this.len) { + const byte = this.buf[this.pos++]; + result |= (byte & 0x7f) << shift; + if ((byte & 0x80) === 0) return result >>> 0; + shift += 7; + if (shift >= 32) throw new Error("Varint exceeds 32 bits"); + } + throw new Error("Unexpected end of protobuf varint"); + } + + int32(): number { + return Number(BigInt.asIntN(32, this.uint64())); + } + + int64(): bigint { + let result = 0n; + let shift = 0n; + while (this.pos < this.len) { + const byte = this.buf[this.pos++]; + result |= BigInt(byte & 0x7f) << shift; + if ((byte & 0x80) === 0) return BigInt.asIntN(64, result); + shift += 7n; + if (shift >= 64n) throw new Error("Varint exceeds 64 bits"); + } + throw new Error("Unexpected end of protobuf 64-bit varint"); + } + + uint64(): bigint { + return BigInt.asUintN(64, this.int64()); + } + + bool(): boolean { + return this.uint32() !== 0; + } + + float(): number { + if (this.pos + 4 > this.len) + throw new Error("Unexpected EOF reading float"); + const view = new DataView( + this.buf.buffer, + this.buf.byteOffset + this.pos, + 4, + ); + this.pos += 4; + return view.getFloat32(0, true); + } + + double(): number { + if (this.pos + 8 > this.len) + throw new Error("Unexpected EOF reading double"); + const view = new DataView( + this.buf.buffer, + this.buf.byteOffset + this.pos, + 8, + ); + this.pos += 8; + return view.getFloat64(0, true); + } + + string(): string { + return textDecoder.decode(this.bytes()); + } + + bytes(): Uint8Array { + const length = this.uint32(); + if (this.pos + length > this.len) + throw new Error("Unexpected EOF reading bytes"); + const result = this.buf.subarray(this.pos, this.pos + length); + this.pos += length; + return result; + } + + slice(start: number, end: number): Uint8Array { + return this.buf.subarray(start, end); + } + + skip(wireType: WireType): void { + switch (wireType) { + case 0: + this.int64(); + return; + case 1: + if (this.pos + 8 > this.len) + throw new Error("Unexpected EOF skipping 64-bit"); + this.pos += 8; + return; + case 2: + this.bytes(); + return; + case 5: + if (this.pos + 4 > this.len) + throw new Error("Unexpected EOF skipping 32-bit"); + this.pos += 4; + return; + } + } +} + +function writeJsonValue(writer: Writer, value: JsonValue): void { + if (value === null) { + writer.tag(1, 0); + writer.uint32(0); + return; + } + if (typeof value === "number") { + writer.tag(2, 1); + writer.double(value); + return; + } + if (typeof value === "string") { + writer.tag(3, 2); + writer.string(value); + return; + } + if (typeof value === "boolean") { + writer.tag(4, 0); + writer.bool(value); + return; + } + if (Array.isArray(value)) { + const listWriter = new Writer(); + for (const item of value) { + listWriter.tag(1, 2); + const itemWriter = new Writer(); + writeJsonValue(itemWriter, item); + listWriter.lengthDelimited(itemWriter.finish()); + } + writer.tag(6, 2); + writer.lengthDelimited(listWriter.finish()); + return; + } + if (isRecord(value)) { + const structWriter = new Writer(); + for (const key in value) { + const item = value[key]; + const entryWriter = new Writer(); + entryWriter.tag(1, 2); + entryWriter.string(key); + entryWriter.tag(2, 2); + const valueWriter = new Writer(); + writeJsonValue(valueWriter, item); + entryWriter.lengthDelimited(valueWriter.finish()); + structWriter.tag(1, 2); + structWriter.lengthDelimited(entryWriter.finish()); + } + writer.tag(5, 2); + writer.lengthDelimited(structWriter.finish()); + } +} + +function readJsonValue(reader: Reader): JsonValue { + let value: JsonValue = null; + while (reader.pos < reader.len) { + const tag = reader.uint32(); + const fieldNumber = tag >>> 3; + const wireType = tag & 7; + if (!isWireType(wireType)) { + throw new Error( + `Unsupported wire type ${wireType} in google.protobuf.Value`, + ); + } + switch (fieldNumber) { + case 1: + assertWireType(wireType, 0); + reader.uint32(); + value = null; + break; + case 2: + assertWireType(wireType, 1); + value = reader.double(); + break; + case 3: + assertWireType(wireType, 2); + value = reader.string(); + break; + case 4: + assertWireType(wireType, 0); + value = reader.bool(); + break; + case 5: + assertWireType(wireType, 2); + value = readJsonStruct(new Reader(reader.bytes())); + break; + case 6: + assertWireType(wireType, 2); + value = readJsonList(new Reader(reader.bytes())); + break; + default: + reader.skip(wireType); + break; + } + } + return value; +} + +function readJsonStruct(reader: Reader): { [key: string]: JsonValue } { + const output: { [key: string]: JsonValue } = {}; + while (reader.pos < reader.len) { + const tag = reader.uint32(); + const fieldNumber = tag >>> 3; + const wireType = tag & 7; + if (!isWireType(wireType)) { + throw new Error(`Unsupported wire type ${wireType} in Struct`); + } + if (fieldNumber === 1) { + assertWireType(wireType, 2); + const entryReader = new Reader(reader.bytes()); + let entryKey = ""; + let entryVal: JsonValue = null; + while (entryReader.pos < entryReader.len) { + const entryTag = entryReader.uint32(); + const entryNo = entryTag >>> 3; + const entryWire = entryTag & 7; + if (entryNo === 1) { + entryKey = entryReader.string(); + } else if (entryNo === 2) { + entryVal = readJsonValue(new Reader(entryReader.bytes())); + } else if (isWireType(entryWire)) { + entryReader.skip(entryWire); + } + } + output[entryKey] = entryVal; + } else { + reader.skip(wireType); + } + } + return output; +} + +function readJsonList(reader: Reader): JsonValue[] { + const list: JsonValue[] = []; + while (reader.pos < reader.len) { + const tag = reader.uint32(); + const fieldNumber = tag >>> 3; + const wireType = tag & 7; + if (!isWireType(wireType)) { + throw new Error(`Unsupported wire type ${wireType} in ListValue`); + } + if (fieldNumber === 1) { + assertWireType(wireType, 2); + list.push(readJsonValue(new Reader(reader.bytes()))); + } else { + reader.skip(wireType); + } + } + return list; +} diff --git a/extensions/ai-providers/cursor/provider.ts b/extensions/ai-providers/cursor/provider.ts new file mode 100644 index 00000000..cc40f08c --- /dev/null +++ b/extensions/ai-providers/cursor/provider.ts @@ -0,0 +1,1168 @@ +import { createHash, randomUUID } from "node:crypto"; +import * as http2 from "node:http2"; +import type { + Api, + AssistantMessage, + AssistantMessageEventStream, + Context, + ImageContent, + Message, + Model, + SimpleStreamOptions, + TextContent, +} from "@earendil-works/pi-ai/compat"; +import { createAssistantMessageEventStream } from "@earendil-works/pi-ai/compat"; +import { + CURSOR_API_URL, + CURSOR_CLIENT_VERSION, + CURSOR_RUN_PATH, +} from "./constants.ts"; +import { + AgentClientMessageSchema, + AgentConversationTurnStructureSchema, + type AgentRunRequest, + AgentRunRequestSchema, + AgentServerMessageSchema, + AssistantMessageSchema, + ClientHeartbeatSchema, + ConversationActionSchema, + type ConversationStateStructure, + ConversationStateStructureSchema, + ConversationStepSchema, + ConversationTurnStructureSchema, + type CursorRule, + CursorRuleSchema, + CursorRuleTypeGlobalSchema, + CursorRuleTypeSchema, + ExecClientControlMessageSchema, + ExecClientMessageSchema, + ExecClientStreamCloseSchema, + ExecClientThrowSchema, + GetBlobResultSchema, + type InteractionUpdate, + KvClientMessageSchema, + type KvServerMessage, + KvServerMessageSchema, + type ModelDetails, + ModelDetailsSchema, + RequestContextResultSchema, + RequestContextSchema, + RequestContextSuccessSchema, + type RequestedModel_ModelParameterbytes, + RequestedModel_ModelParameterbytesSchema, + RequestedModelSchema, + ResumeActionSchema, + SelectedContextSchema, + SelectedImageSchema, + SetBlobResultSchema, + UserMessageActionSchema, + UserMessageSchema, +} from "./proto.ts"; +import { create, fromBinary, toBinary } from "./protobuf.ts"; +import { connectCursorHttp2 } from "./proxy.ts"; + +const CONNECT_END_STREAM_FLAG = 0b00000010; +const CONNECT_COMPRESSED_FLAG = 0b00000001; +const MAX_CONNECT_FRAME_BYTES = 16 * 1024 * 1024; +const HEARTBEAT_INTERVAL_MS = 5_000; +const PROXY_TUNNEL_TIMEOUT_MS = 30_000; + +export const CURSOR_CHAT_ONLY_SYSTEM_PROMPT = + "This Cursor provider is running in chat-only mode. No filesystem, shell, code modification, MCP, web, or user-interaction tools are available. Never emit tool calls or interaction queries. Images attached to the user message are already available for direct analysis. If required information is unavailable, explain the limitation in text instead of attempting a tool."; + +const HTTP2_FORBIDDEN_HEADERS = new Set([ + "connection", + "keep-alive", + "proxy-connection", + "transfer-encoding", + "upgrade", + "http2-settings", +]); + +const CURSOR_RESERVED_HEADERS = new Set([ + "content-type", + "connect-protocol-version", + "te", + "authorization", + "x-ghost-mode", + "x-cursor-client-version", + "x-cursor-client-type", + "x-request-id", + "host", + "content-length", +]); + +type CursorBlobStore = Map; + +export interface CursorRequestBuild { + request: AgentRunRequest; + requestBytes: Uint8Array; + blobStore: CursorBlobStore; + conversationState: ConversationStateStructure; +} + +/** Connect's five-byte big-endian envelope. */ +export function frameConnectMessage(data: Uint8Array, flags = 0): Buffer { + const frame = Buffer.allocUnsafe(5 + data.length); + frame[0] = flags; + frame.writeUInt32BE(data.length, 1); + frame.set(data, 5); + return frame; +} + +function createBlobId(data: Uint8Array): Uint8Array { + return new Uint8Array(createHash("sha256").update(data).digest()); +} + +function storeBlob(store: CursorBlobStore, data: Uint8Array): Uint8Array { + const id = createBlobId(data); + store.set(Buffer.from(id).toString("hex"), data); + return id; +} + +function textFromContent( + content: string | (TextContent | ImageContent)[], +): string { + if (typeof content === "string") return content.trim(); + return content + .filter((item): item is TextContent => item.type === "text") + .map((item) => item.text) + .join("\n") + .trim(); +} + +function imagesFromContent(content: string | (TextContent | ImageContent)[]) { + if (typeof content === "string") return []; + return content + .filter((item): item is ImageContent => item.type === "image") + .map((item) => + create(SelectedImageSchema, { + uuid: randomUUID(), + path: "", + mimeType: item.mimeType, + dataOrBlobId: { + case: "data", + value: Uint8Array.from(Buffer.from(item.data, "base64")), + }, + }), + ); +} + +function userMessageFromContent( + content: string | (TextContent | ImageContent)[], + messageId = randomUUID(), +) { + const text = textFromContent(content); + const images = imagesFromContent(content); + return create(UserMessageSchema, { + text, + messageId, + ...(images.length > 0 + ? { + selectedContext: create(SelectedContextSchema, { + selectedImages: images, + }), + } + : {}), + }); +} + +function rootPromptContent( + content: string | (TextContent | ImageContent)[], +): Array< + | { type: "text"; text: string } + | { type: "image"; image: string; mediaType: string } +> { + if (typeof content === "string") { + const text = content.trim(); + return text ? [{ type: "text", text }] : []; + } + const parts: Array< + | { type: "text"; text: string } + | { type: "image"; image: string; mediaType: string } + > = []; + for (const item of content) { + if (item.type === "text") { + const text = item.text.trim(); + if (text) parts.push({ type: "text", text }); + } else { + parts.push({ + type: "image", + image: `data:${item.mimeType};base64,${item.data}`, + mediaType: item.mimeType, + }); + } + } + return parts; +} + +function assistantRootContent( + message: Extract, +) { + const content: Array> = []; + for (const item of message.content) { + if (item.type === "text" && item.text) { + content.push({ type: "text", text: item.text }); + } + } + return content; +} + +function buildHistoryRootPrompt( + messages: Message[], + store: CursorBlobStore, + activeUserIndex: number, +): Uint8Array[] { + const entries: Uint8Array[] = []; + for (let index = 0; index < messages.length; index++) { + if (index === activeUserIndex) break; + const message = messages[index]; + let value: unknown; + if (message.role === "user") { + const content = rootPromptContent(message.content); + if (content.length === 0) continue; + value = { role: "user", content }; + } else if (message.role === "assistant") { + const content = assistantRootContent(message); + if (content.length === 0) continue; + value = { role: "assistant", content }; + } else { + // Chat-only mode never replays assistant tool calls. Replaying only the + // matching tool result would create an invalid orphan in Cursor history. + continue; + } + entries.push( + storeBlob(store, new TextEncoder().encode(JSON.stringify(value))), + ); + } + return entries; +} + +function buildSystemPrompt( + systemPrompt: Context["systemPrompt"], + store: CursorBlobStore, +): Uint8Array[] { + const prompts = systemPrompt + ? Array.isArray(systemPrompt) + ? systemPrompt + : [systemPrompt] + : ["You are a helpful assistant."]; + return [...prompts, CURSOR_CHAT_ONLY_SYSTEM_PROMPT].map((prompt) => + storeBlob( + store, + new TextEncoder().encode( + JSON.stringify({ role: "system", content: prompt }), + ), + ), + ); +} + +/** + * Cursor asks for these rules over the exec channel before generating text. + * They are global rules only; the chat-only provider intentionally returns an + * empty MCP tool list and never forwards `context.tools`. + */ +export function buildCursorRequestContextRules( + systemPrompt: Context["systemPrompt"], +): CursorRule[] { + const rules: CursorRule[] = systemPrompt?.trim() + ? [ + create(CursorRuleSchema, { + fullPath: "/pi/system-prompt.mdc", + content: systemPrompt, + source: 2, + type: create(CursorRuleTypeSchema, { + type: { + case: "global", + value: create(CursorRuleTypeGlobalSchema, {}), + }, + }), + }), + ] + : []; + rules.push( + create(CursorRuleSchema, { + fullPath: "/pi/cursor-chat-only.mdc", + content: CURSOR_CHAT_ONLY_SYSTEM_PROMPT, + source: 2, + type: create(CursorRuleTypeSchema, { + type: { case: "global", value: create(CursorRuleTypeGlobalSchema, {}) }, + }), + }), + ); + return rules; +} + +function buildHistoryTurns( + messages: Message[], + store: CursorBlobStore, + activeUserIndex: number, +): Uint8Array[] { + const turns: Uint8Array[] = []; + const end = activeUserIndex >= 0 ? activeUserIndex : messages.length; + let index = 0; + while (index < end) { + const user = messages[index]; + if (user.role !== "user") { + index++; + continue; + } + const userMessage = storeBlob( + store, + toBinary(UserMessageSchema, userMessageFromContent(user.content)), + ); + const steps: Uint8Array[] = []; + index++; + while (index < end && messages[index]?.role !== "user") { + const message = messages[index]; + if (message.role === "assistant") { + for (const item of message.content) { + if (item.type === "text" && item.text) { + steps.push( + storeBlob( + store, + toBinary( + ConversationStepSchema, + create(ConversationStepSchema, { + message: { + case: "assistantMessage", + value: create(AssistantMessageSchema, { + text: item.text, + }), + }, + }), + ), + ), + ); + } + } + } + index++; + } + const turn = create(ConversationTurnStructureSchema, { + turn: { + case: "agentConversationTurn", + value: create(AgentConversationTurnStructureSchema, { + userMessage, + steps, + }), + }, + }); + turns.push( + storeBlob(store, toBinary(ConversationTurnStructureSchema, turn)), + ); + } + return turns; +} + +function lastUserIndex(messages: Message[]): number { + for (let index = messages.length - 1; index >= 0; index--) { + const role = messages[index]?.role; + if (role === "user") return index; + } + return -1; +} + +type CursorModelWithOptions = Model & { cursorMaxMode?: boolean }; + +function hasCursorMaxMode(model: Model): model is CursorModelWithOptions { + return Object.hasOwn(model, "cursorMaxMode"); +} + +function cursorMaxMode(model: Model): boolean { + return hasCursorMaxMode(model) && model.cursorMaxMode === true; +} + +function resolveWireModel(model: Model): { + modelId: string; + parameters: RequestedModel_ModelParameterbytes[]; +} { + const id = model.id; + // Cursor resolves the bare Composer 2.5 id to its Fast lane unless the + // Standard tier is requested explicitly. + if (id === "composer-2.5") { + return { + modelId: id, + parameters: [ + create(RequestedModel_ModelParameterbytesSchema, { + id: "fast", + value: "false", + }), + ], + }; + } + const match = /^(.*)-(minimal|low|medium|high|xhigh|max)(-fast)?$/.exec(id); + if (!match?.[1] || !/(?:gpt|codex|o\d)/i.test(match[1])) { + return { modelId: id, parameters: [] }; + } + return { + modelId: `${match[1]}${match[3] ?? ""}`, + parameters: [ + create(RequestedModel_ModelParameterbytesSchema, { + id: "reasoning", + value: match[2]!, + }), + ], + }; +} + +/** Build the protobuf Run request and retain blobs for the same Connect stream. */ +export async function buildCursorRequest( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): Promise { + const store: CursorBlobStore = new Map(); + const activeIndex = lastUserIndex(context.messages); + const active = activeIndex >= 0 ? context.messages[activeIndex] : undefined; + const activeContent = active?.role === "user" ? active.content : undefined; + const rootPromptMessagesJson = [ + ...buildSystemPrompt(context.systemPrompt, store), + ...buildHistoryRootPrompt(context.messages, store, activeIndex), + ]; + const state = create(ConversationStateStructureSchema, { + rootPromptMessagesJson, + turns: buildHistoryTurns(context.messages, store, activeIndex), + pendingToolCalls: [], + }); + const conversationId = options?.sessionId ?? randomUUID(); + const action = create(ConversationActionSchema, { + action: + activeContent !== undefined && + (textFromContent(activeContent).length > 0 || + imagesFromContent(activeContent).length > 0) + ? { + case: "userMessageAction", + value: create(UserMessageActionSchema, { + userMessage: userMessageFromContent(activeContent), + }), + } + : { case: "resumeAction", value: create(ResumeActionSchema, {}) }, + }); + const wire = resolveWireModel(model); + let request = create(AgentRunRequestSchema, { + conversationState: state, + action, + modelDetails: create(ModelDetailsSchema, { + modelId: wire.modelId, + displayModelId: model.id, + displayName: model.name, + displayNameShort: model.name, + aliases: [], + ...(cursorMaxMode(model) ? { maxMode: true } : {}), + }), + requestedModel: create(RequestedModelSchema, { + modelId: wire.modelId, + maxMode: cursorMaxMode(model), + parameters: wire.parameters, + }), + conversationId, + }); + const replacement = await options?.onPayload?.(request, model); + if (replacement !== undefined) request = replacement as AgentRunRequest; + const clientMessage = create(AgentClientMessageSchema, { + message: { case: "runRequest", value: request }, + }); + return { + request, + requestBytes: toBinary(AgentClientMessageSchema, clientMessage), + blobStore: store, + conversationState: state, + }; +} + +function sanitizeCallerHeaders( + headers: SimpleStreamOptions["headers"], +): Record { + const result: Record = {}; + for (const [name, value] of Object.entries(headers ?? {})) { + if (value === null) continue; + const field = name.toLowerCase(); + if (field.startsWith(":")) continue; + if ( + HTTP2_FORBIDDEN_HEADERS.has(field) || + CURSOR_RESERVED_HEADERS.has(field) + ) + continue; + result[field] = value; + } + return result; +} + +function cursorHeaders( + apiKey: string, + options: SimpleStreamOptions | undefined, +): Record { + return { + ...sanitizeCallerHeaders(options?.headers), + ":method": "POST", + ":path": CURSOR_RUN_PATH, + "content-type": "application/connect+proto", + "connect-protocol-version": "1", + te: "trailers", + authorization: `Bearer ${apiKey}`, + "x-ghost-mode": "true", + "x-cursor-client-version": CURSOR_CLIENT_VERSION, + "x-cursor-client-type": "cli", + "x-request-id": randomUUID(), + }; +} + +function emptyUsage(): AssistantMessage["usage"] { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function headerRecord( + headers: http2.IncomingHttpHeaders, +): Record { + const result: Record = {}; + for (const [key, value] of Object.entries(headers)) { + if (typeof value === "string") result[key] = value; + else if (Array.isArray(value)) result[key] = value.join(", "); + } + return result; +} + +function errorFromEndStream(data: Uint8Array): Error | undefined { + try { + const parsed: unknown = JSON.parse(new TextDecoder().decode(data)); + if (parsed && typeof parsed === "object" && "error" in parsed) { + const error = parsed.error; + if (error && typeof error === "object") { + const message = + "message" in error && typeof error.message === "string" + ? error.message + : "Cursor Connect error"; + const code = + "code" in error && typeof error.code === "string" + ? error.code + : "unknown"; + return new Error(`Connect error ${code}: ${message}`); + } + } + return undefined; + } catch { + return new Error("Failed to parse Cursor Connect end-stream envelope"); + } +} + +function isAbortError( + error: unknown, + signal: AbortSignal | undefined, +): boolean { + return ( + Boolean(signal?.aborted) || + (error instanceof Error && + /aborted|cancelled|canceled/i.test(error.message)) + ); +} + +/** Cursor AgentService/Run, deliberately chat-only (no context.tools advertisement or Pi tool calls). */ +export function streamCursor( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): AssistantMessageEventStream { + const stream = createAssistantMessageEventStream(); + (async () => { + const output: AssistantMessage = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: emptyUsage(), + stopReason: "pending", + timestamp: Date.now(), + }; + let h2Client: http2.ClientHttp2Session | undefined; + let h2Request: http2.ClientHttp2Stream | undefined; + let heartbeat: ReturnType | undefined; + let idleTimer: ReturnType | undefined; + let removeAbortListener: (() => void) | undefined; + let currentText: + | Extract + | undefined; + let currentThinking: + | Extract + | undefined; + let turnEnded = false; + let terminalError: Error | undefined; + let finished = false; + + const closeBlocks = () => { + if (currentText) { + const index = output.content.indexOf(currentText); + stream.push({ + type: "text_end", + contentIndex: index, + content: currentText.text, + partial: output, + }); + currentText = undefined; + } + if (currentThinking) { + const index = output.content.indexOf(currentThinking); + stream.push({ + type: "thinking_end", + contentIndex: index, + content: currentThinking.thinking, + partial: output, + }); + currentThinking = undefined; + } + }; + + const finishError = (error: unknown) => { + if (finished) return; + finished = true; + closeBlocks(); + output.stopReason = isAbortError(error, options?.signal) + ? "aborted" + : "error"; + output.errorMessage = + error instanceof Error ? error.message : String(error); + stream.push({ + type: "error", + reason: output.stopReason, + error: output, + }); + stream.end(); + }; + + try { + const apiKey = options?.apiKey?.trim(); + if (!apiKey) + throw new Error("Cursor API key is required — run /login cursor"); + if (options?.fetch) { + throw new Error( + "Cursor uses an HTTP/2 transport and does not support options.fetch", + ); + } + if (options?.signal?.aborted) throw new Error("Cursor request aborted"); + const timeoutMs = options?.timeoutMs; + if ( + timeoutMs !== undefined && + (!Number.isFinite(timeoutMs) || timeoutMs < 0) + ) { + throw new Error(`Invalid timeoutMs: ${String(timeoutMs)}`); + } + const built = await buildCursorRequest(model, context, options); + const baseUrl = model.baseUrl || CURSOR_API_URL; + const completion = Promise.withResolvers(); + let completionSettled = false; + const settle = (error?: unknown) => { + if (completionSettled) return; + completionSettled = true; + if (error !== undefined) completion.reject(error); + else if (terminalError) completion.reject(terminalError); + else if (!turnEnded) + completion.reject(new Error("Cursor stream ended before turnEnded")); + else completion.resolve(); + }; + // Abort can reject completion while we are still awaiting response + // headers; keep the rejection observed so Node does not report it as + // unhandled when the catch path never reaches `await completion.promise`. + void completion.promise.catch(() => {}); + const responseReady = Promise.withResolvers(); + let responseSeen = false; + let responseStatus = 0; + let responseHeaders: Record = {}; + let responseReadySettled = false; + const rejectResponseReady = (error: unknown) => { + if (responseReadySettled) return; + responseReadySettled = true; + responseReady.reject(error); + }; + const resolveResponseReady = () => { + if (responseReadySettled) return; + responseReadySettled = true; + responseReady.resolve(); + }; + const clearIdleTimer = () => { + if (idleTimer) clearTimeout(idleTimer); + idleTimer = undefined; + }; + const armIdleTimer = () => { + clearIdleTimer(); + if (timeoutMs === undefined || timeoutMs === 0) return; + idleTimer = setTimeout(() => { + const error = new Error( + `Cursor request idle timeout after ${Math.floor(timeoutMs)}ms`, + ); + rejectResponseReady(error); + settle(error); + h2Request?.close(http2.constants.NGHTTP2_CANCEL); + }, Math.floor(timeoutMs)); + }; + let frameBuffer: Buffer = Buffer.alloc(0); + const processFrame = (flags: number, bytes: Uint8Array) => { + if ((flags & CONNECT_COMPRESSED_FLAG) !== 0) { + throw new Error("Compressed Cursor Connect frames are unsupported"); + } + if ((flags & CONNECT_END_STREAM_FLAG) !== 0) { + terminalError = errorFromEndStream(bytes); + if (terminalError) h2Request?.close(); + return; + } + const message = fromBinary(AgentServerMessageSchema, bytes); + if (message.message.case === "execServerMessage") { + const exec = message.message.value; + if (exec.message.case === "requestContextArgs") { + const result = create(RequestContextResultSchema, { + result: { + case: "success", + value: create(RequestContextSuccessSchema, { + requestContext: create(RequestContextSchema, { + rules: buildCursorRequestContextRules(context.systemPrompt), + tools: [], + }), + }), + }, + }); + const response = create(ExecClientMessageSchema, { + id: exec.id, + execId: exec.execId, + message: { case: "requestContextResult", value: result }, + }); + const envelope = create(AgentClientMessageSchema, { + message: { case: "execClientMessage", value: response }, + }); + h2Request?.write( + frameConnectMessage(toBinary(AgentClientMessageSchema, envelope)), + ); + return; + } + const throwReply = create(AgentClientMessageSchema, { + message: { + case: "execClientControlMessage", + value: create(ExecClientControlMessageSchema, { + message: { + case: "throw", + value: create(ExecClientThrowSchema, { + id: exec.id, + error: + "Cursor tools are not available in this chat-only provider", + errorCode: "UNIMPLEMENTED", + }), + }, + }), + }, + }); + const closeReply = create(AgentClientMessageSchema, { + message: { + case: "execClientControlMessage", + value: create(ExecClientControlMessageSchema, { + message: { + case: "streamClose", + value: create(ExecClientStreamCloseSchema, { id: exec.id }), + }, + }), + }, + }); + const error = new Error( + "Cursor requested a tool that is unavailable in chat-only mode", + ); + terminalError = error; + if (!h2Request) { + settle(error); + return; + } + h2Request.write( + frameConnectMessage(toBinary(AgentClientMessageSchema, throwReply)), + ); + h2Request.write( + frameConnectMessage(toBinary(AgentClientMessageSchema, closeReply)), + () => settle(error), + ); + return; + } + if (message.message.case === "kvServerMessage") { + sendKvReply(message.message.value, built.blobStore, h2Request); + return; + } + if (message.message.case === "interactionQuery") { + throw new Error( + `Cursor interaction query ${message.message.value.query.case ?? "unknown"} is unavailable in chat-only mode`, + ); + } + if (message.message.case !== "interactionUpdate") return; + processInteraction( + message.message.value, + output, + stream, + () => { + turnEnded = true; + }, + { + setText(value) { + currentText = value; + }, + getText() { + return currentText; + }, + setThinking(value) { + currentThinking = value; + }, + getThinking() { + return currentThinking; + }, + closeBlocks, + }, + ); + }; + const processData = (chunk: Buffer) => { + frameBuffer = + frameBuffer.length === 0 + ? chunk + : Buffer.concat([frameBuffer, chunk]); + while (frameBuffer.length >= 5) { + const size = frameBuffer.readUInt32BE(1); + if (size > MAX_CONNECT_FRAME_BYTES) { + throw new Error( + `Cursor Connect frame exceeds ${MAX_CONNECT_FRAME_BYTES} bytes`, + ); + } + if (frameBuffer.length < size + 5) return; + const flags = frameBuffer[0]!; + const data = frameBuffer.subarray(5, size + 5); + frameBuffer = frameBuffer.subarray(size + 5); + processFrame(flags, data); + } + }; + + h2Client = await connectCursorHttp2(baseUrl, { + signal: options?.signal, + timeoutMs: PROXY_TUNNEL_TIMEOUT_MS, + }); + h2Client.once("error", (error) => { + rejectResponseReady(error); + settle(error); + }); + h2Request = h2Client.request(cursorHeaders(apiKey, options)); + h2Request.once("response", (headers) => { + armIdleTimer(); + responseSeen = true; + responseStatus = Number(headers[":status"] ?? 0); + responseHeaders = headerRecord(headers); + resolveResponseReady(); + }); + h2Request.on("trailers", (trailers) => { + const status = String(trailers["grpc-status"] ?? "0"); + if (status !== "0") { + terminalError = new Error( + `Cursor gRPC error ${status}: ${decodeURIComponent(String(trailers["grpc-message"] ?? ""))}`, + ); + } + }); + const responseCallback = responseReady.promise.then(async () => { + await options?.onResponse?.( + { status: responseStatus, headers: responseHeaders }, + model, + ); + if (responseStatus < 200 || responseStatus >= 300) { + throw new Error( + `Cursor AgentService request failed with HTTP ${responseStatus}`, + ); + } + stream.push({ type: "start", partial: output }); + }); + // Keep rejected transport/callback promises observed even when the peer + // closes immediately after a malformed or unsupported interaction. + void responseReady.promise.catch(() => {}); + void responseCallback.catch(() => {}); + let dataChain = Promise.resolve(); + h2Request.on("data", (chunk: Buffer) => { + armIdleTimer(); + dataChain = dataChain + .then(() => responseCallback) + .then(() => processData(chunk)) + .catch((error) => { + settle(error); + }); + }); + h2Request.once("end", () => { + clearIdleTimer(); + if (!responseSeen) { + rejectResponseReady( + new Error("Cursor response headers were not received"), + ); + } + void dataChain + .then(() => responseCallback) + .then(() => { + if (!responseSeen) + throw new Error("Cursor response headers were not received"); + if (frameBuffer.length !== 0) + throw new Error("Incomplete Cursor Connect frame"); + settle(); + }) + .catch((error) => settle(error)); + }); + h2Request.once("error", (error) => { + rejectResponseReady(error); + settle(error); + }); + h2Request.once("aborted", () => { + const error = new Error("Cursor response aborted"); + rejectResponseReady(error); + settle(error); + }); + const sendHeartbeat = () => { + if (!h2Request || h2Request.closed || h2Request.destroyed) return; + const message = create(AgentClientMessageSchema, { + message: { + case: "clientHeartbeat", + value: create(ClientHeartbeatSchema, {}), + }, + }); + try { + h2Request.write( + frameConnectMessage(toBinary(AgentClientMessageSchema, message)), + ); + } catch { + // The terminal request/error handler owns stream completion. + } + }; + heartbeat = setInterval(sendHeartbeat, HEARTBEAT_INTERVAL_MS); + if (options?.signal) { + const onAbort = () => { + const error = new Error("Cursor request aborted"); + rejectResponseReady(error); + h2Request?.close(http2.constants.NGHTTP2_CANCEL); + settle(error); + }; + if (options.signal.aborted) onAbort(); + else { + options.signal.addEventListener("abort", onAbort, { once: true }); + removeAbortListener = () => + options.signal?.removeEventListener("abort", onAbort); + } + } + armIdleTimer(); + h2Request.write(frameConnectMessage(built.requestBytes)); + await responseCallback; + await completion.promise; + if (heartbeat) clearInterval(heartbeat); + clearIdleTimer(); + removeAbortListener?.(); + h2Request.close(); + h2Client.close(); + closeBlocks(); + if (output.stopReason === "pending") output.stopReason = "stop"; + output.usage.totalTokens = output.usage.input + output.usage.output; + stream.push({ + type: "done", + reason: output.stopReason === "length" ? "length" : "stop", + message: output, + }); + stream.end(); + } catch (error) { + if (heartbeat) clearInterval(heartbeat); + if (idleTimer) clearTimeout(idleTimer); + removeAbortListener?.(); + h2Request?.close(); + h2Client?.close(); + finishError(error); + } + })().catch((error) => { + // The body above handles all expected failures; this guard also protects + // the event stream from an unexpected asynchronous callback rejection. + stream.push({ + type: "error", + reason: "error", + error: { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: emptyUsage(), + stopReason: "error", + errorMessage: error instanceof Error ? error.message : String(error), + timestamp: Date.now(), + }, + }); + stream.end(); + }); + return stream; +} + +interface InteractionState { + setText( + value: + | Extract + | undefined, + ): void; + getText(): + | Extract + | undefined; + setThinking( + value: + | Extract + | undefined, + ): void; + getThinking(): + | Extract + | undefined; + closeBlocks(): void; +} + +function processInteraction( + update: InteractionUpdate, + output: AssistantMessage, + stream: AssistantMessageEventStream, + onTurnEnded: () => void, + state: InteractionState, +): void { + switch (update.message.case) { + case "textDelta": { + const thinking = state.getThinking(); + if (thinking) { + const index = output.content.indexOf(thinking); + stream.push({ + type: "thinking_end", + contentIndex: index, + content: thinking.thinking, + partial: output, + }); + state.setThinking(undefined); + } + const delta = update.message.value.text; + if (!delta) return; + let block = state.getText(); + if (!block) { + block = { type: "text", text: "" }; + output.content.push(block); + state.setText(block); + stream.push({ + type: "text_start", + contentIndex: output.content.length - 1, + partial: output, + }); + } + block.text += delta; + stream.push({ + type: "text_delta", + contentIndex: output.content.indexOf(block), + delta, + partial: output, + }); + break; + } + case "thinkingDelta": { + const delta = update.message.value.text; + if (!delta) return; + const text = state.getText(); + if (text) { + const index = output.content.indexOf(text); + stream.push({ + type: "text_end", + contentIndex: index, + content: text.text, + partial: output, + }); + state.setText(undefined); + } + let block = state.getThinking(); + if (!block) { + block = { type: "thinking", thinking: "" }; + output.content.push(block); + state.setThinking(block); + stream.push({ + type: "thinking_start", + contentIndex: output.content.length - 1, + partial: output, + }); + } + block.thinking += delta; + stream.push({ + type: "thinking_delta", + contentIndex: output.content.indexOf(block), + delta, + partial: output, + }); + break; + } + case "thinkingCompleted": { + const block = state.getThinking(); + if (!block) return; + const index = output.content.indexOf(block); + stream.push({ + type: "thinking_end", + contentIndex: index, + content: block.thinking, + partial: output, + }); + state.setThinking(undefined); + break; + } + case "partialToolCall": + case "toolCallDelta": + case "toolCallStarted": + case "toolCallCompleted": + throw new Error( + `Cursor ${update.message.case} is unavailable in chat-only mode`, + ); + case "tokenDelta": { + // Cursor only reports generated tokens here, not the complete context + // usage Pi needs for context accounting. Keep the usage block empty; + // Pi 0.84.3+ estimates the full history for threshold compaction. + break; + } + case "turnEnded": + onTurnEnded(); + break; + case "heartbeat": + case undefined: + break; + } +} + +function sendKvReply( + message: KvServerMessage, + store: CursorBlobStore, + request: http2.ClientHttp2Stream | undefined, +): void { + if (!request || request.closed || request.destroyed) return; + let reply; + if (message.message.case === "getBlobArgs") { + const key = Buffer.from(message.message.value.blobId).toString("hex"); + reply = create(KvClientMessageSchema, { + id: message.id, + message: { + case: "getBlobResult", + value: create(GetBlobResultSchema, { blobData: store.get(key) }), + }, + }); + } else if (message.message.case === "setBlobArgs") { + const args = message.message.value; + store.set(Buffer.from(args.blobId).toString("hex"), args.blobData); + reply = create(KvClientMessageSchema, { + id: message.id, + message: { + case: "setBlobResult", + value: create(SetBlobResultSchema, {}), + }, + }); + } else { + return; + } + const envelope = create(AgentClientMessageSchema, { + message: { case: "kvClientMessage", value: reply }, + }); + try { + request.write( + frameConnectMessage(toBinary(AgentClientMessageSchema, envelope)), + ); + } catch { + // The owning stream listener reports the transport failure. + } +} diff --git a/extensions/ai-providers/cursor/proxy.ts b/extensions/ai-providers/cursor/proxy.ts new file mode 100644 index 00000000..8811bd01 --- /dev/null +++ b/extensions/ai-providers/cursor/proxy.ts @@ -0,0 +1,203 @@ +import * as http2 from "node:http2"; +import * as net from "node:net"; +import * as tls from "node:tls"; + +export interface CursorHttp2ConnectOptions { + signal?: AbortSignal; + timeoutMs?: number; +} + +function isLocalOrMetadataHost(hostname: string): boolean { + const host = hostname.toLowerCase().replace(/^\[|\]$/g, ""); + if ( + host === "localhost" || + host.endsWith(".localhost") || + host === "metadata.google.internal" + ) { + return true; + } + if (host === "::" || host === "::1" || /^f[cd]/.test(host)) return true; + const ipv4 = /^(\d{1,3})\.(\d{1,3})\./.exec(host); + if (!ipv4) return false; + const first = Number(ipv4[1]); + const second = Number(ipv4[2]); + return ( + first === 0 || + first === 10 || + first === 127 || + (first === 169 && second === 254) || + (first === 172 && second >= 16 && second <= 31) || + (first === 192 && second === 168) + ); +} + +function shouldBypassProxy(target: URL): boolean { + if (isLocalOrMetadataHost(target.hostname)) return true; + const noProxy = process.env.NO_PROXY || process.env.no_proxy; + if (!noProxy) return false; + const targetHost = target.hostname.toLowerCase().replace(/^\[|\]$/g, ""); + const targetPort = + target.port || (target.protocol === "https:" ? "443" : "80"); + for (const rawRule of noProxy.split(/[,\s]+/)) { + let rule = rawRule.trim().toLowerCase(); + if (!rule) continue; + if (rule === "*") return true; + let rulePort: string | undefined; + const portMatch = /^(\[[^\]]+\]|[^:]+):(\d+)$/.exec(rule); + if (portMatch) { + rule = portMatch[1]!; + rulePort = portMatch[2]; + } + if (rulePort && rulePort !== targetPort) continue; + rule = rule.replace(/^\[|\]$/g, "").replace(/^\./, ""); + if (targetHost === rule || targetHost.endsWith(`.${rule}`)) return true; + } + return false; +} + +/** Resolve Cursor's provider override first, then the standard proxy variables. */ +export function resolveCursorProxy(target: URL): string | undefined { + if (shouldBypassProxy(target)) return undefined; + const protocolProxy = + target.protocol === "https:" + ? process.env.HTTPS_PROXY || process.env.https_proxy + : process.env.HTTP_PROXY || process.env.http_proxy; + return [ + process.env.PI_PROXY_CURSOR, + process.env.PI_PROXY, + protocolProxy, + process.env.ALL_PROXY || process.env.all_proxy, + ] + .map((value) => value?.trim()) + .find((value): value is string => Boolean(value)); +} + +function connectProxyTunnel( + proxyUrl: URL, + targetUrl: URL, + options: CursorHttp2ConnectOptions, +): Promise { + if (!["http:", "https:"].includes(proxyUrl.protocol)) { + return Promise.reject( + new Error(`Unsupported Cursor proxy protocol: ${proxyUrl.protocol}`), + ); + } + if (options.signal?.aborted) { + return Promise.reject(new Error("Cursor proxy tunnel aborted")); + } + const proxyTls = proxyUrl.protocol === "https:"; + const proxyPort = Number(proxyUrl.port || (proxyTls ? 443 : 80)); + const targetPort = Number( + targetUrl.port || (targetUrl.protocol === "https:" ? 443 : 80), + ); + const targetAuthority = `${targetUrl.hostname}:${targetPort}`; + const { promise, resolve, reject } = Promise.withResolvers(); + let rawSocket: net.Socket | undefined; + let targetSocket: net.Socket | undefined; + let timer: ReturnType | undefined; + let response = Buffer.alloc(0); + let settled = false; + + const cleanup = () => { + if (timer) clearTimeout(timer); + options.signal?.removeEventListener("abort", onAbort); + rawSocket?.removeListener("error", onError); + rawSocket?.removeListener(proxyTls ? "secureConnect" : "connect", onReady); + rawSocket?.removeListener("data", onData); + targetSocket?.removeListener("error", onError); + targetSocket?.removeListener("secureConnect", onTargetReady); + }; + const fail = (error: Error) => { + if (settled) return; + settled = true; + cleanup(); + targetSocket?.destroy(); + rawSocket?.destroy(); + reject(error); + }; + const succeed = (socket: net.Socket) => { + if (settled) return; + settled = true; + cleanup(); + resolve(socket); + }; + const onAbort = () => fail(new Error("Cursor proxy tunnel aborted")); + const onError = (error: Error) => fail(error); + const onTargetReady = () => { + if (targetSocket) succeed(targetSocket); + }; + const onData = (chunk: Buffer) => { + if (!rawSocket) return; + response = Buffer.concat([response, chunk]); + if (response.length > 64 * 1024) { + fail(new Error("Cursor proxy response headers exceed 64 KiB")); + return; + } + const headerEnd = response.indexOf("\r\n\r\n"); + if (headerEnd === -1) return; + const statusLine = response + .subarray(0, headerEnd) + .toString("latin1") + .split("\r\n")[0]; + if (!/^HTTP\/1\.[01] 200\b/.test(statusLine ?? "")) { + fail( + new Error( + `Cursor proxy tunnel failed: ${statusLine || "invalid response"}`, + ), + ); + return; + } + rawSocket.removeListener("data", onData); + if (targetUrl.protocol !== "https:") { + succeed(rawSocket); + return; + } + targetSocket = tls.connect({ + socket: rawSocket, + servername: targetUrl.hostname, + ALPNProtocols: ["h2"], + }); + targetSocket.once("error", onError); + targetSocket.once("secureConnect", onTargetReady); + }; + const onReady = () => { + if (!rawSocket) return; + let request = `CONNECT ${targetAuthority} HTTP/1.1\r\nHost: ${targetAuthority}\r\n`; + if (proxyUrl.username || proxyUrl.password) { + const credentials = Buffer.from( + `${decodeURIComponent(proxyUrl.username)}:${decodeURIComponent(proxyUrl.password)}`, + ).toString("base64"); + request += `Proxy-Authorization: Basic ${credentials}\r\n`; + } + rawSocket.on("data", onData); + rawSocket.write(`${request}\r\n`); + }; + + options.signal?.addEventListener("abort", onAbort, { once: true }); + if (options.timeoutMs !== undefined && options.timeoutMs > 0) { + const timeoutMs = Math.floor(options.timeoutMs); + timer = setTimeout( + () => + fail(new Error(`Cursor proxy tunnel timed out after ${timeoutMs}ms`)), + timeoutMs, + ); + } + rawSocket = proxyTls + ? tls.connect({ host: proxyUrl.hostname, port: proxyPort }) + : net.connect({ host: proxyUrl.hostname, port: proxyPort }); + rawSocket.once("error", onError); + rawSocket.once(proxyTls ? "secureConnect" : "connect", onReady); + return promise; +} + +/** Open Cursor's HTTP/2 session directly or through an HTTP CONNECT proxy. */ +export async function connectCursorHttp2( + baseUrl: string, + options: CursorHttp2ConnectOptions = {}, +): Promise { + const target = new URL(baseUrl); + const proxy = resolveCursorProxy(target); + if (!proxy) return http2.connect(target); + const socket = await connectProxyTunnel(new URL(proxy), target, options); + return http2.connect(target, { createConnection: () => socket }); +} diff --git a/extensions/ai-providers/cursor/with-resolvers.d.ts b/extensions/ai-providers/cursor/with-resolvers.d.ts new file mode 100644 index 00000000..49d3055e --- /dev/null +++ b/extensions/ai-providers/cursor/with-resolvers.d.ts @@ -0,0 +1,12 @@ +/** Node >=22 provides Promise.withResolvers; the repo's ES2022 lib needs a shim. */ +declare global { + interface PromiseConstructor { + withResolvers(): { + promise: Promise; + resolve: (value: T | PromiseLike) => void; + reject: (reason?: unknown) => void; + }; + } +} + +export {}; diff --git a/extensions/ai-providers/index.ts b/extensions/ai-providers/index.ts new file mode 100644 index 00000000..f96438e6 --- /dev/null +++ b/extensions/ai-providers/index.ts @@ -0,0 +1,86 @@ +/** + * ai-providers — OAuth-backed model providers for pi. + * + * Adds OAuth-backed Google Antigravity and Cursor model providers. Both are + * inert until the user logs in and selects one of their models. Cursor uses + * AgentService/Run in deliberately chat-only mode: Cursor-native coding tools + * are not exposed or executed by this extension. + * + * Wire protocol: Cloud Code Assist `v1internal:streamGenerateContent` over + * SSE (see antigravity/provider.ts). Reference implementation: oh-my-pi's + * google-gemini-cli provider (shared google-gemini-cli/google-antigravity). + */ + +import { createProvider, type ProviderStreams } from "@earendil-works/pi-ai"; +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { createOAuthAuth } from "./oauth-adapter.ts"; +import { encodeApiKey } from "./antigravity/credentials.ts"; +import { fetchAntigravityModels } from "./antigravity/discovery.ts"; +import { + ANTIGRAVITY_API_URL, + ANTIGRAVITY_MODELS, +} from "./antigravity/models.ts"; +import { + loginAntigravity, + refreshAntigravityToken, +} from "./antigravity/oauth.ts"; +import { streamAntigravity } from "./antigravity/provider.ts"; +import { getCursorApiKey } from "./cursor/credentials.ts"; +import { fetchCursorModels } from "./cursor/discovery.ts"; +import { transformCursorImageInput } from "./cursor/input-images.ts"; +import { CURSOR_MODELS } from "./cursor/models.ts"; +import { loginCursor, refreshCursorToken } from "./cursor/oauth.ts"; +import { streamCursor } from "./cursor/provider.ts"; + +function providerStreams( + streamSimple: ProviderStreams["streamSimple"], +): ProviderStreams { + return { + stream: (model, context, options) => streamSimple(model, context, options), + streamSimple, + }; +} + +export default function authProviders(pi: ExtensionAPI) { + pi.on("input", transformCursorImageInput); + + pi.registerProvider( + createProvider({ + id: "google-antigravity", + name: "Google Antigravity", + baseUrl: ANTIGRAVITY_API_URL, + api: providerStreams(streamAntigravity), + auth: { + oauth: createOAuthAuth({ + name: "Google (Antigravity)", + isSubscription: true, + login: loginAntigravity, + refreshToken: refreshAntigravityToken, + getApiKey: encodeApiKey, + }), + }, + models: ANTIGRAVITY_MODELS, + fetchModels: fetchAntigravityModels, + }), + ); + + pi.registerProvider( + createProvider({ + id: "cursor", + name: "Cursor", + baseUrl: "https://api2.cursor.sh", + api: providerStreams(streamCursor), + auth: { + oauth: createOAuthAuth({ + name: "Cursor", + isSubscription: true, + login: loginCursor, + refreshToken: refreshCursorToken, + getApiKey: getCursorApiKey, + }), + }, + models: CURSOR_MODELS, + fetchModels: fetchCursorModels, + }), + ); +} diff --git a/extensions/ai-providers/oauth-adapter.ts b/extensions/ai-providers/oauth-adapter.ts new file mode 100644 index 00000000..0162b946 --- /dev/null +++ b/extensions/ai-providers/oauth-adapter.ts @@ -0,0 +1,81 @@ +import type { + ModelAuth, + OAuthAuth, + OAuthCredential, + OAuthCredentials, + OAuthLoginCallbacks, + ProviderAuthInteraction, +} from "@earendil-works/pi-ai"; + +interface LegacyOAuthImplementation { + name: string; + isSubscription?: boolean; + login(callbacks: CancellableOAuthLoginCallbacks): Promise; + refreshToken( + credential: OAuthCredentials, + signal: AbortSignal, + ): Promise; + getApiKey(credential: OAuthCredentials): string | Promise; +} + +export type CancellableOAuthLoginCallbacks = Omit< + OAuthLoginCallbacks, + "onManualCodeInput" +> & { + onManualCodeInput?(signal?: AbortSignal): Promise; +}; + +function legacyCallbacks( + interaction: ProviderAuthInteraction, +): CancellableOAuthLoginCallbacks { + return { + signal: interaction.signal, + onAuth: (info) => interaction.notify({ type: "auth_url", ...info }), + onDeviceCode: (info) => + interaction.notify({ type: "device_code", ...info }), + onProgress: (message) => interaction.notify({ type: "progress", message }), + onPrompt: (prompt) => + interaction.prompt({ + type: "text", + message: prompt.message, + placeholder: prompt.placeholder, + }), + onManualCodeInput: (signal) => + interaction.prompt({ + type: "manual_code", + message: "Paste the authorization callback URL or code", + signal, + }), + onSelect: (prompt) => + interaction.prompt({ + type: "select", + message: prompt.message, + options: prompt.options, + }), + }; +} + +function canonicalCredential(credentials: OAuthCredentials): OAuthCredential { + return { ...credentials, type: "oauth" }; +} + +/** Adapt pi's retained extension OAuth callbacks to the native Provider API. */ +export function createOAuthAuth( + implementation: LegacyOAuthImplementation, +): OAuthAuth { + return { + name: implementation.name, + isSubscription: implementation.isSubscription, + login: async (interaction) => + canonicalCredential( + await implementation.login(legacyCallbacks(interaction)), + ), + refresh: async (credential, signal) => + canonicalCredential( + await implementation.refreshToken(credential, signal), + ), + toAuth: async (credential): Promise => ({ + apiKey: await implementation.getApiKey(credential), + }), + }; +} diff --git a/tests/extensions/ai-providers/antigravity.test.ts b/tests/extensions/ai-providers/antigravity.test.ts new file mode 100644 index 00000000..c05b8cad --- /dev/null +++ b/tests/extensions/ai-providers/antigravity.test.ts @@ -0,0 +1,1418 @@ +/** + * Behavioral tests for the Antigravity provider: request envelope shape, + * SSE-to-event mapping, endpoint failover, and credential codec. fetch is + * mocked; no network access. + */ + +import assert from "node:assert/strict"; +import * as http from "node:http"; +import { after, before, test } from "node:test"; +import type { + Api, + AssistantMessageEventStream, + Context, + Model, +} from "@earendil-works/pi-ai/compat"; +import { + type AntigravityCredentials, + decodeApiKey, + encodeApiKey, +} from "../../../extensions/ai-providers/antigravity/credentials.ts"; +import { fetchAntigravityModels } from "../../../extensions/ai-providers/antigravity/discovery.ts"; +import { + convertMessages, + isThinkingPart, + mapStopReasonString, + retainThoughtSignature, +} from "../../../extensions/ai-providers/antigravity/google-conversion.ts"; +import { ANTIGRAVITY_MODELS } from "../../../extensions/ai-providers/antigravity/models.ts"; +import { + loginAntigravity, + refreshAntigravityToken, +} from "../../../extensions/ai-providers/antigravity/oauth.ts"; +import { + buildRequestBody, + sanitizeSchemaForCca, + streamAntigravity, +} from "../../../extensions/ai-providers/antigravity/provider.ts"; +import { collapseAntigravityModels } from "../../../extensions/ai-providers/antigravity/routing.ts"; + +const GEMINI_MODEL: Model = { + id: "gemini-3.1-pro", + name: "Gemini 3.1 Pro (Antigravity)", + api: "antigravity-cloudcode", + provider: "google-antigravity", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + reasoning: true, + input: ["text", "image"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 64_000, +}; + +const CLAUDE_MODEL: Model = { + ...GEMINI_MODEL, + id: "claude-sonnet-4-6", + name: "Claude Sonnet 4.6 (Antigravity)", +}; + +const SIMPLE_CONTEXT: Context = { + systemPrompt: "You are helpful.", + messages: [{ role: "user", content: "hello", timestamp: 0 }], +}; + +const API_KEY = encodeApiKey({ + refresh: "r", + access: "tok", + expires: 0, + projectId: "proj-1", +}); + +const ZERO_USAGE = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +function sseResponse(events: unknown[]): Response { + const payload = events + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(payload)); + controller.close(); + }, + }); + return new Response(body, { status: 200 }); +} + +async function collectEvents(stream: AssistantMessageEventStream) { + const events = []; + for await (const event of stream) events.push(event); + return events; +} + +function getStatus(url: URL): Promise { + const { promise, resolve, reject } = Promise.withResolvers(); + const request = http.get(url, (response) => { + response.resume(); + response.once("end", () => resolve(response.statusCode ?? 0)); + }); + request.once("error", reject); + return promise; +} + +// --- credentials ----------------------------------------------------------- + +test("credentials codec round-trips and tolerates bare tokens", () => { + const encoded = encodeApiKey({ + refresh: "r", + access: "a", + expires: 1, + projectId: "p", + }); + assert.deepEqual(decodeApiKey(encoded), { token: "a", projectId: "p" }); + assert.deepEqual(decodeApiKey("bare-token"), { token: "bare-token" }); + assert.deepEqual(decodeApiKey("{not json"), { token: "{not json" }); +}); + +test("local Google conversion preserves only valid same-model signatures", () => { + const contents = convertMessages(GEMINI_MODEL, { + messages: [ + { + role: "assistant", + api: GEMINI_MODEL.api, + provider: GEMINI_MODEL.provider, + model: GEMINI_MODEL.id, + content: [ + { type: "text", text: "", textSignature: "YWJjZA==" }, + { + type: "thinking", + thinking: "reasoning", + thinkingSignature: "not-base64", + }, + ], + usage: ZERO_USAGE, + stopReason: "stop", + timestamp: 0, + }, + ], + }); + + assert.deepEqual(contents, [ + { + role: "model", + parts: [ + { text: "", thoughtSignature: "YWJjZA==" }, + { thought: true, text: "reasoning" }, + ], + }, + ]); + + const crossModel = convertMessages(GEMINI_MODEL, { + messages: [ + { + role: "assistant", + api: GEMINI_MODEL.api, + provider: GEMINI_MODEL.provider, + model: "another-model", + content: [ + { + type: "thinking", + thinking: "reasoning", + thinkingSignature: "YWJjZA==", + }, + { + type: "toolCall", + id: "call|with spaces", + name: "read_file", + arguments: {}, + thoughtSignature: "YWJjZA==", + }, + ], + usage: ZERO_USAGE, + stopReason: "toolUse", + timestamp: 0, + }, + { + role: "toolResult", + toolCallId: "call|with spaces", + toolName: "read_file", + content: [{ type: "text", text: "ok" }], + isError: false, + timestamp: 0, + }, + ], + }); + assert.deepEqual(crossModel, [ + { + role: "model", + parts: [ + { text: "reasoning" }, + { + functionCall: { + id: "call_with_spaces", + name: "read_file", + args: {}, + }, + }, + ], + }, + { + role: "user", + parts: [ + { + functionResponse: { + id: "call_with_spaces", + name: "read_file", + response: { output: "ok" }, + }, + }, + ], + }, + ]); +}); + +test("local Google conversion repairs orphaned calls and routes tool images", () => { + const orphaned = convertMessages(GEMINI_MODEL, { + messages: [ + { + role: "assistant", + api: GEMINI_MODEL.api, + provider: GEMINI_MODEL.provider, + model: GEMINI_MODEL.id, + content: [ + { type: "toolCall", id: "call-1", name: "capture", arguments: {} }, + ], + usage: ZERO_USAGE, + stopReason: "toolUse", + timestamp: 0, + }, + { role: "user", content: "continue", timestamp: 1 }, + ], + }); + assert.deepEqual(orphaned[1], { + role: "user", + parts: [ + { + functionResponse: { + id: "call-1", + name: "capture", + response: { error: "No result provided" }, + }, + }, + ], + }); + + const imageResult = { + role: "toolResult" as const, + toolCallId: "call-2", + toolName: "capture", + content: [{ type: "image" as const, mimeType: "image/png", data: "AA==" }], + isError: false, + timestamp: 0, + }; + const gemini3 = convertMessages(GEMINI_MODEL, { messages: [imageResult] }); + assert.deepEqual(gemini3[0]?.parts[0]?.functionResponse?.parts, [ + { inlineData: { mimeType: "image/png", data: "AA==" } }, + ]); + + const gemini2 = convertMessages( + { ...GEMINI_MODEL, id: "gemini-2.5-pro" }, + { messages: [imageResult] }, + ); + assert.equal(gemini2.length, 2); + assert.equal(gemini2[0]?.parts[0]?.functionResponse?.parts, undefined); + assert.deepEqual(gemini2[1], { + role: "user", + parts: [ + { text: "Tool result image:" }, + { inlineData: { mimeType: "image/png", data: "AA==" } }, + ], + }); +}); + +test("local Google stream helpers preserve protocol semantics", () => { + assert.equal( + isThinkingPart({ thought: true, thoughtSignature: "sig" }), + true, + ); + assert.equal(isThinkingPart({ thoughtSignature: "sig" }), false); + assert.equal(retainThoughtSignature("old", undefined), "old"); + assert.equal(retainThoughtSignature("old", "new"), "new"); + assert.equal(mapStopReasonString("STOP"), "stop"); + assert.equal(mapStopReasonString("MAX_TOKENS"), "length"); + assert.equal(mapStopReasonString("SAFETY"), "error"); +}); + +test("OAuth callback ignores a wrong state without consuming the real waiter", async () => { + globalThis.fetch = (async (input: string | URL | Request) => { + const url = String(input); + if (url.includes("oauth2.googleapis.com/token")) { + return Response.json({ + access_token: "access", + refresh_token: "refresh", + expires_in: 3_600, + }); + } + if (url.includes("v1internal:loadCodeAssist")) { + return Response.json({ + currentTier: { id: "free-tier" }, + cloudaicompanionProject: "project", + }); + } + if (url.includes("googleapis.com/oauth2/v1/userinfo")) { + return Response.json({ email: "user@example.test" }); + } + throw new Error(`Unexpected OAuth fetch: ${url}`); + }) as typeof fetch; + + const authReady = Promise.withResolvers(); + const manualReady = Promise.withResolvers(); + const login = loginAntigravity({ + onAuth(info) { + authReady.resolve(info.url); + }, + onDeviceCode() {}, + async onPrompt() { + return ""; + }, + async onSelect() { + return undefined; + }, + onManualCodeInput(signal) { + assert.ok(signal); + manualReady.resolve(signal); + return new Promise((_resolve, reject) => { + const onAbort = () => reject(new Error("manual prompt cancelled")); + if (signal.aborted) onAbort(); + else signal.addEventListener("abort", onAbort, { once: true }); + }); + }, + }); + const authUrl = new URL(await authReady.promise); + const manualSignal = await manualReady.promise; + const state = authUrl.searchParams.get("state"); + const redirect = authUrl.searchParams.get("redirect_uri"); + assert.ok(state && redirect); + assert.equal(manualSignal.aborted, false); + + const wrong = new URL(redirect); + wrong.searchParams.set("code", "wrong-code"); + wrong.searchParams.set("state", "wrong-state"); + assert.equal(await getStatus(wrong), 400); + assert.equal(manualSignal.aborted, false); + + const valid = new URL(redirect); + valid.searchParams.set("code", "valid-code"); + valid.searchParams.set("state", state); + assert.equal(await getStatus(valid), 200); + const credentials = await login; + assert.equal(manualSignal.aborted, true); + assert.equal(credentials.access, "access"); + assert.equal(credentials.refresh, "refresh"); + assert.equal(credentials.projectId, "project"); +}); + +test("OAuth refresh preserves provider metadata and the old refresh token", async () => { + const originalFetch = globalThis.fetch; + try { + globalThis.fetch = (async (_input, init) => { + assert.equal(init?.method, "POST"); + const body = init?.body; + assert.ok(body instanceof URLSearchParams); + assert.equal(body.get("refresh_token"), "old-refresh"); + return new Response( + JSON.stringify({ access_token: "new-access", expires_in: 3_600 }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + }) as typeof fetch; + const refreshed = await refreshAntigravityToken( + { + access: "old-access", + refresh: "old-refresh", + expires: 0, + projectId: "project-1", + email: "user@example.com", + } satisfies AntigravityCredentials, + new AbortController().signal, + ); + assert.equal(refreshed.access, "new-access"); + assert.equal(refreshed.refresh, "old-refresh"); + assert.equal((refreshed as { projectId?: string }).projectId, "project-1"); + assert.equal((refreshed as { email?: string }).email, "user@example.com"); + } finally { + globalThis.fetch = originalFetch; + } +}); + +// --- schema sanitization --------------------------------------------------- + +test("sanitizeSchemaForCca strips CCA-rejected keywords recursively", () => { + const schema = { + type: "object", + properties: { + path: { type: "string", minLength: 1, pattern: "^/" }, + lines: { type: "array", items: { type: "string", format: "uri" } }, + }, + $schema: "http://json-schema.org/draft-07/schema#", + additionalProperties: false, + }; + assert.deepEqual(sanitizeSchemaForCca(schema), { + type: "object", + properties: { + path: { type: "string", description: '{minLength: 1, pattern: "^/"}' }, + lines: { + type: "array", + items: { type: "string", description: '{format: "uri"}' }, + }, + }, + }); +}); + +test("sanitizeSchemaForCca preserves property names that match schema keywords", () => { + assert.deepEqual( + sanitizeSchemaForCca({ + type: "object", + properties: { + pattern: { + type: "string", + description: "Search pattern", + pattern: "^[a-z]+$", + }, + format: { type: "string" }, + }, + required: ["pattern"], + }), + { + type: "object", + properties: { + pattern: { + type: "string", + description: 'Search pattern\n\n{pattern: "^[a-z]+$"}', + }, + format: { type: "string" }, + }, + required: ["pattern"], + }, + ); +}); + +test("buildRequestBody strips CCA-rejected keywords and spills constraints into description", () => { + const contextWithTools = { + ...SIMPLE_CONTEXT, + tools: [ + { + name: "write_files", + description: "d", + parameters: { + type: "object", + properties: { + files: { + type: "array", + description: "paths", + items: { type: "string", deprecated: true }, + uniqueItems: true, + }, + mode: { type: "string", readOnly: false }, + pattern: { type: "string", description: "Search pattern" }, + $id: { type: "string", description: "User-defined property" }, + }, + }, + }, + ], + }; + const body = buildRequestBody( + CLAUDE_MODEL, + contextWithTools as never, + undefined, + "p", + ) as { + request: { + tools: { + functionDeclarations: { parameters: Record }[]; + }[]; + }; + }; + const parameters = body.request.tools[0]!.functionDeclarations[0]! + .parameters as { + properties: Record>; + }; + assert.deepEqual(parameters.properties.files, { + type: "array", + description: "paths\n\n{uniqueItems: true}", + items: { type: "string" }, + }); + assert.deepEqual(parameters.properties.mode, { type: "string" }); + assert.deepEqual(parameters.properties.pattern, { + type: "string", + description: "Search pattern", + }); + assert.deepEqual(parameters.properties.$id, { + type: "string", + description: "User-defined property", + }); +}); + +// --- request envelope ------------------------------------------------------ + +test("buildRequestBody produces the Antigravity envelope", () => { + const body = buildRequestBody( + CLAUDE_MODEL, + SIMPLE_CONTEXT, + { reasoning: "medium" }, + "proj-1", + ) as { + project: string; + model: string; + userAgent: string; + requestType: string; + request: { + contents: unknown[]; + systemInstruction: { role: string; parts: { text: string }[] }; + toolConfig: { functionCallingConfig: { mode: string } }; + generationConfig: { + maxOutputTokens: number; + thinkingConfig: Record; + }; + labels: Record; + sessionId: string; + }; + }; + assert.equal(body.project, "proj-1"); + assert.equal(body.model, "claude-sonnet-4-6"); + assert.equal(body.userAgent, "antigravity"); + assert.equal(body.requestType, "agent"); + assert.equal(body.request.systemInstruction.role, "user"); + assert.equal( + body.request.systemInstruction.parts[0].text, + "You are helpful.", + ); + assert.ok(Array.isArray(body.request.contents)); + assert.ok(body.request.contents.length > 0); + // Claude routes force VALIDATED even with no tools. + assert.equal(body.request.toolConfig.functionCallingConfig.mode, "VALIDATED"); + // Claude reasoning models take a token budget, not a thinking level. + assert.equal( + body.request.generationConfig.thinkingConfig.includeThoughts, + true, + ); + assert.equal( + typeof body.request.generationConfig.thinkingConfig.thinkingBudget, + "number", + ); + assert.equal(body.request.labels.used_claude, "true"); + assert.ok(body.request.sessionId.length > 0); +}); + +test("buildRequestBody routes Gemini 3.1 high away from the broken deployment", () => { + const body = buildRequestBody( + GEMINI_MODEL, + SIMPLE_CONTEXT, + { reasoning: "high" }, + "p", + ) as { + model: string; + request: { generationConfig: { thinkingConfig: Record } }; + }; + assert.equal(body.model, "gemini-pro-agent"); + assert.equal( + body.request.generationConfig.thinkingConfig.thinkingBudget, + 10_001, + ); +}); + +test("buildRequestBody keeps mandatory Gemini 3.7 thinking enabled", () => { + const body = buildRequestBody( + { ...GEMINI_MODEL, id: "gemini-3.7-flash" }, + SIMPLE_CONTEXT, + undefined, + "p", + ) as { + model: string; + request: { generationConfig: { thinkingConfig: Record } }; + }; + assert.equal(body.model, "gemini-3.7-flash-low"); + assert.equal( + body.request.generationConfig.thinkingConfig.thinkingLevel, + "LOW", + ); +}); + +test("buildRequestBody explicitly suppresses optional Claude thinking when off", () => { + const body = buildRequestBody( + CLAUDE_MODEL, + SIMPLE_CONTEXT, + undefined, + "p", + ) as { + request: { generationConfig: { thinkingConfig: Record } }; + }; + assert.deepEqual(body.request.generationConfig.thinkingConfig, { + includeThoughts: false, + thinkingBudget: 0, + }); +}); + +test("buildRequestBody honors disabled and forced tool choices", () => { + const contextWithTools = { + ...SIMPLE_CONTEXT, + tools: [ + { + name: "read_file", + description: "Read a file", + parameters: { type: "object", properties: {} }, + }, + ], + } as Context; + const disabled = buildRequestBody( + GEMINI_MODEL, + contextWithTools, + { toolChoice: "none" } as never, + "p", + ) as { request: Record }; + assert.equal(disabled.request.tools, undefined); + assert.equal(disabled.request.toolConfig, undefined); + + const forced = buildRequestBody( + GEMINI_MODEL, + contextWithTools, + { + toolChoice: { mode: "ANY", allowedFunctionNames: ["read_file"] }, + } as never, + "p", + ) as { + request: { + contents: { parts?: { text?: string }[] }[]; + toolConfig: { + functionCallingConfig: { + mode: string; + allowedFunctionNames: string[]; + }; + }; + }; + }; + assert.deepEqual(forced.request.toolConfig.functionCallingConfig, { + mode: "ANY", + allowedFunctionNames: ["read_file"], + }); + assert.match( + forced.request.contents.at(-1)?.parts?.[0]?.text ?? "", + /TOOL-ONLY TURN/, + ); +}); + +test("discovery collapses wire variants and validates advertised capabilities", async () => { + globalThis.fetch = (async () => + new Response( + JSON.stringify({ + models: { + "gemini-3.1-pro-low": { + displayName: "Gemini 3.1 Pro Low", + supportsThinking: true, + supportsImages: true, + maxTokens: -1, + maxOutputTokens: 0, + }, + "gemini-3.1-pro-high": { + displayName: "Broken high deployment", + supportsThinking: true, + }, + "text-only": { supportsImages: false }, + }, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + )) as typeof fetch; + + const models = await fetchAntigravityModels({ + allowNetwork: true, + credential: { + type: "oauth", + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }, + publish: async () => true, + signal: new AbortController().signal, + }); + assert.deepEqual( + models.map((model) => model.id), + ["gemini-3.1-pro", "text-only"], + ); + assert.deepEqual(models[0]?.input, ["text", "image"]); + assert.equal(models[0]?.contextWindow, 200_000); + assert.equal(models[0]?.maxTokens, 64_000); + assert.deepEqual(models[1]?.input, ["text"]); +}); + +test("discovery routes through the only live family member", async () => { + globalThis.fetch = (async () => + new Response( + JSON.stringify({ + models: { + "gemini-3.7-flash-high": { + supportsThinking: true, + supportsImages: true, + }, + }, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + )) as typeof fetch; + const [discovered] = await fetchAntigravityModels({ + allowNetwork: true, + credential: { + type: "oauth", + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }, + publish: async () => true, + signal: new AbortController().signal, + }); + assert.ok(discovered); + const body = buildRequestBody( + { ...GEMINI_MODEL, ...discovered }, + SIMPLE_CONTEXT, + { reasoning: "low" }, + "p", + ); + assert.equal(body.model, "gemini-3.7-flash-high"); +}); + +test("variant collapse drops a family when discovery only exposes a retired member", () => { + const collapsed = collapseAntigravityModels([ + { + id: "gemini-3.1-pro-high", + name: "retired", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 64_000, + }, + ]); + assert.deepEqual(collapsed, []); +}); + +test("Claude discovery and request output tokens are capped at 64000", async () => { + globalThis.fetch = (async () => + new Response( + JSON.stringify({ + models: { + "claude-sonnet-4-6": { + supportsThinking: true, + maxOutputTokens: 65_536, + }, + }, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + )) as typeof fetch; + + const [discovered] = await fetchAntigravityModels({ + allowNetwork: true, + credential: { + type: "oauth", + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }, + publish: async () => true, + signal: new AbortController().signal, + }); + assert.ok(discovered); + assert.equal(discovered.maxTokens, 64_000); + + const body = buildRequestBody( + { ...CLAUDE_MODEL, ...discovered, maxTokens: 65_536 }, + SIMPLE_CONTEXT, + { maxTokens: 65_536, reasoning: "medium" }, + "p", + ); + assert.equal( + (body.request as { generationConfig: { maxOutputTokens: number } }) + .generationConfig.maxOutputTokens, + 64_000, + ); +}); + +test("Antigravity discovery failure throws instead of replacing cached models", async () => { + globalThis.fetch = (async () => + new Response("unavailable", { status: 503 })) as typeof fetch; + + await assert.rejects( + fetchAntigravityModels({ + allowNetwork: true, + credential: { + type: "oauth", + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }, + publish: async () => true, + signal: new AbortController().signal, + }), + /failed on all endpoints/, + ); +}); + +test("static GPT-OSS uses its medium wire deployment", () => { + const model = ANTIGRAVITY_MODELS.find((entry) => entry.id === "gpt-oss-120b"); + assert.ok(model); + assert.deepEqual(model.input, ["text"]); + assert.equal(model.contextWindow, 131_072); + assert.equal(model.maxTokens, 32_768); + const body = buildRequestBody( + { ...GEMINI_MODEL, ...model }, + SIMPLE_CONTEXT, + undefined, + "p", + ); + assert.equal(body.model, "gpt-oss-120b-medium"); + assert.equal( + (body.request as { generationConfig: { maxOutputTokens: number } }) + .generationConfig.maxOutputTokens, + 32_768, + ); +}); + +// --- streaming ------------------------------------------------------------- + +let originalFetch: typeof globalThis.fetch; +let originalVersion: string | undefined; + +before(() => { + originalFetch = globalThis.fetch; + originalVersion = process.env.OPENPI_ANTIGRAVITY_VERSION; + process.env.OPENPI_ANTIGRAVITY_VERSION = "2.8.0"; +}); + +after(() => { + globalThis.fetch = originalFetch; + if (originalVersion === undefined) { + delete process.env.OPENPI_ANTIGRAVITY_VERSION; + } else { + process.env.OPENPI_ANTIGRAVITY_VERSION = originalVersion; + } +}); + +test("streamAntigravity maps SSE parts to pi events", async () => { + const requests: { url: string; body: string }[] = []; + globalThis.fetch = (async (input: unknown, init?: { body?: string }) => { + requests.push({ url: String(input), body: String(init?.body) }); + return sseResponse([ + { + response: { + candidates: [ + { + content: { + parts: [ + { text: "Thinking", thought: true, thoughtSignature: "c2ln" }, + { text: "Hi there" }, + { + functionCall: { name: "bash", args: { command: "ls" } }, + thoughtSignature: "c2ln", + }, + ], + }, + finishReason: "STOP", + }, + ], + usageMetadata: { + promptTokenCount: 120, + cachedContentTokenCount: 20, + candidatesTokenCount: 5, + thoughtsTokenCount: 10, + totalTokenCount: 135, + }, + }, + }, + ]); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(CLAUDE_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + + const types = events.map((event) => event.type); + assert.deepEqual(types, [ + "start", + "thinking_start", + "thinking_delta", + "thinking_end", + "text_start", + "text_delta", + "text_end", + "toolcall_start", + "toolcall_delta", + "toolcall_end", + "done", + ]); + + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done"); + assert.equal(done.reason, "toolUse"); + const message = done.message; + assert.equal(message.provider, "google-antigravity"); + assert.equal(message.stopReason, "toolUse"); + const toolCall = message.content.find((b) => b.type === "toolCall"); + assert.ok(toolCall && toolCall.type === "toolCall"); + assert.equal(toolCall.name, "bash"); + assert.deepEqual(toolCall.arguments, { command: "ls" }); + assert.equal(message.usage.input, 100); + assert.equal(message.usage.cacheRead, 20); + assert.equal(message.usage.output, 15); + + assert.equal(requests.length, 1); + assert.match(requests[0].url, /daily-cloudcode-pa\.googleapis\.com/); + const sent = JSON.parse(requests[0].body); + assert.equal(sent.project, "proj-1"); +}); + +test("streamAntigravity fails over to the sandbox endpoint on 5xx", async () => { + const urls: string[] = []; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + if (urls.length === 1) { + return new Response("boom", { status: 500 }); + } + return sseResponse([ + { + response: { + candidates: [ + { content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }, + ], + }, + }, + ]); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(urls.length, 2); + assert.match(urls[1], /sandbox/); + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done" && done.reason === "stop"); +}); + +test("streamAntigravity bounds a stalled non-2xx body and fails over", async () => { + let requests = 0; + const server = http.createServer((_request, response) => { + requests++; + if (requests === 1) { + response.writeHead(503, { "Content-Type": "text/plain" }); + response.write("partial error"); + return; + } + response.writeHead(200, { "Content-Type": "text/event-stream" }); + response.end( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"fallback"}]},"finishReason":"STOP"}]}}\n\n', + ); + }); + const listening = Promise.withResolvers(); + server.once("error", listening.reject); + server.listen(0, "127.0.0.1", listening.resolve); + await listening.promise; + const address = server.address(); + assert.ok(address && typeof address === "object"); + const localUrl = `http://127.0.0.1:${address.port}/stream`; + const startedAt = Date.now(); + let events; + try { + events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + timeoutMs: 50, + fetch: (_input, init) => originalFetch(localUrl, init), + }), + ); + } finally { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); + } + assert.ok(Date.now() - startedAt < 1_000); + assert.equal(requests, 2); + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done" && done.reason === "stop"); +}); + +test("streamAntigravity truncates and cancels an oversized error body", async () => { + let cancellations = 0; + globalThis.fetch = (async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("x".repeat(128 * 1024))); + }, + cancel() { + cancellations++; + }, + }), + { status: 400 }, + )) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(cancellations, 1); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.match(error.error.errorMessage ?? "", /\[truncated\]$/); + assert.ok((error.error.errorMessage?.length ?? 0) < 70 * 1024); +}); + +test("streamAntigravity fails over after a transient error in a 200 stream", async () => { + const urls: string[] = []; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + if (urls.length === 1) { + return sseResponse([ + { response: { usageMetadata: { promptTokenCount: 2 } } }, + { error: { code: 408, message: "upstream timeout" } }, + ]); + } + return sseResponse([ + { + response: { + candidates: [ + { content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }, + ], + }, + }, + ]); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(urls.length, 2); + assert.match(urls[1], /sandbox/); + assert.equal(events.filter((event) => event.type === "start").length, 1); + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done" && done.reason === "stop"); +}); + +test("streamAntigravity fails over after an empty completed stream", async () => { + const urls: string[] = []; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + if (urls.length === 1) { + return sseResponse([ + { response: { candidates: [{ finishReason: "STOP" }] } }, + ]); + } + return sseResponse([ + { + response: { + candidates: [ + { + content: { parts: [{ text: "fallback" }] }, + finishReason: "STOP", + }, + ], + }, + }, + ]); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(urls.length, 2); + assert.match(urls[1], /sandbox/); + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done" && done.reason === "stop"); +}); + +test("streamAntigravity treats thought-only completion as empty and retries", async () => { + const urls: string[] = []; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + if (urls.length === 1) { + return sseResponse([ + { + response: { + candidates: [ + { + content: { + parts: [{ text: "internal only", thought: true }], + }, + finishReason: "STOP", + }, + ], + }, + }, + ]); + } + return sseResponse([ + { + response: { + candidates: [ + { + content: { parts: [{ text: "final answer" }] }, + finishReason: "STOP", + }, + ], + }, + }, + ]); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(urls.length, 2); + assert.match(urls[1], /sandbox/); + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done"); + assert.equal(done.message.content.length, 1); + assert.equal(done.message.content[0]?.type, "text"); + assert.equal( + done.message.content[0]?.type === "text" + ? done.message.content[0].text + : undefined, + "final answer", + ); +}); + +test("streamAntigravity treats whitespace-only completion as empty and retries", async () => { + let calls = 0; + globalThis.fetch = (async () => { + calls++; + return sseResponse([ + { + response: { + candidates: [ + { + content: { + parts: [{ text: calls === 1 ? " " : "answer" }], + }, + finishReason: "STOP", + }, + ], + }, + }, + ]); + }) as typeof fetch; + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(calls, 2); + const done = events.find((event) => event.type === "done"); + assert.ok(done && done.type === "done"); + assert.equal(done.message.content[0]?.type, "text"); +}); + +test("streamAntigravity bounds the first event wait and releases both bodies", async () => { + const urls: string[] = []; + let cancellations = 0; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(": keepalive\n\n")); + }, + cancel() { + cancellations++; + }, + }), + { status: 200 }, + ); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + timeoutMs: 5, + }), + ); + assert.equal(urls.length, 2); + assert.equal(cancellations, 2); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.match(error.error.errorMessage ?? "", /first SSE event/); +}); + +test("streamAntigravity bounds the full SSE lifetime after the first event", async () => { + const urls: string[] = []; + let cancellations = 0; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + new TextEncoder().encode( + 'data: {"response":{"usageMetadata":{"promptTokenCount":1}}}\n\n', + ), + ); + }, + cancel() { + cancellations++; + }, + }), + { status: 200 }, + ); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + timeoutMs: 5, + }), + ); + assert.equal(urls.length, 2); + assert.equal(cancellations, 2); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.match(error.error.errorMessage ?? "", /next SSE event/); +}); + +test("streamAntigravity bounds the response-header wait and fails over", async () => { + let calls = 0; + globalThis.fetch = (async () => { + calls++; + return new Promise(() => {}); + }) as typeof fetch; + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + timeoutMs: 5, + }), + ); + assert.equal(calls, 2); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.match(error.error.errorMessage ?? "", /response headers/); +}); + +test("streamAntigravity carries session execution state across turns", async () => { + const bodies: Record[] = []; + let responseIndex = 0; + globalThis.fetch = (async (_input: unknown, init?: RequestInit) => { + bodies.push(JSON.parse(String(init?.body))); + responseIndex++; + return sseResponse([ + { + response: { + ...(responseIndex === 1 ? { responseId: "execution-1" } : {}), + candidates: [ + { content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }, + ], + }, + }, + ]); + }) as typeof fetch; + + const options = { apiKey: API_KEY, sessionId: "session-state-test" }; + await collectEvents(streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, options)); + await collectEvents(streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, options)); + await collectEvents(streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, options)); + assert.equal(bodies.length, 3); + const first = bodies[0] as { + requestId: string; + request: { labels: Record }; + }; + const second = bodies[1] as { + requestId: string; + request: { labels: Record }; + }; + const third = bodies[2] as { + requestId: string; + request: { labels: Record }; + }; + assert.equal(first.request.labels.last_step_index, "1"); + assert.equal(second.request.labels.last_step_index, "2"); + assert.equal(second.request.labels.last_execution_id, "execution-1"); + assert.equal(third.request.labels.last_execution_id, undefined); + assert.equal(first.requestId.split("/")[1], second.requestId.split("/")[1]); + assert.equal(first.requestId.split("/")[3], second.requestId.split("/")[3]); +}); + +test("streamAntigravity aborts and cancels after a metadata-only event", async () => { + const controller = new AbortController(); + let cancellations = 0; + globalThis.fetch = (async () => + new Response( + new ReadableStream({ + start(streamController) { + streamController.enqueue( + new TextEncoder().encode( + 'data: {"response":{"usageMetadata":{"promptTokenCount":1}}}\n\n', + ), + ); + }, + cancel() { + cancellations++; + }, + }), + { status: 200 }, + )) as typeof fetch; + setTimeout(() => controller.abort("test cancellation"), 5); + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + signal: controller.signal, + }), + ); + assert.equal(cancellations, 1); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.equal(error.reason, "aborted"); +}); + +test("streamAntigravity honors provider request lifecycle options", async () => { + let payloadCalls = 0; + let responseCalls = 0; + let sentBody: Record | undefined; + let sentHeaders: Headers | undefined; + let manifestCalls = 0; + const customFetch: typeof fetch = async (input, init) => { + if (String(input).includes("manifest/latest-arm64-mac.yml")) { + manifestCalls++; + return new Response("version: 2.8.1\n", { status: 200 }); + } + sentBody = JSON.parse(String(init?.body)); + sentHeaders = new Headers(init?.headers); + return sseResponse([ + { + response: { + candidates: [ + { content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }, + ], + }, + }, + ]); + }; + globalThis.fetch = (async () => { + throw new Error("global fetch should not be used"); + }) as typeof fetch; + + const pinnedVersion = process.env.OPENPI_ANTIGRAVITY_VERSION; + delete process.env.OPENPI_ANTIGRAVITY_VERSION; + let events; + try { + events = await collectEvents( + streamAntigravity(CLAUDE_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + fetch: customFetch, + headers: { "X-Probe": "present", "anthropic-beta": null }, + onPayload(payload) { + payloadCalls++; + return { + ...(payload as Record), + project: "override", + }; + }, + onResponse(response) { + responseCalls++; + assert.equal(response.status, 200); + }, + }), + ); + } finally { + if (pinnedVersion === undefined) { + delete process.env.OPENPI_ANTIGRAVITY_VERSION; + } else { + process.env.OPENPI_ANTIGRAVITY_VERSION = pinnedVersion; + } + } + assert.ok(events.some((event) => event.type === "done")); + assert.equal(manifestCalls, 1); + assert.equal(payloadCalls, 1); + assert.equal(responseCalls, 1); + assert.equal(sentBody?.project, "override"); + assert.equal(sentHeaders?.get("x-probe"), "present"); + assert.equal(sentHeaders?.get("authorization"), "Bearer tok"); + assert.equal(sentHeaders?.has("anthropic-beta"), false); +}); + +test("streamAntigravity surfaces 4xx as an error event without failover", async () => { + const urls: string[] = []; + globalThis.fetch = (async (input: unknown) => { + urls.push(String(input)); + return new Response("bad request", { status: 400 }); + }) as typeof fetch; + + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { apiKey: API_KEY }), + ); + assert.equal(urls.length, 1); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.match(error.error.errorMessage ?? "", /400/); +}); + +test("streamAntigravity requires a project id", async () => { + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: encodeApiKey({ refresh: "r", access: "tok", expires: 0 }), + }), + ); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.match(error.error.errorMessage ?? "", /project id/); +}); + +test("streamAntigravity reports an already-aborted request as aborted", async () => { + const controller = new AbortController(); + controller.abort("test cancellation"); + globalThis.fetch = (async (_input: unknown, init?: RequestInit) => { + assert.equal(init?.signal?.aborted, true); + throw new Error("request aborted"); + }) as typeof fetch; + const events = await collectEvents( + streamAntigravity(GEMINI_MODEL, SIMPLE_CONTEXT, { + apiKey: API_KEY, + signal: controller.signal, + }), + ); + const error = events.find((event) => event.type === "error"); + assert.ok(error && error.type === "error"); + assert.equal(error.reason, "aborted"); +}); diff --git a/tests/extensions/ai-providers/cursor.test.ts b/tests/extensions/ai-providers/cursor.test.ts new file mode 100644 index 00000000..1ba91617 --- /dev/null +++ b/tests/extensions/ai-providers/cursor.test.ts @@ -0,0 +1,1049 @@ +/** Targeted Cursor chat-only provider tests; all transport fixtures are local h2c. */ + +import assert from "node:assert/strict"; +import { once } from "node:events"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { createServer, type ServerHttp2Stream } from "node:http2"; +import { createServer as createNetServer } from "node:net"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, test } from "node:test"; +import type { + Api, + AssistantMessage, + AssistantMessageEvent, + AssistantMessageEventStream, + Context, + Model, +} from "@earendil-works/pi-ai/compat"; +import { + AgentSession, + type ExtensionContext, + VERSION as PI_VERSION, +} from "@earendil-works/pi-coding-agent"; +import { + fetchCursorModels, + fetchCursorUsableModels, +} from "../../../extensions/ai-providers/cursor/discovery.ts"; +import { transformCursorImageInput } from "../../../extensions/ai-providers/cursor/input-images.ts"; +import { CURSOR_MODELS } from "../../../extensions/ai-providers/cursor/models.ts"; +import { + generateCursorAuthParams, + getCursorTokenExpiry, + refreshCursorToken, +} from "../../../extensions/ai-providers/cursor/oauth.ts"; +import { + type AgentClientMessage, + AgentClientMessageSchema, + AgentServerMessageSchema, + ExecServerMessageSchema, + GetUsableModelsResponseSchema, + InteractionQueryPayloadSchema, + InteractionQuerySchema, + InteractionUpdateSchema, + ModelDetailsSchema, + TextDeltaUpdateSchema, + ThinkingCompletedUpdateSchema, + ThinkingDeltaUpdateSchema, + ThinkingDetailsSchema, + TokenDeltaUpdateSchema, + TurnEndedUpdateSchema, +} from "../../../extensions/ai-providers/cursor/proto.ts"; +import { + create, + fromBinary, + toBinary, +} from "../../../extensions/ai-providers/cursor/protobuf.ts"; +import { + buildCursorRequest, + CURSOR_CHAT_ONLY_SYSTEM_PROMPT, + frameConnectMessage, + streamCursor, +} from "../../../extensions/ai-providers/cursor/provider.ts"; + +const MODEL: Model = { + ...CURSOR_MODELS[0]!, + baseUrl: "", +}; + +const CONTEXT: Context = { + systemPrompt: "Follow the system rule.", + messages: [{ role: "user", content: "hello", timestamp: 0 }], +}; + +function piVersionAtLeast( + version: string, + minimum: readonly number[], +): boolean { + const current = version.split(".").map((part) => Number.parseInt(part, 10)); + for (let index = 0; index < minimum.length; index++) { + const currentPart = current[index] ?? 0; + const minimumPart = minimum[index] ?? 0; + if (currentPart !== minimumPart) return currentPart > minimumPart; + } + return true; +} + +function localModel(baseUrl: string): Model { + return { ...MODEL, baseUrl }; +} + +function frameServerMessage(message: Parameters[1]): Buffer { + return frameConnectMessage( + toBinary(AgentServerMessageSchema, message as never), + ); +} + +function responseUpdate(message: Parameters[1]): Buffer { + return frameServerMessage( + create(AgentServerMessageSchema, { + message: { + case: "interactionUpdate", + value: message as never, + }, + }), + ); +} + +function collectEvents( + stream: AssistantMessageEventStream, +): Promise { + return (async () => { + const events: AssistantMessageEvent[] = []; + for await (const event of stream) events.push(event); + return events; + })(); +} + +async function startServer( + handler: ( + stream: ServerHttp2Stream, + headers: Record, + ) => void, +): Promise<{ baseUrl: string; close(): Promise }> { + const server = createServer(); + server.on("stream", (stream, headers) => { + handler( + stream as ServerHttp2Stream, + headers as Record, + ); + }); + server.listen(0, "127.0.0.1"); + await once(server, "listening"); + const address = server.address(); + assert.ok(address && typeof address === "object"); + return { + baseUrl: `http://127.0.0.1:${address.port}`, + close: async () => { + if (server.listening) server.close(); + await once(server, "close").catch(() => undefined); + }, + }; +} + +function appendChunk(buffer: Buffer, chunk: Buffer | string): Buffer { + return Buffer.concat([ + buffer, + typeof chunk === "string" ? Buffer.from(chunk) : chunk, + ]); +} + +const servers: Array<{ close(): Promise }> = []; +const tempDirectories: string[] = []; +afterEach(async () => { + await Promise.all(servers.splice(0).map((server) => server.close())); + await Promise.all( + tempDirectories + .splice(0) + .map((directory) => rm(directory, { recursive: true, force: true })), + ); +}); + +test("Cursor OAuth uses PKCE and preserves refresh token when refresh response omits it", async () => { + const auth = await generateCursorAuthParams(); + const loginUrl = new URL(auth.loginUrl); + assert.equal(loginUrl.origin, "https://cursor.com"); + assert.equal(loginUrl.pathname, "/loginDeepControl"); + assert.equal(loginUrl.searchParams.get("challenge"), auth.challenge); + assert.equal(loginUrl.searchParams.get("uuid"), auth.uuid); + assert.equal(loginUrl.searchParams.get("mode"), "login"); + assert.equal(loginUrl.searchParams.get("redirectTarget"), "cli"); + + const now = Math.floor(Date.now() / 1_000); + const token = `eyJhbGciOiJub25lIn0.${Buffer.from(JSON.stringify({ exp: now + 3_600 })).toString("base64url")}.sig`; + assert.equal(getCursorTokenExpiry(token), (now + 3_600) * 1_000 - 300_000); + + const originalFetch = globalThis.fetch; + try { + globalThis.fetch = (async (_input, init) => { + assert.equal(init?.method, "POST"); + assert.equal( + new Headers(init?.headers).get("authorization"), + "Bearer old-refresh", + ); + assert.equal(init?.body, "{}"); + return new Response(JSON.stringify({ accessToken: token }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; + const refreshed = await refreshCursorToken( + { access: "old-access", refresh: "old-refresh", expires: 0 }, + new AbortController().signal, + ); + assert.equal(refreshed.access, token); + assert.equal(refreshed.refresh, "old-refresh"); + assert.equal(refreshed.expires, getCursorTokenExpiry(token)); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("Cursor request encodes image content in the selected image protobuf", async () => { + const bytes = Uint8Array.from([0, 1, 2, 250]); + const built = await buildCursorRequest(MODEL, { + messages: [ + { + role: "user", + content: [ + { type: "text", text: "describe this" }, + { + type: "image", + data: Buffer.from(bytes).toString("base64"), + mimeType: "image/png", + }, + ], + timestamp: 0, + }, + ], + }); + const client = fromBinary(AgentClientMessageSchema, built.requestBytes); + assert.ok(client.message.case === "runRequest"); + const action = client.message.value.action?.action; + assert.ok(action?.case === "userMessageAction"); + const selectedImage = + action.value.userMessage?.selectedContext?.selectedImages[0]; + assert.ok(selectedImage?.dataOrBlobId.case === "data"); + assert.deepEqual([...selectedImage.dataOrBlobId.value], [...bytes]); + assert.equal(selectedImage.mimeType, "image/png"); +}); + +test("Cursor pins bare Composer 2.5 to the Standard lane", async () => { + const standard = await buildCursorRequest( + { ...MODEL, id: "composer-2.5" }, + CONTEXT, + ); + assert.equal(standard.request.requestedModel?.modelId, "composer-2.5"); + assert.deepEqual( + standard.request.requestedModel?.parameters.map(({ id, value }) => ({ + id, + value, + })), + [{ id: "fast", value: "false" }], + ); + + const fast = await buildCursorRequest( + { ...MODEL, id: "composer-2.5-fast" }, + CONTEXT, + ); + assert.deepEqual(fast.request.requestedModel?.parameters, []); +}); + +test("Cursor interactive input turns an explicit leading image path into ImageContent", async () => { + const directory = await mkdtemp(join(tmpdir(), "openpi-cursor-image-")); + tempDirectories.push(directory); + const imagePath = join(directory, "pasted.png"); + const png = Buffer.from([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, + ]); + await writeFile(imagePath, png); + + const result = await transformCursorImageInput( + { + type: "input", + text: `${imagePath} 这个图中说了什么`, + source: "interactive", + }, + { model: { provider: "cursor" } } as ExtensionContext, + ); + + assert.equal(result.action, "transform"); + assert.ok(result.action === "transform"); + assert.doesNotMatch(result.text, new RegExp(directory)); + assert.match(result.text, /Attached image: "pasted\.png"/); + assert.match(result.text, /这个图中说了什么/); + assert.equal(result.images?.[0]?.mimeType, "image/png"); + assert.equal(result.images?.[0]?.data, png.toString("base64")); +}); + +test("Cursor image-path conversion is scoped to interactive Cursor input", async () => { + const event = { + type: "input" as const, + text: "/does/not/exist.png describe this", + source: "interactive" as const, + }; + assert.deepEqual( + await transformCursorImageInput(event, { + model: { provider: "google-antigravity" }, + } as ExtensionContext), + { action: "continue" }, + ); + assert.deepEqual( + await transformCursorImageInput({ ...event, source: "rpc" }, { + model: { provider: "cursor" }, + } as ExtensionContext), + { action: "continue" }, + ); +}); + +test("Cursor multi-turn request omits prior thinking outside OMP's Kimi-only replay", async () => { + const built = await buildCursorRequest(MODEL, { + messages: [ + { role: "user", content: "first", timestamp: 0 }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "private-thought" }, + { type: "text", text: "visible-answer" }, + ], + api: MODEL.api, + provider: MODEL.provider, + model: MODEL.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }, + stopReason: "stop", + timestamp: 1, + }, + { + role: "toolResult", + toolCallId: "old-call", + toolName: "read", + content: [{ type: "text", text: "orphan-tool-result" }], + isError: false, + timestamp: 1, + }, + { role: "user", content: "second", timestamp: 2 }, + ], + }); + const stored = [...built.blobStore.values()] + .map((value) => new TextDecoder().decode(value)) + .join("\n"); + assert.match(stored, /This Cursor provider is running in chat-only mode/); + assert.match(stored, /visible-answer/); + assert.doesNotMatch(stored, /private-thought/); + assert.doesNotMatch(stored, /orphan-tool-result/); +}); + +test("Cursor discovery decodes Connect unary models and preserves the HTTP/2 endpoint", async () => { + const seenHeaders: Record[] = []; + const server = await startServer((stream, headers) => { + seenHeaders.push(headers); + stream.respond({ + ":status": 200, + "content-type": "application/proto", + }); + stream.on("data", () => undefined); + const response = create(GetUsableModelsResponseSchema, { + models: [ + create(ModelDetailsSchema, { + modelId: "claude-sonnet-1m", + displayName: "Claude Sonnet 1M", + thinkingDetails: create(ThinkingDetailsSchema, {}), + }), + create(ModelDetailsSchema, { + modelId: "text-only", + displayName: "Text Only", + }), + ], + }); + stream.end(toBinary(GetUsableModelsResponseSchema, response)); + }); + servers.push(server); + const models = await fetchCursorUsableModels({ + apiKey: "discovery-token", + baseUrl: server.baseUrl, + }); + assert.ok(models); + assert.deepEqual( + models.map((model) => model.id), + ["claude-sonnet-1m", "text-only"], + ); + assert.equal(models[0]?.baseUrl, server.baseUrl); + assert.equal(models[0]?.reasoning, true); + assert.deepEqual(models[0]?.input, ["text", "image"]); + assert.equal(models[0]?.contextWindow, 1_000_000); + assert.equal(models[1]?.input[0], "text"); + assert.equal( + seenHeaders[0]?.[":path"], + "/agent.v1.AgentService/GetUsableModels", + ); + assert.equal(seenHeaders[0]?.authorization, "Bearer discovery-token"); + assert.equal(seenHeaders[0]?.["content-type"], "application/proto"); + assert.equal(seenHeaders[0]?.["connect-protocol-version"], undefined); +}); + +test("Cursor discovery aligns Max Mode and uses conservative context fallbacks", async () => { + const server = await startServer((stream) => { + stream.respond({ + ":status": 200, + "content-type": "application/proto", + }); + stream.on("data", () => undefined); + stream.end( + toBinary( + GetUsableModelsResponseSchema, + create(GetUsableModelsResponseSchema, { + models: [ + create(ModelDetailsSchema, { + modelId: "claude-opus-5-fast", + displayName: "Claude Opus 5 Fast", + maxMode: true, + }), + create(ModelDetailsSchema, { + modelId: "cursor-composer-max", + displayName: "Cursor Composer Max", + maxMode: true, + }), + create(ModelDetailsSchema, { + modelId: "composer-2.5", + displayName: "Composer 2.5", + }), + create(ModelDetailsSchema, { + modelId: "gemini-3.1-pro", + displayName: "Gemini 3.1 Pro", + maxMode: true, + }), + create(ModelDetailsSchema, { + modelId: "gpt-5.6-sol", + displayName: "GPT-5.6 Sol", + maxMode: true, + }), + create(ModelDetailsSchema, { + modelId: "claude-fable-5", + displayName: "Claude Fable 5", + }), + create(ModelDetailsSchema, { + modelId: "gpt-5.6-terra", + displayName: "GPT-5.6 Terra", + }), + create(ModelDetailsSchema, { + modelId: "cursor-grok-4.6-high", + displayName: "Grok 4.6 High", + }), + create(ModelDetailsSchema, { + modelId: "moonshotai/kimi-k3", + displayName: "Kimi K3", + }), + create(ModelDetailsSchema, { + modelId: "z-ai/glm-5.2-turbo", + displayName: "GLM 5.2 Turbo", + }), + create(ModelDetailsSchema, { + modelId: "z-ai/glm-5.2-flash", + displayName: "GLM 5.2 Flash", + }), + ], + }), + ), + ); + }); + servers.push(server); + + const models = await fetchCursorUsableModels({ + apiKey: "discovery-token", + baseUrl: server.baseUrl, + }); + assert.ok(models); + const byId = new Map(models.map((model) => [model.id, model])); + const maxClaude = byId.get("claude-opus-5-fast"); + assert.equal(maxClaude?.cursorMaxMode, true); + assert.equal(maxClaude?.contextWindow, 1_000_000); + assert.equal(byId.get("cursor-composer-max")?.contextWindow, 200_000); + assert.deepEqual(byId.get("composer-2.5")?.input, ["text", "image"]); + assert.equal(byId.get("gemini-3.1-pro")?.contextWindow, 1_000_000); + assert.equal(byId.get("gpt-5.6-sol")?.contextWindow, 1_000_000); + assert.equal(byId.get("claude-fable-5")?.contextWindow, 200_000); + assert.equal(byId.get("gpt-5.6-terra")?.contextWindow, 200_000); + assert.equal(byId.get("cursor-grok-4.6-high")?.contextWindow, 200_000); + assert.deepEqual(byId.get("cursor-grok-4.6-high")?.input, ["text", "image"]); + assert.equal(byId.get("moonshotai/kimi-k3")?.contextWindow, 1_000_000); + assert.deepEqual(byId.get("moonshotai/kimi-k3")?.input, ["text", "image"]); + assert.equal(byId.get("z-ai/glm-5.2-turbo")?.contextWindow, 1_000_000); + assert.equal(byId.get("z-ai/glm-5.2-flash")?.contextWindow, 200_000); + assert.ok(maxClaude); + const built = await buildCursorRequest(maxClaude as Model, CONTEXT); + assert.equal(built.request.modelDetails?.maxMode, true); + assert.equal(built.request.requestedModel?.maxMode, true); +}); + +test("Cursor discovery and chat HTTP/2 transports honor configured proxies", async () => { + const proxyVariables = [ + "PI_PROXY_CURSOR", + "PI_PROXY", + "HTTPS_PROXY", + "https_proxy", + "HTTP_PROXY", + "http_proxy", + "ALL_PROXY", + "all_proxy", + "NO_PROXY", + "no_proxy", + ] as const; + const originalEnvironment = new Map( + proxyVariables.map((name) => [name, process.env[name]]), + ); + const restoreEnvironment = () => { + for (const [name, value] of originalEnvironment) { + if (value === undefined) delete process.env[name]; + else process.env[name] = value; + } + }; + + try { + for (const variable of ["HTTPS_PROXY", "PI_PROXY_CURSOR"] as const) { + for (const name of proxyVariables) delete process.env[name]; + const connectTargets: string[] = []; + const proxy = createNetServer((socket) => { + socket.once("data", (chunk) => { + const firstLine = chunk.toString("utf8").split("\r\n")[0] ?? ""; + const match = /^CONNECT\s+(\S+)\s+HTTP\/1\.1$/.exec(firstLine); + if (match?.[1]) connectTargets.push(match[1]); + socket.end("HTTP/1.1 403 Forbidden\r\n\r\n"); + }); + }); + proxy.listen(0, "127.0.0.1"); + await once(proxy, "listening"); + const address = proxy.address(); + assert.ok(address && typeof address === "object"); + process.env[variable] = `http://127.0.0.1:${address.port}`; + + try { + const baseUrl = "https://198.51.100.7:8443"; + assert.equal( + await fetchCursorUsableModels({ + apiKey: "token", + baseUrl, + timeoutMs: 1_000, + }), + null, + ); + const events = await collectEvents( + streamCursor(localModel(baseUrl), CONTEXT, { apiKey: "token" }), + ); + assert.equal(events.at(-1)?.type, "error"); + assert.deepEqual(connectTargets, [ + "198.51.100.7:8443", + "198.51.100.7:8443", + ]); + } finally { + proxy.close(); + await once(proxy, "close").catch(() => undefined); + } + } + } finally { + restoreEnvironment(); + } +}); + +test("Cursor stream sends required headers and maps Connect text/thinking/done frames", async () => { + const seenHeaders: Record[] = []; + const server = await startServer((stream, headers) => { + seenHeaders.push(headers); + stream.respond({ + ":status": 200, + "content-type": "application/connect+proto", + }); + let requestBytes: Buffer = Buffer.alloc(0); + stream.on("data", (chunk) => { + requestBytes = appendChunk(requestBytes, chunk); + if (requestBytes.length < 5) return; + const length = requestBytes.readUInt32BE(1); + if (requestBytes.length < length + 5) return; + const client = fromBinary( + AgentClientMessageSchema, + requestBytes.subarray(5, length + 5), + ); + requestBytes = requestBytes.subarray(length + 5); + if (client.message.case !== "runRequest") return; + + const frames = [ + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "thinkingDelta", + value: create(ThinkingDeltaUpdateSchema, { text: "think" }), + }, + }), + ), + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "thinkingCompleted", + value: create(ThinkingCompletedUpdateSchema, { + thinkingDurationMs: 7, + }), + }, + }), + ), + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "textDelta", + value: create(TextDeltaUpdateSchema, { text: "hello" }), + }, + }), + ), + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "turnEnded", + value: create(TurnEndedUpdateSchema, {}), + }, + }), + ), + ]; + const payload = Buffer.concat(frames); + stream.write(payload.subarray(0, 3)); + setTimeout(() => { + stream.write(payload.subarray(3)); + stream.end(); + }, 2); + }); + }); + servers.push(server); + + const onResponses: Array<{ + status: number; + headers: Record; + }> = []; + const events = await collectEvents( + streamCursor(localModel(server.baseUrl), CONTEXT, { + apiKey: "access-token", + headers: { "x-trace-id": "trace", authorization: "caller-must-not-win" }, + onResponse(response) { + onResponses.push(response); + }, + }), + ); + + assert.equal(seenHeaders.length, 1); + const headers = seenHeaders[0]!; + assert.equal(headers[":path"], "/agent.v1.AgentService/Run"); + assert.equal(headers["content-type"], "application/connect+proto"); + assert.equal(headers["connect-protocol-version"], "1"); + assert.equal(headers.te, "trailers"); + assert.equal(headers.authorization, "Bearer access-token"); + assert.equal(headers["x-ghost-mode"], "true"); + assert.equal(headers["x-cursor-client-type"], "cli"); + assert.equal(headers["x-cursor-client-version"], "cli-2026.07.23-e383d2b"); + assert.equal(headers["x-trace-id"], "trace"); + assert.match(String(headers["x-request-id"]), /^[0-9a-f-]{36}$/); + assert.deepEqual( + onResponses.map((response) => response.status), + [200], + ); + assert.deepEqual( + events.map((event) => event.type), + [ + "start", + "thinking_start", + "thinking_delta", + "thinking_end", + "text_start", + "text_delta", + "text_end", + "done", + ], + ); + const done = events.at(-1); + assert.ok(done?.type === "done"); + assert.deepEqual(done.message.content, [ + { type: "thinking", thinking: "think" }, + { type: "text", text: "hello" }, + ]); + assert.equal( + events.some((event) => event.type.startsWith("toolcall")), + false, + ); +}); + +test("Cursor output-only token deltas preserve Pi's real compaction boundary", async () => { + const server = await startServer((stream) => { + stream.respond({ + ":status": 200, + "content-type": "application/connect+proto", + }); + let requestBytes: Buffer = Buffer.alloc(0); + stream.on("data", (chunk) => { + requestBytes = appendChunk(requestBytes, chunk); + if (requestBytes.length < 5) return; + const length = requestBytes.readUInt32BE(1); + if (requestBytes.length < length + 5) return; + const client = fromBinary( + AgentClientMessageSchema, + requestBytes.subarray(5, length + 5), + ); + requestBytes = requestBytes.subarray(length + 5); + if (client.message.case !== "runRequest") return; + stream.end( + Buffer.concat([ + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "textDelta", + value: create(TextDeltaUpdateSchema, { text: "OK" }), + }, + }), + ), + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "tokenDelta", + value: create(TokenDeltaUpdateSchema, { tokens: 2 }), + }, + }), + ), + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "turnEnded", + value: create(TurnEndedUpdateSchema, {}), + }, + }), + ), + ]), + ); + }); + }); + servers.push(server); + + const user = { + role: "user" as const, + content: "x".repeat(800_000), + timestamp: 0, + }; + const events = await collectEvents( + streamCursor( + localModel(server.baseUrl), + { systemPrompt: "", messages: [user] }, + { apiKey: "token" }, + ), + ); + const done = events.at(-1); + assert.ok(done?.type === "done"); + assert.deepEqual(done.message.usage, { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }); + + const compactionCalls: Array<{ reason: string; willRetry: boolean }> = []; + const session = { + settingsManager: { + getCompactionSettings: () => ({ + enabled: true, + reserveTokens: 16_384, + keepRecentTokens: 20_000, + }), + }, + model: { ...MODEL, contextWindow: 200_000 }, + sessionManager: { getBranch: () => [] }, + agent: { state: { messages: [user, done.message] } }, + _overflowRecoveryAttempted: false, + _emit: () => {}, + _runAutoCompaction: async (reason: string, willRetry: boolean) => { + compactionCalls.push({ reason, willRetry }); + return true; + }, + }; + const checkCompaction = ( + AgentSession.prototype as unknown as { + _checkCompaction: ( + this: unknown, + message: AssistantMessage, + ) => Promise; + } + )._checkCompaction; + const compacted = await checkCompaction.call(session, done.message); + const hostSupportsZeroUsageCompaction = piVersionAtLeast( + PI_VERSION, + [0, 84, 3], + ); + + assert.equal(compacted, hostSupportsZeroUsageCompaction); + assert.deepEqual( + compactionCalls, + hostSupportsZeroUsageCompaction + ? [{ reason: "threshold", willRetry: false }] + : [], + ); +}); + +test("Cursor request_context succeeds with global rules and empty tools; other exec is thrown", { + timeout: 1_000, +}, async () => { + const replies: AgentClientMessage[] = []; + const streamCloseReceived = Promise.withResolvers(); + const server = await startServer((stream) => { + stream.respond({ + ":status": 200, + "content-type": "application/connect+proto", + }); + let requestBytes: Buffer = Buffer.alloc(0); + let sent = false; + stream.on("data", (chunk) => { + requestBytes = appendChunk(requestBytes, chunk); + while (requestBytes.length >= 5) { + const length = requestBytes.readUInt32BE(1); + if (requestBytes.length < length + 5) return; + const client = fromBinary( + AgentClientMessageSchema, + requestBytes.subarray(5, length + 5), + ); + requestBytes = requestBytes.subarray(length + 5); + if (client.message.case === "runRequest" && !sent) { + sent = true; + const contextExec = create(ExecServerMessageSchema, { + id: 11, + execId: "ctx", + message: { + case: "requestContextArgs", + value: { $typeName: "agent.v1.RequestContextArgs" }, + }, + }); + const unsupportedExec = create(ExecServerMessageSchema, { + id: 12, + execId: "tool", + message: { case: undefined }, + }); + for (const exec of [contextExec, unsupportedExec]) { + stream.write( + frameServerMessage( + create(AgentServerMessageSchema, { + message: { case: "execServerMessage", value: exec }, + }), + ), + ); + } + } else if ( + client.message.case === "execClientMessage" || + client.message.case === "execClientControlMessage" + ) { + replies.push(client); + if ( + client.message.case === "execClientControlMessage" && + client.message.value.message.case === "streamClose" + ) { + streamCloseReceived.resolve(); + stream.write( + responseUpdate( + create(InteractionUpdateSchema, { + message: { + case: "turnEnded", + value: create(TurnEndedUpdateSchema, {}), + }, + }), + ), + ); + setTimeout(() => stream.end(), 2); + } + } + } + }); + }); + servers.push(server); + + const events = await collectEvents( + streamCursor(localModel(server.baseUrl), CONTEXT, { apiKey: "token" }), + ); + await streamCloseReceived.promise; + const contextReply = replies.find( + (reply) => + reply.message.case === "execClientMessage" && + reply.message.value.message.case === "requestContextResult", + ); + assert.ok(contextReply?.message.case === "execClientMessage"); + const contextMessage = contextReply.message.value.message; + assert.ok(contextMessage.case === "requestContextResult"); + assert.equal(contextMessage.value.result.case, "success"); + const requestContext = contextMessage.value.result.value.requestContext; + assert.deepEqual(requestContext?.tools, []); + assert.equal(requestContext?.rules[0]?.content, CONTEXT.systemPrompt); + assert.equal(requestContext?.rules[0]?.type?.type.case, "global"); + assert.equal( + requestContext?.rules[1]?.content, + CURSOR_CHAT_ONLY_SYSTEM_PROMPT, + ); + assert.equal(requestContext?.rules[1]?.fullPath, "/pi/cursor-chat-only.mdc"); + const throwReply = replies.find( + (reply) => + reply.message.case === "execClientControlMessage" && + reply.message.value.message.case === "throw", + ); + assert.ok(throwReply?.message.case === "execClientControlMessage"); + const throwMessage = throwReply.message.value.message; + assert.ok(throwMessage.case === "throw"); + assert.equal(throwMessage.value.errorCode, "UNIMPLEMENTED"); + const terminal = events.at(-1); + assert.ok(terminal?.type === "error"); + assert.match( + terminal.error.errorMessage ?? "", + /unavailable in chat-only mode/, + ); + assert.equal( + events.some((event) => event.type === "done"), + false, + ); + assert.equal( + events.some((event) => event.type.startsWith("toolcall")), + false, + ); +}); + +test("Cursor tool interaction and interactionQuery fail explicitly without a Pi toolCall", async () => { + const cases: Array<"tool" | "query"> = ["tool", "query"]; + for (const kind of cases) { + const server = await startServer((stream) => { + stream.respond({ + ":status": 200, + "content-type": "application/connect+proto", + }); + stream.on("data", (chunk) => { + const bytes = typeof chunk === "string" ? Buffer.from(chunk) : chunk; + if (bytes.length < 5) return; + const length = bytes.readUInt32BE(1); + if (bytes.length < length + 5) return; + const client = fromBinary( + AgentClientMessageSchema, + bytes.subarray(5, length + 5), + ); + if (client.message.case !== "runRequest") return; + const message = + kind === "tool" + ? create(AgentServerMessageSchema, { + message: { + case: "interactionUpdate", + value: create(InteractionUpdateSchema, { + message: { case: "toolCallStarted", value: {} }, + }), + }, + }) + : create(AgentServerMessageSchema, { + message: { + case: "interactionQuery", + value: create(InteractionQuerySchema, { + id: 1, + query: { + case: "askQuestionInteractionQuery", + value: create(InteractionQueryPayloadSchema, {}), + }, + }), + }, + }); + stream.end(frameServerMessage(message)); + }); + }); + servers.push(server); + const events = await collectEvents( + streamCursor(localModel(server.baseUrl), CONTEXT, { apiKey: "token" }), + ); + const error = events.find((event) => event.type === "error"); + assert.ok(error?.type === "error"); + assert.match( + error.error.errorMessage ?? "", + /unavailable in chat-only mode/, + ); + assert.equal( + events.some((event) => event.type.startsWith("toolcall")), + false, + ); + } +}); + +test("Cursor abort produces an aborted error and native fetch skips network offline", async () => { + const controller = new AbortController(); + let requestStarted!: () => void; + const started = new Promise((resolve) => { + requestStarted = resolve; + }); + const server = await startServer((stream) => { + stream.respond({ + ":status": 200, + "content-type": "application/connect+proto", + }); + requestStarted(); + stream.on("data", () => undefined); + }); + servers.push(server); + const eventsPromise = collectEvents( + streamCursor(localModel(server.baseUrl), CONTEXT, { + apiKey: "token", + signal: controller.signal, + }), + ); + await started; + controller.abort(); + const events = await eventsPromise; + const error = events.at(-1); + assert.ok(error?.type === "error"); + assert.equal(error.reason, "aborted"); + + const models = await fetchCursorModels({ + allowNetwork: false, + publish: async () => true, + signal: new AbortController().signal, + }); + assert.deepEqual(models, []); + assert.equal(CURSOR_MODELS[0]?.id, "default"); +}); + +test("Cursor rejects custom fetch and bounds an idle HTTP/2 stream", async () => { + const customFetchEvents = await collectEvents( + streamCursor(MODEL, CONTEXT, { + apiKey: "token", + fetch: async () => new Response(), + }), + ); + const customFetchError = customFetchEvents.at(-1); + assert.ok(customFetchError?.type === "error"); + assert.match( + customFetchError.error.errorMessage ?? "", + /does not support options\.fetch/, + ); + + const server = await startServer((stream) => { + stream.respond({ + ":status": 200, + "content-type": "application/connect+proto", + }); + stream.on("data", () => undefined); + }); + servers.push(server); + const timeoutEvents = await collectEvents( + streamCursor(localModel(server.baseUrl), CONTEXT, { + apiKey: "token", + timeoutMs: 10, + }), + ); + const timeoutError = timeoutEvents.at(-1); + assert.ok(timeoutError?.type === "error"); + assert.match( + timeoutError.error.errorMessage ?? "", + /idle timeout after 10ms/, + ); +}); diff --git a/tests/extensions/ai-providers/registration.test.ts b/tests/extensions/ai-providers/registration.test.ts new file mode 100644 index 00000000..48f26bcc --- /dev/null +++ b/tests/extensions/ai-providers/registration.test.ts @@ -0,0 +1,363 @@ +import assert from "node:assert/strict"; +import { test } from "node:test"; +import type { + Api, + AssistantMessageEventStream, + Model, + ModelsStoreEntry, + OAuthCredentials, + Provider, +} from "@earendil-works/pi-ai"; +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { decodeApiKey } from "../../../extensions/ai-providers/antigravity/credentials.ts"; +import authProviders from "../../../extensions/ai-providers/index.ts"; +import { createOAuthAuth } from "../../../extensions/ai-providers/oauth-adapter.ts"; + +function loadProviders(): { + providers: Provider[]; + inputHandlerRegistered: boolean; +} { + const providers: Provider[] = []; + let inputHandlerRegistered = false; + const pi = { + on(event: string) { + if (event === "input") inputHandlerRegistered = true; + }, + registerProvider(...args: unknown[]) { + assert.equal(args.length, 1); + assert.equal(typeof args[0], "object"); + providers.push(args[0] as Provider); + }, + } as unknown as ExtensionAPI; + + authProviders(pi); + return { providers, inputHandlerRegistered }; +} + +async function collectEvents(stream: AssistantMessageEventStream) { + const events = []; + for await (const event of stream) events.push(event); + return events; +} + +test("ai-providers registers complete native providers", () => { + const { providers, inputHandlerRegistered } = loadProviders(); + assert.equal(inputHandlerRegistered, true); + assert.deepEqual( + providers.map((provider) => provider.id), + ["google-antigravity", "cursor"], + ); + + for (const provider of providers) { + assert.ok(provider.auth.oauth); + assert.ok(provider.refreshModels); + assert.equal(typeof provider.stream, "function"); + assert.equal(typeof provider.streamSimple, "function"); + const models = provider.getModels(); + assert.ok(models.length > 0); + for (const model of models) { + assert.equal(model.provider, provider.id); + assert.ok(model.api.length > 0); + assert.ok(model.baseUrl.length > 0); + } + } +}); + +test("Antigravity OAuth cancellation during version discovery stays bounded", async () => { + const originalFetch = globalThis.fetch; + const originalVersion = process.env.OPENPI_ANTIGRAVITY_VERSION; + delete process.env.OPENPI_ANTIGRAVITY_VERSION; + const manifestStarted = Promise.withResolvers(); + const notifications: unknown[] = []; + let prompts = 0; + try { + globalThis.fetch = ((_input, init) => { + manifestStarted.resolve(); + return new Promise((_resolve, reject) => { + const signal = init?.signal; + const onAbort = () => reject(signal?.reason ?? new Error("aborted")); + if (signal?.aborted) onAbort(); + else signal?.addEventListener("abort", onAbort, { once: true }); + }); + }) as typeof fetch; + const { providers } = loadProviders(); + const provider = providers.find( + (entry) => entry.id === "google-antigravity", + ); + const oauth = provider?.auth.oauth; + assert.ok(oauth); + const controller = new AbortController(); + const login = oauth.login({ + signal: controller.signal, + notify: (event) => notifications.push(event), + prompt: async () => { + prompts++; + return ""; + }, + }); + await manifestStarted.promise; + controller.abort("test cancellation"); + let timer: ReturnType | undefined; + const bounded = Promise.race([ + login, + new Promise((_, reject) => { + timer = setTimeout( + () => reject(new Error("login cancellation remained pending")), + 250, + ); + }), + ]); + try { + await assert.rejects(bounded, /cancelled/); + } finally { + if (timer) clearTimeout(timer); + } + assert.deepEqual(notifications, []); + assert.equal(prompts, 0); + } finally { + globalThis.fetch = originalFetch; + if (originalVersion === undefined) { + delete process.env.OPENPI_ANTIGRAVITY_VERSION; + } else { + process.env.OPENPI_ANTIGRAVITY_VERSION = originalVersion; + } + } +}); + +test("native providers restore stored dynamic models without replacing baselines", async () => { + const { providers } = loadProviders(); + for (const provider of providers) { + const baseline = provider.getModels()[0]; + assert.ok(baseline); + const overridden: Model = { + ...baseline, + name: `${baseline.name} (cached)`, + }; + const accountModel: Model = { + ...baseline, + id: "account-only-model", + name: "Account Only Model", + }; + let publications = 0; + + await provider.refreshModels?.({ + stored: { models: [overridden, accountModel], checkedAt: 123 }, + allowNetwork: false, + signal: new AbortController().signal, + publish: async (publication) => { + publications += 1; + assert.equal(publication.persist, undefined); + publication.update?.(); + return true; + }, + }); + + const restored = provider.getModels(); + assert.equal(publications, 1); + assert.equal( + restored.filter((model) => model.id === baseline.id).length, + 1, + ); + assert.equal( + restored.find((model) => model.id === baseline.id)?.name, + overridden.name, + ); + assert.ok(restored.some((model) => model.id === accountModel.id)); + } +}); + +test("Antigravity native refresh persists success and retains it on later failure", async () => { + const originalFetch = globalThis.fetch; + const originalVersion = process.env.OPENPI_ANTIGRAVITY_VERSION; + process.env.OPENPI_ANTIGRAVITY_VERSION = "2.8.0"; + const { providers } = loadProviders(); + const provider = providers.find((entry) => entry.id === "google-antigravity"); + assert.ok(provider?.refreshModels); + let stored: ModelsStoreEntry | undefined; + + try { + globalThis.fetch = (async () => + new Response( + JSON.stringify({ + models: { + "account-model": { + displayName: "Account Model", + supportsImages: true, + }, + }, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + )) as typeof fetch; + await provider.refreshModels({ + credential: { + type: "oauth", + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }, + allowNetwork: true, + signal: new AbortController().signal, + publish: async (publication) => { + if (publication.persist) stored = publication.persist; + publication.update?.(); + return true; + }, + }); + assert.ok(stored); + assert.equal(stored.models[0]?.id, "account-model"); + assert.ok(Number.isFinite(stored.checkedAt)); + assert.ok( + provider.getModels().some((model) => model.id === "account-model"), + ); + + globalThis.fetch = (async () => + new Response("unavailable", { status: 503 })) as typeof fetch; + let unexpectedPersist = false; + await assert.rejects( + provider.refreshModels({ + stored, + credential: { + type: "oauth", + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }, + allowNetwork: true, + signal: new AbortController().signal, + publish: async (publication) => { + if (publication.persist !== undefined) unexpectedPersist = true; + publication.update?.(); + return true; + }, + }), + /failed on all endpoints/, + ); + assert.equal(unexpectedPersist, false); + assert.ok( + provider.getModels().some((model) => model.id === "account-model"), + ); + } finally { + globalThis.fetch = originalFetch; + if (originalVersion === undefined) { + delete process.env.OPENPI_ANTIGRAVITY_VERSION; + } else { + process.env.OPENPI_ANTIGRAVITY_VERSION = originalVersion; + } + } +}); + +test("native stream and streamSimple both dispatch to provider implementations", async () => { + const { providers } = loadProviders(); + for (const provider of providers) { + const model = provider.getModels()[0]; + assert.ok(model); + for (const stream of [ + provider.stream(model, { messages: [] }), + provider.streamSimple(model, { messages: [] }), + ]) { + const events = await collectEvents(stream); + const final = events.at(-1); + assert.equal(final?.type, "error"); + if (final?.type === "error") { + assert.match( + final.error.errorMessage ?? "", + /OAuth credentials|access token|API key/, + ); + } + } + } +}); + +test("OAuth adapter preserves events, prompts, credentials, and request auth", async () => { + const signal = new AbortController().signal; + const notifications: unknown[] = []; + const promptTypes: string[] = []; + let refreshSignal: AbortSignal | undefined; + let requestedManualSignal: AbortSignal | undefined; + let receivedManualSignal: AbortSignal | undefined; + const oauth = createOAuthAuth({ + name: "Test OAuth", + isSubscription: true, + async login(callbacks) { + callbacks.onAuth({ + url: "https://example.test/login", + instructions: "Sign in", + }); + callbacks.onProgress?.("Waiting"); + requestedManualSignal = new AbortController().signal; + const code = await callbacks.onManualCodeInput?.(requestedManualSignal); + const choice = await callbacks.onSelect({ + message: "Account", + options: [{ id: "one", label: "One" }], + }); + return { + refresh: "refresh-token", + access: "access-token", + expires: 123, + metadata: `${code}:${choice}`, + }; + }, + async refreshToken(credential, receivedSignal) { + refreshSignal = receivedSignal; + return { ...credential, access: "refreshed-access" }; + }, + getApiKey: (credential) => credential.access, + }); + + const credential = await oauth.login({ + signal, + notify: (event) => notifications.push(event), + prompt: async (prompt) => { + promptTypes.push(prompt.type); + if (prompt.type === "manual_code") { + receivedManualSignal = prompt.signal; + } + return prompt.type === "select" ? "one" : "manual-code"; + }, + }); + assert.equal(credential.type, "oauth"); + assert.equal(credential.metadata, "manual-code:one"); + assert.deepEqual( + notifications.map((event) => (event as { type: string }).type), + ["auth_url", "progress"], + ); + assert.deepEqual(promptTypes, ["manual_code", "select"]); + assert.equal(receivedManualSignal, requestedManualSignal); + + const refreshed = await oauth.refresh(credential, signal); + assert.equal(refreshSignal, signal); + assert.equal(refreshed.type, "oauth"); + assert.equal(refreshed.access, "refreshed-access"); + assert.deepEqual(await oauth.toAuth(refreshed), { + apiKey: "refreshed-access", + }); +}); + +test("registered OAuth adapters derive the existing provider credentials", async () => { + const { providers } = loadProviders(); + const antigravity = providers.find( + (provider) => provider.id === "google-antigravity", + ); + const cursor = providers.find((provider) => provider.id === "cursor"); + assert.ok(antigravity?.auth.oauth); + assert.ok(cursor?.auth.oauth); + + const base: OAuthCredentials = { + refresh: "refresh", + access: "access", + expires: Date.now() + 60_000, + }; + const antigravityAuth = await antigravity.auth.oauth.toAuth({ + ...base, + type: "oauth", + projectId: "project", + }); + assert.ok(antigravityAuth.apiKey); + assert.deepEqual(decodeApiKey(antigravityAuth.apiKey), { + token: "access", + projectId: "project", + }); + assert.deepEqual(await cursor.auth.oauth.toAuth({ ...base, type: "oauth" }), { + apiKey: "access", + }); +});