mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
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:
+34
-29
@@ -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 "".
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user