Merge branch 'main' into feat/image-outputs
This commit is contained in:
Cristina Poncela Cubeiro
@@ -113,12 +113,19 @@ function getApiKeyEnvVars(provider: string): readonly string[] | undefined {
|
||||
mistral: "MISTRAL_API_KEY",
|
||||
minimax: "MINIMAX_API_KEY",
|
||||
"minimax-cn": "MINIMAX_CN_API_KEY",
|
||||
moonshotai: "MOONSHOT_API_KEY",
|
||||
"moonshotai-cn": "MOONSHOT_API_KEY",
|
||||
huggingface: "HF_TOKEN",
|
||||
fireworks: "FIREWORKS_API_KEY",
|
||||
opencode: "OPENCODE_API_KEY",
|
||||
"opencode-go": "OPENCODE_API_KEY",
|
||||
"kimi-coding": "KIMI_API_KEY",
|
||||
"cloudflare-workers-ai": "CLOUDFLARE_API_KEY",
|
||||
"cloudflare-ai-gateway": "CLOUDFLARE_API_KEY",
|
||||
xiaomi: "XIAOMI_API_KEY",
|
||||
"xiaomi-token-plan-cn": "XIAOMI_TOKEN_PLAN_CN_API_KEY",
|
||||
"xiaomi-token-plan-ams": "XIAOMI_TOKEN_PLAN_AMS_API_KEY",
|
||||
"xiaomi-token-plan-sgp": "XIAOMI_TOKEN_PLAN_SGP_API_KEY",
|
||||
};
|
||||
|
||||
const envVar = envMap[provider];
|
||||
|
||||
@@ -12,16 +12,21 @@ export type { AnthropicEffort, AnthropicOptions, AnthropicThinkingDisplay } from
|
||||
export type { AzureOpenAIResponsesOptions } from "./providers/azure-openai-responses.js";
|
||||
export * from "./providers/faux.js";
|
||||
export type { GoogleOptions } from "./providers/google.js";
|
||||
export type { GoogleGeminiCliOptions, GoogleThinkingLevel } from "./providers/google-gemini-cli.js";
|
||||
export type { GoogleThinkingLevel } from "./providers/google-shared.js";
|
||||
export type { GoogleVertexOptions } from "./providers/google-vertex.js";
|
||||
export type { MistralOptions } from "./providers/mistral.js";
|
||||
export type { OpenAICodexResponsesOptions } from "./providers/openai-codex-responses.js";
|
||||
export type {
|
||||
OpenAICodexResponsesOptions,
|
||||
OpenAICodexWebSocketDebugStats,
|
||||
} from "./providers/openai-codex-responses.js";
|
||||
export type { OpenAICompletionsOptions } from "./providers/openai-completions.js";
|
||||
export type { OpenAIResponsesOptions } from "./providers/openai-responses.js";
|
||||
export * from "./providers/register-builtins.js";
|
||||
export * from "./providers/register-images-builtins.js";
|
||||
export * from "./session-resources.js";
|
||||
export * from "./stream.js";
|
||||
export * from "./types.js";
|
||||
export * from "./utils/diagnostics.js";
|
||||
export * from "./utils/event-stream.js";
|
||||
export * from "./utils/json-parse.js";
|
||||
export type {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,5 +1,5 @@
|
||||
import { MODELS } from "./models.generated.js";
|
||||
import type { Api, KnownProvider, Model, Usage } from "./types.js";
|
||||
import type { Api, KnownProvider, Model, ModelThinkingLevel, Usage } from "./types.js";
|
||||
|
||||
const modelRegistry: Map<string, Map<string, Model<Api>>> = new Map();
|
||||
|
||||
@@ -45,35 +45,38 @@ export function calculateCost<TApi extends Api>(model: Model<TApi>, usage: Usage
|
||||
return usage.cost;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a model supports xhigh thinking level.
|
||||
*
|
||||
* Supported today:
|
||||
* - GPT-5.2 / GPT-5.3 / GPT-5.4 / GPT-5.5 model families
|
||||
* - DeepSeek V4 Pro
|
||||
* - Opus 4.6+ models (xhigh maps to adaptive effort "max" on Anthropic-compatible providers)
|
||||
*/
|
||||
export function supportsXhigh<TApi extends Api>(model: Model<TApi>): boolean {
|
||||
if (
|
||||
model.id.includes("gpt-5.2") ||
|
||||
model.id.includes("gpt-5.3") ||
|
||||
model.id.includes("gpt-5.4") ||
|
||||
model.id.includes("gpt-5.5") ||
|
||||
model.id.includes("deepseek-v4-pro")
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
const EXTENDED_THINKING_LEVELS: ModelThinkingLevel[] = ["off", "minimal", "low", "medium", "high", "xhigh"];
|
||||
|
||||
if (
|
||||
model.id.includes("opus-4-6") ||
|
||||
model.id.includes("opus-4.6") ||
|
||||
model.id.includes("opus-4-7") ||
|
||||
model.id.includes("opus-4.7")
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
export function getSupportedThinkingLevels<TApi extends Api>(model: Model<TApi>): ModelThinkingLevel[] {
|
||||
if (!model.reasoning) return ["off"];
|
||||
|
||||
return false;
|
||||
return EXTENDED_THINKING_LEVELS.filter((level) => {
|
||||
const mapped = model.thinkingLevelMap?.[level];
|
||||
if (mapped === null) return false;
|
||||
if (level === "xhigh") return mapped !== undefined;
|
||||
return true;
|
||||
});
|
||||
}
|
||||
|
||||
export function clampThinkingLevel<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
level: ModelThinkingLevel,
|
||||
): ModelThinkingLevel {
|
||||
const availableLevels = getSupportedThinkingLevels(model);
|
||||
if (availableLevels.includes(level)) return level;
|
||||
|
||||
const requestedIndex = EXTENDED_THINKING_LEVELS.indexOf(level);
|
||||
if (requestedIndex === -1) return availableLevels[0] ?? "off";
|
||||
|
||||
for (let i = requestedIndex; i < EXTENDED_THINKING_LEVELS.length; i++) {
|
||||
const candidate = EXTENDED_THINKING_LEVELS[i];
|
||||
if (availableLevels.includes(candidate)) return candidate;
|
||||
}
|
||||
for (let i = requestedIndex - 1; i >= 0; i--) {
|
||||
const candidate = EXTENDED_THINKING_LEVELS[i];
|
||||
if (availableLevels.includes(candidate)) return candidate;
|
||||
}
|
||||
return availableLevels[0] ?? "off";
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -481,25 +481,33 @@ function handleContentBlockStop(
|
||||
* Checks both model ID and model name to support application inference profiles
|
||||
* whose ARNs don't contain the model name.
|
||||
*/
|
||||
function getModelMatchCandidates(modelId: string, modelName?: string): string[] {
|
||||
const values = modelName ? [modelId, modelName] : [modelId];
|
||||
return values.flatMap((value) => {
|
||||
const lower = value.toLowerCase();
|
||||
return [lower, lower.replace(/[\s_.:]+/g, "-")];
|
||||
});
|
||||
}
|
||||
|
||||
function supportsAdaptiveThinking(modelId: string, modelName?: string): boolean {
|
||||
const candidates = modelName ? [modelId, modelName] : [modelId];
|
||||
return candidates.some(
|
||||
(s) =>
|
||||
s.includes("opus-4-6") ||
|
||||
s.includes("opus-4.6") ||
|
||||
s.includes("opus-4-7") ||
|
||||
s.includes("opus-4.7") ||
|
||||
s.includes("sonnet-4-6") ||
|
||||
s.includes("sonnet-4.6"),
|
||||
);
|
||||
const candidates = getModelMatchCandidates(modelId, modelName);
|
||||
return candidates.some((s) => s.includes("opus-4-6") || s.includes("opus-4-7") || s.includes("sonnet-4-6"));
|
||||
}
|
||||
|
||||
function supportsNativeXhighEffort(model: Model<"bedrock-converse-stream">): boolean {
|
||||
const candidates = getModelMatchCandidates(model.id, model.name);
|
||||
return candidates.some((s) => s.includes("opus-4-7"));
|
||||
}
|
||||
|
||||
function mapThinkingLevelToEffort(
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
level: SimpleStreamOptions["reasoning"],
|
||||
modelId: string,
|
||||
modelName?: string,
|
||||
): "low" | "medium" | "high" | "xhigh" | "max" {
|
||||
const candidates = modelName ? [modelId, modelName] : [modelId];
|
||||
if (level === "xhigh" && supportsNativeXhighEffort(model)) return "xhigh";
|
||||
|
||||
const mapped = level ? model.thinkingLevelMap?.[level] : undefined;
|
||||
if (typeof mapped === "string") return mapped as "low" | "medium" | "high" | "xhigh" | "max";
|
||||
|
||||
switch (level) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
@@ -508,14 +516,6 @@ function mapThinkingLevelToEffort(
|
||||
return "medium";
|
||||
case "high":
|
||||
return "high";
|
||||
case "xhigh":
|
||||
if (candidates.some((s) => s.includes("opus-4-6") || s.includes("opus-4.6"))) {
|
||||
return "max";
|
||||
}
|
||||
if (candidates.some((s) => s.includes("opus-4-7") || s.includes("opus-4.7"))) {
|
||||
return "xhigh";
|
||||
}
|
||||
return "high";
|
||||
default:
|
||||
return "high";
|
||||
}
|
||||
@@ -565,10 +565,7 @@ function isAnthropicClaudeModel(model: Model<"bedrock-converse-stream">): boolea
|
||||
* Amazon Nova models have automatic caching and don't need explicit cache points.
|
||||
*/
|
||||
function supportsPromptCaching(model: Model<"bedrock-converse-stream">): boolean {
|
||||
const candidates = [model.id.toLowerCase()];
|
||||
if (model.name) {
|
||||
candidates.push(model.name.toLowerCase());
|
||||
}
|
||||
const candidates = getModelMatchCandidates(model.id, model.name);
|
||||
|
||||
const hasClaudeRef = candidates.some((s) => s.includes("claude"));
|
||||
if (!hasClaudeRef) {
|
||||
@@ -578,7 +575,7 @@ function supportsPromptCaching(model: Model<"bedrock-converse-stream">): boolean
|
||||
return false;
|
||||
}
|
||||
// Claude 4.x models (opus-4, sonnet-4, haiku-4)
|
||||
if (candidates.some((s) => s.includes("-4-") || s.includes("-4."))) return true;
|
||||
if (candidates.some((s) => s.includes("-4-"))) return true;
|
||||
// Claude 3.7 Sonnet
|
||||
if (candidates.some((s) => s.includes("claude-3-7-sonnet"))) return true;
|
||||
// Claude 3.5 Haiku
|
||||
@@ -895,7 +892,7 @@ function buildAdditionalModelRequestFields(
|
||||
const result: Record<string, any> = supportsAdaptiveThinking(model.id, model.name)
|
||||
? {
|
||||
thinking: { type: "adaptive", ...(display !== undefined ? { display } : {}) },
|
||||
output_config: { effort: mapThinkingLevelToEffort(options.reasoning, model.id, model.name) },
|
||||
output_config: { effort: mapThinkingLevelToEffort(model, options.reasoning) },
|
||||
}
|
||||
: (() => {
|
||||
const defaultBudgets: Record<ThinkingLevel, number> = {
|
||||
|
||||
@@ -32,6 +32,7 @@ import { headersToRecord } from "../utils/headers.js";
|
||||
import { parseJsonWithRepair, parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
|
||||
import { resolveCloudflareBaseUrl } from "./cloudflare.js";
|
||||
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
|
||||
import { adjustMaxTokensForThinking, buildBaseOptions } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
@@ -215,8 +216,8 @@ export interface AnthropicOptions extends StreamOptions {
|
||||
client?: Anthropic;
|
||||
}
|
||||
|
||||
function mergeHeaders(...headerSources: (Record<string, string> | undefined)[]): Record<string, string> {
|
||||
const merged: Record<string, string> = {};
|
||||
function mergeHeaders(...headerSources: (Record<string, string | null> | undefined)[]): Record<string, string | null> {
|
||||
const merged: Record<string, string | null> = {};
|
||||
for (const headers of headerSources) {
|
||||
if (headers) {
|
||||
Object.assign(merged, headers);
|
||||
@@ -384,6 +385,9 @@ async function* iterateAnthropicEvents(
|
||||
throw new Error("Attempted to iterate over an Anthropic response with no body");
|
||||
}
|
||||
|
||||
let sawMessageStart = false;
|
||||
let sawMessageEnd = false;
|
||||
|
||||
for await (const sse of iterateSseMessages(response.body, signal)) {
|
||||
if (sse.event === "error") {
|
||||
throw new Error(sse.data);
|
||||
@@ -394,7 +398,13 @@ async function* iterateAnthropicEvents(
|
||||
}
|
||||
|
||||
try {
|
||||
yield parseJsonWithRepair<RawMessageStreamEvent>(sse.data);
|
||||
const event = parseJsonWithRepair<RawMessageStreamEvent>(sse.data);
|
||||
if (event.type === "message_start") {
|
||||
sawMessageStart = true;
|
||||
} else if (event.type === "message_stop") {
|
||||
sawMessageEnd = true;
|
||||
}
|
||||
yield event;
|
||||
} catch (error) {
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
throw new Error(
|
||||
@@ -402,6 +412,10 @@ async function* iterateAnthropicEvents(
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (sawMessageStart && !sawMessageEnd) {
|
||||
throw new Error("Anthropic stream ended before message_stop");
|
||||
}
|
||||
}
|
||||
|
||||
export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOptions> = (
|
||||
@@ -680,24 +694,21 @@ function supportsAdaptiveThinking(modelId: string): boolean {
|
||||
* Map ThinkingLevel to Anthropic effort levels for adaptive thinking.
|
||||
* Note: effort "max" is only valid on Opus 4.6, while Opus 4.7 supports "xhigh".
|
||||
*/
|
||||
function mapThinkingLevelToEffort(level: SimpleStreamOptions["reasoning"], modelId: string): AnthropicEffort {
|
||||
function mapThinkingLevelToEffort(
|
||||
model: Model<"anthropic-messages">,
|
||||
level: SimpleStreamOptions["reasoning"],
|
||||
): AnthropicEffort {
|
||||
const mapped = level ? model.thinkingLevelMap?.[level] : undefined;
|
||||
if (typeof mapped === "string") return mapped as AnthropicEffort;
|
||||
|
||||
switch (level) {
|
||||
case "minimal":
|
||||
return "low";
|
||||
case "low":
|
||||
return "low";
|
||||
case "medium":
|
||||
return "medium";
|
||||
case "high":
|
||||
return "high";
|
||||
case "xhigh":
|
||||
if (modelId.includes("opus-4-6") || modelId.includes("opus-4.6")) {
|
||||
return "max";
|
||||
}
|
||||
if (modelId.includes("opus-4-7") || modelId.includes("opus-4.7")) {
|
||||
return "xhigh";
|
||||
}
|
||||
return "high";
|
||||
default:
|
||||
return "high";
|
||||
}
|
||||
@@ -721,7 +732,7 @@ export const streamSimpleAnthropic: StreamFunction<"anthropic-messages", SimpleS
|
||||
// For Opus 4.6 and Sonnet 4.6: use adaptive thinking with effort level
|
||||
// For older models: use budget-based thinking
|
||||
if (supportsAdaptiveThinking(model.id)) {
|
||||
const effort = mapThinkingLevelToEffort(options.reasoning, model.id);
|
||||
const effort = mapThinkingLevelToEffort(model, options.reasoning);
|
||||
return streamAnthropic(model, context, {
|
||||
...base,
|
||||
thinkingEnabled: true,
|
||||
@@ -759,17 +770,39 @@ function createClient(
|
||||
// Adaptive thinking models (Opus 4.6, Sonnet 4.6) have interleaved thinking built-in.
|
||||
// The beta header is deprecated on Opus 4.6 and redundant on Sonnet 4.6, so skip it.
|
||||
const needsInterleavedBeta = interleavedThinking && !supportsAdaptiveThinking(model.id);
|
||||
const betaFeatures: string[] = [];
|
||||
if (useFineGrainedToolStreamingBeta) {
|
||||
betaFeatures.push(FINE_GRAINED_TOOL_STREAMING_BETA);
|
||||
}
|
||||
if (needsInterleavedBeta) {
|
||||
betaFeatures.push(INTERLEAVED_THINKING_BETA);
|
||||
}
|
||||
|
||||
if (model.provider === "cloudflare-ai-gateway") {
|
||||
const client = new Anthropic({
|
||||
apiKey: null,
|
||||
authToken: null,
|
||||
baseURL: resolveCloudflareBaseUrl(model),
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: mergeHeaders(
|
||||
{
|
||||
accept: "application/json",
|
||||
"anthropic-dangerous-direct-browser-access": "true",
|
||||
"cf-aig-authorization": `Bearer ${apiKey}`,
|
||||
"x-api-key": null,
|
||||
Authorization: null,
|
||||
...(betaFeatures.length > 0 ? { "anthropic-beta": betaFeatures.join(",") } : {}),
|
||||
},
|
||||
model.headers,
|
||||
optionsHeaders,
|
||||
),
|
||||
});
|
||||
|
||||
return { client, isOAuthToken: false };
|
||||
}
|
||||
|
||||
// Copilot: Bearer auth, selective betas.
|
||||
if (model.provider === "github-copilot") {
|
||||
const betaFeatures: string[] = [];
|
||||
if (useFineGrainedToolStreamingBeta) {
|
||||
betaFeatures.push(FINE_GRAINED_TOOL_STREAMING_BETA);
|
||||
}
|
||||
if (needsInterleavedBeta) {
|
||||
betaFeatures.push(INTERLEAVED_THINKING_BETA);
|
||||
}
|
||||
|
||||
const client = new Anthropic({
|
||||
apiKey: null,
|
||||
authToken: apiKey,
|
||||
@@ -790,14 +823,6 @@ function createClient(
|
||||
return { client, isOAuthToken: false };
|
||||
}
|
||||
|
||||
const betaFeatures: string[] = [];
|
||||
if (useFineGrainedToolStreamingBeta) {
|
||||
betaFeatures.push(FINE_GRAINED_TOOL_STREAMING_BETA);
|
||||
}
|
||||
if (needsInterleavedBeta) {
|
||||
betaFeatures.push(INTERLEAVED_THINKING_BETA);
|
||||
}
|
||||
|
||||
// OAuth: Bearer auth, Claude Code identity headers
|
||||
if (isOAuthToken(apiKey)) {
|
||||
const client = new Anthropic({
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { AzureOpenAI } from "openai";
|
||||
import type { ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { supportsXhigh } from "../models.js";
|
||||
import { clampThinkingLevel } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
@@ -14,7 +14,7 @@ import type {
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
|
||||
const DEFAULT_AZURE_API_VERSION = "v1";
|
||||
const AZURE_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode", "azure-openai-responses"]);
|
||||
@@ -139,7 +139,8 @@ export const streamSimpleAzureOpenAIResponses: StreamFunction<"azure-openai-resp
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
|
||||
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
|
||||
|
||||
return streamAzureOpenAIResponses(model, context, {
|
||||
...base,
|
||||
@@ -261,13 +262,18 @@ function buildParams(
|
||||
|
||||
if (model.reasoning) {
|
||||
if (options?.reasoningEffort || options?.reasoningSummary) {
|
||||
const effort = options?.reasoningEffort
|
||||
? (model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort)
|
||||
: "medium";
|
||||
params.reasoning = {
|
||||
effort: options?.reasoningEffort || "medium",
|
||||
effort: effort as NonNullable<typeof params.reasoning>["effort"],
|
||||
summary: options?.reasoningSummary || "auto",
|
||||
};
|
||||
params.include = ["reasoning.encrypted_content"];
|
||||
} else {
|
||||
params.reasoning = { effort: "none" };
|
||||
} else if (model.thinkingLevelMap?.off !== null) {
|
||||
params.reasoning = {
|
||||
effort: (model.thinkingLevelMap?.off ?? "none") as NonNullable<typeof params.reasoning>["effort"],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,22 +1,35 @@
|
||||
import type { Model } from "../types.js";
|
||||
import type { Api, Model } from "../types.js";
|
||||
|
||||
/** Workers AI endpoint. `{CLOUDFLARE_ACCOUNT_ID}` is substituted at request time. */
|
||||
/** Workers AI direct endpoint. */
|
||||
export const CLOUDFLARE_WORKERS_AI_BASE_URL =
|
||||
"https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/v1";
|
||||
|
||||
/** AI Gateway Unified API. https://developers.cloudflare.com/ai-gateway/usage/unified-api/ */
|
||||
export const CLOUDFLARE_AI_GATEWAY_COMPAT_BASE_URL =
|
||||
"https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/compat";
|
||||
|
||||
/** AI Gateway → OpenAI passthrough. Used until /compat supports /v1/responses. */
|
||||
export const CLOUDFLARE_AI_GATEWAY_OPENAI_BASE_URL =
|
||||
"https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/openai";
|
||||
|
||||
/** AI Gateway → Anthropic passthrough. */
|
||||
export const CLOUDFLARE_AI_GATEWAY_ANTHROPIC_BASE_URL =
|
||||
"https://gateway.ai.cloudflare.com/v1/{CLOUDFLARE_ACCOUNT_ID}/{CLOUDFLARE_GATEWAY_ID}/anthropic";
|
||||
|
||||
export function isCloudflareProvider(provider: string): boolean {
|
||||
return provider === "cloudflare-workers-ai";
|
||||
return provider === "cloudflare-workers-ai" || provider === "cloudflare-ai-gateway";
|
||||
}
|
||||
|
||||
/** Substitute `{VAR}` placeholders in a Cloudflare baseUrl from process.env. */
|
||||
export function resolveCloudflareBaseUrl(model: Model<"openai-completions">): string {
|
||||
export function resolveCloudflareBaseUrl(model: Model<Api>): string {
|
||||
const url = model.baseUrl;
|
||||
if (!url.includes("{")) return url;
|
||||
return url.replace(/\{([A-Z_][A-Z0-9_]*)\}/g, (_match, name: string) => {
|
||||
const baseUrl = url.replace(/\{([A-Z_][A-Z0-9_]*)\}/g, (_match, name: string) => {
|
||||
const value = process.env[name];
|
||||
if (!value) {
|
||||
throw new Error(`${name} is required for provider ${model.provider} but is not set.`);
|
||||
}
|
||||
return value;
|
||||
});
|
||||
return baseUrl;
|
||||
}
|
||||
|
||||
@@ -1,996 +0,0 @@
|
||||
/**
|
||||
* Google Gemini CLI / Antigravity provider.
|
||||
* Shared implementation for both google-gemini-cli and google-antigravity providers.
|
||||
* Uses the Cloud Code Assist API endpoint to access Gemini and Claude models.
|
||||
*/
|
||||
|
||||
import type { Content, ThinkingConfig } from "@google/genai";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingBudgets,
|
||||
ThinkingContent,
|
||||
ThinkingLevel,
|
||||
ToolCall,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import {
|
||||
convertMessages,
|
||||
convertTools,
|
||||
isThinkingPart,
|
||||
mapStopReasonString,
|
||||
mapToolChoice,
|
||||
retainThoughtSignature,
|
||||
} from "./google-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
/**
|
||||
* Thinking level for Gemini 3 models.
|
||||
* Mirrors Google's ThinkingLevel enum values.
|
||||
*/
|
||||
export type GoogleThinkingLevel = "THINKING_LEVEL_UNSPECIFIED" | "MINIMAL" | "LOW" | "MEDIUM" | "HIGH";
|
||||
|
||||
export interface GoogleGeminiCliOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any";
|
||||
/**
|
||||
* Thinking/reasoning configuration.
|
||||
* - Gemini 2.x models: use `budgetTokens` to set the thinking budget
|
||||
* - Gemini 3 models (gemini-3-pro-*, gemini-3-flash-*): use `level` instead
|
||||
*
|
||||
* When using `streamSimple`, this is handled automatically based on the model.
|
||||
*/
|
||||
thinking?: {
|
||||
enabled: boolean;
|
||||
/** Thinking budget in tokens. Use for Gemini 2.x models. */
|
||||
budgetTokens?: number;
|
||||
/** Thinking level. Use for Gemini 3 models (LOW/HIGH for Pro, MINIMAL/LOW/MEDIUM/HIGH for Flash). */
|
||||
level?: GoogleThinkingLevel;
|
||||
};
|
||||
projectId?: string;
|
||||
}
|
||||
|
||||
const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com";
|
||||
const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";
|
||||
const ANTIGRAVITY_AUTOPUSH_ENDPOINT = "https://autopush-cloudcode-pa.sandbox.googleapis.com";
|
||||
const ANTIGRAVITY_ENDPOINT_FALLBACKS = [
|
||||
ANTIGRAVITY_DAILY_ENDPOINT,
|
||||
ANTIGRAVITY_AUTOPUSH_ENDPOINT,
|
||||
DEFAULT_ENDPOINT,
|
||||
] as const;
|
||||
// Headers for Gemini CLI (prod endpoint)
|
||||
const GEMINI_CLI_HEADERS = {
|
||||
"User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
"X-Goog-Api-Client": "gl-node/22.17.0",
|
||||
"Client-Metadata": JSON.stringify({
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
}),
|
||||
};
|
||||
|
||||
// Headers for Antigravity (sandbox endpoint) - requires specific User-Agent
|
||||
const DEFAULT_ANTIGRAVITY_VERSION = "1.21.9";
|
||||
|
||||
function getAntigravityHeaders() {
|
||||
const version = process.env.PI_AI_ANTIGRAVITY_VERSION || DEFAULT_ANTIGRAVITY_VERSION;
|
||||
return {
|
||||
"User-Agent": `antigravity/${version} darwin/arm64`,
|
||||
};
|
||||
}
|
||||
|
||||
// Antigravity system instruction (compact version from CLIProxyAPI).
|
||||
const ANTIGRAVITY_SYSTEM_INSTRUCTION =
|
||||
"You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding." +
|
||||
"You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question." +
|
||||
"**Absolute paths only**" +
|
||||
"**Proactiveness**";
|
||||
|
||||
// Counter for generating unique tool call IDs
|
||||
let toolCallCounter = 0;
|
||||
|
||||
// Retry configuration
|
||||
const MAX_RETRIES = 3;
|
||||
const BASE_DELAY_MS = 1000;
|
||||
const MAX_EMPTY_STREAM_RETRIES = 2;
|
||||
const EMPTY_STREAM_BASE_DELAY_MS = 500;
|
||||
const CLAUDE_THINKING_BETA_HEADER = "interleaved-thinking-2025-05-14";
|
||||
|
||||
/**
|
||||
* Extract retry delay from Gemini error response (in milliseconds).
|
||||
* Checks headers first (Retry-After, x-ratelimit-reset, x-ratelimit-reset-after),
|
||||
* then parses body patterns like:
|
||||
* - "Your quota will reset after 39s"
|
||||
* - "Your quota will reset after 18h31m10s"
|
||||
* - "Please retry in Xs" or "Please retry in Xms"
|
||||
* - "retryDelay": "34.074824224s" (JSON field)
|
||||
*/
|
||||
export function extractRetryDelay(errorText: string, response?: Response | Headers): number | undefined {
|
||||
const normalizeDelay = (ms: number): number | undefined => (ms > 0 ? Math.ceil(ms + 1000) : undefined);
|
||||
|
||||
const headers = response instanceof Headers ? response : response?.headers;
|
||||
if (headers) {
|
||||
const retryAfter = headers.get("retry-after");
|
||||
if (retryAfter) {
|
||||
const retryAfterSeconds = Number(retryAfter);
|
||||
if (Number.isFinite(retryAfterSeconds)) {
|
||||
const delay = normalizeDelay(retryAfterSeconds * 1000);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
const retryAfterDate = new Date(retryAfter);
|
||||
const retryAfterMs = retryAfterDate.getTime();
|
||||
if (!Number.isNaN(retryAfterMs)) {
|
||||
const delay = normalizeDelay(retryAfterMs - Date.now());
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const rateLimitReset = headers.get("x-ratelimit-reset");
|
||||
if (rateLimitReset) {
|
||||
const resetSeconds = Number.parseInt(rateLimitReset, 10);
|
||||
if (!Number.isNaN(resetSeconds)) {
|
||||
const delay = normalizeDelay(resetSeconds * 1000 - Date.now());
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const rateLimitResetAfter = headers.get("x-ratelimit-reset-after");
|
||||
if (rateLimitResetAfter) {
|
||||
const resetAfterSeconds = Number(rateLimitResetAfter);
|
||||
if (Number.isFinite(resetAfterSeconds)) {
|
||||
const delay = normalizeDelay(resetAfterSeconds * 1000);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pattern 1: "Your quota will reset after ..." (formats: "18h31m10s", "10m15s", "6s", "39s")
|
||||
const durationMatch = errorText.match(/reset after (?:(\d+)h)?(?:(\d+)m)?(\d+(?:\.\d+)?)s/i);
|
||||
if (durationMatch) {
|
||||
const hours = durationMatch[1] ? parseInt(durationMatch[1], 10) : 0;
|
||||
const minutes = durationMatch[2] ? parseInt(durationMatch[2], 10) : 0;
|
||||
const seconds = parseFloat(durationMatch[3]);
|
||||
if (!Number.isNaN(seconds)) {
|
||||
const totalMs = ((hours * 60 + minutes) * 60 + seconds) * 1000;
|
||||
const delay = normalizeDelay(totalMs);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pattern 2: "Please retry in X[ms|s]"
|
||||
const retryInMatch = errorText.match(/Please retry in ([0-9.]+)(ms|s)/i);
|
||||
if (retryInMatch?.[1]) {
|
||||
const value = parseFloat(retryInMatch[1]);
|
||||
if (!Number.isNaN(value) && value > 0) {
|
||||
const ms = retryInMatch[2].toLowerCase() === "ms" ? value : value * 1000;
|
||||
const delay = normalizeDelay(ms);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pattern 3: "retryDelay": "34.074824224s" (JSON field in error details)
|
||||
const retryDelayMatch = errorText.match(/"retryDelay":\s*"([0-9.]+)(ms|s)"/i);
|
||||
if (retryDelayMatch?.[1]) {
|
||||
const value = parseFloat(retryDelayMatch[1]);
|
||||
if (!Number.isNaN(value) && value > 0) {
|
||||
const ms = retryDelayMatch[2].toLowerCase() === "ms" ? value : value * 1000;
|
||||
const delay = normalizeDelay(ms);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function needsClaudeThinkingBetaHeader(model: Model<"google-gemini-cli">): boolean {
|
||||
return model.provider === "google-antigravity" && model.id.startsWith("claude-") && model.reasoning;
|
||||
}
|
||||
|
||||
function isGemini3ProModel(modelId: string): boolean {
|
||||
return /gemini-3(?:\.1)?-pro/.test(modelId.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3FlashModel(modelId: string): boolean {
|
||||
return /gemini-3(?:\.1)?-flash/.test(modelId.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3Model(modelId: string): boolean {
|
||||
return isGemini3ProModel(modelId) || isGemini3FlashModel(modelId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an error is retryable (rate limit, server error, network error, etc.)
|
||||
*/
|
||||
function isRetryableError(status: number, errorText: string): boolean {
|
||||
if (status === 429 || status === 500 || status === 502 || status === 503 || status === 504) {
|
||||
return true;
|
||||
}
|
||||
return /resource.?exhausted|rate.?limit|overloaded|service.?unavailable|other.?side.?closed/i.test(errorText);
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract a clean, user-friendly error message from Google API error response.
|
||||
* Parses JSON error responses and returns just the message field.
|
||||
*/
|
||||
function extractErrorMessage(errorText: string): string {
|
||||
try {
|
||||
const parsed = JSON.parse(errorText) as { error?: { message?: string } };
|
||||
if (parsed.error?.message) {
|
||||
return parsed.error.message;
|
||||
}
|
||||
} catch {
|
||||
// Not JSON, return as-is
|
||||
}
|
||||
return errorText;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sleep for a given number of milliseconds, respecting abort signal.
|
||||
*/
|
||||
function sleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(new Error("Request was aborted"));
|
||||
return;
|
||||
}
|
||||
const timeout = setTimeout(resolve, ms);
|
||||
signal?.addEventListener("abort", () => {
|
||||
clearTimeout(timeout);
|
||||
reject(new Error("Request was aborted"));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
interface CloudCodeAssistRequest {
|
||||
project: string;
|
||||
model: string;
|
||||
request: {
|
||||
contents: Content[];
|
||||
sessionId?: string;
|
||||
systemInstruction?: { role?: string; parts: { text: string }[] };
|
||||
generationConfig?: {
|
||||
maxOutputTokens?: number;
|
||||
temperature?: number;
|
||||
thinkingConfig?: ThinkingConfig;
|
||||
};
|
||||
tools?: ReturnType<typeof convertTools>;
|
||||
toolConfig?: {
|
||||
functionCallingConfig: {
|
||||
mode: ReturnType<typeof mapToolChoice>;
|
||||
};
|
||||
};
|
||||
};
|
||||
requestType?: string;
|
||||
userAgent?: string;
|
||||
requestId?: string;
|
||||
}
|
||||
|
||||
interface CloudCodeAssistResponseChunk {
|
||||
response?: {
|
||||
candidates?: Array<{
|
||||
content?: {
|
||||
role: string;
|
||||
parts?: Array<{
|
||||
text?: string;
|
||||
thought?: boolean;
|
||||
thoughtSignature?: string;
|
||||
functionCall?: {
|
||||
name: string;
|
||||
args: Record<string, unknown>;
|
||||
id?: string;
|
||||
};
|
||||
}>;
|
||||
};
|
||||
finishReason?: string;
|
||||
}>;
|
||||
usageMetadata?: {
|
||||
promptTokenCount?: number;
|
||||
candidatesTokenCount?: number;
|
||||
thoughtsTokenCount?: number;
|
||||
totalTokenCount?: number;
|
||||
cachedContentTokenCount?: number;
|
||||
};
|
||||
modelVersion?: string;
|
||||
responseId?: string;
|
||||
};
|
||||
traceId?: string;
|
||||
}
|
||||
|
||||
export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGeminiCliOptions> = (
|
||||
model: Model<"google-gemini-cli">,
|
||||
context: Context,
|
||||
options?: GoogleGeminiCliOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "google-gemini-cli" as 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: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
// apiKey is JSON-encoded: { token, projectId }
|
||||
const apiKeyRaw = options?.apiKey;
|
||||
if (!apiKeyRaw) {
|
||||
throw new Error("Google Cloud Code Assist requires OAuth authentication. Use /login to authenticate.");
|
||||
}
|
||||
|
||||
let accessToken: string;
|
||||
let projectId: string;
|
||||
|
||||
try {
|
||||
const parsed = JSON.parse(apiKeyRaw) as { token: string; projectId: string };
|
||||
accessToken = parsed.token;
|
||||
projectId = parsed.projectId;
|
||||
} catch {
|
||||
throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate.");
|
||||
}
|
||||
|
||||
if (!accessToken || !projectId) {
|
||||
throw new Error("Missing token or projectId in Google Cloud credentials. Use /login to re-authenticate.");
|
||||
}
|
||||
|
||||
const isAntigravity = model.provider === "google-antigravity";
|
||||
const baseUrl = model.baseUrl?.trim();
|
||||
const endpoints = baseUrl ? [baseUrl] : isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT];
|
||||
|
||||
let requestBody = buildRequest(model, context, projectId, options, isAntigravity);
|
||||
const nextRequestBody = await options?.onPayload?.(requestBody, model);
|
||||
if (nextRequestBody !== undefined) {
|
||||
requestBody = nextRequestBody as CloudCodeAssistRequest;
|
||||
}
|
||||
const headers = isAntigravity ? getAntigravityHeaders() : GEMINI_CLI_HEADERS;
|
||||
|
||||
const requestHeaders = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
Accept: "text/event-stream",
|
||||
...headers,
|
||||
...(needsClaudeThinkingBetaHeader(model) ? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER } : {}),
|
||||
...options?.headers,
|
||||
};
|
||||
const requestBodyJson = JSON.stringify(requestBody);
|
||||
|
||||
// Fetch with retry logic for rate limits, transient errors, and endpoint fallbacks.
|
||||
// On 403/404, immediately try the next endpoint (no delay).
|
||||
// On 429/5xx, retry with backoff on the same or next endpoint.
|
||||
let response: Response | undefined;
|
||||
let lastError: Error | undefined;
|
||||
let requestUrl: string | undefined;
|
||||
let endpointIndex = 0;
|
||||
|
||||
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
try {
|
||||
const endpoint = endpoints[endpointIndex];
|
||||
requestUrl = `${endpoint}/v1internal:streamGenerateContent?alt=sse`;
|
||||
response = await fetch(requestUrl, {
|
||||
method: "POST",
|
||||
headers: requestHeaders,
|
||||
body: requestBodyJson,
|
||||
signal: options?.signal,
|
||||
});
|
||||
await options?.onResponse?.(
|
||||
{ status: response.status, headers: headersToRecord(response.headers) },
|
||||
model,
|
||||
);
|
||||
|
||||
if (response.ok) {
|
||||
break; // Success, exit retry loop
|
||||
}
|
||||
|
||||
const errorText = await response.text();
|
||||
|
||||
// On 403/404, cascade to the next endpoint immediately (no delay)
|
||||
if ((response.status === 403 || response.status === 404) && endpointIndex < endpoints.length - 1) {
|
||||
endpointIndex++;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if retryable (429, 5xx, network patterns)
|
||||
if (attempt < MAX_RETRIES && isRetryableError(response.status, errorText)) {
|
||||
// Advance endpoint if possible
|
||||
if (endpointIndex < endpoints.length - 1) {
|
||||
endpointIndex++;
|
||||
}
|
||||
|
||||
// Use server-provided delay or exponential backoff
|
||||
const serverDelay = extractRetryDelay(errorText, response);
|
||||
const delayMs = serverDelay ?? BASE_DELAY_MS * 2 ** attempt;
|
||||
|
||||
// Check if server delay exceeds max allowed (default: 60s)
|
||||
const maxDelayMs = options?.maxRetryDelayMs ?? 60000;
|
||||
if (maxDelayMs > 0 && serverDelay && serverDelay > maxDelayMs) {
|
||||
const delaySeconds = Math.ceil(serverDelay / 1000);
|
||||
throw new Error(
|
||||
`Server requested ${delaySeconds}s retry delay (max: ${Math.ceil(maxDelayMs / 1000)}s). ${extractErrorMessage(errorText)}`,
|
||||
);
|
||||
}
|
||||
|
||||
await sleep(delayMs, options?.signal);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Not retryable or max retries exceeded
|
||||
throw new Error(`Cloud Code Assist API error (${response.status}): ${extractErrorMessage(errorText)}`);
|
||||
} catch (error) {
|
||||
// Check for abort - fetch throws AbortError, our code throws "Request was aborted"
|
||||
if (error instanceof Error) {
|
||||
if (error.name === "AbortError" || error.message === "Request was aborted") {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
}
|
||||
// Extract detailed error message from fetch errors (Node includes cause)
|
||||
lastError = error instanceof Error ? error : new Error(String(error));
|
||||
if (lastError.message === "fetch failed" && lastError.cause instanceof Error) {
|
||||
lastError = new Error(`Network error: ${lastError.cause.message}`);
|
||||
}
|
||||
// Network errors are retryable
|
||||
if (attempt < MAX_RETRIES) {
|
||||
const delayMs = BASE_DELAY_MS * 2 ** attempt;
|
||||
await sleep(delayMs, options?.signal);
|
||||
continue;
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
}
|
||||
|
||||
if (!response || !response.ok) {
|
||||
throw lastError ?? new Error("Failed to get response after retries");
|
||||
}
|
||||
|
||||
let started = false;
|
||||
const ensureStarted = () => {
|
||||
if (!started) {
|
||||
stream.push({ type: "start", partial: output });
|
||||
started = true;
|
||||
}
|
||||
};
|
||||
|
||||
const resetOutput = () => {
|
||||
output.content = [];
|
||||
output.usage = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
output.stopReason = "stop";
|
||||
output.errorMessage = undefined;
|
||||
output.timestamp = Date.now();
|
||||
started = false;
|
||||
};
|
||||
|
||||
const streamResponse = async (activeResponse: Response): Promise<boolean> => {
|
||||
if (!activeResponse.body) {
|
||||
throw new Error("No response body");
|
||||
}
|
||||
|
||||
let hasContent = false;
|
||||
let currentBlock: TextContent | ThinkingContent | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
|
||||
// Read SSE stream
|
||||
const reader = activeResponse.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
// Set up abort handler to cancel reader when signal fires
|
||||
const abortHandler = () => {
|
||||
void reader.cancel().catch(() => {});
|
||||
};
|
||||
options?.signal?.addEventListener("abort", abortHandler);
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
// Check abort signal before each read
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop() || "";
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith("data:")) continue;
|
||||
|
||||
const jsonStr = line.slice(5).trim();
|
||||
if (!jsonStr) continue;
|
||||
|
||||
let chunk: CloudCodeAssistResponseChunk;
|
||||
try {
|
||||
chunk = JSON.parse(jsonStr);
|
||||
} catch {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Unwrap the response
|
||||
const responseData = chunk.response;
|
||||
if (!responseData) continue;
|
||||
// Cloud Code Assist mirrors Gemini's responseId field. Keep the first non-empty one.
|
||||
// A single streamed response should retain the same ID across chunks.
|
||||
output.responseId ||= responseData.responseId;
|
||||
|
||||
const candidate = responseData.candidates?.[0];
|
||||
if (candidate?.content?.parts) {
|
||||
for (const part of candidate.content.parts) {
|
||||
if (part.text !== undefined) {
|
||||
hasContent = true;
|
||||
const isThinking = isThinkingPart(part);
|
||||
if (
|
||||
!currentBlock ||
|
||||
(isThinking && currentBlock.type !== "thinking") ||
|
||||
(!isThinking && currentBlock.type !== "text")
|
||||
) {
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blocks.length - 1,
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
if (isThinking) {
|
||||
currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
|
||||
output.content.push(currentBlock);
|
||||
ensureStarted();
|
||||
stream.push({
|
||||
type: "thinking_start",
|
||||
contentIndex: blockIndex(),
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
ensureStarted();
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
}
|
||||
if (currentBlock.type === "thinking") {
|
||||
currentBlock.thinking += part.text;
|
||||
currentBlock.thinkingSignature = retainThoughtSignature(
|
||||
currentBlock.thinkingSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
currentBlock.text += part.text;
|
||||
currentBlock.textSignature = retainThoughtSignature(
|
||||
currentBlock.textSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (part.functionCall) {
|
||||
hasContent = true;
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
currentBlock = null;
|
||||
}
|
||||
|
||||
const providedId = part.functionCall.id;
|
||||
const needsNewId =
|
||||
!providedId ||
|
||||
output.content.some((b) => b.type === "toolCall" && b.id === providedId);
|
||||
const toolCallId = needsNewId
|
||||
? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
|
||||
: providedId;
|
||||
|
||||
const toolCall: ToolCall = {
|
||||
type: "toolCall",
|
||||
id: toolCallId,
|
||||
name: part.functionCall.name || "",
|
||||
arguments: (part.functionCall.args as Record<string, unknown>) ?? {},
|
||||
...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
|
||||
};
|
||||
|
||||
output.content.push(toolCall);
|
||||
ensureStarted();
|
||||
stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: JSON.stringify(toolCall.arguments),
|
||||
partial: output,
|
||||
});
|
||||
stream.push({
|
||||
type: "toolcall_end",
|
||||
contentIndex: blockIndex(),
|
||||
toolCall,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (candidate?.finishReason) {
|
||||
output.stopReason = mapStopReasonString(candidate.finishReason);
|
||||
if (output.content.some((b) => b.type === "toolCall")) {
|
||||
output.stopReason = "toolUse";
|
||||
}
|
||||
}
|
||||
|
||||
if (responseData.usageMetadata) {
|
||||
// promptTokenCount includes cachedContentTokenCount, so subtract to get fresh input
|
||||
const promptTokens = responseData.usageMetadata.promptTokenCount || 0;
|
||||
const cacheReadTokens = responseData.usageMetadata.cachedContentTokenCount || 0;
|
||||
output.usage = {
|
||||
input: promptTokens - cacheReadTokens,
|
||||
output:
|
||||
(responseData.usageMetadata.candidatesTokenCount || 0) +
|
||||
(responseData.usageMetadata.thoughtsTokenCount || 0),
|
||||
cacheRead: cacheReadTokens,
|
||||
cacheWrite: 0,
|
||||
totalTokens: responseData.usageMetadata.totalTokenCount || 0,
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
total: 0,
|
||||
},
|
||||
};
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
options?.signal?.removeEventListener("abort", abortHandler);
|
||||
}
|
||||
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return hasContent;
|
||||
};
|
||||
|
||||
let receivedContent = false;
|
||||
let currentResponse = response;
|
||||
|
||||
for (let emptyAttempt = 0; emptyAttempt <= MAX_EMPTY_STREAM_RETRIES; emptyAttempt++) {
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (emptyAttempt > 0) {
|
||||
const backoffMs = EMPTY_STREAM_BASE_DELAY_MS * 2 ** (emptyAttempt - 1);
|
||||
await sleep(backoffMs, options?.signal);
|
||||
|
||||
if (!requestUrl) {
|
||||
throw new Error("Missing request URL");
|
||||
}
|
||||
|
||||
currentResponse = await fetch(requestUrl, {
|
||||
method: "POST",
|
||||
headers: requestHeaders,
|
||||
body: requestBodyJson,
|
||||
signal: options?.signal,
|
||||
});
|
||||
await options?.onResponse?.(
|
||||
{ status: currentResponse.status, headers: headersToRecord(currentResponse.headers) },
|
||||
model,
|
||||
);
|
||||
|
||||
if (!currentResponse.ok) {
|
||||
const retryErrorText = await currentResponse.text();
|
||||
throw new Error(`Cloud Code Assist API error (${currentResponse.status}): ${retryErrorText}`);
|
||||
}
|
||||
}
|
||||
|
||||
const streamed = await streamResponse(currentResponse);
|
||||
if (streamed) {
|
||||
receivedContent = true;
|
||||
break;
|
||||
}
|
||||
|
||||
if (emptyAttempt < MAX_EMPTY_STREAM_RETRIES) {
|
||||
resetOutput();
|
||||
}
|
||||
}
|
||||
|
||||
if (!receivedContent) {
|
||||
throw new Error("Cloud Code Assist API returned an empty response");
|
||||
}
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
if ("index" in block) {
|
||||
delete (block as { index?: number }).index;
|
||||
}
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleGoogleGeminiCli: StreamFunction<"google-gemini-cli", SimpleStreamOptions> = (
|
||||
model: Model<"google-gemini-cli">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey;
|
||||
if (!apiKey) {
|
||||
throw new Error("Google Cloud Code Assist requires OAuth authentication. Use /login to authenticate.");
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
if (!options?.reasoning) {
|
||||
return streamGoogleGeminiCli(model, context, {
|
||||
...base,
|
||||
thinking: { enabled: false },
|
||||
} satisfies GoogleGeminiCliOptions);
|
||||
}
|
||||
|
||||
const effort = clampReasoning(options.reasoning)!;
|
||||
if (isGemini3Model(model.id)) {
|
||||
return streamGoogleGeminiCli(model, context, {
|
||||
...base,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
level: getGeminiCliThinkingLevel(effort, model.id),
|
||||
},
|
||||
} satisfies GoogleGeminiCliOptions);
|
||||
}
|
||||
|
||||
const defaultBudgets: ThinkingBudgets = {
|
||||
minimal: 1024,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 16384,
|
||||
};
|
||||
const budgets = { ...defaultBudgets, ...options.thinkingBudgets };
|
||||
|
||||
const minOutputTokens = 1024;
|
||||
let thinkingBudget = budgets[effort]!;
|
||||
const maxTokens = Math.min((base.maxTokens || 0) + thinkingBudget, model.maxTokens);
|
||||
|
||||
if (maxTokens <= thinkingBudget) {
|
||||
thinkingBudget = Math.max(0, maxTokens - minOutputTokens);
|
||||
}
|
||||
|
||||
return streamGoogleGeminiCli(model, context, {
|
||||
...base,
|
||||
maxTokens,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
budgetTokens: thinkingBudget,
|
||||
},
|
||||
} satisfies GoogleGeminiCliOptions);
|
||||
};
|
||||
|
||||
export function buildRequest(
|
||||
model: Model<"google-gemini-cli">,
|
||||
context: Context,
|
||||
projectId: string,
|
||||
options: GoogleGeminiCliOptions = {},
|
||||
isAntigravity = false,
|
||||
): CloudCodeAssistRequest {
|
||||
const contents = convertMessages(model, context);
|
||||
|
||||
const generationConfig: CloudCodeAssistRequest["request"]["generationConfig"] = {};
|
||||
if (options.temperature !== undefined) {
|
||||
generationConfig.temperature = options.temperature;
|
||||
}
|
||||
if (options.maxTokens !== undefined) {
|
||||
generationConfig.maxOutputTokens = options.maxTokens;
|
||||
}
|
||||
|
||||
// Thinking config
|
||||
if (options.thinking?.enabled && model.reasoning) {
|
||||
generationConfig.thinkingConfig = {
|
||||
includeThoughts: true,
|
||||
};
|
||||
// Gemini 3 models use thinkingLevel, older models use thinkingBudget
|
||||
if (options.thinking.level !== undefined) {
|
||||
// Cast to any since our GoogleThinkingLevel mirrors Google's ThinkingLevel enum values
|
||||
generationConfig.thinkingConfig.thinkingLevel = options.thinking.level as any;
|
||||
} else if (options.thinking.budgetTokens !== undefined) {
|
||||
generationConfig.thinkingConfig.thinkingBudget = options.thinking.budgetTokens;
|
||||
}
|
||||
} else if (model.reasoning && options.thinking && !options.thinking.enabled) {
|
||||
generationConfig.thinkingConfig = getDisabledThinkingConfig(model.id);
|
||||
}
|
||||
|
||||
const request: CloudCodeAssistRequest["request"] = {
|
||||
contents,
|
||||
};
|
||||
|
||||
request.sessionId = options.sessionId;
|
||||
|
||||
// System instruction must be object with parts, not plain string
|
||||
if (context.systemPrompt) {
|
||||
request.systemInstruction = {
|
||||
parts: [{ text: sanitizeSurrogates(context.systemPrompt) }],
|
||||
};
|
||||
}
|
||||
|
||||
if (Object.keys(generationConfig).length > 0) {
|
||||
request.generationConfig = generationConfig;
|
||||
}
|
||||
|
||||
if (context.tools && context.tools.length > 0) {
|
||||
// Claude models on Cloud Code Assist need the legacy `parameters` field;
|
||||
// the API translates it into Anthropic's `input_schema`.
|
||||
const useParameters = model.id.startsWith("claude-");
|
||||
request.tools = convertTools(context.tools, useParameters);
|
||||
if (options.toolChoice) {
|
||||
request.toolConfig = {
|
||||
functionCallingConfig: {
|
||||
mode: mapToolChoice(options.toolChoice),
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if (isAntigravity) {
|
||||
const existingParts = request.systemInstruction?.parts ?? [];
|
||||
request.systemInstruction = {
|
||||
role: "user",
|
||||
parts: [
|
||||
{ text: ANTIGRAVITY_SYSTEM_INSTRUCTION },
|
||||
{ text: `Please ignore following [ignore]${ANTIGRAVITY_SYSTEM_INSTRUCTION}[/ignore]` },
|
||||
...existingParts,
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
project: projectId,
|
||||
model: model.id,
|
||||
request,
|
||||
...(isAntigravity ? { requestType: "agent" } : {}),
|
||||
userAgent: isAntigravity ? "antigravity" : "pi-coding-agent",
|
||||
requestId: `${isAntigravity ? "agent" : "pi"}-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`,
|
||||
};
|
||||
}
|
||||
|
||||
type ClampedThinkingLevel = Exclude<ThinkingLevel, "xhigh">;
|
||||
|
||||
function getDisabledThinkingConfig(modelId: string): ThinkingConfig {
|
||||
// Google docs: Gemini 3.1 Pro cannot disable thinking, and Gemini 3 Flash / Flash-Lite
|
||||
// do not support full thinking-off either. For Gemini 3 models, use the lowest supported
|
||||
// thinkingLevel without includeThoughts so hidden thinking remains invisible to pi.
|
||||
if (isGemini3ProModel(modelId)) {
|
||||
return { thinkingLevel: "LOW" as any };
|
||||
}
|
||||
if (isGemini3FlashModel(modelId)) {
|
||||
return { thinkingLevel: "MINIMAL" as any };
|
||||
}
|
||||
|
||||
// Gemini 2.x supports disabling via thinkingBudget = 0.
|
||||
return { thinkingBudget: 0 };
|
||||
}
|
||||
|
||||
function getGeminiCliThinkingLevel(effort: ClampedThinkingLevel, modelId: string): GoogleThinkingLevel {
|
||||
if (isGemini3ProModel(modelId)) {
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
return "MINIMAL";
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
return "MEDIUM";
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,5 @@
|
||||
/**
|
||||
* Shared utilities for Google Generative AI and Google Cloud Code Assist providers.
|
||||
* Shared utilities for Google Generative AI and Google Vertex providers.
|
||||
*/
|
||||
|
||||
import { type Content, FinishReason, FunctionCallingConfigMode, type Part } from "@google/genai";
|
||||
@@ -7,7 +7,13 @@ import type { Context, ImageContent, Model, StopReason, TextContent, Tool } from
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
type GoogleApiType = "google-generative-ai" | "google-gemini-cli" | "google-vertex";
|
||||
type GoogleApiType = "google-generative-ai" | "google-vertex";
|
||||
|
||||
/**
|
||||
* Thinking level for Gemini 3 models.
|
||||
* Mirrors Google's ThinkingLevel enum values.
|
||||
*/
|
||||
export type GoogleThinkingLevel = "THINKING_LEVEL_UNSPECIFIED" | "MINIMAL" | "LOW" | "MEDIUM" | "HIGH";
|
||||
|
||||
/**
|
||||
* Determines whether a streamed Gemini `Part` should be treated as "thinking".
|
||||
@@ -45,11 +51,6 @@ export function retainThoughtSignature(existing: string | undefined, incoming: s
|
||||
// Thought signatures must be base64 for Google APIs (TYPE_BYTES).
|
||||
const base64SignaturePattern = /^[A-Za-z0-9+/]+={0,2}$/;
|
||||
|
||||
// Sentinel value that tells the Gemini API to skip thought signature validation.
|
||||
// Used for unsigned function call parts (e.g. replayed from providers without thought signatures).
|
||||
// See: https://ai.google.dev/gemini-api/docs/thought-signatures
|
||||
const SKIP_THOUGHT_SIGNATURE = "skip_thought_signature_validator";
|
||||
|
||||
function isValidThoughtSignature(signature: string | undefined): boolean {
|
||||
if (!signature) return false;
|
||||
if (signature.length % 4 !== 0) return false;
|
||||
@@ -129,7 +130,7 @@ export function convertMessages<T extends GoogleApiType>(model: Model<T>, contex
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "text") {
|
||||
// Skip empty text blocks - they can cause issues with some models (e.g. Claude via Antigravity)
|
||||
// Skip empty text blocks
|
||||
if (!block.text || block.text.trim() === "") continue;
|
||||
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.textSignature);
|
||||
parts.push({
|
||||
@@ -155,18 +156,13 @@ export function convertMessages<T extends GoogleApiType>(model: Model<T>, contex
|
||||
}
|
||||
} else if (block.type === "toolCall") {
|
||||
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.thoughtSignature);
|
||||
// Gemini 3 requires thoughtSignature on all function calls when thinking mode is enabled.
|
||||
// Use the skip_thought_signature_validator sentinel for unsigned function calls
|
||||
// (e.g. replayed from providers without thought signatures like Claude via Antigravity).
|
||||
const isGemini3 = model.id.toLowerCase().includes("gemini-3");
|
||||
const effectiveSignature = thoughtSignature || (isGemini3 ? SKIP_THOUGHT_SIGNATURE : undefined);
|
||||
const part: Part = {
|
||||
functionCall: {
|
||||
name: block.name,
|
||||
args: block.arguments ?? {},
|
||||
...(requiresToolCallId(model.id) ? { id: block.id } : {}),
|
||||
},
|
||||
...(effectiveSignature && { thoughtSignature: effectiveSignature }),
|
||||
...(thoughtSignature && { thoughtSignature }),
|
||||
};
|
||||
parts.push(part);
|
||||
}
|
||||
@@ -190,7 +186,7 @@ export function convertMessages<T extends GoogleApiType>(model: Model<T>, contex
|
||||
|
||||
// Gemini 3+ models support multimodal function responses with images nested inside
|
||||
// functionResponse.parts. Claude and other non-Gemini models behind Cloud Code Assist /
|
||||
// Antigravity also accept this shape. Gemini < 3 still needs a separate user image turn.
|
||||
// Gemini < 3 still needs a separate user image turn.
|
||||
const modelSupportsMultimodalFunctionResponse = supportsMultimodalFunctionResponse(model.id);
|
||||
|
||||
// Use "output" key for success, "error" key for errors as per SDK documentation
|
||||
|
||||
@@ -7,7 +7,7 @@ import {
|
||||
type ThinkingConfig,
|
||||
ThinkingLevel,
|
||||
} from "@google/genai";
|
||||
import { calculateCost } from "../models.js";
|
||||
import { calculateCost, clampThinkingLevel } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
@@ -24,7 +24,7 @@ import type {
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import type { GoogleThinkingLevel } from "./google-gemini-cli.js";
|
||||
import type { GoogleThinkingLevel } from "./google-shared.js";
|
||||
import {
|
||||
convertMessages,
|
||||
convertTools,
|
||||
@@ -33,7 +33,7 @@ import {
|
||||
mapToolChoice,
|
||||
retainThoughtSignature,
|
||||
} from "./google-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
|
||||
export interface GoogleVertexOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any";
|
||||
@@ -305,7 +305,8 @@ export const streamSimpleGoogleVertex: StreamFunction<"google-vertex", SimpleStr
|
||||
} satisfies GoogleVertexOptions);
|
||||
}
|
||||
|
||||
const effort = clampReasoning(options.reasoning)!;
|
||||
const clampedReasoning = clampThinkingLevel(model, options.reasoning);
|
||||
const effort = (clampedReasoning === "off" ? "high" : clampedReasoning) as ClampedThinkingLevel;
|
||||
const geminiModel = model as unknown as Model<"google-generative-ai">;
|
||||
|
||||
if (isGemini3ProModel(geminiModel) || isGemini3FlashModel(geminiModel)) {
|
||||
|
||||
@@ -5,7 +5,7 @@ import {
|
||||
type ThinkingConfig,
|
||||
} from "@google/genai";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost } from "../models.js";
|
||||
import { calculateCost, clampThinkingLevel } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
@@ -22,7 +22,7 @@ import type {
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import type { GoogleThinkingLevel } from "./google-gemini-cli.js";
|
||||
import type { GoogleThinkingLevel } from "./google-shared.js";
|
||||
import {
|
||||
convertMessages,
|
||||
convertTools,
|
||||
@@ -31,7 +31,7 @@ import {
|
||||
mapToolChoice,
|
||||
retainThoughtSignature,
|
||||
} from "./google-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
|
||||
export interface GoogleOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any";
|
||||
@@ -290,7 +290,8 @@ export const streamSimpleGoogle: StreamFunction<"google-generative-ai", SimpleSt
|
||||
return streamGoogle(model, context, { ...base, thinking: { enabled: false } } satisfies GoogleOptions);
|
||||
}
|
||||
|
||||
const effort = clampReasoning(options.reasoning)!;
|
||||
const clampedReasoning = clampThinkingLevel(model, options.reasoning);
|
||||
const effort = (clampedReasoning === "off" ? "high" : clampedReasoning) as ClampedThinkingLevel;
|
||||
const googleModel = model as Model<"google-generative-ai">;
|
||||
|
||||
if (isGemini3ProModel(googleModel) || isGemini3FlashModel(googleModel) || isGemma4Model(googleModel)) {
|
||||
|
||||
@@ -7,7 +7,7 @@ import type {
|
||||
FunctionTool,
|
||||
} from "@mistralai/mistralai/models/components";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost } from "../models.js";
|
||||
import { calculateCost, clampThinkingLevel } from "../models.js";
|
||||
import type {
|
||||
AssistantMessage,
|
||||
Context,
|
||||
@@ -26,7 +26,7 @@ import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { shortHash } from "../utils/hash.js";
|
||||
import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
const MISTRAL_TOOL_CALL_ID_LENGTH = 9;
|
||||
@@ -119,13 +119,15 @@ export const streamSimpleMistral: StreamFunction<"mistral-conversations", Simple
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoning = clampReasoning(options?.reasoning);
|
||||
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
|
||||
const reasoning = clampedReasoning === "off" ? undefined : clampedReasoning;
|
||||
const shouldUseReasoning = model.reasoning && reasoning !== undefined;
|
||||
|
||||
return streamMistral(model, context, {
|
||||
...base,
|
||||
promptMode: shouldUseReasoning && usesPromptModeReasoning(model) ? "reasoning" : undefined,
|
||||
reasoningEffort: shouldUseReasoning && usesReasoningEffort(model) ? mapReasoningEffort(reasoning) : undefined,
|
||||
reasoningEffort:
|
||||
shouldUseReasoning && usesReasoningEffort(model) ? mapReasoningEffort(model, reasoning) : undefined,
|
||||
} satisfies MistralOptions);
|
||||
};
|
||||
|
||||
@@ -587,15 +589,18 @@ function buildToolResultText(text: string, hasImages: boolean, supportsImages: b
|
||||
}
|
||||
|
||||
function usesReasoningEffort(model: Model<"mistral-conversations">): boolean {
|
||||
return model.id === "mistral-small-2603" || model.id === "mistral-small-latest";
|
||||
return model.id === "mistral-small-2603" || model.id === "mistral-small-latest" || model.id === "mistral-medium-3.5";
|
||||
}
|
||||
|
||||
function usesPromptModeReasoning(model: Model<"mistral-conversations">): boolean {
|
||||
return model.reasoning && !usesReasoningEffort(model);
|
||||
}
|
||||
|
||||
function mapReasoningEffort(_level: Exclude<SimpleStreamOptions["reasoning"], undefined>): MistralReasoningEffort {
|
||||
return "high";
|
||||
function mapReasoningEffort(
|
||||
model: Model<"mistral-conversations">,
|
||||
level: Exclude<SimpleStreamOptions["reasoning"], undefined>,
|
||||
): MistralReasoningEffort {
|
||||
return (model.thinkingLevelMap?.[level] ?? "high") as MistralReasoningEffort;
|
||||
}
|
||||
|
||||
function mapToolChoice(
|
||||
|
||||
@@ -21,7 +21,8 @@ if (typeof process !== "undefined" && (process.versions?.node || process.version
|
||||
}
|
||||
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { supportsXhigh } from "../models.js";
|
||||
import { clampThinkingLevel } from "../models.js";
|
||||
import { registerSessionResourceCleanup } from "../session-resources.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
@@ -32,10 +33,15 @@ import type {
|
||||
StreamOptions,
|
||||
Usage,
|
||||
} from "../types.js";
|
||||
import {
|
||||
appendAssistantMessageDiagnostic,
|
||||
createAssistantMessageDiagnostic,
|
||||
formatThrownValue,
|
||||
} from "../utils/diagnostics.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
|
||||
// ============================================================================
|
||||
// Configuration
|
||||
@@ -46,6 +52,7 @@ const JWT_CLAIM_PATH = "https://api.openai.com/auth" as const;
|
||||
const MAX_RETRIES = 3;
|
||||
const BASE_DELAY_MS = 1000;
|
||||
const CODEX_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
|
||||
const WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE = 1009;
|
||||
|
||||
const CODEX_RESPONSE_STATUSES = new Set<CodexResponseStatus>([
|
||||
"completed",
|
||||
@@ -74,6 +81,7 @@ interface RequestBody {
|
||||
store?: boolean;
|
||||
stream?: boolean;
|
||||
instructions?: string;
|
||||
previous_response_id?: string;
|
||||
input?: ResponseInput;
|
||||
tools?: OpenAITool[];
|
||||
tool_choice?: "auto";
|
||||
@@ -164,9 +172,13 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
|
||||
websocketRequestId,
|
||||
);
|
||||
const bodyJson = JSON.stringify(body);
|
||||
const transport = options?.transport || "sse";
|
||||
const transport = options?.transport || "auto";
|
||||
const websocketDisabledForSession = transport !== "sse" && isWebSocketSseFallbackActive(options?.sessionId);
|
||||
if (websocketDisabledForSession) {
|
||||
recordWebSocketSseFallback(options?.sessionId);
|
||||
}
|
||||
|
||||
if (transport !== "sse") {
|
||||
if (transport !== "sse" && !websocketDisabledForSession) {
|
||||
let websocketStarted = false;
|
||||
try {
|
||||
await processWebSocketStream(
|
||||
@@ -193,9 +205,25 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
|
||||
stream.end();
|
||||
return;
|
||||
} catch (error) {
|
||||
if (transport === "websocket" || websocketStarted) {
|
||||
const aborted = options?.signal?.aborted;
|
||||
if (aborted || isCodexNonTransportError(error)) {
|
||||
throw error;
|
||||
}
|
||||
appendAssistantMessageDiagnostic(
|
||||
output,
|
||||
createAssistantMessageDiagnostic("provider_transport_failure", error, {
|
||||
configuredTransport: transport,
|
||||
fallbackTransport: websocketStarted ? undefined : "sse",
|
||||
eventsEmitted: websocketStarted,
|
||||
phase: websocketStarted ? "after_message_stream_start" : "before_message_stream_start",
|
||||
requestBytes: new TextEncoder().encode(bodyJson).byteLength,
|
||||
}),
|
||||
);
|
||||
recordWebSocketFailure(options?.sessionId, error);
|
||||
if (websocketStarted) {
|
||||
throw error;
|
||||
}
|
||||
recordWebSocketSseFallback(options?.sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -298,7 +326,8 @@ export const streamSimpleOpenAICodexResponses: StreamFunction<"openai-codex-resp
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
|
||||
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
|
||||
|
||||
return streamOpenAICodexResponses(model, context, {
|
||||
...base,
|
||||
@@ -345,27 +374,21 @@ function buildRequestBody(
|
||||
}
|
||||
|
||||
if (options?.reasoningEffort !== undefined) {
|
||||
body.reasoning = {
|
||||
effort: clampReasoningEffort(model.id, options.reasoningEffort),
|
||||
summary: options.reasoningSummary ?? "auto",
|
||||
};
|
||||
const effort =
|
||||
options.reasoningEffort === "none"
|
||||
? (model.thinkingLevelMap?.off ?? "none")
|
||||
: (model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort);
|
||||
if (effort !== null) {
|
||||
body.reasoning = {
|
||||
effort,
|
||||
summary: options.reasoningSummary ?? "auto",
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
return body;
|
||||
}
|
||||
|
||||
function clampReasoningEffort(modelId: string, effort: string): string {
|
||||
const id = modelId.includes("/") ? modelId.split("/").pop()! : modelId;
|
||||
if (
|
||||
(id.startsWith("gpt-5.2") || id.startsWith("gpt-5.3") || id.startsWith("gpt-5.4") || id.startsWith("gpt-5.5")) &&
|
||||
effort === "minimal"
|
||||
)
|
||||
return "low";
|
||||
if (id === "gpt-5.1" && effort === "xhigh") return "high";
|
||||
if (id === "gpt-5.1-codex-mini") return effort === "high" || effort === "xhigh" ? "high" : "medium";
|
||||
return effort;
|
||||
}
|
||||
|
||||
function getServiceTierCostMultiplier(
|
||||
model: Pick<Model<"openai-codex-responses">, "id">,
|
||||
serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
|
||||
@@ -438,6 +461,34 @@ async function processStream(
|
||||
});
|
||||
}
|
||||
|
||||
class CodexApiError extends Error {
|
||||
readonly code?: string;
|
||||
readonly payload?: Record<string, unknown>;
|
||||
|
||||
constructor(message: string, options?: { code?: string; payload?: Record<string, unknown>; cause?: unknown }) {
|
||||
super(message);
|
||||
this.name = "CodexApiError";
|
||||
this.code = options?.code;
|
||||
this.payload = options?.payload;
|
||||
this.cause = options?.cause;
|
||||
}
|
||||
}
|
||||
|
||||
class CodexProtocolError extends Error {
|
||||
readonly payload?: unknown;
|
||||
|
||||
constructor(message: string, options?: { payload?: unknown; cause?: unknown }) {
|
||||
super(message);
|
||||
this.name = "CodexProtocolError";
|
||||
this.payload = options?.payload;
|
||||
this.cause = options?.cause;
|
||||
}
|
||||
}
|
||||
|
||||
function isCodexNonTransportError(error: unknown): boolean {
|
||||
return error instanceof CodexApiError || error instanceof CodexProtocolError;
|
||||
}
|
||||
|
||||
async function* mapCodexEvents(events: AsyncIterable<Record<string, unknown>>): AsyncGenerator<ResponseStreamEvent> {
|
||||
for await (const event of events) {
|
||||
const type = typeof event.type === "string" ? event.type : undefined;
|
||||
@@ -446,12 +497,17 @@ async function* mapCodexEvents(events: AsyncIterable<Record<string, unknown>>):
|
||||
if (type === "error") {
|
||||
const code = (event as { code?: string }).code || "";
|
||||
const message = (event as { message?: string }).message || "";
|
||||
throw new Error(`Codex error: ${message || code || JSON.stringify(event)}`);
|
||||
throw new CodexApiError(`Codex error: ${message || code || JSON.stringify(event)}`, {
|
||||
code: code || undefined,
|
||||
payload: event,
|
||||
});
|
||||
}
|
||||
|
||||
if (type === "response.failed") {
|
||||
const msg = (event as { response?: { error?: { message?: string } } }).response?.error?.message;
|
||||
throw new Error(msg || "Codex response failed");
|
||||
const response = (event as { response?: { error?: { code?: string; message?: string } } }).response;
|
||||
const code = response?.error?.code;
|
||||
const message = response?.error?.message;
|
||||
throw new CodexApiError(message || "Codex response failed", { code, payload: event });
|
||||
}
|
||||
|
||||
if (type === "response.done" || type === "response.completed" || type === "response.incomplete") {
|
||||
@@ -502,8 +558,13 @@ async function* parseSSE(response: Response): AsyncGenerator<Record<string, unkn
|
||||
const data = dataLines.join("\n").trim();
|
||||
if (data && data !== "[DONE]") {
|
||||
try {
|
||||
yield JSON.parse(data);
|
||||
} catch {}
|
||||
yield JSON.parse(data) as Record<string, unknown>;
|
||||
} catch (cause) {
|
||||
throw new CodexProtocolError(`Invalid Codex SSE JSON: ${formatThrownValue(cause)}`, {
|
||||
cause,
|
||||
payload: data,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
idx = buffer.indexOf("\n\n");
|
||||
@@ -536,13 +597,114 @@ interface WebSocketLike {
|
||||
removeEventListener(type: WebSocketEventType, listener: WebSocketListener): void;
|
||||
}
|
||||
|
||||
interface CachedWebSocketContinuationState {
|
||||
lastRequestBody: RequestBody;
|
||||
lastResponseId: string;
|
||||
lastResponseItems: ResponseInput;
|
||||
}
|
||||
|
||||
interface CachedWebSocketConnection {
|
||||
socket: WebSocketLike;
|
||||
busy: boolean;
|
||||
idleTimer?: ReturnType<typeof setTimeout>;
|
||||
continuation?: CachedWebSocketContinuationState;
|
||||
}
|
||||
|
||||
export interface OpenAICodexWebSocketDebugStats {
|
||||
requests: number;
|
||||
connectionsCreated: number;
|
||||
connectionsReused: number;
|
||||
cachedContextRequests: number;
|
||||
storeTrueRequests: number;
|
||||
fullContextRequests: number;
|
||||
deltaRequests: number;
|
||||
lastInputItems: number;
|
||||
lastDeltaInputItems?: number;
|
||||
lastPreviousResponseId?: string;
|
||||
websocketFailures: number;
|
||||
sseFallbacks: number;
|
||||
websocketFallbackActive?: boolean;
|
||||
lastWebSocketError?: string;
|
||||
}
|
||||
|
||||
const websocketSessionCache = new Map<string, CachedWebSocketConnection>();
|
||||
const websocketDebugStats = new Map<string, OpenAICodexWebSocketDebugStats>();
|
||||
const websocketSseFallbackSessions = new Set<string>();
|
||||
|
||||
function getOrCreateWebSocketDebugStats(sessionId: string): OpenAICodexWebSocketDebugStats {
|
||||
let stats = websocketDebugStats.get(sessionId);
|
||||
if (!stats) {
|
||||
stats = {
|
||||
requests: 0,
|
||||
connectionsCreated: 0,
|
||||
connectionsReused: 0,
|
||||
cachedContextRequests: 0,
|
||||
storeTrueRequests: 0,
|
||||
fullContextRequests: 0,
|
||||
deltaRequests: 0,
|
||||
lastInputItems: 0,
|
||||
websocketFailures: 0,
|
||||
sseFallbacks: 0,
|
||||
};
|
||||
websocketDebugStats.set(sessionId, stats);
|
||||
}
|
||||
return stats;
|
||||
}
|
||||
|
||||
export function getOpenAICodexWebSocketDebugStats(sessionId: string): OpenAICodexWebSocketDebugStats | undefined {
|
||||
const stats = websocketDebugStats.get(sessionId);
|
||||
return stats ? { ...stats } : undefined;
|
||||
}
|
||||
|
||||
export function resetOpenAICodexWebSocketDebugStats(sessionId?: string): void {
|
||||
if (sessionId) {
|
||||
websocketDebugStats.delete(sessionId);
|
||||
websocketSseFallbackSessions.delete(sessionId);
|
||||
return;
|
||||
}
|
||||
websocketDebugStats.clear();
|
||||
websocketSseFallbackSessions.clear();
|
||||
}
|
||||
|
||||
export function closeOpenAICodexWebSocketSessions(sessionId?: string): void {
|
||||
const closeEntry = (entry: CachedWebSocketConnection) => {
|
||||
if (entry.idleTimer) clearTimeout(entry.idleTimer);
|
||||
closeWebSocketSilently(entry.socket, 1000, "debug_close");
|
||||
};
|
||||
if (sessionId) {
|
||||
const entry = websocketSessionCache.get(sessionId);
|
||||
if (entry) closeEntry(entry);
|
||||
websocketSessionCache.delete(sessionId);
|
||||
return;
|
||||
}
|
||||
for (const entry of websocketSessionCache.values()) {
|
||||
closeEntry(entry);
|
||||
}
|
||||
websocketSessionCache.clear();
|
||||
}
|
||||
|
||||
registerSessionResourceCleanup(closeOpenAICodexWebSocketSessions);
|
||||
|
||||
function isWebSocketSseFallbackActive(sessionId: string | undefined): boolean {
|
||||
return sessionId ? websocketSseFallbackSessions.has(sessionId) : false;
|
||||
}
|
||||
|
||||
function recordWebSocketSseFallback(sessionId: string | undefined): void {
|
||||
if (!sessionId) return;
|
||||
const stats = getOrCreateWebSocketDebugStats(sessionId);
|
||||
stats.sseFallbacks++;
|
||||
stats.websocketFallbackActive = isWebSocketSseFallbackActive(sessionId);
|
||||
}
|
||||
|
||||
function recordWebSocketFailure(sessionId: string | undefined, error: unknown): void {
|
||||
if (!sessionId) return;
|
||||
websocketSseFallbackSessions.add(sessionId);
|
||||
|
||||
const stats = getOrCreateWebSocketDebugStats(sessionId);
|
||||
stats.websocketFailures++;
|
||||
stats.lastWebSocketError = formatThrownValue(error);
|
||||
stats.websocketFallbackActive = true;
|
||||
}
|
||||
|
||||
type WebSocketConstructor = new (
|
||||
url: string,
|
||||
@@ -555,6 +717,20 @@ function getWebSocketConstructor(): WebSocketConstructor | null {
|
||||
return ctor as unknown as WebSocketConstructor;
|
||||
}
|
||||
|
||||
class WebSocketCloseError extends Error {
|
||||
readonly code?: number;
|
||||
readonly reason?: string;
|
||||
readonly wasClean?: boolean;
|
||||
|
||||
constructor(message: string, options?: { code?: number; reason?: string; wasClean?: boolean }) {
|
||||
super(message);
|
||||
this.name = "WebSocketCloseError";
|
||||
this.code = options?.code;
|
||||
this.reason = options?.reason;
|
||||
this.wasClean = options?.wasClean;
|
||||
}
|
||||
}
|
||||
|
||||
function getWebSocketReadyState(socket: WebSocketLike): number | undefined {
|
||||
const readyState = (socket as { readyState?: unknown }).readyState;
|
||||
return typeof readyState === "number" ? readyState : undefined;
|
||||
@@ -650,11 +826,17 @@ async function acquireWebSocket(
|
||||
headers: Headers,
|
||||
sessionId: string | undefined,
|
||||
signal?: AbortSignal,
|
||||
): Promise<{ socket: WebSocketLike; release: (options?: { keep?: boolean }) => void }> {
|
||||
): Promise<{
|
||||
socket: WebSocketLike;
|
||||
entry?: CachedWebSocketConnection;
|
||||
reused: boolean;
|
||||
release: (options?: { keep?: boolean }) => void;
|
||||
}> {
|
||||
if (!sessionId) {
|
||||
const socket = await connectWebSocket(url, headers, signal);
|
||||
return {
|
||||
socket,
|
||||
reused: false,
|
||||
release: ({ keep } = {}) => {
|
||||
if (keep === false) {
|
||||
closeWebSocketSilently(socket);
|
||||
@@ -675,6 +857,8 @@ async function acquireWebSocket(
|
||||
cached.busy = true;
|
||||
return {
|
||||
socket: cached.socket,
|
||||
entry: cached,
|
||||
reused: true,
|
||||
release: ({ keep } = {}) => {
|
||||
if (!keep || !isWebSocketReusable(cached.socket)) {
|
||||
closeWebSocketSilently(cached.socket);
|
||||
@@ -690,6 +874,7 @@ async function acquireWebSocket(
|
||||
const socket = await connectWebSocket(url, headers, signal);
|
||||
return {
|
||||
socket,
|
||||
reused: false,
|
||||
release: () => {
|
||||
closeWebSocketSilently(socket);
|
||||
},
|
||||
@@ -706,6 +891,8 @@ async function acquireWebSocket(
|
||||
websocketSessionCache.set(sessionId, entry);
|
||||
return {
|
||||
socket,
|
||||
entry,
|
||||
reused: false,
|
||||
release: ({ keep } = {}) => {
|
||||
if (!keep || !isWebSocketReusable(entry.socket)) {
|
||||
closeWebSocketSilently(entry.socket);
|
||||
@@ -722,11 +909,22 @@ async function acquireWebSocket(
|
||||
}
|
||||
|
||||
function extractWebSocketError(event: unknown): Error {
|
||||
if (event && typeof event === "object" && "message" in event) {
|
||||
const message = (event as { message?: unknown }).message;
|
||||
if (event && typeof event === "object") {
|
||||
const message = "message" in event ? (event as { message?: unknown }).message : undefined;
|
||||
if (typeof message === "string" && message.length > 0) {
|
||||
return new Error(message);
|
||||
}
|
||||
|
||||
const nestedError = "error" in event ? (event as { error?: unknown }).error : undefined;
|
||||
if (nestedError instanceof Error && nestedError.message.length > 0) {
|
||||
return nestedError;
|
||||
}
|
||||
if (nestedError && typeof nestedError === "object" && "message" in nestedError) {
|
||||
const nestedMessage = (nestedError as { message?: unknown }).message;
|
||||
if (typeof nestedMessage === "string" && nestedMessage.length > 0) {
|
||||
return new Error(nestedMessage);
|
||||
}
|
||||
}
|
||||
}
|
||||
return new Error("WebSocket error");
|
||||
}
|
||||
@@ -735,9 +933,17 @@ function extractWebSocketCloseError(event: unknown): Error {
|
||||
if (event && typeof event === "object") {
|
||||
const code = "code" in event ? (event as { code?: unknown }).code : undefined;
|
||||
const reason = "reason" in event ? (event as { reason?: unknown }).reason : undefined;
|
||||
const wasClean = "wasClean" in event ? (event as { wasClean?: unknown }).wasClean : undefined;
|
||||
const codeText = typeof code === "number" ? ` ${code}` : "";
|
||||
const reasonText = typeof reason === "string" && reason.length > 0 ? ` ${reason}` : "";
|
||||
return new Error(`WebSocket closed${codeText}${reasonText}`.trim());
|
||||
let reasonText = typeof reason === "string" && reason.length > 0 ? ` ${reason}` : "";
|
||||
if (!reasonText && code === WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE) {
|
||||
reasonText = " message too big";
|
||||
}
|
||||
return new WebSocketCloseError(`WebSocket closed${codeText}${reasonText}`.trim(), {
|
||||
code: typeof code === "number" ? code : undefined,
|
||||
reason: typeof reason === "string" && reason.length > 0 ? reason : undefined,
|
||||
wasClean: typeof wasClean === "boolean" ? wasClean : undefined,
|
||||
});
|
||||
}
|
||||
return new Error("WebSocket closed");
|
||||
}
|
||||
@@ -775,10 +981,11 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy
|
||||
|
||||
const onMessage: WebSocketListener = (event) => {
|
||||
void (async () => {
|
||||
if (!event || typeof event !== "object" || !("data" in event)) return;
|
||||
const text = await decodeWebSocketData((event as { data?: unknown }).data);
|
||||
if (!text) return;
|
||||
let text: string | null = null;
|
||||
try {
|
||||
if (!event || typeof event !== "object" || !("data" in event)) return;
|
||||
text = await decodeWebSocketData((event as { data?: unknown }).data);
|
||||
if (!text) return;
|
||||
const parsed = JSON.parse(text) as Record<string, unknown>;
|
||||
const type = typeof parsed.type === "string" ? parsed.type : "";
|
||||
if (type === "response.completed" || type === "response.done" || type === "response.incomplete") {
|
||||
@@ -787,7 +994,14 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy
|
||||
}
|
||||
queue.push(parsed);
|
||||
wake();
|
||||
} catch {}
|
||||
} catch (cause) {
|
||||
failed = new CodexProtocolError(`Invalid Codex WebSocket JSON: ${formatThrownValue(cause)}`, {
|
||||
cause,
|
||||
payload: text,
|
||||
});
|
||||
done = true;
|
||||
wake();
|
||||
}
|
||||
})();
|
||||
};
|
||||
|
||||
@@ -850,6 +1064,77 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy
|
||||
}
|
||||
}
|
||||
|
||||
function requestBodyWithoutInput(body: RequestBody): RequestBody {
|
||||
const { input: _input, previous_response_id: _previousResponseId, ...rest } = body;
|
||||
return rest;
|
||||
}
|
||||
|
||||
function responseInputsEqual(a: ResponseInput | undefined, b: ResponseInput | undefined): boolean {
|
||||
return JSON.stringify(a ?? []) === JSON.stringify(b ?? []);
|
||||
}
|
||||
|
||||
function requestBodiesMatchExceptInput(a: RequestBody, b: RequestBody): boolean {
|
||||
return JSON.stringify(requestBodyWithoutInput(a)) === JSON.stringify(requestBodyWithoutInput(b));
|
||||
}
|
||||
|
||||
function getCachedWebSocketInputDelta(
|
||||
body: RequestBody,
|
||||
continuation: CachedWebSocketContinuationState,
|
||||
): ResponseInput | undefined {
|
||||
if (!requestBodiesMatchExceptInput(body, continuation.lastRequestBody)) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const currentInput = body.input ?? [];
|
||||
const baseline = [...(continuation.lastRequestBody.input ?? []), ...continuation.lastResponseItems];
|
||||
if (currentInput.length < baseline.length) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const prefix = currentInput.slice(0, baseline.length);
|
||||
if (!responseInputsEqual(prefix, baseline)) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return currentInput.slice(baseline.length);
|
||||
}
|
||||
|
||||
function buildCachedWebSocketRequestBody(entry: CachedWebSocketConnection, body: RequestBody): RequestBody {
|
||||
const continuation = entry.continuation;
|
||||
if (!continuation) {
|
||||
return body;
|
||||
}
|
||||
|
||||
const delta = getCachedWebSocketInputDelta(body, continuation);
|
||||
if (!delta || !continuation.lastResponseId) {
|
||||
entry.continuation = undefined;
|
||||
return body;
|
||||
}
|
||||
|
||||
return {
|
||||
...body,
|
||||
previous_response_id: continuation.lastResponseId,
|
||||
input: delta,
|
||||
};
|
||||
}
|
||||
|
||||
async function* startWebSocketOutputOnFirstEvent(
|
||||
events: AsyncIterable<ResponseStreamEvent>,
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
onStart: () => void,
|
||||
): AsyncGenerator<ResponseStreamEvent> {
|
||||
let started = false;
|
||||
for await (const event of events) {
|
||||
if (!started) {
|
||||
started = true;
|
||||
onStart();
|
||||
stream.push({ type: "start", partial: output });
|
||||
}
|
||||
yield event;
|
||||
}
|
||||
}
|
||||
|
||||
async function processWebSocketStream(
|
||||
url: string,
|
||||
body: RequestBody,
|
||||
@@ -860,21 +1145,65 @@ async function processWebSocketStream(
|
||||
onStart: () => void,
|
||||
options?: OpenAICodexResponsesOptions,
|
||||
): Promise<void> {
|
||||
const { socket, release } = await acquireWebSocket(url, headers, options?.sessionId, options?.signal);
|
||||
const { socket, entry, reused, release } = await acquireWebSocket(url, headers, options?.sessionId, options?.signal);
|
||||
let keepConnection = true;
|
||||
const useCachedContext = options?.transport === "websocket-cached" || options?.transport === "auto";
|
||||
// ChatGPT Codex Responses rejects `store: true` ("Store must be set to false").
|
||||
// WebSocket continuation still works via connection-scoped previous_response_id state.
|
||||
const fullBody = body;
|
||||
const requestBody = useCachedContext && entry ? buildCachedWebSocketRequestBody(entry, fullBody) : fullBody;
|
||||
const stats = options?.sessionId ? getOrCreateWebSocketDebugStats(options.sessionId) : undefined;
|
||||
if (stats) {
|
||||
stats.requests++;
|
||||
if (reused) stats.connectionsReused++;
|
||||
else stats.connectionsCreated++;
|
||||
if (useCachedContext) stats.cachedContextRequests++;
|
||||
if (requestBody.store === true) stats.storeTrueRequests++;
|
||||
stats.lastInputItems = requestBody.input?.length ?? 0;
|
||||
if (requestBody.previous_response_id) {
|
||||
stats.deltaRequests++;
|
||||
stats.lastDeltaInputItems = requestBody.input?.length ?? 0;
|
||||
stats.lastPreviousResponseId = requestBody.previous_response_id;
|
||||
} else {
|
||||
stats.fullContextRequests++;
|
||||
stats.lastDeltaInputItems = undefined;
|
||||
stats.lastPreviousResponseId = undefined;
|
||||
}
|
||||
}
|
||||
try {
|
||||
socket.send(JSON.stringify({ type: "response.create", ...body }));
|
||||
onStart();
|
||||
stream.push({ type: "start", partial: output });
|
||||
await processResponsesStream(mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, model, {
|
||||
serviceTier: options?.serviceTier,
|
||||
resolveServiceTier: resolveCodexServiceTier,
|
||||
applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model),
|
||||
});
|
||||
socket.send(JSON.stringify({ type: "response.create", ...requestBody }));
|
||||
await processResponsesStream(
|
||||
startWebSocketOutputOnFirstEvent(
|
||||
mapCodexEvents(parseWebSocket(socket, options?.signal)),
|
||||
output,
|
||||
stream,
|
||||
onStart,
|
||||
),
|
||||
output,
|
||||
stream,
|
||||
model,
|
||||
{
|
||||
serviceTier: options?.serviceTier,
|
||||
resolveServiceTier: resolveCodexServiceTier,
|
||||
applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model),
|
||||
},
|
||||
);
|
||||
if (options?.signal?.aborted) {
|
||||
keepConnection = false;
|
||||
} else if (useCachedContext && entry && output.responseId) {
|
||||
const responseItems = convertResponsesMessages(model, { messages: [output] }, CODEX_TOOL_CALL_PROVIDERS, {
|
||||
includeSystemPrompt: false,
|
||||
}).filter((item) => item.type !== "function_call_output");
|
||||
entry.continuation = {
|
||||
lastRequestBody: fullBody,
|
||||
lastResponseId: output.responseId,
|
||||
lastResponseItems: responseItems,
|
||||
};
|
||||
}
|
||||
} catch (error) {
|
||||
if (entry) {
|
||||
entry.continuation = undefined;
|
||||
}
|
||||
keepConnection = false;
|
||||
throw error;
|
||||
} finally {
|
||||
|
||||
@@ -11,7 +11,7 @@ import type {
|
||||
ChatCompletionToolMessageParam,
|
||||
} from "openai/resources/chat/completions.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost, supportsXhigh } from "../models.js";
|
||||
import { calculateCost, clampThinkingLevel } from "../models.js";
|
||||
import type {
|
||||
AssistantMessage,
|
||||
CacheRetention,
|
||||
@@ -36,7 +36,7 @@ import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { isCloudflareProvider, resolveCloudflareBaseUrl } from "./cloudflare.js";
|
||||
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
/**
|
||||
@@ -207,6 +207,9 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
|
||||
// OpenAI documents ChatCompletionChunk.id as the unique chat completion identifier,
|
||||
// and each chunk in a streamed completion carries the same id.
|
||||
output.responseId ||= chunk.id;
|
||||
if (typeof chunk.model === "string" && chunk.model.length > 0 && chunk.model !== model.id) {
|
||||
output.responseModel ||= chunk.model;
|
||||
}
|
||||
if (chunk.usage) {
|
||||
output.usage = parseChunkUsage(chunk.usage, model);
|
||||
}
|
||||
@@ -407,7 +410,8 @@ export const streamSimpleOpenAICompletions: StreamFunction<"openai-completions",
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
|
||||
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
|
||||
const toolChoice = (options as OpenAICompletionsOptions | undefined)?.toolChoice;
|
||||
|
||||
return streamOpenAICompletions(model, context, {
|
||||
@@ -455,11 +459,20 @@ function createClient(
|
||||
Object.assign(headers, optionsHeaders);
|
||||
}
|
||||
|
||||
const defaultHeaders =
|
||||
model.provider === "cloudflare-ai-gateway"
|
||||
? {
|
||||
...headers,
|
||||
Authorization: headers.Authorization ?? null,
|
||||
"cf-aig-authorization": `Bearer ${apiKey}`,
|
||||
}
|
||||
: headers;
|
||||
|
||||
return new OpenAI({
|
||||
apiKey,
|
||||
baseURL: isCloudflareProvider(model.provider) ? resolveCloudflareBaseUrl(model) : model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: headers,
|
||||
defaultHeaders,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -535,21 +548,27 @@ function buildParams(
|
||||
} else if (compat.thinkingFormat === "deepseek" && model.reasoning) {
|
||||
(params as any).thinking = { type: options?.reasoningEffort ? "enabled" : "disabled" };
|
||||
if (options?.reasoningEffort) {
|
||||
(params as any).reasoning_effort = mapReasoningEffort(options.reasoningEffort, compat.reasoningEffortMap);
|
||||
(params as any).reasoning_effort =
|
||||
model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort;
|
||||
}
|
||||
} else if (compat.thinkingFormat === "openrouter" && model.reasoning) {
|
||||
// OpenRouter normalizes reasoning across providers via a nested reasoning object.
|
||||
const openRouterParams = params as typeof params & { reasoning?: { effort?: string } };
|
||||
if (options?.reasoningEffort) {
|
||||
openRouterParams.reasoning = {
|
||||
effort: mapReasoningEffort(options.reasoningEffort, compat.reasoningEffortMap),
|
||||
effort: model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort,
|
||||
};
|
||||
} else {
|
||||
openRouterParams.reasoning = { effort: "none" };
|
||||
} else if (model.thinkingLevelMap?.off !== null) {
|
||||
openRouterParams.reasoning = { effort: model.thinkingLevelMap?.off ?? "none" };
|
||||
}
|
||||
} else if (options?.reasoningEffort && model.reasoning && compat.supportsReasoningEffort) {
|
||||
// OpenAI-style reasoning_effort
|
||||
(params as any).reasoning_effort = mapReasoningEffort(options.reasoningEffort, compat.reasoningEffortMap);
|
||||
(params as any).reasoning_effort = model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort;
|
||||
} else if (!options?.reasoningEffort && model.reasoning && compat.supportsReasoningEffort) {
|
||||
const offValue = model.thinkingLevelMap?.off;
|
||||
if (typeof offValue === "string") {
|
||||
(params as any).reasoning_effort = offValue;
|
||||
}
|
||||
}
|
||||
|
||||
// OpenRouter provider routing preferences
|
||||
@@ -571,13 +590,6 @@ function buildParams(
|
||||
return params;
|
||||
}
|
||||
|
||||
function mapReasoningEffort(
|
||||
effort: NonNullable<OpenAICompletionsOptions["reasoningEffort"]>,
|
||||
reasoningEffortMap: Partial<Record<NonNullable<OpenAICompletionsOptions["reasoningEffort"]>, string>>,
|
||||
): string {
|
||||
return reasoningEffortMap[effort] ?? effort;
|
||||
}
|
||||
|
||||
function getCompatCacheControl(
|
||||
compat: ResolvedOpenAICompletionsCompat,
|
||||
cacheRetention: CacheRetention,
|
||||
@@ -956,12 +968,13 @@ function parseChunkUsage(
|
||||
rawUsage: {
|
||||
prompt_tokens?: number;
|
||||
completion_tokens?: number;
|
||||
prompt_cache_hit_tokens?: number;
|
||||
prompt_tokens_details?: { cached_tokens?: number; cache_write_tokens?: number };
|
||||
},
|
||||
model: Model<"openai-completions">,
|
||||
): AssistantMessage["usage"] {
|
||||
const promptTokens = rawUsage.prompt_tokens || 0;
|
||||
const reportedCachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
|
||||
const reportedCachedTokens = rawUsage.prompt_tokens_details?.cached_tokens ?? rawUsage.prompt_cache_hit_tokens ?? 0;
|
||||
const cacheWriteTokens = rawUsage.prompt_tokens_details?.cache_write_tokens || 0;
|
||||
|
||||
// Normalize to pi-ai semantics:
|
||||
@@ -1023,7 +1036,9 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
const baseUrl = model.baseUrl;
|
||||
|
||||
const isZai = provider === "zai" || baseUrl.includes("api.z.ai");
|
||||
const isMoonshot = provider === "moonshotai" || provider === "moonshotai-cn" || baseUrl.includes("api.moonshot.");
|
||||
const isCloudflareWorkersAI = provider === "cloudflare-workers-ai" || baseUrl.includes("api.cloudflare.com");
|
||||
const isCloudflareAiGateway = provider === "cloudflare-ai-gateway" || baseUrl.includes("gateway.ai.cloudflare.com");
|
||||
|
||||
const isNonStandard =
|
||||
provider === "cerebras" ||
|
||||
@@ -1033,39 +1048,22 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
baseUrl.includes("chutes.ai") ||
|
||||
baseUrl.includes("deepseek.com") ||
|
||||
isZai ||
|
||||
isMoonshot ||
|
||||
provider === "opencode" ||
|
||||
baseUrl.includes("opencode.ai") ||
|
||||
isCloudflareWorkersAI;
|
||||
isCloudflareWorkersAI ||
|
||||
isCloudflareAiGateway;
|
||||
|
||||
const useMaxTokens = baseUrl.includes("chutes.ai");
|
||||
const useMaxTokens = baseUrl.includes("chutes.ai") || isMoonshot || isCloudflareAiGateway;
|
||||
|
||||
const isGrok = provider === "xai" || baseUrl.includes("api.x.ai");
|
||||
const isGroq = provider === "groq" || baseUrl.includes("groq.com");
|
||||
const isDeepSeek = provider === "deepseek" || baseUrl.includes("deepseek.com");
|
||||
const cacheControlFormat = provider === "openrouter" && model.id.startsWith("anthropic/") ? "anthropic" : undefined;
|
||||
|
||||
const reasoningEffortMap = isDeepSeek
|
||||
? {
|
||||
minimal: "high",
|
||||
low: "high",
|
||||
medium: "high",
|
||||
high: "high",
|
||||
xhigh: "max",
|
||||
}
|
||||
: isGroq && model.id === "qwen/qwen3-32b"
|
||||
? {
|
||||
minimal: "default",
|
||||
low: "default",
|
||||
medium: "default",
|
||||
high: "default",
|
||||
xhigh: "default",
|
||||
}
|
||||
: {};
|
||||
return {
|
||||
supportsStore: !isNonStandard,
|
||||
supportsDeveloperRole: !isNonStandard,
|
||||
supportsReasoningEffort: !isGrok && !isZai,
|
||||
reasoningEffortMap,
|
||||
supportsReasoningEffort: !isGrok && !isZai && !isMoonshot && !isCloudflareAiGateway,
|
||||
supportsUsageInStreaming: true,
|
||||
maxTokensField: useMaxTokens ? "max_tokens" : "max_completion_tokens",
|
||||
requiresToolResultName: false,
|
||||
@@ -1082,10 +1080,10 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
openRouterRouting: {},
|
||||
vercelGatewayRouting: {},
|
||||
zaiToolStream: false,
|
||||
supportsStrictMode: true,
|
||||
supportsStrictMode: !isMoonshot && !isCloudflareAiGateway,
|
||||
cacheControlFormat,
|
||||
sendSessionAffinityHeaders: false,
|
||||
supportsLongCacheRetention: true,
|
||||
supportsLongCacheRetention: !(isCloudflareWorkersAI || isCloudflareAiGateway),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -1101,7 +1099,6 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
|
||||
supportsStore: model.compat.supportsStore ?? detected.supportsStore,
|
||||
supportsDeveloperRole: model.compat.supportsDeveloperRole ?? detected.supportsDeveloperRole,
|
||||
supportsReasoningEffort: model.compat.supportsReasoningEffort ?? detected.supportsReasoningEffort,
|
||||
reasoningEffortMap: model.compat.reasoningEffortMap ?? detected.reasoningEffortMap,
|
||||
supportsUsageInStreaming: model.compat.supportsUsageInStreaming ?? detected.supportsUsageInStreaming,
|
||||
maxTokensField: model.compat.maxTokensField ?? detected.maxTokensField,
|
||||
requiresToolResultName: model.compat.requiresToolResultName ?? detected.requiresToolResultName,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import OpenAI from "openai";
|
||||
import type { ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { supportsXhigh } from "../models.js";
|
||||
import { clampThinkingLevel } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
@@ -16,9 +16,10 @@ import type {
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { isCloudflareProvider, resolveCloudflareBaseUrl } from "./cloudflare.js";
|
||||
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
|
||||
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { buildBaseOptions } from "./simple-options.js";
|
||||
|
||||
const OPENAI_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
|
||||
|
||||
@@ -149,7 +150,8 @@ export const streamSimpleOpenAIResponses: StreamFunction<"openai-responses", Sim
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
|
||||
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
|
||||
|
||||
return streamOpenAIResponses(model, context, {
|
||||
...base,
|
||||
@@ -196,11 +198,20 @@ function createClient(
|
||||
Object.assign(headers, optionsHeaders);
|
||||
}
|
||||
|
||||
const defaultHeaders =
|
||||
model.provider === "cloudflare-ai-gateway"
|
||||
? {
|
||||
...headers,
|
||||
Authorization: headers.Authorization ?? null,
|
||||
"cf-aig-authorization": `Bearer ${apiKey}`,
|
||||
}
|
||||
: headers;
|
||||
|
||||
return new OpenAI({
|
||||
apiKey,
|
||||
baseURL: model.baseUrl,
|
||||
baseURL: isCloudflareProvider(model.provider) ? resolveCloudflareBaseUrl(model) : model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: headers,
|
||||
defaultHeaders,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -236,13 +247,18 @@ function buildParams(model: Model<"openai-responses">, context: Context, options
|
||||
|
||||
if (model.reasoning) {
|
||||
if (options?.reasoningEffort || options?.reasoningSummary) {
|
||||
const effort = options?.reasoningEffort
|
||||
? (model.thinkingLevelMap?.[options.reasoningEffort] ?? options.reasoningEffort)
|
||||
: "medium";
|
||||
params.reasoning = {
|
||||
effort: options?.reasoningEffort || "medium",
|
||||
effort: effort as NonNullable<typeof params.reasoning>["effort"],
|
||||
summary: options?.reasoningSummary || "auto",
|
||||
};
|
||||
params.include = ["reasoning.encrypted_content"];
|
||||
} else if (model.provider !== "github-copilot") {
|
||||
params.reasoning = { effort: "none" };
|
||||
} else if (model.provider !== "github-copilot" && model.thinkingLevelMap?.off !== null) {
|
||||
params.reasoning = {
|
||||
effort: (model.thinkingLevelMap?.off ?? "none") as NonNullable<typeof params.reasoning>["effort"],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -14,7 +14,6 @@ import type { BedrockOptions } from "./amazon-bedrock.js";
|
||||
import type { AnthropicOptions } from "./anthropic.js";
|
||||
import type { AzureOpenAIResponsesOptions } from "./azure-openai-responses.js";
|
||||
import type { GoogleOptions } from "./google.js";
|
||||
import type { GoogleGeminiCliOptions } from "./google-gemini-cli.js";
|
||||
import type { GoogleVertexOptions } from "./google-vertex.js";
|
||||
import type { MistralOptions } from "./mistral.js";
|
||||
import type { OpenAICodexResponsesOptions } from "./openai-codex-responses.js";
|
||||
@@ -49,11 +48,6 @@ interface GoogleProviderModule {
|
||||
streamSimpleGoogle: StreamFunction<"google-generative-ai", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface GoogleGeminiCliProviderModule {
|
||||
streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGeminiCliOptions>;
|
||||
streamSimpleGoogleGeminiCli: StreamFunction<"google-gemini-cli", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface GoogleVertexProviderModule {
|
||||
streamGoogleVertex: StreamFunction<"google-vertex", GoogleVertexOptions>;
|
||||
streamSimpleGoogleVertex: StreamFunction<"google-vertex", SimpleStreamOptions>;
|
||||
@@ -103,9 +97,6 @@ let azureOpenAIResponsesProviderModulePromise:
|
||||
let googleProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"google-generative-ai", GoogleOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let googleGeminiCliProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"google-gemini-cli", GoogleGeminiCliOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let googleVertexProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"google-vertex", GoogleVertexOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
@@ -248,19 +239,6 @@ function loadGoogleProviderModule(): Promise<
|
||||
return googleProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadGoogleGeminiCliProviderModule(): Promise<
|
||||
LazyProviderModule<"google-gemini-cli", GoogleGeminiCliOptions, SimpleStreamOptions>
|
||||
> {
|
||||
googleGeminiCliProviderModulePromise ||= import("./google-gemini-cli.js").then((module) => {
|
||||
const provider = module as GoogleGeminiCliProviderModule;
|
||||
return {
|
||||
stream: provider.streamGoogleGeminiCli,
|
||||
streamSimple: provider.streamSimpleGoogleGeminiCli,
|
||||
};
|
||||
});
|
||||
return googleGeminiCliProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadGoogleVertexProviderModule(): Promise<
|
||||
LazyProviderModule<"google-vertex", GoogleVertexOptions, SimpleStreamOptions>
|
||||
> {
|
||||
@@ -348,8 +326,6 @@ export const streamAzureOpenAIResponses = createLazyStream(loadAzureOpenAIRespon
|
||||
export const streamSimpleAzureOpenAIResponses = createLazySimpleStream(loadAzureOpenAIResponsesProviderModule);
|
||||
export const streamGoogle = createLazyStream(loadGoogleProviderModule);
|
||||
export const streamSimpleGoogle = createLazySimpleStream(loadGoogleProviderModule);
|
||||
export const streamGoogleGeminiCli = createLazyStream(loadGoogleGeminiCliProviderModule);
|
||||
export const streamSimpleGoogleGeminiCli = createLazySimpleStream(loadGoogleGeminiCliProviderModule);
|
||||
export const streamGoogleVertex = createLazyStream(loadGoogleVertexProviderModule);
|
||||
export const streamSimpleGoogleVertex = createLazySimpleStream(loadGoogleVertexProviderModule);
|
||||
export const streamMistral = createLazyStream(loadMistralProviderModule);
|
||||
@@ -406,12 +382,6 @@ export function registerBuiltInApiProviders(): void {
|
||||
streamSimple: streamSimpleGoogle,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "google-gemini-cli",
|
||||
stream: streamGoogleGeminiCli,
|
||||
streamSimple: streamSimpleGoogleGeminiCli,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "google-vertex",
|
||||
stream: streamGoogleVertex,
|
||||
|
||||
@@ -6,6 +6,7 @@ export function buildBaseOptions(model: Model<Api>, options?: SimpleStreamOption
|
||||
maxTokens: options?.maxTokens ?? (model.maxTokens > 0 ? Math.min(model.maxTokens, 32000) : undefined),
|
||||
signal: options?.signal,
|
||||
apiKey: apiKey || options?.apiKey,
|
||||
transport: options?.transport,
|
||||
cacheRetention: options?.cacheRetention,
|
||||
sessionId: options?.sessionId,
|
||||
headers: options?.headers,
|
||||
|
||||
24
packages/ai/src/session-resources.ts
Normal file
24
packages/ai/src/session-resources.ts
Normal file
@@ -0,0 +1,24 @@
|
||||
export type SessionResourceCleanup = (sessionId?: string) => void;
|
||||
|
||||
const sessionResourceCleanups = new Set<SessionResourceCleanup>();
|
||||
|
||||
export function registerSessionResourceCleanup(cleanup: SessionResourceCleanup): () => void {
|
||||
sessionResourceCleanups.add(cleanup);
|
||||
return () => {
|
||||
sessionResourceCleanups.delete(cleanup);
|
||||
};
|
||||
}
|
||||
|
||||
export function cleanupSessionResources(sessionId?: string): void {
|
||||
const errors: unknown[] = [];
|
||||
for (const cleanup of sessionResourceCleanups) {
|
||||
try {
|
||||
cleanup(sessionId);
|
||||
} catch (error) {
|
||||
errors.push(error);
|
||||
}
|
||||
}
|
||||
if (errors.length > 0) {
|
||||
throw new AggregateError(errors, "Failed to cleanup session resources");
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
import type { AssistantMessageDiagnostic } from "./utils/diagnostics.js";
|
||||
import type { AssistantImagesEventStream, AssistantMessageEventStream } from "./utils/event-stream.js";
|
||||
|
||||
export type {
|
||||
@@ -14,7 +15,6 @@ export type KnownApi =
|
||||
| "anthropic-messages"
|
||||
| "bedrock-converse-stream"
|
||||
| "google-generative-ai"
|
||||
| "google-gemini-cli"
|
||||
| "google-vertex";
|
||||
|
||||
export type Api = KnownApi | (string & {});
|
||||
@@ -27,8 +27,6 @@ export type KnownProvider =
|
||||
| "amazon-bedrock"
|
||||
| "anthropic"
|
||||
| "google"
|
||||
| "google-gemini-cli"
|
||||
| "google-antigravity"
|
||||
| "google-vertex"
|
||||
| "openai"
|
||||
| "azure-openai-responses"
|
||||
@@ -44,12 +42,19 @@ export type KnownProvider =
|
||||
| "mistral"
|
||||
| "minimax"
|
||||
| "minimax-cn"
|
||||
| "moonshotai"
|
||||
| "moonshotai-cn"
|
||||
| "huggingface"
|
||||
| "fireworks"
|
||||
| "opencode"
|
||||
| "opencode-go"
|
||||
| "kimi-coding"
|
||||
| "cloudflare-workers-ai";
|
||||
| "cloudflare-workers-ai"
|
||||
| "cloudflare-ai-gateway"
|
||||
| "xiaomi"
|
||||
| "xiaomi-token-plan-cn"
|
||||
| "xiaomi-token-plan-ams"
|
||||
| "xiaomi-token-plan-sgp";
|
||||
export type Provider = KnownProvider | string;
|
||||
|
||||
export type KnownImagesProvider = "openrouter";
|
||||
@@ -57,6 +62,8 @@ export type KnownImagesProvider = "openrouter";
|
||||
export type ImagesProvider = KnownImagesProvider | string;
|
||||
|
||||
export type ThinkingLevel = "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
export type ModelThinkingLevel = "off" | ThinkingLevel;
|
||||
export type ThinkingLevelMap = Partial<Record<ModelThinkingLevel, string | null>>;
|
||||
|
||||
/** Token budgets for each thinking level (token-based providers only) */
|
||||
export interface ThinkingBudgets {
|
||||
@@ -69,7 +76,7 @@ export interface ThinkingBudgets {
|
||||
// Base options all providers share
|
||||
export type CacheRetention = "none" | "short" | "long";
|
||||
|
||||
export type Transport = "sse" | "websocket" | "auto";
|
||||
export type Transport = "sse" | "websocket" | "websocket-cached" | "auto";
|
||||
|
||||
export interface ProviderResponse {
|
||||
status: number;
|
||||
@@ -275,7 +282,9 @@ export interface AssistantMessage {
|
||||
api: Api;
|
||||
provider: Provider;
|
||||
model: string;
|
||||
responseModel?: string; // Concrete `chunk.model` when different from the requested `model` (e.g. OpenRouter `auto` -> `anthropic/...`)
|
||||
responseId?: string; // Provider-specific response/message identifier when the upstream API exposes one
|
||||
diagnostics?: AssistantMessageDiagnostic[]; // Redacted provider/runtime diagnostics for failures and recoveries.
|
||||
usage: Usage;
|
||||
stopReason: StopReason;
|
||||
errorMessage?: string;
|
||||
@@ -369,8 +378,6 @@ export interface OpenAICompletionsCompat {
|
||||
supportsDeveloperRole?: boolean;
|
||||
/** Whether the provider supports `reasoning_effort`. Default: auto-detected from URL. */
|
||||
supportsReasoningEffort?: boolean;
|
||||
/** Optional mapping from pi-ai reasoning levels to provider/model-specific `reasoning_effort` values. */
|
||||
reasoningEffortMap?: Partial<Record<ThinkingLevel, string>>;
|
||||
/** Whether the provider supports `stream_options: { include_usage: true }` for token usage in streaming responses. Default: true. */
|
||||
supportsUsageInStreaming?: boolean;
|
||||
/** Which field to use for max tokens. Default: auto-detected from URL. */
|
||||
@@ -518,6 +525,11 @@ export interface Model<TApi extends Api> {
|
||||
provider: Provider;
|
||||
baseUrl: string;
|
||||
reasoning: boolean;
|
||||
/**
|
||||
* Maps pi thinking levels to provider/model-specific values.
|
||||
* Missing keys use provider defaults. null marks a level as unsupported.
|
||||
*/
|
||||
thinkingLevelMap?: ThinkingLevelMap;
|
||||
input: ("text" | "image")[];
|
||||
cost: {
|
||||
input: number; // $/million tokens
|
||||
|
||||
45
packages/ai/src/utils/diagnostics.ts
Normal file
45
packages/ai/src/utils/diagnostics.ts
Normal file
@@ -0,0 +1,45 @@
|
||||
export interface DiagnosticErrorInfo {
|
||||
name?: string;
|
||||
message: string;
|
||||
stack?: string;
|
||||
code?: string | number;
|
||||
}
|
||||
|
||||
export interface AssistantMessageDiagnostic {
|
||||
type: string;
|
||||
timestamp: number;
|
||||
error?: DiagnosticErrorInfo;
|
||||
details?: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export function formatThrownValue(value: unknown): string {
|
||||
if (value instanceof Error) return value.message || value.name;
|
||||
if (typeof value === "string") return value;
|
||||
return String(value);
|
||||
}
|
||||
|
||||
export function extractDiagnosticError(error: unknown): DiagnosticErrorInfo {
|
||||
if (!(error instanceof Error)) return { name: "ThrownValue", message: formatThrownValue(error) };
|
||||
const code = (error as Error & { code?: unknown }).code;
|
||||
return {
|
||||
name: error.name || undefined,
|
||||
message: error.message || error.name,
|
||||
stack: error.stack,
|
||||
code: typeof code === "string" || typeof code === "number" ? code : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
export function createAssistantMessageDiagnostic(
|
||||
type: string,
|
||||
error: unknown,
|
||||
details?: Record<string, unknown>,
|
||||
): AssistantMessageDiagnostic {
|
||||
return { type, timestamp: Date.now(), error: extractDiagnosticError(error), details };
|
||||
}
|
||||
|
||||
export function appendAssistantMessageDiagnostic<T extends { diagnostics?: AssistantMessageDiagnostic[] }>(
|
||||
message: T,
|
||||
diagnostic: AssistantMessageDiagnostic,
|
||||
): void {
|
||||
message.diagnostics = [...(message.diagnostics ?? []), diagnostic];
|
||||
}
|
||||
@@ -1,455 +0,0 @@
|
||||
/**
|
||||
* Antigravity OAuth flow (Gemini 3, Claude, GPT-OSS via Google Cloud)
|
||||
* Uses different OAuth credentials than google-gemini-cli for access to additional models.
|
||||
*
|
||||
* NOTE: This module uses Node.js http.createServer for the OAuth callback.
|
||||
* It is only intended for CLI use, not browser environments.
|
||||
*/
|
||||
|
||||
import type { Server } from "node:http";
|
||||
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.js";
|
||||
import { generatePKCE } from "./pkce.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
type AntigravityCredentials = OAuthCredentials & {
|
||||
projectId: string;
|
||||
};
|
||||
|
||||
const CALLBACK_HOST = process.env.PI_OAUTH_CALLBACK_HOST || "127.0.0.1";
|
||||
|
||||
let _createServer: typeof import("node:http").createServer | null = null;
|
||||
let _httpImportPromise: Promise<void> | null = null;
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
_httpImportPromise = import("node:http").then((m) => {
|
||||
_createServer = m.createServer;
|
||||
});
|
||||
}
|
||||
|
||||
// Antigravity OAuth credentials (different from Gemini CLI)
|
||||
const decode = (s: string) => atob(s);
|
||||
const CLIENT_ID = decode(
|
||||
"MTA3MTAwNjA2MDU5MS10bWhzc2luMmgyMWxjcmUyMzV2dG9sb2poNGc0MDNlcC5hcHBzLmdvb2dsZXVzZXJjb250ZW50LmNvbQ==",
|
||||
);
|
||||
const CLIENT_SECRET = decode("R09DU1BYLUs1OEZXUjQ4NkxkTEoxbUxCOHNYQzR6NnFEQWY=");
|
||||
const REDIRECT_URI = "http://localhost:51121/oauth-callback";
|
||||
|
||||
// Antigravity requires additional scopes
|
||||
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 AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth";
|
||||
const TOKEN_URL = "https://oauth2.googleapis.com/token";
|
||||
|
||||
// Fallback project ID when discovery fails
|
||||
const DEFAULT_PROJECT_ID = "rising-fact-p41fc";
|
||||
|
||||
type CallbackServerInfo = {
|
||||
server: Server;
|
||||
cancelWait: () => void;
|
||||
waitForCode: () => Promise<{ code: string; state: string } | null>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Start a local HTTP server to receive the OAuth callback
|
||||
*/
|
||||
async function getNodeCreateServer(): Promise<typeof import("node:http").createServer> {
|
||||
if (_createServer) return _createServer;
|
||||
if (_httpImportPromise) {
|
||||
await _httpImportPromise;
|
||||
}
|
||||
if (_createServer) return _createServer;
|
||||
throw new Error("Antigravity OAuth is only available in Node.js environments");
|
||||
}
|
||||
|
||||
async function startCallbackServer(): Promise<CallbackServerInfo> {
|
||||
const createServer = await getNodeCreateServer();
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
let settleWait: ((value: { code: string; state: string } | null) => void) | undefined;
|
||||
const waitForCodePromise = new Promise<{ code: string; state: string } | null>((resolveWait) => {
|
||||
let settled = false;
|
||||
settleWait = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
resolveWait(value);
|
||||
};
|
||||
});
|
||||
|
||||
const server = createServer((req, res) => {
|
||||
const url = new URL(req.url || "", `http://localhost:51121`);
|
||||
|
||||
if (url.pathname === "/oauth-callback") {
|
||||
const code = url.searchParams.get("code");
|
||||
const state = url.searchParams.get("state");
|
||||
const error = url.searchParams.get("error");
|
||||
|
||||
if (error) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Google authentication did not complete.", `Error: ${error}`));
|
||||
return;
|
||||
}
|
||||
|
||||
if (code && state) {
|
||||
res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthSuccessHtml("Google authentication completed. You can close this window."));
|
||||
settleWait?.({ code, state });
|
||||
} else {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Missing code or state parameter."));
|
||||
}
|
||||
} else {
|
||||
res.writeHead(404, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Callback route not found."));
|
||||
}
|
||||
});
|
||||
|
||||
server.on("error", (err) => {
|
||||
reject(err);
|
||||
});
|
||||
|
||||
server.listen(51121, CALLBACK_HOST, () => {
|
||||
resolve({
|
||||
server,
|
||||
cancelWait: () => {
|
||||
settleWait?.(null);
|
||||
},
|
||||
waitForCode: () => waitForCodePromise,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse redirect URL to extract code and state
|
||||
*/
|
||||
function parseRedirectUrl(input: string): { code?: string; state?: string } {
|
||||
const value = input.trim();
|
||||
if (!value) return {};
|
||||
|
||||
try {
|
||||
const url = new URL(value);
|
||||
return {
|
||||
code: url.searchParams.get("code") ?? undefined,
|
||||
state: url.searchParams.get("state") ?? undefined,
|
||||
};
|
||||
} catch {
|
||||
// Not a URL, return empty
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
interface LoadCodeAssistPayload {
|
||||
cloudaicompanionProject?: string | { id?: string };
|
||||
currentTier?: { id?: string };
|
||||
allowedTiers?: Array<{ id?: string; isDefault?: boolean }>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Discover or provision a project for the user
|
||||
*/
|
||||
async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise<string> {
|
||||
const headers = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "google-api-nodejs-client/9.15.1",
|
||||
"X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
"Client-Metadata": JSON.stringify({
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
}),
|
||||
};
|
||||
|
||||
// Try endpoints in order: prod first, then sandbox
|
||||
const endpoints = ["https://cloudcode-pa.googleapis.com", "https://daily-cloudcode-pa.sandbox.googleapis.com"];
|
||||
|
||||
onProgress?.("Checking for existing project...");
|
||||
|
||||
for (const endpoint of endpoints) {
|
||||
try {
|
||||
const loadResponse = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({
|
||||
metadata: {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
},
|
||||
}),
|
||||
});
|
||||
|
||||
if (loadResponse.ok) {
|
||||
const data = (await loadResponse.json()) as LoadCodeAssistPayload;
|
||||
|
||||
// Handle both string and object formats
|
||||
if (typeof data.cloudaicompanionProject === "string" && data.cloudaicompanionProject) {
|
||||
return data.cloudaicompanionProject;
|
||||
}
|
||||
if (
|
||||
data.cloudaicompanionProject &&
|
||||
typeof data.cloudaicompanionProject === "object" &&
|
||||
data.cloudaicompanionProject.id
|
||||
) {
|
||||
return data.cloudaicompanionProject.id;
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Try next endpoint
|
||||
}
|
||||
}
|
||||
|
||||
// Use fallback project ID
|
||||
onProgress?.("Using default project...");
|
||||
return DEFAULT_PROJECT_ID;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user email from the access token
|
||||
*/
|
||||
async function getUserEmail(accessToken: string): Promise<string | undefined> {
|
||||
try {
|
||||
const response = await fetch("https://www.googleapis.com/oauth2/v1/userinfo?alt=json", {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
},
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const data = (await response.json()) as { email?: string };
|
||||
return data.email;
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors, email is optional
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh Antigravity token
|
||||
*/
|
||||
export async function refreshAntigravityToken(refreshToken: string, projectId: string): Promise<OAuthCredentials> {
|
||||
const response = await fetch(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: refreshToken,
|
||||
grant_type: "refresh_token",
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.text();
|
||||
throw new Error(`Antigravity token refresh failed: ${error}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as {
|
||||
access_token: string;
|
||||
expires_in: number;
|
||||
refresh_token?: string;
|
||||
};
|
||||
|
||||
return {
|
||||
refresh: data.refresh_token || refreshToken,
|
||||
access: data.access_token,
|
||||
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
|
||||
projectId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with Antigravity OAuth
|
||||
*
|
||||
* @param onAuth - Callback with URL and optional instructions
|
||||
* @param onProgress - Optional progress callback
|
||||
* @param onManualCodeInput - Optional promise that resolves with user-pasted redirect URL.
|
||||
* Races with browser callback - whichever completes first wins.
|
||||
*/
|
||||
export async function loginAntigravity(
|
||||
onAuth: (info: { url: string; instructions?: string }) => void,
|
||||
onProgress?: (message: string) => void,
|
||||
onManualCodeInput?: () => Promise<string>,
|
||||
): Promise<OAuthCredentials> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
|
||||
// Start local server for callback
|
||||
onProgress?.("Starting local server for OAuth callback...");
|
||||
const server = await startCallbackServer();
|
||||
|
||||
let code: string | undefined;
|
||||
|
||||
try {
|
||||
// Build authorization URL
|
||||
const authParams = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
response_type: "code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
scope: SCOPES.join(" "),
|
||||
code_challenge: challenge,
|
||||
code_challenge_method: "S256",
|
||||
state: verifier,
|
||||
access_type: "offline",
|
||||
prompt: "consent",
|
||||
});
|
||||
|
||||
const authUrl = `${AUTH_URL}?${authParams.toString()}`;
|
||||
|
||||
// Notify caller with URL to open
|
||||
onAuth({
|
||||
url: authUrl,
|
||||
instructions: "Complete the sign-in in your browser.",
|
||||
});
|
||||
|
||||
// Wait for the callback, racing with manual input if provided
|
||||
onProgress?.("Waiting for OAuth callback...");
|
||||
|
||||
if (onManualCodeInput) {
|
||||
// Race between browser callback and manual input
|
||||
let manualInput: string | undefined;
|
||||
let manualError: Error | undefined;
|
||||
const manualPromise = onManualCodeInput()
|
||||
.then((input) => {
|
||||
manualInput = input;
|
||||
server.cancelWait();
|
||||
})
|
||||
.catch((err) => {
|
||||
manualError = err instanceof Error ? err : new Error(String(err));
|
||||
server.cancelWait();
|
||||
});
|
||||
|
||||
const result = await server.waitForCode();
|
||||
|
||||
// If manual input was cancelled, throw that error
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
|
||||
if (result?.code) {
|
||||
// Browser callback won - verify state
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
} else if (manualInput) {
|
||||
// Manual input won
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
|
||||
// If still no code, wait for manual promise and try that
|
||||
if (!code) {
|
||||
await manualPromise;
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
if (manualInput) {
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Original flow: just wait for callback
|
||||
const result = await server.waitForCode();
|
||||
if (result?.code) {
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
}
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
throw new Error("No authorization code received");
|
||||
}
|
||||
|
||||
// Exchange code for tokens
|
||||
onProgress?.("Exchanging authorization code for tokens...");
|
||||
const tokenResponse = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
},
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
code,
|
||||
grant_type: "authorization_code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
code_verifier: verifier,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!tokenResponse.ok) {
|
||||
const error = await tokenResponse.text();
|
||||
throw new Error(`Token exchange failed: ${error}`);
|
||||
}
|
||||
|
||||
const tokenData = (await tokenResponse.json()) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
};
|
||||
|
||||
if (!tokenData.refresh_token) {
|
||||
throw new Error("No refresh token received. Please try again.");
|
||||
}
|
||||
|
||||
// Get user email
|
||||
onProgress?.("Getting user info...");
|
||||
const email = await getUserEmail(tokenData.access_token);
|
||||
|
||||
// Discover project
|
||||
const projectId = await discoverProject(tokenData.access_token, onProgress);
|
||||
|
||||
// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
|
||||
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
|
||||
|
||||
const credentials: OAuthCredentials = {
|
||||
refresh: tokenData.refresh_token,
|
||||
access: tokenData.access_token,
|
||||
expires: expiresAt,
|
||||
projectId,
|
||||
email,
|
||||
};
|
||||
|
||||
return credentials;
|
||||
} finally {
|
||||
server.server.close();
|
||||
}
|
||||
}
|
||||
|
||||
export const antigravityOAuthProvider: OAuthProviderInterface = {
|
||||
id: "google-antigravity",
|
||||
name: "Antigravity (Gemini 3, Claude, GPT-OSS)",
|
||||
usesCallbackServer: true,
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginAntigravity(callbacks.onAuth, callbacks.onProgress, callbacks.onManualCodeInput);
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
const creds = credentials as AntigravityCredentials;
|
||||
if (!creds.projectId) {
|
||||
throw new Error("Antigravity credentials missing projectId");
|
||||
}
|
||||
return refreshAntigravityToken(creds.refresh, creds.projectId);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
const creds = credentials as AntigravityCredentials;
|
||||
return JSON.stringify({ token: creds.access, projectId: creds.projectId });
|
||||
},
|
||||
};
|
||||
@@ -1,597 +0,0 @@
|
||||
/**
|
||||
* Gemini CLI OAuth flow (Google Cloud Code Assist)
|
||||
* Standard Gemini models only (gemini-2.0-flash, gemini-2.5-*)
|
||||
*
|
||||
* NOTE: This module uses Node.js http.createServer for the OAuth callback.
|
||||
* It is only intended for CLI use, not browser environments.
|
||||
*/
|
||||
|
||||
import type { Server } from "node:http";
|
||||
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.js";
|
||||
import { generatePKCE } from "./pkce.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
type GeminiCredentials = OAuthCredentials & {
|
||||
projectId: string;
|
||||
};
|
||||
|
||||
const CALLBACK_HOST = process.env.PI_OAUTH_CALLBACK_HOST || "127.0.0.1";
|
||||
|
||||
let _createServer: typeof import("node:http").createServer | null = null;
|
||||
let _httpImportPromise: Promise<void> | null = null;
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
_httpImportPromise = import("node:http").then((m) => {
|
||||
_createServer = m.createServer;
|
||||
});
|
||||
}
|
||||
|
||||
const decode = (s: string) => atob(s);
|
||||
const CLIENT_ID = decode(
|
||||
"NjgxMjU1ODA5Mzk1LW9vOGZ0Mm9wcmRybnA5ZTNhcWY2YXYzaG1kaWIxMzVqLmFwcHMuZ29vZ2xldXNlcmNvbnRlbnQuY29t",
|
||||
);
|
||||
const CLIENT_SECRET = decode("R09DU1BYLTR1SGdNUG0tMW83U2stZ2VWNkN1NWNsWEZzeGw=");
|
||||
const REDIRECT_URI = "http://localhost:8085/oauth2callback";
|
||||
const SCOPES = [
|
||||
"https://www.googleapis.com/auth/cloud-platform",
|
||||
"https://www.googleapis.com/auth/userinfo.email",
|
||||
"https://www.googleapis.com/auth/userinfo.profile",
|
||||
];
|
||||
const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth";
|
||||
const TOKEN_URL = "https://oauth2.googleapis.com/token";
|
||||
const CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com";
|
||||
|
||||
type CallbackServerInfo = {
|
||||
server: Server;
|
||||
cancelWait: () => void;
|
||||
waitForCode: () => Promise<{ code: string; state: string } | null>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Start a local HTTP server to receive the OAuth callback
|
||||
*/
|
||||
async function getNodeCreateServer(): Promise<typeof import("node:http").createServer> {
|
||||
if (_createServer) return _createServer;
|
||||
if (_httpImportPromise) {
|
||||
await _httpImportPromise;
|
||||
}
|
||||
if (_createServer) return _createServer;
|
||||
throw new Error("Gemini CLI OAuth is only available in Node.js environments");
|
||||
}
|
||||
|
||||
async function startCallbackServer(): Promise<CallbackServerInfo> {
|
||||
const createServer = await getNodeCreateServer();
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
let settleWait: ((value: { code: string; state: string } | null) => void) | undefined;
|
||||
const waitForCodePromise = new Promise<{ code: string; state: string } | null>((resolveWait) => {
|
||||
let settled = false;
|
||||
settleWait = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
resolveWait(value);
|
||||
};
|
||||
});
|
||||
|
||||
const server = createServer((req, res) => {
|
||||
const url = new URL(req.url || "", `http://localhost:8085`);
|
||||
|
||||
if (url.pathname === "/oauth2callback") {
|
||||
const code = url.searchParams.get("code");
|
||||
const state = url.searchParams.get("state");
|
||||
const error = url.searchParams.get("error");
|
||||
|
||||
if (error) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Google authentication did not complete.", `Error: ${error}`));
|
||||
return;
|
||||
}
|
||||
|
||||
if (code && state) {
|
||||
res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthSuccessHtml("Google authentication completed. You can close this window."));
|
||||
settleWait?.({ code, state });
|
||||
} else {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Missing code or state parameter."));
|
||||
}
|
||||
} else {
|
||||
res.writeHead(404, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Callback route not found."));
|
||||
}
|
||||
});
|
||||
|
||||
server.on("error", (err) => {
|
||||
reject(err);
|
||||
});
|
||||
|
||||
server.listen(8085, CALLBACK_HOST, () => {
|
||||
resolve({
|
||||
server,
|
||||
cancelWait: () => {
|
||||
settleWait?.(null);
|
||||
},
|
||||
waitForCode: () => waitForCodePromise,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse redirect URL to extract code and state
|
||||
*/
|
||||
function parseRedirectUrl(input: string): { code?: string; state?: string } {
|
||||
const value = input.trim();
|
||||
if (!value) return {};
|
||||
|
||||
try {
|
||||
const url = new URL(value);
|
||||
return {
|
||||
code: url.searchParams.get("code") ?? undefined,
|
||||
state: url.searchParams.get("state") ?? undefined,
|
||||
};
|
||||
} catch {
|
||||
// Not a URL, return empty
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
interface LoadCodeAssistPayload {
|
||||
cloudaicompanionProject?: string;
|
||||
currentTier?: { id?: string };
|
||||
allowedTiers?: Array<{ id?: string; isDefault?: boolean }>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Long-running operation response from onboardUser
|
||||
*/
|
||||
interface LongRunningOperationResponse {
|
||||
name?: string;
|
||||
done?: boolean;
|
||||
response?: {
|
||||
cloudaicompanionProject?: { id?: string };
|
||||
};
|
||||
}
|
||||
|
||||
// Tier IDs as used by the Cloud Code API
|
||||
const TIER_FREE = "free-tier";
|
||||
const TIER_LEGACY = "legacy-tier";
|
||||
const TIER_STANDARD = "standard-tier";
|
||||
|
||||
interface GoogleRpcErrorResponse {
|
||||
error?: {
|
||||
details?: Array<{ reason?: string }>;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait helper for onboarding retries
|
||||
*/
|
||||
function wait(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
/**
|
||||
* Get default tier from allowed tiers
|
||||
*/
|
||||
function getDefaultTier(allowedTiers?: Array<{ id?: string; isDefault?: boolean }>): { id?: string } {
|
||||
if (!allowedTiers || allowedTiers.length === 0) return { id: TIER_LEGACY };
|
||||
const defaultTier = allowedTiers.find((t) => t.isDefault);
|
||||
return defaultTier ?? { id: TIER_LEGACY };
|
||||
}
|
||||
|
||||
function isVpcScAffectedUser(payload: unknown): boolean {
|
||||
if (!payload || typeof payload !== "object") return false;
|
||||
if (!("error" in payload)) return false;
|
||||
const error = (payload as GoogleRpcErrorResponse).error;
|
||||
if (!error?.details || !Array.isArray(error.details)) return false;
|
||||
return error.details.some((detail) => detail.reason === "SECURITY_POLICY_VIOLATED");
|
||||
}
|
||||
|
||||
/**
|
||||
* Poll a long-running operation until completion
|
||||
*/
|
||||
async function pollOperation(
|
||||
operationName: string,
|
||||
headers: Record<string, string>,
|
||||
onProgress?: (message: string) => void,
|
||||
): Promise<LongRunningOperationResponse> {
|
||||
let attempt = 0;
|
||||
while (true) {
|
||||
if (attempt > 0) {
|
||||
onProgress?.(`Waiting for project provisioning (attempt ${attempt + 1})...`);
|
||||
await wait(5000);
|
||||
}
|
||||
|
||||
const response = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal/${operationName}`, {
|
||||
method: "GET",
|
||||
headers,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to poll operation: ${response.status} ${response.statusText}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as LongRunningOperationResponse;
|
||||
if (data.done) {
|
||||
return data;
|
||||
}
|
||||
|
||||
attempt += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Discover or provision a Google Cloud project for the user
|
||||
*/
|
||||
async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise<string> {
|
||||
// Check for user-provided project ID via environment variable
|
||||
const envProjectId = process.env.GOOGLE_CLOUD_PROJECT || process.env.GOOGLE_CLOUD_PROJECT_ID;
|
||||
|
||||
const headers = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "google-api-nodejs-client/9.15.1",
|
||||
"X-Goog-Api-Client": "gl-node/22.17.0",
|
||||
};
|
||||
|
||||
// Try to load existing project via loadCodeAssist
|
||||
onProgress?.("Checking for existing Cloud Code Assist project...");
|
||||
const loadResponse = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({
|
||||
cloudaicompanionProject: envProjectId,
|
||||
metadata: {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
duetProject: envProjectId,
|
||||
},
|
||||
}),
|
||||
});
|
||||
|
||||
let data: LoadCodeAssistPayload;
|
||||
|
||||
if (!loadResponse.ok) {
|
||||
let errorPayload: unknown;
|
||||
try {
|
||||
errorPayload = await loadResponse.clone().json();
|
||||
} catch {
|
||||
errorPayload = undefined;
|
||||
}
|
||||
|
||||
if (isVpcScAffectedUser(errorPayload)) {
|
||||
data = { currentTier: { id: TIER_STANDARD } };
|
||||
} else {
|
||||
const errorText = await loadResponse.text();
|
||||
throw new Error(`loadCodeAssist failed: ${loadResponse.status} ${loadResponse.statusText}: ${errorText}`);
|
||||
}
|
||||
} else {
|
||||
data = (await loadResponse.json()) as LoadCodeAssistPayload;
|
||||
}
|
||||
|
||||
// If user already has a current tier and project, use it
|
||||
if (data.currentTier) {
|
||||
if (data.cloudaicompanionProject) {
|
||||
return data.cloudaicompanionProject;
|
||||
}
|
||||
// User has a tier but no managed project - they need to provide one via env var
|
||||
if (envProjectId) {
|
||||
return envProjectId;
|
||||
}
|
||||
throw new Error(
|
||||
"This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||
);
|
||||
}
|
||||
|
||||
// User needs to be onboarded - get the default tier
|
||||
const tier = getDefaultTier(data.allowedTiers);
|
||||
const tierId = tier?.id ?? TIER_FREE;
|
||||
|
||||
if (tierId !== TIER_FREE && !envProjectId) {
|
||||
throw new Error(
|
||||
"This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||
);
|
||||
}
|
||||
|
||||
onProgress?.("Provisioning Cloud Code Assist project (this may take a moment)...");
|
||||
|
||||
// Build onboard request - for free tier, don't include project ID (Google provisions one)
|
||||
// For other tiers, include the user's project ID if available
|
||||
const onboardBody: Record<string, unknown> = {
|
||||
tierId,
|
||||
metadata: {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
},
|
||||
};
|
||||
|
||||
if (tierId !== TIER_FREE && envProjectId) {
|
||||
onboardBody.cloudaicompanionProject = envProjectId;
|
||||
(onboardBody.metadata as Record<string, unknown>).duetProject = envProjectId;
|
||||
}
|
||||
|
||||
// Start onboarding - this returns a long-running operation
|
||||
const onboardResponse = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal:onboardUser`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify(onboardBody),
|
||||
});
|
||||
|
||||
if (!onboardResponse.ok) {
|
||||
const errorText = await onboardResponse.text();
|
||||
throw new Error(`onboardUser failed: ${onboardResponse.status} ${onboardResponse.statusText}: ${errorText}`);
|
||||
}
|
||||
|
||||
let lroData = (await onboardResponse.json()) as LongRunningOperationResponse;
|
||||
|
||||
// If the operation isn't done yet, poll until completion
|
||||
if (!lroData.done && lroData.name) {
|
||||
lroData = await pollOperation(lroData.name, headers, onProgress);
|
||||
}
|
||||
|
||||
// Try to get project ID from the response
|
||||
const projectId = lroData.response?.cloudaicompanionProject?.id;
|
||||
if (projectId) {
|
||||
return projectId;
|
||||
}
|
||||
|
||||
// If no project ID from onboarding, fall back to env var
|
||||
if (envProjectId) {
|
||||
return envProjectId;
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
"Could not discover or provision a Google Cloud project. " +
|
||||
"Try setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user email from the access token
|
||||
*/
|
||||
async function getUserEmail(accessToken: string): Promise<string | undefined> {
|
||||
try {
|
||||
const response = await fetch("https://www.googleapis.com/oauth2/v1/userinfo?alt=json", {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
},
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const data = (await response.json()) as { email?: string };
|
||||
return data.email;
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors, email is optional
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh Google Cloud Code Assist token
|
||||
*/
|
||||
export async function refreshGoogleCloudToken(refreshToken: string, projectId: string): Promise<OAuthCredentials> {
|
||||
const response = await fetch(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: refreshToken,
|
||||
grant_type: "refresh_token",
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.text();
|
||||
throw new Error(`Google Cloud token refresh failed: ${error}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as {
|
||||
access_token: string;
|
||||
expires_in: number;
|
||||
refresh_token?: string;
|
||||
};
|
||||
|
||||
return {
|
||||
refresh: data.refresh_token || refreshToken,
|
||||
access: data.access_token,
|
||||
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
|
||||
projectId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with Gemini CLI (Google Cloud Code Assist) OAuth
|
||||
*
|
||||
* @param onAuth - Callback with URL and optional instructions
|
||||
* @param onProgress - Optional progress callback
|
||||
* @param onManualCodeInput - Optional promise that resolves with user-pasted redirect URL.
|
||||
* Races with browser callback - whichever completes first wins.
|
||||
*/
|
||||
export async function loginGeminiCli(
|
||||
onAuth: (info: { url: string; instructions?: string }) => void,
|
||||
onProgress?: (message: string) => void,
|
||||
onManualCodeInput?: () => Promise<string>,
|
||||
): Promise<OAuthCredentials> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
|
||||
// Start local server for callback
|
||||
onProgress?.("Starting local server for OAuth callback...");
|
||||
const server = await startCallbackServer();
|
||||
|
||||
let code: string | undefined;
|
||||
|
||||
try {
|
||||
// Build authorization URL
|
||||
const authParams = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
response_type: "code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
scope: SCOPES.join(" "),
|
||||
code_challenge: challenge,
|
||||
code_challenge_method: "S256",
|
||||
state: verifier,
|
||||
access_type: "offline",
|
||||
prompt: "consent",
|
||||
});
|
||||
|
||||
const authUrl = `${AUTH_URL}?${authParams.toString()}`;
|
||||
|
||||
// Notify caller with URL to open
|
||||
onAuth({
|
||||
url: authUrl,
|
||||
instructions: "Complete the sign-in in your browser.",
|
||||
});
|
||||
|
||||
// Wait for the callback, racing with manual input if provided
|
||||
onProgress?.("Waiting for OAuth callback...");
|
||||
|
||||
if (onManualCodeInput) {
|
||||
// Race between browser callback and manual input
|
||||
let manualInput: string | undefined;
|
||||
let manualError: Error | undefined;
|
||||
const manualPromise = onManualCodeInput()
|
||||
.then((input) => {
|
||||
manualInput = input;
|
||||
server.cancelWait();
|
||||
})
|
||||
.catch((err) => {
|
||||
manualError = err instanceof Error ? err : new Error(String(err));
|
||||
server.cancelWait();
|
||||
});
|
||||
|
||||
const result = await server.waitForCode();
|
||||
|
||||
// If manual input was cancelled, throw that error
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
|
||||
if (result?.code) {
|
||||
// Browser callback won - verify state
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
} else if (manualInput) {
|
||||
// Manual input won
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
|
||||
// If still no code, wait for manual promise and try that
|
||||
if (!code) {
|
||||
await manualPromise;
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
if (manualInput) {
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Original flow: just wait for callback
|
||||
const result = await server.waitForCode();
|
||||
if (result?.code) {
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
}
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
throw new Error("No authorization code received");
|
||||
}
|
||||
|
||||
// Exchange code for tokens
|
||||
onProgress?.("Exchanging authorization code for tokens...");
|
||||
const tokenResponse = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
},
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
code,
|
||||
grant_type: "authorization_code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
code_verifier: verifier,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!tokenResponse.ok) {
|
||||
const error = await tokenResponse.text();
|
||||
throw new Error(`Token exchange failed: ${error}`);
|
||||
}
|
||||
|
||||
const tokenData = (await tokenResponse.json()) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
};
|
||||
|
||||
if (!tokenData.refresh_token) {
|
||||
throw new Error("No refresh token received. Please try again.");
|
||||
}
|
||||
|
||||
// Get user email
|
||||
onProgress?.("Getting user info...");
|
||||
const email = await getUserEmail(tokenData.access_token);
|
||||
|
||||
// Discover project
|
||||
const projectId = await discoverProject(tokenData.access_token, onProgress);
|
||||
|
||||
// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
|
||||
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
|
||||
|
||||
const credentials: OAuthCredentials = {
|
||||
refresh: tokenData.refresh_token,
|
||||
access: tokenData.access_token,
|
||||
expires: expiresAt,
|
||||
projectId,
|
||||
email,
|
||||
};
|
||||
|
||||
return credentials;
|
||||
} finally {
|
||||
server.server.close();
|
||||
}
|
||||
}
|
||||
|
||||
export const geminiCliOAuthProvider: OAuthProviderInterface = {
|
||||
id: "google-gemini-cli",
|
||||
name: "Google Cloud Code Assist (Gemini CLI)",
|
||||
usesCallbackServer: true,
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginGeminiCli(callbacks.onAuth, callbacks.onProgress, callbacks.onManualCodeInput);
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
const creds = credentials as GeminiCredentials;
|
||||
if (!creds.projectId) {
|
||||
throw new Error("Google Cloud credentials missing projectId");
|
||||
}
|
||||
return refreshGoogleCloudToken(creds.refresh, creds.projectId);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
const creds = credentials as GeminiCredentials;
|
||||
return JSON.stringify({ token: creds.access, projectId: creds.projectId });
|
||||
},
|
||||
};
|
||||
@@ -5,8 +5,6 @@
|
||||
* for OAuth-based providers:
|
||||
* - Anthropic (Claude Pro/Max)
|
||||
* - GitHub Copilot
|
||||
* - Google Cloud Code Assist (Gemini CLI)
|
||||
* - Antigravity (Gemini 3, Claude, GPT-OSS via Google Cloud)
|
||||
*/
|
||||
|
||||
// Anthropic
|
||||
@@ -19,10 +17,6 @@ export {
|
||||
normalizeDomain,
|
||||
refreshGitHubCopilotToken,
|
||||
} from "./github-copilot.js";
|
||||
// Google Antigravity
|
||||
export { antigravityOAuthProvider, loginAntigravity, refreshAntigravityToken } from "./google-antigravity.js";
|
||||
// Google Gemini CLI
|
||||
export { geminiCliOAuthProvider, loginGeminiCli, refreshGoogleCloudToken } from "./google-gemini-cli.js";
|
||||
// OpenAI Codex (ChatGPT OAuth)
|
||||
export { loginOpenAICodex, openaiCodexOAuthProvider, refreshOpenAICodexToken } from "./openai-codex.js";
|
||||
|
||||
@@ -34,16 +28,12 @@ export * from "./types.js";
|
||||
|
||||
import { anthropicOAuthProvider } from "./anthropic.js";
|
||||
import { githubCopilotOAuthProvider } from "./github-copilot.js";
|
||||
import { antigravityOAuthProvider } from "./google-antigravity.js";
|
||||
import { geminiCliOAuthProvider } from "./google-gemini-cli.js";
|
||||
import { openaiCodexOAuthProvider } from "./openai-codex.js";
|
||||
import type { OAuthCredentials, OAuthProviderId, OAuthProviderInfo, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
const BUILT_IN_OAUTH_PROVIDERS: OAuthProviderInterface[] = [
|
||||
anthropicOAuthProvider,
|
||||
githubCopilotOAuthProvider,
|
||||
geminiCliOAuthProvider,
|
||||
antigravityOAuthProvider,
|
||||
openaiCodexOAuthProvider,
|
||||
];
|
||||
|
||||
|
||||
@@ -23,6 +23,9 @@ import type { AssistantMessage } from "../types.js";
|
||||
* - Cerebras: "400/413 status code (no body)"
|
||||
* - Mistral: "Prompt contains X tokens ... too large for model with Y maximum context length"
|
||||
* - z.ai: Does NOT error, accepts overflow silently - handled via usage.input > contextWindow
|
||||
* - Xiaomi MiMo: Truncates input to fill contextWindow exactly, then returns finish_reason "length"
|
||||
* with output=0 (no room left to generate). Detected via stopReason "length" + zero output +
|
||||
* input filling the context window.
|
||||
* - Ollama: Some deployments truncate silently, others return errors like "prompt too long; exceeded max context length by X tokens"
|
||||
*/
|
||||
const OVERFLOW_PATTERNS = [
|
||||
@@ -90,6 +93,8 @@ const NON_OVERFLOW_PATTERNS = [
|
||||
* **Unreliable detection:**
|
||||
* - z.ai: Sometimes accepts overflow silently (detectable via usage.input > contextWindow),
|
||||
* sometimes returns rate limit errors. Pass contextWindow param to detect silent overflow.
|
||||
* - Xiaomi MiMo: Truncates input to fit contextWindow then returns stopReason "length" with
|
||||
* output=0. Pass contextWindow param to detect via the "filled context + zero output" signal.
|
||||
* - Ollama: May truncate input silently for some setups, but may also return explicit
|
||||
* overflow errors that match the patterns above. Silent truncation still cannot be
|
||||
* detected here because we do not know the expected token count.
|
||||
@@ -127,6 +132,16 @@ export function isContextOverflow(message: AssistantMessage, contextWindow?: num
|
||||
}
|
||||
}
|
||||
|
||||
// Case 3: Length-stop overflow (Xiaomi MiMo style) - server truncates oversized input
|
||||
// to fit the context window, leaving no room for output. Returns stopReason "length"
|
||||
// with output=0 and input+cacheRead filling the context window.
|
||||
if (contextWindow && message.stopReason === "length" && message.usage.output === 0) {
|
||||
const inputTokens = message.usage.input + message.usage.cacheRead;
|
||||
if (inputTokens >= contextWindow * 0.99) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user