update

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

@@ -207,7 +207,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,
     );
@@ -227,6 +227,8 @@ describe('GeminiAgent', () => {
     expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
       AuthType.LOGIN_WITH_GOOGLE,
       undefined,
+      undefined,
+      undefined,
     );
     expect(mockSettings.setValue).toHaveBeenCalledWith(
       SettingScope.User,
@@ -246,6 +248,8 @@ describe('GeminiAgent', () => {
     expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
       AuthType.USE_GEMINI,
       'test-api-key',
+      undefined,
+      undefined,
     );
     expect(mockSettings.setValue).toHaveBeenCalledWith(
       SettingScope.User,
@@ -267,6 +271,37 @@ describe('GeminiAgent', () => {
       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,

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

@@ -117,14 +117,14 @@ export class GeminiAgent {
         },
       },
       {
-        id: 'gemini-custom-url',
-        name: 'Gemini API (Custom URL)',
-        description: 'Use an API key and custom base URL',
+        id: 'gateway',
+        name: 'AI API Gateway',
+        description: 'Use a custom AI API Gateway',
         _meta: {
-          'api-key': {
-            provider: 'google',
+          gateway: {
+            protocol: 'google',
+            restartRequired: 'false',
           },
-          'base-url': {},
         },
       },
       {
@@ -177,8 +177,28 @@ export class GeminiAgent {
     const meta = hasMeta(req) ? req._meta : undefined;
     const apiKey =
       typeof meta?.['api-key'] === 'string' ? meta['api-key'] : undefined;
-    const baseUrl =
+    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,
@@ -194,6 +214,7 @@ export class GeminiAgent {
         method,
         apiKey ?? this.apiKey,
         baseUrl ?? this.baseUrl,
+        headers,
       );
     } catch (e) {
       throw new acp.RequestError(-32000, getAcpErrorMessage(e));