From d84eeefed1ca5d7f4575da1a29edd9ed4bdf99d5 Mon Sep 17 00:00:00 2001
From: Shreya Keshive <shreyakeshive@google.com>
Date: Thu, 5 Mar 2026 11:50:19 -0500
Subject: [PATCH] fix

---
 .../cli/src/zed-integration/zedIntegration.ts | 20 ++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts
index cd383f4258..306355adf4 100644
--- a/packages/cli/src/zed-integration/zedIntegration.ts
+++ b/packages/cli/src/zed-integration/zedIntegration.ts
@@ -99,6 +99,7 @@ export class GeminiAgent {
   private clientCapabilities: acp.ClientCapabilities | undefined;
   private apiKey: string | undefined;
   private baseUrl: string | undefined;
+  private headers: Record<string, string> | undefined;
 
   constructor(
     private config: Config,
@@ -173,7 +174,7 @@ export class GeminiAgent {
   async authenticate(req: acp.AuthenticateRequest): Promise<void> {
     const methodId = req.methodId;
     let method = AuthType.USE_GEMINI;
-    if (methodId === 'gemini-custom-url') {
+    if (methodId === 'gemini-custom-url' || methodId === 'gateway') {
       method = AuthType.USE_GEMINI;
     } else {
       method = z.nativeEnum(AuthType).parse(methodId);
@@ -221,6 +222,9 @@ export class GeminiAgent {
       if (baseUrl) {
         this.baseUrl = baseUrl;
       }
+      if (headers) {
+        this.headers = headers;
+      }
       await this.config.refreshAuth(
         method,
         apiKey ?? this.apiKey,
@@ -256,7 +260,12 @@ export class GeminiAgent {
     let isAuthenticated = false;
     let authErrorMessage = '';
     try {
-      await config.refreshAuth(authType, this.apiKey, this.baseUrl);
+      await config.refreshAuth(
+        authType,
+        this.apiKey,
+        this.baseUrl,
+        this.headers,
+      );
       isAuthenticated = true;
 
       // Extra validation for Gemini API key
@@ -418,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();