Merge branch 'main' into acp-base-url

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

@@ -64,6 +64,31 @@ vi.mock('node:path', async (importOriginal) => {
   };
 });
 
+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('../ui/commands/memoryCommand.js', () => ({
   memoryCommand: {
     name: 'memory',
@@ -112,6 +137,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',
@@ -208,7 +240,7 @@ describe('GeminiAgent', () => {
     });
 
     expect(response.protocolVersion).toBe(acp.PROTOCOL_VERSION);
-    expect(response.authMethods).toHaveLength(3);
+    expect(response.authMethods).toHaveLength(4);
     const geminiAuth = response.authMethods?.find(
       (m) => m.id === AuthType.USE_GEMINI,
     );
@@ -228,6 +260,8 @@ describe('GeminiAgent', () => {
     expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
       AuthType.LOGIN_WITH_GOOGLE,
       undefined,
+      undefined,
+      undefined,
     );
     expect(mockSettings.setValue).toHaveBeenCalledWith(
       SettingScope.User,
@@ -247,6 +281,60 @@ describe('GeminiAgent', () => {
     expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
       AuthType.USE_GEMINI,
       'test-api-key',
+      undefined,
+      undefined,
+    );
+    expect(mockSettings.setValue).toHaveBeenCalledWith(
+      SettingScope.User,
+      'security.auth.selectedType',
+      AuthType.USE_GEMINI,
+    );
+  });
+
+  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',
+      undefined,
+    );
+    expect(mockSettings.setValue).toHaveBeenCalledWith(
+      SettingScope.User,
+      'security.auth.selectedType',
+      AuthType.USE_GEMINI,
+    );
+  });
+
+  it('should authenticate correctly with gateway method', async () => {
+    await agent.authenticate({
+      methodId: 'gateway',
+      _meta: {
+        gateway: {
+          baseUrl: 'https://gateway.example.com',
+          headers: {
+            Authorization: 'Bearer test-token',
+            'X-Custom-Header': 'custom-value',
+          },
+        },
+      },
+    } as unknown as acp.AuthenticateRequest);
+
+    expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
+      AuthType.USE_GEMINI,
+      undefined,
+      'https://gateway.example.com',
+      {
+        Authorization: 'Bearer test-token',
+        'X-Custom-Header': 'custom-value',
+      },
     );
     expect(mockSettings.setValue).toHaveBeenCalledWith(
       SettingScope.User,
@@ -524,6 +612,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();
+  });
+
   it('should delegate setModel to session (unstable)', async () => {
     await agent.newSession({ cwd: '/tmp', mcpServers: [] });
     const session = (

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

@@ -98,6 +98,8 @@ export class GeminiAgent {
   private sessions: Map<string, Session> = new Map();
   private clientCapabilities: acp.ClientCapabilities | undefined;
   private apiKey: string | undefined;
+  private baseUrl: string | undefined;
+  private headers: Record<string, string> | undefined;
 
   constructor(
     private config: Config,
@@ -126,6 +128,17 @@ export class GeminiAgent {
           },
         },
       },
+      {
+        id: 'gateway',
+        name: 'AI API Gateway',
+        description: 'Use a custom AI API Gateway',
+        _meta: {
+          gateway: {
+            protocol: 'google',
+            restartRequired: 'false',
+          },
+        },
+      },
       {
         id: AuthType.USE_VERTEX_AI,
         name: 'Vertex AI',
@@ -159,8 +172,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' || methodId === 'gateway') {
+      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
@@ -171,6 +189,28 @@ export class GeminiAgent {
     const meta = hasMeta(req) ? req._meta : undefined;
     const apiKey =
       typeof meta?.['api-key'] === 'string' ? meta['api-key'] : undefined;
+    let baseUrl =
+      typeof meta?.['base-url'] === 'string' ? meta['base-url'] : undefined;
+    let headers: Record<string, string> | undefined;
+
+    if (methodId === 'gateway') {
+      method = AuthType.USE_GEMINI;
+      // Gateway specific handling
+      const gatewaySchema = z
+        .object({
+          baseUrl: z.string().optional(),
+          headers: z.record(z.string()).optional(),
+        })
+        .optional();
+      const gatewayParams = gatewaySchema.parse(meta?.['gateway']);
+
+      if (gatewayParams?.baseUrl) {
+        baseUrl = gatewayParams.baseUrl;
+      }
+      if (gatewayParams?.headers) {
+        headers = gatewayParams.headers;
+      }
+    }
 
     // Refresh auth with the requested method
     // This will reuse existing credentials if they're valid,
@@ -179,7 +219,18 @@ export class GeminiAgent {
       if (apiKey) {
         this.apiKey = apiKey;
       }
-      await this.config.refreshAuth(method, apiKey ?? this.apiKey);
+      if (baseUrl) {
+        this.baseUrl = baseUrl;
+      }
+      if (headers) {
+        this.headers = headers;
+      }
+      await this.config.refreshAuth(
+        method,
+        apiKey ?? this.apiKey,
+        baseUrl ?? this.baseUrl,
+        headers,
+      );
     } catch (e) {
       throw new acp.RequestError(-32000, getAcpErrorMessage(e));
     }
@@ -209,7 +260,12 @@ export class GeminiAgent {
     let isAuthenticated = false;
     let authErrorMessage = '';
     try {
-      await config.refreshAuth(authType, this.apiKey);
+      await config.refreshAuth(
+        authType,
+        this.apiKey,
+        this.baseUrl,
+        this.headers,
+      );
       isAuthenticated = true;
 
       // Extra validation for Gemini API key
@@ -371,7 +427,12 @@ export class GeminiAgent {
     // This satisfies the security requirement to verify the user before executing
     // potentially unsafe server definitions.
     try {
-      await config.refreshAuth(selectedAuthType, this.apiKey);
+      await config.refreshAuth(
+        selectedAuthType,
+        this.apiKey,
+        this.baseUrl,
+        this.headers,
+      );
     } catch (e) {
       debugLogger.error(`Authentication failed: ${e}`);
       throw acp.RequestError.authRequired();

packages/core/src/config/config.test.ts

@@ -512,6 +512,8 @@ describe('Server Config (config.ts)', () => {
         config,
         authType,
         undefined,
+        undefined,
+        undefined,
       );
       // Verify that contentGeneratorConfig is updated
       expect(config.getContentGeneratorConfig()).toEqual(mockContentConfig);

packages/core/src/config/config.ts

@@ -1207,7 +1207,12 @@ export class Config implements McpContext {
     return this.contentGenerator;
   }
 
-  async refreshAuth(authMethod: AuthType, apiKey?: string) {
+  async refreshAuth(
+    authMethod: AuthType,
+    apiKey?: string,
+    baseUrl?: string,
+    headers?: Record<string, string>,
+  ) {
     // Reset availability service when switching auth
     this.modelAvailabilityService.reset();
 
@@ -1234,6 +1239,8 @@ export class Config implements McpContext {
       this,
       authMethod,
       apiKey,
+      baseUrl,
+      headers,
     );
     this.contentGenerator = await createContentGenerator(
       newContentGeneratorConfig,

packages/core/src/core/contentGenerator.test.ts

@@ -174,7 +174,6 @@ describe('createContentGenerator', () => {
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
           'X-Test-Header': 'test-value',
-          'Another-Header': 'another value',
         }),
       },
       AuthType.LOGIN_WITH_GOOGLE,
@@ -475,6 +474,64 @@ describe('createContentGenerator', () => {
       apiVersion: 'v1alpha',
     });
   });
+
+  it('should include config.headers for Gemini auth but not for others', async () => {
+    const mockConfig = {
+      getModel: vi.fn().mockReturnValue('gemini-pro'),
+      getProxy: vi.fn().mockReturnValue(undefined),
+      getUsageStatisticsEnabled: () => false,
+    } as unknown as Config;
+
+    const mockGenerator = {
+      models: {},
+    } as unknown as GoogleGenAI;
+    vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
+
+    // Test USE_GEMINI gets headers
+    await createContentGenerator(
+      {
+        apiKey: 'test-api-key',
+        authType: AuthType.USE_GEMINI,
+        headers: { 'X-Custom-Config-Header': 'custom-value' },
+      },
+      mockConfig,
+    );
+
+    expect(GoogleGenAI).toHaveBeenCalledWith(
+      expect.objectContaining({
+        httpOptions: expect.objectContaining({
+          headers: expect.objectContaining({
+            'X-Custom-Config-Header': 'custom-value',
+          }),
+        }),
+      }),
+    );
+
+    // Test LOGIN_WITH_GOOGLE does NOT get headers (unless they were in baseHeaders, which they are not)
+    const mockCodeAssistGenerator = {} as unknown as ContentGenerator;
+    vi.mocked(createCodeAssistContentGenerator).mockResolvedValue(
+      mockCodeAssistGenerator as never,
+    );
+
+    await createContentGenerator(
+      {
+        authType: AuthType.LOGIN_WITH_GOOGLE,
+        headers: { 'X-Should-Not-Be-Here': 'nope' },
+      },
+      mockConfig,
+    );
+
+    expect(createCodeAssistContentGenerator).toHaveBeenCalledWith(
+      expect.objectContaining({
+        headers: expect.not.objectContaining({
+          'X-Should-Not-Be-Here': 'nope',
+        }),
+      }),
+      AuthType.LOGIN_WITH_GOOGLE,
+      mockConfig,
+      undefined,
+    );
+  });
 });
 
 describe('createContentGeneratorConfig', () => {

packages/core/src/core/contentGenerator.ts

@@ -90,15 +90,19 @@ export function getAuthTypeFromEnv(): AuthType | undefined {
 
 export type ContentGeneratorConfig = {
   apiKey?: string;
+  baseUrl?: string;
   vertexai?: boolean;
   authType?: AuthType;
   proxy?: string;
+  headers?: Record<string, string>;
 };
 
 export async function createContentGeneratorConfig(
   config: Config,
   authType: AuthType | undefined,
   apiKey?: string,
+  baseUrl?: string,
+  headers?: Record<string, string>,
 ): Promise<ContentGeneratorConfig> {
   const geminiApiKey =
     apiKey ||
@@ -115,6 +119,7 @@ export async function createContentGeneratorConfig(
   const contentGeneratorConfig: ContentGeneratorConfig = {
     authType,
     proxy: config?.getProxy(),
+    headers,
   };
 
   // If we are using Google auth or we are in Cloud Shell, there is nothing else to validate for now
@@ -127,6 +132,7 @@ export async function createContentGeneratorConfig(
 
   if (authType === AuthType.USE_GEMINI && geminiApiKey) {
     contentGeneratorConfig.apiKey = geminiApiKey;
+    contentGeneratorConfig.baseUrl = baseUrl;
     contentGeneratorConfig.vertexai = false;
 
     return contentGeneratorConfig;
@@ -205,7 +211,10 @@ export async function createContentGenerator(
       config.authType === AuthType.USE_GEMINI ||
       config.authType === AuthType.USE_VERTEX_AI
     ) {
-      let headers: Record<string, string> = { ...baseHeaders };
+      let headers: Record<string, string> = {
+        ...baseHeaders,
+        ...config.headers,
+      };
       if (gcConfig?.getUsageStatisticsEnabled()) {
         const installationManager = new InstallationManager();
         const installationId = installationManager.getInstallationId();
@@ -214,7 +223,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,