diff --git a/packages/cli/src/zed-integration/zedIntegration.test.ts b/packages/cli/src/zed-integration/zedIntegration.test.ts index 37da3035c3..45e115ca9e 100644 --- a/packages/cli/src/zed-integration/zedIntegration.test.ts +++ b/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', () => { diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index d4f1b27b92..8cf630dc3b 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -87,6 +87,7 @@ export class GeminiAgent { private sessions: Map = 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 { - 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 diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 7297693b8e..7751daa269 100644 --- a/packages/core/src/config/config.ts +++ b/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, diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 98d8d50020..b96eb4f8fe 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/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 { 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; baseUrl?: string } = + { headers }; + if (config.baseUrl) { + httpOptions.baseUrl = config.baseUrl; + } const googleGenAI = new GoogleGenAI({ apiKey: config.apiKey === '' ? undefined : config.apiKey,