Merge pull request #2681 from afjcjsbx/fix/gemini-mcp-schema-sanitization

fix(mcp): sanitize MCP tool schemas for Gemini function calling
This commit is contained in:
Mauro
2026-05-03 20:25:35 +02:00
committed by GitHub
25 changed files with 1663 additions and 248 deletions
+34 -29
View File
@@ -35,14 +35,15 @@ 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"`
CustomHeaders map[string]string `json:"custom_headers,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"`
ToolSchemaTransform string `json:"tool_schema_transform,omitempty"`
ExtraBody map[string]any `json:"extra_body,omitempty"`
CustomHeaders map[string]string `json:"custom_headers,omitempty"`
// Meta
Enabled bool `json:"enabled"`
Available bool `json:"available"`
@@ -78,27 +79,28 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) {
for i, m := range cfg.ModelList {
provider, modelID := providers.ExtractProtocol(m)
models = append(models, modelResponse{
Index: i,
ModelName: m.ModelName,
Provider: provider,
Model: modelID,
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,
CustomHeaders: m.CustomHeaders,
Enabled: m.Enabled,
Available: modelStatuses[i].Available,
Status: modelStatuses[i].Status,
IsDefault: m.ModelName == defaultModel,
IsVirtual: m.IsVirtual(),
Index: i,
ModelName: m.ModelName,
Provider: provider,
Model: modelID,
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,
ToolSchemaTransform: m.ToolSchemaTransform,
ExtraBody: m.ExtraBody,
CustomHeaders: m.CustomHeaders,
Enabled: m.Enabled,
Available: modelStatuses[i].Available,
Status: modelStatuses[i].Status,
IsDefault: m.ModelName == defaultModel,
IsVirtual: m.IsVirtual(),
})
}
@@ -237,6 +239,9 @@ func (h *Handler) handleUpdateModel(w http.ResponseWriter, r *http.Request) {
} else if len(mc.CustomHeaders) == 0 {
mc.CustomHeaders = nil
}
if _, ok := rawFields["tool_schema_transform"]; !ok {
mc.ToolSchemaTransform = cfg.ModelList[idx].ToolSchemaTransform
}
// Preserve the existing Provider when the caller omits it. This keeps the
// update API backward-compatible for clients that haven't started sending
// the new field yet, while still allowing explicit clearing via "".
+94
View File
@@ -584,6 +584,37 @@ func TestHandleAddModel_PersistsCustomHeaders(t *testing.T) {
}
}
func TestHandleAddModel_PersistsToolSchemaTransform(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":"new-model-transform",
"model":"openai/gpt-4o-mini",
"tool_schema_transform":"simple"
}`))
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)
}
added := cfg.ModelList[len(cfg.ModelList)-1]
if got := added.ToolSchemaTransform; got != "simple" {
t.Fatalf("tool_schema_transform = %q, want %q", got, "simple")
}
}
func TestHandleUpdateModel_CustomHeadersPreserveAndClear(t *testing.T) {
configPath, cleanup := setupOAuthTestEnv(t)
defer cleanup()
@@ -649,6 +680,69 @@ func TestHandleUpdateModel_CustomHeadersPreserveAndClear(t *testing.T) {
}
}
func TestHandleUpdateModel_ToolSchemaTransformPreserveAndClear(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: "editable",
Model: "openai/gpt-4o-mini",
APIKeys: config.SimpleSecureStrings("sk-existing"),
ToolSchemaTransform: "simple",
}}
err = config.SaveConfig(configPath, cfg)
if err != nil {
t.Fatalf("SaveConfig() error = %v", err)
}
h := NewHandler(configPath)
mux := http.NewServeMux()
h.RegisterRoutes(mux)
recPreserve := httptest.NewRecorder()
reqPreserve := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{
"model_name":"editable",
"model":"openai/gpt-4o-mini"
}`))
reqPreserve.Header.Set("Content-Type", "application/json")
mux.ServeHTTP(recPreserve, reqPreserve)
if recPreserve.Code != http.StatusOK {
t.Fatalf("preserve status = %d, want %d, body=%s", recPreserve.Code, http.StatusOK, recPreserve.Body.String())
}
afterPreserve, err := config.LoadConfig(configPath)
if err != nil {
t.Fatalf("LoadConfig() after preserve error = %v", err)
}
if got := afterPreserve.ModelList[0].ToolSchemaTransform; got != "simple" {
t.Fatalf("preserved tool_schema_transform = %q, want %q", got, "simple")
}
recClear := httptest.NewRecorder()
reqClear := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{
"model_name":"editable",
"model":"openai/gpt-4o-mini",
"tool_schema_transform":""
}`))
reqClear.Header.Set("Content-Type", "application/json")
mux.ServeHTTP(recClear, reqClear)
if recClear.Code != http.StatusOK {
t.Fatalf("clear status = %d, want %d, body=%s", recClear.Code, http.StatusOK, recClear.Body.String())
}
afterClear, err := config.LoadConfig(configPath)
if err != nil {
t.Fatalf("LoadConfig() after clear error = %v", err)
}
if afterClear.ModelList[0].ToolSchemaTransform != "" {
t.Fatalf("tool_schema_transform = %q, want empty", afterClear.ModelList[0].ToolSchemaTransform)
}
}
func TestHandleUpdateModel_PersistsProvider(t *testing.T) {
configPath, cleanup := setupOAuthTestEnv(t)
defer cleanup()