feat(model): support per-model extra HTTP headers for openai-compatible providers
This commit is contained in:
parent
cbe92286e9
commit
02427474ad
11 changed files with 410 additions and 32 deletions
|
|
@ -586,11 +586,12 @@ type ModelConfig struct {
|
||||||
Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers
|
Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers
|
||||||
|
|
||||||
// Optional optimizations
|
// Optional optimizations
|
||||||
RPM int `json:"rpm,omitempty"` // Requests per minute limit
|
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")
|
MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens")
|
||||||
RequestTimeout int `json:"request_timeout,omitempty"`
|
RequestTimeout int `json:"request_timeout,omitempty"`
|
||||||
ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive
|
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
|
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)
|
APIKeys SecureStrings `json:"api_keys,omitzero" yaml:"api_keys,omitempty"` // API authentication keys (multiple keys for failover)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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) {
|
func TestDefaultConfig_MinimaxExtraBody(t *testing.T) {
|
||||||
cfg := DefaultConfig()
|
cfg := DefaultConfig()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -98,6 +98,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
||||||
cfg.MaxTokensField,
|
cfg.MaxTokensField,
|
||||||
cfg.RequestTimeout,
|
cfg.RequestTimeout,
|
||||||
cfg.ExtraBody,
|
cfg.ExtraBody,
|
||||||
|
cfg.ExtraHeaders,
|
||||||
), modelID, nil
|
), modelID, nil
|
||||||
|
|
||||||
case "azure", "azure-openai":
|
case "azure", "azure-openai":
|
||||||
|
|
@ -174,6 +175,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
||||||
cfg.MaxTokensField,
|
cfg.MaxTokensField,
|
||||||
cfg.RequestTimeout,
|
cfg.RequestTimeout,
|
||||||
cfg.ExtraBody,
|
cfg.ExtraBody,
|
||||||
|
cfg.ExtraHeaders,
|
||||||
), modelID, nil
|
), modelID, nil
|
||||||
|
|
||||||
case "minimax":
|
case "minimax":
|
||||||
|
|
@ -199,6 +201,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
||||||
cfg.MaxTokensField,
|
cfg.MaxTokensField,
|
||||||
cfg.RequestTimeout,
|
cfg.RequestTimeout,
|
||||||
extraBody,
|
extraBody,
|
||||||
|
cfg.ExtraHeaders,
|
||||||
), modelID, nil
|
), modelID, nil
|
||||||
|
|
||||||
case "anthropic":
|
case "anthropic":
|
||||||
|
|
@ -225,6 +228,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
||||||
cfg.MaxTokensField,
|
cfg.MaxTokensField,
|
||||||
cfg.RequestTimeout,
|
cfg.RequestTimeout,
|
||||||
cfg.ExtraBody,
|
cfg.ExtraBody,
|
||||||
|
cfg.ExtraHeaders,
|
||||||
), modelID, nil
|
), modelID, nil
|
||||||
|
|
||||||
case "anthropic-messages":
|
case "anthropic-messages":
|
||||||
|
|
|
||||||
|
|
@ -24,13 +24,14 @@ func NewHTTPProvider(apiKey, apiBase, proxy string) *HTTPProvider {
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField 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(
|
func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
|
||||||
apiKey, apiBase, proxy, maxTokensField string,
|
apiKey, apiBase, proxy, maxTokensField string,
|
||||||
requestTimeoutSeconds int,
|
requestTimeoutSeconds int,
|
||||||
extraBody map[string]any,
|
extraBody map[string]any,
|
||||||
|
extraHeaders map[string]string,
|
||||||
) *HTTPProvider {
|
) *HTTPProvider {
|
||||||
return &HTTPProvider{
|
return &HTTPProvider{
|
||||||
delegate: openai_compat.NewProvider(
|
delegate: openai_compat.NewProvider(
|
||||||
|
|
@ -40,6 +41,7 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
|
||||||
openai_compat.WithMaxTokensField(maxTokensField),
|
openai_compat.WithMaxTokensField(maxTokensField),
|
||||||
openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
|
openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second),
|
||||||
openai_compat.WithExtraBody(extraBody),
|
openai_compat.WithExtraBody(extraBody),
|
||||||
|
openai_compat.WithExtraHeaders(extraHeaders),
|
||||||
),
|
),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -36,6 +36,7 @@ type Provider struct {
|
||||||
maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
|
maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models)
|
||||||
httpClient *http.Client
|
httpClient *http.Client
|
||||||
extraBody map[string]any // Additional fields to inject into request body
|
extraBody map[string]any // Additional fields to inject into request body
|
||||||
|
extraHeaders map[string]string
|
||||||
}
|
}
|
||||||
|
|
||||||
type Option func(*Provider)
|
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 {
|
func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
|
||||||
p := &Provider{
|
p := &Provider{
|
||||||
apiKey: apiKey,
|
apiKey: apiKey,
|
||||||
|
|
@ -183,6 +190,9 @@ func (p *Provider) Chat(
|
||||||
if p.apiKey != "" {
|
if p.apiKey != "" {
|
||||||
req.Header.Set("Authorization", "Bearer "+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)
|
resp, err := p.httpClient.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
@ -229,6 +239,9 @@ func (p *Provider) ChatStream(
|
||||||
if p.apiKey != "" {
|
if p.apiKey != "" {
|
||||||
req.Header.Set("Authorization", "Bearer "+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
|
// Use a client without Timeout for streaming — the http.Client.Timeout covers
|
||||||
// the entire request lifecycle including body reads, which would kill long streams.
|
// the entire request lifecycle including body reads, which would kill long streams.
|
||||||
|
|
|
||||||
|
|
@ -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)
|
type roundTripperFunc func(*http.Request) (*http.Response, error)
|
||||||
|
|
||||||
func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||||
|
|
|
||||||
|
|
@ -22,7 +22,7 @@ func (h *Handler) registerModelRoutes(mux *http.ServeMux) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// modelResponse is the JSON structure returned for each model in the list.
|
// 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 {
|
type modelResponse struct {
|
||||||
Index int `json:"index"`
|
Index int `json:"index"`
|
||||||
ModelName string `json:"model_name"`
|
ModelName string `json:"model_name"`
|
||||||
|
|
@ -32,13 +32,14 @@ type modelResponse struct {
|
||||||
Proxy string `json:"proxy,omitempty"`
|
Proxy string `json:"proxy,omitempty"`
|
||||||
AuthMethod string `json:"auth_method,omitempty"`
|
AuthMethod string `json:"auth_method,omitempty"`
|
||||||
// Advanced fields
|
// Advanced fields
|
||||||
ConnectMode string `json:"connect_mode,omitempty"`
|
ConnectMode string `json:"connect_mode,omitempty"`
|
||||||
Workspace string `json:"workspace,omitempty"`
|
Workspace string `json:"workspace,omitempty"`
|
||||||
RPM int `json:"rpm,omitempty"`
|
RPM int `json:"rpm,omitempty"`
|
||||||
MaxTokensField string `json:"max_tokens_field,omitempty"`
|
MaxTokensField string `json:"max_tokens_field,omitempty"`
|
||||||
RequestTimeout int `json:"request_timeout,omitempty"`
|
RequestTimeout int `json:"request_timeout,omitempty"`
|
||||||
ThinkingLevel string `json:"thinking_level,omitempty"`
|
ThinkingLevel string `json:"thinking_level,omitempty"`
|
||||||
ExtraBody map[string]any `json:"extra_body,omitempty"`
|
ExtraBody map[string]any `json:"extra_body,omitempty"`
|
||||||
|
HasExtraHeaders bool `json:"has_extra_headers"`
|
||||||
// Meta
|
// Meta
|
||||||
Enabled bool `json:"enabled"`
|
Enabled bool `json:"enabled"`
|
||||||
Configured bool `json:"configured"`
|
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))
|
models := make([]modelResponse, 0, len(cfg.ModelList))
|
||||||
for i, m := range cfg.ModelList {
|
for i, m := range cfg.ModelList {
|
||||||
models = append(models, modelResponse{
|
models = append(models, modelResponse{
|
||||||
Index: i,
|
Index: i,
|
||||||
ModelName: m.ModelName,
|
ModelName: m.ModelName,
|
||||||
Model: m.Model,
|
Model: m.Model,
|
||||||
APIBase: m.APIBase,
|
APIBase: m.APIBase,
|
||||||
APIKey: maskAPIKey(m.APIKey()),
|
APIKey: maskAPIKey(m.APIKey()),
|
||||||
Proxy: m.Proxy,
|
Proxy: m.Proxy,
|
||||||
AuthMethod: m.AuthMethod,
|
AuthMethod: m.AuthMethod,
|
||||||
ConnectMode: m.ConnectMode,
|
ConnectMode: m.ConnectMode,
|
||||||
Workspace: m.Workspace,
|
Workspace: m.Workspace,
|
||||||
RPM: m.RPM,
|
RPM: m.RPM,
|
||||||
MaxTokensField: m.MaxTokensField,
|
MaxTokensField: m.MaxTokensField,
|
||||||
RequestTimeout: m.RequestTimeout,
|
RequestTimeout: m.RequestTimeout,
|
||||||
ThinkingLevel: m.ThinkingLevel,
|
ThinkingLevel: m.ThinkingLevel,
|
||||||
ExtraBody: m.ExtraBody,
|
ExtraBody: m.ExtraBody,
|
||||||
Enabled: m.Enabled,
|
HasExtraHeaders: len(m.ExtraHeaders) > 0,
|
||||||
Configured: configured[i],
|
Enabled: m.Enabled,
|
||||||
IsDefault: m.ModelName == defaultModel,
|
Configured: configured[i],
|
||||||
IsVirtual: m.IsVirtual(),
|
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 {
|
} else if len(mc.ExtraBody) == 0 {
|
||||||
mc.ExtraBody = nil
|
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
|
cfg.ModelList[idx] = &mc.ModelConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
// TestHandleSetDefaultModel_RejectsNonexistentModel tests that setting a non-existent
|
||||||
// model as default returns 404. This covers the case where virtual models (which are
|
// model as default returns 404. This covers the case where virtual models (which are
|
||||||
// filtered by SaveConfig) cannot be set as default.
|
// filtered by SaveConfig) cannot be set as default.
|
||||||
|
|
|
||||||
|
|
@ -19,6 +19,8 @@ export interface ModelInfo {
|
||||||
request_timeout?: number
|
request_timeout?: number
|
||||||
thinking_level?: string
|
thinking_level?: string
|
||||||
extra_body?: Record<string, unknown>
|
extra_body?: Record<string, unknown>
|
||||||
|
extra_headers?: Record<string, string>
|
||||||
|
has_extra_headers?: boolean
|
||||||
// Meta
|
// Meta
|
||||||
configured: boolean
|
configured: boolean
|
||||||
is_default: boolean
|
is_default: boolean
|
||||||
|
|
|
||||||
|
|
@ -36,6 +36,7 @@ interface AddForm {
|
||||||
requestTimeout: string
|
requestTimeout: string
|
||||||
thinkingLevel: string
|
thinkingLevel: string
|
||||||
extraBody: string
|
extraBody: string
|
||||||
|
extraHeaders: string
|
||||||
}
|
}
|
||||||
|
|
||||||
const EMPTY_ADD_FORM: AddForm = {
|
const EMPTY_ADD_FORM: AddForm = {
|
||||||
|
|
@ -52,6 +53,7 @@ const EMPTY_ADD_FORM: AddForm = {
|
||||||
requestTimeout: "",
|
requestTimeout: "",
|
||||||
thinkingLevel: "",
|
thinkingLevel: "",
|
||||||
extraBody: "",
|
extraBody: "",
|
||||||
|
extraHeaders: "",
|
||||||
}
|
}
|
||||||
|
|
||||||
interface AddModelSheetProps {
|
interface AddModelSheetProps {
|
||||||
|
|
@ -136,6 +138,9 @@ export function AddModelSheet({
|
||||||
extra_body: form.extraBody.trim()
|
extra_body: form.extraBody.trim()
|
||||||
? JSON.parse(form.extraBody.trim())
|
? JSON.parse(form.extraBody.trim())
|
||||||
: undefined,
|
: undefined,
|
||||||
|
extra_headers: form.extraHeaders.trim()
|
||||||
|
? JSON.parse(form.extraHeaders.trim())
|
||||||
|
: undefined,
|
||||||
})
|
})
|
||||||
if (setAsDefault) {
|
if (setAsDefault) {
|
||||||
await setDefaultModel(modelName)
|
await setDefaultModel(modelName)
|
||||||
|
|
@ -324,6 +329,18 @@ export function AddModelSheet({
|
||||||
rows={3}
|
rows={3}
|
||||||
/>
|
/>
|
||||||
</Field>
|
</Field>
|
||||||
|
|
||||||
|
<Field
|
||||||
|
label="Extra Headers (JSON)"
|
||||||
|
hint='Optional HTTP headers sent with model requests, for example {"X-API-Key":"..."}'
|
||||||
|
>
|
||||||
|
<Textarea
|
||||||
|
value={form.extraHeaders}
|
||||||
|
onChange={setField("extraHeaders")}
|
||||||
|
placeholder='{"X-API-Key": "secondary-key"}'
|
||||||
|
rows={3}
|
||||||
|
/>
|
||||||
|
</Field>
|
||||||
</AdvancedSection>
|
</AdvancedSection>
|
||||||
|
|
||||||
{serverError && (
|
{serverError && (
|
||||||
|
|
|
||||||
|
|
@ -34,6 +34,7 @@ interface EditForm {
|
||||||
requestTimeout: string
|
requestTimeout: string
|
||||||
thinkingLevel: string
|
thinkingLevel: string
|
||||||
extraBody: string
|
extraBody: string
|
||||||
|
extraHeaders: string
|
||||||
}
|
}
|
||||||
|
|
||||||
interface EditModelSheetProps {
|
interface EditModelSheetProps {
|
||||||
|
|
@ -62,9 +63,11 @@ export function EditModelSheet({
|
||||||
requestTimeout: "",
|
requestTimeout: "",
|
||||||
thinkingLevel: "",
|
thinkingLevel: "",
|
||||||
extraBody: "",
|
extraBody: "",
|
||||||
|
extraHeaders: "",
|
||||||
})
|
})
|
||||||
const [saving, setSaving] = useState(false)
|
const [saving, setSaving] = useState(false)
|
||||||
const [setAsDefault, setSetAsDefault] = useState(false)
|
const [setAsDefault, setSetAsDefault] = useState(false)
|
||||||
|
const [clearExtraHeaders, setClearExtraHeaders] = useState(false)
|
||||||
const [error, setError] = useState("")
|
const [error, setError] = useState("")
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
|
|
@ -85,8 +88,10 @@ export function EditModelSheet({
|
||||||
extraBody: model.extra_body
|
extraBody: model.extra_body
|
||||||
? JSON.stringify(model.extra_body, null, 2)
|
? JSON.stringify(model.extra_body, null, 2)
|
||||||
: "",
|
: "",
|
||||||
|
extraHeaders: "",
|
||||||
})
|
})
|
||||||
setSetAsDefault(model.is_default)
|
setSetAsDefault(model.is_default)
|
||||||
|
setClearExtraHeaders(false)
|
||||||
setError("")
|
setError("")
|
||||||
}
|
}
|
||||||
}, [model])
|
}, [model])
|
||||||
|
|
@ -119,6 +124,11 @@ export function EditModelSheet({
|
||||||
extra_body: form.extraBody.trim()
|
extra_body: form.extraBody.trim()
|
||||||
? JSON.parse(form.extraBody.trim())
|
? JSON.parse(form.extraBody.trim())
|
||||||
: {},
|
: {},
|
||||||
|
extra_headers: clearExtraHeaders
|
||||||
|
? {}
|
||||||
|
: form.extraHeaders.trim()
|
||||||
|
? JSON.parse(form.extraHeaders.trim())
|
||||||
|
: undefined,
|
||||||
})
|
})
|
||||||
if (setAsDefault && !model.is_default) {
|
if (setAsDefault && !model.is_default) {
|
||||||
await setDefaultModel(model.model_name)
|
await setDefaultModel(model.model_name)
|
||||||
|
|
@ -295,6 +305,29 @@ export function EditModelSheet({
|
||||||
rows={3}
|
rows={3}
|
||||||
/>
|
/>
|
||||||
</Field>
|
</Field>
|
||||||
|
|
||||||
|
<Field
|
||||||
|
label="Extra Headers (JSON)"
|
||||||
|
hint='Write-only HTTP headers for model requests. Leave blank to keep current headers, enter JSON to replace, or use the clear switch below.'
|
||||||
|
>
|
||||||
|
<Textarea
|
||||||
|
value={form.extraHeaders}
|
||||||
|
onChange={setField("extraHeaders")}
|
||||||
|
placeholder='{"X-API-Key": "secondary-key"}'
|
||||||
|
rows={3}
|
||||||
|
/>
|
||||||
|
</Field>
|
||||||
|
|
||||||
|
<SwitchCardField
|
||||||
|
label="Clear Existing Extra Headers"
|
||||||
|
hint={
|
||||||
|
model?.has_extra_headers
|
||||||
|
? "This model currently has saved extra headers. Enable to clear all of them."
|
||||||
|
: "Enable to send an explicit clear request for extra headers."
|
||||||
|
}
|
||||||
|
checked={clearExtraHeaders}
|
||||||
|
onCheckedChange={setClearExtraHeaders}
|
||||||
|
/>
|
||||||
</AdvancedSection>
|
</AdvancedSection>
|
||||||
|
|
||||||
{error && (
|
{error && (
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue