@link-assistant/agent 0.6.3 → 0.8.1

package/src/provider/google-cloudcode.ts

@@ -0,0 +1,384 @@
+/**
+ * Google Cloud Code API Client
+ *
+ * This module provides a client for Google's Cloud Code API (cloudcode-pa.googleapis.com),
+ * which is used by the official Gemini CLI for OAuth-authenticated requests.
+ *
+ * The Cloud Code API:
+ * 1. Accepts `cloud-platform` OAuth scope (unlike generativelanguage.googleapis.com)
+ * 2. Handles subscription tier validation (FREE, STANDARD, etc.)
+ * 3. Proxies requests to the Generative Language API internally
+ *
+ * @see https://github.com/google-gemini/gemini-cli/blob/main/packages/core/src/code_assist/server.ts
+ * @see https://github.com/link-assistant/agent/issues/100
+ */
+
+import { Log } from '../util/log';
+import { Auth } from '../auth';
+
+const log = Log.create({ service: 'google-cloudcode' });
+
+// Cloud Code API endpoints (from gemini-cli)
+const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com';
+const CODE_ASSIST_API_VERSION = 'v1internal';
+
+// Google OAuth endpoints
+const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token';
+const GOOGLE_OAUTH_CLIENT_ID =
+  '681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com';
+const GOOGLE_OAUTH_CLIENT_SECRET = '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';
+
+/**
+ * User tier from Cloud Code API
+ */
+export enum UserTierId {
+  FREE = 'FREE',
+  STANDARD = 'STANDARD',
+  LEGACY = 'LEGACY',
+}
+
+/**
+ * Cloud Code API request format
+ */
+interface CloudCodeRequest {
+  model: string;
+  project?: string;
+  user_prompt_id?: string;
+  request: {
+    contents: Array<{
+      role: string;
+      parts: Array<{ text?: string; [key: string]: unknown }>;
+    }>;
+    systemInstruction?: {
+      role: string;
+      parts: Array<{ text?: string }>;
+    };
+    tools?: unknown[];
+    toolConfig?: unknown;
+    generationConfig?: {
+      temperature?: number;
+      topP?: number;
+      topK?: number;
+      maxOutputTokens?: number;
+      candidateCount?: number;
+      stopSequences?: string[];
+      responseMimeType?: string;
+      responseSchema?: unknown;
+      thinkingConfig?: {
+        thinkingBudget?: number;
+      };
+    };
+    safetySettings?: unknown[];
+  };
+}
+
+/**
+ * Cloud Code API response format
+ */
+interface CloudCodeResponse {
+  response: {
+    candidates: Array<{
+      content: {
+        role: string;
+        parts: Array<{
+          text?: string;
+          thought?: boolean;
+          functionCall?: unknown;
+          functionResponse?: unknown;
+        }>;
+      };
+      finishReason?: string;
+      safetyRatings?: unknown[];
+    }>;
+    usageMetadata?: {
+      promptTokenCount?: number;
+      candidatesTokenCount?: number;
+      totalTokenCount?: number;
+      thoughtsTokenCount?: number;
+      cachedContentTokenCount?: number;
+    };
+    modelVersion?: string;
+  };
+  traceId?: string;
+}
+
+/**
+ * Load Code Assist response (for user setup)
+ */
+interface LoadCodeAssistResponse {
+  cloudaicompanionProject?: string;
+  currentTier?: {
+    id: UserTierId;
+    name: string;
+    description: string;
+  };
+  allowedTiers?: Array<{
+    id: UserTierId;
+    name: string;
+    description: string;
+    isDefault?: boolean;
+    userDefinedCloudaicompanionProject?: boolean;
+  }>;
+}
+
+/**
+ * Onboard user response (for user setup)
+ */
+interface OnboardUserResponse {
+  done: boolean;
+  response?: {
+    cloudaicompanionProject?: {
+      id: string;
+    };
+  };
+}
+
+/**
+ * Google Cloud Code API client
+ *
+ * This client implements the same API used by the official Gemini CLI.
+ * It wraps OAuth authentication and provides methods for:
+ * - Setting up user (onboarding to FREE/STANDARD tier)
+ * - Generating content (streaming and non-streaming)
+ * - Counting tokens
+ */
+export class CloudCodeClient {
+  private accessToken: string;
+  private refreshToken: string;
+  private tokenExpiry: number;
+  private projectId?: string;
+  private userTier?: UserTierId;
+
+  constructor(
+    private auth: {
+      access: string;
+      refresh: string;
+      expires: number;
+    },
+    projectId?: string
+  ) {
+    this.accessToken = auth.access;
+    this.refreshToken = auth.refresh;
+    this.tokenExpiry = auth.expires;
+    this.projectId = projectId;
+  }
+
+  /**
+   * Refresh the OAuth access token if expired
+   */
+  private async ensureValidToken(): Promise<void> {
+    const FIVE_MIN_MS = 5 * 60 * 1000;
+    if (this.tokenExpiry > Date.now() + FIVE_MIN_MS) {
+      return; // Token is still valid
+    }
+
+    log.info(() => ({
+      message: 'refreshing google oauth token for cloud code',
+    }));
+
+    const response = await fetch(GOOGLE_TOKEN_URL, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/x-www-form-urlencoded',
+      },
+      body: new URLSearchParams({
+        client_id: GOOGLE_OAUTH_CLIENT_ID,
+        client_secret: GOOGLE_OAUTH_CLIENT_SECRET,
+        refresh_token: this.refreshToken,
+        grant_type: 'refresh_token',
+      }),
+    });
+
+    if (!response.ok) {
+      throw new Error(`Token refresh failed: ${response.status}`);
+    }
+
+    const json = await response.json();
+    this.accessToken = json.access_token;
+    this.tokenExpiry = Date.now() + json.expires_in * 1000;
+
+    // Update stored auth
+    await Auth.set('google', {
+      type: 'oauth',
+      refresh: this.refreshToken,
+      access: this.accessToken,
+      expires: this.tokenExpiry,
+    });
+  }
+
+  /**
+   * Make an authenticated request to the Cloud Code API
+   */
+  private async request<T>(
+    method: string,
+    body: unknown,
+    options: { stream?: boolean } = {}
+  ): Promise<T | Response> {
+    await this.ensureValidToken();
+
+    const url = `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`;
+
+    const response = await fetch(url, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${this.accessToken}`,
+        'x-goog-api-client': 'agent/0.6.3',
+      },
+      body: JSON.stringify(body),
+    });
+
+    if (!response.ok) {
+      const errorText = await response.text();
+      log.error(() => ({
+        message: 'cloud code api error',
+        status: response.status,
+        error: errorText,
+      }));
+      throw new Error(
+        `Cloud Code API error: ${response.status} - ${errorText}`
+      );
+    }
+
+    if (options.stream) {
+      return response;
+    }
+
+    return response.json() as Promise<T>;
+  }
+
+  /**
+   * Load code assist to check user tier and project
+   */
+  async loadCodeAssist(): Promise<LoadCodeAssistResponse> {
+    const body = {
+      cloudaicompanionProject: this.projectId,
+      metadata: {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI',
+        duetProject: this.projectId,
+      },
+    };
+
+    return this.request<LoadCodeAssistResponse>('loadCodeAssist', body);
+  }
+
+  /**
+   * Onboard user to a tier (FREE or STANDARD)
+   */
+  async onboardUser(tierId: UserTierId): Promise<OnboardUserResponse> {
+    const body = {
+      tierId,
+      cloudaicompanionProject:
+        tierId === UserTierId.FREE ? undefined : this.projectId,
+      metadata: {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI',
+        ...(tierId !== UserTierId.FREE && { duetProject: this.projectId }),
+      },
+    };
+
+    return this.request<OnboardUserResponse>('onboardUser', body);
+  }
+
+  /**
+   * Setup user - check tier and onboard if necessary
+   */
+  async setupUser(): Promise<{ projectId?: string; userTier: UserTierId }> {
+    const loadRes = await this.loadCodeAssist();
+
+    // If user already has a tier, use it
+    if (loadRes.currentTier) {
+      this.userTier = loadRes.currentTier.id;
+      this.projectId = loadRes.cloudaicompanionProject || this.projectId;
+      return {
+        projectId: this.projectId,
+        userTier: this.userTier,
+      };
+    }
+
+    // Find the default tier to onboard to
+    let targetTier = UserTierId.FREE;
+    for (const tier of loadRes.allowedTiers || []) {
+      if (tier.isDefault) {
+        targetTier = tier.id;
+        break;
+      }
+    }
+
+    log.info(() => ({ message: 'onboarding user', tier: targetTier }));
+
+    // Poll onboardUser until done
+    let lroRes = await this.onboardUser(targetTier);
+    while (!lroRes.done) {
+      await new Promise((resolve) => setTimeout(resolve, 5000));
+      lroRes = await this.onboardUser(targetTier);
+    }
+
+    this.userTier = targetTier;
+    if (lroRes.response?.cloudaicompanionProject?.id) {
+      this.projectId = lroRes.response.cloudaicompanionProject.id;
+    }
+
+    return {
+      projectId: this.projectId,
+      userTier: this.userTier,
+    };
+  }
+
+  /**
+   * Generate content (non-streaming)
+   */
+  async generateContent(request: CloudCodeRequest): Promise<CloudCodeResponse> {
+    return this.request<CloudCodeResponse>('generateContent', request);
+  }
+
+  /**
+   * Generate content (streaming)
+   */
+  async generateContentStream(request: CloudCodeRequest): Promise<Response> {
+    return this.request<Response>('streamGenerateContent', request, {
+      stream: true,
+    }) as Promise<Response>;
+  }
+
+  /**
+   * Get project ID (may be set during setup)
+   */
+  getProjectId(): string | undefined {
+    return this.projectId;
+  }
+
+  /**
+   * Get user tier
+   */
+  getUserTier(): UserTierId | undefined {
+    return this.userTier;
+  }
+}
+
+/**
+ * Create a Cloud Code client from stored Google OAuth credentials
+ */
+export async function createCloudCodeClient(): Promise<CloudCodeClient | null> {
+  const auth = await Auth.get('google');
+  if (!auth || auth.type !== 'oauth') {
+    return null;
+  }
+
+  const projectId =
+    process.env['GOOGLE_CLOUD_PROJECT'] ||
+    process.env['GOOGLE_CLOUD_PROJECT_ID'];
+
+  return new CloudCodeClient(auth, projectId);
+}
+
+/**
+ * Check if Cloud Code API is available (user has OAuth credentials)
+ */
+export async function isCloudCodeAvailable(): Promise<boolean> {
+  const auth = await Auth.get('google');
+  return auth?.type === 'oauth';
+}

package/src/session/message-v2.ts

@@ -42,6 +42,22 @@ export namespace MessageV2 {
   );
   export type APIError = z.infer<typeof APIError.Schema>;
 
+  /**
+   * Socket connection error - typically caused by Bun's 10-second idle timeout
+   * in Bun.serve() context. These errors are transient and should be retried.
+   * See: https://github.com/oven-sh/bun/issues/14439
+   */
+  export const SocketConnectionError = NamedError.create(
+    'SocketConnectionError',
+    z.object({
+      message: z.string(),
+      isRetryable: z.literal(true),
+    })
+  );
+  export type SocketConnectionError = z.infer<
+    typeof SocketConnectionError.Schema
+  >;
+
   const PartBase = z.object({
     id: z.string(),
     sessionID: z.string(),
@@ -787,11 +803,24 @@ export namespace MessageV2 {
           },
           { cause: e }
         ).toObject();
-      case e instanceof Error:
+      case e instanceof Error: {
+        const message = e.message || e.toString();
+        // Detect Bun socket connection errors (known Bun issue with 10s idle timeout)
+        // See: https://github.com/oven-sh/bun/issues/14439
+        const isSocketError =
+          message.includes('socket connection was closed') ||
+          message.includes('closed unexpectedly');
+        if (isSocketError) {
+          return new MessageV2.SocketConnectionError(
+            { message, isRetryable: true },
+            { cause: e }
+          ).toObject();
+        }
         return new NamedError.Unknown(
           { message: e.toString() },
           { cause: e }
         ).toObject();
+      }
       default:
         return new NamedError.Unknown(
           { message: JSON.stringify(e) },

package/src/session/processor.ts

@@ -314,9 +314,28 @@ export namespace SessionProcessor {
             const error = MessageV2.fromError(e, {
               providerID: input.providerID,
             });
-            if (error?.name === 'APIError' && error.data.isRetryable) {
+
+            // Check if error is retryable (APIError or SocketConnectionError)
+            const isRetryableAPIError =
+              error?.name === 'APIError' && error.data.isRetryable;
+            const isRetryableSocketError =
+              error?.name === 'SocketConnectionError' &&
+              error.data.isRetryable &&
+              attempt < SessionRetry.SOCKET_ERROR_MAX_RETRIES;
+
+            if (isRetryableAPIError || isRetryableSocketError) {
               attempt++;
-              const delay = SessionRetry.delay(error, attempt);
+              // Use socket-specific delay for socket errors
+              const delay =
+                error?.name === 'SocketConnectionError'
+                  ? SessionRetry.socketErrorDelay(attempt)
+                  : SessionRetry.delay(error, attempt);
+              log.info(() => ({
+                message: 'retrying',
+                errorType: error?.name,
+                attempt,
+                delay,
+              }));
               SessionStatus.set(input.sessionID, {
                 type: 'retry',
                 attempt,

package/src/session/prompt.ts

@@ -34,6 +34,7 @@ import { mergeDeep, pipe } from 'remeda';
 import { ToolRegistry } from '../tool/registry';
 import { Wildcard } from '../util/wildcard';
 import { MCP } from '../mcp';
+import { withTimeout } from '../util/timeout';
 import { ReadTool } from '../tool/read';
 import { ListTool } from '../tool/ls';
 import { FileTime } from '../file/time';
@@ -877,7 +878,28 @@ export namespace SessionPrompt {
       const execute = item.execute;
       if (!execute) continue;
       item.execute = async (args, opts) => {
-        const result = await execute(args, opts);
+        // Get timeout for this specific tool
+        const timeout = await MCP.getToolTimeout(key);
+
+        // Wrap the execute call with timeout to prevent indefinite hangs
+        let result;
+        try {
+          result = await withTimeout(execute(args, opts), timeout);
+        } catch (error) {
+          // Check if it's a timeout error
+          if (
+            error instanceof Error &&
+            error.message.includes('timed out after')
+          ) {
+            const timeoutSec = Math.round(timeout / 1000);
+            throw new Error(
+              `MCP tool "${key}" timed out after ${timeoutSec} seconds. ` +
+                `The operation did not complete within the configured timeout. ` +
+                `You can increase the timeout in the MCP server configuration using tool_call_timeout or tool_timeouts.`
+            );
+          }
+          throw error;
+        }
 
         const textParts: string[] = [];
         const attachments: MessageV2.FilePart[] = [];

package/src/session/retry.ts

@@ -6,6 +6,13 @@ export namespace SessionRetry {
   export const RETRY_BACKOFF_FACTOR = 2;
   export const RETRY_MAX_DELAY_NO_HEADERS = 30_000; // 30 seconds
 
+  // Socket connection error retry configuration
+  // Bun's fetch() has a known 10-second idle timeout issue
+  // See: https://github.com/oven-sh/bun/issues/14439
+  export const SOCKET_ERROR_MAX_RETRIES = 3;
+  export const SOCKET_ERROR_INITIAL_DELAY = 1000; // 1 second
+  export const SOCKET_ERROR_BACKOFF_FACTOR = 2;
+
   export async function sleep(ms: number, signal: AbortSignal): Promise<void> {
     return new Promise((resolve, reject) => {
       const timeout = setTimeout(resolve, ms);
@@ -53,4 +60,15 @@ export namespace SessionRetry {
       RETRY_MAX_DELAY_NO_HEADERS
     );
   }
+
+  /**
+   * Calculate delay for socket connection error retries.
+   * Uses exponential backoff: 1s, 2s, 4s, etc.
+   */
+  export function socketErrorDelay(attempt: number): number {
+    return (
+      SOCKET_ERROR_INITIAL_DELAY *
+      Math.pow(SOCKET_ERROR_BACKOFF_FACTOR, attempt - 1)
+    );
+  }
 }