From be6bf9f6c6405ecccef8ce7d72205d321cdabb4b Mon Sep 17 00:00:00 2001 From: uiyzzi Date: Tue, 24 Mar 2026 23:56:45 +0800 Subject: [PATCH] Add virtual model support for multi-key expansion Virtual models generated from multi-key expansion are now marked and filtered during config persistence. Virtual models display with a badge in the UI and cannot be set as default. --- pkg/config/config.go | 23 ++++++ pkg/config/config_test.go | 59 +++++++++++++++ pkg/config/multikey_test.go | 72 +++++++++++++++++++ web/backend/api/models.go | 10 ++- web/backend/api/models_test.go | 40 +++++++++++ web/frontend/src/api/models.ts | 1 + .../src/components/models/model-card.tsx | 7 +- web/frontend/src/i18n/locales/en.json | 3 +- web/frontend/src/i18n/locales/zh.json | 3 +- 9 files changed, 214 insertions(+), 4 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index 4e903d140..47bb7e8f1 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -857,6 +857,10 @@ type ModelConfig struct { secModelName string apiKeys []string secDirty bool + + // isVirtual marks this model as a virtual model generated from multi-key expansion. + // Virtual models should not be persisted to config files. + isVirtual bool } // APIKey returns the first API key from apiKeys @@ -867,6 +871,11 @@ func (c *ModelConfig) APIKey() string { return "" } +// IsVirtual returns true if this model was generated from multi-key expansion. +func (c *ModelConfig) IsVirtual() bool { + return c.isVirtual +} + // Validate checks if the ModelConfig has all required fields. func (c *ModelConfig) Validate() error { if c.ModelName == "" { @@ -1800,7 +1809,20 @@ func SaveConfig(path string, cfg *Config) error { return err } + // Filter out virtual models before serializing to config file + nonVirtualModels := make([]*ModelConfig, 0, len(cfg.ModelList)) + for _, m := range cfg.ModelList { + if !m.isVirtual { + nonVirtualModels = append(nonVirtualModels, m) + } + } + // Temporarily replace ModelList with filtered version for serialization + originalModelList := cfg.ModelList + cfg.ModelList = nonVirtualModels + data, err := json.MarshalIndent(cfg, "", " ") + // Restore original ModelList after serialization + cfg.ModelList = originalModelList if err != nil { return err } @@ -2018,6 +2040,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ExtraBody: m.ExtraBody, + isVirtual: true, } expanded = append(expanded, additionalEntry) fallbackNames = append(fallbackNames, expandedName) diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 16a758bef..bedd46f6e 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -391,6 +391,65 @@ func TestSaveConfig_PreservesDisabledTelegramPlaceholder(t *testing.T) { } } +// TestSaveConfig_FiltersVirtualModels verifies that SaveConfig does not write +// virtual models (generated by expandMultiKeyModels) to the config file. +func TestSaveConfig_FiltersVirtualModels(t *testing.T) { + tmpDir := t.TempDir() + path := filepath.Join(tmpDir, "config.json") + + cfg := DefaultConfig() + + // Manually add a virtual model to ModelList (simulating what expandMultiKeyModels does) + primaryModel := &ModelConfig{ + ModelName: "gpt-4", + Model: "openai/gpt-4o", + apiKeys: []string{"key1"}, + } + virtualModel := &ModelConfig{ + ModelName: "gpt-4__key_1", + Model: "openai/gpt-4o", + apiKeys: []string{"key2"}, + isVirtual: true, + } + cfg.ModelList = []*ModelConfig{primaryModel, virtualModel} + + // SaveConfig should filter out virtual models + if err := SaveConfig(path, cfg); err != nil { + t.Fatalf("SaveConfig failed: %v", err) + } + + // Reload and verify + reloaded, err := LoadConfig(path) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + // Should only have the primary model, not the virtual one + if len(reloaded.ModelList) != 1 { + t.Fatalf("expected 1 model after reload, got %d", len(reloaded.ModelList)) + } + + if reloaded.ModelList[0].ModelName != "gpt-4" { + t.Errorf("expected model_name 'gpt-4', got %q", reloaded.ModelList[0].ModelName) + } + + // Verify virtual model was not persisted + for _, m := range reloaded.ModelList { + if m.ModelName == "gpt-4__key_1" { + t.Errorf("virtual model gpt-4__key_1 should not have been saved") + } + } + + // Verify the saved file does not contain the virtual model name + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if strings.Contains(string(data), "gpt-4__key_1") { + t.Errorf("saved config should not contain virtual model name 'gpt-4__key_1'") + } +} + // TestConfig_Complete verifies all config fields are set func TestConfig_Complete(t *testing.T) { cfg := DefaultConfig() diff --git a/pkg/config/multikey_test.go b/pkg/config/multikey_test.go index cc529905c..c17fcc53b 100644 --- a/pkg/config/multikey_test.go +++ b/pkg/config/multikey_test.go @@ -232,6 +232,78 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { } } +func TestExpandMultiKeyModels_IsVirtualFlag(t *testing.T) { + models := []*ModelConfig{ + { + ModelName: "gpt-4", + Model: "openai/gpt-4o", + apiKeys: []string{"key1", "key2", "key3"}, + }, + } + + result := expandMultiKeyModels(models) + + // Should expand to 3 models + if len(result) != 3 { + t.Fatalf("expected 3 models, got %d", len(result)) + } + + // Primary model should NOT be virtual + primary := result[2] + if primary.isVirtual { + t.Errorf("primary model should not be virtual") + } + if primary.ModelName != "gpt-4" { + t.Errorf("expected primary model_name 'gpt-4', got %q", primary.ModelName) + } + + // Virtual models should have isVirtual = true + virtual1 := result[0] + if !virtual1.isVirtual { + t.Errorf("gpt-4__key_1 should be virtual") + } + if virtual1.ModelName != "gpt-4__key_1" { + t.Errorf("expected virtual model_name 'gpt-4__key_1', got %q", virtual1.ModelName) + } + + virtual2 := result[1] + if !virtual2.isVirtual { + t.Errorf("gpt-4__key_2 should be virtual") + } + if virtual2.ModelName != "gpt-4__key_2" { + t.Errorf("expected virtual model_name 'gpt-4__key_2', got %q", virtual2.ModelName) + } + + // IsVirtual() method should work + if !virtual1.IsVirtual() { + t.Errorf("IsVirtual() should return true for virtual model") + } + if primary.IsVirtual() { + t.Errorf("IsVirtual() should return false for primary model") + } +} + +func TestExpandMultiKeyModels_SingleKey_NotVirtual(t *testing.T) { + models := []*ModelConfig{ + { + ModelName: "gpt-4", + Model: "openai/gpt-4o", + apiKeys: []string{"single-key"}, + }, + } + + result := expandMultiKeyModels(models) + + if len(result) != 1 { + t.Fatalf("expected 1 model, got %d", len(result)) + } + + // Single key model should NOT be virtual + if result[0].isVirtual { + t.Errorf("single key model should not be virtual") + } +} + func TestMergeAPIKeys(t *testing.T) { tests := []struct { name string diff --git a/web/backend/api/models.go b/web/backend/api/models.go index 48babd8cd..ce7719906 100644 --- a/web/backend/api/models.go +++ b/web/backend/api/models.go @@ -42,6 +42,7 @@ type modelResponse struct { // Meta Configured bool `json:"configured"` IsDefault bool `json:"is_default"` + IsVirtual bool `json:"is_virtual"` } // handleListModels returns all model_list entries with masked API keys. @@ -86,6 +87,7 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) { ExtraBody: m.ExtraBody, Configured: configured[i], IsDefault: m.ModelName == defaultModel, + IsVirtual: m.IsVirtual(), }) } @@ -288,11 +290,13 @@ func (h *Handler) handleSetDefaultModel(w http.ResponseWriter, r *http.Request) return } - // Verify the model_name exists in model_list + // Verify the model_name exists in model_list and is not a virtual model found := false + isVirtual := false for _, m := range cfg.ModelList { if m.ModelName == req.ModelName { found = true + isVirtual = m.IsVirtual() break } } @@ -300,6 +304,10 @@ func (h *Handler) handleSetDefaultModel(w http.ResponseWriter, r *http.Request) http.Error(w, fmt.Sprintf("Model %q not found in model_list", req.ModelName), http.StatusNotFound) return } + if isVirtual { + http.Error(w, fmt.Sprintf("Cannot set virtual model %q as default", req.ModelName), http.StatusBadRequest) + return + } cfg.Agents.Defaults.ModelName = req.ModelName diff --git a/web/backend/api/models_test.go b/web/backend/api/models_test.go index 9d3e72bd3..c80527fe3 100644 --- a/web/backend/api/models_test.go +++ b/web/backend/api/models_test.go @@ -356,6 +356,46 @@ func TestHandleAddModel_PersistsAPIKey(t *testing.T) { } } +// TestHandleSetDefaultModel_RejectsNonexistentModel tests that setting a non-existent +// model as default returns 404. This covers the case where virtual models (which are +// filtered by SaveConfig) cannot be set as default. +func TestHandleSetDefaultModel_RejectsNonexistentModel(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + // First save a valid config with a primary model + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{ + {ModelName: "gpt-4", Model: "openai/gpt-4o"}, + } + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + // Try to set a non-existent model (like a virtual model name) as default + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/models/default", bytes.NewBufferString(`{ + "model_name": "gpt-4__key_1" + }`)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + + // Should return 404 because the virtual model doesn't exist in the persisted config + if rec.Code != http.StatusNotFound { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusNotFound, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "not found") { + t.Fatalf("error message should mention 'not found', got: %s", rec.Body.String()) + } +} + func TestMaskAPIKey(t *testing.T) { tests := []struct { name string diff --git a/web/frontend/src/api/models.ts b/web/frontend/src/api/models.ts index 2fd042593..aa66a7389 100644 --- a/web/frontend/src/api/models.ts +++ b/web/frontend/src/api/models.ts @@ -21,6 +21,7 @@ export interface ModelInfo { // Meta configured: boolean is_default: boolean + is_virtual: boolean } interface ModelsListResponse { diff --git a/web/frontend/src/components/models/model-card.tsx b/web/frontend/src/components/models/model-card.tsx index 316e05e4d..319cb11a3 100644 --- a/web/frontend/src/components/models/model-card.tsx +++ b/web/frontend/src/components/models/model-card.tsx @@ -28,7 +28,7 @@ export function ModelCard({ }: ModelCardProps) { const { t } = useTranslation() const isOAuth = model.auth_method === "oauth" - const canSetDefault = model.configured && !model.is_default + const canSetDefault = model.configured && !model.is_default && !model.is_virtual return (
)} + {model.is_virtual && ( + + {t("models.badge.virtual")} + + )}
diff --git a/web/frontend/src/i18n/locales/en.json b/web/frontend/src/i18n/locales/en.json index 68105fcea..e966219e6 100644 --- a/web/frontend/src/i18n/locales/en.json +++ b/web/frontend/src/i18n/locales/en.json @@ -154,7 +154,8 @@ "unconfigured": "Not configured" }, "badge": { - "default": "Default" + "default": "Default", + "virtual": "Virtual" }, "action": { "edit": "Edit API key", diff --git a/web/frontend/src/i18n/locales/zh.json b/web/frontend/src/i18n/locales/zh.json index 69ca939cb..9657d78b1 100644 --- a/web/frontend/src/i18n/locales/zh.json +++ b/web/frontend/src/i18n/locales/zh.json @@ -154,7 +154,8 @@ "unconfigured": "未配置" }, "badge": { - "default": "默认" + "default": "默认", + "virtual": "虚拟" }, "action": { "edit": "编辑 API Key",