Skip to content

fix(mcp): preserve complete input schemas #86

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mudler merged 1 commit into from
Oct 9, 2026
Merged

fix(mcp): preserve complete input schemas #86

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 107 additions & 0 deletions clients/schema_serialization_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
package clients

import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"reflect"
"testing"
"time"

"github.com/mudler/cogito"
"github.com/sashabaranov/go-openai"
)

// MCP discovery supplies RawMessage parameters. Exercise the actual HTTP
// serializers, including LocalAI's extensions and native-media message rewrite.
func TestClientsPreserveRawToolSchema(t *testing.T) {
const schema = `{"type":"object","$defs":{"id":{"type":"integer"}},"properties":{"status":{"anyOf":[{"enum":["pending","shipped","delivered","cancelled"],"type":"string"},{"type":"null"}],"default":null},"limit":{"type":"integer","default":20},"id":{"$ref":"#/$defs/id"},"never":false,"flag":{"type":"boolean","default":false},"tags":{"type":["array","null"],"items":true}},"allOf":[{"required":["id"]}],"additionalProperties":false}`
factories := map[string]func(string) cogito.StreamingLLM{
"OpenAI": func(url string) cogito.StreamingLLM { return NewOpenAILLM("test-model", "", url) },
"LocalAI": func(url string) cogito.StreamingLLM { return NewLocalAILLM("test-model", "", url) },
"LocalAI extensions": func(url string) cogito.StreamingLLM {
c := NewLocalAILLM("test-model", "", url)
c.SetMetadata(map[string]string{"test": "true"})
return c
},
"LocalAI native media": func(url string) cogito.StreamingLLM {
c := NewLocalAILLM("test-model", "", url)
c.SetPendingNativeParts([]cogito.NativePart{{Kind: cogito.MediaAudio, Data: "AAAA", Format: "wav"}})
return c
},
}
for name, factory := range factories {
for _, stream := range []bool{false, true} {
mode := "completion"
if stream {
mode = "stream"
}
t.Run(name+"/"+mode, func(t *testing.T) {
captured := make(chan json.RawMessage, 1)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var request struct {
Tools []struct {
Function struct {
Parameters json.RawMessage `json:"parameters"`
} `json:"function"`
} `json:"tools"`
}
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Error(err)
http.Error(w, "invalid JSON", 400)
return
}
if len(request.Tools) != 1 {
t.Errorf("got %d tools, want 1", len(request.Tools))
http.Error(w, "missing tool", 400)
return
}
captured <- request.Tools[0].Function.Parameters
if stream {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"))
} else {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`))
}
}))
defer srv.Close()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
client := factory(srv.URL + "/v1")
req := openai.ChatCompletionRequest{Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "test"}}, Tools: []openai.Tool{{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{Name: "example", Parameters: json.RawMessage(schema)}}}}
if stream {
events, err := client.CreateChatCompletionStream(ctx, req)
if err != nil {
t.Fatal(err)
}
for event := range events {
if event.Error != nil {
t.Fatal(event.Error)
}
}
} else {
if _, _, err := client.CreateChatCompletion(ctx, req); err != nil {
t.Fatal(err)
}
}
select {
case got := <-captured:
var wantValue, gotValue any
if err := json.Unmarshal([]byte(schema), &wantValue); err != nil {
t.Fatal(err)
}
if err := json.Unmarshal(got, &gotValue); err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(wantValue, gotValue) {
t.Fatalf("schema changed on HTTP boundary:\nwant %s\n got %s", schema, got)
}
case <-ctx.Done():
t.Fatal("no request captured")
}
})
}
}
}
240 changes: 6 additions & 234 deletions mcp.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,36 +9,22 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/mudler/xlog"
"github.com/sashabaranov/go-openai"
"github.com/tmc/langchaingo/jsonschema"
)

type mcpTool struct {
name, description string
inputSchema toolInputSchema
inputSchema json.RawMessage
session *mcp.ClientSession
ctx context.Context
props map[string]jsonschema.Definition
}

func (t *mcpTool) Tool() openai.Tool {
// A tool without arguments has no properties, and the unmarshal leaves
// the map nil. jsonschema.Definition encoded that as {}, but this
// struct would send null, which vLLM and others reject.
props := t.props
if props == nil {
props = map[string]jsonschema.Definition{}
}
return openai.Tool{
Type: openai.ToolTypeFunction,
Function: &openai.FunctionDefinition{
Name: t.name,
Description: t.description,
Parameters: mcpToolParameters{
Type: jsonschema.Object,
Properties: props,
Required: t.inputSchema.Required,
AdditionalProperties: t.inputSchema.AdditionalProperties,
},
Parameters: t.inputSchema,
},
}
}
Expand Down Expand Up @@ -105,194 +91,6 @@ func (t *mcpTool) Close() {
}
}

type toolInputSchema struct {
Type string `json:"type"`
Properties map[string]interface{} `json:"properties,omitempty"`
Required []string `json:"required,omitempty"`
// AdditionalProperties is kept as the server declared it. A server whose
// schema forbids unknown arguments (the go-sdk infers false for a Go
// struct) rejects a call that has one, so the model should see the same
// rule, and a strict or grammar-constrained backend can enforce it.
AdditionalProperties any `json:"additionalProperties,omitempty"`
}

// mcpToolParameters is the parameters schema an MCP tool is advertised
// with. It is not a jsonschema.Definition because that type has no
// additionalProperties field, and embedding it would not help: its value
// MarshalJSON would replace this struct's encoding.
type mcpToolParameters struct {
Type jsonschema.DataType `json:"type"`
Properties map[string]jsonschema.Definition `json:"properties"`
Required []string `json:"required,omitempty"`
AdditionalProperties any `json:"additionalProperties,omitempty"`
}

// CoerceNullableTypes is an exported alias for the same workaround so
// downstream tests can verify their own MCP servers stay compatible
// with the import path. Most callers won't need it.
func CoerceNullableTypes(props map[string]any) { coerceNullableTypes(props) }

// coerceNullableTypes recursively walks a property bag and normalizes the
// JSON-Schema 2020-12 constructs that langchaingo/jsonschema.Definition
// cannot represent. Both rewrites exist for the same reason: Definition is
// the struct we unmarshal every listed tool's properties into, and ANY
// failure there drops the whole tool from the model's tool list with a
// single log line — the server stays healthy and its tool simply never
// reaches the model.
//
// 1. "type": ["null", "X"] becomes "type": "X". Definition.Type is a single
// string. Picks the first non-null member; falls back to the first
// member if all are null.
// 2. A boolean schema becomes its object equivalent: `true` (allow
// anything) becomes {}, `false` (allow nothing) becomes {"not": {}}.
// 2020-12 permits a boolean wherever a schema is allowed, and
// google/jsonschema-go emits exactly that for a Go `any` — an empty
// schema marshals as `true`, so a [][]any field yields "items": true.
// Definition models nested schemas as *Definition, so a bare bool fails
// the unmarshal.
//
// Recurses into every nested schema location: properties, items,
// oneOf/anyOf/allOf members, prefixItems, $defs/definitions,
// additionalProperties, patternProperties, contains, not, if/then/else,
// propertyNames.
func coerceNullableTypes(props map[string]any) {
if props == nil {
return
}
// Range with the key so a property that is ITSELF a boolean schema can be
// replaced in place; the value-only loop could not rewrite it.
for name, raw := range props {
if b, ok := raw.(bool); ok {
props[name] = boolSchema(b)
continue
}
coerceSchema(raw)
}
}

// boolSchema returns the object form of a JSON-Schema boolean schema,
// preserving its meaning: `true` allows anything, `false` allows nothing.
// Collapsing both to {} would silently widen a deliberately-closed schema.
func boolSchema(allow bool) map[string]any {
if allow {
return map[string]any{}
}
return map[string]any{"not": map[string]any{}}
}

// schemaValuedKeys are the keywords whose value is a single schema, so a
// boolean there is a boolean schema rather than an ordinary flag.
var schemaValuedKeys = []string{
"items", "additionalProperties", "contains", "not",
"if", "then", "else", "propertyNames",
}

// schemaListKeys are the keywords whose value is an array of schemas.
var schemaListKeys = []string{"oneOf", "anyOf", "allOf", "prefixItems", "items"}

// normalizeBoolSchemas replaces every boolean schema directly under obj with
// its object equivalent. Nested bags (properties, $defs, …) are reached by
// coerceNullableTypes, which performs the same replacement for their members.
func normalizeBoolSchemas(obj map[string]any) {
for _, key := range schemaValuedKeys {
if b, ok := obj[key].(bool); ok {
obj[key] = boolSchema(b)
}
}
for _, key := range schemaListKeys {
arr, ok := obj[key].([]any)
if !ok {
continue
}
for i, member := range arr {
if b, ok := member.(bool); ok {
arr[i] = boolSchema(b)
}
}
}
}

// coerceSchema applies the type-array → string and boolean-schema →
// object rewrites to a single schema node, then recurses into every
// nested schema location.
func coerceSchema(node any) {
obj, ok := node.(map[string]any)
if !ok {
return
}

// Run first, so the objects it creates are recursed into below like any
// other nested schema.
normalizeBoolSchemas(obj)

if t, ok := obj["type"].([]any); ok {
pick := ""
for _, m := range t {
s, ok := m.(string)
if !ok || s == "null" {
continue
}
pick = s
break
}
if pick == "" && len(t) > 0 {
if s, ok := t[0].(string); ok {
pick = s
}
}
if pick != "" {
obj["type"] = pick
}
}

// Object properties — map of name → schema.
if nested, ok := obj["properties"].(map[string]any); ok {
coerceNullableTypes(nested)
}
// patternProperties — same shape as properties, just regex-keyed.
if nested, ok := obj["patternProperties"].(map[string]any); ok {
coerceNullableTypes(nested)
}
// $defs / definitions — JSON-Schema named schema bag.
if nested, ok := obj["$defs"].(map[string]any); ok {
coerceNullableTypes(nested)
}
if nested, ok := obj["definitions"].(map[string]any); ok {
coerceNullableTypes(nested)
}
// Single nested schema fields.
coerceSchema(obj["items"])
coerceSchema(obj["additionalProperties"])
coerceSchema(obj["contains"])
coerceSchema(obj["not"])
coerceSchema(obj["if"])
coerceSchema(obj["then"])
coerceSchema(obj["else"])
coerceSchema(obj["propertyNames"])
// Composition keywords — arrays of schemas.
for _, key := range []string{"oneOf", "anyOf", "allOf"} {
arr, ok := obj[key].([]any)
if !ok {
continue
}
for _, member := range arr {
coerceSchema(member)
}
}
// "items" can also be an array (tuple validation in older drafts).
if arr, ok := obj["items"].([]any); ok {
for _, member := range arr {
coerceSchema(member)
}
}
// "prefixItems" (2020-12 tuple).
if arr, ok := obj["prefixItems"].([]any); ok {
for _, member := range arr {
coerceSchema(member)
}
}
}

func mcpPromptsFromTransport(ctx context.Context, session *mcp.ClientSession, arguments map[string]string) ([]openai.ChatCompletionMessage, error) {
prompts, err := session.ListPrompts(ctx, nil)
if err != nil {
Expand Down Expand Up @@ -340,47 +138,21 @@ func mcpToolsFromTransport(ctx context.Context, session *mcp.ClientSession, filt
if filter != nil && !filter(session, tool.Name) {
continue
}
// Keep the official SDK schema intact, including keywords and boolean
// schemas that a restricted schema struct cannot represent. Do not
// normalize types, add properties, or apply defaults here.
dat, err := json.Marshal(tool.InputSchema)
if err != nil {
xlog.Error("Error marshalling input schema: %v", err)
continue
}

var inputSchema toolInputSchema
err = json.Unmarshal(dat, &inputSchema)
if err != nil {
xlog.Error("Error unmarshalling input schema: %v", err)
continue
}

// Some MCP servers (e.g. modelcontextprotocol/go-sdk v1.4+)
// emit JSON Schema 2020-12 "type": ["null", "array"] for
// nullable fields like Go []string slices. langchaingo's
// jsonschema.Definition.Type is a single string, so the
// unmarshal would fail and silently drop the entire tool.
// Coerce any type-array to its non-null string member before
// unmarshalling so those tools stay discoverable.
coerceNullableTypes(inputSchema.Properties)

props := map[string]jsonschema.Definition{}
dat, err = json.Marshal(inputSchema.Properties)
if err != nil {
xlog.Error("Error marshalling input schema: %v", err)
continue
}
err = json.Unmarshal(dat, &props)
if err != nil {
xlog.Error("Error unmarshalling input schema properties: %v", err)
continue
}

allTools = append(allTools, &mcpTool{
name: tool.Name,
description: tool.Description,
session: session,
ctx: ctx,
props: props,
inputSchema: inputSchema,
inputSchema: json.RawMessage(dat),
})
}

Expand Down
Loading
Loading