@jameslovespancakes/pi-plus 1.0.15 → 1.0.16

package/src/core/gemini/request.ts

@@ -0,0 +1,259 @@
+import type { Api, Model, ModelThinkingLevel, ToolChoice, TranscriptContext } from "@earendil-works/pi-ai";
+import { sanitizeSurrogates } from "@earendil-works/pi-ai/utils/sanitize-unicode";
+import { getCurrentSystemPrompt, getCurrentTools } from "@earendil-works/pi-ai/utils/transcript";
+import { stableUuid } from "./client.ts";
+import { runtimeModelId, thinkingConfig } from "./models.ts";
+import { bridgeSchema, selfContainedSchema } from "./schema.ts";
+
+/**
+ * Builds a Gemini `streamGenerateContent` request.
+ *
+ * Inside the envelope the body is ordinary Gemini, so message conversion,
+ * thought-signature validation, tool-call ids and function-calling mode all
+ * come from pi's own Google adapter. What is added here is only what this
+ * backend demands beyond the public Gemini API: the runtime model id, its
+ * thinking budget, the Claude/GPT-OSS schema bridge, a few conversation-shape
+ * repairs it enforces, and the agent envelope it expects.
+ */
+
+type GoogleShared = typeof import("@earendil-works/pi-ai/api/google-shared");
+let googleSharedModule: Promise<GoogleShared> | undefined;
+
+/** Loaded on first request, as pi loads its own Google adapter: it pulls in `@google/genai`. */
+export function googleShared(): Promise<GoogleShared> {
+  return googleSharedModule ??= import("@earendil-works/pi-ai/api/google-shared");
+}
+
+export interface Part {
+  text?: string;
+  thought?: boolean;
+  thoughtSignature?: string;
+  inlineData?: { mimeType?: string; data?: string };
+  functionCall?: { name?: string; args?: Record<string, unknown>; id?: string };
+  functionResponse?: { name?: string; id?: string; response?: Record<string, unknown>; parts?: Part[] };
+}
+
+export interface Content {
+  role: "user" | "model";
+  parts: Part[];
+}
+
+export interface RequestOptions {
+  /** Level after pi's clamp; undefined means thinking off. */
+  reasoning?: ModelThinkingLevel;
+  temperature?: number;
+  maxTokens?: number;
+  toolChoice?: ToolChoice | "any";
+  sessionId?: string;
+}
+
+export interface GeminiRequest {
+  project: string;
+  model: string;
+  request: Record<string, unknown>;
+  requestType: "agent";
+  userAgent: "antigravity";
+  requestId: string;
+}
+
+/** The only text this module ever adds to a conversation; see {@link repairContents}. */
+export const CONTINUE_TEXT = "Continue the active task using the available instructions and context.";
+
+const isText = (part: Part) => typeof part.text === "string" && part.text.trim().length > 0 && !part.thought;
+const hasFunctionCall = (turn: Content) => turn.parts.some((part) => part.functionCall);
+
+/** Gemini 3+ rejects a replayed function call that lacks its thought signature. */
+export function requiresThoughtSignatures(runtimeId: string): boolean {
+  if (!runtimeId.startsWith("gemini-")) return false;
+  const major = /^gemini-(\d+)/.exec(runtimeId)?.[1];
+  // Unversioned agent runtimes (`gemini-pro-agent`) are current-generation.
+  return major === undefined || Number(major) >= 3;
+}
+
+/** Claude and GPT-OSS are served through the custom-tool bridge. */
+export function usesToolBridge(runtimeId: string): boolean {
+  return runtimeId.startsWith("claude-") || runtimeId.startsWith("gpt-oss-");
+}
+
+function observationText(name: string, args: Record<string, unknown> | undefined, response: Part["functionResponse"]): string {
+  const argsText = args && Object.keys(args).length > 0 ? ` (${JSON.stringify(args)})` : "";
+  const payload = response?.response ?? {};
+  const failed = "error" in payload;
+  const value = failed ? payload.error : "output" in payload ? payload.output : payload;
+  const body = typeof value === "string" ? value : JSON.stringify(value);
+  return `[${failed ? "Failed observation" : "Observation"} from \`${name}\`${argsText}:\n${body}]`;
+}
+
+/**
+ * History from another model carries tool calls without this model's thought
+ * signature, which Gemini 3 rejects outright. Such a call and its result are
+ * replayed as plain text instead, so the model keeps what happened without an
+ * unverifiable call in its own voice. Gemini only checks the first call of a
+ * turn, so a signed first call keeps the whole turn intact.
+ */
+function observeUnsignedCalls(contents: Content[]): Content[] {
+  const pending = new Map<string, { name: string; args?: Record<string, unknown> }>();
+  const keyOf = (call: { id?: string; name?: string }) => call.id || `name:${call.name ?? ""}`;
+
+  return contents.map((turn) => {
+    if (turn.role === "model") {
+      const calls = turn.parts.filter((part) => part.functionCall);
+      if (calls.length === 0 || calls[0].thoughtSignature) return turn;
+      for (const { functionCall } of calls) {
+        pending.set(keyOf(functionCall!), { name: functionCall!.name ?? "tool", args: functionCall!.args });
+      }
+      return { ...turn, parts: turn.parts.filter((part) => !part.functionCall) };
+    }
+    if (pending.size === 0) return turn;
+    return {
+      ...turn,
+      parts: turn.parts.flatMap((part): Part[] => {
+        const response = part.functionResponse;
+        const call = response && pending.get(keyOf(response));
+        if (!response || !call) return [part];
+        pending.delete(keyOf(response));
+        return [{ text: sanitizeSurrogates(observationText(call.name, call.args, response)) }, ...(response.parts ?? [])];
+      }),
+    };
+  });
+}
+
+/** Adjacent turns with the same role become one; empty turns disappear. */
+function mergeTurns(contents: Content[]): Content[] {
+  const merged: Content[] = [];
+  for (const turn of contents) {
+    if (turn.parts.length === 0) continue;
+    const last = merged.at(-1);
+    if (last?.role === turn.role) last.parts.push(...turn.parts);
+    else merged.push({ role: turn.role, parts: [...turn.parts] });
+  }
+  return merged;
+}
+
+/**
+ * Conversation shapes the public Gemini API tolerates but this backend
+ * rejects with a 400. Compaction and model switches produce every one of
+ * them, so they are repaired rather than surfaced:
+ *
+ *   - an unsigned tool call on Gemini 3 (see {@link observeUnsignedCalls});
+ *   - a conversation that does not open with a user turn;
+ *   - no natural-language user text anywhere, e.g. a tool-only continuation;
+ *   - a request that ends on a model turn.
+ */
+export function repairContents(contents: Content[], requireSignatures: boolean): Content[] {
+  const turns = mergeTurns(requireSignatures ? observeUnsignedCalls(contents) : contents);
+  const bridge = (): Part => ({ text: CONTINUE_TEXT });
+
+  if (turns.length === 0 || turns[0].role === "model") turns.unshift({ role: "user", parts: [bridge()] });
+
+  if (!turns.some((turn) => turn.role === "user" && turn.parts.some(isText))) {
+    turns.find((turn) => turn.role === "user")!.parts.push(bridge());
+  }
+
+  const last = turns.at(-1)!;
+  if (last.role === "model") {
+    if (hasFunctionCall(last)) throw new Error("Gemini request ends on a tool call with no result.");
+    turns.push({ role: "user", parts: [bridge()] });
+  }
+  return turns;
+}
+
+function toolDeclarations(declared: { functionDeclarations: Record<string, unknown>[] }[], bridge: boolean) {
+  return declared.map((group) => ({
+    functionDeclarations: group.functionDeclarations.map(({ parametersJsonSchema, parameters, ...rest }) => ({
+      ...rest,
+      ...(bridge
+        ? { parameters: bridgeSchema(parametersJsonSchema ?? parameters) }
+        : { parametersJsonSchema: selfContainedSchema(parametersJsonSchema ?? parameters) }),
+    })),
+  }));
+}
+
+/** Random signed 64-bit decimal, the shape the Antigravity CLI uses for session ids. */
+function randomSessionId(): string {
+  const bytes = crypto.getRandomValues(new Uint8Array(8));
+  return new DataView(bytes.buffer).getBigInt64(0, true).toString();
+}
+
+/**
+ * The agent envelope the Antigravity CLI sends. Ids are derived from the pi
+ * session, so one conversation keeps one trajectory across requests and
+ * restarts without any state held here.
+ */
+function envelope(context: TranscriptContext, contents: Content[], runtimeId: string, sessionId?: string) {
+  const first = context.messages[0];
+  const seed = sessionId ?? (first ? `${first.role}:${first.timestamp ?? ""}` : crypto.randomUUID());
+  const conversationId = stableUuid(`antigravity:conv:${seed}`);
+  const trajectoryId = stableUuid(`antigravity:traj:${seed}`);
+  const step = Math.max(1, contents.length);
+  const turn = context.messages.filter((message) =>
+    message.role === "assistant" && message.stopReason !== "error" && message.stopReason !== "aborted").length;
+  const claude = String(runtimeId.startsWith("claude-"));
+
+  return {
+    sessionId: sessionId ?? randomSessionId(),
+    requestId: `agent/${conversationId}/${Date.now()}/${trajectoryId}/${step}`,
+    labels: {
+      last_step_index: String(step - 1),
+      request_id: `${trajectoryId}-${turn}`,
+      trajectory_id: trajectoryId,
+      used_claude: claude,
+      used_claude_conservative: claude,
+      used_non_gemini_model: String(!runtimeId.startsWith("gemini-")),
+    },
+  };
+}
+
+export async function buildRequest(
+  model: Model<Api>,
+  context: TranscriptContext,
+  projectId: string,
+  options: RequestOptions = {},
+): Promise<GeminiRequest> {
+  const google = await googleShared();
+  // pi's Google helpers are typed against its own Google APIs; the body inside
+  // the envelope has exactly the same model semantics.
+  const googleModel = model as unknown as Model<"google-generative-ai">;
+  const runtimeId = runtimeModelId(model, options.reasoning);
+
+  const contents = repairContents(
+    google.convertMessages(googleModel, context) as unknown as Content[],
+    requiresThoughtSignatures(runtimeId),
+  );
+
+  // The system prompt and tools live in the transcript's system messages,
+  // never on the context object; reading them any other way sends neither.
+  const systemPrompt = getCurrentSystemPrompt(context.messages);
+  const tools = getCurrentTools(context.messages);
+  // Strict tool sampling (Gemini's VALIDATED mode) is not offered by this backend.
+  const declared = google.convertTools(tools, false, false);
+  const mode = tools.length > 0
+    ? google.resolveGoogleFunctionCallingMode(tools, options.toolChoice, false)
+    : undefined;
+
+  const thinking = thinkingConfig(runtimeId, options.reasoning);
+  const generationConfig = {
+    ...(options.temperature !== undefined && { temperature: options.temperature }),
+    maxOutputTokens: Math.min(options.maxTokens ?? model.maxTokens, model.maxTokens),
+    ...(thinking && { thinkingConfig: thinking }),
+  };
+
+  const { sessionId, requestId, labels } = envelope(context, contents, runtimeId, options.sessionId);
+
+  return {
+    project: projectId,
+    model: runtimeId,
+    request: {
+      contents,
+      ...(systemPrompt && { systemInstruction: { role: "user", parts: [{ text: sanitizeSurrogates(systemPrompt) }] } }),
+      generationConfig,
+      ...(declared && { tools: toolDeclarations(declared, usesToolBridge(runtimeId)) }),
+      ...(mode !== undefined && { toolConfig: { functionCallingConfig: { mode } } }),
+      sessionId,
+      labels,
+    },
+    requestType: "agent",
+    userAgent: "antigravity",
+    requestId,
+  };
+}

package/src/core/gemini/schema.ts

@@ -0,0 +1,142 @@
+/**
+ * Tool schemas in the two shapes the Antigravity backend accepts.
+ *
+ *   - Gemini reads full JSON Schema from `parametersJsonSchema`, but only
+ *     self-contained: `$ref` must already be resolved and `$defs` removed.
+ *   - Claude and GPT-OSS go through a custom-tool bridge that reads the
+ *     legacy protobuf `parameters` field. It rejects every keyword it does not
+ *     know (`nullable`, `anyOf`, `format`, `const`, …) with a 400, so the
+ *     schema is reduced to an allowlist rather than a denylist: a new JSON
+ *     Schema keyword can never break a request. pi still validates the
+ *     arguments the model returns against the original schema.
+ *
+ * Tool schemas can come from MCP servers, so expansion is bounded. A
+ * reference that cannot be resolved — missing, circular, or past a bound —
+ * becomes the unconstrained schema `{}` instead of failing the whole request.
+ */
+
+type Json = Record<string, unknown>;
+
+const MAX_DEPTH = 32;
+const MAX_NODES = 10_000;
+
+const METADATA = new Set(["$schema", "$id", "$anchor", "$dynamicAnchor", "$vocabulary", "$comment", "$defs", "definitions"]);
+/** Keywords whose value is a map of *names* to schemas; the names are never keywords. */
+const SCHEMA_MAPS = new Set(["properties", "patternProperties", "dependentSchemas"]);
+/** Keywords whose value is a schema, or an array of schemas. */
+const SCHEMA_VALUES = new Set([
+  "items", "prefixItems", "additionalItems", "additionalProperties", "unevaluatedItems", "unevaluatedProperties",
+  "contains", "propertyNames", "not", "if", "then", "else", "contentSchema", "allOf", "anyOf", "oneOf",
+]);
+
+const BRIDGE_KEYWORDS = new Set(["type", "description", "properties", "required", "items", "enum"]);
+
+function isRecord(value: unknown): value is Json {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+/** RFC 6901 pointer into the root schema; only local (`#…`) references resolve. */
+function resolvePointer(root: unknown, ref: string): unknown {
+  if (ref === "#") return root;
+  if (!ref.startsWith("#/")) return undefined;
+  let node: unknown = root;
+  for (const raw of ref.slice(2).split("/")) {
+    const key = raw.replace(/~1/g, "/").replace(/~0/g, "~");
+    if (Array.isArray(node)) {
+      if (!/^(0|[1-9]\d*)$/.test(key)) return undefined;
+      node = node[Number(key)];
+    } else if (isRecord(node) && Object.hasOwn(node, key)) {
+      node = node[key];
+    } else {
+      return undefined;
+    }
+  }
+  return node;
+}
+
+/** A JSON Schema with every local `$ref` inlined and schema metadata removed. */
+export function selfContainedSchema(schema: unknown): Json {
+  let nodes = 0;
+
+  const walk = (node: unknown, refs: ReadonlySet<string>, depth: number): unknown => {
+    if (++nodes > MAX_NODES || depth > MAX_DEPTH) return {};
+    if (Array.isArray(node)) return node.map((item) => walk(item, refs, depth + 1));
+    if (!isRecord(node)) return node;
+
+    if (typeof node.$ref === "string") {
+      const { $ref: ref, ...siblings } = node;
+      const target = refs.has(ref) ? undefined : resolvePointer(schema, ref);
+      const resolved = target === undefined ? {} : walk(target, new Set([...refs, ref]), depth + 1);
+      const rest = walk(siblings, refs, depth + 1);
+      return { ...(isRecord(resolved) ? resolved : {}), ...(isRecord(rest) ? rest : {}) };
+    }
+
+    const out: Json = {};
+    for (const [key, value] of Object.entries(node)) {
+      if (METADATA.has(key)) continue;
+      if (SCHEMA_MAPS.has(key) && isRecord(value)) {
+        out[key] = Object.fromEntries(Object.entries(value).map(([name, sub]) => [name, walk(sub, refs, depth + 1)]));
+      } else if (SCHEMA_VALUES.has(key)) {
+        out[key] = walk(value, refs, depth + 1);
+      } else {
+        // `enum`, `default`, `examples`… are data, not schemas: copied verbatim.
+        out[key] = value;
+      }
+    }
+    return out;
+  };
+
+  return asObjectRoot(walk(schema, new Set(), 0));
+}
+
+/** Function declarations must describe an object, even for a tool with no parameters. */
+function asObjectRoot(schema: unknown): Json {
+  if (!isRecord(schema)) return { type: "object", properties: {} };
+  return schema.type ? schema : { ...schema, type: "object", properties: schema.properties ?? {} };
+}
+
+/** `["string", "null"]` → `"string"`: the bridge takes a single type. */
+function singleType(value: unknown): string | undefined {
+  if (typeof value === "string") return value;
+  return Array.isArray(value) ? value.find((entry): entry is string => typeof entry === "string" && entry !== "null") : undefined;
+}
+
+/**
+ * A union the bridge cannot express is narrowed to its first non-null branch,
+ * which keeps the common `T | null` shape typed instead of unconstrained.
+ */
+function firstBranch(node: Json): Json | undefined {
+  const branches = Array.isArray(node.anyOf) ? node.anyOf : Array.isArray(node.oneOf) ? node.oneOf : undefined;
+  return branches?.find((branch): branch is Json => isRecord(branch) && branch.type !== "null");
+}
+
+function toBridge(node: unknown): unknown {
+  if (!isRecord(node)) return node;
+  const branch = node.type === undefined ? firstBranch(node) : undefined;
+  const source = branch ? { ...branch, ...(node.description !== undefined && { description: node.description }) } : node;
+
+  const out: Json = {};
+  for (const [key, value] of Object.entries(source)) {
+    if (!BRIDGE_KEYWORDS.has(key)) continue;
+    if (key === "type") {
+      const type = singleType(value);
+      if (type) out.type = type;
+    } else if (key === "properties" && isRecord(value)) {
+      out.properties = Object.fromEntries(Object.entries(value).map(([name, sub]) => [name, toBridge(sub)]));
+    } else if (key === "enum") {
+      // The bridge's enum is string-only; a mixed enum is dropped rather than coerced.
+      if (Array.isArray(value) && value.every((entry) => typeof entry === "string")) out.enum = value;
+    } else if (key === "items") {
+      out.items = toBridge(Array.isArray(value) ? value[0] : value);
+    } else {
+      out[key] = value;
+    }
+  }
+  if (typeof source.const === "string" && out.enum === undefined) out.enum = [source.const];
+  return out;
+}
+
+/** The subset of a tool schema Gemini's Claude/GPT-OSS bridge accepts. */
+export function bridgeSchema(schema: unknown): Json {
+  return asObjectRoot(toBridge(selfContainedSchema(schema)));
+}