diff --git a/src/cli/commands/proxy/__tests__/connect-orchestrator.test.ts b/src/cli/commands/proxy/__tests__/connect-orchestrator.test.ts index 85cbf0090..b3e63b4e8 100644 --- a/src/cli/commands/proxy/__tests__/connect-orchestrator.test.ts +++ b/src/cli/commands/proxy/__tests__/connect-orchestrator.test.ts @@ -302,7 +302,7 @@ describe('connectTargets — per-target dispatch, summary, partial-failure seman const { writeDesktopConfig } = await import('../connectors/desktop.js'); const { writeVsCodeLanguageModelsConfig } = await import('../connectors/vscode.js'); vi.mocked(writeDesktopConfig).mockResolvedValue('/desktop/config.json'); - vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/vscode/models.json', requiresSecretConfiguration: false }); + vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/vscode/models.json', requiresSecretConfiguration: false, modelIds: ['model-a'] }); const { connectTargets } = await import('../connect-orchestrator.js'); await connectTargets({ targets: { claudeDesktop: true, vscode: true } }); @@ -319,7 +319,7 @@ describe('connectTargets — per-target dispatch, summary, partial-failure seman const { writeDesktopConfig } = await import('../connectors/desktop.js'); const { writeVsCodeLanguageModelsConfig } = await import('../connectors/vscode.js'); vi.mocked(writeDesktopConfig).mockRejectedValue(new Error('disk full')); - vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/vscode/models.json', requiresSecretConfiguration: false }); + vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/vscode/models.json', requiresSecretConfiguration: false, modelIds: ['model-a'] }); const { connectTargets } = await import('../connect-orchestrator.js'); await connectTargets({ targets: { claudeDesktop: true, vscode: true } }); @@ -381,7 +381,7 @@ describe('connectTargets — per-target dispatch, summary, partial-failure seman it('single successful target still prints the per-target summary (spec §3.4)', async () => { await withFreshDaemon(); const { writeVsCodeLanguageModelsConfig } = await import('../connectors/vscode.js'); - vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/vscode/models.json', requiresSecretConfiguration: false }); + vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/vscode/models.json', requiresSecretConfiguration: false, modelIds: ['model-a'] }); const { connectTargets } = await import('../connect-orchestrator.js'); await connectTargets({ targets: { vscode: true } }); diff --git a/src/cli/commands/proxy/__tests__/index.test.ts b/src/cli/commands/proxy/__tests__/index.test.ts index c6cc9c125..998f6e1c9 100644 --- a/src/cli/commands/proxy/__tests__/index.test.ts +++ b/src/cli/commands/proxy/__tests__/index.test.ts @@ -400,6 +400,7 @@ describe('proxy connect vscode', () => { vi.mocked(writeVsCodeLanguageModelsConfig).mockResolvedValue({ configPath: '/mock/chatLanguageModels.json', requiresSecretConfiguration: false, + modelIds: ['gpt-4.1'], }); }); @@ -450,6 +451,7 @@ describe('proxy connect vscode', () => { expect(daemonOptions).not.toHaveProperty('model'); expect(writeVsCodeLanguageModelsConfig).toHaveBeenCalledWith( 'http://127.0.0.1:4001', + 'local-key', false ); }); @@ -491,6 +493,7 @@ describe('proxy connect vscode', () => { expect(spawnDaemon).not.toHaveBeenCalled(); expect(writeVsCodeLanguageModelsConfig).toHaveBeenCalledWith( 'http://127.0.0.1:4001', + 'local-key', false ); }); diff --git a/src/cli/commands/proxy/connect-orchestrator.ts b/src/cli/commands/proxy/connect-orchestrator.ts index 8fcc3c241..354040d5b 100644 --- a/src/cli/commands/proxy/connect-orchestrator.ts +++ b/src/cli/commands/proxy/connect-orchestrator.ts @@ -36,7 +36,6 @@ import { import { fetchManagedMcpServers } from './connectors/managed-mcp-remote.js'; import { writeVsCodeClaudeCodeConfig } from './connectors/vscode-claude-code.js'; import { writeVsCodeLanguageModelsConfig } from './connectors/vscode.js'; -import { VS_CODE_SUPPORTED_MODELS } from './connectors/vscode-models.js'; import { checkProxyHealth } from './health-check.js'; import { discoverCodexModels, @@ -442,7 +441,7 @@ async function runVscodeByok( verbose: boolean ): Promise { try { - const result = await writeVsCodeLanguageModelsConfig(state.url, insiders); + const result = await writeVsCodeLanguageModelsConfig(state.url, state.gatewayKey, insiders); logger.info( '[proxy] VS Code BYOK configuration written', ...sanitizeLogArgs({ @@ -450,7 +449,7 @@ async function runVscodeByok( gatewayUrl: state.url, profile: state.profile, project: state.project, - modelCount: VS_CODE_SUPPORTED_MODELS.length, + modelCount: result.modelIds.length, clientType: state.clientType, requiresSecretConfiguration: result.requiresSecretConfiguration, }) @@ -459,7 +458,7 @@ async function runVscodeByok( if (verbose) { console.log(` Config: ${result.configPath}`); console.log(` Gateway: ${state.url}`); - console.log(` Models: ${VS_CODE_SUPPORTED_MODELS.length}`); + console.log(` Models: ${result.modelIds.length}`); console.log(` Project: ${config.codeMieProject || '(not configured)'}`); } diff --git a/src/cli/commands/proxy/connectors/__tests__/vscode.test.ts b/src/cli/commands/proxy/connectors/__tests__/vscode.test.ts index 42afa7407..a0e4bfb56 100644 --- a/src/cli/commands/proxy/connectors/__tests__/vscode.test.ts +++ b/src/cli/commands/proxy/connectors/__tests__/vscode.test.ts @@ -74,7 +74,11 @@ describe('writeVsCodeLanguageModelsConfigAtPath', () => { const provider = providers[0]; const models = provider.models as Array>; - expect(result).toEqual({ configPath, requiresSecretConfiguration: true }); + expect(result).toEqual({ + configPath, + requiresSecretConfiguration: true, + modelIds: [...EXPECTED_MODEL_IDS], + }); expect(provider).toMatchObject({ name: 'CodeMie', vendor: 'customendpoint', diff --git a/src/cli/commands/proxy/connectors/gateway-models.ts b/src/cli/commands/proxy/connectors/gateway-models.ts new file mode 100644 index 000000000..084f9024e --- /dev/null +++ b/src/cli/commands/proxy/connectors/gateway-models.ts @@ -0,0 +1,77 @@ +import { ConfigurationError } from '@/utils/errors.js'; +import { logger } from '@/utils/logger.js'; +import { sanitizeLogArgs } from '@/utils/security.js'; + +/** SSO-backed catalog endpoint; `include_all` returns every model the tenant may use, not just the defaults. */ +const MODELS_PATH = '/v1/llm_models?include_all=true'; + +interface ModelsListResponse { + data?: Array<{ id?: string }>; +} + +interface CodeMieLlmModel { + id?: string; + base_name?: string; + deployment_name?: string; +} + +function extractModelIds(json: ModelsListResponse | CodeMieLlmModel[]): string[] { + const ids = Array.isArray(json) + ? json.map((model) => model.id || model.base_name || model.deployment_name) + : (json.data ?? []).map((model) => model.id); + return [...new Set(ids.filter((id): id is string => typeof id === 'string' && id.length > 0))]; +} + +/** + * Fetch every model ID the gateway serves for the active profile. + * + * Unlike the Claude Desktop path this applies no family filter — callers decide + * what to expose. Always throws {@link ConfigurationError} on failure so callers + * can pick their own fallback instead of silently receiving an empty catalog. + */ +export async function fetchGatewayModelIds( + proxyUrl: string, + gatewayKey: string +): Promise { + const endpoint = new URL(MODELS_PATH, proxyUrl).toString(); + + let response: Response; + try { + response = await fetch(endpoint, { + headers: { Authorization: `Bearer ${gatewayKey}` }, + }); + } catch (error) { + throw new ConfigurationError( + `Local proxy model discovery could not reach ${endpoint}. ` + + `Reason: ${error instanceof Error ? error.message : String(error)}` + ); + } + + if (!response.ok) { + throw new ConfigurationError( + response.status === 401 + ? `Local proxy model discovery was rejected with 401 Unauthorized at ${endpoint}. ` + + 'The local gateway key was not accepted by the proxy or was forwarded upstream incorrectly.' + : `Local proxy model discovery failed at ${endpoint}: ${response.status} ${response.statusText}` + ); + } + + const contentType = response.headers.get('content-type') ?? ''; + if (!contentType.includes('application/json')) { + throw new ConfigurationError( + `Local proxy model discovery received an unexpected response (${contentType || 'no content-type'}) from ${endpoint}. ` + + 'Your SSO session may have expired — run `codemie proxy stop && codemie profile login` to re-authenticate.' + ); + } + + const ids = extractModelIds(await response.json() as ModelsListResponse | CodeMieLlmModel[]); + logger.info( + '[proxy] Gateway model discovery completed', + ...sanitizeLogArgs({ + endpoint, + modelCount: ids.length, + models: ids, + }) + ); + return ids; +} diff --git a/src/cli/commands/proxy/connectors/vscode.ts b/src/cli/commands/proxy/connectors/vscode.ts index bb972dfcd..9f37a14f5 100644 --- a/src/cli/commands/proxy/connectors/vscode.ts +++ b/src/cli/commands/proxy/connectors/vscode.ts @@ -3,14 +3,27 @@ import { mkdir, readFile, rename, stat, unlink, writeFile } from 'node:fs/promis import { homedir } from 'node:os'; import { dirname, join } from 'node:path'; import { ConfigurationError } from '@/utils/errors.js'; +import { logger } from '@/utils/logger.js'; +import { sanitizeLogArgs } from '@/utils/security.js'; +import { fetchGatewayModelIds } from './gateway-models.js'; import { VS_CODE_SUPPORTED_MODELS, type VsCodeApiType, + type VsCodeModelDefinition, type VsCodeReasoningEffort, } from './vscode-models.js'; const SECRET_REFERENCE_PATTERN = /^\$\{input:chat\.lm\.secret\.[^}]+\}$/; +/** Capabilities assumed for a gateway model the curated catalog does not describe. */ +const UNCATALOGED_MODEL_DEFAULTS: Omit = { + apiType: 'chat-completions', + vision: false, + thinking: false, + maxInputTokens: 128000, + maxOutputTokens: 16384, +}; + interface VsCodeLanguageModelProvider { [key: string]: unknown; name?: string; @@ -46,6 +59,8 @@ interface VsCodeManagedModel { export interface WriteVsCodeConfigResult { configPath: string; requiresSecretConfiguration: boolean; + /** Model IDs written under the CodeMie provider, in the order VS Code will see them. */ + modelIds: string[]; } function isRecord(value: unknown): value is Record { @@ -104,41 +119,95 @@ function getApiPath(apiType: VsCodeApiType): string { return '/v1/chat/completions'; } -function buildManagedModels(proxyUrl: string): VsCodeManagedModel[] { - return VS_CODE_SUPPORTED_MODELS.map(definition => { - const model: VsCodeManagedModel = { - id: definition.id, - name: definition.id, - url: new URL(getApiPath(definition.apiType), proxyUrl).toString(), - apiType: definition.apiType, - toolCalling: true, - vision: definition.vision, - streaming: true, - thinking: definition.thinking, - maxInputTokens: definition.maxInputTokens, - maxOutputTokens: definition.maxOutputTokens, - }; - - if (definition.adaptiveThinking) model.adaptiveThinking = true; - if (definition.zeroDataRetentionEnabled !== undefined) { - model.zeroDataRetentionEnabled = definition.zeroDataRetentionEnabled; - } - if (definition.modelOptions) model.modelOptions = definition.modelOptions; - if (definition.requestHeaders) model.requestHeaders = definition.requestHeaders; - if (definition.supportsReasoningEffort) { - model.supportsReasoningEffort = definition.supportsReasoningEffort; - } - if (definition.reasoningEffortFormat) { - model.reasoningEffortFormat = definition.reasoningEffortFormat; - } +const CATALOG_BY_ID = new Map(VS_CODE_SUPPORTED_MODELS.map(definition => [definition.id, definition])); +const CATALOG_ORDER = new Map(VS_CODE_SUPPORTED_MODELS.map((definition, index) => [definition.id, index])); + +/** + * Match a gateway model ID against the curated catalog, tolerating the suffixes + * the gateway adds to its registrations (`-YYYYMMDD` snapshots, `-vertex` + * deployments). Returns undefined when the catalog says nothing about the model, + * which is not an error: the model is still exposed, with default capabilities. + */ +export function resolveModelDefinition(id: string): VsCodeModelDefinition | undefined { + const stripVertex = (value: string): string => value.replace(/-vertex$/i, ''); + const stripSnapshot = (value: string): string => value.replace(/-\d{6,10}$/, ''); + const candidates = new Set([ + id, + stripVertex(id), + stripSnapshot(id), + stripSnapshot(stripVertex(id)), + ]); + + for (const candidate of candidates) { + const definition = CATALOG_BY_ID.get(candidate); + if (definition) return definition; + } + return undefined; +} - return model; - }); +/** Cataloged models first (in catalog order) so the familiar picker ordering survives. */ +function orderModelIds(availableIds: readonly string[]): string[] { + const known: Array<{ id: string; rank: number }> = []; + const unknown: string[] = []; + + for (const id of new Set(availableIds)) { + const rank = CATALOG_ORDER.get(resolveModelDefinition(id)?.id ?? ''); + if (rank === undefined) unknown.push(id); + else known.push({ id, rank }); + } + + known.sort((a, b) => a.rank - b.rank || a.id.localeCompare(b.id)); + unknown.sort((a, b) => a.localeCompare(b)); + return [...known.map(entry => entry.id), ...unknown]; +} + +function toManagedModel(proxyUrl: string, id: string): VsCodeManagedModel { + const definition = resolveModelDefinition(id) ?? UNCATALOGED_MODEL_DEFAULTS; + const model: VsCodeManagedModel = { + id, + name: id, + url: new URL(getApiPath(definition.apiType), proxyUrl).toString(), + apiType: definition.apiType, + toolCalling: true, + vision: definition.vision, + streaming: true, + thinking: definition.thinking, + maxInputTokens: definition.maxInputTokens, + maxOutputTokens: definition.maxOutputTokens, + }; + + if (definition.adaptiveThinking) model.adaptiveThinking = true; + if (definition.zeroDataRetentionEnabled !== undefined) { + model.zeroDataRetentionEnabled = definition.zeroDataRetentionEnabled; + } + if (definition.modelOptions) model.modelOptions = definition.modelOptions; + if (definition.requestHeaders) model.requestHeaders = definition.requestHeaders; + if (definition.supportsReasoningEffort) { + model.supportsReasoningEffort = definition.supportsReasoningEffort; + } + if (definition.reasoningEffortFormat) { + model.reasoningEffortFormat = definition.reasoningEffortFormat; + } + + return model; +} + +/** + * Build the provider's model list. `availableIds` is what the gateway actually + * serves; omitting it falls back to the curated catalog, which is only correct + * when discovery could not run. + */ +function buildManagedModels(proxyUrl: string, availableIds?: readonly string[]): VsCodeManagedModel[] { + const ids = availableIds + ? orderModelIds(availableIds) + : VS_CODE_SUPPORTED_MODELS.map(definition => definition.id); + return ids.map(id => toManagedModel(proxyUrl, id)); } function mergeManagedProviders( providers: VsCodeLanguageModelProvider[], - proxyUrl: string + proxyUrl: string, + availableIds?: readonly string[] ): { provider: VsCodeLanguageModelProvider; requiresSecretConfiguration: boolean } { const existingProvider = Object.assign({}, ...providers); const existingSettings = Object.assign( @@ -154,7 +223,7 @@ function mergeManagedProviders( name: 'CodeMie', vendor: 'customendpoint', apiType: 'chat-completions', - models: buildManagedModels(proxyUrl), + models: buildManagedModels(proxyUrl, availableIds), }; // VS Code owns effort selections. Preserve them instead of racing with the editor. @@ -223,19 +292,47 @@ export async function writeAtomically(configPath: string, content: string): Prom } } +/** + * Resolve what the gateway serves, falling back to the curated catalog when + * discovery fails so a transient backend outage cannot leave VS Code with no + * CodeMie models at all. + */ +async function discoverModelIds( + proxyUrl: string, + gatewayKey: string +): Promise { + try { + const ids = await fetchGatewayModelIds(proxyUrl, gatewayKey); + if (ids.length > 0) return ids; + logger.warn('[proxy] Gateway returned no models — falling back to the curated VS Code catalog'); + } catch (error) { + logger.warn( + '[proxy] Gateway model discovery failed — falling back to the curated VS Code catalog', + ...sanitizeLogArgs({ + error: error instanceof Error ? error.message : String(error), + }) + ); + } + return undefined; +} + export async function writeVsCodeLanguageModelsConfig( proxyUrl: string, + gatewayKey: string, insiders = false ): Promise { + const configPath = getVsCodeLanguageModelsPath(insiders); return writeVsCodeLanguageModelsConfigAtPath( - getVsCodeLanguageModelsPath(insiders), - proxyUrl + configPath, + proxyUrl, + await discoverModelIds(proxyUrl, gatewayKey) ); } export async function writeVsCodeLanguageModelsConfigAtPath( configPath: string, - proxyUrl: string + proxyUrl: string, + availableModelIds?: readonly string[] ): Promise { const providers = await readProviders(configPath); const managedProviderIndexes = providers @@ -245,7 +342,7 @@ export async function writeVsCodeLanguageModelsConfigAtPath( .map(index => providers[index]) .filter(isManagedProvider); const { provider: managedProvider, requiresSecretConfiguration } = - mergeManagedProviders(managedProviders, proxyUrl); + mergeManagedProviders(managedProviders, proxyUrl, availableModelIds); const firstManagedProviderIndex = managedProviderIndexes[0] ?? providers.length; const managedProviderIndexSet = new Set(managedProviderIndexes); const reconciledProviders = providers.flatMap((provider, index) => { @@ -264,5 +361,18 @@ export async function writeVsCodeLanguageModelsConfigAtPath( ); } - return { configPath, requiresSecretConfiguration }; + const modelIds = (managedProvider.models as VsCodeManagedModel[]).map(model => model.id); + logger.info( + '[proxy] VS Code managed model catalog resolved', + ...sanitizeLogArgs({ + configPath, + proxyUrl, + discoveredModelCount: availableModelIds?.length ?? 0, + usedCuratedFallback: availableModelIds === undefined, + writtenModelCount: modelIds.length, + writtenModels: modelIds, + }) + ); + + return { configPath, requiresSecretConfiguration, modelIds }; }