@tt-a1i/openpi 0.5.0 → 0.6.0

package/extensions/ai-providers/antigravity/routing.ts

@@ -0,0 +1,340 @@
+import type { Model, ThinkingLevel } from "@earendil-works/pi-ai/compat";
+
+export type AntigravityModelDefinition = Omit<
+  Model<string>,
+  "api" | "provider" | "baseUrl"
+> & {
+  /** Provider-private fields preserved by pi's extension model composer. */
+  requestModelId?: string;
+  antigravityEffortRouting?: Partial<Record<ThinkingLevel, string>>;
+};
+
+type Family = {
+  id: string;
+  name: string;
+  members: readonly string[];
+  defaultWireId: string;
+  routes?: Partial<Record<ThinkingLevel, string>>;
+  mode: "budget" | "google-level";
+  budgets?: Partial<Record<ThinkingLevel, number>>;
+  mandatory?: boolean;
+  retiredMembers?: readonly string[];
+  preserveAbsentEffortRoutes?: boolean;
+};
+
+function thinkingPairs(
+  pairs: readonly (readonly [id: string, name: string])[],
+): Family[] {
+  return pairs.map(([id, name]) => ({
+    id,
+    name,
+    members: [id, `${id}-thinking`],
+    defaultWireId: id,
+    routes: {
+      minimal: `${id}-thinking`,
+      low: `${id}-thinking`,
+      medium: `${id}-thinking`,
+      high: `${id}-thinking`,
+      xhigh: `${id}-thinking`,
+      max: `${id}-thinking`,
+    },
+    mode: "budget",
+    preserveAbsentEffortRoutes: true,
+  }));
+}
+
+const FAMILIES: readonly Family[] = [
+  {
+    id: "gemini-3.7-flash",
+    name: "Gemini 3.7 Flash",
+    members: [
+      "gemini-3.7-flash-low",
+      "gemini-3.7-flash-medium",
+      "gemini-3.7-flash-high",
+    ],
+    defaultWireId: "gemini-3.7-flash-low",
+    routes: {
+      minimal: "gemini-3.7-flash-low",
+      low: "gemini-3.7-flash-low",
+      medium: "gemini-3.7-flash-medium",
+      high: "gemini-3.7-flash-high",
+      xhigh: "gemini-3.7-flash-high",
+      max: "gemini-3.7-flash-high",
+    },
+    mode: "google-level",
+    mandatory: true,
+  },
+  {
+    id: "gemini-3.6-flash",
+    name: "Gemini 3.6 Flash",
+    members: [
+      "gemini-3.6-flash-low",
+      "gemini-3.6-flash-medium",
+      "gemini-3.6-flash-high",
+      "gemini-3.6-flash-tiered",
+    ],
+    defaultWireId: "gemini-3.6-flash-low",
+    routes: {
+      minimal: "gemini-3.6-flash-low",
+      low: "gemini-3.6-flash-low",
+      medium: "gemini-3.6-flash-medium",
+      high: "gemini-3.6-flash-high",
+      xhigh: "gemini-3.6-flash-high",
+      max: "gemini-3.6-flash-high",
+    },
+    mode: "google-level",
+    mandatory: true,
+  },
+  {
+    id: "gemini-3.5-flash",
+    name: "Gemini 3.5 Flash",
+    members: [
+      "gemini-3.5-flash-extra-low",
+      "gemini-3.5-flash-low",
+      "gemini-3-flash-agent",
+    ],
+    defaultWireId: "gemini-3.5-flash-extra-low",
+    routes: {
+      minimal: "gemini-3.5-flash-extra-low",
+      low: "gemini-3.5-flash-extra-low",
+      medium: "gemini-3.5-flash-low",
+      high: "gemini-3-flash-agent",
+      xhigh: "gemini-3-flash-agent",
+      max: "gemini-3-flash-agent",
+    },
+    mode: "budget",
+    budgets: {
+      minimal: 1_000,
+      low: 1_000,
+      medium: 4_000,
+      high: 10_000,
+      xhigh: 10_000,
+      max: 10_000,
+    },
+  },
+  {
+    id: "gemini-3.1-pro",
+    name: "Gemini 3.1 Pro",
+    members: [
+      "gemini-3.1-pro-low",
+      "gemini-pro-agent",
+      // Discovery still publishes this deployment, but requests always fail.
+      "gemini-3.1-pro-high",
+    ],
+    defaultWireId: "gemini-3.1-pro-low",
+    routes: {
+      minimal: "gemini-3.1-pro-low",
+      low: "gemini-3.1-pro-low",
+      medium: "gemini-3.1-pro-low",
+      high: "gemini-pro-agent",
+      xhigh: "gemini-pro-agent",
+      max: "gemini-pro-agent",
+    },
+    mode: "budget",
+    budgets: {
+      minimal: 1_001,
+      low: 1_001,
+      medium: 1_001,
+      high: 10_001,
+      xhigh: 10_001,
+      max: 10_001,
+    },
+    retiredMembers: ["gemini-3.1-pro-high"],
+  },
+  {
+    id: "gemini-3-pro",
+    name: "Gemini 3 Pro",
+    members: ["gemini-3-pro-low", "gemini-3-pro-high"],
+    defaultWireId: "gemini-3-pro-low",
+    routes: {
+      minimal: "gemini-3-pro-low",
+      low: "gemini-3-pro-low",
+      medium: "gemini-3-pro-low",
+      high: "gemini-3-pro-high",
+      xhigh: "gemini-3-pro-high",
+      max: "gemini-3-pro-high",
+    },
+    mode: "google-level",
+  },
+  {
+    id: "gpt-oss-120b",
+    name: "GPT-OSS 120B",
+    members: ["gpt-oss-120b-medium"],
+    defaultWireId: "gpt-oss-120b-medium",
+    mode: "budget",
+  },
+  {
+    id: "claude-sonnet-4-6",
+    name: "Claude Sonnet 4.6",
+    members: ["claude-sonnet-4-6", "claude-sonnet-4-6-thinking"],
+    defaultWireId: "claude-sonnet-4-6",
+    mode: "budget",
+    retiredMembers: ["claude-sonnet-4-6-thinking"],
+  },
+  {
+    id: "claude-opus-4-6",
+    name: "Claude Opus 4.6",
+    members: ["claude-opus-4-6-thinking", "claude-opus-4-6"],
+    defaultWireId: "claude-opus-4-6-thinking",
+    mode: "budget",
+    retiredMembers: ["claude-opus-4-6"],
+  },
+  ...thinkingPairs([
+    ["claude-sonnet-4-5", "Claude Sonnet 4.5"],
+    ["claude-opus-4-5", "Claude Opus 4.5"],
+    ["gemini-2.5-flash", "Gemini 2.5 Flash"],
+  ]),
+];
+
+const familyById = new Map<string, Family>();
+for (const family of FAMILIES) {
+  familyById.set(family.id, family);
+  for (const member of family.members) familyById.set(member, family);
+}
+
+function clampLevel(
+  level: ThinkingLevel,
+): "minimal" | "low" | "medium" | "high" {
+  if (level === "xhigh" || level === "max") return "high";
+  return level;
+}
+
+export function collapseAntigravityModels<T extends AntigravityModelDefinition>(
+  models: readonly T[],
+): T[] {
+  const byId = new Map(models.map((model) => [model.id, model]));
+  const consumed = new Set<string>();
+  const collapsed: T[] = [];
+
+  for (const model of models) {
+    if (consumed.has(model.id)) continue;
+    const family = familyById.get(model.id);
+    if (!family) {
+      collapsed.push(model);
+      continue;
+    }
+    if (collapsed.some((entry) => entry.id === family.id)) continue;
+    for (const member of family.members) consumed.add(member);
+    consumed.add(family.id);
+    const logical = byId.get(family.id);
+    const source =
+      logical ?? family.members.map((id) => byId.get(id)).find(Boolean);
+    if (!source) continue;
+    const retired = new Set(family.retiredMembers ?? []);
+    const liveWireIds = new Set(
+      family.members.filter((id) => byId.has(id) && !retired.has(id)),
+    );
+    if (logical?.requestModelId && !retired.has(logical.requestModelId)) {
+      liveWireIds.add(logical.requestModelId);
+    } else if (logical && family.members.includes(family.id)) {
+      liveWireIds.add(family.id);
+    }
+    // A discovery response containing only a retired deployment is unusable;
+    // consume it without publishing a logical model that cannot be requested.
+    if (liveWireIds.size === 0) continue;
+    const requestModelId = liveWireIds.has(family.defaultWireId)
+      ? family.defaultWireId
+      : liveWireIds.values().next().value;
+    if (!requestModelId) continue;
+    const effortRouting: Partial<Record<ThinkingLevel, string>> = {};
+    for (const [effort, target] of Object.entries(family.routes ?? {}) as [
+      ThinkingLevel,
+      string,
+    ][]) {
+      if (
+        !retired.has(target) &&
+        (liveWireIds.has(target) || family.preserveAbsentEffortRoutes)
+      ) {
+        effortRouting[effort] = target;
+      }
+    }
+    collapsed.push({
+      ...source,
+      id: family.id,
+      name: family.name,
+      reasoning: true,
+      ...(requestModelId !== family.id ? { requestModelId } : {}),
+      ...(Object.keys(effortRouting).length > 0
+        ? { antigravityEffortRouting: effortRouting }
+        : {}),
+    });
+  }
+
+  return collapsed;
+}
+
+export function routeAntigravityModel(
+  modelId: string,
+  reasoning: ThinkingLevel | undefined,
+  thinkingBudgets: Partial<Record<ThinkingLevel, number>> | undefined,
+  overrides?: Pick<
+    AntigravityModelDefinition,
+    "requestModelId" | "antigravityEffortRouting"
+  >,
+): { wireModelId: string; thinkingConfig?: Record<string, unknown> } {
+  const family = familyById.get(modelId);
+  if (!family) {
+    if (!reasoning) return { wireModelId: modelId };
+    if (modelId.toLowerCase().includes("claude")) {
+      const level = clampLevel(reasoning);
+      const budget =
+        thinkingBudgets?.[level] ??
+        {
+          minimal: 1_024,
+          low: 8_192,
+          medium: 16_384,
+          high: 32_768,
+        }[level];
+      return {
+        wireModelId: modelId,
+        thinkingConfig: { includeThoughts: true, thinkingBudget: budget },
+      };
+    }
+    return {
+      wireModelId: modelId,
+      thinkingConfig: {
+        includeThoughts: true,
+        thinkingLevel: clampLevel(reasoning).toUpperCase(),
+      },
+    };
+  }
+
+  const effective = reasoning ?? (family.mandatory ? "minimal" : undefined);
+  const discoveredRouting = overrides?.antigravityEffortRouting;
+  const wireModelId = effective
+    ? ((discoveredRouting
+        ? (discoveredRouting[effective] ?? overrides?.requestModelId)
+        : (family.routes?.[effective] ?? overrides?.requestModelId)) ??
+      family.defaultWireId)
+    : (overrides?.requestModelId ?? family.defaultWireId);
+  if (!effective) {
+    return {
+      wireModelId,
+      thinkingConfig: { includeThoughts: false, thinkingBudget: 0 },
+    };
+  }
+  if (family.mode === "google-level") {
+    const level = clampLevel(effective);
+    return {
+      wireModelId,
+      thinkingConfig: {
+        includeThoughts: true,
+        thinkingLevel: level === "minimal" ? "LOW" : level.toUpperCase(),
+      },
+    };
+  }
+  const level = clampLevel(effective);
+  const budget =
+    thinkingBudgets?.[level] ??
+    family.budgets?.[effective] ??
+    {
+      minimal: 1_024,
+      low: 8_192,
+      medium: 16_384,
+      high: 32_768,
+    }[level];
+  return {
+    wireModelId,
+    thinkingConfig: { includeThoughts: true, thinkingBudget: budget },
+  };
+}

package/extensions/ai-providers/antigravity/with-resolvers.d.ts

@@ -0,0 +1,19 @@
+/**
+ * Promise.withResolvers type shim.
+ *
+ * The repo tsconfig targets ES2022 libs, but openpi runs on Node >=22.19
+ * (package.json engines), where Promise.withResolvers is available since
+ * Node 22.0. This declaration matches lib.es2024.promise.withresolvers.d.ts.
+ */
+
+declare global {
+  interface PromiseConstructor {
+    withResolvers<T>(): {
+      promise: Promise<T>;
+      resolve: (value: T | PromiseLike<T>) => void;
+      reject: (reason?: unknown) => void;
+    };
+  }
+}
+
+export {};

package/extensions/ai-providers/cursor/constants.ts

@@ -0,0 +1,5 @@
+/** Wire constants shared by Cursor OAuth, discovery, and AgentService. */
+export const CURSOR_API_URL = "https://api2.cursor.sh";
+export const CURSOR_CLIENT_VERSION = "cli-2026.07.23-e383d2b";
+export const CURSOR_RUN_PATH = "/agent.v1.AgentService/Run";
+export const CURSOR_MODELS_PATH = "/agent.v1.AgentService/GetUsableModels";

package/extensions/ai-providers/cursor/credentials.ts

@@ -0,0 +1,14 @@
+import type { OAuthCredentials } from "@earendil-works/pi-ai/compat";
+
+/** Cursor stores the short-lived access JWT and its refresh token together. */
+export type CursorCredentials = OAuthCredentials;
+
+/** pi passes the return value of this function to `streamSimple.apiKey`. */
+export function getCursorApiKey(credentials: CursorCredentials): string {
+  return credentials.access;
+}
+
+/** Accept a bare access token for hand-written or older auth.json entries. */
+export function decodeCursorApiKey(value: string | undefined): string {
+  return value?.trim() ?? "";
+}

package/extensions/ai-providers/cursor/discovery.ts

@@ -0,0 +1,291 @@
+import { randomUUID } from "node:crypto";
+import * as http2 from "node:http2";
+import type { Api, Model, RefreshModelsContext } from "@earendil-works/pi-ai";
+import {
+  CURSOR_API_URL,
+  CURSOR_CLIENT_VERSION,
+  CURSOR_MODELS_PATH,
+} from "./constants.ts";
+import type { CursorModelDefinition } from "./models.ts";
+import {
+  GetUsableModelsRequestSchema,
+  GetUsableModelsResponseSchema,
+  type ModelDetails,
+} from "./proto.ts";
+import { create, fromBinary, toBinary } from "./protobuf.ts";
+import { connectCursorHttp2 } from "./proxy.ts";
+
+export interface CursorModelDiscoveryOptions {
+  apiKey: string;
+  baseUrl?: string;
+  clientVersion?: string;
+  timeoutMs?: number;
+  signal?: AbortSignal;
+  customModelIds?: string[];
+}
+
+const FALLBACK_CONTEXT_WINDOW = 200_000;
+const DEFAULT_MAX_TOKENS = 64_000;
+const MAX_DISCOVERY_RESPONSE_BYTES = 16 * 1024 * 1024;
+const ONE_MILLION_CONTEXT_WINDOW = 1_000_000;
+
+/**
+ * GetUsableModels has no numeric context-window field. Recover 1M only from
+ * signals Cursor does send; use 200k as the conservative unknown-model fallback.
+ */
+function resolveContextWindow(details: ModelDetails, id: string): number {
+  const labels = [
+    id,
+    details.displayName,
+    details.displayNameShort,
+    details.displayModelId,
+    ...details.aliases,
+  ].join(" ");
+  if (/\b1m\b/i.test(labels)) return ONE_MILLION_CONTEXT_WINDOW;
+  if (details.maxMode && /claude|gemini|gpt-5\.6-sol/i.test(id)) {
+    return ONE_MILLION_CONTEXT_WINDOW;
+  }
+  if (isNativeOneMillionModel(id)) return ONE_MILLION_CONTEXT_WINDOW;
+  return FALLBACK_CONTEXT_WINDOW;
+}
+
+/** Cursor serves these coding families with a native, unlabeled 1M window. */
+function isNativeOneMillionModel(id: string): boolean {
+  const bareId = id.split("/").at(-1)?.toLowerCase() ?? id.toLowerCase();
+  if (/^(?:kimi-)?k3$/.test(bareId)) return true;
+
+  const glm =
+    /^glm-(\d{1,2})(?:\.(\d+))?(v)?(?:-(air|turbo|flashx|flash|preview))?$/.exec(
+      bareId,
+    );
+  if (!glm || glm[3]) return false;
+  const variant = glm[4];
+  if (variant && variant !== "air" && variant !== "turbo") return false;
+  const major = Number(glm[1]);
+  const minor = Number(glm[2] ?? 0);
+  return major > 5 || (major === 5 && minor >= 2);
+}
+
+/**
+ * Fetch account-specific models over Cursor's HTTP/2 Connect endpoint.
+ * `null` means transport/protocol failure; an empty successful response is
+ * deliberately represented as `[]` so callers can choose their fallback policy.
+ */
+export async function fetchCursorUsableModels(
+  options: CursorModelDiscoveryOptions,
+): Promise<CursorModelDefinition[] | null> {
+  const token = options.apiKey.trim();
+  if (!token || options.signal?.aborted) return null;
+  const baseUrl = (options.baseUrl ?? CURSOR_API_URL).replace(/\/+$/, "");
+  const request = create(GetUsableModelsRequestSchema, {
+    customModelIds: normalizeModelIds(options.customModelIds),
+  });
+  const requestBytes = toBinary(GetUsableModelsRequestSchema, request);
+  const responseBytes = await requestHttp2(
+    baseUrl,
+    requestBytes,
+    token,
+    options,
+  );
+  if (!responseBytes) return null;
+
+  const payload = decodeUnaryPayload(responseBytes);
+  if (!payload) return null;
+  let response;
+  try {
+    response = fromBinary(GetUsableModelsResponseSchema, payload);
+  } catch {
+    return null;
+  }
+
+  const models: CursorModelDefinition[] = [];
+  const seen = new Set<string>();
+  for (const details of response.models) {
+    const normalized = normalizeCursorModel(details, baseUrl);
+    if (!normalized || seen.has(normalized.id)) continue;
+    seen.add(normalized.id);
+    models.push(normalized);
+  }
+  models.sort((left, right) => left.id.localeCompare(right.id));
+  return models;
+}
+
+/** Provider extension hook: discovery always uses the credential passed by pi. */
+export async function fetchCursorModels(
+  context: RefreshModelsContext,
+): Promise<CursorModelDefinition[]> {
+  if (!context.allowNetwork) return [];
+  context.signal.throwIfAborted();
+  const credential = context.credential;
+  const apiKey =
+    credential?.type === "oauth"
+      ? credential.access
+      : credential?.type === "api_key"
+        ? credential.key
+        : undefined;
+  if (!apiKey) return [];
+  const discovered = await fetchCursorUsableModels({
+    apiKey,
+    signal: context.signal,
+  });
+  context.signal.throwIfAborted();
+  if (discovered === null) {
+    throw new Error("Cursor model discovery failed");
+  }
+  return discovered;
+}
+
+function normalizeModelIds(ids: readonly string[] | undefined): string[] {
+  if (!ids) return [];
+  const result = new Set<string>();
+  for (const id of ids) {
+    if (typeof id !== "string") continue;
+    const value = id.trim();
+    if (value) result.add(value);
+  }
+  return [...result];
+}
+
+function normalizeCursorModel(
+  details: ModelDetails,
+  baseUrl: string,
+): CursorModelDefinition | null {
+  const id = details.modelId.trim();
+  if (!id) return null;
+  const name =
+    [
+      details.displayName,
+      details.displayNameShort,
+      details.displayModelId,
+      ...details.aliases,
+    ]
+      .map((value) => value.trim())
+      .find(Boolean) ?? id;
+  const multimodal = supportsCursorImages(id);
+  return {
+    id,
+    name,
+    api: "cursor-agent",
+    provider: "cursor",
+    baseUrl,
+    reasoning: Boolean(details.thinkingDetails) || /^cursor-grok-\d/i.test(id),
+    input: multimodal ? ["text", "image"] : ["text"],
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+    contextWindow: resolveContextWindow(details, id),
+    maxTokens: DEFAULT_MAX_TOKENS,
+    ...(details.maxMode ? { cursorMaxMode: true } : {}),
+  };
+}
+
+/** GetUsableModels omits modality metadata for Cursor-native image families. */
+function supportsCursorImages(id: string): boolean {
+  const lower = id.toLowerCase();
+  if (/claude|gemini|gpt-|codex/.test(lower)) return true;
+  const bareId = lower.split("/").at(-1) ?? lower;
+  return (
+    /^(?:kimi-)?k3(?:$|[._:-])/.test(bareId) ||
+    /^cursor-grok-4(?:$|[._:-])/.test(bareId) ||
+    /^(?:cursor-)?composer-2\.5(?:$|[._:-])/.test(bareId)
+  );
+}
+
+function buildHeaders(
+  apiKey: string,
+  clientVersion: string,
+): Record<string, string> {
+  return {
+    "content-type": "application/proto",
+    te: "trailers",
+    authorization: `Bearer ${apiKey}`,
+    "x-ghost-mode": "true",
+    "x-cursor-client-version": clientVersion,
+    "x-cursor-client-type": "cli",
+    "x-request-id": randomUUID(),
+  };
+}
+
+async function requestHttp2(
+  baseUrl: string,
+  body: Uint8Array,
+  apiKey: string,
+  options: CursorModelDiscoveryOptions,
+): Promise<Uint8Array | null> {
+  const timeoutMs = options.timeoutMs ?? 5_000;
+  let client: http2.ClientHttp2Session;
+  try {
+    client = await connectCursorHttp2(baseUrl, {
+      signal: options.signal,
+      timeoutMs,
+    });
+  } catch {
+    return null;
+  }
+  const { promise, resolve } = Promise.withResolvers<Uint8Array | null>();
+  let settled = false;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  let removeAbortListener: (() => void) | undefined;
+  const finish = (result: Uint8Array | null, destroy = false) => {
+    if (settled) return;
+    settled = true;
+    if (timer) clearTimeout(timer);
+    removeAbortListener?.();
+    if (destroy) client.destroy();
+    else client.close();
+    resolve(result);
+  };
+  timer = setTimeout(() => finish(null, true), timeoutMs);
+  client.once("error", () => finish(null, true));
+  const req = client.request({
+    ":method": "POST",
+    ":path": CURSOR_MODELS_PATH,
+    ...buildHeaders(apiKey, options.clientVersion ?? CURSOR_CLIENT_VERSION),
+  });
+  const chunks: Buffer[] = [];
+  let responseBytes = 0;
+  req.on("response", (headers) => {
+    const status = Number(headers[":status"] ?? 0);
+    if (status < 200 || status >= 300) finish(null, true);
+  });
+  req.on("data", (chunk: Buffer) => {
+    responseBytes += chunk.length;
+    if (responseBytes > MAX_DISCOVERY_RESPONSE_BYTES) {
+      req.close(http2.constants.NGHTTP2_CANCEL);
+      finish(null, true);
+      return;
+    }
+    chunks.push(chunk);
+  });
+  req.on("end", () => finish(new Uint8Array(Buffer.concat(chunks))));
+  req.once("error", () => finish(null, true));
+  const onAbort = () => {
+    options.signal?.removeEventListener("abort", onAbort);
+    req.close(http2.constants.NGHTTP2_CANCEL);
+    finish(null, true);
+  };
+  if (options.signal) {
+    if (options.signal.aborted) onAbort();
+    else {
+      options.signal.addEventListener("abort", onAbort, { once: true });
+      removeAbortListener = () =>
+        options.signal?.removeEventListener("abort", onAbort);
+    }
+  }
+  req.end(Buffer.from(body));
+  return promise;
+}
+
+/** Decode the first uncompressed Connect data frame, or accept raw unary proto. */
+export function decodeUnaryPayload(body: Uint8Array): Uint8Array | null {
+  if (body.length === 0) return body;
+  if (body.length < 5) return body;
+  const flags = body[0]!;
+  const size = new DataView(body.buffer, body.byteOffset, 5).getUint32(1);
+  const end = 5 + size;
+  // Unary Connect responses are normally one data frame followed by an
+  // optional end-stream frame. If the prefix is not a valid uncompressed frame,
+  // accept the raw protobuf response used by older Cursor deployments.
+  if (flags > 3 || end > body.length || (flags & 1) !== 0) return body;
+  const data = body.subarray(5, end);
+  if ((flags & 2) !== 0) return null;
+  return data;
+}