pi-provider-antigravity 0.14.0 → 0.15.0

package/README.md

@@ -39,7 +39,9 @@ pi -e npm:pi-provider-antigravity
 
 ### Legacy Pi Support
 
-For Pi 0.80.0 ~ 0.85.x, install the compatible `0.13.0` release:
+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

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,6 +1412,7 @@ var QuotaStatusCoordinator = class {
 
 // src/stream.ts
 import {
+  calculateCost,
   createAssistantMessageEventStream
 } from "@earendil-works/pi-ai";
 
@@ -1801,6 +1830,10 @@ 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;
@@ -1821,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) {
@@ -1851,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}`,
@@ -1899,36 +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();
@@ -1936,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();
@@ -1970,7 +2025,7 @@ 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,
@@ -1979,13 +2034,51 @@ function streamAntigravity(model, context, options, catalog) {
         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",
@@ -1995,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();
     }
@@ -2293,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