@benzsiangco/jarvis 1.0.0 → 1.1.0

package/src/providers/antigravity-loader.ts

@@ -0,0 +1,233 @@
+// Antigravity plugin loader adapter for Jarvis
+// This uses Jarvis's own account storage instead of the plugin's storage
+import { loadAccounts as loadJarvisAccounts, saveAccounts as saveJarvisAccounts, refreshAccountToken as jarvisRefreshToken } from './antigravity.js';
+import type { AntigravityAccount, AccountsConfig } from './antigravity.js';
+import { join } from 'path';
+import { homedir } from 'os';
+import crypto from 'crypto';
+
+// Import only the request transformation functions, not the full plugin
+import { prepareAntigravityRequest, transformAntigravityResponse } from 'opencode-antigravity-auth/dist/src/plugin/request.js';
+
+// Ensure accounts have required fields
+function ensureAccountFields(accounts: AccountsConfig): AccountsConfig {
+  const updated = { ...accounts };
+  updated.accounts = accounts.accounts.map((acc: any, index) => ({
+    id: acc.id || `acc_${Date.now()}_${index}`,
+    email: acc.email,
+    refreshToken: acc.refreshToken,
+    accessToken: acc.accessToken,
+    expiresAt: acc.expiresAt,
+    projectId: acc.projectId,
+    rateLimitedUntil: acc.rateLimitedUntil,
+    requestCount: acc.requestCount || 0,
+    lastUsed: acc.lastUsed || Date.now(),
+  }));
+  return updated;
+}
+
+// Get next available account with token refresh
+async function getAccountWithToken(): Promise<AntigravityAccount | null> {
+  const accountsConfig = ensureAccountFields(loadJarvisAccounts());
+  
+  if (accountsConfig.accounts.length === 0) {
+    return null;
+  }
+  
+  const account = accountsConfig.accounts[accountsConfig.activeIndex || 0];
+  
+  if (!account) {
+    return null;
+  }
+  
+  // If no access token or expired, refresh it
+  if (!account.accessToken || (account.expiresAt && account.expiresAt < Date.now() + 5 * 60 * 1000)) {
+    const refreshed = await jarvisRefreshToken(account);
+    if (refreshed) {
+      return refreshed;
+    }
+  }
+  
+  return account;
+}
+
+// Track usage
+function trackUsage(accountId: string) {
+  const accounts = loadJarvisAccounts();
+  const account = accounts.accounts.find(a => (a as any).id === accountId);
+  if (account) {
+    (account as any).requestCount = ((account as any).requestCount || 0) + 1;
+    (account as any).lastUsed = Date.now();
+    saveJarvisAccounts(accounts);
+  }
+}
+
+// Mark rate limited
+function markRateLimited(accountId: string, durationMs: number) {
+  const accounts = loadJarvisAccounts();
+  const account = accounts.accounts.find(a => (a as any).id === accountId);
+  if (account) {
+    (account as any).rateLimitedUntil = Date.now() + durationMs;
+    saveJarvisAccounts(accounts);
+  }
+}
+
+// Create a custom fetch function using Jarvis's accounts
+export async function createJarvisAntigravityFetch() {
+  // Generate a session ID for signature caching
+  const sessionId = `jarvis-${crypto.randomUUID()}`;
+  
+  return async (url: string | URL | Request, init?: RequestInit): Promise<Response> => {
+    // Get current account from Jarvis's storage
+    let account = await getAccountWithToken();
+    
+    if (!account || !account.accessToken) {
+      throw new Error('No Antigravity account available. Run "jarvis auth login"');
+    }
+
+    const urlString = typeof url === 'string' ? url : url instanceof URL ? url.toString() : url.url;
+    
+    // Only process Google Generative Language API requests
+    if (!urlString.includes('generativelanguage.googleapis.com')) {
+      return fetch(url, init);
+    }
+
+    // Force thinking recovery for all models that may have function calls
+    // The thought_signature error happens when function calls are made with thinking models
+    // We need to enable forceThinkingRecovery to properly handle the thought signatures
+    const bodyStr = typeof init?.body === 'string' ? init.body : '';
+    const hasToolsOrFunctions = bodyStr.includes('"tools"') || 
+                                 bodyStr.includes('"functions"') || 
+                                 bodyStr.includes('"function_declarations"') ||
+                                 bodyStr.includes('"functionDeclarations"');
+    
+    // Gemini 3 models and any model with tools need thinking recovery
+    const isGemini3 = urlString.includes('gemini-3') || urlString.includes('gemini-2');
+    const needsThinkingRecovery = hasToolsOrFunctions && isGemini3;
+    
+    // Extract model name from URL
+    const modelMatch = urlString.match(/models\/([^:\/]+)/);
+    const requestedModel = modelMatch ? modelMatch[1] : 'unknown';
+    const endpoint = 'https://daily-cloudcode-pa.sandbox.googleapis.com';
+    const projectId = account.projectId || 'jarvis-default-project';
+    
+    // Use plugin's request preparation
+    const prepared = prepareAntigravityRequest(
+      urlString,
+      init,
+      account.accessToken!,
+      projectId,
+      endpoint,
+      'antigravity',
+      needsThinkingRecovery, // Enable forceThinkingRecovery when tools are used
+      {} // Options
+    );
+    
+    const response = await fetch(prepared.request, prepared.init);
+    
+    if (!response.ok) {
+      const errorText = await response.text();
+      let errorMessage = `Antigravity error ${response.status}`;
+      try {
+        const errorJson = JSON.parse(errorText);
+        errorMessage = errorJson.error?.message || errorMessage;
+      } catch (e) {
+        errorMessage = errorText || errorMessage;
+      }
+      throw new Error(errorMessage);
+    }
+    
+    // Handle rate limiting
+    if (response.status === 429) {
+      const retryAfter = response.headers.get('Retry-After');
+      const duration = retryAfter ? parseInt(retryAfter) * 1000 : 60000;
+      markRateLimited(account.id, duration);
+      
+      // Try to get another account
+      const nextAccount = await getAccountWithToken();
+      if (nextAccount && nextAccount.id !== account.id && nextAccount.accessToken) {
+        const nextPrepared = prepareAntigravityRequest(
+          urlString,
+          init,
+          nextAccount.accessToken,
+          nextAccount.projectId || 'jarvis-default-project',
+          endpoint,
+          'antigravity',
+          needsThinkingRecovery,
+          {}
+        );
+        const retryResponse = await fetch(nextPrepared.request, nextPrepared.init);
+        return transformAntigravityResponse(
+          retryResponse,
+          prepared.streaming,
+          null, // debugContext
+          requestedModel,
+          projectId,
+          endpoint,
+          requestedModel, // effectiveModel
+          sessionId
+        );
+      }
+    }
+    
+    // Handle auth errors
+    if (response.status === 401 || response.status === 403) {
+      const refreshed = await jarvisRefreshToken(account);
+      if (refreshed) {
+        const refreshedPrepared = prepareAntigravityRequest(
+          urlString,
+          init,
+          refreshed.accessToken!,
+          refreshed.projectId || 'jarvis-default-project',
+          endpoint,
+          'antigravity',
+          needsThinkingRecovery,
+          {}
+        );
+        const retryResponse = await fetch(refreshedPrepared.request, refreshedPrepared.init);
+        return transformAntigravityResponse(
+          retryResponse,
+          prepared.streaming,
+          null,
+          requestedModel,
+          projectId,
+          endpoint,
+          requestedModel,
+          sessionId
+        );
+      }
+    }
+    
+    // Track usage
+    if (response.ok) {
+      trackUsage(account.id);
+    }
+    
+    // Transform the response using plugin's logic with full parameters
+    return transformAntigravityResponse(
+      response, 
+      prepared.streaming, 
+      null, // debugContext
+      requestedModel,
+      projectId,
+      endpoint,
+      requestedModel, // effectiveModel
+      sessionId
+    );
+  };
+}
+
+// Initialize with Jarvis's accounts only
+export async function initializeAntigravityLoader() {
+  try {
+    // Ensure accounts have proper format
+    const accounts = ensureAccountFields(loadJarvisAccounts());
+    saveJarvisAccounts(accounts);
+    
+    // Return the custom fetch function that uses Jarvis accounts
+    return await createJarvisAntigravityFetch();
+  } catch (error) {
+    console.error('Failed to initialize Antigravity loader:', error);
+    throw error;
+  }
+}