@theia/ai-google 1.76.0-next.28 → 1.76.0-next.34

package/src/node/google-language-model.ts

@@ -156,6 +156,20 @@ function toGoogleRole(message: LanguageModelMessage): 'user' | 'model' {
  * Implements the Gemini language model integration for Theia. Reasoning-level
  * translation lives in {@link googleReasoningFor}.
  */
+/** Options for {@link createGoogleClient}. */
+export interface GoogleClientOptions {
+    readonly apiKey: string;
+}
+
+/**
+ * The single place a Gemini SDK client is built, so that a chat request, a model lookup and the model
+ * discovery all reach the provider the same way.
+ */
+export function createGoogleClient(options: GoogleClientOptions): GoogleGenAI {
+    // TODO test vertexai
+    return new GoogleGenAI({ apiKey: options.apiKey, vertexai: false });
+}
+
 export class GoogleModel implements LanguageModel {
 
     /** Provider identifier, used to key per-provider settings (e.g. server tool selections) and the capabilities UI. */
@@ -580,8 +594,7 @@ export class GoogleModel implements LanguageModel {
             throw new Error('Please provide GOOGLE_API_KEY in preferences or via environment variable');
         }
 
-        // TODO test vertexai
-        return new GoogleGenAI({ apiKey, vertexai: false });
+        return createGoogleClient({ apiKey });
     }
 
     /**

package/src/node/google-language-models-manager-impl.spec.ts

@@ -16,10 +16,11 @@
 
 import { expect } from 'chai';
 import type { Model } from '@google/genai';
-import { ReasoningApi } from '@theia/ai-core';
+import { DiscoveredModel, ReasoningApi } from '@theia/ai-core';
 import { GoogleLanguageModelsManagerImpl, reasoningApiFromModelId } from './google-language-models-manager-impl';
 import { GoogleModelDescription } from '../common';
 import { MockLogger } from '@theia/core/lib/common/test/mock-logger';
+import { TestModelDiscoveryFetcher } from '@theia/ai-core/lib/node/test/test-model-discovery-fetcher';
 
 class TestableGoogleManager extends GoogleLanguageModelsManagerImpl {
     public retrieveCalls: string[] = [];
@@ -46,6 +47,36 @@ class TestableGoogleManager extends GoogleLanguageModelsManagerImpl {
         }
         return this.stubbedInfo;
     }
+
+    public stubbedModels: Model[] = [];
+    public listCalls = 0;
+    /** Throw {@link failWith} for the first `failTimes` list calls, then return {@link stubbedModels}. */
+    public failTimes = 0;
+    public failWith: Error = new Error('boom');
+    /** The real discovery fetcher, snapshotting in memory and retrying without waiting. */
+    public readonly testFetcher = new TestModelDiscoveryFetcher();
+    protected override readonly discoveryFetcher = this.testFetcher;
+
+    /** In-memory stand-in for the on-disk snapshot. */
+    public get snapshot(): DiscoveredModel[] | undefined {
+        return this.testFetcher.snapshot;
+    }
+
+    public set snapshot(models: DiscoveredModel[] | undefined) {
+        this.testFetcher.snapshot = models;
+    }
+
+    protected override async listModels(_apiKey: string): Promise<Model[]> {
+        this.listCalls++;
+        if (this.listCalls <= this.failTimes) {
+            throw this.failWith;
+        }
+        return this.stubbedModels;
+    }
+}
+
+function listedModel(name: string, supportedActions?: string[], displayName?: string, modelDescription?: string): Model {
+    return ({ name, supportedActions, displayName, description: modelDescription }) as unknown as Model;
 }
 
 function description(model: string): GoogleModelDescription {
@@ -184,3 +215,135 @@ describe('GoogleLanguageModelsManagerImpl - fetchModelInfo cache', () => {
         expect(manager.retrieveCalls).to.deep.equal(['gemini-3-pro']);
     });
 });
+
+describe('GoogleLanguageModelsManagerImpl - fetchAvailableModels', () => {
+    let manager: TestableGoogleManager;
+
+    beforeEach(() => {
+        manager = new TestableGoogleManager();
+        (manager as unknown as { logger: MockLogger }).logger = new MockLogger();
+        manager.setApiKey('key');
+    });
+
+    it('returns an empty result and skips the network when no API key is set', async () => {
+        const previous = { google: process.env.GOOGLE_API_KEY, gemini: process.env.GEMINI_API_KEY };
+        delete process.env.GOOGLE_API_KEY;
+        delete process.env.GEMINI_API_KEY;
+        try {
+            manager.setApiKey(undefined);
+            expect(await manager.fetchAvailableModels()).to.deep.equal({ models: [], fromCache: false });
+            expect(manager.listCalls).to.equal(0);
+        } finally {
+            if (previous.google !== undefined) { process.env.GOOGLE_API_KEY = previous.google; }
+            if (previous.gemini !== undefined) { process.env.GEMINI_API_KEY = previous.gemini; }
+        }
+    });
+
+    it('strips the "models/" prefix, drops models that cannot generate content, and caches the result', async () => {
+        manager.stubbedModels = [
+            listedModel('models/gemini-3-pro', ['generateContent', 'countTokens'], 'Gemini 3 Pro', 'The best one.'),
+            listedModel('models/text-embedding-004', ['embedContent']),
+            // Older SDKs omit the field; such models are kept rather than silently dropped.
+            listedModel('models/gemini-3-flash'),
+            listedModel('models/gemini-3-pro', ['generateContent'])
+        ];
+        const result = await manager.fetchAvailableModels();
+        expect(result.fromCache).to.equal(false);
+        expect(result.models.map(model => model.id)).to.deep.equal(['gemini-3-pro', 'gemini-3-flash']);
+        expect(result.models[0].label).to.equal('Gemini 3 Pro');
+        expect(result.models[0].description).to.equal('The best one.');
+        expect(manager.snapshot).to.deep.equal(result.models);
+    });
+
+    it('keeps the Gemini chat models and drops what the endpoint otherwise carries', async () => {
+        manager.stubbedModels = [
+            listedModel('models/gemini-flash-latest'),
+            listedModel('models/gemini-3.8-flash'),
+            listedModel('models/gemini-3.1-pro-preview'),
+            // Chat-capable, but not a model of this family: its parameter count would read as a version.
+            listedModel('models/gemma-3-27b-it'),
+            listedModel('models/learnlm-2.0-experimental'),
+            // Gemini, and reported as generating content, but not chat: they answer through an API of
+            // their own, or they generate something other than text.
+            listedModel('models/gemini-3.5-pro-deep-research'),
+            listedModel('models/gemini-2.5-flash-preview-tts'),
+            listedModel('models/gemini-2.5-flash-image'),
+            listedModel('models/gemini-2.5-computer-use-preview'),
+            // The Live API and embedding models need no id term: they report what they can do.
+            listedModel('models/gemini-2.5-flash-live', ['bidiGenerateContent']),
+            listedModel('models/gemini-embedding-001', ['embedContent']),
+            listedModel('models/imagen-4.0-generate-001')
+        ];
+        const result = await manager.fetchAvailableModels();
+        expect(result.models.map(model => model.id)).to.deep.equal([
+            'gemini-flash-latest',
+            'gemini-3.8-flash',
+            'gemini-3.1-pro-preview'
+        ]);
+    });
+
+    it('retries transient network errors and then succeeds', async () => {
+        manager.stubbedModels = [listedModel('models/gemini-3-pro')];
+        manager.failTimes = 2;
+        manager.failWith = new Error('fetch failed');
+        const result = await manager.fetchAvailableModels();
+        expect(result.models.map(model => model.id)).to.deep.equal(['gemini-3-pro']);
+        expect(manager.listCalls).to.equal(3);
+    });
+
+    it('falls back to the cached snapshot when the fetch keeps failing', async () => {
+        manager.snapshot = [{ id: 'gemini-3-pro' }];
+        manager.failTimes = 99;
+        manager.failWith = new Error('ETIMEDOUT');
+        const result = await manager.fetchAvailableModels();
+        expect(result).to.deep.equal({ models: [{ id: 'gemini-3-pro' }], fromCache: true, error: 'ETIMEDOUT' });
+        expect(manager.listCalls).to.equal(3);
+    });
+
+    it('does not retry auth errors and throws without a snapshot', async () => {
+        manager.failTimes = 99;
+        manager.failWith = new Error('403 permission denied');
+        let threw = false;
+        try {
+            await manager.fetchAvailableModels();
+        } catch (error) {
+            threw = true;
+            expect((error as Error).message).to.equal('403 permission denied');
+        }
+        expect(threw).to.be.true;
+        expect(manager.listCalls).to.equal(1);
+    });
+});
+
+describe('GoogleLanguageModelsManagerImpl - environment API key consent', () => {
+
+    it('leaves an environment key unused until it is allowed, wherever a key would be read', async () => {
+        const previous = { GOOGLE_API_KEY: process.env.GOOGLE_API_KEY, GEMINI_API_KEY: process.env.GEMINI_API_KEY };
+        delete process.env.GOOGLE_API_KEY;
+        delete process.env.GEMINI_API_KEY;
+        process.env.GOOGLE_API_KEY = 'from-the-environment';
+        try {
+            const manager = new TestableGoogleManager();
+            // The gate sits on the key, so a custom endpoint or a configured model cannot reach past it either.
+            expect(manager.apiKey).to.equal(undefined);
+            // Discovery still needs to know the key is there, in order to ask for it.
+            expect(await manager.getApiKeySource()).to.equal('environment');
+
+            manager.setAllowEnvironmentApiKey(true);
+            expect(manager.apiKey).to.equal('from-the-environment');
+
+            // Withdrawing the consent stops it being used at once.
+            manager.setAllowEnvironmentApiKey(false);
+            expect(manager.apiKey).to.equal(undefined);
+        } finally {
+            if (previous.GOOGLE_API_KEY !== undefined) { process.env.GOOGLE_API_KEY = previous.GOOGLE_API_KEY; } else { delete process.env.GOOGLE_API_KEY; }
+            if (previous.GEMINI_API_KEY !== undefined) { process.env.GEMINI_API_KEY = previous.GEMINI_API_KEY; } else { delete process.env.GEMINI_API_KEY; }
+        }
+    });
+
+    it('uses a key set in the preferences whatever the environment says', () => {
+        const manager = new TestableGoogleManager();
+        manager.setApiKey('from-the-preference');
+        expect(manager.apiKey).to.equal('from-the-preference');
+    });
+});

package/src/node/google-language-models-manager-impl.ts

@@ -14,14 +14,19 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { LanguageModelRegistry, LanguageModelStatus, ReasoningApi, ReasoningSupport } from '@theia/ai-core';
+import {
+    ApiKeySource, DiscoveredModel, DiscoveredModels, LanguageModelRegistry, LanguageModelStatus, ModelDiscoveryResult, ReasoningApi, ReasoningSupport
+} from '@theia/ai-core';
+import { ModelDiscoveryFetcher } from '@theia/ai-core/lib/node';
 import { inject, injectable, named } from '@theia/core/shared/inversify';
-import { GoogleGenAI, Model } from '@google/genai';
-import { GoogleModel } from './google-language-model';
+import { Model } from '@google/genai';
+import { createGoogleClient, GoogleModel } from './google-language-model';
 import { GOOGLE_SERVER_TOOLS } from './google-server-tools';
 import { GoogleLanguageModelsManager, GoogleModelDescription } from '../common';
 import { ILogger } from '@theia/core';
 
+const GOOGLE_SNAPSHOT_FILE = 'google-models.json';
+
 export interface GoogleLanguageModelRetrySettings {
     maxRetriesOnErrors: number;
     retryDelayOnRateLimitError: number;
@@ -56,6 +61,12 @@ interface ResolvedModelMetadata {
 @injectable()
 export class GoogleLanguageModelsManagerImpl implements GoogleLanguageModelsManager {
     protected _apiKey: string | undefined;
+    /**
+     * Whether a key found in the environment may be used. Withheld until the user confirms it, so the
+     * gate sits on the key itself: every path that reaches for one — discovery, a custom endpoint, a
+     * manually configured model — is covered, and revoking the consent takes effect at once.
+     */
+    protected _allowEnvironmentApiKey = false;
     protected retrySettings: GoogleLanguageModelRetrySettings = {
         maxRetriesOnErrors: 3,
         retryDelayOnRateLimitError: 60,
@@ -71,8 +82,97 @@ export class GoogleLanguageModelsManagerImpl implements GoogleLanguageModelsMana
     @inject(ILogger) @named('ai-google:GoogleLanguageModelsManagerImpl')
     protected readonly logger: ILogger;
 
+    @inject(ModelDiscoveryFetcher)
+    protected readonly discoveryFetcher: ModelDiscoveryFetcher;
+
     get apiKey(): string | undefined {
-        return this._apiKey ?? process.env.GOOGLE_API_KEY ?? process.env.GEMINI_API_KEY;
+        return this._apiKey ?? (this._allowEnvironmentApiKey ? process.env.GOOGLE_API_KEY ?? process.env.GEMINI_API_KEY : undefined);
+    }
+
+    async getApiKeySource(): Promise<ApiKeySource> {
+        if (this._apiKey) {
+            return 'preference';
+        }
+        if (process.env.GOOGLE_API_KEY || process.env.GEMINI_API_KEY) {
+            return 'environment';
+        }
+        return 'none';
+    }
+
+    async fetchAvailableModels(): Promise<ModelDiscoveryResult> {
+        const apiKey = this.apiKey;
+        if (!apiKey) {
+            return { models: [], fromCache: false };
+        }
+        return this.discoveryFetcher.fetch({
+            snapshotFile: GOOGLE_SNAPSHOT_FILE,
+            providerLabel: 'Google',
+            listModels: async () => this.toDiscoveredModels(await this.listModels(apiKey)),
+            isRetryable: error => this.isRetryableError(error)
+        });
+    }
+
+    /**
+     * Maps the endpoint's entries onto {@link DiscoveredModel}s. Gemini ids carry no release date
+     * today, so collapsing dated variants is a no-op here; it is applied all the same, so a provider
+     * that starts pinning releases does not quietly multiply the list.
+     */
+    protected toDiscoveredModels(models: Model[]): DiscoveredModel[] {
+        const byId = new Map<string, DiscoveredModel>();
+        for (const model of models) {
+            // Only models usable for chat/content generation; older SDKs may omit the field.
+            if (model.supportedActions && !model.supportedActions.includes('generateContent')) {
+                continue;
+            }
+            const id = (model.name ?? '').replace(/^models\//, '');
+            if (id.length > 0 && this.isChatModelId(id) && !byId.has(id)) {
+                byId.set(id, { id, label: model.displayName, description: model.description });
+            }
+        }
+        return DiscoveredModels.withUndatedAliases([...byId.values()]);
+    }
+
+    /**
+     * Heuristic for the Gemini chat models among the entries that are left once {@link toDiscoveredModels}
+     * has dropped what cannot generate content at all. Two things remain to be excluded, and both
+     * report `generateContent` like a chat model: a model that generates something other than text
+     * (an image, speech), and one that answers through an agentic API of its own (deep research,
+     * computer use, robotics). Everything outside the `gemini-*` family goes too — Gemma, LearnLM —
+     * since being able to generate content does not make a model one of the models this provider is
+     * about. Anything this misses can be configured as a custom endpoint.
+     *
+     * The Live API models need no term of their own: they report `bidiGenerateContent`, so the
+     * capability check has already left them out.
+     *
+     * Without a release date to go by, the ranking that decides which models the chat input offers
+     * falls back to the numbers in the id, and those only mean a version within this family: a
+     * parameter count (`gemma-3-27b-it`) or an experiment date would otherwise read as the newest
+     * model there is.
+     */
+    protected isChatModelId(id: string): boolean {
+        if (!/^gemini-/.test(id)) {
+            return false;
+        }
+        return !/(image|tts|deep-research|computer-use|robotics)/.test(id);
+    }
+
+    /**
+     * Retry transient network errors; auth/quota errors fail fast. The SDK reports them as plain
+     * errors, so there is nothing to match on but the message.
+     */
+    protected isRetryableError(error: unknown): boolean {
+        const message = (error instanceof Error ? `${error.name} ${error.message}` : String(error)).toLowerCase();
+        return /econn|etimedout|enotfound|network|fetch failed|socket|timeout|aborted/.test(message);
+    }
+
+    /** Iterates the (auto-paginated) `/v1beta/models` endpoint. Overridable for testing. */
+    protected async listModels(apiKey: string): Promise<Model[]> {
+        const genAI = createGoogleClient({ apiKey });
+        const models: Model[] = [];
+        for await (const model of await genAI.models.list()) {
+            models.push(model);
+        }
+        return models;
     }
 
     protected calculateStatus(effectiveApiKey: string | undefined): LanguageModelStatus {
@@ -168,7 +268,7 @@ export class GoogleLanguageModelsManagerImpl implements GoogleLanguageModelsMana
     }
 
     protected retrieveModelInfo(modelDescription: GoogleModelDescription, apiKey: string): Promise<Model> {
-        const genAI = new GoogleGenAI({ apiKey, vertexai: false });
+        const genAI = createGoogleClient({ apiKey });
         return genAI.models.get({ model: modelDescription.model });
     }
 
@@ -184,6 +284,10 @@ export class GoogleLanguageModelsManagerImpl implements GoogleLanguageModelsMana
         this.languageModelRegistry.removeLanguageModels(modelIds);
     }
 
+    setAllowEnvironmentApiKey(allowed: boolean): void {
+        this._allowEnvironmentApiKey = allowed;
+    }
+
     setApiKey(apiKey: string | undefined): void {
         if (apiKey) {
             this._apiKey = apiKey;