diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 501dea50b82..a3cc65d5919 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -206,6 +206,42 @@ vi.mock('@qwen-code/qwen-code-core', () => ({ provider.ownsModel ?? ((model: { envKey?: string }) => model.envKey === provider.envKey), ), + findExistingProviderModels: vi.fn( + ( + provider: { + envKey?: string | ((...args: unknown[]) => string); + protocol: string; + protocolOptions?: string[]; + ownsModel?: (model: { envKey?: string }) => boolean; + }, + modelProviders: Record | undefined, + ) => { + const ownsModel = + provider.ownsModel ?? + (typeof provider.envKey === 'string' + ? (model: { envKey?: string }) => model.envKey === provider.envKey + : undefined); + if (!ownsModel || !modelProviders) return undefined; + const protocols = + provider.protocolOptions && provider.protocolOptions.length > 0 + ? provider.protocolOptions + : [provider.protocol]; + for (const protocol of protocols) { + const raw = modelProviders[protocol]; + if (!Array.isArray(raw)) continue; + const models = raw + .filter( + (m): m is { id: string; envKey?: string } => + typeof m === 'object' && + m !== null && + typeof (m as { id?: unknown }).id === 'string', + ) + .filter(ownsModel); + if (models.length > 0) return { protocol, models }; + } + return undefined; + }, + ), ExtensionManager: vi.fn().mockImplementation(() => ({ refreshCache: mockExtensionManagerState.refreshCache, getLoadedExtensions: vi.fn(() => mockExtensionManagerState.extensions), diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index df6292db961..0179f627fc8 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -39,7 +39,7 @@ import { MCPServerStatus, McpTransportPool, POOLED_TRANSPORTS_DEFAULT, - resolveOwnsModel, + findExistingProviderModels, ExtensionManager, ExtensionSettingScope, HookEventName, @@ -1431,11 +1431,6 @@ function resolveProviderDocumentationUrl( return undefined; } -function isProviderModelConfig(value: unknown): value is ProviderModelConfig { - const record = toRecord(value); - return typeof record['id'] === 'string'; -} - function readSettingsEnv( settings: LoadedSettings, envKey: string | undefined, @@ -1446,35 +1441,6 @@ function readSettingsEnv( return typeof value === 'string' && value.length > 0 ? value : undefined; } -function readProviderModels( - settings: LoadedSettings, - protocol: string, -): ProviderModelConfig[] { - const modelProviders = toRecord( - (settings.merged as Record)['modelProviders'], - ); - const models = modelProviders[protocol]; - return Array.isArray(models) ? models.filter(isProviderModelConfig) : []; -} - -function findExistingProviderModels( - config: ProviderConfig, - settings: LoadedSettings, -): - | { protocol: ProviderConfig['protocol']; models: ProviderModelConfig[] } - | undefined { - const ownsModel = resolveOwnsModel(config); - if (!ownsModel) return undefined; - const protocols = config.protocolOptions?.length - ? config.protocolOptions - : [config.protocol]; - for (const protocol of protocols) { - const models = readProviderModels(settings, protocol).filter(ownsModel); - if (models.length > 0) return { protocol, models }; - } - return undefined; -} - function resolveProviderEnvKey( config: ProviderConfig, protocol: ProviderConfig['protocol'], @@ -1508,7 +1474,10 @@ function readExistingProviderConfig( config: ProviderConfig, settings: LoadedSettings, ): Record | undefined { - const existing = findExistingProviderModels(config, settings); + const existing = findExistingProviderModels( + config, + toRecord((settings.merged as Record)['modelProviders']), + ); const firstModel = existing?.models[0]; const protocol = existing?.protocol ?? config.protocol; const baseUrl = diff --git a/packages/cli/src/ui/auth/AuthDialog.test.tsx b/packages/cli/src/ui/auth/AuthDialog.test.tsx index 86dfa086055..4239f18085f 100644 --- a/packages/cli/src/ui/auth/AuthDialog.test.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.test.tsx @@ -1255,6 +1255,78 @@ describe('AuthDialog', { timeout: 15000 }, () => { }, ); + itWhenTuiInputReliable( + 'should pre-fill the Model IDs step with previously saved custom model IDs', + async () => { + // User previously saved a custom model ID for Token Plan in settings. + const savedSettings = { + security: { auth: { selectedType: undefined } }, + ui: { customThemes: {} }, + mcpServers: {}, + modelProviders: { + openai: [ + { + id: 'my-custom-token-model', + name: '[ModelStudio Token Plan] my-custom-token-model', + baseUrl: 'https://dashscope.aliyuncs.com/compatible-mode/v1', + envKey: 'BAILIAN_TOKEN_PLAN_API_KEY', + }, + ], + }, + }; + const settings: LoadedSettings = new LoadedSettings( + { + settings: { ui: { customThemes: {} }, mcpServers: {} }, + originalSettings: { ui: { customThemes: {} }, mcpServers: {} }, + path: '', + }, + { + settings: {}, + originalSettings: {}, + path: '', + }, + { + settings: savedSettings, + originalSettings: savedSettings, + path: '', + }, + { + settings: { ui: { customThemes: {} }, mcpServers: {} }, + originalSettings: { ui: { customThemes: {} }, mcpServers: {} }, + path: '', + }, + true, + new Set(), + ); + + const { stdin, lastFrame, unmount } = renderAuthDialog(settings); + + await waitForSelectedOption(lastFrame, 'Alibaba ModelStudio'); + stdin.write('\r'); + await waitForSelectedOption(lastFrame, 'Coding Plan'); + await moveDownAndWaitForSelection(stdin, lastFrame, 'Token Plan'); + await pressEnterAndWaitFor( + stdin, + lastFrame, + 'Alibaba ModelStudio · Step 1/2 · API Key', + ); + + await typeText(stdin, 'sk-token-plan'); + + await pressEnterAndWaitFor( + stdin, + lastFrame, + 'Alibaba ModelStudio · Step 2/2 · Model IDs', + ); + + // The Model IDs input is pre-filled with the saved custom model id + // (which only exists in settings, never among the built-in defaults). + expect(lastFrame()).toContain('my-custom-token-model'); + + unmount(); + }, + ); + itWhenTuiInputReliable( 'should return from Token Plan API key input to Token Plan selection', async () => { diff --git a/packages/cli/src/ui/auth/AuthDialog.tsx b/packages/cli/src/ui/auth/AuthDialog.tsx index 69616891996..7817b648bfa 100644 --- a/packages/cli/src/ui/auth/AuthDialog.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.tsx @@ -19,6 +19,7 @@ import { t } from '../../i18n/index.js'; import { findProviderById, findProviderByCredentials, + findExistingProviderModels, customProvider, ALIBABA_PROVIDERS, THIRD_PARTY_PROVIDERS, @@ -171,11 +172,25 @@ export function AuthDialog(): React.JSX.Element { const existingEnv = (settings.merged.env ?? {}) as Record; + // Model IDs already saved for this provider in settings.json (including any + // custom ones), so re-entering the wizard pre-fills them instead of resetting + // to the built-in defaults and overwriting them on submit. + const existingModelIds = (providerConfig: ProviderConfig): string[] => + findExistingProviderModels( + providerConfig, + settings.merged.modelProviders as Record | undefined, + )?.models.map((model) => model.id) ?? []; + const handleProviderSelect = (providerId: string) => { clearErrors(); const providerConfig = findProviderById(providerId); if (!providerConfig) return; - setupFlow.start(providerConfig, undefined, existingEnv); + setupFlow.start( + providerConfig, + undefined, + existingEnv, + existingModelIds(providerConfig), + ); pushView('provider-setup'); }; @@ -228,7 +243,12 @@ export function AuthDialog(): React.JSX.Element { pushView('thirdparty-select'); break; case 'CUSTOM_PROVIDER': - setupFlow.start(customProvider, undefined, existingEnv); + setupFlow.start( + customProvider, + undefined, + existingEnv, + existingModelIds(customProvider), + ); pushView('provider-setup'); break; default: diff --git a/packages/cli/src/ui/auth/useProviderSetupFlow.ts b/packages/cli/src/ui/auth/useProviderSetupFlow.ts index cb750cac767..34e6d43dd2a 100644 --- a/packages/cli/src/ui/auth/useProviderSetupFlow.ts +++ b/packages/cli/src/ui/auth/useProviderSetupFlow.ts @@ -130,6 +130,7 @@ export function useProviderSetupFlow( config: ProviderConfig, initialProtocol?: AuthType, existingEnv?: Record, + existingModelIds?: string[], ) => { setProvider(config); const steps = getVisibleSteps(config); @@ -160,7 +161,14 @@ export function useProviderSetupFlow( setApiKey(prefillKey); setApiKeyError(null); - setModelIds(getDefaultModelIds(config).join(', ')); + // Pre-fill with the user's previously saved model IDs (including custom + // ones) when present, so re-entering the wizard doesn't reset to — and + // later overwrite with — the provider's built-in defaults. + const initialModelIds = + existingModelIds && existingModelIds.length > 0 + ? existingModelIds + : getDefaultModelIds(config); + setModelIds(initialModelIds.join(', ')); setModelIdsError(null); setThinkingEnabled(false); setModalityEnabled(false); diff --git a/packages/core/src/providers/__tests__/provider-config.test.ts b/packages/core/src/providers/__tests__/provider-config.test.ts index 88217dac748..9cf902c8d88 100644 --- a/packages/core/src/providers/__tests__/provider-config.test.ts +++ b/packages/core/src/providers/__tests__/provider-config.test.ts @@ -10,6 +10,7 @@ import { buildInstallPlan, buildProviderTemplate, computeModelListVersion, + findExistingProviderModels, findProviderByCredentials, getAllProviderBaseUrls, getDefaultModelIds, @@ -365,6 +366,68 @@ describe('getDefaultModelIds', () => { }); }); +describe('findExistingProviderModels', () => { + const config = makeConfig({ modelNamePrefix: '', envKey: 'TEST_API_KEY' }); + + it('returns the user-saved models owned by the provider', () => { + const result = findExistingProviderModels(config, { + [AuthType.USE_OPENAI]: [ + { id: 'custom-model', envKey: 'TEST_API_KEY' }, + { id: 'default-model', envKey: 'TEST_API_KEY' }, + { id: 'other-provider-model', envKey: 'OTHER_API_KEY' }, + ], + }); + expect(result).toEqual({ + protocol: AuthType.USE_OPENAI, + models: [ + { id: 'custom-model', envKey: 'TEST_API_KEY' }, + { id: 'default-model', envKey: 'TEST_API_KEY' }, + ], + }); + }); + + it('returns undefined when no saved models are owned by the provider', () => { + expect( + findExistingProviderModels(config, { + [AuthType.USE_OPENAI]: [{ id: 'x', envKey: 'OTHER_API_KEY' }], + }), + ).toBeUndefined(); + }); + + it('returns undefined when modelProviders is empty or missing', () => { + expect(findExistingProviderModels(config, {})).toBeUndefined(); + expect(findExistingProviderModels(config, undefined)).toBeUndefined(); + }); + + it('returns undefined when ownership cannot be resolved (function envKey)', () => { + const customConfig = makeConfig({ + envKey: () => 'DYNAMIC_KEY', + modelNamePrefix: '', + }); + expect( + findExistingProviderModels(customConfig, { + [AuthType.USE_OPENAI]: [{ id: 'x', envKey: 'DYNAMIC_KEY' }], + }), + ).toBeUndefined(); + }); + + it('scans protocolOptions in order and picks the first with owned models', () => { + const multiProtocol = makeConfig({ + modelNamePrefix: '', + envKey: 'TEST_API_KEY', + protocolOptions: [AuthType.USE_ANTHROPIC, AuthType.USE_OPENAI], + }); + const result = findExistingProviderModels(multiProtocol, { + [AuthType.USE_OPENAI]: [{ id: 'openai-model', envKey: 'TEST_API_KEY' }], + [AuthType.USE_ANTHROPIC]: [ + { id: 'anthropic-model', envKey: 'TEST_API_KEY' }, + ], + }); + expect(result?.protocol).toBe(AuthType.USE_ANTHROPIC); + expect(result?.models.map((m) => m.id)).toEqual(['anthropic-model']); + }); +}); + describe('shouldShowStep', () => { it('shows protocol step only when multiple options', () => { const single = makeConfig({ diff --git a/packages/core/src/providers/index.ts b/packages/core/src/providers/index.ts index 00fd811d8c9..30472e85508 100644 --- a/packages/core/src/providers/index.ts +++ b/packages/core/src/providers/index.ts @@ -23,6 +23,7 @@ export { buildInstallPlan, buildProviderTemplate, computeModelListVersion, + findExistingProviderModels, getDefaultBaseUrlForProtocol, getDefaultModelIds, providerMatchesCredentials, diff --git a/packages/core/src/providers/provider-config.ts b/packages/core/src/providers/provider-config.ts index ad7a0cfba3a..d1f81682b0d 100644 --- a/packages/core/src/providers/provider-config.ts +++ b/packages/core/src/providers/provider-config.ts @@ -355,6 +355,41 @@ export function getDefaultModelIds(config: ProviderConfig): string[] { return config.models?.map((s) => s.id) ?? []; } +function isProviderModelConfig(value: unknown): value is ProviderModelConfig { + return ( + typeof value === 'object' && + value !== null && + typeof (value as { id?: unknown }).id === 'string' + ); +} + +/** + * Find the model entries a user has already saved for `config` under the + * `modelProviders` map in settings. Returns the first protocol (in the + * provider's own preference order) that owns stored models, or `undefined` + * when none are saved. Used to pre-fill the auth wizard / connect form with + * existing model IDs instead of resetting to the provider's built-in defaults. + */ +export function findExistingProviderModels( + config: ProviderConfig, + modelProviders: Record | undefined, +): + | { protocol: ProviderConfig['protocol']; models: ProviderModelConfig[] } + | undefined { + const ownsModel = resolveOwnsModel(config); + if (!ownsModel || !modelProviders) return undefined; + const protocols = config.protocolOptions?.length + ? config.protocolOptions + : [config.protocol]; + for (const protocol of protocols) { + const raw = modelProviders[protocol]; + if (!Array.isArray(raw)) continue; + const models = raw.filter(isProviderModelConfig).filter(ownsModel); + if (models.length > 0) return { protocol, models }; + } + return undefined; +} + // --------------------------------------------------------------------------- // Check if a step should be shown in the UI // ---------------------------------------------------------------------------