diff --git a/pkg/config/config.go b/pkg/config/config.go index 3dc3422fb..3d88a136b 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -586,11 +586,12 @@ type ModelConfig struct { Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers // Optional optimizations - RPM int `json:"rpm,omitempty"` // Requests per minute limit - MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive - ExtraBody map[string]any `json:"extra_body,omitempty"` // Additional fields to inject into request body + RPM int `json:"rpm,omitempty"` // Requests per minute limit + MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive + ExtraBody map[string]any `json:"extra_body,omitempty"` // Additional fields to inject into request body + ExtraHeaders map[string]string `json:"extra_headers,omitempty"` // Additional HTTP headers to inject into provider requests APIKeys SecureStrings `json:"api_keys,omitzero" yaml:"api_keys,omitempty"` // API authentication keys (multiple keys for failover) diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 6734257f4..0ad2b809a 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -1454,6 +1454,42 @@ func TestModelConfig_ExtraBodyRoundTrip(t *testing.T) { } } +func TestModelConfig_ExtraHeadersRoundTrip(t *testing.T) { + dir := t.TempDir() + cfgPath := filepath.Join(dir, "config.json") + + cfg := &Config{ + Version: CurrentVersion, + ModelList: []*ModelConfig{ + { + ModelName: "test-model", + Model: "openai/test", + APIKeys: SimpleSecureStrings("sk-test"), + ExtraHeaders: map[string]string{"X-API-Key": "test-key", "X-Tenant": "demo"}, + }, + }, + } + + if err := SaveConfig(cfgPath, cfg); err != nil { + t.Fatalf("SaveConfig error: %v", err) + } + + loaded, err := LoadConfig(cfgPath) + if err != nil { + t.Fatalf("LoadConfig error: %v", err) + } + + if loaded.ModelList[0].ExtraHeaders == nil { + t.Fatal("ExtraHeaders should not be nil after round-trip") + } + if got := loaded.ModelList[0].ExtraHeaders["X-API-Key"]; got != "test-key" { + t.Errorf("ExtraHeaders[X-API-Key] = %v, want test-key", got) + } + if got := loaded.ModelList[0].ExtraHeaders["X-Tenant"]; got != "demo" { + t.Errorf("ExtraHeaders[X-Tenant] = %v, want demo", got) + } +} + func TestDefaultConfig_MinimaxExtraBody(t *testing.T) { cfg := DefaultConfig() diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index 962e6ae19..be8971565 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -98,6 +98,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err cfg.MaxTokensField, cfg.RequestTimeout, cfg.ExtraBody, + cfg.ExtraHeaders, ), modelID, nil case "azure", "azure-openai": @@ -174,6 +175,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err cfg.MaxTokensField, cfg.RequestTimeout, cfg.ExtraBody, + cfg.ExtraHeaders, ), modelID, nil case "minimax": @@ -199,6 +201,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err cfg.MaxTokensField, cfg.RequestTimeout, extraBody, + cfg.ExtraHeaders, ), modelID, nil case "anthropic": @@ -225,6 +228,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err cfg.MaxTokensField, cfg.RequestTimeout, cfg.ExtraBody, + cfg.ExtraHeaders, ), modelID, nil case "anthropic-messages": diff --git a/pkg/providers/http_provider.go b/pkg/providers/http_provider.go index f2ff52f1d..f3a2034a2 100644 --- a/pkg/providers/http_provider.go +++ b/pkg/providers/http_provider.go @@ -24,13 +24,14 @@ func NewHTTPProvider(apiKey, apiBase, proxy string) *HTTPProvider { } func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string) *HTTPProvider { - return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, 0, nil) + return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, 0, nil, nil) } func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout( apiKey, apiBase, proxy, maxTokensField string, requestTimeoutSeconds int, extraBody map[string]any, + extraHeaders map[string]string, ) *HTTPProvider { return &HTTPProvider{ delegate: openai_compat.NewProvider( @@ -40,6 +41,7 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout( openai_compat.WithMaxTokensField(maxTokensField), openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second), openai_compat.WithExtraBody(extraBody), + openai_compat.WithExtraHeaders(extraHeaders), ), } } diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 90bc683b8..9313936a1 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -36,6 +36,7 @@ type Provider struct { maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models) httpClient *http.Client extraBody map[string]any // Additional fields to inject into request body + extraHeaders map[string]string } type Option func(*Provider) @@ -62,6 +63,12 @@ func WithExtraBody(extraBody map[string]any) Option { } } +func WithExtraHeaders(extraHeaders map[string]string) Option { + return func(p *Provider) { + p.extraHeaders = extraHeaders + } +} + func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { p := &Provider{ apiKey: apiKey, @@ -183,6 +190,9 @@ func (p *Provider) Chat( if p.apiKey != "" { req.Header.Set("Authorization", "Bearer "+p.apiKey) } + for key, value := range p.extraHeaders { + req.Header.Set(key, value) + } resp, err := p.httpClient.Do(req) if err != nil { @@ -229,6 +239,9 @@ func (p *Provider) ChatStream( if p.apiKey != "" { req.Header.Set("Authorization", "Bearer "+p.apiKey) } + for key, value := range p.extraHeaders { + req.Header.Set(key, value) + } // Use a client without Timeout for streaming — the http.Client.Timeout covers // the entire request lifecycle including body reads, which would kill long streams. diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index ab632ccf3..4b157ec37 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -694,6 +694,104 @@ func TestProviderChat_ExtraBodyOverridesOptions(t *testing.T) { } } +func TestProviderChat_ExtraHeadersInjected(t *testing.T) { + var capturedAuthorization string + var capturedAPIKey string + var capturedTenant string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedAuthorization = r.Header.Get("Authorization") + capturedAPIKey = r.Header.Get("X-API-Key") + capturedTenant = r.Header.Get("X-Tenant") + + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "ok"}, + "finish_reason": "stop", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("bearer-token", server.URL, "", WithExtraHeaders(map[string]string{ + "X-API-Key": "secondary-key", + "X-Tenant": "tenant-a", + })) + + _, err := p.Chat( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "openai/gpt-4o-mini", + nil, + ) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if capturedAuthorization != "Bearer bearer-token" { + t.Fatalf("Authorization = %q, want %q", capturedAuthorization, "Bearer bearer-token") + } + if capturedAPIKey != "secondary-key" { + t.Fatalf("X-API-Key = %q, want %q", capturedAPIKey, "secondary-key") + } + if capturedTenant != "tenant-a" { + t.Fatalf("X-Tenant = %q, want %q", capturedTenant, "tenant-a") + } +} + +func TestProviderChatStream_ExtraHeadersInjected(t *testing.T) { + var capturedAuthorization string + var capturedAPIKey string + var capturedTenant string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedAuthorization = r.Header.Get("Authorization") + capturedAPIKey = r.Header.Get("X-API-Key") + capturedTenant = r.Header.Get("X-Tenant") + + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer server.Close() + + p := NewProvider("bearer-token", server.URL, "", WithExtraHeaders(map[string]string{ + "X-API-Key": "secondary-key", + "X-Tenant": "tenant-a", + })) + + resp, err := p.ChatStream( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "openai/gpt-4o-mini", + nil, + nil, + ) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + + if resp.Content != "ok" { + t.Fatalf("stream content = %q, want %q", resp.Content, "ok") + } + if capturedAuthorization != "Bearer bearer-token" { + t.Fatalf("Authorization = %q, want %q", capturedAuthorization, "Bearer bearer-token") + } + if capturedAPIKey != "secondary-key" { + t.Fatalf("X-API-Key = %q, want %q", capturedAPIKey, "secondary-key") + } + if capturedTenant != "tenant-a" { + t.Fatalf("X-Tenant = %q, want %q", capturedTenant, "tenant-a") + } +} + type roundTripperFunc func(*http.Request) (*http.Response, error) func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { diff --git a/web/backend/api/models.go b/web/backend/api/models.go index fd3cd85b7..372a20485 100644 --- a/web/backend/api/models.go +++ b/web/backend/api/models.go @@ -22,7 +22,7 @@ func (h *Handler) registerModelRoutes(mux *http.ServeMux) { } // modelResponse is the JSON structure returned for each model in the list. -// All ModelConfig fields are included so the frontend can display and edit them. +// Sensitive fields are masked and write-only fields are represented by metadata flags. type modelResponse struct { Index int `json:"index"` ModelName string `json:"model_name"` @@ -32,13 +32,14 @@ type modelResponse struct { Proxy string `json:"proxy,omitempty"` AuthMethod string `json:"auth_method,omitempty"` // Advanced fields - ConnectMode string `json:"connect_mode,omitempty"` - Workspace string `json:"workspace,omitempty"` - RPM int `json:"rpm,omitempty"` - MaxTokensField string `json:"max_tokens_field,omitempty"` - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` - ExtraBody map[string]any `json:"extra_body,omitempty"` + ConnectMode string `json:"connect_mode,omitempty"` + Workspace string `json:"workspace,omitempty"` + RPM int `json:"rpm,omitempty"` + MaxTokensField string `json:"max_tokens_field,omitempty"` + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` + ExtraBody map[string]any `json:"extra_body,omitempty"` + HasExtraHeaders bool `json:"has_extra_headers"` // Meta Enabled bool `json:"enabled"` Configured bool `json:"configured"` @@ -72,24 +73,25 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) { models := make([]modelResponse, 0, len(cfg.ModelList)) for i, m := range cfg.ModelList { models = append(models, modelResponse{ - Index: i, - ModelName: m.ModelName, - Model: m.Model, - APIBase: m.APIBase, - APIKey: maskAPIKey(m.APIKey()), - Proxy: m.Proxy, - AuthMethod: m.AuthMethod, - ConnectMode: m.ConnectMode, - Workspace: m.Workspace, - RPM: m.RPM, - MaxTokensField: m.MaxTokensField, - RequestTimeout: m.RequestTimeout, - ThinkingLevel: m.ThinkingLevel, - ExtraBody: m.ExtraBody, - Enabled: m.Enabled, - Configured: configured[i], - IsDefault: m.ModelName == defaultModel, - IsVirtual: m.IsVirtual(), + Index: i, + ModelName: m.ModelName, + Model: m.Model, + APIBase: m.APIBase, + APIKey: maskAPIKey(m.APIKey()), + Proxy: m.Proxy, + AuthMethod: m.AuthMethod, + ConnectMode: m.ConnectMode, + Workspace: m.Workspace, + RPM: m.RPM, + MaxTokensField: m.MaxTokensField, + RequestTimeout: m.RequestTimeout, + ThinkingLevel: m.ThinkingLevel, + ExtraBody: m.ExtraBody, + HasExtraHeaders: len(m.ExtraHeaders) > 0, + Enabled: m.Enabled, + Configured: configured[i], + IsDefault: m.ModelName == defaultModel, + IsVirtual: m.IsVirtual(), }) } @@ -214,6 +216,13 @@ func (h *Handler) handleUpdateModel(w http.ResponseWriter, r *http.Request) { } else if len(mc.ExtraBody) == 0 { mc.ExtraBody = nil } + // Preserve existing ExtraHeaders when omitted (nil), but clear them when + // the frontend sends an empty object {}. + if mc.ExtraHeaders == nil { + mc.ExtraHeaders = cfg.ModelList[idx].ExtraHeaders + } else if len(mc.ExtraHeaders) == 0 { + mc.ExtraHeaders = nil + } cfg.ModelList[idx] = &mc.ModelConfig diff --git a/web/backend/api/models_test.go b/web/backend/api/models_test.go index 97f153a80..a9308ba96 100644 --- a/web/backend/api/models_test.go +++ b/web/backend/api/models_test.go @@ -352,6 +352,169 @@ func TestHandleAddModel_PersistsAPIKey(t *testing.T) { } } +func TestHandleAddModel_PersistsExtraHeaders(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/models", bytes.NewBufferString(`{ + "model_name":"dual-header-model", + "model":"openai/gpt-4o-mini", + "api_key":"sk-new-model-key", + "extra_headers":{"X-API-Key":"secondary-key","X-Tenant":"tenant-a"} + }`)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + if len(cfg.ModelList) != 2 { + t.Fatalf("len(model_list) = %d, want 2", len(cfg.ModelList)) + } + + added := cfg.ModelList[1] + if got := added.ExtraHeaders["X-API-Key"]; got != "secondary-key" { + t.Fatalf("extra_headers[X-API-Key] = %q, want %q", got, "secondary-key") + } + if got := added.ExtraHeaders["X-Tenant"]; got != "tenant-a" { + t.Fatalf("extra_headers[X-Tenant] = %q, want %q", got, "tenant-a") + } +} + +func TestHandleUpdateModel_PreservesAndClearsExtraHeaders(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "primary", + Model: "openai/gpt-4o-mini", + APIKeys: config.SimpleSecureStrings("sk-primary"), + ExtraHeaders: map[string]string{"X-API-Key": "secondary-key"}, + }, + } + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + // Omit extra_headers: should preserve existing value. + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{ + "model_name":"primary", + "model":"openai/gpt-4o-mini", + "api_base":"https://api.example.com/v1" + }`)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("preserve status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + cfg, err = config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() after preserve error = %v", err) + } + if got := cfg.ModelList[0].ExtraHeaders["X-API-Key"]; got != "secondary-key" { + t.Fatalf("preserve extra_headers[X-API-Key] = %q, want %q", got, "secondary-key") + } + + // Send empty object: should clear existing value. + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{ + "model_name":"primary", + "model":"openai/gpt-4o-mini", + "api_base":"https://api.example.com/v1", + "extra_headers":{} + }`)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("clear status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + cfg, err = config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() after clear error = %v", err) + } + if cfg.ModelList[0].ExtraHeaders != nil { + t.Fatalf("extra_headers = %#v, want nil after clear", cfg.ModelList[0].ExtraHeaders) + } +} + +func TestHandleListModels_DoesNotLeakExtraHeaderSecrets(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "dual-header-model", + Model: "openai/gpt-4o-mini", + APIKeys: config.SimpleSecureStrings("sk-primary"), + ExtraHeaders: map[string]string{"X-API-Key": "secondary-secret"}, + }, + } + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/models", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + var resp struct { + Models []struct { + ModelName string `json:"model_name"` + HasExtraHeaders bool `json:"has_extra_headers"` + ExtraHeaders map[string]string `json:"extra_headers,omitempty"` + } `json:"models"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if len(resp.Models) != 1 { + t.Fatalf("len(models) = %d, want 1", len(resp.Models)) + } + if !resp.Models[0].HasExtraHeaders { + t.Fatalf("has_extra_headers = false, want true") + } + if resp.Models[0].ExtraHeaders != nil { + t.Fatalf("extra_headers should be omitted from list response, got %#v", resp.Models[0].ExtraHeaders) + } + if strings.Contains(rec.Body.String(), "secondary-secret") { + t.Fatalf("raw extra header secret leaked in response body: %s", rec.Body.String()) + } +} + // 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. diff --git a/web/frontend/src/api/models.ts b/web/frontend/src/api/models.ts index d75b3ec3c..14781d5fa 100644 --- a/web/frontend/src/api/models.ts +++ b/web/frontend/src/api/models.ts @@ -19,6 +19,8 @@ export interface ModelInfo { request_timeout?: number thinking_level?: string extra_body?: Record + extra_headers?: Record + has_extra_headers?: boolean // Meta configured: boolean is_default: boolean diff --git a/web/frontend/src/components/models/add-model-sheet.tsx b/web/frontend/src/components/models/add-model-sheet.tsx index de9481391..b2d8c1b12 100644 --- a/web/frontend/src/components/models/add-model-sheet.tsx +++ b/web/frontend/src/components/models/add-model-sheet.tsx @@ -36,6 +36,7 @@ interface AddForm { requestTimeout: string thinkingLevel: string extraBody: string + extraHeaders: string } const EMPTY_ADD_FORM: AddForm = { @@ -52,6 +53,7 @@ const EMPTY_ADD_FORM: AddForm = { requestTimeout: "", thinkingLevel: "", extraBody: "", + extraHeaders: "", } interface AddModelSheetProps { @@ -136,6 +138,9 @@ export function AddModelSheet({ extra_body: form.extraBody.trim() ? JSON.parse(form.extraBody.trim()) : undefined, + extra_headers: form.extraHeaders.trim() + ? JSON.parse(form.extraHeaders.trim()) + : undefined, }) if (setAsDefault) { await setDefaultModel(modelName) @@ -324,6 +329,18 @@ export function AddModelSheet({ rows={3} /> + + +