@theia/ai-google 1.76.0-next.6 → 1.76.0

package/src/common/google-preferences.ts

@@ -18,7 +18,8 @@ import { AI_CORE_PREFERENCES_TITLE, MODEL_PROVIDER_TYPE_DETAIL, ModelProviderTyp
 import { LINUX_ENV_HINT, nls, PreferenceSchema } from '@theia/core';
 
 export const API_KEY_PREF = 'ai-features.google.apiKey';
-export const MODELS_PREF = 'ai-features.google.models';
+export const ALLOW_ENV_API_KEY_PREF = 'ai-features.google.allowEnvironmentApiKey';
+export const MODEL_OVERRIDES_PREF = 'ai-features.google.modelOverrides';
 export const MAX_RETRIES = 'ai-features.google.maxRetriesOnErrors';
 export const RETRY_DELAY_RATE_LIMIT = 'ai-features.google.retryDelayOnRateLimitError';
 export const RETRY_DELAY_OTHER_ERRORS = 'ai-features.google.retryDelayOnOtherErrors';
@@ -33,14 +34,24 @@ export const GooglePreferencesSchema: PreferenceSchema = {
             on the machine running Theia. Use the environment variable `GOOGLE_API_KEY` to set the key securely.') + LINUX_ENV_HINT,
             title: AI_CORE_PREFERENCES_TITLE,
         },
-        [MODELS_PREF]: {
-            type: 'array',
-            description: nls.localize('theia/ai/google/models/description', 'Official Google Gemini models to use'),
+        [ALLOW_ENV_API_KEY_PREF]: {
+            type: 'boolean',
+            default: false,
             title: AI_CORE_PREFERENCES_TITLE,
-            default: ['gemini-3.1-pro-preview', 'gemini-3.7-flash', 'gemini-3.6-flash', 'gemini-3.5-flash', 'gemini-3.5-flash-lite'],
+            markdownDescription: nls.localize('theia/ai/google/allowEnvApiKey/description',
+                'Allow Theia to use a Google AI (Gemini) API key found in the environment (`GOOGLE_API_KEY` / `GEMINI_API_KEY`). '
+                + 'You are asked to confirm this once before the key is used; set it back to `false` to revoke consent.'),
+        },
+        [MODEL_OVERRIDES_PREF]: {
+            type: 'array',
+            default: [],
             items: {
                 type: 'string'
-            }
+            },
+            title: AI_CORE_PREFERENCES_TITLE,
+            markdownDescription: nls.localize('theia/ai/google/modelOverrides/description',
+                'Override the models discovered from Google Gemini. When empty (default), the available models are discovered from the provider. '
+                + 'Set explicit model ids to use exactly those instead; discovery is then not used at all.')
         },
         [MAX_RETRIES]: {
             type: 'integer',

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

@@ -15,7 +15,8 @@
 // *****************************************************************************
 
 import { expect } from 'chai';
-import { LanguageModelRequest, ReasoningApi, ReasoningSupport } from '@theia/ai-core';
+import { LanguageModelRequest, LanguageModelTextResponse, ReasoningApi, ReasoningSupport } from '@theia/ai-core';
+import type { GoogleGenAI } from '@google/genai';
 import { GoogleModel } from './google-language-model';
 
 const GEMINI_REASONING_SUPPORT: ReasoningSupport = {
@@ -92,3 +93,33 @@ describe('GoogleModel reasoning translation', () => {
         });
     });
 });
+
+describe('GoogleModel non-streaming requests', () => {
+    /** Model whose generateContent() resolves to the given response instead of calling the API. */
+    class NonStreamingGoogleModel extends GoogleModel {
+        constructor(protected readonly generateContentResponse: object) {
+            super(
+                'test-id', 'gemini-3-pro', { status: 'ready' }, false,
+                () => 'test-key',
+                () => ({ maxRetriesOnErrors: 0, retryDelayOnRateLimitError: -1, retryDelayOnOtherErrors: -1 }),
+                GEMINI_REASONING_SUPPORT, 'effort'
+            );
+        }
+        protected override initializeGemini(): GoogleGenAI {
+            return { models: { generateContent: async () => this.generateContentResponse } } as unknown as GoogleGenAI;
+        }
+    }
+
+    it('excludes thought parts from the response text', async () => {
+        const model = new NonStreamingGoogleModel({
+            candidates: [{ content: { role: 'model', parts: [{ text: 'Weighing the options.', thought: true }, { text: 'Answer' }] } }],
+            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 }
+        });
+        const response = await model.request({
+            messages: [{ actor: 'user', type: 'text', text: 'hello' }],
+            reasoning: { level: 'medium' },
+            agentId: 'test', sessionId: 'session', requestId: 'req'
+        });
+        expect((response as LanguageModelTextResponse).text).to.equal('Answer');
+    });
+});

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

@@ -30,6 +30,7 @@ import {
     ReasoningSupport,
     ServerToolCall,
     ServerToolDescriptor,
+    TokenUsageParams,
     ToolCallResult,
     ToolInvocationContext,
     UserRequest
@@ -155,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. */
@@ -246,6 +261,7 @@ export class GoogleModel implements LanguageModel {
                 let latestUrlContextMetadata: UrlContextMetadata | undefined;
                 let latestGroundingMetadata: GroundingMetadata | undefined;
                 try {
+                    let tokenUsage: TokenUsageParams | undefined = undefined;
                     for await (const chunk of stream) {
                         if (cancellationToken?.isCancellationRequested) {
                             break;
@@ -345,16 +361,25 @@ export class GoogleModel implements LanguageModel {
                             yield { content: chunk.text };
                         }
 
-                        // Report token usage if available
+                        // Remember the token usage as Gemini's metadata is cumulative
                         if (chunk.usageMetadata) {
                             const promptTokens = chunk.usageMetadata.promptTokenCount;
                             const completionTokens = chunk.usageMetadata.candidatesTokenCount;
                             if (promptTokens !== undefined && completionTokens !== undefined) {
-                                yield { input_tokens: promptTokens, output_tokens: completionTokens };
+                                tokenUsage = {
+                                    inputTokens: promptTokens,
+                                    outputTokens: completionTokens,
+                                    requestId: request.requestId
+                                };
                             }
                         }
                     }
 
+                    // Report token usage if available
+                    if (tokenUsage !== undefined && that.id) {
+                        yield { input_tokens: tokenUsage.inputTokens, output_tokens: tokenUsage.outputTokens };
+                    }
+
                     // Surface any server tools (url_context / google_search) that the provider executed.
                     const serverToolCalls = that.buildServerToolCalls(latestUrlContextMetadata, latestGroundingMetadata);
                     if (serverToolCalls.length > 0) {
@@ -538,10 +563,11 @@ export class GoogleModel implements LanguageModel {
 
         try {
             let responseText = '';
-            // For non streaming requests we are always only interested in text parts
+            // For non streaming requests we are always only interested in text parts; thought summaries
+            // (parts flagged `thought`, present when includeThoughts is set) are not part of the answer.
             if (model.candidates?.[0]?.content?.parts) {
                 for (const part of model.candidates[0].content.parts) {
-                    if (part.text) {
+                    if (part.text && !part.thought) {
                         responseText += part.text;
                     }
                 }
@@ -569,8 +595,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;