diff --git a/README.md b/README.md index bd65dfe8..317e8c20 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ NeoCode 是一个在终端中运行的 AI 编码助手，采用 ReAct（Reason-A - **流式输出** — 实时展示模型思考过程 - **工具系统** — 文件操作、代码执行、搜索等内置工具 -- **多 Provider 支持** — OpenAI、Gemini 等主流模型，易于扩展 +- **内建 Provider 支持** — OpenAI、Gemini、OpenLL、Qiniu，模型列表支持动态发现与缓存 - **终端原生体验** — 基于 Bubble Tea 的现代化 TUI ## 一键安装 @@ -66,7 +66,7 @@ go run ./cmd/neocode 使用 slash 命令快速切换配置： -- `/provider` — 切换模型提供商 +- `/provider` — 切换内建模型提供商 - `/model` — 切换模型 ## 架构概览 @@ -89,8 +89,8 @@ go run ./cmd/neocode 核心模块职责： - **`internal/config`** — 配置管理、环境变量、YAML 加载 -- **`internal/context`** — 多 section system prompt 组装、消息裁剪与上下文构建 -- **`internal/provider`** — LLM 提供商抽象，抹平厂商差异 +- **`internal/context`** — system prompt、消息裁剪与上下文构建 +- **`internal/provider`** — 内建 LLM Provider 抽象、模型元数据归一化、动态发现与 catalog 缓存 - **`internal/runtime`** — ReAct 主循环、事件流、会话管理 - **`internal/tools`** — 工具注册表与具体工具实现 - **`internal/tui`** — 终端 UI、交互体验、事件桥接 @@ -107,8 +107,8 @@ go run ./cmd/neocode ├── internal │ ├── app # 应用装配 │ ├── config # 配置管理 -│ ├── context # system prompt 组装与上下文构建 -│ ├── provider # Provider 抽象与实现 +│ ├── context # 上下文构建 +│ ├── provider # Provider 抽象、模型发现与元数据缓存 │ ├── runtime # ReAct 循环与事件流 │ ├── tools # 工具系统 │ └── tui # 终端 UI diff --git a/docs/guides/adding-providers.md b/docs/guides/adding-providers.md index 8880a882..db3a0995 100644 --- a/docs/guides/adding-providers.md +++ b/docs/guides/adding-providers.md @@ -10,12 +10,12 @@ NeoCode 的 provider 架构采用**集中式配置 + 驱动复用**的设计： - **配置层**（`internal/config/builtin_providers.go`）：集中管理所有 provider 的元数据 - Provider 名称、Driver 名称 - - Base URL、默认模型、模型列表 - - API Key 环境变量名 + - Base URL、默认模型、API Key 环境变量名 - **驱动层**（`internal/provider/openai/` 等）：负责实际的 API 协议实现 - 请求构造、响应解析 - 流式输出、Tool Call 处理 + - 动态模型发现（如 `GET /models`）与结果归一化 ### 驱动复用 @@ -46,11 +46,6 @@ const ( DeepSeekDefaultAPIKeyEnv = "DEEPSEEK_API_KEY" ) -var deepSeekModels = []string{ - DeepSeekDefaultModel, - "deepseek-coder", -} - // DeepSeekProvider 返回 DeepSeek provider 的默认配置。 func DeepSeekProvider() ProviderConfig { return ProviderConfig{ @@ -58,7 +53,6 @@ func DeepSeekProvider() ProviderConfig { Driver: "openai", // 复用 openai 驱动 BaseURL: DeepSeekDefaultBaseURL, Model: DeepSeekDefaultModel, - Models: append([]string(nil), deepSeekModels...), APIKeyEnv: DeepSeekDefaultAPIKeyEnv, } } @@ -182,19 +176,12 @@ const ( AnthropicDefaultAPIKeyEnv = "ANTHROPIC_API_KEY" ) -var anthropicModels = []string{ - AnthropicDefaultModel, - "claude-opus-4-20250514", - "claude-3-5-sonnet-20241022", -} - func AnthropicProvider() ProviderConfig { return ProviderConfig{ Name: AnthropicName, Driver: "anthropic", // 使用新的 anthropic 驱动 BaseURL: AnthropicDefaultBaseURL, Model: AnthropicDefaultModel, - Models: append([]string(nil), anthropicModels...), APIKeyEnv: AnthropicDefaultAPIKeyEnv, } } @@ -226,7 +213,7 @@ func DefaultProviders() []ProviderConfig { | `ChatRequest` | `internal/provider/types.go` | 请求：`Model`、`SystemPrompt`、`Messages`、`Tools` | | `ChatResponse` | `internal/provider/types.go` | 响应：`Message`、`FinishReason`、`Usage` | | `StreamEvent` | `internal/provider/types.go` | 流式事件：`TextDelta`、`ToolCallStart` | -| `ProviderConfig` | `internal/config/model.go` | 配置：`Name`、`Driver`、`BaseURL`、`Model`、`Models`、`APIKeyEnv` | +| `ProviderConfig` | `internal/config/model.go` | 配置：`Name`、`Driver`、`BaseURL`、`Model`、`APIKeyEnv` | ## 设计约束 @@ -242,7 +229,7 @@ func DefaultProviders() []ProviderConfig { ✅ **驱动职责清晰** - 驱动只负责协议构造与响应解析 -- 不持有模型列表、base_url 等配置信息 +- 不持有 provider 元数据；模型目录由 driver 发现，缓存与合并由 service 处理 ✅ **架构分层** - 厂商差异收敛在 `internal/provider/` 内 @@ -274,9 +261,10 @@ func DefaultProviders() []ProviderConfig { OpenAIProvider(), // OpenAI 官方 API GeminiProvider(), // Google Gemini (OpenAI-compatible) OpenLLProvider(), // OpenLL 服务 (OpenAI-compatible) + QiniuProvider(), // 七牛云推理服务 (OpenAI-compatible) } } ``` -所有内置 provider 都复用 `openai` 驱动，配置集中在 `builtin_providers.go`。 +所有内置 provider 都通过代码集中注册。模型选择器展示的候选模型由默认模型、动态发现结果和本地缓存共同组成。 diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 226b9e5b..b668c8be 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -16,6 +16,7 @@ NeoCode 采用"内置 Provider 优先"的策略： - `config.yaml` 不再持久化完整 `providers` 列表 - 只保存当前选择状态和通用运行配置 - 运行时的 `providers` 完全来自代码内置定义 +- 用户不能通过 YAML 注入新的 provider ✅ **安全第一** - API Key 只从环境变量读取，永不写入 YAML @@ -28,6 +29,7 @@ NeoCode 采用"内置 Provider 优先"的策略： | `openai` | `openai` | OpenAI 官方 API | | `gemini` | `openai` | Google Gemini (OpenAI-compatible API) | | `openll` | `openai` | OpenLL 服务 (OpenAI-compatible API) | +| `qiniu` | `openai` | 七牛云推理服务 (OpenAI-compatible API) | 所有内置 provider 都复用 `openai` 驱动，支持流式输出和 Tool Call。 @@ -39,6 +41,7 @@ NeoCode 采用"内置 Provider 优先"的策略： - ✅ 未来代码新增 provider 时，用户无需修改 YAML - ✅ 老配置文件中的 `providers` / `provider_overrides` 会在加载时被清理 - ✅ 配置文件始终保持简洁，只包含必要的运行时状态 +- ✅ 如需新增 provider，只能通过代码扩展内建列表完成 ## 配置文件 @@ -106,6 +109,7 @@ tools: | `openai` | `OPENAI_API_KEY` | [OpenAI Platform](https://platform.openai.com/api-keys) | | `gemini` | `GEMINI_API_KEY` | [Google AI Studio](https://aistudio.google.com/apikey) | | `openll` | `AI_API_KEY` | OpenLL 服务提供商 | +| `qiniu` | `QINIU_API_KEY` | 七牛云推理服务 | ### 环境变量管理 @@ -156,6 +160,7 @@ NeoCode 提供以下 slash 命令用于快速切换配置： ❯ openai gemini openll + qiniu ``` ### /model - 模型选择器 @@ -166,6 +171,8 @@ NeoCode 提供以下 slash 命令用于快速切换配置： 打开当前 provider 的模型选择器。 +模型列表会合并当前 provider 的默认模型、服务端动态发现结果和本地缓存结果。 + **界面示例**（选择 openai provider 后）： ``` ? Select a model: @@ -208,9 +215,9 @@ NeoCode 提供以下 slash 命令用于快速切换配置： 2. 立即持久化到 `config.yaml` 3. 下次启动自动恢复选择状态 -## 添加自定义 Provider +## 扩展内建 Provider -如需添加自定义 provider（如企业内部服务），请参考： +如需添加新的内建 provider（包括企业内部服务），请参考： 👉 [adding-providers.md](./adding-providers.md) diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 3acc7af7..1ccd1462 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -22,22 +22,26 @@ import ( func NewProgram(ctx context.Context) (*tea.Program, error) { loader := config.NewLoader("", builtin.DefaultConfig()) manager := config.NewManager(loader) - cfg, err := manager.Load(ctx) - if err != nil { + if _, err := manager.Load(ctx); err != nil { return nil, err } - toolRegistry := buildToolRegistry(cfg) - toolManager, err := buildToolManager(toolRegistry) + providerRegistry, err := builtin.NewRegistry() if err != nil { return nil, err } + providerService := provider.NewService(manager, providerRegistry, nil) + if _, err := providerService.EnsureSelection(ctx); err != nil { + return nil, err + } - providerRegistry, err := builtin.NewRegistry() + cfg := manager.Get() + + toolRegistry := buildToolRegistry(cfg) + toolManager, err := buildToolManager(toolRegistry) if err != nil { return nil, err } - providerService := provider.NewService(manager, providerRegistry) sessionStore := agentruntime.NewSessionStore(loader.BaseDir()) runtimeSvc := agentruntime.NewWithFactory( diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 8b2194bb..1743a0a6 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -7,6 +7,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strings" "testing" "neo-code/internal/config" @@ -14,6 +15,8 @@ import ( ) func TestNewProgram(t *testing.T) { + disableBuiltinProviderAPIKeys(t) + home := t.TempDir() t.Setenv("HOME", home) t.Setenv("USERPROFILE", home) @@ -32,6 +35,41 @@ func TestNewProgram(t *testing.T) { } } +func TestNewProgramNormalizesInvalidCurrentModelOnStartup(t *testing.T) { + disableBuiltinProviderAPIKeys(t) + + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + configDir := filepath.Join(home, ".neocode") + if err := os.MkdirAll(configDir, 0o755); err != nil { + t.Fatalf("mkdir config dir: %v", err) + } + + configPath := filepath.Join(configDir, "config.yaml") + raw := []byte("selected_provider: openai\ncurrent_model: unsupported-current\nworkdir: .\nshell: powershell\n") + if err := os.WriteFile(configPath, raw, 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + program, err := NewProgram(context.Background()) + if err != nil { + t.Fatalf("NewProgram() error = %v", err) + } + if program == nil { + t.Fatalf("expected tea program") + } + + data, err := os.ReadFile(configPath) + if err != nil { + t.Fatalf("read config: %v", err) + } + if !strings.Contains(string(data), "current_model: "+config.OpenAIDefaultModel) { + t.Fatalf("expected startup normalization to rewrite current_model, got:\n%s", string(data)) + } +} + func TestBuildToolRegistryUsesWebFetchConfig(t *testing.T) { t.Parallel() @@ -126,3 +164,11 @@ func (s stubToolForBootstrap) Schema() map[string]any { return map[string]any{"t func (s stubToolForBootstrap) Execute(ctx context.Context, call tools.ToolCallInput) (tools.ToolResult, error) { return tools.ToolResult{Name: s.name, Content: s.content}, nil } + +func disableBuiltinProviderAPIKeys(t *testing.T) { + t.Helper() + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "") + t.Setenv(config.GeminiDefaultAPIKeyEnv, "") + t.Setenv(config.OpenLLDefaultAPIKeyEnv, "") + t.Setenv(config.QiniuDefaultAPIKeyEnv, "") +} diff --git a/internal/config/builtin_providers.go b/internal/config/builtin_providers.go index 90916086..57d1a2eb 100644 --- a/internal/config/builtin_providers.go +++ b/internal/config/builtin_providers.go @@ -22,78 +22,51 @@ const ( QiniuDefaultAPIKeyEnv = "QINIU_API_KEY" ) -var openAIModels = []string{ - OpenAIDefaultModel, - "gpt-4o", - "gpt-5.4", - "gpt-5.3-codex", -} - -var geminiModels = []string{ - GeminiDefaultModel, - "gemini-2.5-pro", - "gemini-2.0-flash", -} - -var openLLModels = []string{ - OpenLLDefaultModel, - "gpt-5.3-codex", - "gpt-5.3-turbo", -} - -var qiniuModels = []string{ - QiniuDefaultModel, - "openai/gpt-5.2-codex", -} - -// OpenAIProvider 返回 OpenAI provider 的默认配置。 +// OpenAIProvider returns the builtin OpenAI provider definition. func OpenAIProvider() ProviderConfig { return ProviderConfig{ Name: OpenAIName, Driver: "openai", BaseURL: OpenAIDefaultBaseURL, Model: OpenAIDefaultModel, - Models: append([]string(nil), openAIModels...), APIKeyEnv: OpenAIDefaultAPIKeyEnv, } } -// GeminiProvider 返回 Gemini provider 的默认配置。 +// GeminiProvider returns the builtin Gemini provider definition. func GeminiProvider() ProviderConfig { return ProviderConfig{ Name: GeminiName, Driver: "openai", BaseURL: GeminiDefaultBaseURL, Model: GeminiDefaultModel, - Models: append([]string(nil), geminiModels...), APIKeyEnv: GeminiDefaultAPIKeyEnv, } } -// OpenLLProvider 返回 OpenLL provider 的默认配置。 +// OpenLLProvider returns the builtin OpenLL provider definition. func OpenLLProvider() ProviderConfig { return ProviderConfig{ Name: OpenLLName, Driver: "openai", BaseURL: OpenLLDefaultBaseURL, Model: OpenLLDefaultModel, - Models: append([]string(nil), openLLModels...), APIKeyEnv: OpenLLDefaultAPIKeyEnv, } } +// QiniuProvider returns the builtin Qiniu provider definition. func QiniuProvider() ProviderConfig { return ProviderConfig{ Name: QiniuName, Driver: "openai", BaseURL: QiniuDefaultBaseURL, Model: QiniuDefaultModel, - Models: append([]string(nil), qiniuModels...), APIKeyEnv: QiniuDefaultAPIKeyEnv, } } -// DefaultProviders 返回所有内置 provider 配置列表。 +// DefaultProviders returns all builtin provider definitions. func DefaultProviders() []ProviderConfig { return []ProviderConfig{ OpenAIProvider(), diff --git a/internal/config/builtin_providers_test.go b/internal/config/builtin_providers_test.go index 5234e4f7..6309a671 100644 --- a/internal/config/builtin_providers_test.go +++ b/internal/config/builtin_providers_test.go @@ -40,12 +40,6 @@ func TestOpenAIProviderConfig(t *testing.T) { if provider.APIKeyEnv != OpenAIDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", OpenAIDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) == 0 { - t.Fatal("expected non-empty models list") - } - if !ContainsModelID(provider.Models, OpenAIDefaultModel) { - t.Fatalf("expected models to contain default model %q", OpenAIDefaultModel) - } } func TestGeminiProviderConfig(t *testing.T) { @@ -68,12 +62,6 @@ func TestGeminiProviderConfig(t *testing.T) { if provider.APIKeyEnv != GeminiDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", GeminiDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) == 0 { - t.Fatal("expected non-empty models list") - } - if !ContainsModelID(provider.Models, GeminiDefaultModel) { - t.Fatalf("expected models to contain default model %q", GeminiDefaultModel) - } } func TestOpenLLProviderConfig(t *testing.T) { @@ -96,27 +84,20 @@ func TestOpenLLProviderConfig(t *testing.T) { if provider.APIKeyEnv != OpenLLDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", OpenLLDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) == 0 { - t.Fatal("expected non-empty models list") - } - if !ContainsModelID(provider.Models, OpenLLDefaultModel) { - t.Fatalf("expected models to contain default model %q", OpenLLDefaultModel) - } } -func TestProviderModelsAreImmutable(t *testing.T) { +func TestProviderDefaultsAreIndependent(t *testing.T) { t.Parallel() - // Verify that modifying returned models slice doesn't affect future calls provider1 := OpenAIProvider() - provider1.Models[0] = "modified-model" + provider1.Model = "modified-model" provider2 := OpenAIProvider() - if provider2.Models[0] == "modified-model" { - t.Fatal("expected models slice to be independent between calls") + if provider2.Model == "modified-model" { + t.Fatal("expected builtin provider defaults to be independent between calls") } - if provider2.Models[0] != OpenAIDefaultModel { - t.Fatalf("expected first model %q, got %q", OpenAIDefaultModel, provider2.Models[0]) + if provider2.Model != OpenAIDefaultModel { + t.Fatalf("expected default model %q, got %q", OpenAIDefaultModel, provider2.Model) } } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 97a677b9..cc0d4ead 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -22,7 +22,6 @@ func testDefaultProviderConfig() ProviderConfig { Driver: testProviderName, BaseURL: testBaseURL, Model: testModel, - Models: []string{testModel, "gpt-4o", "gpt-5.4", "gpt-5.3-codex"}, APIKeyEnv: testAPIKeyEnv, } } @@ -120,11 +119,11 @@ providers: }, }, { - name: "legacy format", + name: "legacy fields are ignored", data: ` selected_provider: openai current_model: gpt-4o -workspace_root: . +workspace_root: ./definitely-legacy-root shell: bash max_loop: 5 providers: @@ -137,8 +136,8 @@ providers: `, assert: func(t *testing.T, cfg *Config) { t.Helper() - if cfg.MaxLoops != 5 { - t.Fatalf("expected max loops 5, got %d", cfg.MaxLoops) + if cfg.MaxLoops != DefaultMaxLoops { + t.Fatalf("expected legacy max_loop to be ignored, got %d", cfg.MaxLoops) } provider, err := cfg.SelectedProviderConfig() if err != nil { @@ -150,8 +149,8 @@ providers: if cfg.CurrentModel != "gpt-4o" { t.Fatalf("expected current model %q, got %q", "gpt-4o", cfg.CurrentModel) } - if len(provider.Models) != len(testDefaultProviderConfig().Models) { - t.Fatalf("expected builtin provider models, got %+v", provider.Models) + if strings.Contains(filepath.ToSlash(cfg.Workdir), "definitely-legacy-root") { + t.Fatalf("expected legacy workspace_root to be ignored, got %q", cfg.Workdir) } if cfg.Tools.WebFetch.MaxResponseBytes != DefaultWebFetchMaxResponseBytes { t.Fatalf("expected default max_response_bytes %d, got %d", DefaultWebFetchMaxResponseBytes, cfg.Tools.WebFetch.MaxResponseBytes) @@ -167,7 +166,7 @@ providers: tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - cfg, err := parseConfig([]byte(tt.data), *testDefaultConfig()) + cfg, err := parseConfig([]byte(tt.data)) if err != nil { t.Fatalf("parseConfig() error = %v", err) } @@ -296,7 +295,6 @@ func TestManagerConcurrentAccess(t *testing.T) { for k := range next.Providers { if next.Providers[k].Name == next.SelectedProvider { next.Providers[k].Model = model - next.Providers[k].Models = append([]string(nil), models...) } } return nil @@ -434,13 +432,19 @@ func TestConfigValidateFailures(t *testing.T) { expectErr: "supported_content_types[0] is empty", }, { - name: "current model unsupported by selected provider", + name: "duplicate provider endpoints after normalization", config: func() *Config { cfg := validConfig.Clone() - cfg.CurrentModel = "unsupported-model" + cfg.Providers = append(cfg.Providers, ProviderConfig{ + Name: "openai-shadow", + Driver: "OPENAI", + BaseURL: "https://API.OPENAI.COM/v1/", + Model: "shadow-model", + APIKeyEnv: "OPENAI_SHADOW_KEY", + }) return &cfg }(), - expectErr: "current_model \"unsupported-model\" is not supported", + expectErr: "duplicate provider endpoint", }, } @@ -503,18 +507,6 @@ func TestProviderConfigValidateFailures(t *testing.T) { }, expectErr: "api_key_env is empty", }, - { - name: "default model missing from models", - provider: ProviderConfig{ - Name: testProviderName, - Driver: testProviderName, - BaseURL: testBaseURL, - Model: testModel, - Models: []string{"gpt-4o"}, - APIKeyEnv: testAPIKeyEnv, - }, - expectErr: "default model \"gpt-4.1\" is not in models", - }, } for _, tt := range tests { @@ -568,7 +560,6 @@ func TestLoaderLoadAndSaveRoundTrip(t *testing.T) { } cfg.CurrentModel = "gpt-5.4" - cfg.Providers[0].Models = []string{"gpt-4.1", "gpt-5.4"} cfg.Providers[0].BaseURL = "https://ignored.example/v1" cfg.Tools.WebFetch.MaxResponseBytes = 1024 cfg.Tools.WebFetch.SupportedContentTypes = []string{"text/html", "application/json"} @@ -608,9 +599,6 @@ func TestLoaderLoadAndSaveRoundTrip(t *testing.T) { if provider.BaseURL != testBaseURL { t.Fatalf("expected provider base url to come from builtin definition, got %q", provider.BaseURL) } - if len(provider.Models) != len(testDefaultProviderConfig().Models) { - t.Fatalf("expected provider models to come from builtin definition, got %+v", provider.Models) - } if reloaded.Tools.WebFetch.MaxResponseBytes != 1024 { t.Fatalf("expected max_response_bytes %d, got %d", 1024, reloaded.Tools.WebFetch.MaxResponseBytes) } @@ -650,25 +638,54 @@ func TestLoaderUsesUpdatedBuiltinProviderWhenUserHasNoOverride(t *testing.T) { } } -func TestApplyProviderDefaultsMatchesByNameOnly(t *testing.T) { - defaults := []ProviderConfig{testDefaultProviderConfig()} +func TestApplyDefaultsReplacesProvidersWithBuiltinSnapshot(t *testing.T) { + t.Parallel() + current := Config{ - Providers: []ProviderConfig{ - { - Name: "openai-alt", - Driver: testProviderName, - }, - }, + Providers: []ProviderConfig{{ + Name: "openai-alt", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "custom-model", + APIKeyEnv: "CUSTOM_API_KEY", + }}, + SelectedProvider: "openai-alt", + CurrentModel: "custom-model", } - current.ApplyDefaultsFrom(Config{Providers: defaults}) + current.ApplyDefaultsFrom(*testDefaultConfig()) - provider, err := current.ProviderByName("openai-alt") + if len(current.Providers) != 1 { + t.Fatalf("expected builtin provider snapshot, got %+v", current.Providers) + } + if _, err := current.ProviderByName("openai-alt"); err == nil { + t.Fatalf("expected custom provider to be dropped, got %+v", current.Providers) + } + provider, err := current.ProviderByName(testProviderName) if err != nil { - t.Fatalf("ProviderByName(openai-alt) error = %v", err) + t.Fatalf("ProviderByName(%s) error = %v", testProviderName, err) } - if provider.BaseURL != "" || provider.Model != "" || provider.APIKeyEnv != "" || len(provider.Models) != 0 { - t.Fatalf("expected provider metadata to stay empty when only driver matches, got %+v", provider) + if provider.BaseURL != testBaseURL || provider.Model != testModel || provider.APIKeyEnv != testAPIKeyEnv { + t.Fatalf("expected builtin provider metadata, got %+v", provider) + } + if current.SelectedProvider != testProviderName { + t.Fatalf("expected selected provider to reset to builtin %q, got %q", testProviderName, current.SelectedProvider) + } + if current.CurrentModel != testModel { + t.Fatalf("expected current model to reset with selected builtin provider, got %q", current.CurrentModel) + } +} + +func TestApplyDefaultsPreservesDynamicCurrentModel(t *testing.T) { + t.Parallel() + + cfg := testDefaultConfig() + cfg.CurrentModel = "server-discovered-model" + + cfg.ApplyDefaultsFrom(*testDefaultConfig()) + + if cfg.CurrentModel != "server-discovered-model" { + t.Fatalf("expected dynamic current model to be preserved, got %q", cfg.CurrentModel) } } @@ -738,13 +755,13 @@ func TestNormalizeWorkdirAndClone(t *testing.T) { cfg := testDefaultConfig() cloned := cfg.Clone() cloned.CurrentModel = "modified" - cloned.Providers[0].Models[0] = "modified-model" + cloned.Providers[0].BaseURL = "https://modified.example/v1" cloned.Tools.WebFetch.SupportedContentTypes[0] = "application/json" if cfg.CurrentModel == cloned.CurrentModel { t.Fatalf("expected clone to be independent from source") } - if cfg.Providers[0].Models[0] == cloned.Providers[0].Models[0] { - t.Fatalf("expected provider models to be cloned") + if cfg.Providers[0].BaseURL == cloned.Providers[0].BaseURL { + t.Fatalf("expected providers to be cloned") } if cfg.Tools.WebFetch.SupportedContentTypes[0] == cloned.Tools.WebFetch.SupportedContentTypes[0] { t.Fatalf("expected webfetch supported content types to be cloned") diff --git a/internal/config/loader.go b/internal/config/loader.go index 5973f662..c3563828 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -84,7 +84,7 @@ func (l *Loader) Load(ctx context.Context) (*Config, error) { return nil, fmt.Errorf("config: read config file: %w", err) } - cfg, err := parseConfig(data, l.defaults) + cfg, err := parseConfig(data) if err != nil { return nil, fmt.Errorf("config: parse config file: %w", err) } @@ -92,7 +92,11 @@ func (l *Loader) Load(ctx context.Context) (*Config, error) { if err := cfg.Validate(); err != nil { return nil, err } - if requiresConfigRewrite(data) { + needsRewrite, err := persistedConfigDiffers(data, *cfg) + if err != nil { + return nil, err + } + if needsRewrite { if err := l.Save(ctx, cfg); err != nil { return nil, err } @@ -115,23 +119,9 @@ func (l *Loader) Save(ctx context.Context, cfg *Config) error { return err } - file := persistedConfig{ - SelectedProvider: snapshot.SelectedProvider, - CurrentModel: snapshot.CurrentModel, - Workdir: snapshot.Workdir, - Shell: snapshot.Shell, - MaxLoops: snapshot.MaxLoops, - ToolTimeoutSec: snapshot.ToolTimeoutSec, - Tools: snapshot.Tools, - } - - data, err := yaml.Marshal(&file) + data, err := marshalPersistedConfig(snapshot) if err != nil { - return fmt.Errorf("config: marshal config: %w", err) - } - - if len(data) == 0 || data[len(data)-1] != '\n' { - data = append(data, '\n') + return err } if err := os.WriteFile(l.ConfigPath(), data, 0o644); err != nil { @@ -149,35 +139,20 @@ func defaultBaseDir() string { return filepath.Join(home, dirName) } -func parseConfig(data []byte, defaults Config) (*Config, error) { +func parseConfig(data []byte) (*Config, error) { if len(bytes.TrimSpace(data)) == 0 { return &Config{}, nil } - return parseCurrentConfig(data, defaults) + return parseCurrentConfig(data) } -type aliasConfig struct { - MaxLoop int `yaml:"max_loop"` - WorkspaceRoot string `yaml:"workspace_root"` -} - -func parseCurrentConfig(data []byte, _ Config) (*Config, error) { +func parseCurrentConfig(data []byte) (*Config, error) { var file persistedConfig if err := yaml.Unmarshal(data, &file); err != nil { return nil, err } - var aliases aliasConfig - if err := yaml.Unmarshal(data, &aliases); err == nil { - if file.MaxLoops == 0 && aliases.MaxLoop > 0 { - file.MaxLoops = aliases.MaxLoop - } - if strings.TrimSpace(file.Workdir) == "" && strings.TrimSpace(aliases.WorkspaceRoot) != "" { - file.Workdir = aliases.WorkspaceRoot - } - } - cfg := &Config{ SelectedProvider: strings.TrimSpace(file.SelectedProvider), CurrentModel: strings.TrimSpace(file.CurrentModel), @@ -191,18 +166,31 @@ func parseCurrentConfig(data []byte, _ Config) (*Config, error) { return cfg, nil } -func requiresConfigRewrite(data []byte) bool { - text := strings.TrimSpace(string(data)) - switch { - case text == "": - return false - case strings.Contains(text, "\nproviders:") || strings.HasPrefix(text, "providers:"): - return true - case strings.Contains(text, "provider_overrides:"): - return true - case strings.Contains(text, "workspace_root:") || strings.Contains(text, "max_loop:"): - return true - default: - return false +func marshalPersistedConfig(snapshot Config) ([]byte, error) { + file := persistedConfig{ + SelectedProvider: snapshot.SelectedProvider, + CurrentModel: snapshot.CurrentModel, + Workdir: snapshot.Workdir, + Shell: snapshot.Shell, + MaxLoops: snapshot.MaxLoops, + ToolTimeoutSec: snapshot.ToolTimeoutSec, + Tools: snapshot.Tools, + } + + data, err := yaml.Marshal(&file) + if err != nil { + return nil, fmt.Errorf("config: marshal config: %w", err) + } + if len(data) == 0 || data[len(data)-1] != '\n' { + data = append(data, '\n') + } + return data, nil +} + +func persistedConfigDiffers(data []byte, cfg Config) (bool, error) { + canonical, err := marshalPersistedConfig(cfg) + if err != nil { + return false, err } + return !bytes.Equal(bytes.TrimSpace(data), bytes.TrimSpace(canonical)), nil } diff --git a/internal/config/loader_test.go b/internal/config/loader_test.go index 1b54453a..054e9db5 100644 --- a/internal/config/loader_test.go +++ b/internal/config/loader_test.go @@ -79,9 +79,6 @@ providers: type: openai base_url: https://example.com/v1 model: gpt-5.4 - models: - - gpt-4.1 - - gpt-5.4 api_key_env: OPENAI_API_KEY ` if err := os.WriteFile(loader.ConfigPath(), []byte(strings.TrimSpace(legacy)+"\n"), 0o644); err != nil { @@ -102,9 +99,6 @@ providers: if cfg.CurrentModel != "gpt-5.4" { t.Fatalf("expected current model to stay %q, got %q", "gpt-5.4", cfg.CurrentModel) } - if len(provider.Models) != len(testDefaultProviderConfig().Models) { - t.Fatalf("expected builtin provider models, got %+v", provider.Models) - } rewritten, err := os.ReadFile(loader.ConfigPath()) if err != nil { @@ -121,3 +115,82 @@ providers: t.Fatalf("expected rewritten config to omit provider metadata, got:\n%s", text) } } + +func TestLoaderRewritesNormalizedSelectionStateOnLoad(t *testing.T) { + t.Parallel() + + loader := NewLoader(t.TempDir(), testDefaultConfig()) + if err := os.MkdirAll(loader.BaseDir(), 0o755); err != nil { + t.Fatalf("mkdir base dir: %v", err) + } + + raw := ` +selected_provider: missing-provider +workdir: . +shell: powershell +` + if err := os.WriteFile(loader.ConfigPath(), []byte(strings.TrimSpace(raw)+"\n"), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if cfg.SelectedProvider != testProviderName { + t.Fatalf("expected selected provider %q, got %q", testProviderName, cfg.SelectedProvider) + } + if cfg.CurrentModel != testModel { + t.Fatalf("expected current model %q, got %q", testModel, cfg.CurrentModel) + } + + rewritten, err := os.ReadFile(loader.ConfigPath()) + if err != nil { + t.Fatalf("read rewritten config: %v", err) + } + text := string(rewritten) + if !strings.Contains(text, "selected_provider: "+testProviderName) { + t.Fatalf("expected rewritten config to persist selected provider, got:\n%s", text) + } + if !strings.Contains(text, "current_model: "+testModel) { + t.Fatalf("expected rewritten config to persist current model, got:\n%s", text) + } +} + +func TestLoaderRewritesMissingCurrentModelOnLoad(t *testing.T) { + t.Parallel() + + loader := NewLoader(t.TempDir(), testDefaultConfig()) + if err := os.MkdirAll(loader.BaseDir(), 0o755); err != nil { + t.Fatalf("mkdir base dir: %v", err) + } + + raw := ` +selected_provider: openai +workdir: . +shell: powershell +` + if err := os.WriteFile(loader.ConfigPath(), []byte(strings.TrimSpace(raw)+"\n"), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if cfg.SelectedProvider != testProviderName { + t.Fatalf("expected selected provider %q, got %q", testProviderName, cfg.SelectedProvider) + } + if cfg.CurrentModel != testModel { + t.Fatalf("expected current model %q, got %q", testModel, cfg.CurrentModel) + } + + rewritten, err := os.ReadFile(loader.ConfigPath()) + if err != nil { + t.Fatalf("read rewritten config: %v", err) + } + text := string(rewritten) + if !strings.Contains(text, "current_model: "+testModel) { + t.Fatalf("expected rewritten config to persist current model, got:\n%s", text) + } +} diff --git a/internal/config/model.go b/internal/config/model.go index c0b57f0b..f8c2bd52 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -38,12 +38,11 @@ type Config struct { } type ProviderConfig struct { - Name string `yaml:"name"` - Driver string `yaml:"driver"` - BaseURL string `yaml:"base_url"` - Model string `yaml:"model"` - Models []string `yaml:"models,omitempty"` - APIKeyEnv string `yaml:"api_key_env"` + Name string `yaml:"name"` + Driver string `yaml:"driver"` + BaseURL string `yaml:"base_url"` + Model string `yaml:"model"` + APIKeyEnv string `yaml:"api_key_env"` } type ResolvedProviderConfig struct { @@ -92,25 +91,29 @@ func (c *Config) ApplyDefaultsFrom(defaults Config) { return } - if len(c.Providers) == 0 { + if len(defaults.Providers) > 0 { c.Providers = cloneProviders(defaults.Providers) - } else { - c.Providers = applyProviderDefaults(c.Providers, defaults.Providers) } - if strings.TrimSpace(c.SelectedProvider) == "" { - c.SelectedProvider = defaults.SelectedProvider + fallbackProvider := defaultSelectedProviderName(c.Providers, defaults.SelectedProvider) + selectedReset := false + switch current := strings.TrimSpace(c.SelectedProvider); { + case current == "": + c.SelectedProvider = fallbackProvider + selectedReset = true + case !containsProviderName(c.Providers, current): + c.SelectedProvider = fallbackProvider + selectedReset = true + default: + c.SelectedProvider = current } - if strings.TrimSpace(c.CurrentModel) == "" { + if strings.TrimSpace(c.CurrentModel) == "" || selectedReset { if selected, err := c.SelectedProviderConfig(); err == nil { c.CurrentModel = selected.Model } else if strings.TrimSpace(defaults.CurrentModel) != "" { c.CurrentModel = defaults.CurrentModel } } - if selected, err := c.SelectedProviderConfig(); err == nil && !ContainsModelID(selected.SupportedModels(), c.CurrentModel) { - c.CurrentModel = selected.Model - } if strings.TrimSpace(c.Workdir) == "" { c.Workdir = defaults.Workdir } @@ -137,16 +140,31 @@ func (c *Config) Validate() error { } seen := make(map[string]struct{}, len(c.Providers)) + seenEndpoints := make(map[string]string, len(c.Providers)) for i, provider := range c.Providers { if err := provider.Validate(); err != nil { return fmt.Errorf("config: provider[%d]: %w", i, err) } - key := strings.ToLower(strings.TrimSpace(provider.Name)) + key := NormalizeProviderName(provider.Name) if _, exists := seen[key]; exists { return fmt.Errorf("config: duplicate provider name %q", provider.Name) } seen[key] = struct{}{} + + identity, err := provider.Identity() + if err != nil { + return fmt.Errorf("config: provider[%d]: %w", i, err) + } + if existingName, exists := seenEndpoints[identity.Key()]; exists { + return fmt.Errorf( + "config: duplicate provider endpoint %q for providers %q and %q", + identity.BaseURL, + existingName, + provider.Name, + ) + } + seenEndpoints[identity.Key()] = provider.Name } if strings.TrimSpace(c.SelectedProvider) == "" { @@ -168,9 +186,6 @@ func (c *Config) Validate() error { if strings.TrimSpace(selected.Model) == "" { return fmt.Errorf("config: selected provider %q has empty model", selected.Name) } - if !ContainsModelID(selected.SupportedModels(), c.CurrentModel) { - return fmt.Errorf("config: current_model %q is not supported by provider %q", c.CurrentModel, selected.Name) - } if err := c.Tools.Validate(); err != nil { return fmt.Errorf("config: tools: %w", err) } @@ -190,9 +205,9 @@ func (c *Config) ProviderByName(name string) (ProviderConfig, error) { return ProviderConfig{}, errors.New("config: config is nil") } - target := strings.ToLower(strings.TrimSpace(name)) + target := NormalizeProviderName(name) for _, provider := range c.Providers { - if strings.ToLower(strings.TrimSpace(provider.Name)) == target { + if NormalizeProviderName(provider.Name) == target { return provider, nil } } @@ -216,28 +231,14 @@ func (p ProviderConfig) Validate() error { if strings.TrimSpace(p.APIKeyEnv) == "" { return fmt.Errorf("provider %q api_key_env is empty", p.Name) } - if models := normalizeModelIDs(p.Models); len(p.Models) > 0 { - if len(models) == 0 { - return fmt.Errorf("provider %q models is empty", p.Name) - } - if !ContainsModelID(models, p.Model) { - return fmt.Errorf("provider %q default model %q is not in models", p.Name, strings.TrimSpace(p.Model)) - } + if _, err := p.Identity(); err != nil { + return fmt.Errorf("provider %q: %w", p.Name, err) } return nil } -func (p ProviderConfig) SupportedModels() []string { - models := normalizeModelIDs(p.Models) - if len(models) > 0 { - return models - } - - model := strings.TrimSpace(p.Model) - if model == "" { - return nil - } - return []string{model} +func (p ProviderConfig) Identity() (ProviderIdentity, error) { + return NewProviderIdentity(p.Driver, p.BaseURL) } func (p ProviderConfig) ResolveAPIKey() (string, error) { @@ -266,57 +267,6 @@ func (p ProviderConfig) Resolve() (ResolvedProviderConfig, error) { }, nil } -func applyProviderDefaults(providers []ProviderConfig, defaults []ProviderConfig) []ProviderConfig { - out := make([]ProviderConfig, 0, len(providers)) - for _, provider := range providers { - out = append(out, mergeProviderDefaults(provider, defaults)) - } - return out -} - -func mergeProviderDefaults(provider ProviderConfig, defaults []ProviderConfig) ProviderConfig { - base, ok := matchDefaultProvider(provider, defaults) - if !ok { - return provider - } - - if strings.TrimSpace(provider.Name) == "" { - provider.Name = base.Name - } - if strings.TrimSpace(provider.Driver) == "" { - provider.Driver = base.Driver - } - if strings.TrimSpace(provider.BaseURL) == "" { - provider.BaseURL = base.BaseURL - } - if strings.TrimSpace(provider.Model) == "" { - provider.Model = base.Model - } - if len(provider.Models) == 0 { - provider.Models = cloneModelIDs(base.Models) - } - if strings.TrimSpace(provider.APIKeyEnv) == "" { - provider.APIKeyEnv = base.APIKeyEnv - } - - return provider -} - -func matchDefaultProvider(provider ProviderConfig, defaults []ProviderConfig) (ProviderConfig, bool) { - name := strings.ToLower(strings.TrimSpace(provider.Name)) - if name == "" { - return ProviderConfig{}, false - } - - for _, candidate := range defaults { - if strings.ToLower(candidate.Name) == name { - return candidate, true - } - } - - return ProviderConfig{}, false -} - func normalizeWorkdir(workdir string) string { workdir = strings.TrimSpace(workdir) if workdir == "" { @@ -455,50 +405,41 @@ func cloneProviders(providers []ProviderConfig) []ProviderConfig { cloned := make([]ProviderConfig, 0, len(providers)) for _, provider := range providers { - next := provider - next.Models = cloneModelIDs(provider.Models) - cloned = append(cloned, next) + cloned = append(cloned, provider) } return cloned } -func cloneModelIDs(models []string) []string { - normalized := normalizeModelIDs(models) - if len(normalized) == 0 { - return nil +func defaultSelectedProviderName(providers []ProviderConfig, fallback string) string { + if containsProviderName(providers, fallback) { + return strings.TrimSpace(fallback) + } + if len(providers) == 0 { + return "" } - return append([]string(nil), normalized...) + return strings.TrimSpace(providers[0].Name) } -func normalizeModelIDs(models []string) []string { - if len(models) == 0 { - return nil +func containsProviderName(providers []ProviderConfig, name string) bool { + target := NormalizeProviderName(name) + if target == "" { + return false } - - normalized := make([]string, 0, len(models)) - seen := make(map[string]struct{}, len(models)) - for _, model := range models { - id := strings.TrimSpace(model) - if id == "" { - continue - } - key := strings.ToLower(id) - if _, exists := seen[key]; exists { - continue + for _, provider := range providers { + if NormalizeProviderName(provider.Name) == target { + return true } - seen[key] = struct{}{} - normalized = append(normalized, id) } - return normalized + return false } func ContainsModelID(models []string, model string) bool { - target := strings.ToLower(strings.TrimSpace(model)) + target := NormalizeKey(model) if target == "" { return false } - for _, candidate := range normalizeModelIDs(models) { - if strings.EqualFold(candidate, target) { + for _, candidate := range models { + if NormalizeKey(candidate) == target { return true } } diff --git a/internal/config/provider_identity.go b/internal/config/provider_identity.go new file mode 100644 index 00000000..93f01207 --- /dev/null +++ b/internal/config/provider_identity.go @@ -0,0 +1,76 @@ +package config + +import ( + "fmt" + "net/url" + "path" + "strings" +) + +// ProviderIdentity identifies a concrete provider endpoint for discovery and caching. +type ProviderIdentity struct { + Driver string `json:"driver"` + BaseURL string `json:"base_url"` +} + +func (i ProviderIdentity) Key() string { + return i.Driver + "|" + i.BaseURL +} + +func NormalizeKey(s string) string { + return strings.ToLower(strings.TrimSpace(s)) +} + +func NormalizeProviderName(name string) string { + return NormalizeKey(name) +} + +func NormalizeProviderDriver(driver string) string { + return NormalizeKey(driver) +} + +func NormalizeProviderBaseURL(raw string) (string, error) { + value := strings.TrimSpace(raw) + if value == "" { + return "", fmt.Errorf("provider base_url is empty") + } + + parsed, err := url.Parse(value) + if err != nil { + return "", fmt.Errorf("provider base_url %q is invalid: %w", raw, err) + } + if parsed.Scheme == "" || parsed.Host == "" { + return "", fmt.Errorf("provider base_url %q must include scheme and host", raw) + } + + parsed.Scheme = strings.ToLower(parsed.Scheme) + parsed.Host = strings.ToLower(parsed.Host) + parsed.RawQuery = "" + parsed.Fragment = "" + parsed.User = nil + + if cleaned := path.Clean(strings.TrimSpace(parsed.Path)); cleaned == "." { + parsed.Path = "" + } else { + parsed.Path = strings.TrimRight(cleaned, "/") + } + + return parsed.String(), nil +} + +func NewProviderIdentity(driver string, baseURL string) (ProviderIdentity, error) { + normalizedDriver := NormalizeProviderDriver(driver) + if normalizedDriver == "" { + return ProviderIdentity{}, fmt.Errorf("provider driver is empty") + } + + normalizedBaseURL, err := NormalizeProviderBaseURL(baseURL) + if err != nil { + return ProviderIdentity{}, err + } + + return ProviderIdentity{ + Driver: normalizedDriver, + BaseURL: normalizedBaseURL, + }, nil +} diff --git a/internal/provider/discovery/discovery.go b/internal/provider/discovery/discovery.go new file mode 100644 index 00000000..4311117c --- /dev/null +++ b/internal/provider/discovery/discovery.go @@ -0,0 +1,50 @@ +package discovery + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +type openAIModelsResponse struct { + Data []map[string]any `json:"data"` +} + +// FetchOpenAICompatibleModels fetches raw model objects from an OpenAI-compatible /models endpoint. +func FetchOpenAICompatibleModels(ctx context.Context, client *http.Client, baseURL string, apiKey string) ([]map[string]any, error) { + endpoint := strings.TrimRight(strings.TrimSpace(baseURL), "/") + "/models" + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return nil, fmt.Errorf("provider discovery: build request: %w", err) + } + req.Header.Set("Accept", "application/json") + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(apiKey)) + } + + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("provider discovery: send request: %w", err) + } + defer func(body io.ReadCloser) { + _ = body.Close() + }(resp.Body) + + if resp.StatusCode >= http.StatusBadRequest { + data, _ := io.ReadAll(resp.Body) + body := strings.TrimSpace(string(data)) + if body == "" { + body = resp.Status + } + return nil, fmt.Errorf("provider discovery: %s", body) + } + + var payload openAIModelsResponse + if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil { + return nil, fmt.Errorf("provider discovery: decode response: %w", err) + } + return payload.Data, nil +} diff --git a/internal/provider/discovery/discovery_test.go b/internal/provider/discovery/discovery_test.go new file mode 100644 index 00000000..50380a4c --- /dev/null +++ b/internal/provider/discovery/discovery_test.go @@ -0,0 +1,51 @@ +package discovery + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestFetchOpenAICompatibleModels(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/models" { + t.Fatalf("unexpected path %q", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Fatalf("unexpected auth header %q", got) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"id":"gpt-test","context_window":128000,"extra":"kept"}]}`)) + })) + defer server.Close() + + models, err := FetchOpenAICompatibleModels(context.Background(), server.Client(), server.URL, "test-key") + if err != nil { + t.Fatalf("FetchOpenAICompatibleModels() error = %v", err) + } + if len(models) != 1 || models[0]["id"] != "gpt-test" { + t.Fatalf("unexpected models payload: %+v", models) + } + if models[0]["extra"] != "kept" { + t.Fatalf("expected unknown fields to remain in raw payload, got %+v", models[0]) + } +} + +func TestFetchOpenAICompatibleModelsHTTPError(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadGateway) + _, _ = w.Write([]byte("gateway failed")) + })) + defer server.Close() + + _, err := FetchOpenAICompatibleModels(context.Background(), server.Client(), server.URL, "test-key") + if err == nil || !strings.Contains(err.Error(), "gateway failed") { + t.Fatalf("expected gateway failure error, got %v", err) + } +} diff --git a/internal/provider/errors.go b/internal/provider/errors.go index 083fb39f..2280487a 100644 --- a/internal/provider/errors.go +++ b/internal/provider/errors.go @@ -1,10 +1,21 @@ package provider import ( + "errors" "fmt" "net/http" ) +// 通用领域错误。 +var ( + ErrProviderNotFound = errors.New("provider not found") + ErrModelNotFound = errors.New("model not found") + ErrDriverNotFound = errors.New("provider driver not found") + ErrDriverAlreadyRegistered = errors.New("provider: driver already registered") + ErrServiceManagerNil = errors.New("provider: config manager is nil") + ErrServiceRegistryNil = errors.New("provider: registry is nil") +) + type ProviderErrorCode string const ( diff --git a/internal/provider/metadata.go b/internal/provider/metadata.go new file mode 100644 index 00000000..319a87f9 --- /dev/null +++ b/internal/provider/metadata.go @@ -0,0 +1,210 @@ +package provider + +import ( + "strings" + + "neo-code/internal/config" +) + +// DescriptorFromRawModel normalizes a raw provider model object into a ModelDescriptor. +func DescriptorFromRawModel(raw map[string]any) (ModelDescriptor, bool) { + id := firstNonEmptyString( + stringValue(raw["id"]), + stringValue(raw["model"]), + stringValue(raw["name"]), + ) + if id == "" { + return ModelDescriptor{}, false + } + + descriptor := ModelDescriptor{ + ID: id, + Name: firstNonEmptyString(stringValue(raw["name"]), stringValue(raw["display_name"]), id), + Description: stringValue(raw["description"]), + ContextWindow: firstPositiveInt(raw["context_window"], raw["contextLength"], raw["input_token_limit"], raw["max_context_tokens"]), + MaxOutputTokens: firstPositiveInt(raw["max_output_tokens"], raw["output_token_limit"], raw["max_tokens"]), + Capabilities: boolMapValue(raw["capabilities"]), + } + return normalizeModelDescriptor(descriptor), true +} + +func modelDescriptorsFromIDs(modelIDs []string) []ModelDescriptor { + if len(modelIDs) == 0 { + return nil + } + + descriptors := make([]ModelDescriptor, 0, len(modelIDs)) + for _, modelID := range modelIDs { + id := strings.TrimSpace(modelID) + if id == "" { + continue + } + descriptors = append(descriptors, ModelDescriptor{ + ID: id, + Name: id, + }) + } + if len(descriptors) == 0 { + return nil + } + return descriptors +} + +func MergeModelDescriptors(sources ...[]ModelDescriptor) []ModelDescriptor { + if len(sources) == 0 { + return nil + } + + merged := make([]ModelDescriptor, 0) + indexByID := make(map[string]int) + + for _, source := range sources { + for _, candidate := range source { + normalized := normalizeModelDescriptor(candidate) + key := normalizeModelID(normalized.ID) + if key == "" { + continue + } + + if index, exists := indexByID[key]; exists { + merged[index] = mergeModelDescriptor(merged[index], normalized) + continue + } + + indexByID[key] = len(merged) + merged = append(merged, normalized) + } + } + + if len(merged) == 0 { + return nil + } + return merged +} + +func normalizeModelDescriptor(descriptor ModelDescriptor) ModelDescriptor { + descriptor.ID = strings.TrimSpace(descriptor.ID) + descriptor.Name = strings.TrimSpace(descriptor.Name) + descriptor.Description = strings.TrimSpace(descriptor.Description) + if descriptor.Name == "" { + descriptor.Name = descriptor.ID + } + descriptor.Capabilities = cloneStringBoolMap(descriptor.Capabilities) + return descriptor +} + +func mergeModelDescriptor(primary ModelDescriptor, secondary ModelDescriptor) ModelDescriptor { + if strings.TrimSpace(primary.Name) == "" { + primary.Name = secondary.Name + } + if strings.TrimSpace(primary.Description) == "" { + primary.Description = secondary.Description + } + if primary.ContextWindow <= 0 { + primary.ContextWindow = secondary.ContextWindow + } + if primary.MaxOutputTokens <= 0 { + primary.MaxOutputTokens = secondary.MaxOutputTokens + } + primary.Capabilities = mergeStringBoolMaps(primary.Capabilities, secondary.Capabilities) + return normalizeModelDescriptor(primary) +} + +func mergeStringBoolMaps(primary map[string]bool, secondary map[string]bool) map[string]bool { + if len(primary) == 0 && len(secondary) == 0 { + return nil + } + + merged := cloneStringBoolMap(primary) + if merged == nil { + merged = make(map[string]bool, len(secondary)) + } + for key, value := range secondary { + if _, exists := merged[key]; exists { + continue + } + merged[key] = value + } + return merged +} + +func boolMapValue(value any) map[string]bool { + raw, ok := value.(map[string]any) + if !ok || len(raw) == 0 { + return nil + } + + normalized := make(map[string]bool, len(raw)) + for key, value := range raw { + boolValue, ok := value.(bool) + if !ok { + continue + } + normalized[strings.TrimSpace(key)] = boolValue + } + if len(normalized) == 0 { + return nil + } + return normalized +} + +func stringValue(value any) string { + switch typed := value.(type) { + case string: + return strings.TrimSpace(typed) + default: + return "" + } +} + +func firstNonEmptyString(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} + +func firstPositiveInt(values ...any) int { + for _, value := range values { + switch typed := value.(type) { + case int: + if typed > 0 { + return typed + } + case int32: + if typed > 0 { + return int(typed) + } + case int64: + if typed > 0 { + return int(typed) + } + case float32: + if typed > 0 { + return int(typed) + } + case float64: + if typed > 0 { + return int(typed) + } + } + } + return 0 +} + +func cloneStringBoolMap(source map[string]bool) map[string]bool { + if len(source) == 0 { + return nil + } + cloned := make(map[string]bool, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + +func normalizeModelID(id string) string { + return config.NormalizeKey(id) +} diff --git a/internal/provider/metadata_test.go b/internal/provider/metadata_test.go new file mode 100644 index 00000000..d205e196 --- /dev/null +++ b/internal/provider/metadata_test.go @@ -0,0 +1,72 @@ +package provider + +import "testing" + +func TestDescriptorFromRawModelNormalizesUsefulFields(t *testing.T) { + t.Parallel() + + raw := map[string]any{ + "id": "gpt-test", + "name": "GPT Test", + "description": "custom model", + "context_window": float64(128000), + "max_output_tokens": float64(8192), + "capabilities": map[string]any{ + "tool_call": true, + "vision": false, + "notes": "ignored", + }, + "experimental": map[string]any{ + "tier": "beta", + }, + } + + descriptor, ok := DescriptorFromRawModel(raw) + if !ok { + t.Fatal("expected descriptor to be normalized") + } + if descriptor.ID != "gpt-test" || descriptor.Name != "GPT Test" { + t.Fatalf("unexpected descriptor identity: %+v", descriptor) + } + if descriptor.ContextWindow != 128000 || descriptor.MaxOutputTokens != 8192 { + t.Fatalf("expected token metadata to be normalized, got %+v", descriptor) + } + if !descriptor.Capabilities["tool_call"] { + t.Fatalf("expected tool_call capability, got %+v", descriptor.Capabilities) + } + if _, ok := descriptor.Capabilities["notes"]; ok { + t.Fatalf("expected non-bool capability values to be ignored, got %+v", descriptor.Capabilities) + } +} + +func TestMergeModelDescriptorsPrefersEarlierSourceAndBackfillsUsefulFields(t *testing.T) { + t.Parallel() + + manual := []ModelDescriptor{{ + ID: "gpt-test", + Name: "Manual Name", + }} + discovered := []ModelDescriptor{{ + ID: "gpt-test", + Name: "Discovered Name", + ContextWindow: 64000, + MaxOutputTokens: 4096, + Capabilities: map[string]bool{ + "tool_call": true, + }, + }} + + merged := MergeModelDescriptors(manual, discovered) + if len(merged) != 1 { + t.Fatalf("expected 1 merged model, got %d", len(merged)) + } + if merged[0].Name != "Manual Name" { + t.Fatalf("expected earlier source to win for name, got %+v", merged[0]) + } + if merged[0].ContextWindow != 64000 { + t.Fatalf("expected metadata to be backfilled, got %+v", merged[0]) + } + if !merged[0].Capabilities["tool_call"] { + t.Fatalf("expected capabilities to be backfilled, got %+v", merged[0].Capabilities) + } +} diff --git a/internal/provider/model_catalog_store.go b/internal/provider/model_catalog_store.go new file mode 100644 index 00000000..bca392b9 --- /dev/null +++ b/internal/provider/model_catalog_store.go @@ -0,0 +1,130 @@ +package provider + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "neo-code/internal/config" +) + +const modelCatalogSchemaVersion = 1 + +var ErrModelCatalogNotFound = errors.New("provider: model catalog not found") + +// ModelCatalog stores discovered provider models for a concrete endpoint. +type ModelCatalog struct { + SchemaVersion int `json:"schema_version"` + Identity config.ProviderIdentity `json:"identity"` + FetchedAt time.Time `json:"fetched_at"` + ExpiresAt time.Time `json:"expires_at"` + Models []ModelDescriptor `json:"models"` +} + +func (c ModelCatalog) Expired(now time.Time) bool { + return !c.ExpiresAt.IsZero() && !now.Before(c.ExpiresAt) +} + +// ModelCatalogStore persists model catalogs keyed by normalized provider identity. +type ModelCatalogStore interface { + Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) + Save(ctx context.Context, catalog ModelCatalog) error +} + +type JSONModelCatalogStore struct { + dir string +} + +func NewJSONModelCatalogStore(baseDir string) *JSONModelCatalogStore { + return &JSONModelCatalogStore{ + dir: filepath.Join(strings.TrimSpace(baseDir), "cache", "models"), + } +} + +func (s *JSONModelCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return ModelCatalog{}, err + } + + normalized, err := config.NewProviderIdentity(identity.Driver, identity.BaseURL) + if err != nil { + return ModelCatalog{}, fmt.Errorf("provider: normalize model catalog key: %w", err) + } + + path, err := s.catalogPathFromNormalized(normalized) + if err != nil { + return ModelCatalog{}, err + } + + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return ModelCatalog{}, ErrModelCatalogNotFound + } + return ModelCatalog{}, fmt.Errorf("provider: read model catalog: %w", err) + } + + var catalog ModelCatalog + if err := json.Unmarshal(data, &catalog); err != nil { + return ModelCatalog{}, fmt.Errorf("provider: decode model catalog: %w", err) + } + + catalog.Identity = normalized + return normalizeModelCatalog(catalog), nil +} + +func (s *JSONModelCatalogStore) Save(ctx context.Context, catalog ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + normalized := normalizeModelCatalog(catalog) + + normalizedIdentity, err := config.NewProviderIdentity(normalized.Identity.Driver, normalized.Identity.BaseURL) + if err != nil { + return fmt.Errorf("provider: normalize model catalog key: %w", err) + } + normalized.Identity = normalizedIdentity + + path, err := s.catalogPathFromNormalized(normalized.Identity) + if err != nil { + return err + } + + if err := os.MkdirAll(s.dir, 0o755); err != nil { + return fmt.Errorf("provider: create model catalog dir: %w", err) + } + + data, err := json.MarshalIndent(normalized, "", " ") + if err != nil { + return fmt.Errorf("provider: encode model catalog: %w", err) + } + if len(data) == 0 || data[len(data)-1] != '\n' { + data = append(data, '\n') + } + + if err := os.WriteFile(path, data, 0o644); err != nil { + return fmt.Errorf("provider: write model catalog: %w", err) + } + return nil +} + +func normalizeModelCatalog(catalog ModelCatalog) ModelCatalog { + if catalog.SchemaVersion == 0 { + catalog.SchemaVersion = modelCatalogSchemaVersion + } + catalog.Models = MergeModelDescriptors(catalog.Models) + return catalog +} + +func (s *JSONModelCatalogStore) catalogPathFromNormalized(identity config.ProviderIdentity) (string, error) { + sum := sha256.Sum256([]byte(identity.Key())) + return filepath.Join(s.dir, hex.EncodeToString(sum[:])+".json"), nil +} diff --git a/internal/provider/model_catalog_store_test.go b/internal/provider/model_catalog_store_test.go new file mode 100644 index 00000000..70b17043 --- /dev/null +++ b/internal/provider/model_catalog_store_test.go @@ -0,0 +1,73 @@ +package provider + +import ( + "context" + "errors" + "testing" + "time" + + "neo-code/internal/config" +) + +func TestJSONModelCatalogStoreRoundTrip(t *testing.T) { + t.Parallel() + + store := NewJSONModelCatalogStore(t.TempDir()) + identity, err := config.NewProviderIdentity("OPENAI", "https://API.OPENAI.COM/v1/") + if err != nil { + t.Fatalf("NewProviderIdentity() error = %v", err) + } + + expected := ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + FetchedAt: time.Date(2026, 4, 2, 10, 0, 0, 0, time.UTC), + ExpiresAt: time.Date(2026, 4, 3, 10, 0, 0, 0, time.UTC), + Models: []ModelDescriptor{{ + ID: "gpt-test", + Name: "GPT Test", + Description: "normalized model", + ContextWindow: 128000, + MaxOutputTokens: 8192, + Capabilities: map[string]bool{ + "tool_call": true, + }, + }}, + } + + if err := store.Save(context.Background(), expected); err != nil { + t.Fatalf("Save() error = %v", err) + } + + got, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if got.Identity.BaseURL != "https://api.openai.com/v1" { + t.Fatalf("expected normalized base url, got %q", got.Identity.BaseURL) + } + if len(got.Models) != 1 { + t.Fatalf("expected 1 model, got %+v", got.Models) + } + if got.Models[0].ContextWindow != 128000 || got.Models[0].MaxOutputTokens != 8192 { + t.Fatalf("expected normalized token fields to survive round-trip, got %+v", got.Models[0]) + } + if !got.Models[0].Capabilities["tool_call"] { + t.Fatalf("expected capabilities to survive round-trip, got %+v", got.Models[0].Capabilities) + } +} + +func TestJSONModelCatalogStoreMissingCatalog(t *testing.T) { + t.Parallel() + + store := NewJSONModelCatalogStore(t.TempDir()) + identity, err := config.NewProviderIdentity("openai", "https://api.openai.com/v1") + if err != nil { + t.Fatalf("NewProviderIdentity() error = %v", err) + } + + _, err = store.Load(context.Background(), identity) + if !errors.Is(err, ErrModelCatalogNotFound) { + t.Fatalf("expected ErrModelCatalogNotFound, got %v", err) + } +} diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index 700b50a3..c0f891c3 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -16,6 +16,7 @@ import ( "neo-code/internal/config" domain "neo-code/internal/provider" + modeldiscovery "neo-code/internal/provider/discovery" "neo-code/internal/provider/transport" ) @@ -39,14 +40,24 @@ func WithTransport(rt http.RoundTripper) BuildOption { const DriverName = "openai" +// defaultRetryTransport 返回内置的带重试的 HTTP Transport。 +func defaultRetryTransport() http.RoundTripper { + return transport.NewRetryTransport(http.DefaultTransport, transport.DefaultRetryConfig()) +} + // Driver 返回 OpenAI 协议驱动的定义。 func Driver() domain.DriverDefinition { return domain.DriverDefinition{ Name: DriverName, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (domain.Provider, error) { - baseTransport := http.DefaultTransport - retryTransport := transport.NewRetryTransport(baseTransport, transport.DefaultRetryConfig()) - return New(cfg, WithTransport(retryTransport)) + return New(cfg, WithTransport(defaultRetryTransport())) + }, + Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]domain.ModelDescriptor, error) { + provider, err := New(cfg, WithTransport(defaultRetryTransport())) + if err != nil { + return nil, err + } + return provider.DiscoverModels(ctx) }, } } @@ -75,6 +86,23 @@ func New(cfg config.ResolvedProviderConfig, opts ...BuildOption) (*Provider, err }, nil } +func (p *Provider) DiscoverModels(ctx context.Context) ([]domain.ModelDescriptor, error) { + rawModels, err := modeldiscovery.FetchOpenAICompatibleModels(ctx, p.client, p.cfg.BaseURL, p.cfg.APIKey) + if err != nil { + return nil, err + } + + descriptors := make([]domain.ModelDescriptor, 0, len(rawModels)) + for _, raw := range rawModels { + descriptor, ok := domain.DescriptorFromRawModel(raw) + if !ok { + continue + } + descriptors = append(descriptors, descriptor) + } + return domain.MergeModelDescriptors(descriptors), nil +} + func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan<- domain.StreamEvent) (domain.ChatResponse, error) { payload, err := p.buildRequest(req) if err != nil { diff --git a/internal/provider/registry.go b/internal/provider/registry.go index c017b04c..f5e17324 100644 --- a/internal/provider/registry.go +++ b/internal/provider/registry.go @@ -10,18 +10,18 @@ import ( ) type Builder func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) +type DiscoveryFunc func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) type DriverDefinition struct { - Name string - Build Builder + Name string + Build Builder + Discover DiscoveryFunc } type Registry struct { drivers map[string]DriverDefinition } -var errDriverAlreadyRegistered = errors.New("provider: driver already registered") - func NewRegistry() *Registry { return &Registry{drivers: map[string]DriverDefinition{}} } @@ -34,7 +34,7 @@ func (r *Registry) Register(driver DriverDefinition) error { r.ensureDrivers() driver.Name = strings.TrimSpace(driver.Name) - driverType := normalizeKey(driver.Name) + driverType := config.NormalizeKey(driver.Name) if driverType == "" { return errors.New("provider: driver name is empty") } @@ -42,7 +42,7 @@ func (r *Registry) Register(driver DriverDefinition) error { return fmt.Errorf("provider: driver %q build func is nil", driver.Name) } if _, exists := r.drivers[driverType]; exists { - return fmt.Errorf("%w: %s", errDriverAlreadyRegistered, driver.Name) + return fmt.Errorf("%w: %s", ErrDriverAlreadyRegistered, driver.Name) } r.drivers[driverType] = driver return nil @@ -56,6 +56,17 @@ func (r *Registry) Build(ctx context.Context, cfg config.ResolvedProviderConfig) return driver.Build(ctx, cfg) } +func (r *Registry) DiscoverModels(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + driver, err := r.driver(cfg.Driver) + if err != nil { + return nil, err + } + if driver.Discover == nil { + return nil, nil + } + return driver.Discover(ctx, cfg) +} + func (r *Registry) Supports(driverType string) bool { _, err := r.driver(driverType) return err == nil @@ -65,17 +76,13 @@ func (r *Registry) driver(driverType string) (DriverDefinition, error) { if r == nil { return DriverDefinition{}, ErrDriverNotFound } - driver, ok := r.drivers[normalizeKey(driverType)] + driver, ok := r.drivers[config.NormalizeKey(driverType)] if !ok { return DriverDefinition{}, fmt.Errorf("%w: %s", ErrDriverNotFound, strings.TrimSpace(driverType)) } return driver, nil } -func normalizeKey(value string) string { - return strings.ToLower(strings.TrimSpace(value)) -} - func (r *Registry) ensureDrivers() { if r.drivers == nil { r.drivers = map[string]DriverDefinition{} diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index 16c022ca..7f3cb9df 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -3,6 +3,7 @@ package provider_test import ( "context" "errors" + "sync" "testing" "neo-code/internal/config" @@ -98,72 +99,47 @@ func TestRegistryRejectsDuplicateDriverRegistration(t *testing.T) { } } -func TestServiceListProvidersUsesConfiguredMetadata(t *testing.T) { +func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { t.Parallel() manager := newTestManager(t) - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "unsupported", - Driver: "custom", - BaseURL: "https://example.com", - Model: "custom-model", - Models: []string{"custom-model"}, - APIKeyEnv: "CUSTOM_API_KEY", - }) - return nil - }); err != nil { - t.Fatalf("append provider: %v", err) - } - service := provider.NewService(manager, newTestRegistry(t)) + service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) items, err := service.ListProviders(context.Background()) if err != nil { t.Fatalf("ListProviders() error = %v", err) } - expectedModels := map[string]int{ - config.OpenAIName: len(config.OpenAIProvider().Models), - config.GeminiName: len(config.GeminiProvider().Models), - config.OpenLLName: len(config.OpenLLProvider().Models), - config.QiniuName: len(config.QiniuProvider().Models), + expected := map[string]int{ + config.OpenAIName: 1, + config.GeminiName: 1, + config.OpenLLName: 1, + config.QiniuName: 1, } - if len(items) != len(expectedModels) { - t.Fatalf("expected only supported providers, got %d", len(items)) + if len(items) != len(expected) { + t.Fatalf("expected only builtin providers, got %d", len(items)) } for _, item := range items { - wantModels, ok := expectedModels[item.ID] + wantModels, ok := expected[item.ID] if !ok { - t.Fatalf("unexpected supported provider %q", item.ID) + t.Fatalf("unexpected builtin provider %q", item.ID) } if item.Description != "" { t.Fatalf("expected provider description to stay empty for hidden metadata, got %q", item.Description) } if len(item.Models) != wantModels { - t.Fatalf("expected provider models to come from config, got %+v", item.Models) + t.Fatalf("expected provider models to fall back to the default model, got %+v", item.Models) } - delete(expectedModels, item.ID) + delete(expected, item.ID) } - if len(expectedModels) != 0 { - t.Fatalf("missing supported providers from catalog: %+v", expectedModels) + if len(expected) != 0 { + t.Fatalf("missing builtin providers from catalog: %+v", expected) } } func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { - defaults := builtin.DefaultConfig() - defaults.Providers = append(defaults.Providers, config.ProviderConfig{ - Name: "custom-main", - Driver: "custom", - BaseURL: "https://example.com", - Model: "custom-model", - Models: []string{"custom-model", "custom-alt"}, - APIKeyEnv: "CUSTOM_API_KEY", - }) - manager := newTestManagerWithDefaults(t, defaults) + manager := newTestManager(t) registry := newTestRegistry(t) - if err := registry.Register(stubDriver("custom")); err != nil { - t.Fatalf("register stub driver: %v", err) - } if err := manager.Update(context.Background(), func(cfg *config.Config) error { cfg.CurrentModel = "gpt-5.4" @@ -172,13 +148,28 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { t.Fatalf("seed current model: %v", err) } - service := provider.NewService(manager, registry) + store := newCatalogStoreStub() + identity, err := config.QiniuProvider().Identity() + if err != nil { + t.Fatalf("qiniu provider identity: %v", err) + } + if err := store.Save(context.Background(), provider.ModelCatalog{ + SchemaVersion: 1, + Identity: identity, + Models: []provider.ModelDescriptor{ + {ID: "qiniu-alt", Name: "qiniu-alt"}, + }, + }); err != nil { + t.Fatalf("seed catalog: %v", err) + } + + service := provider.NewService(manager, registry, store) - selection, err := service.SelectProvider(context.Background(), "custom-main") + selection, err := service.SelectProvider(context.Background(), config.QiniuName) if err != nil { t.Fatalf("SelectProvider() error = %v", err) } - if selection.ProviderID != "custom-main" || selection.ModelID != "custom-model" { + if selection.ProviderID != config.QiniuName || selection.ModelID != config.QiniuDefaultModel { t.Fatalf("unexpected selection after switch: %+v", selection) } @@ -186,12 +177,12 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { t.Fatalf("expected ErrModelNotFound, got %v", err) } - selection, err = service.SetCurrentModel(context.Background(), "custom-alt") + selection, err = service.SetCurrentModel(context.Background(), "qiniu-alt") if err != nil { t.Fatalf("SetCurrentModel() error = %v", err) } - if selection.ModelID != "custom-alt" { - t.Fatalf("expected selected model %q, got %+v", "custom-alt", selection) + if selection.ModelID != "qiniu-alt" { + t.Fatalf("expected selected model %q, got %+v", "qiniu-alt", selection) } cfg := manager.Get() @@ -199,10 +190,10 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { if err != nil { t.Fatalf("SelectedProviderConfig() error = %v", err) } - if selected.Name != "custom-main" { - t.Fatalf("expected selected provider %q, got %+v", "custom-main", selected) + if selected.Name != config.QiniuName { + t.Fatalf("expected selected provider %q, got %+v", config.QiniuName, selected) } - if selected.Model != "custom-model" || cfg.CurrentModel != "custom-alt" { + if selected.Model != config.QiniuDefaultModel || cfg.CurrentModel != "qiniu-alt" { t.Fatalf("expected provider default and current model to diverge safely, got provider=%q current=%q", selected.Model, cfg.CurrentModel) } @@ -214,70 +205,58 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { if err != nil { t.Fatalf("SelectedProviderConfig() after reload error = %v", err) } - if selected.Model != "custom-model" || reloaded.CurrentModel != "custom-alt" { + if selected.Model != config.QiniuDefaultModel || reloaded.CurrentModel != "qiniu-alt" { t.Fatalf("expected current model persistence without overriding provider default, got provider=%q current=%q", selected.Model, reloaded.CurrentModel) } } -func TestServiceModelOperationsUseProviderConfigEvenWithoutDriver(t *testing.T) { +func TestServiceModelOperationsUseBuiltinConfigEvenWithoutDriver(t *testing.T) { t.Parallel() - manager := newTestManager(t) - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "broken-provider", - Driver: "missing-driver", - BaseURL: "https://example.com", - Model: "broken-model", - Models: []string{"broken-model", "broken-alt"}, - APIKeyEnv: "BROKEN_API_KEY", - }) - cfg.SelectedProvider = "broken-provider" - cfg.CurrentModel = "broken-model" - return nil - }); err != nil { - t.Fatalf("append broken provider: %v", err) - } + defaults := builtin.DefaultConfig() + defaults.Providers = []config.ProviderConfig{{ + Name: config.OpenAIName, + Driver: "missing-driver", + BaseURL: config.OpenAIDefaultBaseURL, + Model: "broken-model", + APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, + }} + defaults.SelectedProvider = config.OpenAIName + defaults.CurrentModel = "broken-model" + manager := newTestManagerWithDefaults(t, defaults) - service := provider.NewService(manager, newTestRegistry(t)) + service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) models, err := service.ListModels(context.Background()) if err != nil { t.Fatalf("ListModels() error = %v", err) } - if len(models) != 2 || models[1].ID != "broken-alt" { - t.Fatalf("expected models from provider config, got %+v", models) + if len(models) != 1 || models[0].ID != "broken-model" { + t.Fatalf("expected default model fallback, got %+v", models) } - selection, err := service.SetCurrentModel(context.Background(), "broken-alt") - if err != nil { - t.Fatalf("SetCurrentModel() error = %v", err) - } - if selection.ModelID != "broken-alt" { - t.Fatalf("expected current model to update from provider config, got %+v", selection) + if _, err := service.SetCurrentModel(context.Background(), "broken-alt"); !errors.Is(err, provider.ErrModelNotFound) { + t.Fatalf("expected SetCurrentModel() to reject undiscovered model, got %v", err) } } func TestServiceSelectProviderRejectsUnsupportedDriver(t *testing.T) { t.Parallel() - manager := newTestManager(t) - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "broken-provider", - Driver: "missing-driver", - BaseURL: "https://example.com", - Model: "broken-model", - Models: []string{"broken-model"}, - APIKeyEnv: "BROKEN_API_KEY", - }) - return nil - }); err != nil { - t.Fatalf("append broken provider: %v", err) - } + defaults := builtin.DefaultConfig() + defaults.Providers = []config.ProviderConfig{{ + Name: config.OpenAIName, + Driver: "missing-driver", + BaseURL: config.OpenAIDefaultBaseURL, + Model: config.OpenAIDefaultModel, + APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, + }} + defaults.SelectedProvider = config.OpenAIName + defaults.CurrentModel = config.OpenAIDefaultModel + manager := newTestManagerWithDefaults(t, defaults) - service := provider.NewService(manager, newTestRegistry(t)) - if _, err := service.SelectProvider(context.Background(), "broken-provider"); !errors.Is(err, provider.ErrDriverNotFound) { + service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) + if _, err := service.SelectProvider(context.Background(), config.OpenAIName); !errors.Is(err, provider.ErrDriverNotFound) { t.Fatalf("expected SelectProvider() to preserve ErrDriverNotFound, got %v", err) } } @@ -323,10 +302,47 @@ func TestServiceBuildReturnsErrorOnNilRegistry(t *testing.T) { t.Parallel() manager := newTestManager(t) - service := provider.NewService(manager, nil) + service := provider.NewService(manager, nil, newCatalogStoreStub()) _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) if err == nil { t.Fatalf("expected error for nil registry") } } + +type catalogStoreStub struct { + mu sync.Mutex + catalogs map[string]provider.ModelCatalog +} + +func newCatalogStoreStub() *catalogStoreStub { + return &catalogStoreStub{ + catalogs: map[string]provider.ModelCatalog{}, + } +} + +func (s *catalogStoreStub) Load(ctx context.Context, identity config.ProviderIdentity) (provider.ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return provider.ModelCatalog{}, err + } + + s.mu.Lock() + defer s.mu.Unlock() + + catalog, ok := s.catalogs[identity.Key()] + if !ok { + return provider.ModelCatalog{}, provider.ErrModelCatalogNotFound + } + return catalog, nil +} + +func (s *catalogStoreStub) Save(ctx context.Context, catalog provider.ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + s.mu.Lock() + defer s.mu.Unlock() + s.catalogs[catalog.Identity.Key()] = catalog + return nil +} diff --git a/internal/provider/service.go b/internal/provider/service.go index acbac2db..0dbded09 100644 --- a/internal/provider/service.go +++ b/internal/provider/service.go @@ -2,26 +2,43 @@ package provider import ( "context" - "errors" "strings" + "sync" + "time" "neo-code/internal/config" ) +const ( + defaultModelCatalogTTL = 24 * time.Hour + defaultBackgroundRefreshTime = 30 * time.Second +) + type Service struct { - manager *config.Manager - registry *Registry + manager *config.Manager + registry *Registry + catalogs ModelCatalogStore + catalogTTL time.Duration + backgroundTimeout time.Duration + now func() time.Time + + refreshMu sync.Mutex + inFlightByID map[string]struct{} } -var ( - errServiceManagerNil = errors.New("provider: config manager is nil") - errServiceRegistryNil = errors.New("provider: registry is nil") -) +func NewService(manager *config.Manager, registry *Registry, catalogs ModelCatalogStore) *Service { + if catalogs == nil && manager != nil { + catalogs = NewJSONModelCatalogStore(manager.BaseDir()) + } -func NewService(manager *config.Manager, registry *Registry) *Service { return &Service{ - manager: manager, - registry: registry, + manager: manager, + registry: registry, + catalogs: catalogs, + catalogTTL: defaultModelCatalogTTL, + backgroundTimeout: defaultBackgroundRefreshTime, + now: time.Now, + inFlightByID: map[string]struct{}{}, } } @@ -42,7 +59,9 @@ func (s *Service) ListProviders(ctx context.Context) ([]ProviderCatalogItem, err if !s.registry.Supports(providerCfg.Driver) { continue } - items = append(items, catalogItemFromConfig(providerCfg)) + + models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) + items = append(items, catalogItemFromConfig(providerCfg, models)) } return items, nil @@ -53,19 +72,30 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov return ProviderSelection{}, err } - var selection ProviderSelection + cfgSnapshot := s.manager.Get() + providerCfg, err := cfgSnapshot.ProviderByName(providerName) + if err != nil { + return ProviderSelection{}, ErrProviderNotFound + } + if !s.registry.Supports(providerCfg.Driver) { + return ProviderSelection{}, ErrDriverNotFound + } + + models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) - err := s.manager.Update(ctx, func(cfg *config.Config) error { - providerCfg, err := cfg.ProviderByName(providerName) + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + selected, err := cfg.ProviderByName(providerName) if err != nil { return ErrProviderNotFound } - if !s.registry.Supports(providerCfg.Driver) { + if !s.registry.Supports(selected.Driver) { return ErrDriverNotFound } - cfg.SelectedProvider = providerCfg.Name - cfg.CurrentModel = selectModel(cfg.CurrentModel, providerCfg.SupportedModels(), providerCfg.Model) + cfg.SelectedProvider = selected.Name + nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) + cfg.CurrentModel = nextModel selection = selectionFromConfig(*cfg) return nil }) @@ -73,6 +103,7 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov return ProviderSelection{}, err } + s.queueRefresh(providerCfg) return selection, nil } @@ -90,7 +121,10 @@ func (s *Service) ListModels(ctx context.Context) ([]ModelDescriptor, error) { return nil, err } - return modelDescriptors(selected.SupportedModels()), nil + return s.modelsForProvider(ctx, selected, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }), nil } func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (ProviderSelection, error) { @@ -103,18 +137,72 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider return ProviderSelection{}, ErrModelNotFound } + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } + + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }) + if !containsModelDescriptorID(models, modelID) { + return ProviderSelection{}, ErrModelNotFound + } + var selection ProviderSelection - err := s.manager.Update(ctx, func(cfg *config.Config) error { - selected, err := cfg.SelectedProviderConfig() - if err != nil { + err = s.manager.Update(ctx, func(cfg *config.Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { return err } + cfg.CurrentModel = modelID + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err + } - if !config.ContainsModelID(selected.SupportedModels(), modelID) { - return ErrModelNotFound - } + return selection, nil +} - cfg.CurrentModel = modelID +func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { + if err := s.validate(); err != nil { + return nil, err + } + return s.registry.Build(ctx, cfg) +} + +func (s *Service) EnsureSelection(ctx context.Context) (ProviderSelection, error) { + if err := s.validate(); err != nil { + return ProviderSelection{}, err + } + if err := ctx.Err(); err != nil { + return ProviderSelection{}, err + } + + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } + + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }) + nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) + if !changed { + return selectionFromConfig(cfgSnapshot), nil + } + + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { + return err + } + cfg.CurrentModel = nextModel selection = selectionFromConfig(*cfg) return nil }) @@ -125,6 +213,139 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider return selection, nil } +func (s *Service) validate() error { + if s == nil || s.manager == nil { + return ErrServiceManagerNil + } + if s.registry == nil { + return ErrServiceRegistryNil + } + return nil +} + +type modelQueryOptions struct { + allowSyncRefresh bool + queueRefresh bool +} + +func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.ProviderConfig, options modelQueryOptions) []ModelDescriptor { + defaultModels := modelDescriptorsFromIDs([]string{providerCfg.Model}) + + cached, cachedOK := s.loadCatalogModels(ctx, providerCfg) + if !cachedOK && options.allowSyncRefresh { + discovered, ok := s.discoverAndPersist(ctx, providerCfg) + if ok { + cached = discovered + cachedOK = true + } + } + + if cachedOK && options.queueRefresh && s.catalogExpired(ctx, providerCfg) { + s.queueRefresh(providerCfg) + } + + return MergeModelDescriptors(cached, defaultModels) +} + +func (s *Service) catalogExpired(ctx context.Context, providerCfg config.ProviderConfig) bool { + catalog, err := s.loadCatalog(ctx, providerCfg) + if err != nil { + return false + } + return catalog.Expired(s.now()) +} + +func (s *Service) loadCatalogModels(ctx context.Context, providerCfg config.ProviderConfig) ([]ModelDescriptor, bool) { + catalog, err := s.loadCatalog(ctx, providerCfg) + if err != nil { + return nil, false + } + return catalog.Models, true +} + +func (s *Service) loadCatalog(ctx context.Context, providerCfg config.ProviderConfig) (ModelCatalog, error) { + if s.catalogs == nil { + return ModelCatalog{}, ErrModelCatalogNotFound + } + + identity, err := providerCfg.Identity() + if err != nil { + return ModelCatalog{}, err + } + return s.catalogs.Load(ctx, identity) +} + +func (s *Service) discoverAndPersist(ctx context.Context, providerCfg config.ProviderConfig) ([]ModelDescriptor, bool) { + if s.registry == nil { + return nil, false + } + if !s.registry.Supports(providerCfg.Driver) { + return nil, false + } + + resolved, err := providerCfg.Resolve() + if err != nil { + return nil, false + } + + discovered, err := s.registry.DiscoverModels(ctx, resolved) + if err != nil { + return nil, false + } + + discovered = MergeModelDescriptors(discovered) + if s.catalogs == nil { + return discovered, true + } + + identity, err := providerCfg.Identity() + if err != nil { + return discovered, true + } + + now := s.now() + _ = s.catalogs.Save(ctx, ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + FetchedAt: now, + ExpiresAt: now.Add(s.catalogTTL), + Models: discovered, + }) + return discovered, true +} + +func (s *Service) queueRefresh(providerCfg config.ProviderConfig) { + if s.catalogs == nil || s.registry == nil || !s.registry.Supports(providerCfg.Driver) { + return + } + + identity, err := providerCfg.Identity() + if err != nil { + return + } + + key := identity.Key() + s.refreshMu.Lock() + if _, exists := s.inFlightByID[key]; exists { + s.refreshMu.Unlock() + return + } + s.inFlightByID[key] = struct{}{} + s.refreshMu.Unlock() + + go func() { + defer func() { + s.refreshMu.Lock() + delete(s.inFlightByID, key) + s.refreshMu.Unlock() + }() + + ctx, cancel := context.WithTimeout(context.Background(), s.backgroundTimeout) + defer cancel() + _, _ = s.discoverAndPersist(ctx, providerCfg) + }() +} + func selectionFromConfig(cfg config.Config) ProviderSelection { return ProviderSelection{ ProviderID: cfg.SelectedProvider, @@ -132,55 +353,38 @@ func selectionFromConfig(cfg config.Config) ProviderSelection { } } -func selectModel(currentModel string, models []string, fallback string) string { - if config.ContainsModelID(models, currentModel) { - return strings.TrimSpace(currentModel) +func resolveCurrentModel(currentModel string, models []ModelDescriptor, fallback string) (string, bool) { + currentModel = strings.TrimSpace(currentModel) + if containsModelDescriptorID(models, currentModel) { + return currentModel, false + } + + fallback = strings.TrimSpace(fallback) + if fallback != "" && containsModelDescriptorID(models, fallback) { + return fallback, currentModel != fallback } - return strings.TrimSpace(fallback) + + return currentModel, false } -func catalogItemFromConfig(cfg config.ProviderConfig) ProviderCatalogItem { +func catalogItemFromConfig(cfg config.ProviderConfig, models []ModelDescriptor) ProviderCatalogItem { return ProviderCatalogItem{ ID: strings.TrimSpace(cfg.Name), Name: strings.TrimSpace(cfg.Name), - Models: modelDescriptors(cfg.SupportedModels()), + Models: MergeModelDescriptors(models), } } -func modelDescriptors(models []string) []ModelDescriptor { - if len(models) == 0 { - return nil +func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { + target := config.NormalizeKey(modelID) + if target == "" { + return false } - descriptors := make([]ModelDescriptor, 0, len(models)) - for _, modelID := range models { - modelID = strings.TrimSpace(modelID) - if modelID == "" { - continue + for _, model := range models { + if config.NormalizeKey(model.ID) == target { + return true } - descriptors = append(descriptors, ModelDescriptor{ - ID: modelID, - Name: modelID, - }) } - return descriptors -} - -// Build 根据已解析的 provider 配置构建 Provider 实例。 -// 使 Service 满足 runtime.ProviderFactory 接口，简化装配层依赖。 -func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { - if err := s.validate(); err != nil { - return nil, err - } - return s.registry.Build(ctx, cfg) -} - -func (s *Service) validate() error { - if s == nil || s.manager == nil { - return errServiceManagerNil - } - if s.registry == nil { - return errServiceRegistryNil - } - return nil + return false } diff --git a/internal/provider/service_test.go b/internal/provider/service_test.go index 248dd31f..6a74067a 100644 --- a/internal/provider/service_test.go +++ b/internal/provider/service_test.go @@ -2,7 +2,9 @@ package provider import ( "context" + "sync" "testing" + "time" "neo-code/internal/config" ) @@ -10,428 +12,399 @@ import ( func TestNewService(t *testing.T) { t.Parallel() - t.Run("success", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - registry := NewRegistry() - service := NewService(manager, registry) - if service == nil { - t.Fatal("expected non-nil service") - } - }) - - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - registry := NewRegistry() - service := NewService(nil, registry) - if service == nil { - t.Fatal("expected non-nil service even with nil manager") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + registry := NewRegistry() + store := newMemoryCatalogStore() - t.Run("nil registry", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - service := NewService(manager, nil) - if service == nil { - t.Fatal("expected non-nil service even with nil registry") - } - }) + service := NewService(manager, registry, store) + if service == nil { + t.Fatal("expected non-nil service") + } + if service.catalogs != store { + t.Fatal("expected explicit catalog store to be used") + } } -func TestServiceListProviders(t *testing.T) { +func TestServiceListProvidersReturnsBuiltinCatalog(t *testing.T) { t.Parallel() - t.Run("success", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } - registry := newServiceTestRegistry(t) - service := NewService(manager, registry) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register test driver: %v", err) + } - items, err := service.ListProviders(context.Background()) - if err != nil { - t.Fatalf("ListProviders() error = %v", err) - } - if len(items) == 0 { - t.Fatal("expected at least one provider") - } - if items[0].ID != testProviderName { - t.Fatalf("expected first provider ID %q, got %q", testProviderName, items[0].ID) + service := NewService(manager, registry, newMemoryCatalogStore()) + items, err := service.ListProviders(context.Background()) + if err != nil { + t.Fatalf("ListProviders() error = %v", err) + } + expected := config.DefaultProviders() + if len(items) != len(expected) { + t.Fatalf("expected %d builtin providers, got %d", len(expected), len(items)) + } + expectedNames := make(map[string]struct{}, len(expected)) + for _, providerCfg := range expected { + expectedNames[providerCfg.Name] = struct{}{} + } + for _, item := range items { + if _, ok := expectedNames[item.ID]; !ok { + t.Fatalf("unexpected provider in catalog: %+v", item) } - }) - - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.ListProviders(context.Background()) - if err == nil { - t.Fatal("expected error for nil service") + if len(item.Models) != 1 { + t.Fatalf("expected default fallback model for %q, got %+v", item.ID, item.Models) } - }) + } +} - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - service := &Service{} - _, err := service.ListProviders(context.Background()) - if err == nil { - t.Fatal("expected error for nil manager") - } - }) +func TestServiceSelectProviderFallsBackToProviderDefault(t *testing.T) { + t.Parallel() - t.Run("nil registry", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - service := &Service{manager: manager} - _, err := service.ListProviders(context.Background()) - if err == nil { - t.Fatal("expected error for nil registry") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.CurrentModel = "gpt-5.4" + return nil + }); err != nil { + t.Fatalf("seed current model: %v", err) + } - t.Run("canceled context", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register openai driver: %v", err) + } - ctx, cancel := context.WithCancel(context.Background()) - cancel() - _, err := service.ListProviders(ctx) - if err == nil { - t.Fatal("expected error for canceled context") - } - }) + service := NewService(manager, registry, newMemoryCatalogStore()) + selection, err := service.SelectProvider(context.Background(), config.QiniuName) + if err != nil { + t.Fatalf("SelectProvider() error = %v", err) + } + if selection.ProviderID != config.QiniuName || selection.ModelID != config.QiniuDefaultModel { + t.Fatalf("unexpected selection: %+v", selection) + } } -func TestServiceSelectProvider(t *testing.T) { - t.Run("success", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") +func TestServiceEnsureSelectionRepairsInvalidCurrentModel(t *testing.T) { + t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.CurrentModel = "unsupported-current" + return nil + }); err != nil { + t.Fatalf("seed invalid current model: %v", err) + } - registry := newServiceTestRegistry(t) - service := NewService(manager, registry) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register openai driver: %v", err) + } - selection, err := service.SelectProvider(context.Background(), testProviderName) - if err != nil { - t.Fatalf("SelectProvider() error = %v", err) - } - if selection.ProviderID != testProviderName { - t.Fatalf("expected provider ID %q, got %q", testProviderName, selection.ProviderID) - } - }) + service := NewService(manager, registry, newMemoryCatalogStore()) + selection, err := service.EnsureSelection(context.Background()) + if err != nil { + t.Fatalf("EnsureSelection() error = %v", err) + } + if selection.ProviderID != config.OpenAIName || selection.ModelID != config.OpenAIDefaultModel { + t.Fatalf("unexpected normalized selection: %+v", selection) + } - t.Run("provider not found", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + reloaded, err := manager.Reload(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + if reloaded.CurrentModel != config.OpenAIDefaultModel { + t.Fatalf("expected rewritten current model %q, got %q", config.OpenAIDefaultModel, reloaded.CurrentModel) + } +} - _, err := service.SelectProvider(context.Background(), "nonexistent") - if err == nil { - t.Fatal("expected error for nonexistent provider") - } - }) +func TestServiceListModelsFallsBackToDefaultModelWithoutDiscovery(t *testing.T) { + t.Parallel() - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.SelectProvider(context.Background(), testProviderName) - if err == nil { - t.Fatal("expected error for nil service") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + service := NewService(manager, NewRegistry(), newMemoryCatalogStore()) + models, err := service.ListModels(context.Background()) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + if len(models) != 1 || models[0].ID != config.OpenAIDefaultModel { + t.Fatalf("expected default model fallback, got %+v", models) + } } -func TestServiceListModels(t *testing.T) { - t.Run("success", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") +func TestServiceSetCurrentModelUsesDiscoveredModels(t *testing.T) { + t.Setenv(testAPIKeyEnv, "test-key") - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } - registry := NewRegistry() - service := NewService(manager, registry) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + return []ModelDescriptor{ + {ID: cfg.Model, Name: cfg.Model}, + {ID: "gpt-4o", Name: "GPT-4o"}, + }, nil + })); err != nil { + t.Fatalf("register discovery driver: %v", err) + } - models, err := service.ListModels(context.Background()) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if len(models) == 0 { - t.Fatal("expected at least one model") - } - }) + service := NewService(manager, registry, newMemoryCatalogStore()) + selection, err := service.SetCurrentModel(context.Background(), "gpt-4o") + if err != nil { + t.Fatalf("SetCurrentModel() error = %v", err) + } + if selection.ModelID != "gpt-4o" { + t.Fatalf("expected selected model %q, got %+v", "gpt-4o", selection) + } +} - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.ListModels(context.Background()) - if err == nil { - t.Fatal("expected error for nil service") - } - }) +func TestServiceListModelsDiscoversAndCachesOnMiss(t *testing.T) { + t.Setenv(testAPIKeyEnv, "test-key") - t.Run("canceled context", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + store := newMemoryCatalogStore() + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + return []ModelDescriptor{{ + ID: "server-model", + Name: "Server Model", + ContextWindow: 32000, + MaxOutputTokens: 4096, + }}, nil + })); err != nil { + t.Fatalf("register discovery driver: %v", err) + } - ctx, cancel := context.WithCancel(context.Background()) - cancel() - _, err := service.ListModels(ctx) - if err == nil { - t.Fatal("expected error for canceled context") - } - }) + service := NewService(manager, registry, store) + models, err := service.ListModels(context.Background()) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + if !containsModelDescriptorID(models, "server-model") { + t.Fatalf("expected discovered model in result, got %+v", models) + } + + identity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } + catalog, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() cached catalog error = %v", err) + } + if !containsModelDescriptorID(catalog.Models, "server-model") { + t.Fatalf("expected cached discovered model, got %+v", catalog.Models) + } } -func TestServiceSetCurrentModel(t *testing.T) { - t.Run("success", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") +func TestServiceListModelsReturnsStaleCacheAndRefreshesInBackground(t *testing.T) { + t.Setenv(testAPIKeyEnv, "test-key") - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } - registry := NewRegistry() - service := NewService(manager, registry) + identity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } - newModel := "gpt-4o" - selection, err := service.SetCurrentModel(context.Background(), newModel) - if err != nil { - t.Fatalf("SetCurrentModel() error = %v", err) - } - if selection.ModelID != newModel { - t.Fatalf("expected model ID %q, got %q", newModel, selection.ModelID) - } - }) + store := newMemoryCatalogStore() + now := time.Date(2026, 4, 2, 12, 0, 0, 0, time.UTC) + if err := store.Save(context.Background(), ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + FetchedAt: now.Add(-48 * time.Hour), + ExpiresAt: now.Add(-24 * time.Hour), + Models: []ModelDescriptor{{ + ID: "stale-model", + Name: "Stale Model", + }}, + }); err != nil { + t.Fatalf("seed stale catalog: %v", err) + } - t.Run("empty model ID", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + refreshed := make(chan struct{}, 1) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + select { + case refreshed <- struct{}{}: + default: + } + return []ModelDescriptor{{ID: "fresh-model", Name: "Fresh Model"}}, nil + })); err != nil { + t.Fatalf("register custom driver: %v", err) + } - _, err := service.SetCurrentModel(context.Background(), "") - if err == nil { - t.Fatal("expected error for empty model ID") - } - }) + service := NewService(manager, registry, store) + service.now = func() time.Time { return now } + service.backgroundTimeout = time.Second - t.Run("model not found", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + models, err := service.ListModels(context.Background()) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + if !containsModelDescriptorID(models, "stale-model") { + t.Fatalf("expected stale cached model to be returned immediately, got %+v", models) + } - _, err := service.SetCurrentModel(context.Background(), "nonexistent-model") - if err == nil { - t.Fatal("expected error for nonexistent model") - } - }) + select { + case <-refreshed: + case <-time.After(2 * time.Second): + t.Fatal("expected background refresh to run") + } - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.SetCurrentModel(context.Background(), "model") - if err == nil { - t.Fatal("expected error for nil service") + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + catalog, err := store.Load(context.Background(), identity) + if err == nil && containsModelDescriptorID(catalog.Models, "fresh-model") { + return } - }) + time.Sleep(20 * time.Millisecond) + } + + catalog, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() refreshed catalog error = %v", err) + } + t.Fatalf("expected refreshed catalog to contain fresh-model, got %+v", catalog.Models) } -func TestServiceBuild(t *testing.T) { +func TestServiceBuildAndValidate(t *testing.T) { t.Parallel() - t.Run("nil service", func(t *testing.T) { + t.Run("build delegates to registry", func(t *testing.T) { t.Parallel() - var service *Service - _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) - if err == nil { - t.Fatal("expected error for nil service") - } - }) - - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - service := &Service{} - _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) - if err == nil { - t.Fatal("expected error for nil manager") - } - }) - t.Run("nil registry", func(t *testing.T) { - t.Parallel() manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - service := &Service{manager: manager} - _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) - if err == nil { - t.Fatal("expected error for nil registry") - } - }) -} - -func TestServiceValidate(t *testing.T) { - t.Parallel() - - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - err := service.validate() - if err == nil { - t.Fatal("expected error for nil service") + registry := NewRegistry() + if err := registry.Register(testDriverDefinition("custom", nil)); err != nil { + t.Fatalf("register driver: %v", err) + } + + service := NewService(manager, registry, newMemoryCatalogStore()) + providerInstance, err := service.Build(context.Background(), config.ResolvedProviderConfig{ + ProviderConfig: config.ProviderConfig{ + Name: "custom", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "model", + APIKeyEnv: "CUSTOM_API_KEY", + }, + APIKey: "test-key", + }) + if err != nil { + t.Fatalf("Build() error = %v", err) } - }) - - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - service := &Service{} - err := service.validate() - if err == nil { - t.Fatal("expected error for nil manager") + if _, ok := providerInstance.(serviceTestProvider); !ok { + t.Fatalf("expected serviceTestProvider, got %T", providerInstance) } }) - t.Run("nil registry", func(t *testing.T) { + t.Run("nil service fails validate", func(t *testing.T) { t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - service := &Service{manager: manager} - err := service.validate() - if err == nil { - t.Fatal("expected error for nil registry") - } - }) - t.Run("success", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - registry := NewRegistry() - service := NewService(manager, registry) - err := service.validate() - if err != nil { - t.Fatalf("expected no error, got %v", err) + var service *Service + if err := service.validate(); err == nil { + t.Fatal("expected validate error for nil service") } }) } -func TestSelectModelHelper(t *testing.T) { +func TestResolveCurrentModelHelper(t *testing.T) { t.Parallel() tests := []struct { name string currentModel string - models []string + models []ModelDescriptor fallback string expected string + changed bool }{ { name: "current model in list", currentModel: "gpt-4o", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4o", + models: []ModelDescriptor{ + {ID: "gpt-4.1"}, + {ID: "gpt-4o"}, + {ID: "gpt-5.4"}, + }, + fallback: "gpt-4.1", + expected: "gpt-4o", + changed: false, }, { - name: "current model not in list", + name: "current model falls back to provider default", currentModel: "unknown-model", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4.1", + models: []ModelDescriptor{ + {ID: "gpt-4.1"}, + {ID: "gpt-4o"}, + {ID: "gpt-5.4"}, + }, + fallback: "gpt-4.1", + expected: "gpt-4.1", + changed: true, }, { - name: "empty current model", - currentModel: "", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4.1", - }, - { - name: "current model with whitespace", - currentModel: " gpt-4o ", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4o", + name: "missing fallback keeps current model unchanged", + currentModel: "unknown-model", + models: []ModelDescriptor{ + {ID: "gpt-4o"}, + }, + fallback: "gpt-4.1", + expected: "unknown-model", + changed: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - result := selectModel(tt.currentModel, tt.models, tt.fallback) - if result != tt.expected { - t.Fatalf("expected %q, got %q", tt.expected, result) + got, changed := resolveCurrentModel(tt.currentModel, tt.models, tt.fallback) + if got != tt.expected || changed != tt.changed { + t.Fatalf( + "resolveCurrentModel() = (%q, %v), want (%q, %v)", + got, + changed, + tt.expected, + tt.changed, + ) } }) } } -func TestModelDescriptorsHelper(t *testing.T) { +func TestModelDescriptorsFromIDsHelper(t *testing.T) { t.Parallel() - t.Run("empty models", func(t *testing.T) { - t.Parallel() - result := modelDescriptors(nil) - if result != nil { - t.Fatalf("expected nil for empty models, got %v", result) - } - }) - - t.Run("non-empty models", func(t *testing.T) { - t.Parallel() - models := []string{"gpt-4.1", "gpt-4o", "gpt-5.4"} - result := modelDescriptors(models) - if len(result) != 3 { - t.Fatalf("expected 3 descriptors, got %d", len(result)) - } - for i, desc := range result { - if desc.ID != models[i] { - t.Fatalf("expected ID %q, got %q", models[i], desc.ID) - } - if desc.Name != models[i] { - t.Fatalf("expected name %q, got %q", models[i], desc.Name) - } - } - }) - - t.Run("models with empty strings", func(t *testing.T) { - t.Parallel() - models := []string{"gpt-4.1", "", "gpt-4o", " ", "gpt-5.4"} - result := modelDescriptors(models) - if len(result) != 3 { - t.Fatalf("expected 3 descriptors after filtering empty strings, got %d", len(result)) - } - }) + models := modelDescriptorsFromIDs([]string{"gpt-4.1", "", "gpt-4o"}) + if len(models) != 2 { + t.Fatalf("expected 2 descriptors, got %d", len(models)) + } + if models[0].ID != "gpt-4.1" || models[1].ID != "gpt-4o" { + t.Fatalf("unexpected descriptors: %+v", models) + } } func testDefaultConfig() *config.Config { @@ -445,19 +418,14 @@ func testDefaultConfig() *config.Config { return cfg } -func newServiceTestRegistry(t *testing.T) *Registry { - t.Helper() - - registry := NewRegistry() - if err := registry.Register(DriverDefinition{ - Name: testProviderName, +func testDriverDefinition(name string, discover DiscoveryFunc) DriverDefinition { + return DriverDefinition{ + Name: name, + Discover: discover, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { return serviceTestProvider{}, nil }, - }); err != nil { - t.Fatalf("register test driver: %v", err) } - return registry } type serviceTestProvider struct{} @@ -466,6 +434,44 @@ func (serviceTestProvider) Chat(ctx context.Context, req ChatRequest, events cha return ChatResponse{}, nil } +type memoryCatalogStore struct { + mu sync.Mutex + catalogs map[string]ModelCatalog +} + +func newMemoryCatalogStore() *memoryCatalogStore { + return &memoryCatalogStore{ + catalogs: map[string]ModelCatalog{}, + } +} + +func (s *memoryCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return ModelCatalog{}, err + } + + s.mu.Lock() + defer s.mu.Unlock() + + catalog, ok := s.catalogs[identity.Key()] + if !ok { + return ModelCatalog{}, ErrModelCatalogNotFound + } + return catalog, nil +} + +func (s *memoryCatalogStore) Save(ctx context.Context, catalog ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + s.mu.Lock() + defer s.mu.Unlock() + + s.catalogs[catalog.Identity.Key()] = catalog + return nil +} + const ( testProviderName = "openai" testAPIKeyEnv = "OPENAI_API_KEY" diff --git a/internal/provider/types.go b/internal/provider/types.go index 2a08b2f0..ee5b8443 100644 --- a/internal/provider/types.go +++ b/internal/provider/types.go @@ -1,9 +1,5 @@ package provider -import ( - "errors" -) - const ( RoleSystem = "system" RoleUser = "user" @@ -50,16 +46,13 @@ type Usage struct { TotalTokens int `json:"total_tokens"` } -var ( - ErrProviderNotFound = errors.New("provider not found") - ErrModelNotFound = errors.New("model not found") - ErrDriverNotFound = errors.New("provider driver not found") -) - type ModelDescriptor struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description,omitempty"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + ContextWindow int `json:"context_window,omitempty"` + MaxOutputTokens int `json:"max_output_tokens,omitempty"` + Capabilities map[string]bool `json:"capabilities,omitempty"` } type ProviderCatalogItem struct { diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 40207d75..6942e988 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -956,10 +956,12 @@ func TestAppUpdateModelPickerEnterAppliesSelection(t *testing.T) { t.Fatalf("expected picker to close after selection") } - for _, msg := range collectTeaMessages(cmd) { - model, follow := app.Update(msg) - app = model.(App) - _ = collectTeaMessages(follow) + if cmd != nil { + if msg := cmd(); msg != nil { + model, follow := app.Update(msg) + app = model.(App) + _ = collectTeaMessages(follow) + } } cfg := manager.Get() @@ -971,18 +973,10 @@ func TestAppUpdateModelPickerEnterAppliesSelection(t *testing.T) { func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { manager := newTestConfigManager(t) if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "openai-alt", - Driver: "openai", - BaseURL: config.OpenAIDefaultBaseURL, - Model: "gpt-4o", - Models: []string{"gpt-4o"}, - APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, - }) cfg.CurrentModel = "unsupported-current" return nil }); err != nil { - t.Fatalf("append provider: %v", err) + t.Fatalf("set unsupported current model: %v", err) } runtime := newStubRuntime() @@ -1002,14 +996,14 @@ func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { if !ok { continue } - if candidate.id == "openai-alt" { + if candidate.id == config.QiniuName { selectedIndex = idx selected = candidate.id break } } if selectedIndex < 0 { - t.Fatalf("expected provider picker to include openai-alt") + t.Fatalf("expected provider picker to include %s", config.QiniuName) } app.providerPicker.Select(selectedIndex) @@ -1029,7 +1023,7 @@ func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { if cfg.SelectedProvider != selected { t.Fatalf("expected selected provider %q, got %q", selected, cfg.SelectedProvider) } - if cfg.CurrentModel != "gpt-4o" { + if cfg.CurrentModel != config.QiniuDefaultModel { t.Fatalf("expected current model to follow provider default, got %q", cfg.CurrentModel) } } @@ -1586,11 +1580,55 @@ func newTestConfigManager(t *testing.T) *config.Manager { func newTestProviderService(t *testing.T, manager *config.Manager) *provider.Service { t.Helper() - registry, err := builtin.NewRegistry() + + registry := provider.NewRegistry() + err := registry.Register(provider.DriverDefinition{ + Name: config.OpenAIName, + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { + return tuiTestProvider{}, nil + }, + }) if err != nil { t.Fatalf("register provider drivers: %v", err) } - return provider.NewService(manager, registry) + return provider.NewService(manager, registry, newTUITestCatalogStore()) +} + +type tuiTestProvider struct{} + +func (tuiTestProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { + return provider.ChatResponse{}, nil +} + +type tuiTestCatalogStore struct { + catalogs map[string]provider.ModelCatalog +} + +func newTUITestCatalogStore() *tuiTestCatalogStore { + return &tuiTestCatalogStore{ + catalogs: map[string]provider.ModelCatalog{}, + } +} + +func (s *tuiTestCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (provider.ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return provider.ModelCatalog{}, err + } + + catalog, ok := s.catalogs[identity.Key()] + if !ok { + return provider.ModelCatalog{}, provider.ErrModelCatalogNotFound + } + return catalog, nil +} + +func (s *tuiTestCatalogStore) Save(ctx context.Context, catalog provider.ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + s.catalogs[catalog.Identity.Key()] = catalog + return nil } func collectTeaMessages(cmd tea.Cmd) []tea.Msg { @@ -1605,7 +1643,7 @@ func collectTeaMessages(cmd tea.Cmd) []tea.Msg { var msg tea.Msg select { case msg = <-msgCh: - case <-time.After(25 * time.Millisecond): + case <-time.After(250 * time.Millisecond): return nil } if msg == nil {