pi-provider-antigravity 0.13.0 → 0.15.0

package/README.md

@@ -22,7 +22,7 @@ Use Google Antigravity models directly in [Pi](https://pi.dev).
 
 ## Requirements
 
-- Pi >= 0.80.0
+- Pi >= 0.86.0 (For Pi 0.80.0 ~ 0.85.x, see [Legacy Pi Support](#legacy-pi-support))
 - A Google account with Antigravity access
 
 ## Installation
@@ -37,6 +37,22 @@ Or try it without installing:
 pi -e npm:pi-provider-antigravity
 ```
 
+### Legacy Pi Support
+
+We strongly recommend upgrading to Pi >= 0.86.0 to take advantage of the latest SDK capabilities, streaming reliability improvements, and full feature parity.
+
+However, if you are unable to upgrade your Pi installation just yet (for environment or workflow reasons), you can pin to the compatible `0.13.0` release for Pi 0.80.0 ~ 0.85.x:
+
+```bash
+pi install npm:pi-provider-antigravity@0.13.0
+```
+
+Or try it without installing:
+
+```bash
+pi -e npm:pi-provider-antigravity@0.13.0
+```
+
 ## Quick Start
 
 1. Authenticate:

package/dist/index.js

@@ -28,13 +28,15 @@ function resolveToken(auth) {
 async function postAntigravityRaw({
   auth,
   path: path2,
+  endpoint = DEFAULT_ENDPOINT,
+  headers,
   body,
   signal
 }) {
   const token = resolveToken(auth);
-  return fetch(`${DEFAULT_ENDPOINT}/${path2}`, {
+  return fetch(`${endpoint}/${path2}`, {
     method: "POST",
-    headers: buildAntigravityHeaders(token),
+    headers: { ...buildAntigravityHeaders(token), ...headers },
     body: JSON.stringify(body),
     signal
   });
@@ -56,7 +58,7 @@ async function postAntigravityStream(params) {
   if (!res.body) {
     throw new Error(`No response stream received from Antigravity for ${params.path}.`);
   }
-  return res.body;
+  return { response: res, stream: res.body };
 }
 function buildAntigravityHeaders(token, userAgent = DEFAULT_USER_AGENT) {
   return {
@@ -122,7 +124,8 @@ async function fetchProjectId(token, signal) {
     if (data.cloudaicompanionProject) {
       return data.cloudaicompanionProject;
     }
-  } catch {
+  } catch (err) {
+    if (signal?.aborted) throw err;
   }
   return FALLBACK_PROJECT_ID;
 }
@@ -142,16 +145,37 @@ async function loginAntigravity(callbacks) {
   });
   const authUrl = `${AUTH_URL}?${authParams.toString()}`;
   callbacks.onAuth({ url: authUrl });
+  let timer;
+  const abortPromise = new Promise((_, rej) => {
+    if (callbacks.signal?.aborted) {
+      rej(callbacks.signal.reason ?? new Error("Login cancelled"));
+      return;
+    }
+    callbacks.signal?.addEventListener(
+      "abort",
+      () => rej(callbacks.signal?.reason ?? new Error("Login cancelled")),
+      { once: true }
+    );
+  });
   const manualPromptPromise = (async () => {
     const input = await callbacks.onPrompt({
       message: "Paste authorization code from browser:"
     });
     return extractCodeFromInput(input);
   })();
-  const timeout = new Promise(
-    (_, rej) => setTimeout(() => rej(new Error("Timed out waiting for authorization code (5 minutes)")), OAUTH_CALLBACK_TIMEOUT_MS)
-  );
-  const code = await Promise.race([manualPromptPromise, timeout]);
+  const timeoutPromise = new Promise((_, rej) => {
+    timer = setTimeout(
+      () => rej(new Error("Timed out waiting for authorization code (5 minutes)")),
+      OAUTH_CALLBACK_TIMEOUT_MS
+    );
+  });
+  let code;
+  try {
+    code = await Promise.race([manualPromptPromise, timeoutPromise, abortPromise]);
+  } finally {
+    if (timer) clearTimeout(timer);
+  }
+  callbacks.onProgress?.("Exchanging authorization code for tokens...");
   const tokenParams = new URLSearchParams({
     client_id: CLIENT_ID,
     client_secret: CLIENT_SECRET,
@@ -163,14 +187,16 @@ async function loginAntigravity(callbacks) {
   const tokenRes = await fetch(TOKEN_URL, {
     method: "POST",
     headers: { "Content-Type": "application/x-www-form-urlencoded" },
-    body: tokenParams.toString()
+    body: tokenParams.toString(),
+    signal: callbacks.signal
   });
   if (!tokenRes.ok) {
     const errText = await tokenRes.text();
     throw new Error(`Failed to exchange token (${tokenRes.status}): ${errText}`);
   }
   const tokens = await tokenRes.json();
-  const projectId = await fetchProjectId(tokens.access_token);
+  callbacks.onProgress?.("Resolving Code Assist project ID...");
+  const projectId = await fetchProjectId(tokens.access_token, callbacks.signal);
   return {
     refresh: tokens.refresh_token || "",
     access: JSON.stringify({ token: tokens.access_token, projectId }),
@@ -527,6 +553,7 @@ function synthesizeDynamicModel(baseId, items) {
   const isGpt = family === "gpt";
   const defaultContext = isFlash ? 1048576 : isClaude ? 25e4 : isGpt ? 128e3 : 1048576;
   const defaultMaxOutput = isClaude ? 64e3 : isGpt ? 32768 : 65536;
+  const promptCache = isClaude ? { short: 300, long: 3600 } : { short: 300 };
   return {
     id: baseId,
     name: formatModelDisplayName(baseId, repItem?.displayName),
@@ -537,6 +564,7 @@ function synthesizeDynamicModel(baseId, items) {
     thinkingLevelMap,
     input: repItem?.supportsImages ? ["text", "image"] : ["text"],
     cost: estimateModelCost(baseId),
+    promptCache,
     contextWindow: repItem?.maxTokens || defaultContext,
     maxTokens: repItem?.maxOutputTokens || defaultMaxOutput
   };
@@ -1384,11 +1412,19 @@ var QuotaStatusCoordinator = class {
 
 // src/stream.ts
 import {
+  calculateCost,
   createAssistantMessageEventStream
 } from "@earendil-works/pi-ai";
 
 // src/builder.ts
 import { createHash as createHash2, randomUUID } from "node:crypto";
+import {
+  collapseSystemMessages,
+  getCurrentSystemPrompt,
+  getCurrentTools,
+  normalizeContext,
+  withoutInitialSystemMessage
+} from "@earendil-works/pi-ai";
 var BASE64_SIGNATURE_PATTERN = /^[A-Za-z0-9+/_-]+={0,2}$/;
 var SKIP_THOUGHT_SIGNATURE_VALIDATOR = "skip_thought_signature_validator";
 function isValidThoughtSignature(sig) {
@@ -1400,11 +1436,11 @@ function resolveThoughtSignature(isSameProviderAndModel, sig) {
   if (!isSameProviderAndModel) return void 0;
   return isValidThoughtSignature(sig) ? sig : void 0;
 }
-function deriveSessionId(context, explicitSessionId) {
+function deriveSessionId(conversation, explicitSessionId) {
   if (explicitSessionId && typeof explicitSessionId === "string" && explicitSessionId.trim().length > 0) {
     return explicitSessionId.trim();
   }
-  const firstMsg = context.messages?.[0];
+  const firstMsg = conversation?.[0];
   if (!firstMsg) {
     return randomUUID();
   }
@@ -1497,9 +1533,9 @@ function convertTools(tools) {
     }
   ];
 }
-function translateTurnTrace(context, runtimeModelId) {
+function translateTurnTrace(conversation, runtimeModelId) {
   const contents = [];
-  const messages = context.messages || [];
+  const messages = conversation || [];
   const isGeminiRequest = classifyModelFamily(runtimeModelId) === "gemini";
   const isGptRequest = classifyModelFamily(runtimeModelId) === "gpt";
   for (const msg of messages) {
@@ -1690,9 +1726,13 @@ function buildAntigravityRequestBody(params) {
     maxOutputTokens = 65536
   } = params;
   const runtimeModelId = plan.runtimeModelId;
-  const sessionId = deriveSessionId(context, params.sessionId);
+  const transcript = collapseSystemMessages(context);
+  const systemPrompt = getCurrentSystemPrompt(transcript.messages);
+  const tools = convertTools(getCurrentTools(transcript.messages));
+  const conversation = withoutInitialSystemMessage(transcript.messages);
+  const sessionId = deriveSessionId(conversation, params.sessionId);
   const trajectoryId = params.trajectoryId && typeof params.trajectoryId === "string" && params.trajectoryId.trim().length > 0 ? params.trajectoryId.trim() : resolveSessionTrajectory(sessionId);
-  const contents = translateTurnTrace(context, runtimeModelId);
+  const contents = translateTurnTrace(conversation, runtimeModelId);
   const requestIndex = contents.filter(
     (c) => c.role === "model" && !c.parts?.some((p) => p.functionResponse)
   ).length;
@@ -1714,13 +1754,12 @@ function buildAntigravityRequestBody(params) {
   const request = {
     contents
   };
-  if (context.systemPrompt) {
+  if (systemPrompt) {
     request.systemInstruction = {
       role: "user",
-      parts: [{ text: context.systemPrompt }]
+      parts: [{ text: systemPrompt }]
     };
   }
-  const tools = convertTools(context.tools);
   if (tools) {
     request.tools = tools;
   }
@@ -1791,7 +1830,14 @@ function processLine(rawLine, state, onEvent) {
   } catch {
     return;
   }
+  const streamError = payload?.error ?? payload?.response?.error;
+  if (streamError) {
+    throw new Error(streamError.message ?? JSON.stringify(streamError));
+  }
   const response = payload.response || payload;
+  if (response.responseId && !state.responseId) {
+    state.responseId = response.responseId;
+  }
   if (response.usageMetadata) {
     const um = response.usageMetadata;
     const promptTokens = typeof um.promptTokenCount === "number" ? um.promptTokenCount : 0;
@@ -1808,8 +1854,14 @@ function processLine(rawLine, state, onEvent) {
   if (!candidate) return;
   if (candidate.finishReason === "MAX_TOKENS") {
     state.stopReason = "length";
-  } else if (typeof candidate.finishReason === "string" && candidate.finishReason !== "STOP") {
+    state.rawStopReason = "MAX_TOKENS";
+  } else if (candidate.finishReason === "STOP") {
+    if (state.stopReason === "pending") {
+      state.stopReason = "stop";
+    }
+  } else if (typeof candidate.finishReason === "string") {
     state.stopReason = "error";
+    state.rawStopReason = candidate.finishReason;
   }
   const parts = candidate.content?.parts || [];
   for (const part of parts) {
@@ -1838,7 +1890,9 @@ function processLine(rawLine, state, onEvent) {
       });
     } else if (part.functionCall) {
       closeOpenBlock(state, onEvent);
-      state.stopReason = "toolUse";
+      if (state.stopReason === "pending" || state.stopReason === "stop") {
+        state.stopReason = "toolUse";
+      }
       const block = {
         type: "toolCall",
         id: part.functionCall.id || `call_${state.content.length}`,
@@ -1886,33 +1940,43 @@ function finishStream(state, onEvent) {
   closeOpenBlock(state, onEvent);
   attachLoneSignature(state.content, state.lastThoughtSignature);
 }
-async function consumeAntigravityStream(streamBody, output, stream, modelCost) {
+var NON_OVERFLOW_PATTERN = /rate limit|too many requests|quota exceeded|resource_exhausted/i;
+var CONTEXT_OVERFLOW_PATTERN = /(?:(?:input|prompt) token count.*exceeds|prompt (?:is )?too long)/i;
+function normalizeOverflowError(rawMessage) {
+  if (rawMessage.includes("context_length_exceeded")) return rawMessage;
+  if (!NON_OVERFLOW_PATTERN.test(rawMessage) && CONTEXT_OVERFLOW_PATTERN.test(rawMessage)) {
+    return `context_length_exceeded: ${rawMessage}`;
+  }
+  return rawMessage;
+}
+async function consumeAntigravityStream(streamBody, output, stream, model, signal) {
   const state = {
     content: output.content,
     usage: { input: 0, output: 0, cacheRead: 0, reasoning: 0, total: 0 },
-    stopReason: "stop",
+    stopReason: "pending",
     openType: null,
     buffer: ""
   };
+  let lastTotalTokens = -1;
   const updateUsageAndCost = () => {
     output.usage.input = state.usage.input;
     output.usage.cacheRead = state.usage.cacheRead;
     output.usage.output = state.usage.output;
     output.usage.reasoning = state.usage.reasoning;
     output.usage.totalTokens = state.usage.total;
-    if (modelCost) {
-      const inputCost = output.usage.input * (modelCost.input || 0) / 1e6;
-      const outputCost = output.usage.output * (modelCost.output || 0) / 1e6;
-      const cacheCost = output.usage.cacheRead * (modelCost.cacheRead || 0) / 1e6;
-      output.usage.cost = {
-        input: inputCost,
-        output: outputCost,
-        cacheRead: cacheCost,
-        cacheWrite: 0,
-        total: inputCost + outputCost + cacheCost
-      };
+    if (state.usage.total !== lastTotalTokens) {
+      lastTotalTokens = state.usage.total;
+      if (model.cost) {
+        calculateCost(model, output.usage);
+      }
+    }
+    if (state.responseId) {
+      output.responseId ||= state.responseId;
     }
     output.stopReason = state.stopReason;
+    if (state.rawStopReason) {
+      output.rawStopReason = state.rawStopReason;
+    }
   };
   const dispatchEvent = (ev) => {
     updateUsageAndCost();
@@ -1920,13 +1984,20 @@ async function consumeAntigravityStream(streamBody, output, stream, modelCost) {
   };
   const reader = streamBody.getReader();
   const decoder = new TextDecoder();
-  while (true) {
-    const { done, value } = await reader.read();
-    if (done) break;
-    feedChunk(decoder.decode(value, { stream: true }), state, dispatchEvent);
+  try {
+    while (true) {
+      if (signal?.aborted) {
+        throw new Error("Request was aborted");
+      }
+      const { done, value } = await reader.read();
+      if (done) break;
+      feedChunk(decoder.decode(value, { stream: true }), state, dispatchEvent);
+    }
+    finishStream(state, dispatchEvent);
+    updateUsageAndCost();
+  } finally {
+    reader.releaseLock();
   }
-  finishStream(state, dispatchEvent);
-  updateUsageAndCost();
 }
 function streamAntigravity(model, context, options, catalog) {
   const stream = createAssistantMessageEventStream();
@@ -1954,22 +2025,60 @@ function streamAntigravity(model, context, options, catalog) {
       const effort = options?.reasoning;
       const trajectoryId = options && "trajectoryId" in options && typeof options.trajectoryId === "string" ? options.trajectoryId : void 0;
       const plan = catalog.resolvePlan(model.id, effort);
-      const requestBody = buildAntigravityRequestBody({
+      let requestBody = buildAntigravityRequestBody({
         projectId,
         plan,
         context,
         sessionId: options?.sessionId,
         trajectoryId,
-        maxOutputTokens: model.maxTokens,
+        maxOutputTokens: options?.maxTokens ?? model.maxTokens,
         toolChoice: options?.toolChoice
       });
-      const streamBody = await postAntigravityStream({
+      if (options?.onPayload) {
+        const replacement = await options.onPayload(requestBody, model);
+        if (replacement !== void 0) {
+          requestBody = replacement;
+        }
+      }
+      const endpoint = model.baseUrl || DEFAULT_ENDPOINT;
+      const mergedHeaders = {};
+      if (model.headers) {
+        for (const [k, v] of Object.entries(model.headers)) {
+          if (typeof v === "string") mergedHeaders[k] = v;
+        }
+      }
+      if (options?.headers) {
+        for (const [k, v] of Object.entries(options.headers)) {
+          if (typeof v === "string") mergedHeaders[k] = v;
+        }
+      }
+      const { response, stream: streamBody } = await postAntigravityStream({
         auth: token,
+        endpoint,
         path: "v1internal:streamGenerateContent?alt=sse",
+        headers: mergedHeaders,
         body: requestBody,
         signal: options?.signal
       });
-      await consumeAntigravityStream(streamBody, output, stream, model.cost);
+      if (options?.onResponse) {
+        const resHeaders = {};
+        response.headers.forEach((val, key) => {
+          resHeaders[key] = val;
+        });
+        await options.onResponse({ status: response.status, headers: resHeaders }, model);
+      }
+      await consumeAntigravityStream(streamBody, output, stream, model, options?.signal);
+      if (options?.signal?.aborted) {
+        throw new Error("Request was aborted");
+      }
+      if (output.stopReason === "pending") {
+        throw new Error("Provider stream ended without a stop reason");
+      }
+      if (output.stopReason === "error" || output.stopReason === "aborted") {
+        throw new Error(
+          output.rawStopReason ? `Provider stopped with: ${output.rawStopReason}` : output.errorMessage || "An unknown error occurred"
+        );
+      }
       const doneReason = output.stopReason === "toolUse" || output.stopReason === "length" ? output.stopReason : "stop";
       stream.push({
         type: "done",
@@ -1979,7 +2088,8 @@ function streamAntigravity(model, context, options, catalog) {
       stream.end();
     } catch (error) {
       output.stopReason = options?.signal?.aborted ? "aborted" : "error";
-      output.errorMessage = error instanceof Error ? error.message : String(error);
+      const rawMessage = error instanceof Error ? error.message : String(error);
+      output.errorMessage = normalizeOverflowError(rawMessage);
       stream.push({ type: "error", reason: output.stopReason, error: output });
       stream.end();
     }
@@ -2277,6 +2387,7 @@ function index_default(pi) {
     // Purely dynamic provider per Pi SDK architecture
     oauth: {
       name: PROVIDER_NAME,
+      isSubscription: true,
       login: loginAntigravity,
       refreshToken: refreshAntigravityToken,
       getApiKey