diff --git a/extensions/vertex-ai/package.json b/extensions/vertex-ai/package.json index 99c01ad5eb..dcc74bc415 100644 --- a/extensions/vertex-ai/package.json +++ b/extensions/vertex-ai/package.json @@ -34,6 +34,31 @@ "format": "file", "scope": "InferenceProviderConnectionFactory", "description": "Path to gcloud credentials file (e.g. ~/.config/gcloud/application_default_credentials.json)" + }, + "vertex-ai.connection._type": { + "type": "string", + "scope": "InferenceProviderConnection", + "description": "Vertex AI provider type", + "hidden": true + }, + "vertex-ai.connection.token": { + "type": "string", + "scope": "InferenceProviderConnection", + "description": "Credentials secret reference", + "format": "password", + "hidden": true + }, + "vertex-ai.connection.VERTEX_AI_PROJECT_ID": { + "type": "string", + "scope": "InferenceProviderConnection", + "description": "Google Cloud project ID", + "hidden": true + }, + "vertex-ai.connection.VERTEX_AI_REGION": { + "type": "string", + "scope": "InferenceProviderConnection", + "description": "Google Cloud region", + "hidden": true } } } diff --git a/extensions/vertex-ai/src/extension.ts b/extensions/vertex-ai/src/extension.ts index bc324cee6a..c74f39a138 100644 --- a/extensions/vertex-ai/src/extension.ts +++ b/extensions/vertex-ai/src/extension.ts @@ -17,7 +17,7 @@ ***********************************************************************/ import type { ExtensionContext } from '@openkaiden/api'; -import { provider } from '@openkaiden/api'; +import { configuration, provider } from '@openkaiden/api'; import { VertexAi } from './vertex-ai'; @@ -26,7 +26,7 @@ let vertexAi: VertexAi | undefined; export async function activate(extensionContext: ExtensionContext): Promise { console.log('starting vertex-ai extension'); - vertexAi = new VertexAi(provider, extensionContext.secrets); + vertexAi = new VertexAi(provider, extensionContext.secrets, configuration); extensionContext.subscriptions.push(vertexAi); await vertexAi.init(); diff --git a/extensions/vertex-ai/src/vertex-ai.spec.ts b/extensions/vertex-ai/src/vertex-ai.spec.ts index 135cff6f91..5ebaab7305 100644 --- a/extensions/vertex-ai/src/vertex-ai.spec.ts +++ b/extensions/vertex-ai/src/vertex-ai.spec.ts @@ -22,12 +22,21 @@ import { homedir } from 'node:os'; import { join } from 'node:path'; import { createVertexAnthropic } from '@ai-sdk/google-vertex/anthropic'; -import type { Disposable, Logger, Provider, provider as ProviderAPI, SecretStorage } from '@openkaiden/api'; +import type { + Configuration, + configuration as ConfigurationAPI, + Disposable, + Logger, + Provider, + provider as ProviderAPI, + SecretStorage, +} from '@openkaiden/api'; import { assert, beforeEach, describe, expect, test, vi } from 'vitest'; import { CONNECTIONS_KEY, FALLBACK_MODELS, + PROVIDER_ID, type StoredConnection, VertexAi, type VertexAiConnectionConfig, @@ -58,6 +67,19 @@ const SECRET_STORAGE_MOCK: SecretStorage = { onDidChange: vi.fn(), }; +const CONFIG_UPDATE_MOCK = vi.fn(); + +const CONFIGURATION_MOCK: Configuration = { + get: vi.fn(), + has: vi.fn(), + update: CONFIG_UPDATE_MOCK, +} as unknown as Configuration; + +const CONFIGURATION_API_MOCK: typeof ConfigurationAPI = { + getConfiguration: vi.fn().mockReturnValue(CONFIGURATION_MOCK), + onDidChangeConfiguration: vi.fn(), +} as unknown as typeof ConfigurationAPI; + const VALID_CREDENTIALS = JSON.stringify({ client_id: 'test-client-id.apps.googleusercontent.com', client_secret: 'test-client-secret', @@ -107,12 +129,13 @@ beforeEach(() => { vi.mocked(homedir).mockReturnValue('/home/testuser'); vi.mocked(readFile).mockResolvedValue(VALID_CREDENTIALS); vi.mocked(access).mockResolvedValue(undefined); + vi.mocked(CONFIGURATION_API_MOCK.getConfiguration).mockReturnValue(CONFIGURATION_MOCK); mockFetchResponses(); }); function createVertexAi(): VertexAi { - return new VertexAi(PROVIDER_API_MOCK, SECRET_STORAGE_MOCK); + return new VertexAi(PROVIDER_API_MOCK, SECRET_STORAGE_MOCK, CONFIGURATION_API_MOCK); } describe('init', () => { @@ -522,7 +545,6 @@ describe('factory', () => { 'vertex-ai.factory.credentialsFile': '/home/user/.config/gcloud/application_default_credentials.json', }); - expect(SECRET_STORAGE_MOCK.store).toHaveBeenCalledOnce(); const expected: StoredConnection[] = [ { id: 'fake-uuid-1', @@ -590,20 +612,77 @@ describe('connection delete lifecycle', () => { mDelete = lifecycle.delete; }); - test('should remove config from storage on delete', async () => { + test('calling delete should update secrets, clear configuration, and dispose provider inference connection', async () => { await mDelete(); - expect(SECRET_STORAGE_MOCK.store).toHaveBeenCalledTimes(2); - // Second store call is the removal (filtered list) - expect(SECRET_STORAGE_MOCK.store).toHaveBeenNthCalledWith(2, CONNECTIONS_KEY, '[]'); - }); + expect(SECRET_STORAGE_MOCK.delete).toHaveBeenCalledWith(`${PROVIDER_ID}:fake-uuid-1:token`); + + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection._type', undefined); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.token', undefined); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_PROJECT_ID', undefined); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_REGION', undefined); - test('should dispose provider inference connection on delete', async () => { - await mDelete(); expect(disposeMock).toHaveBeenCalledOnce(); }); }); +describe('workspace configuration', () => { + beforeEach(async () => { + vi.mocked(PROVIDER_MOCK.registerInferenceProviderConnection).mockReturnValue({ + dispose: vi.fn(), + }); + }); + + test('should store per-connection secret and set configuration after registration', async () => { + const vertexAi = createVertexAi(); + await vertexAi.init(); + + const mock = vi.mocked(PROVIDER_MOCK.setInferenceProviderConnectionFactory); + const create = mock.mock.calls[0][0].create; + + await create({ + 'vertex-ai.factory.projectId': 'my-project', + 'vertex-ai.factory.region': 'us-east5', + 'vertex-ai.factory.credentialsFile': '/home/user/.config/gcloud/application_default_credentials.json', + }); + + expect(SECRET_STORAGE_MOCK.store).toHaveBeenCalledWith( + `${PROVIDER_ID}:fake-uuid-1:token`, + '/home/user/.config/gcloud/application_default_credentials.json', + ); + + const connection = vi.mocked(PROVIDER_MOCK.registerInferenceProviderConnection).mock.calls[0][0]; + expect(CONFIGURATION_API_MOCK.getConfiguration).toHaveBeenCalledWith(undefined, connection); + + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection._type', PROVIDER_ID); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.token', `${PROVIDER_ID}:fake-uuid-1:token`); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_PROJECT_ID', 'my-project'); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_REGION', 'us-east5'); + }); + + test('should set workspace configuration for each restored connection', async () => { + const stored: StoredConnection[] = [ + { id: 'id-1', projectId: 'proj-a', region: 'us-east5', credentialsFile: '/path/a' }, + { id: 'id-2', projectId: 'proj-b', region: 'europe-west1', credentialsFile: '/path/b' }, + ]; + vi.mocked(SECRET_STORAGE_MOCK.get).mockResolvedValue(JSON.stringify(stored)); + + const vertexAi = createVertexAi(); + await vertexAi.init(); + + expect(SECRET_STORAGE_MOCK.store).toHaveBeenCalledWith(`${PROVIDER_ID}:id-1:token`, '/path/a'); + expect(SECRET_STORAGE_MOCK.store).toHaveBeenCalledWith(`${PROVIDER_ID}:id-2:token`, '/path/b'); + + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection._type', PROVIDER_ID); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.token', `${PROVIDER_ID}:id-1:token`); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.token', `${PROVIDER_ID}:id-2:token`); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_PROJECT_ID', 'proj-a'); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_PROJECT_ID', 'proj-b'); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_REGION', 'us-east5'); + expect(CONFIG_UPDATE_MOCK).toHaveBeenCalledWith('vertex-ai.connection.VERTEX_AI_REGION', 'europe-west1'); + }); +}); + describe('dispose', () => { test('should dispose provider and all connections', async () => { const disposeMock = vi.fn(); diff --git a/extensions/vertex-ai/src/vertex-ai.ts b/extensions/vertex-ai/src/vertex-ai.ts index 35ab763bbb..48b0fe1854 100644 --- a/extensions/vertex-ai/src/vertex-ai.ts +++ b/extensions/vertex-ai/src/vertex-ai.ts @@ -23,8 +23,10 @@ import { join } from 'node:path'; import { createVertexAnthropic } from '@ai-sdk/google-vertex/anthropic'; import type { + configuration as ConfigurationAPI, Disposable, InferenceModel, + InferenceProviderConnection, Provider, provider as ProviderAPI, ProviderConnectionStatus, @@ -32,6 +34,7 @@ import type { } from '@openkaiden/api'; export const CONNECTIONS_KEY = 'vertex-ai:connections'; +export const PROVIDER_ID = 'vertex-ai'; const FETCH_TIMEOUT_MS = 30_000; const TOKEN_ENDPOINT = 'https://oauth2.googleapis.com/token'; export interface VertexAiConnectionConfig { @@ -101,6 +104,7 @@ export class VertexAi implements Disposable { constructor( private readonly providerAPI: typeof ProviderAPI, private readonly secrets: SecretStorage, + private readonly configurationAPI: typeof ConfigurationAPI, ) {} async init(): Promise { @@ -268,10 +272,35 @@ export class VertexAi implements Disposable { } } - /** - * Registers a connection using pre-validated models (factory path) - * or by fetching them fresh with graceful degradation (restore path). - */ + private getSecretName(connectionId: string): string { + return `${PROVIDER_ID}:${connectionId}:token`; + } + + private async setConnectionConfiguration( + connection: InferenceProviderConnection, + config: VertexAiConnectionConfig, + ): Promise { + const secretName = this.getSecretName(connection.id); + await this.secrets.store(secretName, config.credentialsFile); + + const cfg = this.configurationAPI.getConfiguration(undefined, connection); + await cfg.update('vertex-ai.connection._type', PROVIDER_ID); + await cfg.update('vertex-ai.connection.token', secretName); + await cfg.update('vertex-ai.connection.VERTEX_AI_PROJECT_ID', config.projectId); + await cfg.update('vertex-ai.connection.VERTEX_AI_REGION', config.region); + } + + private async clearConnectionConfiguration(connection: InferenceProviderConnection): Promise { + const secretName = this.getSecretName(connection.id); + await this.secrets.delete(secretName); + + const cfg = this.configurationAPI.getConfiguration(undefined, connection); + await cfg.update('vertex-ai.connection._type', undefined); + await cfg.update('vertex-ai.connection.token', undefined); + await cfg.update('vertex-ai.connection.VERTEX_AI_PROJECT_ID', undefined); + await cfg.update('vertex-ai.connection.VERTEX_AI_REGION', undefined); + } + private async registerInferenceProviderConnection( id: string, config: VertexAiConnectionConfig, @@ -289,12 +318,6 @@ export class VertexAi implements Disposable { }, }); - const clean = async (): Promise => { - this.connections.get(id)?.dispose(); - this.connections.delete(id); - await this.removeConnection(id); - }; - const status: ProviderConnectionStatus = 'unknown'; let models: InferenceModel[]; @@ -318,7 +341,7 @@ export class VertexAi implements Disposable { } } - const connectionDisposable = this.provider.registerInferenceProviderConnection({ + const connection: InferenceProviderConnection = { id, name: `${config.projectId} (${config.region})`, type: 'cloud', @@ -330,7 +353,12 @@ export class VertexAi implements Disposable { return status; }, lifecycle: { - delete: clean.bind(this), + delete: async (): Promise => { + await this.clearConnectionConfiguration(connection); + this.connections.get(id)?.dispose(); + this.connections.delete(id); + await this.removeConnection(id); + }, }, models, credentials(): Record { @@ -340,8 +368,18 @@ export class VertexAi implements Disposable { credentialsFile: config.credentialsFile, }; }, - }); + }; + + const connectionDisposable = this.provider.registerInferenceProviderConnection(connection); this.connections.set(id, connectionDisposable); + + try { + await this.setConnectionConfiguration(connection, config); + } catch (error) { + connectionDisposable.dispose(); + this.connections.delete(id); + throw error; + } } /**