diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 6c8a4037..d87241ea 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -10,7 +10,6 @@ import ( agentcontext "neo-code/internal/context" "neo-code/internal/provider/builtin" providercatalog "neo-code/internal/provider/catalog" - providerselection "neo-code/internal/provider/selection" agentruntime "neo-code/internal/runtime" "neo-code/internal/security" "neo-code/internal/tools" @@ -38,8 +37,8 @@ func ensureConsoleUTF8() { func NewProgram(ctx context.Context) (*tea.Program, error) { ensureConsoleUTF8() - - loader := config.NewLoader("", config.DefaultConfig()) + + loader := config.NewLoader("", config.DefaultConfig()) manager := config.NewManager(loader) if _, err := manager.Load(ctx); err != nil { return nil, err @@ -50,7 +49,7 @@ func NewProgram(ctx context.Context) (*tea.Program, error) { return nil, err } modelCatalogs := providercatalog.NewService(manager.BaseDir(), providerRegistry, nil) - providerSelection := providerselection.NewService(manager, providerRegistry, modelCatalogs) + providerSelection := config.NewSelectionService(manager, providerRegistry, providerRegistry, modelCatalogs) if _, err := providerSelection.EnsureSelection(ctx); err != nil { return nil, err } diff --git a/internal/config/model.go b/internal/config/model.go index 9c0040be..e8d1d3c9 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -535,16 +535,3 @@ func containsProviderName(providers []ProviderConfig, name string) bool { } return false } - -func ContainsModelID(models []string, model string) bool { - target := NormalizeKey(model) - if target == "" { - return false - } - for _, candidate := range models { - if NormalizeKey(candidate) == target { - return true - } - } - return false -} diff --git a/internal/config/selection.go b/internal/config/selection.go new file mode 100644 index 00000000..54963402 --- /dev/null +++ b/internal/config/selection.go @@ -0,0 +1,538 @@ +package config + +import ( + "context" + "errors" + "strings" +) + +// Selection 领域错误。 +var ( + ErrProviderNotFound = errors.New("provider not found") + ErrModelNotFound = errors.New("model not found") + ErrDriverUnsupported = errors.New("provider driver not supported by current runtime") +) + +// ---- 数据类型定义 ---- + +// ModelDescriptor 表示模型元数据描述符。 +type ModelDescriptor struct { + 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"` +} + +// ProviderCatalogItem 表示一个已配置的 provider 及其可用模型列表（用于 UI 展示）。 +type ProviderCatalogItem struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Models []ModelDescriptor `json:"models,omitempty"` +} + +// ProviderSelection 表示当前选中的 provider 和 model。 +type ProviderSelection struct { + ProviderID string `json:"provider_id"` + ModelID string `json:"model_id"` +} + +// ---- 接口定义（避免循环依赖）---- + +// DriverSupporter 检查给定驱动类型是否被支持。 +type DriverSupporter interface { + Supports(driverType string) bool +} + +// ModelDiscovery 定义模型发现能力，由 provider.Registry 提供。 +type ModelDiscovery interface { + DiscoverModels(ctx context.Context, cfg ResolvedProviderConfig) ([]ModelDescriptor, error) +} + +// ---- 模型描述符工具函数 ---- + +// DescriptorFromRawModel 将原始 provider 模型对象标准化为 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 +} + +// MergeModelDescriptors 合并多个 ModelDescriptor 来源，按 ID 去重， +// 优先保留较早来源的字段值，后续来源用于回填空字段。 +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 := NormalizeKey(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 +} + +// DescriptorsFromIDs 从模型 ID 字符串列表构建最小化的 ModelDescriptor 列表。 +func DescriptorsFromIDs(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 +} + +// ---- 选择服务 ---- + +// ModelCatalog 定义模型目录查询接口，用于获取 provider 可用模型列表。 +type ModelCatalog interface { + ListProviderModels(ctx context.Context, providerCfg ProviderConfig) ([]ModelDescriptor, error) + ListProviderModelsSnapshot(ctx context.Context, providerCfg ProviderConfig) ([]ModelDescriptor, error) + ListProviderModelsCached(ctx context.Context, providerCfg ProviderConfig) ([]ModelDescriptor, error) +} + +// SelectionService 管理 provider/模型选择状态，负责切换当前 provider、切换当前模型、 +// 校验和修复无效的 selection 等。所有变更通过 ConfigManager 持久化到配置中。 +// +// 通过 DriverSupporter 和 ModelDiscovery 接口注入 provider 能力，避免 config -> provider 循环依赖。 +type SelectionService struct { + manager *Manager + supporters DriverSupporter + discovery ModelDiscovery + catalogs ModelCatalog +} + +// NewSelectionService 创建选择服务实例。 +func NewSelectionService(manager *Manager, supporters DriverSupporter, discovery ModelDiscovery, catalogs ModelCatalog) *SelectionService { + return &SelectionService{ + manager: manager, + supporters: supporters, + discovery: discovery, + catalogs: catalogs, + } +} + +// ListProviders 枚举所有已配置且驱动支持的 provider 及其可用模型。 +func (s *SelectionService) ListProviders(ctx context.Context) ([]ProviderCatalogItem, error) { + if err := s.validate(); err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + + cfg := s.manager.Get() + items := make([]ProviderCatalogItem, 0, len(cfg.Providers)) + for _, providerCfg := range cfg.Providers { + if err := ctx.Err(); err != nil { + return nil, err + } + if !s.supporters.Supports(providerCfg.Driver) { + continue + } + + models, err := s.catalogs.ListProviderModelsCached(ctx, providerCfg) + if err != nil { + return nil, err + } + items = append(items, providerCatalogItem(providerCfg, models)) + } + + return items, nil +} + +// SelectProvider 切换当前选中 provider，同时将 current_model 解析为有效值。 +func (s *SelectionService) SelectProvider(ctx context.Context, providerName string) (ProviderSelection, error) { + if err := s.validate(); err != nil { + return ProviderSelection{}, err + } + + cfgSnapshot := s.manager.Get() + providerCfg, err := cfgSnapshot.ProviderByName(providerName) + if err != nil { + return ProviderSelection{}, ErrProviderNotFound + } + if !s.supporters.Supports(providerCfg.Driver) { + return ProviderSelection{}, ErrDriverUnsupported + } + + models, err := s.catalogs.ListProviderModelsCached(ctx, providerCfg) + if err != nil { + return ProviderSelection{}, err + } + + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *Config) error { + selected, err := cfg.ProviderByName(providerName) + if err != nil { + return ErrProviderNotFound + } + if !s.supporters.Supports(selected.Driver) { + return ErrDriverUnsupported + } + + cfg.SelectedProvider = selected.Name + nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err + } + + return selection, nil +} + +// ListModels 获取当前选中 provider 的模型列表（可触发远程刷新）。 +func (s *SelectionService) ListModels(ctx context.Context) ([]ModelDescriptor, error) { + if err := s.validate(); err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + + selected, err := s.selectedProviderConfig() + if err != nil { + return nil, err + } + return s.catalogs.ListProviderModels(ctx, selected) +} + +// ListModelsSnapshot 获取当前选中 provider 的缓存模型列表。 +func (s *SelectionService) ListModelsSnapshot(ctx context.Context) ([]ModelDescriptor, error) { + if err := s.validate(); err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } + + selected, err := s.selectedProviderConfig() + if err != nil { + return nil, err + } + return s.catalogs.ListProviderModelsSnapshot(ctx, selected) +} + +// SetCurrentModel 切换当前使用的模型（必须在当前 provider 的可用模型列表中）。 +func (s *SelectionService) SetCurrentModel(ctx context.Context, modelID string) (ProviderSelection, error) { + if err := s.validate(); err != nil { + return ProviderSelection{}, err + } + + modelID = strings.TrimSpace(modelID) + if modelID == "" { + return ProviderSelection{}, ErrModelNotFound + } + + selected, err := s.selectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } + + models, err := s.catalogs.ListProviderModelsSnapshot(ctx, selected) + if err != nil { + return ProviderSelection{}, err + } + if !containsModelDescriptorID(models, modelID) { + return ProviderSelection{}, ErrModelNotFound + } + + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { + return err + } + cfg.CurrentModel = modelID + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err + } + + return selection, nil +} + +// EnsureSelection 确保当前 selection 有效：如果 current_model 不在可用模型列表中则自动修正。 +func (s *SelectionService) 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, err := s.catalogs.ListProviderModelsSnapshot(ctx, selected) + if err != nil { + return ProviderSelection{}, err + } + 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) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { + return err + } + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err + } + + return selection, nil +} + +func (s *SelectionService) validate() error { + if s == nil { + return errors.New("selection: service is nil") + } + if s.manager == nil { + return errors.New("selection: config manager is nil") + } + if s.supporters == nil { + return errors.New("selection: driver supporter is nil") + } + if s.catalogs == nil { + return errors.New("selection: catalog service is nil") + } + return nil +} + +func (s *SelectionService) selectedProviderConfig() (ProviderConfig, error) { + cfg := s.manager.Get() + return cfg.SelectedProviderConfig() +} + +// ---- 内部辅助函数 ---- + +func selectionFromConfig(cfg Config) ProviderSelection { + return ProviderSelection{ + ProviderID: cfg.SelectedProvider, + ModelID: cfg.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 currentModel, false +} + +func providerCatalogItem(cfg ProviderConfig, models []ModelDescriptor) ProviderCatalogItem { + return ProviderCatalogItem{ + ID: strings.TrimSpace(cfg.Name), + Name: strings.TrimSpace(cfg.Name), + Models: MergeModelDescriptors(models), + } +} + +func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { + target := NormalizeKey(modelID) + if target == "" { + return false + } + + for _, model := range models { + if NormalizeKey(model.ID) == target { + return true + } + } + return false +} + +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 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 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 +} diff --git a/internal/config/selection_test.go b/internal/config/selection_test.go new file mode 100644 index 00000000..4dd40946 --- /dev/null +++ b/internal/config/selection_test.go @@ -0,0 +1,382 @@ +package config + +import ( + "context" + "errors" + "strings" + "testing" +) + +func TestSelectionServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { + t.Parallel() + + manager := newSelectionTestManager(t, DefaultConfig()) + service := NewSelectionService(manager, newDriverSupporterStub(), nil, newCatalogStub()) + + items, err := service.ListProviders(context.Background()) + if err != nil { + t.Fatalf("ListProviders() error = %v", err) + } + expected := map[string]int{ + OpenAIName: 2, + GeminiName: 2, + OpenLLName: 2, + QiniuName: 2, + } + if len(items) != len(expected) { + t.Fatalf("expected only builtin providers, got %d", len(items)) + } + + for _, item := range items { + wantModels, ok := expected[item.ID] + if !ok { + t.Fatalf("unexpected builtin provider %q", item.ID) + } + if len(item.Models) != wantModels { + t.Fatalf("expected provider models to fall back to the default model, got %+v", item.Models) + } + delete(expected, item.ID) + } + if len(expected) != 0 { + t.Fatalf("missing builtin providers from catalog: %+v", expected) + } +} + +func TestSelectionServiceSelectProviderAndSetCurrentModel(t *testing.T) { + manager := newSelectionTestManager(t, DefaultConfig()) + + if err := manager.Update(context.Background(), func(cfg *Config) error { + cfg.CurrentModel = "gpt-5.4" + return nil + }); err != nil { + t.Fatalf("seed current model: %v", err) + } + + service := NewSelectionService(manager, newDriverSupporterStub(), nil, newCatalogStub()) + + selection, err := service.SelectProvider(context.Background(), QiniuName) + if err != nil { + t.Fatalf("SelectProvider() error = %v", err) + } + if selection.ProviderID != QiniuName || selection.ModelID != QiniuDefaultModel { + t.Fatalf("unexpected selection after switch: %+v", selection) + } + + if _, err := service.SetCurrentModel(context.Background(), "missing-model"); !errors.Is(err, ErrModelNotFound) { + t.Fatalf("expected ErrModelNotFound, got %v", err) + } + + selection, err = service.SetCurrentModel(context.Background(), QiniuDefaultModel+"-alt") + if err != nil { + t.Fatalf("SetCurrentModel() error = %v", err) + } + if selection.ModelID != QiniuDefaultModel+"-alt" { + t.Fatalf("expected selected model %q, got %+v", QiniuDefaultModel+"-alt", selection) + } + + cfg := manager.Get() + selected, err := cfg.SelectedProviderConfig() + if err != nil { + t.Fatalf("SelectedProviderConfig() error = %v", err) + } + if selected.Name != QiniuName { + t.Fatalf("expected selected provider %q, got %+v", QiniuName, selected) + } + if selected.Model != QiniuDefaultModel || cfg.CurrentModel != QiniuDefaultModel+"-alt" { + t.Fatalf("expected provider default and current model to diverge safely, got provider=%q current=%q", selected.Model, cfg.CurrentModel) + } +} + +func TestSelectionServiceEnsureSelectionRepairsInvalidCurrentModel(t *testing.T) { + t.Parallel() + + manager := newSelectionTestManager(t, DefaultConfig()) + if err := manager.Update(context.Background(), func(cfg *Config) error { + cfg.CurrentModel = "unsupported-current" + return nil + }); err != nil { + t.Fatalf("seed invalid current model: %v", err) + } + + service := NewSelectionService(manager, newDriverSupporterStub(), nil, newCatalogStub()) + + selection, err := service.EnsureSelection(context.Background()) + if err != nil { + t.Fatalf("EnsureSelection() error = %v", err) + } + if selection.ProviderID != OpenAIName || selection.ModelID != OpenAIDefaultModel { + t.Fatalf("unexpected normalized selection: %+v", selection) + } + + reloaded, err := manager.Reload(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + if reloaded.CurrentModel != OpenAIDefaultModel { + t.Fatalf("expected rewritten current model %q, got %q", OpenAIDefaultModel, reloaded.CurrentModel) + } +} + +func TestSelectionServiceSelectProviderRejectsUnsupportedDriver(t *testing.T) { + t.Parallel() + + defaults := DefaultConfig() + defaults.Providers = []ProviderConfig{{ + Name: OpenAIName, + Driver: "missing-driver", + BaseURL: OpenAIDefaultBaseURL, + Model: OpenAIDefaultModel, + APIKeyEnv: OpenAIDefaultAPIKeyEnv, + }} + defaults.SelectedProvider = OpenAIName + defaults.CurrentModel = OpenAIDefaultModel + + manager := newSelectionTestManager(t, defaults) + supporters := &selectiveDriverSupporter{supported: map[string]bool{"openai": true}} + service := NewSelectionService(manager, supporters, nil, newCatalogStub()) + + if _, err := service.SelectProvider(context.Background(), OpenAIName); !errors.Is(err, ErrDriverUnsupported) { + t.Fatalf("expected SelectProvider() to preserve driver error, got %v", err) + } +} + +func TestSelectionServiceValidateErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + service *SelectionService + errMsg string + }{ + { + name: "nil service", + service: nil, + errMsg: "selection: service is nil", + }, + { + name: "nil manager", + service: &SelectionService{ + manager: nil, + supporters: &driverSupporterAll{}, + catalogs: newCatalogStub(), + }, + errMsg: "selection: config manager is nil", + }, + { + name: "nil supporters", + service: &SelectionService{ + manager: NewManager(NewLoader(t.TempDir(), DefaultConfig())), + supporters: nil, + catalogs: newCatalogStub(), + }, + errMsg: "selection: driver supporter is nil", + }, + { + name: "nil catalogs", + service: &SelectionService{ + manager: NewManager(NewLoader(t.TempDir(), DefaultConfig())), + supporters: &driverSupporterAll{}, + catalogs: nil, + }, + errMsg: "selection: catalog service is nil", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + _, err := tt.service.ListProviders(context.Background()) + if err == nil || err.Error() != tt.errMsg { + t.Fatalf("expected error %q, got %v", tt.errMsg, err) + } + }) + } +} + +func TestSelectionServiceOperationsWithCanceledContext(t *testing.T) { + t.Parallel() + + manager := newSelectionTestManager(t, DefaultConfig()) + service := NewSelectionService(manager, newDriverSupporterStub(), nil, newCatalogStub()) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + tests := []struct { + name string + fn func(context.Context) error + }{ + { + name: "ListProviders", + fn: func(ctx context.Context) error { + _, err := service.ListProviders(ctx) + return err + }, + }, + { + name: "ListModels", + fn: func(ctx context.Context) error { + _, err := service.ListModels(ctx) + return err + }, + }, + { + name: "ListModelsSnapshot", + fn: func(ctx context.Context) error { + _, err := service.ListModelsSnapshot(ctx) + return err + }, + }, + { + name: "EnsureSelection", + fn: func(ctx context.Context) error { + _, err := service.EnsureSelection(ctx) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := tt.fn(ctx) + if err == nil { + t.Fatalf("expected error for canceled context") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + } +} + +func TestSelectionServiceSetCurrentModelEmptyModelID(t *testing.T) { + t.Parallel() + + manager := newSelectionTestManager(t, DefaultConfig()) + service := NewSelectionService(manager, newDriverSupporterStub(), nil, newCatalogStub()) + + _, err := service.SetCurrentModel(context.Background(), "") + if !errors.Is(err, ErrModelNotFound) { + t.Fatalf("expected ErrModelNotFound for empty model ID, got %v", err) + } + + _, err = service.SetCurrentModel(context.Background(), " ") + if !errors.Is(err, ErrModelNotFound) { + t.Fatalf("expected ErrModelNotFound for whitespace model ID, got %v", err) + } +} + +func TestResolveCurrentModelHelper(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + currentModel string + models []ModelDescriptor + fallback string + expected string + changed bool + }{ + { + name: "current model in list", + currentModel: "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 falls back to provider default", + currentModel: "unknown-model", + models: []ModelDescriptor{ + {ID: "gpt-4.1"}, + {ID: "gpt-4o"}, + {ID: "gpt-5.4"}, + }, + fallback: "gpt-4.1", + expected: "gpt-4.1", + changed: true, + }, + { + 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() + + 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) + } + }) + } +} + +// ---- 测试辅助 ---- + +type driverSupporterAll struct{} + +func (*driverSupporterAll) Supports(_ string) bool { return true } + +type selectiveDriverSupporter struct { + supported map[string]bool +} + +func (s *selectiveDriverSupporter) Supports(driverType string) bool { + return s.supported[NormalizeKey(driverType)] +} + +type catalogStub struct{} + +func (catalogStub) ListProviderModels(_ context.Context, cfg ProviderConfig) ([]ModelDescriptor, error) { + return defaultModelsForProvider(cfg), nil +} + +func (catalogStub) ListProviderModelsSnapshot(_ context.Context, cfg ProviderConfig) ([]ModelDescriptor, error) { + return defaultModelsForProvider(cfg), nil +} + +func (catalogStub) ListProviderModelsCached(_ context.Context, cfg ProviderConfig) ([]ModelDescriptor, error) { + return defaultModelsForProvider(cfg), nil +} + +// defaultModelsForProvider 为给定 provider 返回包含默认模型和额外测试模型的列表。 +func defaultModelsForProvider(cfg ProviderConfig) []ModelDescriptor { + model := strings.TrimSpace(cfg.Model) + if model == "" { + return nil + } + return []ModelDescriptor{ + {ID: model, Name: model}, + {ID: model + "-alt", Name: model + "-alt"}, + } +} + +func newCatalogStub() ModelCatalog { return catalogStub{} } +func newDriverSupporterStub() DriverSupporter { return &driverSupporterAll{} } + +func newSelectionTestManager(t *testing.T, defaults *Config) *Manager { + t.Helper() + + manager := NewManager(NewLoader(t.TempDir(), defaults)) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("load config: %v", err) + } + return manager +} diff --git a/internal/provider/catalog/service.go b/internal/provider/catalog/service.go index 730ba9f1..cb229990 100644 --- a/internal/provider/catalog/service.go +++ b/internal/provider/catalog/service.go @@ -42,34 +42,28 @@ func NewService(baseDir string, registry *provider.Registry, store Store) *Servi } } -func (s *Service) ListProviderModels(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, error) { - if err := s.validate(); err != nil { - return nil, err - } - if err := ctx.Err(); err != nil { - return nil, err - } - - return s.modelsForProvider(ctx, providerCfg, queryOptions{ +func (s *Service) ListProviderModels(ctx context.Context, providerCfg config.ProviderConfig) ([]config.ModelDescriptor, error) { + return s.listProviderModels(ctx, providerCfg, queryOptions{ allowSyncRefresh: true, queueRefresh: true, - }), nil + }) } -func (s *Service) ListProviderModelsSnapshot(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, error) { - if err := s.validate(); err != nil { - return nil, err - } - if err := ctx.Err(); err != nil { - return nil, err - } - - return s.modelsForProvider(ctx, providerCfg, queryOptions{ +func (s *Service) ListProviderModelsSnapshot(ctx context.Context, providerCfg config.ProviderConfig) ([]config.ModelDescriptor, error) { + return s.listProviderModels(ctx, providerCfg, queryOptions{ queueRefresh: true, - }), nil + }) +} + +func (s *Service) ListProviderModelsCached(ctx context.Context, providerCfg config.ProviderConfig) ([]config.ModelDescriptor, error) { + return s.listProviderModels(ctx, providerCfg, queryOptions{}) } -func (s *Service) ListProviderModelsCached(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, error) { +func (s *Service) listProviderModels( + ctx context.Context, + providerCfg config.ProviderConfig, + options queryOptions, +) ([]config.ModelDescriptor, error) { if err := s.validate(); err != nil { return nil, err } @@ -77,7 +71,7 @@ func (s *Service) ListProviderModelsCached(ctx context.Context, providerCfg conf return nil, err } - return s.modelsForProvider(ctx, providerCfg, queryOptions{}), nil + return s.modelsForProvider(ctx, providerCfg, options), nil } func (s *Service) validate() error { @@ -95,39 +89,50 @@ type queryOptions struct { queueRefresh bool } -func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.ProviderConfig, options queryOptions) []provider.ModelDescriptor { - defaultModels := descriptorsFromIDs([]string{providerCfg.Model}) +type catalogSnapshot struct { + models []config.ModelDescriptor + ok bool + expired bool + identity config.ProviderIdentity +} + +func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.ProviderConfig, options queryOptions) []config.ModelDescriptor { + defaultModels := config.DescriptorsFromIDs([]string{providerCfg.Model}) + snapshot := s.catalogSnapshot(ctx, providerCfg) - cached, cachedOK := s.loadCatalogModels(ctx, providerCfg) - if !cachedOK && options.allowSyncRefresh { + models := snapshot.models + modelsOK := snapshot.ok + if !modelsOK && options.allowSyncRefresh { discovered, ok := s.discoverAndPersist(ctx, providerCfg) if ok { - cached = discovered - cachedOK = true + models = discovered + modelsOK = true } } - if options.queueRefresh && (!cachedOK || s.catalogExpired(ctx, providerCfg)) { - s.queueRefresh(providerCfg) + if options.queueRefresh && (!modelsOK || snapshot.expired) { + s.queueRefresh(providerCfg, snapshot.identity) } - return provider.MergeModelDescriptors(cached, defaultModels) + return config.MergeModelDescriptors(models, defaultModels) } -func (s *Service) catalogExpired(ctx context.Context, providerCfg config.ProviderConfig) bool { - modelCatalog, err := s.loadCatalog(ctx, providerCfg) +func (s *Service) catalogSnapshot(ctx context.Context, providerCfg config.ProviderConfig) catalogSnapshot { + identity, err := providerCfg.Identity() if err != nil { - return false + return catalogSnapshot{} } - return modelCatalog.Expired(s.now()) -} -func (s *Service) loadCatalogModels(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, bool) { modelCatalog, err := s.loadCatalog(ctx, providerCfg) if err != nil { - return nil, false + return catalogSnapshot{identity: identity} + } + return catalogSnapshot{ + models: modelCatalog.Models, + ok: true, + expired: modelCatalog.Expired(s.now()), + identity: identity, } - return modelCatalog.Models, true } func (s *Service) loadCatalog(ctx context.Context, providerCfg config.ProviderConfig) (ModelCatalog, error) { @@ -142,7 +147,7 @@ func (s *Service) loadCatalog(ctx context.Context, providerCfg config.ProviderCo return s.store.Load(ctx, identity) } -func (s *Service) discoverAndPersist(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, bool) { +func (s *Service) discoverAndPersist(ctx context.Context, providerCfg config.ProviderConfig) ([]config.ModelDescriptor, bool) { if !s.registry.Supports(providerCfg.Driver) { return nil, false } @@ -157,7 +162,7 @@ func (s *Service) discoverAndPersist(ctx context.Context, providerCfg config.Pro return nil, false } - discovered = provider.MergeModelDescriptors(discovered) + discovered = config.MergeModelDescriptors(discovered) if s.store == nil { return discovered, true } @@ -178,13 +183,11 @@ func (s *Service) discoverAndPersist(ctx context.Context, providerCfg config.Pro return discovered, true } -func (s *Service) queueRefresh(providerCfg config.ProviderConfig) { +func (s *Service) queueRefresh(providerCfg config.ProviderConfig, identity config.ProviderIdentity) { if s.store == nil || !s.registry.Supports(providerCfg.Driver) { return } - - identity, err := providerCfg.Identity() - if err != nil { + if identity.Driver == "" || identity.BaseURL == "" { return } @@ -209,25 +212,3 @@ func (s *Service) queueRefresh(providerCfg config.ProviderConfig) { _, _ = s.discoverAndPersist(ctx, providerCfg) }() } - -func descriptorsFromIDs(modelIDs []string) []provider.ModelDescriptor { - if len(modelIDs) == 0 { - return nil - } - - descriptors := make([]provider.ModelDescriptor, 0, len(modelIDs)) - for _, modelID := range modelIDs { - id := strings.TrimSpace(modelID) - if id == "" { - continue - } - descriptors = append(descriptors, provider.ModelDescriptor{ - ID: id, - Name: id, - }) - } - if len(descriptors) == 0 { - return nil - } - return descriptors -} diff --git a/internal/provider/catalog/service_test.go b/internal/provider/catalog/service_test.go index 3e330e3d..7aaa5c34 100644 --- a/internal/provider/catalog/service_test.go +++ b/internal/provider/catalog/service_test.go @@ -42,12 +42,12 @@ func TestListProviderModelsSnapshotReturnsDefaultAndRefreshesInBackgroundOnMiss( t.Setenv(testAPIKeyEnv, "test-key") refreshed := make(chan struct{}, 1) - registry := newRegistry(t, config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]provider.ModelDescriptor, error) { + registry := newRegistry(t, config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) { select { case refreshed <- struct{}{}: default: } - return []provider.ModelDescriptor{{ID: "gpt-4o", Name: "GPT-4o"}}, nil + return []config.ModelDescriptor{{ID: "gpt-4o", Name: "GPT-4o"}}, nil }) store := newMemoryStore() @@ -92,8 +92,8 @@ func TestListProviderModelsSnapshotReturnsDefaultAndRefreshesInBackgroundOnMiss( func TestListProviderModelsDiscoversAndCachesOnMiss(t *testing.T) { t.Setenv(testAPIKeyEnv, "test-key") - registry := newRegistry(t, config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]provider.ModelDescriptor, error) { - return []provider.ModelDescriptor{{ + registry := newRegistry(t, config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) { + return []config.ModelDescriptor{{ ID: "server-model", Name: "Server Model", ContextWindow: 32000, @@ -140,21 +140,20 @@ func TestListProviderModelsReturnsStaleCacheAndRefreshesInBackground(t *testing. Identity: identity, FetchedAt: now.Add(-48 * time.Hour), ExpiresAt: now.Add(-24 * time.Hour), - Models: []provider.ModelDescriptor{{ - ID: "stale-model", - Name: "Stale Model", - }}, + Models: []config.ModelDescriptor{ + {ID: "stale-model", Name: "Stale Model"}, + }, }); err != nil { t.Fatalf("seed stale catalog: %v", err) } refreshed := make(chan struct{}, 1) - registry := newRegistry(t, config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]provider.ModelDescriptor, error) { + registry := newRegistry(t, config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) { select { case refreshed <- struct{}{}: default: } - return []provider.ModelDescriptor{{ID: "fresh-model", Name: "Fresh Model"}}, nil + return []config.ModelDescriptor{{ID: "fresh-model", Name: "Fresh Model"}}, nil }) service := NewService("", registry, store) @@ -194,7 +193,7 @@ func TestListProviderModelsReturnsStaleCacheAndRefreshesInBackground(t *testing. func TestDescriptorsFromIDsHelper(t *testing.T) { t.Parallel() - models := descriptorsFromIDs([]string{"gpt-4.1", "", "gpt-4o"}) + models := config.DescriptorsFromIDs([]string{"gpt-4.1", "", "gpt-4o"}) if len(models) != 2 { t.Fatalf("expected 2 descriptors, got %d", len(models)) } @@ -219,7 +218,7 @@ func newRegistry(t *testing.T, name string, discover provider.DiscoveryFunc) *pr return registry } -func containsModelDescriptorID(models []provider.ModelDescriptor, modelID string) bool { +func containsModelDescriptorID(models []config.ModelDescriptor, modelID string) bool { target := config.NormalizeKey(modelID) if target == "" { return false diff --git a/internal/provider/catalog/store.go b/internal/provider/catalog/store.go index e46dc5f0..522dce7c 100644 --- a/internal/provider/catalog/store.go +++ b/internal/provider/catalog/store.go @@ -10,10 +10,10 @@ import ( "os" "path/filepath" "strings" + "sync" "time" "neo-code/internal/config" - "neo-code/internal/provider" ) const SchemaVersion = 1 @@ -22,11 +22,11 @@ var ErrCatalogNotFound = errors.New("provider: model catalog not found") // ModelCatalog stores discovered models for a concrete provider 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 []provider.ModelDescriptor `json:"models"` + SchemaVersion int `json:"schema_version"` + Identity config.ProviderIdentity `json:"identity"` + FetchedAt time.Time `json:"fetched_at"` + ExpiresAt time.Time `json:"expires_at"` + Models []config.ModelDescriptor `json:"models"` } func (c ModelCatalog) Expired(now time.Time) bool { @@ -41,6 +41,7 @@ type Store interface { type JSONStore struct { dir string + mu sync.RWMutex } func NewJSONStore(baseDir string) *JSONStore { @@ -60,7 +61,9 @@ func (s *JSONStore) Load(ctx context.Context, identity config.ProviderIdentity) } path := s.catalogPath(normalized) + s.mu.RLock() data, err := os.ReadFile(path) + s.mu.RUnlock() if err != nil { if os.IsNotExist(err) { return ModelCatalog{}, ErrCatalogNotFound @@ -89,10 +92,6 @@ func (s *JSONStore) Save(ctx context.Context, catalog ModelCatalog) error { } normalized.Identity = identity - 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) @@ -101,7 +100,10 @@ func (s *JSONStore) Save(ctx context.Context, catalog ModelCatalog) error { data = append(data, '\n') } - if err := os.WriteFile(s.catalogPath(identity), data, 0o644); err != nil { + s.mu.Lock() + defer s.mu.Unlock() + + if err := s.writeCatalogFile(s.catalogPath(identity), data); err != nil { return fmt.Errorf("provider: write model catalog: %w", err) } return nil @@ -111,7 +113,7 @@ func normalizeCatalog(modelCatalog ModelCatalog) ModelCatalog { if modelCatalog.SchemaVersion == 0 { modelCatalog.SchemaVersion = SchemaVersion } - modelCatalog.Models = provider.MergeModelDescriptors(modelCatalog.Models) + modelCatalog.Models = config.MergeModelDescriptors(modelCatalog.Models) return modelCatalog } @@ -119,3 +121,45 @@ func (s *JSONStore) catalogPath(identity config.ProviderIdentity) string { sum := sha256.Sum256([]byte(identity.Key())) return filepath.Join(s.dir, hex.EncodeToString(sum[:])+".json") } + +func (s *JSONStore) writeCatalogFile(path string, data []byte) error { + if err := os.MkdirAll(s.dir, 0o755); err != nil { + return fmt.Errorf("provider: create model catalog dir: %w", err) + } + + tempFile, err := os.CreateTemp(s.dir, filepath.Base(path)+".*.tmp") + if err != nil { + return fmt.Errorf("provider: create temp model catalog: %w", err) + } + tempPath := tempFile.Name() + cleanupTemp := true + defer func() { + if cleanupTemp { + _ = os.Remove(tempPath) + } + }() + + if _, err := tempFile.Write(data); err != nil { + _ = tempFile.Close() + return fmt.Errorf("provider: write temp model catalog: %w", err) + } + if err := tempFile.Sync(); err != nil { + _ = tempFile.Close() + return fmt.Errorf("provider: sync temp model catalog: %w", err) + } + if err := tempFile.Close(); err != nil { + return fmt.Errorf("provider: close temp model catalog: %w", err) + } + + if err := os.Rename(tempPath, path); err == nil { + cleanupTemp = false + return nil + } + + // Windows 上覆盖已存在文件时 rename 可能失败，退回到同步保护下的直接写入。 + if err := os.WriteFile(path, data, 0o644); err != nil { + return fmt.Errorf("provider: replace model catalog: %w", err) + } + + return nil +} diff --git a/internal/provider/catalog/store_test.go b/internal/provider/catalog/store_test.go index 6b073beb..ed20e435 100644 --- a/internal/provider/catalog/store_test.go +++ b/internal/provider/catalog/store_test.go @@ -3,11 +3,12 @@ package catalog import ( "context" "errors" + "os" + "path/filepath" "testing" "time" "neo-code/internal/config" - "neo-code/internal/provider" ) func TestJSONStoreRoundTrip(t *testing.T) { @@ -24,7 +25,7 @@ func TestJSONStoreRoundTrip(t *testing.T) { 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: []provider.ModelDescriptor{ + Models: []config.ModelDescriptor{ { ID: "gpt-4.1", Name: "GPT-4.1", @@ -80,3 +81,64 @@ func TestJSONStoreMissingCatalog(t *testing.T) { t.Fatalf("expected ErrCatalogNotFound, got %v", err) } } + +func TestJSONStoreSaveReplacesExistingCatalogWithoutTempLeak(t *testing.T) { + t.Parallel() + + baseDir := t.TempDir() + store := NewJSONStore(baseDir) + identity, err := config.NewProviderIdentity("openai", "https://api.openai.com/v1") + if err != nil { + t.Fatalf("NewProviderIdentity() error = %v", err) + } + + first := ModelCatalog{ + SchemaVersion: SchemaVersion, + 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: []config.ModelDescriptor{ + {ID: "gpt-old", Name: "GPT Old"}, + }, + } + second := ModelCatalog{ + SchemaVersion: SchemaVersion, + Identity: identity, + FetchedAt: time.Date(2026, 4, 4, 10, 0, 0, 0, time.UTC), + ExpiresAt: time.Date(2026, 4, 5, 10, 0, 0, 0, time.UTC), + Models: []config.ModelDescriptor{ + {ID: "gpt-new", Name: "GPT New"}, + }, + } + + if err := store.Save(context.Background(), first); err != nil { + t.Fatalf("first Save() error = %v", err) + } + if err := store.Save(context.Background(), second); err != nil { + t.Fatalf("second Save() error = %v", err) + } + + got, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if len(got.Models) != 1 || got.Models[0].ID != "gpt-new" { + t.Fatalf("expected replaced catalog contents, got %+v", got.Models) + } + + matches, err := filepath.Glob(filepath.Join(baseDir, "cache", "models", "*.tmp")) + if err != nil { + t.Fatalf("Glob() error = %v", err) + } + if len(matches) != 0 { + t.Fatalf("expected no temp files to remain, got %+v", matches) + } + + data, err := os.ReadFile(store.catalogPath(identity)) + if err != nil { + t.Fatalf("ReadFile() error = %v", err) + } + if len(data) == 0 || data[len(data)-1] != '\n' { + t.Fatalf("expected persisted catalog to end with newline, got %q", string(data)) + } +} diff --git a/internal/provider/errors.go b/internal/provider/errors.go index 1163c64c..49936b9e 100644 --- a/internal/provider/errors.go +++ b/internal/provider/errors.go @@ -8,8 +8,6 @@ import ( // 通用领域错误。 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") ) diff --git a/internal/provider/metadata.go b/internal/provider/metadata.go deleted file mode 100644 index 319a87f9..00000000 --- a/internal/provider/metadata.go +++ /dev/null @@ -1,210 +0,0 @@ -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 deleted file mode 100644 index 0223d3b6..00000000 --- a/internal/provider/metadata_test.go +++ /dev/null @@ -1,243 +0,0 @@ -package provider - -import "testing" - -func TestModelDescriptorsFromIDs(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - modelIDs []string - want []ModelDescriptor - }{ - { - name: "empty slice", - modelIDs: []string{}, - want: nil, - }, - { - name: "nil slice", - modelIDs: nil, - want: nil, - }, - { - name: "single model ID", - modelIDs: []string{"gpt-4"}, - want: []ModelDescriptor{ - {ID: "gpt-4", Name: "gpt-4"}, - }, - }, - { - name: "multiple model IDs", - modelIDs: []string{"gpt-4", "gpt-3.5-turbo"}, - want: []ModelDescriptor{ - {ID: "gpt-4", Name: "gpt-4"}, - {ID: "gpt-3.5-turbo", Name: "gpt-3.5-turbo"}, - }, - }, - { - name: "IDs with whitespace", - modelIDs: []string{" gpt-4 ", "\tgpt-3.5-turbo\t"}, - want: []ModelDescriptor{ - {ID: "gpt-4", Name: "gpt-4"}, - {ID: "gpt-3.5-turbo", Name: "gpt-3.5-turbo"}, - }, - }, - { - name: "empty strings are skipped", - modelIDs: []string{"gpt-4", "", " ", "gpt-3.5-turbo"}, - want: []ModelDescriptor{ - {ID: "gpt-4", Name: "gpt-4"}, - {ID: "gpt-3.5-turbo", Name: "gpt-3.5-turbo"}, - }, - }, - { - name: "all empty strings", - modelIDs: []string{"", " ", "\t"}, - want: nil, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := modelDescriptorsFromIDs(tt.modelIDs) - if len(got) != len(tt.want) { - t.Fatalf("expected %d descriptors, got %d", len(tt.want), len(got)) - } - for i := range got { - if got[i].ID != tt.want[i].ID || got[i].Name != tt.want[i].Name { - t.Fatalf("descriptor %d: expected %+v, got %+v", i, tt.want[i], got[i]) - } - } - }) - } -} - -func TestFirstPositiveInt(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - values []any - want int - }{ - { - name: "empty values", - values: []any{}, - want: 0, - }, - { - name: "int positive", - values: []any{int(100)}, - want: 100, - }, - { - name: "int negative returns zero", - values: []any{int(-100)}, - want: 0, - }, - { - name: "int zero returns zero", - values: []any{int(0)}, - want: 0, - }, - { - name: "int32 positive", - values: []any{int32(200)}, - want: 200, - }, - { - name: "int32 negative returns zero", - values: []any{int32(-200)}, - want: 0, - }, - { - name: "int64 positive", - values: []any{int64(300)}, - want: 300, - }, - { - name: "int64 negative returns zero", - values: []any{int64(-300)}, - want: 0, - }, - { - name: "float32 positive", - values: []any{float32(400.5)}, - want: 400, - }, - { - name: "float32 negative returns zero", - values: []any{float32(-400.5)}, - want: 0, - }, - { - name: "float64 positive", - values: []any{float64(500.7)}, - want: 500, - }, - { - name: "float64 negative returns zero", - values: []any{float64(-500.7)}, - want: 0, - }, - { - name: "first positive wins", - values: []any{0, -10, 100, 200}, - want: 100, - }, - { - name: "all negative or zero", - values: []any{-1, -2, 0, -3}, - want: 0, - }, - { - name: "mixed types", - values: []any{int(0), int32(-10), int64(1000), float64(2000)}, - want: 1000, - }, - { - name: "non-numeric type ignored", - values: []any{"100", 200}, - want: 200, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := firstPositiveInt(tt.values...) - if got != tt.want { - t.Fatalf("expected %d, got %d", tt.want, got) - } - }) - } -} - -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/openai/openai.go b/internal/provider/openai/openai.go index e19489b0..b1b3599c 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -15,7 +15,7 @@ import ( "time" "neo-code/internal/config" - domain "neo-code/internal/provider" + "neo-code/internal/provider" modeldiscovery "neo-code/internal/provider/discovery" "neo-code/internal/provider/transport" ) @@ -46,13 +46,13 @@ func defaultRetryTransport() http.RoundTripper { } // Driver 返回 OpenAI 协议驱动的定义。 -func Driver() domain.DriverDefinition { - return domain.DriverDefinition{ +func Driver() provider.DriverDefinition { + return provider.DriverDefinition{ Name: DriverName, - Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (domain.Provider, error) { + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { return New(cfg, WithTransport(defaultRetryTransport())) }, - Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]domain.ModelDescriptor, error) { + Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) { provider, err := New(cfg, WithTransport(defaultRetryTransport())) if err != nil { return nil, err @@ -86,38 +86,38 @@ func New(cfg config.ResolvedProviderConfig, opts ...BuildOption) (*Provider, err }, nil } -func (p *Provider) DiscoverModels(ctx context.Context) ([]domain.ModelDescriptor, error) { +func (p *Provider) DiscoverModels(ctx context.Context) ([]config.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)) + descriptors := make([]config.ModelDescriptor, 0, len(rawModels)) for _, raw := range rawModels { - descriptor, ok := domain.DescriptorFromRawModel(raw) + descriptor, ok := config.DescriptorFromRawModel(raw) if !ok { continue } descriptors = append(descriptors, descriptor) } - return domain.MergeModelDescriptors(descriptors), nil + return config.MergeModelDescriptors(descriptors), nil } -func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan<- domain.StreamEvent) (domain.ChatResponse, error) { +func (p *Provider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { payload, err := p.buildRequest(req) if err != nil { - return domain.ChatResponse{}, err + return provider.ChatResponse{}, err } body, err := json.Marshal(payload) if err != nil { - return domain.ChatResponse{}, fmt.Errorf("openai provider: marshal request: %w", err) + return provider.ChatResponse{}, fmt.Errorf("openai provider: marshal request: %w", err) } endpoint := strings.TrimRight(p.cfg.BaseURL, "/") + "/chat/completions" httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) if err != nil { - return domain.ChatResponse{}, fmt.Errorf("openai provider: build request: %w", err) + return provider.ChatResponse{}, fmt.Errorf("openai provider: build request: %w", err) } httpReq.Header.Set("Authorization", "Bearer "+p.cfg.APIKey) httpReq.Header.Set("Content-Type", "application/json") @@ -125,7 +125,7 @@ func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan resp, err := p.client.Do(httpReq) if err != nil { - return domain.ChatResponse{}, fmt.Errorf("openai provider: send request: %w", err) + return provider.ChatResponse{}, fmt.Errorf("openai provider: send request: %w", err) } defer func(Body io.ReadCloser) { err := Body.Close() @@ -135,13 +135,13 @@ func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan }(resp.Body) if resp.StatusCode >= http.StatusBadRequest { - return domain.ChatResponse{}, p.parseError(resp) + return provider.ChatResponse{}, p.parseError(resp) } return p.consumeStream(ctx, resp.Body, events) } -func (p *Provider) buildRequest(req domain.ChatRequest) (chatCompletionRequest, error) { +func (p *Provider) buildRequest(req provider.ChatRequest) (chatCompletionRequest, error) { model := strings.TrimSpace(req.Model) if model == "" { model = strings.TrimSpace(p.cfg.Model) @@ -158,7 +158,7 @@ func (p *Provider) buildRequest(req domain.ChatRequest) (chatCompletionRequest, if strings.TrimSpace(req.SystemPrompt) != "" { payload.Messages = append(payload.Messages, openAIMessage{ - Role: domain.RoleSystem, + Role: provider.RoleSystem, Content: req.SystemPrompt, }) } @@ -185,27 +185,22 @@ func (p *Provider) buildRequest(req domain.ChatRequest) (chatCompletionRequest, return payload, nil } -func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events chan<- domain.StreamEvent) (domain.ChatResponse, error) { +func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { reader := bufio.NewReader(body) var ( contentBuilder strings.Builder finishReason string - usage domain.Usage + usage provider.Usage done bool ) - toolCalls := make(map[int]*domain.ToolCall) + toolCalls := make(map[int]*provider.ToolCall) dataLines := make([]string, 0, 4) - flushEvent := func() error { - if len(dataLines) == 0 { - return nil - } - - payload := strings.Join(dataLines, "\n") - dataLines = dataLines[:0] - + // processChunk 解析单个 SSE data payload，更新累积状态。 + // 返回错误表示应中止流；done 标志通过闭包变量传递。 + processChunk := func(payload string) error { if strings.TrimSpace(payload) == "[DONE]" { done = true return nil @@ -220,68 +215,46 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha return errors.New(chunk.Error.Message) } - if chunk.Usage != nil { - usage = domain.Usage{ - InputTokens: chunk.Usage.PromptTokens, - OutputTokens: chunk.Usage.CompletionTokens, - TotalTokens: chunk.Usage.TotalTokens, - } - } + extractStreamUsage(&usage, chunk.Usage) for _, choice := range chunk.Choices { - if strings.TrimSpace(choice.FinishReason) != "" { + if choice.FinishReason != "" { finishReason = choice.FinishReason } - - if text := choice.Delta.Content; text != "" { - contentBuilder.WriteString(text) - if err := emitTextDelta(ctx, events, text); err != nil { + if choice.Delta.Content != "" { + contentBuilder.WriteString(choice.Delta.Content) + if err := emitTextDelta(ctx, events, choice.Delta.Content); err != nil { return err } } - - // 处理 tool call 增量 for _, delta := range choice.Delta.ToolCalls { - call, exists := toolCalls[delta.Index] - if !exists { - call = &domain.ToolCall{} - toolCalls[delta.Index] = call - } - - // 首次发现且有名称，发送 tool_call_start - if !exists && strings.TrimSpace(delta.Function.Name) != "" { - call.ID = delta.ID - call.Name = delta.Function.Name - if err := emitToolCallStart(ctx, events, delta.Index, delta.ID, delta.Function.Name); err != nil { - return err - } - } else { - // 更新现有 tool call - if strings.TrimSpace(delta.ID) != "" { - call.ID = delta.ID - } - if strings.TrimSpace(delta.Function.Name) != "" { - call.Name = delta.Function.Name - } - } - - // 发送 tool_call_delta - if delta.Function.Arguments != "" { - call.Arguments += delta.Function.Arguments - if err := emitToolCallDelta(ctx, events, delta.Index, delta.Function.Arguments); err != nil { - return err - } + if err := mergeToolCallDelta(ctx, events, toolCalls, delta); err != nil { + return err } } } - return nil } + // finishStream 统一的流结束处理：发送 message_done 事件并组装最终响应。 + finishStream := func() (provider.ChatResponse, error) { + if err := emitMessageDone(ctx, events, finishReason, &usage); err != nil { + return provider.ChatResponse{}, err + } + return finalizeResponse(contentBuilder.String(), toolCalls, finishReason, usage), nil + } + + flushPendingData := func() error { + defer func() { + dataLines = dataLines[:0] + }() + return flushDataLines(dataLines, processChunk) + } + for { line, err := reader.ReadString('\n') if err != nil && !errors.Is(err, io.EOF) { - return domain.ChatResponse{}, fmt.Errorf("openai provider: read stream: %w", err) + return provider.ChatResponse{}, fmt.Errorf("openai provider: read stream: %w", err) } line = strings.TrimRight(line, "\r\n") @@ -291,29 +264,21 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha case strings.HasPrefix(trimmed, "data:"): dataLines = append(dataLines, strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))) case trimmed == "": - if flushErr := flushEvent(); flushErr != nil { - return domain.ChatResponse{}, flushErr + if flushErr := flushPendingData(); flushErr != nil { + return provider.ChatResponse{}, flushErr } if done { - // 发送 message_done 事件 - if err := emitMessageDone(ctx, events, finishReason, &usage); err != nil { - return domain.ChatResponse{}, err - } - return finalizeResponse(contentBuilder.String(), toolCalls, finishReason, usage), nil + return finishStream() } case strings.HasPrefix(trimmed, ":"): - // SSE comment/heartbeat; ignore it. + // SSE comment/heartbeat; ignore. } if errors.Is(err, io.EOF) { - if flushErr := flushEvent(); flushErr != nil { - return domain.ChatResponse{}, flushErr - } - // 发送 message_done 事件 - if err := emitMessageDone(ctx, events, finishReason, &usage); err != nil { - return domain.ChatResponse{}, err + if flushErr := flushPendingData(); flushErr != nil { + return provider.ChatResponse{}, flushErr } - return finalizeResponse(contentBuilder.String(), toolCalls, finishReason, usage), nil + return finishStream() } } } @@ -321,24 +286,24 @@ func (p *Provider) consumeStream(ctx context.Context, body io.Reader, events cha func (p *Provider) parseError(resp *http.Response) error { data, readErr := io.ReadAll(resp.Body) if readErr != nil { - return domain.NewProviderErrorFromStatus(resp.StatusCode, + return provider.NewProviderErrorFromStatus(resp.StatusCode, fmt.Sprintf("openai provider: read error response: %v", readErr)) } var parsed openAIErrorResponse if err := json.Unmarshal(data, &parsed); err == nil && strings.TrimSpace(parsed.Error.Message) != "" { - return domain.NewProviderErrorFromStatus(resp.StatusCode, parsed.Error.Message) + return provider.NewProviderErrorFromStatus(resp.StatusCode, parsed.Error.Message) } bodyText := strings.TrimSpace(string(data)) if bodyText == "" { - return domain.NewProviderErrorFromStatus(resp.StatusCode, resp.Status) + return provider.NewProviderErrorFromStatus(resp.StatusCode, resp.Status) } - return domain.NewProviderErrorFromStatus(resp.StatusCode, bodyText) + return provider.NewProviderErrorFromStatus(resp.StatusCode, bodyText) } -func toOpenAIMessage(message domain.Message) openAIMessage { +func toOpenAIMessage(message provider.Message) openAIMessage { out := openAIMessage{ Role: message.Role, Content: message.Content, @@ -362,85 +327,109 @@ func toOpenAIMessage(message domain.Message) openAIMessage { return out } -func emitTextDelta(ctx context.Context, events chan<- domain.StreamEvent, text string) error { - if events == nil || text == "" { +func emitTextDelta(ctx context.Context, events chan<- provider.StreamEvent, text string) error { + if text == "" { return nil } - - select { - case events <- domain.StreamEvent{ - Type: domain.StreamEventTextDelta, + return emitStreamEvent(ctx, events, provider.StreamEvent{ + Type: provider.StreamEventTextDelta, Text: text, - }: - return nil - case <-ctx.Done(): - return ctx.Err() - } + }) } -func emitToolCallStart(ctx context.Context, events chan<- domain.StreamEvent, index int, id, name string) error { - if events == nil || name == "" { +func emitToolCallStart(ctx context.Context, events chan<- provider.StreamEvent, index int, id, name string) error { + if name == "" { return nil } - - select { - case events <- domain.StreamEvent{ - Type: domain.StreamEventToolCallStart, + return emitStreamEvent(ctx, events, provider.StreamEvent{ + Type: provider.StreamEventToolCallStart, ToolCallID: id, ToolName: name, ToolCallIndex: index, - }: - return nil - case <-ctx.Done(): - return ctx.Err() - } + }) } // emitToolCallDelta 发送工具调用参数增量事件。 -func emitToolCallDelta(ctx context.Context, events chan<- domain.StreamEvent, index int, argumentsDelta string) error { - if events == nil || argumentsDelta == "" { +func emitToolCallDelta(ctx context.Context, events chan<- provider.StreamEvent, index int, argumentsDelta string) error { + if argumentsDelta == "" { return nil } - - select { - case events <- domain.StreamEvent{ - Type: domain.StreamEventToolCallDelta, + return emitStreamEvent(ctx, events, provider.StreamEvent{ + Type: provider.StreamEventToolCallDelta, ToolCallIndex: index, ToolArgumentsDelta: argumentsDelta, - }: - return nil - case <-ctx.Done(): - return ctx.Err() - } + }) } // emitMessageDone 发送消息完成事件。 -func emitMessageDone(ctx context.Context, events chan<- domain.StreamEvent, finishReason string, usage *domain.Usage) error { +func emitMessageDone(ctx context.Context, events chan<- provider.StreamEvent, finishReason string, usage *provider.Usage) error { if events == nil { return nil } - - select { - case events <- domain.StreamEvent{ - Type: domain.StreamEventMessageDone, + return emitStreamEvent(ctx, events, provider.StreamEvent{ + Type: provider.StreamEventMessageDone, FinishReason: finishReason, Usage: usage, - }: - return nil - case <-ctx.Done(): - return ctx.Err() + }) +} + +// extractStreamUsage 从 OpenAI usage 响应提取并覆盖累积的 token 统计。 +func extractStreamUsage(usage *provider.Usage, raw *openAIUsage) { + if raw == nil { + return + } + *usage = provider.Usage{ + InputTokens: raw.PromptTokens, + OutputTokens: raw.CompletionTokens, + TotalTokens: raw.TotalTokens, } } -func finalizeResponse(content string, toolCalls map[int]*domain.ToolCall, finishReason string, usage domain.Usage) domain.ChatResponse { +// mergeToolCallDelta 将单个 tool call delta 累积到 toolCalls map 中。 +// 首次发现带名称的 delta 时发送 tool_call_start 事件； +// 每次收到 arguments 增量时发送 tool_call_delta 事件。 +func mergeToolCallDelta(ctx context.Context, events chan<- provider.StreamEvent, toolCalls map[int]*provider.ToolCall, delta toolCallDelta) error { + call, exists := toolCalls[delta.Index] + if !exists { + call = &provider.ToolCall{} + toolCalls[delta.Index] = call + } + + hadName := strings.TrimSpace(call.Name) != "" + + if id := strings.TrimSpace(delta.ID); id != "" { + call.ID = id + } + if name := strings.TrimSpace(delta.Function.Name); name != "" { + call.Name = name + } + + if !hadName && strings.TrimSpace(call.Name) != "" { + if err := emitToolCallStart(ctx, events, delta.Index, call.ID, call.Name); err != nil { + return err + } + } + + // 发送参数增量事件（同一 chunk 可能同时携带 name 和 arguments） + if args := delta.Function.Arguments; args != "" { + call.Arguments += args + if err := emitToolCallDelta(ctx, events, delta.Index, args); err != nil { + return err + } + } + return nil +} + +// finalizeResponse 将累积的内容、tool calls 和元数据组装为最终 ChatResponse。 +func finalizeResponse(content string, toolCalls map[int]*provider.ToolCall, finishReason string, usage provider.Usage) provider.ChatResponse { ordered := make([]int, 0, len(toolCalls)) for index := range toolCalls { ordered = append(ordered, index) } sort.Ints(ordered) - message := domain.Message{ - Role: domain.RoleAssistant, + message := provider.Message{ + Role: provider.RoleAssistant, Content: content, } @@ -456,13 +445,34 @@ func finalizeResponse(content string, toolCalls map[int]*domain.ToolCall, finish finishReason = "tool_calls" } - return domain.ChatResponse{ + return provider.ChatResponse{ Message: message, FinishReason: finishReason, Usage: usage, } } +func emitStreamEvent(ctx context.Context, events chan<- provider.StreamEvent, event provider.StreamEvent) error { + if events == nil { + return nil + } + + select { + case events <- event: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// flushDataLines 将缓冲的 data lines 合并为单个 payload 并通过 processChunk 处理。 +func flushDataLines(dataLines []string, processChunk func(string) error) error { + if len(dataLines) == 0 { + return nil + } + return processChunk(strings.Join(dataLines, "\n")) +} + type chatCompletionRequest struct { Model string `json:"model"` Messages []openAIMessage `json:"messages"` diff --git a/internal/provider/openai/openai_test.go b/internal/provider/openai/openai_test.go index 4d9d33de..413e5b44 100644 --- a/internal/provider/openai/openai_test.go +++ b/internal/provider/openai/openai_test.go @@ -613,6 +613,60 @@ func TestEmitToolCallStartGuards(t *testing.T) { } } +func TestMergeToolCallDeltaEmitsStartWhenNameArrivesLater(t *testing.T) { + t.Parallel() + + events := make(chan domain.StreamEvent, 4) + toolCalls := make(map[int]*domain.ToolCall) + + if err := mergeToolCallDelta(context.Background(), events, toolCalls, toolCallDelta{ + Index: 0, + ID: "call_late_name", + }); err != nil { + t.Fatalf("mergeToolCallDelta() first delta error = %v", err) + } + + select { + case evt := <-events: + t.Fatalf("expected no event before tool name arrives, got %+v", evt) + default: + } + + if err := mergeToolCallDelta(context.Background(), events, toolCalls, toolCallDelta{ + Index: 0, + Function: openAIFunctionCall{ + Name: "filesystem_edit", + Arguments: `{"path":"main.go"}`, + }, + }); err != nil { + t.Fatalf("mergeToolCallDelta() late-name delta error = %v", err) + } + + start := <-events + if start.Type != domain.StreamEventToolCallStart { + t.Fatalf("expected tool_call_start event, got %+v", start) + } + if start.ToolCallID != "call_late_name" || start.ToolName != "filesystem_edit" { + t.Fatalf("unexpected tool_call_start payload: %+v", start) + } + + delta := <-events + if delta.Type != domain.StreamEventToolCallDelta { + t.Fatalf("expected tool_call_delta event, got %+v", delta) + } + if delta.ToolArgumentsDelta != `{"path":"main.go"}` { + t.Fatalf("unexpected tool arguments delta: %+v", delta) + } + + call := toolCalls[0] + if call == nil { + t.Fatal("expected tool call to be accumulated") + } + if call.ID != "call_late_name" || call.Name != "filesystem_edit" || call.Arguments != `{"path":"main.go"}` { + t.Fatalf("unexpected accumulated tool call: %+v", call) + } +} + func TestProviderChatEmitsToolCallStartEvent(t *testing.T) { t.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key") diff --git a/internal/provider/registry.go b/internal/provider/registry.go index f5e17324..8a2d4e36 100644 --- a/internal/provider/registry.go +++ b/internal/provider/registry.go @@ -10,7 +10,7 @@ import ( ) type Builder func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) -type DiscoveryFunc func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) +type DiscoveryFunc func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) type DriverDefinition struct { Name string @@ -56,7 +56,7 @@ 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) { +func (r *Registry) DiscoverModels(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) { driver, err := r.driver(cfg.Driver) if err != nil { return nil, err diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index 3aefddf7..f650253d 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -87,7 +87,7 @@ func TestRegistryDiscoverModels(t *testing.T) { t.Run("driver with discovery function", func(t *testing.T) { t.Parallel() - expectedModels := []provider.ModelDescriptor{ + expectedModels := []config.ModelDescriptor{ {ID: "model-1", Name: "Model 1"}, {ID: "model-2", Name: "Model 2"}, } @@ -98,7 +98,7 @@ func TestRegistryDiscoverModels(t *testing.T) { Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { return stubProvider{}, nil }, - Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]provider.ModelDescriptor, error) { + Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]config.ModelDescriptor, error) { return expectedModels, nil }, } diff --git a/internal/provider/selection/service.go b/internal/provider/selection/service.go deleted file mode 100644 index a3f7af53..00000000 --- a/internal/provider/selection/service.go +++ /dev/null @@ -1,273 +0,0 @@ -package selection - -import ( - "context" - "errors" - "strings" - - "neo-code/internal/config" - "neo-code/internal/provider" - "neo-code/internal/provider/catalog" -) - -type modelCatalog interface { - ListProviderModels(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, error) - ListProviderModelsSnapshot(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, error) - ListProviderModelsCached(ctx context.Context, providerCfg config.ProviderConfig) ([]provider.ModelDescriptor, error) -} - -type Service struct { - manager *config.Manager - registry *provider.Registry - catalogs modelCatalog -} - -func NewService(manager *config.Manager, registry *provider.Registry, catalogs *catalog.Service) *Service { - return &Service{ - manager: manager, - registry: registry, - catalogs: catalogs, - } -} - -func (s *Service) ListProviders(ctx context.Context) ([]provider.ProviderCatalogItem, error) { - if err := s.validate(); err != nil { - return nil, err - } - if err := ctx.Err(); err != nil { - return nil, err - } - - cfg := s.manager.Get() - items := make([]provider.ProviderCatalogItem, 0, len(cfg.Providers)) - for _, providerCfg := range cfg.Providers { - if err := ctx.Err(); err != nil { - return nil, err - } - if !s.registry.Supports(providerCfg.Driver) { - continue - } - - models, err := s.catalogs.ListProviderModelsCached(ctx, providerCfg) - if err != nil { - return nil, err - } - items = append(items, providerCatalogItem(providerCfg, models)) - } - - return items, nil -} - -func (s *Service) SelectProvider(ctx context.Context, providerName string) (provider.ProviderSelection, error) { - if err := s.validate(); err != nil { - return provider.ProviderSelection{}, err - } - - cfgSnapshot := s.manager.Get() - providerCfg, err := cfgSnapshot.ProviderByName(providerName) - if err != nil { - return provider.ProviderSelection{}, provider.ErrProviderNotFound - } - if !s.registry.Supports(providerCfg.Driver) { - return provider.ProviderSelection{}, provider.ErrDriverNotFound - } - - models, err := s.catalogs.ListProviderModelsCached(ctx, providerCfg) - if err != nil { - return provider.ProviderSelection{}, err - } - - var selection provider.ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - selected, err := cfg.ProviderByName(providerName) - if err != nil { - return provider.ErrProviderNotFound - } - if !s.registry.Supports(selected.Driver) { - return provider.ErrDriverNotFound - } - - cfg.SelectedProvider = selected.Name - nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) - cfg.CurrentModel = nextModel - selection = selectionFromConfig(*cfg) - return nil - }) - if err != nil { - return provider.ProviderSelection{}, err - } - - return selection, nil -} - -func (s *Service) ListModels(ctx context.Context) ([]provider.ModelDescriptor, error) { - if err := s.validate(); err != nil { - return nil, err - } - if err := ctx.Err(); err != nil { - return nil, err - } - - selected, err := s.selectedProviderConfig() - if err != nil { - return nil, err - } - return s.catalogs.ListProviderModels(ctx, selected) -} - -func (s *Service) ListModelsSnapshot(ctx context.Context) ([]provider.ModelDescriptor, error) { - if err := s.validate(); err != nil { - return nil, err - } - if err := ctx.Err(); err != nil { - return nil, err - } - - selected, err := s.selectedProviderConfig() - if err != nil { - return nil, err - } - return s.catalogs.ListProviderModelsSnapshot(ctx, selected) -} - -func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (provider.ProviderSelection, error) { - if err := s.validate(); err != nil { - return provider.ProviderSelection{}, err - } - - modelID = strings.TrimSpace(modelID) - if modelID == "" { - return provider.ProviderSelection{}, provider.ErrModelNotFound - } - - selected, err := s.selectedProviderConfig() - if err != nil { - return provider.ProviderSelection{}, err - } - - models, err := s.catalogs.ListProviderModelsSnapshot(ctx, selected) - if err != nil { - return provider.ProviderSelection{}, err - } - if !containsModelDescriptorID(models, modelID) { - return provider.ProviderSelection{}, provider.ErrModelNotFound - } - - var selection provider.ProviderSelection - 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 provider.ProviderSelection{}, err - } - - return selection, nil -} - -func (s *Service) EnsureSelection(ctx context.Context) (provider.ProviderSelection, error) { - if err := s.validate(); err != nil { - return provider.ProviderSelection{}, err - } - if err := ctx.Err(); err != nil { - return provider.ProviderSelection{}, err - } - - cfgSnapshot := s.manager.Get() - selected, err := cfgSnapshot.SelectedProviderConfig() - if err != nil { - return provider.ProviderSelection{}, err - } - - models, err := s.catalogs.ListProviderModelsSnapshot(ctx, selected) - if err != nil { - return provider.ProviderSelection{}, err - } - nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) - if !changed { - return selectionFromConfig(cfgSnapshot), nil - } - - var selection provider.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 - }) - if err != nil { - return provider.ProviderSelection{}, err - } - - return selection, nil -} - -func (s *Service) validate() error { - if s == nil { - return errors.New("provider selection: service is nil") - } - if s.manager == nil { - return errors.New("provider selection: config manager is nil") - } - if s.registry == nil { - return errors.New("provider selection: registry is nil") - } - if s.catalogs == nil { - return errors.New("provider selection: catalog service is nil") - } - return nil -} - -func (s *Service) selectedProviderConfig() (config.ProviderConfig, error) { - cfg := s.manager.Get() - return cfg.SelectedProviderConfig() -} - -func selectionFromConfig(cfg config.Config) provider.ProviderSelection { - return provider.ProviderSelection{ - ProviderID: cfg.SelectedProvider, - ModelID: cfg.CurrentModel, - } -} - -func resolveCurrentModel(currentModel string, models []provider.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 currentModel, false -} - -func providerCatalogItem(cfg config.ProviderConfig, models []provider.ModelDescriptor) provider.ProviderCatalogItem { - return provider.ProviderCatalogItem{ - ID: strings.TrimSpace(cfg.Name), - Name: strings.TrimSpace(cfg.Name), - Models: provider.MergeModelDescriptors(models), - } -} - -func containsModelDescriptorID(models []provider.ModelDescriptor, modelID string) bool { - target := config.NormalizeKey(modelID) - if target == "" { - return false - } - - for _, model := range models { - if config.NormalizeKey(model.ID) == target { - return true - } - } - return false -} diff --git a/internal/provider/selection/service_test.go b/internal/provider/selection/service_test.go deleted file mode 100644 index 33c68d8a..00000000 --- a/internal/provider/selection/service_test.go +++ /dev/null @@ -1,469 +0,0 @@ -package selection - -import ( - "context" - "errors" - "sync" - "testing" - - "neo-code/internal/config" - "neo-code/internal/provider" - providercatalog "neo-code/internal/provider/catalog" -) - -func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { - t.Parallel() - - manager := newTestManager(t, config.DefaultConfig()) - registry := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - items, err := service.ListProviders(context.Background()) - if err != nil { - t.Fatalf("ListProviders() error = %v", err) - } - expected := map[string]int{ - config.OpenAIName: 1, - config.GeminiName: 1, - config.OpenLLName: 1, - config.QiniuName: 1, - } - if len(items) != len(expected) { - t.Fatalf("expected only builtin providers, got %d", len(items)) - } - - for _, item := range items { - wantModels, ok := expected[item.ID] - if !ok { - t.Fatalf("unexpected builtin provider %q", item.ID) - } - if len(item.Models) != wantModels { - t.Fatalf("expected provider models to fall back to the default model, got %+v", item.Models) - } - delete(expected, item.ID) - } - if len(expected) != 0 { - t.Fatalf("missing builtin providers from catalog: %+v", expected) - } -} - -func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { - manager := newTestManager(t, config.DefaultConfig()) - registry := newRegistry(t, config.OpenAIName, nil) - - 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) - } - - store := newCatalogStoreStub() - identity, err := config.QiniuProvider().Identity() - if err != nil { - t.Fatalf("qiniu provider identity: %v", err) - } - if err := store.Save(context.Background(), providercatalog.ModelCatalog{ - SchemaVersion: providercatalog.SchemaVersion, - Identity: identity, - Models: []provider.ModelDescriptor{ - {ID: "qiniu-alt", Name: "qiniu-alt"}, - }, - }); err != nil { - t.Fatalf("seed catalog: %v", err) - } - - service := NewService(manager, registry, providercatalog.NewService("", registry, store)) - - 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 after switch: %+v", selection) - } - - if _, err := service.SetCurrentModel(context.Background(), "missing-model"); !errors.Is(err, provider.ErrModelNotFound) { - t.Fatalf("expected ErrModelNotFound, got %v", err) - } - - selection, err = service.SetCurrentModel(context.Background(), "qiniu-alt") - if err != nil { - t.Fatalf("SetCurrentModel() error = %v", err) - } - if selection.ModelID != "qiniu-alt" { - t.Fatalf("expected selected model %q, got %+v", "qiniu-alt", selection) - } - - cfg := manager.Get() - selected, err := cfg.SelectedProviderConfig() - if err != nil { - t.Fatalf("SelectedProviderConfig() error = %v", err) - } - if selected.Name != config.QiniuName { - t.Fatalf("expected selected provider %q, got %+v", config.QiniuName, selected) - } - 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) - } -} - -func TestServiceEnsureSelectionRepairsInvalidCurrentModel(t *testing.T) { - t.Parallel() - - manager := newTestManager(t, config.DefaultConfig()) - 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 := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - 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) - } - - 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) - } -} - -func TestServiceModelOperationsUseBuiltinConfigEvenWithoutDriver(t *testing.T) { - t.Parallel() - - defaults := config.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 := newTestManager(t, defaults) - registry := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - models, err := service.ListModels(context.Background()) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if len(models) != 1 || models[0].ID != "broken-model" { - t.Fatalf("expected default model fallback, got %+v", models) - } - - 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() - - defaults := config.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 := newTestManager(t, defaults) - registry := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - if _, err := service.SelectProvider(context.Background(), config.OpenAIName); !errors.Is(err, provider.ErrDriverNotFound) { - t.Fatalf("expected SelectProvider() to preserve ErrDriverNotFound, got %v", err) - } -} - -func TestServiceListModelsSnapshot(t *testing.T) { - t.Parallel() - - manager := newTestManager(t, config.DefaultConfig()) - registry := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - models, err := service.ListModelsSnapshot(context.Background()) - if err != nil { - t.Fatalf("ListModelsSnapshot() error = %v", err) - } - if len(models) != 1 || models[0].ID != config.OpenAIDefaultModel { - t.Fatalf("expected default model fallback, got %+v", models) - } -} - -func TestServiceValidateErrors(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - service *Service - errMsg string - }{ - { - name: "nil service", - service: nil, - errMsg: "provider selection: service is nil", - }, - { - name: "nil manager", - service: &Service{ - manager: nil, - registry: provider.NewRegistry(), - catalogs: providercatalog.NewService("", provider.NewRegistry(), newCatalogStoreStub()), - }, - errMsg: "provider selection: config manager is nil", - }, - { - name: "nil registry", - service: &Service{ - manager: config.NewManager(config.NewLoader(t.TempDir(), config.DefaultConfig())), - registry: nil, - catalogs: providercatalog.NewService("", provider.NewRegistry(), newCatalogStoreStub()), - }, - errMsg: "provider selection: registry is nil", - }, - { - name: "nil catalogs", - service: &Service{ - manager: config.NewManager(config.NewLoader(t.TempDir(), config.DefaultConfig())), - registry: provider.NewRegistry(), - catalogs: nil, - }, - errMsg: "provider selection: catalog service is nil", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - _, err := tt.service.ListProviders(context.Background()) - if err == nil || err.Error() != tt.errMsg { - t.Fatalf("expected error %q, got %v", tt.errMsg, err) - } - }) - } -} - -func TestServiceOperationsWithCanceledContext(t *testing.T) { - t.Parallel() - - manager := newTestManager(t, config.DefaultConfig()) - registry := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - tests := []struct { - name string - fn func(context.Context) error - }{ - { - name: "ListProviders", - fn: func(ctx context.Context) error { - _, err := service.ListProviders(ctx) - return err - }, - }, - { - name: "ListModels", - fn: func(ctx context.Context) error { - _, err := service.ListModels(ctx) - return err - }, - }, - { - name: "ListModelsSnapshot", - fn: func(ctx context.Context) error { - _, err := service.ListModelsSnapshot(ctx) - return err - }, - }, - { - name: "EnsureSelection", - fn: func(ctx context.Context) error { - _, err := service.EnsureSelection(ctx) - return err - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - err := tt.fn(ctx) - if err == nil { - t.Fatalf("expected error for canceled context") - } - if !errors.Is(err, context.Canceled) { - t.Fatalf("expected context.Canceled, got %v", err) - } - }) - } -} - -func TestServiceSetCurrentModelEmptyModelID(t *testing.T) { - t.Parallel() - - manager := newTestManager(t, config.DefaultConfig()) - registry := newRegistry(t, config.OpenAIName, nil) - service := NewService(manager, registry, providercatalog.NewService("", registry, newCatalogStoreStub())) - - _, err := service.SetCurrentModel(context.Background(), "") - if !errors.Is(err, provider.ErrModelNotFound) { - t.Fatalf("expected ErrModelNotFound for empty model ID, got %v", err) - } - - _, err = service.SetCurrentModel(context.Background(), " ") - if !errors.Is(err, provider.ErrModelNotFound) { - t.Fatalf("expected ErrModelNotFound for whitespace model ID, got %v", err) - } -} - -func TestResolveCurrentModelHelper(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - currentModel string - models []provider.ModelDescriptor - fallback string - expected string - changed bool - }{ - { - name: "current model in list", - currentModel: "gpt-4o", - models: []provider.ModelDescriptor{ - {ID: "gpt-4.1"}, - {ID: "gpt-4o"}, - {ID: "gpt-5.4"}, - }, - fallback: "gpt-4.1", - expected: "gpt-4o", - changed: false, - }, - { - name: "current model falls back to provider default", - currentModel: "unknown-model", - models: []provider.ModelDescriptor{ - {ID: "gpt-4.1"}, - {ID: "gpt-4o"}, - {ID: "gpt-5.4"}, - }, - fallback: "gpt-4.1", - expected: "gpt-4.1", - changed: true, - }, - { - name: "missing fallback keeps current model unchanged", - currentModel: "unknown-model", - models: []provider.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() - - 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 newTestManager(t *testing.T, defaults *config.Config) *config.Manager { - t.Helper() - - manager := config.NewManager(config.NewLoader(t.TempDir(), defaults)) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("load config: %v", err) - } - return manager -} - -func newRegistry(t *testing.T, name string, discover provider.DiscoveryFunc) *provider.Registry { - t.Helper() - - registry := provider.NewRegistry() - if err := registry.Register(provider.DriverDefinition{ - Name: name, - Discover: discover, - Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { - return selectionTestProvider{}, nil - }, - }); err != nil { - t.Fatalf("register driver: %v", err) - } - return registry -} - -type selectionTestProvider struct{} - -func (selectionTestProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { - return provider.ChatResponse{}, nil -} - -type catalogStoreStub struct { - mu sync.Mutex - catalogs map[string]providercatalog.ModelCatalog -} - -func newCatalogStoreStub() *catalogStoreStub { - return &catalogStoreStub{ - catalogs: map[string]providercatalog.ModelCatalog{}, - } -} - -func (s *catalogStoreStub) Load(ctx context.Context, identity config.ProviderIdentity) (providercatalog.ModelCatalog, error) { - if err := ctx.Err(); err != nil { - return providercatalog.ModelCatalog{}, err - } - - s.mu.Lock() - defer s.mu.Unlock() - - catalog, ok := s.catalogs[identity.Key()] - if !ok { - return providercatalog.ModelCatalog{}, providercatalog.ErrCatalogNotFound - } - return catalog, nil -} - -func (s *catalogStoreStub) Save(ctx context.Context, catalog providercatalog.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/transport/retry.go b/internal/provider/transport/retry.go index dca6a23b..d0c68c0c 100644 --- a/internal/provider/transport/retry.go +++ b/internal/provider/transport/retry.go @@ -1,10 +1,12 @@ package transport import ( + "context" "errors" "io" "log" "math/rand/v2" + "net" "net/http" "strings" "time" @@ -53,22 +55,31 @@ func DefaultRetryableFunc(err error) bool { } // isNetworkError 判断是否为网络层错误（不包含 TLS 错误，避免重试证书问题）。 +// 优先通过类型断言检测 *net.OpError 和 context.DeadlineExceeded， +// 再通过消息文本匹配兜底（兼容包装过的错误）。 func isNetworkError(err error) bool { if err == nil { return false } + // 类型断言：Go 标准库网络操作返回的底层错误 + var opErr *net.OpError + if errors.As(err, &opErr) { + return true + } + // context 超时通常由网络操作触发 + if errors.Is(err, context.DeadlineExceeded) { + return true + } + // 消息匹配：兼容被 fmt.Errorf 包装过的网络错误 msg := err.Error() - // 常见网络错误关键字 - networkIndicators := []string{ + for _, indicator := range []string{ "connection refused", "connection reset", "connection timed out", - "timeout", "no such host", - "i/o timeout", "dial tcp", - } - for _, indicator := range networkIndicators { + "timeout", + } { if strings.Contains(msg, indicator) { return true } @@ -126,15 +137,23 @@ func (rt *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) { resp, err := rt.base.RoundTrip(req) if err == nil { - return resp, nil - } + if !provider.IsRetryableStatus(resp.StatusCode) { + return resp, nil + } - lastErr = err - lastResp = nil + lastResp = resp + lastErr = provider.NewProviderErrorFromStatus(resp.StatusCode, resp.Status) + if !rt.config.RetryableFunc(lastErr) || attempt == rt.config.MaxRetries { + return resp, nil + } + } else { + lastErr = err + lastResp = nil - // 检查是否可重试。 - if !rt.config.RetryableFunc(err) { - return nil, err + // 检查是否可重试。 + if !rt.config.RetryableFunc(err) { + return nil, err + } } // 检查 context 是否已取消。 @@ -143,6 +162,9 @@ func (rt *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) { } } + if lastResp != nil { + return lastResp, nil + } return nil, lastErr } diff --git a/internal/provider/transport/retry_test.go b/internal/provider/transport/retry_test.go index ac5cc7d7..a36be54d 100644 --- a/internal/provider/transport/retry_test.go +++ b/internal/provider/transport/retry_test.go @@ -203,6 +203,72 @@ func TestRoundTrip_RetryThenSuccess(t *testing.T) { } } +func TestRoundTrip_RetryableStatusThenSuccess(t *testing.T) { + t.Parallel() + + wantResp := newSuccessResponse(http.StatusOK, "ok") + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { + return newSuccessResponse(http.StatusTooManyRequests, "retry later"), nil + }, + func() (*http.Response, error) { return wantResp, nil }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if m.callCount != 2 { + t.Fatalf("callCount = %d, want 2", m.callCount) + } +} + +func TestRoundTrip_RetryableStatusExhaustedReturnsLastResponse(t *testing.T) { + t.Parallel() + + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { + return newSuccessResponse(http.StatusTooManyRequests, "retry later"), nil + }, + func() (*http.Response, error) { return newSuccessResponse(http.StatusBadGateway, "still broken"), nil }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 1, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp == nil { + t.Fatal("expected final response after retries are exhausted") + } + if resp.StatusCode != http.StatusBadGateway { + t.Fatalf("status = %d, want 502", resp.StatusCode) + } + if m.callCount != 2 { + t.Fatalf("callCount = %d, want 2", m.callCount) + } +} + func TestRoundTrip_AllRetriesFail(t *testing.T) { t.Parallel() diff --git a/internal/provider/types.go b/internal/provider/types.go index 87b34207..eba95f19 100644 --- a/internal/provider/types.go +++ b/internal/provider/types.go @@ -1,5 +1,6 @@ package provider +// Role 常量定义消息角色标识。 const ( RoleSystem = "system" RoleUser = "user" @@ -7,6 +8,7 @@ const ( RoleTool = "tool" ) +// Message 表示对话中的单条消息。 type Message struct { Role string `json:"role"` Content string `json:"content"` @@ -15,18 +17,21 @@ type Message struct { IsError bool `json:"is_error,omitempty"` } +// ToolCall 表示模型发起的工具调用请求。 type ToolCall struct { ID string `json:"id"` Name string `json:"name"` Arguments string `json:"arguments"` } +// ToolSpec 表示暴露给模型的可调用工具描述。 type ToolSpec struct { Name string `json:"name"` Description string `json:"description"` Schema map[string]any `json:"schema"` } +// ChatRequest 是 provider.Chat() 的请求参数。 type ChatRequest struct { Model string `json:"model"` SystemPrompt string `json:"system_prompt"` @@ -34,39 +39,21 @@ type ChatRequest struct { Tools []ToolSpec `json:"tools,omitempty"` } +// ChatResponse 是 provider.Chat() 的返回结果。 type ChatResponse struct { Message Message `json:"message"` FinishReason string `json:"finish_reason"` Usage Usage `json:"usage"` } +// Usage 记录本次请求的 token 使用统计。 type Usage struct { InputTokens int `json:"input_tokens"` OutputTokens int `json:"output_tokens"` TotalTokens int `json:"total_tokens"` } -type ModelDescriptor struct { - 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 { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description,omitempty"` - Models []ModelDescriptor `json:"models,omitempty"` -} - -type ProviderSelection struct { - ProviderID string `json:"provider_id"` - ModelID string `json:"model_id"` -} - +// StreamEventType 定义流式事件类型。 type StreamEventType string const ( diff --git a/internal/tui/commands.go b/internal/tui/commands.go index fe2897fe..d1b2b4ad 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -9,7 +9,6 @@ import ( tea "github.com/charmbracelet/bubbletea" "neo-code/internal/config" - "neo-code/internal/provider" ) const ( @@ -137,7 +136,7 @@ func newSelectionPicker(items []list.Item) list.Model { return picker } -func newProviderPicker(items []provider.ProviderCatalogItem) list.Model { +func newProviderPicker(items []config.ProviderCatalogItem) list.Model { listItems := make([]list.Item, 0, len(items)) for _, item := range items { listItems = append(listItems, providerItem{ @@ -149,7 +148,7 @@ func newProviderPicker(items []provider.ProviderCatalogItem) list.Model { return newSelectionPicker(listItems) } -func newModelPicker(models []provider.ModelDescriptor) list.Model { +func newModelPicker(models []config.ModelDescriptor) list.Model { items := make([]list.Item, 0, len(models)) for _, option := range models { items = append(items, modelItem{ diff --git a/internal/tui/provider_service.go b/internal/tui/provider_service.go index 256bb0ba..8893c139 100644 --- a/internal/tui/provider_service.go +++ b/internal/tui/provider_service.go @@ -3,13 +3,13 @@ package tui import ( "context" - "neo-code/internal/provider" + "neo-code/internal/config" ) type ProviderController interface { - ListProviders(ctx context.Context) ([]provider.ProviderCatalogItem, error) - SelectProvider(ctx context.Context, providerID string) (provider.ProviderSelection, error) - ListModels(ctx context.Context) ([]provider.ModelDescriptor, error) - ListModelsSnapshot(ctx context.Context) ([]provider.ModelDescriptor, error) - SetCurrentModel(ctx context.Context, modelID string) (provider.ProviderSelection, error) + ListProviders(ctx context.Context) ([]config.ProviderCatalogItem, error) + SelectProvider(ctx context.Context, providerID string) (config.ProviderSelection, error) + ListModels(ctx context.Context) ([]config.ModelDescriptor, error) + ListModelsSnapshot(ctx context.Context) ([]config.ModelDescriptor, error) + SetCurrentModel(ctx context.Context, modelID string) (config.ProviderSelection, error) } diff --git a/internal/tui/update.go b/internal/tui/update.go index 472115a0..f344cca2 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -26,7 +26,7 @@ type RuntimeClosedMsg struct{} type runFinishedMsg struct{ err error } type modelCatalogRefreshMsg struct { providerID string - models []provider.ModelDescriptor + models []config.ModelDescriptor err error } diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 5122bdb1..ffa4025a 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -8,6 +8,7 @@ import ( "path/filepath" "regexp" "strings" + "sync" "testing" "time" @@ -19,7 +20,6 @@ import ( contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" providercatalog "neo-code/internal/provider/catalog" - providerselection "neo-code/internal/provider/selection" agentruntime "neo-code/internal/runtime" "neo-code/internal/tools" ) @@ -867,7 +867,7 @@ func TestTUIStandaloneHelpers(t *testing.T) { t.Fatalf("expected delegate update to return nil") } var buf bytes.Buffer - model := newModelPicker([]provider.ModelDescriptor{{ID: "gpt-4.1", Name: "gpt-4.1"}}) + model := newModelPicker([]config.ModelDescriptor{{ID: "gpt-4.1", Name: "gpt-4.1"}}) sessionList := []list.Item{sItem} listModel := list.New(sessionList, delegate, 30, 10) delegate.Render(&buf, listModel, 0, sItem) @@ -2485,40 +2485,41 @@ func newTestConfigManager(t *testing.T) *config.Manager { return manager } -func newTestProviderService(t *testing.T, manager *config.Manager) *providerselection.Service { +func newTestProviderService(t *testing.T, manager *config.Manager) *config.SelectionService { t.Helper() 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 + return tUItestProvider{}, nil }, }) if err != nil { t.Fatalf("register provider drivers: %v", err) } - modelCatalogs := providercatalog.NewService("", registry, newTUITestCatalogStore()) - return providerselection.NewService(manager, registry, modelCatalogs) + modelCatalogs := providercatalog.NewService("", registry, newTUItestCatalogStore()) + return config.NewSelectionService(manager, registry, registry, modelCatalogs) } -type tuiTestProvider struct{} +type tUItestProvider struct{} -func (tuiTestProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { +func (tUItestProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { return provider.ChatResponse{}, nil } -type tuiTestCatalogStore struct { +type tUItestCatalogStore struct { catalogs map[string]providercatalog.ModelCatalog + mu sync.Mutex } -func newTUITestCatalogStore() *tuiTestCatalogStore { - return &tuiTestCatalogStore{ +func newTUItestCatalogStore() *tUItestCatalogStore { + return &tUItestCatalogStore{ catalogs: map[string]providercatalog.ModelCatalog{}, } } -func (s *tuiTestCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (providercatalog.ModelCatalog, error) { +func (s *tUItestCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (providercatalog.ModelCatalog, error) { if err := ctx.Err(); err != nil { return providercatalog.ModelCatalog{}, err } @@ -2530,7 +2531,7 @@ func (s *tuiTestCatalogStore) Load(ctx context.Context, identity config.Provider return catalog, nil } -func (s *tuiTestCatalogStore) Save(ctx context.Context, catalog providercatalog.ModelCatalog) error { +func (s *tUItestCatalogStore) Save(ctx context.Context, catalog providercatalog.ModelCatalog) error { if err := ctx.Err(); err != nil { return err }