feat(cli): add support for custom base URL in Zed integration

packages/cli/src/zed-integration/zedIntegration.test.ts

@@ -62,7 +62,31 @@ vi.mock('node:path', async (importOriginal) => {
   };
 });
 
-// Mock ReadManyFilesTool
+vi.mock('../utils/sessionUtils.js', () => ({
+  SessionSelector: vi.fn().mockImplementation(() => ({
+    resolveSession: vi.fn().mockResolvedValue({
+      sessionData: {
+        messages: [
+          { type: 'user', content: [{ text: 'Hello' }] },
+          {
+            type: 'gemini',
+            content: [{ text: 'Hi' }],
+            toolCalls: [
+              {
+                id: 'call-1',
+                name: 'test_tool',
+                status: 'success',
+                resultDisplay: 'Tool output',
+              },
+            ],
+          },
+        ],
+      },
+      sessionPath: '/path/to/session.json',
+    }),
+  })),
+}));
+
 vi.mock(
   '@google/gemini-cli-core',
   async (
@@ -84,6 +108,13 @@ vi.mock(
       })),
       logToolCall: vi.fn(),
       isWithinRoot: vi.fn().mockReturnValue(true),
+      convertSessionToClientHistory: vi.fn().mockReturnValue([]),
+      partListUnionToString: vi.fn().mockImplementation((content) => {
+        if (Array.isArray(content)) {
+          return content.map((p) => p.text || '').join('');
+        }
+        return '';
+      }),
       LlmRole: {
         MAIN: 'main',
         SUBAGENT: 'subagent',
@@ -223,6 +254,27 @@ describe('GeminiAgent', () => {
     );
   });
 
+  it('should authenticate correctly with gemini-custom-url and base-url in _meta', async () => {
+    await agent.authenticate({
+      methodId: 'gemini-custom-url',
+      _meta: {
+        'api-key': 'test-api-key',
+        'base-url': 'https://custom.api.endpoint',
+      },
+    } as unknown as acp.AuthenticateRequest);
+
+    expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
+      AuthType.USE_GEMINI,
+      'test-api-key',
+      'https://custom.api.endpoint',
+    );
+    expect(mockSettings.setValue).toHaveBeenCalledWith(
+      SettingScope.User,
+      'security.auth.selectedType',
+      AuthType.USE_GEMINI,
+    );
+  });
+
   it('should create a new session', async () => {
     mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({
       apiKey: 'test-key',
@@ -438,6 +490,17 @@ describe('GeminiAgent', () => {
       }),
     ).rejects.toThrow('Session not found: unknown');
   });
+
+  it('should load an existing session', async () => {
+    const response = await agent.loadSession({
+      sessionId: 'test-session-id',
+      cwd: '/tmp',
+      mcpServers: [],
+    });
+
+    expect(mockConfig.getGeminiClient).toHaveBeenCalled();
+    expect(response.modes).toBeDefined();
+  });
 });
 
 describe('Session', () => {

packages/cli/src/zed-integration/zedIntegration.ts

@@ -87,6 +87,7 @@ export class GeminiAgent {
   private sessions: Map<string, Session> = new Map();
   private clientCapabilities: acp.ClientCapabilities | undefined;
   private apiKey: string | undefined;
+  private baseUrl: string | undefined;
 
   constructor(
     private config: Config,
@@ -115,6 +116,17 @@ export class GeminiAgent {
           },
         },
       },
+      {
+        id: 'gemini-custom-url',
+        name: 'Gemini API (Custom URL)',
+        description: 'Use an API key and custom base URL',
+        _meta: {
+          'api-key': {
+            provider: 'google',
+          },
+          'base-url': {},
+        },
+      },
       {
         id: AuthType.USE_VERTEX_AI,
         name: 'Vertex AI',
@@ -148,8 +160,13 @@ export class GeminiAgent {
   }
 
   async authenticate(req: acp.AuthenticateRequest): Promise<void> {
-    const { methodId } = req;
-    const method = z.nativeEnum(AuthType).parse(methodId);
+    const methodId = req.methodId;
+    let method = AuthType.USE_GEMINI;
+    if (methodId === 'gemini-custom-url') {
+      method = AuthType.USE_GEMINI;
+    } else {
+      method = z.nativeEnum(AuthType).parse(methodId);
+    }
     const selectedAuthType = this.settings.merged.security.auth.selectedType;
 
     // Only clear credentials when switching to a different auth method
@@ -160,6 +177,8 @@ export class GeminiAgent {
     const meta = hasMeta(req) ? req._meta : undefined;
     const apiKey =
       typeof meta?.['api-key'] === 'string' ? meta['api-key'] : undefined;
+    const baseUrl =
+      typeof meta?.['base-url'] === 'string' ? meta['base-url'] : undefined;
 
     // Refresh auth with the requested method
     // This will reuse existing credentials if they're valid,
@@ -168,7 +187,14 @@ export class GeminiAgent {
       if (apiKey) {
         this.apiKey = apiKey;
       }
-      await this.config.refreshAuth(method, apiKey ?? this.apiKey);
+      if (baseUrl) {
+        this.baseUrl = baseUrl;
+      }
+      await this.config.refreshAuth(
+        method,
+        apiKey ?? this.apiKey,
+        baseUrl ?? this.baseUrl,
+      );
     } catch (e) {
       throw new acp.RequestError(-32000, getAcpErrorMessage(e));
     }
@@ -198,7 +224,7 @@ export class GeminiAgent {
     let isAuthenticated = false;
     let authErrorMessage = '';
     try {
-      await config.refreshAuth(authType, this.apiKey);
+      await config.refreshAuth(authType, this.apiKey, this.baseUrl);
       isAuthenticated = true;
 
       // Extra validation for Gemini API key

packages/core/src/config/config.ts

@@ -1126,7 +1126,7 @@ export class Config {
     return this.contentGenerator;
   }
 
-  async refreshAuth(authMethod: AuthType, apiKey?: string) {
+  async refreshAuth(authMethod: AuthType, apiKey?: string, baseUrl?: string) {
     // Reset availability service when switching auth
     this.modelAvailabilityService.reset();
 
@@ -1153,6 +1153,7 @@ export class Config {
       this,
       authMethod,
       apiKey,
+      baseUrl,
     );
     this.contentGenerator = await createContentGenerator(
       newContentGeneratorConfig,

packages/core/src/core/contentGenerator.ts

@@ -82,6 +82,7 @@ export function getAuthTypeFromEnv(): AuthType | undefined {
 
 export type ContentGeneratorConfig = {
   apiKey?: string;
+  baseUrl?: string;
   vertexai?: boolean;
   authType?: AuthType;
   proxy?: string;
@@ -91,6 +92,7 @@ export async function createContentGeneratorConfig(
   config: Config,
   authType: AuthType | undefined,
   apiKey?: string,
+  baseUrl?: string,
 ): Promise<ContentGeneratorConfig> {
   const geminiApiKey =
     apiKey ||
@@ -119,6 +121,7 @@ export async function createContentGeneratorConfig(
 
   if (authType === AuthType.USE_GEMINI && geminiApiKey) {
     contentGeneratorConfig.apiKey = geminiApiKey;
+    contentGeneratorConfig.baseUrl = baseUrl;
     contentGeneratorConfig.vertexai = false;
 
     return contentGeneratorConfig;
@@ -206,7 +209,11 @@ export async function createContentGenerator(
           'x-gemini-api-privileged-user-id': `${installationId}`,
         };
       }
-      const httpOptions = { headers };
+      const httpOptions: { headers: Record<string, string>; baseUrl?: string } =
+        { headers };
+      if (config.baseUrl) {
+        httpOptions.baseUrl = config.baseUrl;
+      }
 
       const googleGenAI = new GoogleGenAI({
         apiKey: config.apiKey === '' ? undefined : config.apiKey,