From bca971d875ac5a2ece18a8496248a6dde5233457 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@localai.io>
Date: Sun, 27 Sep 2026 19:55:37 +0000
Subject: [PATCH] feat(tools): add EnableStrictToolSchemas and keep MCP
 additionalProperties

A tool's schema only described the arguments; nothing made the model
follow it. In nib, models called the bash tool with "command" instead
of "script", the MCP server rejected every call, and the model gave up
on the shell. LocalAI disables its own tool grammar for models that use
the backend's chat template, and llama.cpp's grammar is lazy under
tool_choice auto, so those calls were never constrained.

EnableStrictToolSchemas wraps every tool usableTools returns so it is
sent with "strict": true and a schema that follows strict mode's rules:
additionalProperties false on every object, every property required,
formerly optional ones nullable. The wrapper drops the nulls the model
then sends before it calls the tool. A schema strict mode cannot
express is sent unchanged. LocalAI switches its tool grammar on for a
request with a strict tool; OpenAI enforces the schema.

MCP tools were advertised with a jsonschema.Definition, which has no
additionalProperties field, so a server's "additionalProperties": false
never reached the model. They now carry it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 README.md      |  30 ++++++++
 guidelines.go  |   4 +
 mcp.go         |  25 ++++++-
 options.go     |  11 +++
 strict.go      | 194 +++++++++++++++++++++++++++++++++++++++++++++++++
 strict_test.go | 166 ++++++++++++++++++++++++++++++++++++++++++
 6 files changed, 426 insertions(+), 4 deletions(-)
 create mode 100644 strict.go
 create mode 100644 strict_test.go

diff --git a/README.md b/README.md
index ba48ed4..5bca30f 100644
--- a/README.md
+++ b/README.md
@@ -1389,6 +1389,36 @@ make example-chat
 
 This starts an interactive chat session with tool support including web search capabilities.
 
+### Strict Tool Schemas
+
+A tool's schema tells the model which arguments to send, but on most
+backends nothing makes the model follow it: a call with a misspelled or
+unknown argument reaches the tool, and the tool rejects it. With
+`EnableStrictToolSchemas`, cogito sends each tool with `"strict": true`, so a
+backend that supports it constrains the arguments while the model generates
+them. On LocalAI, a strict tool switches on its tool grammar. OpenAI enforces
+the schema.
+
+```go
+result, err := cogito.ExecuteTools(llm, fragment,
+    cogito.WithMCPs(session),
+    cogito.EnableStrictToolSchemas)
+```
+
+Strict mode has rules of its own, so cogito adjusts each schema it sends:
+
+- Every object gets `"additionalProperties": false`.
+- Every property is required. A property that was optional becomes nullable,
+  and the model sends `null` to leave it out. cogito removes those `null`
+  values before it calls the tool, so the tool sees an omitted argument.
+
+A tool whose schema strict mode cannot express (an object with schema-valued
+or `true` `additionalProperties`, or a `oneOf`) is sent unchanged and not
+strict.
+
+MCP tools are advertised with the `additionalProperties` their server
+declares, with or without this option.
+
 ### Custom Tool Implementation
 
 See `examples/internal/search/search.go` for a complete example of implementing a DuckDuckGo search tool.
diff --git a/guidelines.go b/guidelines.go
index 543b8ee..eb09cc9 100644
--- a/guidelines.go
+++ b/guidelines.go
@@ -207,5 +207,9 @@ func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines,
 		}
 	}
 
+	if o.strictToolSchemas {
+		tools = strictTools(tools)
+	}
+
 	return tools, guidelines, prompts, nil
 }
diff --git a/mcp.go b/mcp.go
index 16a7ce6..2b95887 100644
--- a/mcp.go
+++ b/mcp.go
@@ -26,10 +26,11 @@ func (t *mcpTool) Tool() openai.Tool {
 		Function: &openai.FunctionDefinition{
 			Name:        t.name,
 			Description: t.description,
-			Parameters: jsonschema.Definition{
-				Type:       jsonschema.Object,
-				Properties: t.props,
-				Required:   t.inputSchema.Required,
+			Parameters: mcpToolParameters{
+				Type:                 jsonschema.Object,
+				Properties:           t.props,
+				Required:             t.inputSchema.Required,
+				AdditionalProperties: t.inputSchema.AdditionalProperties,
 			},
 		},
 	}
@@ -101,6 +102,22 @@ 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
diff --git a/options.go b/options.go
index 8340e73..4eda8d1 100644
--- a/options.go
+++ b/options.go
@@ -37,6 +37,7 @@ type Options struct {
 	maxAdjustmentAttempts             int
 	toolCallResultCallback            func(ToolStatus)
 	strictGuidelines                  bool
+	strictToolSchemas                 bool
 	mcpSessions                       []*mcp.ClientSession
 	guidelines                        Guidelines
 	mcpPrompts                        bool
@@ -158,6 +159,16 @@ var (
 		o.strictGuidelines = true
 	}
 
+	// EnableStrictToolSchemas sends every tool that can be expressed in strict
+	// mode with "strict": true and a schema adjusted to the rules strict mode
+	// needs (see strictTool). A backend that honors it constrains the model's
+	// arguments to the schema while it generates them, instead of the tool
+	// rejecting a malformed call afterwards. LocalAI switches its tool grammar
+	// on for a request with a strict tool; OpenAI enforces the schema.
+	EnableStrictToolSchemas Option = func(o *Options) {
+		o.strictToolSchemas = true
+	}
+
 	// EnableAutoPlan enables cogito to automatically use planning if needed
 	EnableAutoPlan Option = func(o *Options) {
 		o.autoPlan = true
diff --git a/strict.go b/strict.go
new file mode 100644
index 0000000..c0486da
--- /dev/null
+++ b/strict.go
@@ -0,0 +1,194 @@
+package cogito
+
+import (
+	"encoding/json"
+	"sort"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+// strictTool advertises inner with "strict": true and a schema that follows
+// the rules strict mode needs:
+//
+//   - every object sets additionalProperties to false;
+//   - every property is listed in required, so a property that was optional
+//     becomes nullable instead, and the model sends null to leave it out.
+//
+// Execute removes those nulls again before it calls inner, so inner sees an
+// omitted argument, as it would without strict mode. A schema that strict
+// mode cannot express (an object whose additionalProperties is a schema or
+// true, oneOf, a non-object top level) is sent unchanged and not strict.
+type strictTool struct {
+	inner ToolDefinitionInterface
+}
+
+// strictTools wraps each tool in strictTool, once.
+func strictTools(tools Tools) Tools {
+	out := make(Tools, len(tools))
+	for i, t := range tools {
+		if _, done := t.(strictTool); done {
+			out[i] = t
+			continue
+		}
+		out[i] = strictTool{inner: t}
+	}
+	return out
+}
+
+func (t strictTool) Tool() openai.Tool {
+	tool := t.inner.Tool()
+	if tool.Function == nil || tool.Function.Parameters == nil {
+		return tool
+	}
+	raw, err := json.Marshal(tool.Function.Parameters)
+	if err != nil {
+		return tool
+	}
+	var schema map[string]any
+	if err := json.Unmarshal(raw, &schema); err != nil {
+		return tool
+	}
+	if schema["type"] != "object" || !strictSchema(schema) {
+		return tool
+	}
+	fn := *tool.Function
+	fn.Parameters = schema
+	fn.Strict = true
+	tool.Function = &fn
+	return tool
+}
+
+func (t strictTool) Execute(args map[string]any) (string, any, error) {
+	return t.inner.Execute(dropNulls(args))
+}
+
+// strictSchema rewrites schema in place to follow strict mode's rules and
+// reports whether it could. On false the caller must discard schema.
+func strictSchema(schema map[string]any) bool {
+	if _, ok := schema["oneOf"]; ok {
+		return false
+	}
+	if isObjectSchema(schema) {
+		switch ap := schema["additionalProperties"].(type) {
+		case nil:
+		case bool:
+			if ap {
+				return false
+			}
+		default:
+			return false // a map type: strict mode cannot describe it
+		}
+		schema["additionalProperties"] = false
+
+		props, _ := schema["properties"].(map[string]any)
+		if props == nil {
+			props = map[string]any{}
+			schema["properties"] = props
+		}
+		required := map[string]bool{}
+		if rs, ok := schema["required"].([]any); ok {
+			for _, r := range rs {
+				if s, ok := r.(string); ok {
+					required[s] = true
+				}
+			}
+		}
+		names := make([]string, 0, len(props))
+		for name, p := range props {
+			ps, ok := p.(map[string]any)
+			if !ok {
+				return false // a boolean schema
+			}
+			if !strictSchema(ps) {
+				return false
+			}
+			if !required[name] {
+				props[name] = nullable(ps)
+			}
+			names = append(names, name)
+		}
+		sort.Strings(names)
+		req := make([]any, len(names))
+		for i, n := range names {
+			req[i] = n
+		}
+		schema["required"] = req
+	}
+	if items, ok := schema["items"].(map[string]any); ok && !strictSchema(items) {
+		return false
+	}
+	for _, key := range []string{"anyOf", "allOf"} {
+		if list, ok := schema[key].([]any); ok {
+			for _, s := range list {
+				if m, ok := s.(map[string]any); ok && !strictSchema(m) {
+					return false
+				}
+			}
+		}
+	}
+	for _, key := range []string{"$defs", "definitions"} {
+		if defs, ok := schema[key].(map[string]any); ok {
+			for _, d := range defs {
+				if m, ok := d.(map[string]any); ok && !strictSchema(m) {
+					return false
+				}
+			}
+		}
+	}
+	return true
+}
+
+func isObjectSchema(schema map[string]any) bool {
+	if schema["type"] == "object" {
+		return true
+	}
+	if ts, ok := schema["type"].([]any); ok {
+		for _, t := range ts {
+			if t == "object" {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// nullable returns schema extended to also accept null.
+func nullable(schema map[string]any) map[string]any {
+	switch t := schema["type"].(type) {
+	case string:
+		schema["type"] = []any{t, "null"}
+	case []any:
+		for _, x := range t {
+			if x == "null" {
+				return schema
+			}
+		}
+		schema["type"] = append(t, "null")
+	default:
+		return map[string]any{"anyOf": []any{schema, map[string]any{"type": "null"}}}
+	}
+	if enum, ok := schema["enum"].([]any); ok {
+		schema["enum"] = append(enum, nil)
+	}
+	return schema
+}
+
+// dropNulls returns args without the null values strict mode makes the
+// model send for an argument it leaves out, in nested objects too.
+func dropNulls(args map[string]any) map[string]any {
+	if args == nil {
+		return nil
+	}
+	out := make(map[string]any, len(args))
+	for k, v := range args {
+		switch x := v.(type) {
+		case nil:
+			continue
+		case map[string]any:
+			out[k] = dropNulls(x)
+		default:
+			out[k] = v
+		}
+	}
+	return out
+}
diff --git a/strict_test.go b/strict_test.go
new file mode 100644
index 0000000..240784b
--- /dev/null
+++ b/strict_test.go
@@ -0,0 +1,166 @@
+package cogito
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/sashabaranov/go-openai"
+)
+
+type shellArgs struct {
+	Script  string `json:"script" jsonschema:"the script to run"`
+	Timeout int    `json:"timeout,omitempty" jsonschema:"seconds"`
+}
+
+// startTypedMCP serves one tool whose schema the go-sdk infers from
+// shellArgs, as nib's bash tool does: script required, timeout optional,
+// additionalProperties false. The server validates every call against it.
+func startTypedMCP() (*mcpsdk.ClientSession, func()) {
+	impl := &mcpsdk.Implementation{Name: "typed", Version: "0.0.1"}
+	srv := mcpsdk.NewServer(impl, nil)
+	mcpsdk.AddTool(srv, &mcpsdk.Tool{Name: "bash", Description: "run a script"},
+		func(_ context.Context, _ *mcpsdk.CallToolRequest, in shellArgs) (*mcpsdk.CallToolResult, any, error) {
+			return &mcpsdk.CallToolResult{Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "ran " + in.Script}}}, nil, nil
+		})
+	srvT, clientT := mcpsdk.NewInMemoryTransports()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	go func() { _ = srv.Run(ctx, srvT) }()
+	sess, err := mcpsdk.NewClient(impl, nil).Connect(ctx, clientT, nil)
+	Expect(err).ToNot(HaveOccurred())
+	return sess, func() { _ = sess.Close(); cancel() }
+}
+
+func paramsOf(t openai.Tool) map[string]any {
+	raw, err := json.Marshal(t.Function.Parameters)
+	Expect(err).ToNot(HaveOccurred())
+	var m map[string]any
+	Expect(json.Unmarshal(raw, &m)).To(Succeed())
+	return m
+}
+
+var _ = Describe("strict tool schemas", func() {
+	var (
+		sess     *mcpsdk.ClientSession
+		teardown func()
+	)
+	AfterEach(func() {
+		if teardown != nil {
+			teardown()
+			teardown = nil
+		}
+	})
+
+	It("advertises an MCP tool with the additionalProperties its server declares", func() {
+		sess, teardown = startTypedMCP()
+		tools, err := mcpToolsFromTransport(context.Background(), sess, nil)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(tools).To(HaveLen(1))
+
+		p := paramsOf(tools[0].Tool())
+		Expect(p).To(HaveKeyWithValue("additionalProperties", false))
+		Expect(p["required"]).To(ConsistOf("script"))
+		Expect(tools[0].Tool().Function.Strict).To(BeFalse(), "strict is opt-in")
+	})
+
+	It("sends a strict schema and drops the nulls it makes the model send", func() {
+		sess, teardown = startTypedMCP()
+		tools, err := mcpToolsFromTransport(context.Background(), sess, nil)
+		Expect(err).ToNot(HaveOccurred())
+		tool := strictTools(tools)[0]
+
+		def := tool.Tool()
+		Expect(def.Function.Strict).To(BeTrue())
+		p := paramsOf(def)
+		Expect(p).To(HaveKeyWithValue("additionalProperties", false))
+		Expect(p["required"]).To(ConsistOf("script", "timeout"))
+		timeout := p["properties"].(map[string]any)["timeout"].(map[string]any)
+		Expect(timeout["type"]).To(ConsistOf("integer", "null"))
+
+		// The server rejects "timeout": null (it wants an integer), so this
+		// only succeeds if the wrapper leaves the null out.
+		out, _, err := tool.Execute(map[string]any{"script": "ls", "timeout": nil})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(out).To(ContainSubstring("ran ls"))
+	})
+
+	It("wraps the tools usableTools returns only when the option is set", func() {
+		sess, teardown = startTypedMCP()
+		plain, _, _, err := usableTools(nil, NewEmptyFragment(), WithMCPs(sess))
+		Expect(err).ToNot(HaveOccurred())
+		Expect(plain[0].Tool().Function.Strict).To(BeFalse())
+
+		strict, _, _, err := usableTools(nil, NewEmptyFragment(), WithMCPs(sess), EnableStrictToolSchemas)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(strict[0].Tool().Function.Strict).To(BeTrue())
+		Expect(strictTools(strict)[0]).To(Equal(strict[0]), "wrapping twice is a no-op")
+	})
+})
+
+var _ = Describe("strictSchema", func() {
+	parse := func(s string) map[string]any {
+		var m map[string]any
+		Expect(json.Unmarshal([]byte(s), &m)).To(Succeed())
+		return m
+	}
+
+	It("makes nested optional properties nullable and closes every object", func() {
+		m := parse(`{"type":"object","properties":{
+			"mode":{"type":"string","enum":["a","b"]},
+			"opts":{"type":"object","properties":{"deep":{"type":"boolean"}}},
+			"list":{"type":"array","items":{"type":"object","properties":{"k":{"type":"string"}},"required":["k"]}}
+		},"required":["opts"]}`)
+		Expect(strictSchema(m)).To(BeTrue())
+		Expect(m["required"]).To(ConsistOf("mode", "opts", "list"))
+
+		props := m["properties"].(map[string]any)
+		mode := props["mode"].(map[string]any)
+		Expect(mode["type"]).To(ConsistOf("string", "null"))
+		Expect(mode["enum"]).To(ContainElement(BeNil()))
+
+		opts := props["opts"].(map[string]any)
+		Expect(opts["type"]).To(Equal("object"), "a required property stays non-null")
+		Expect(opts).To(HaveKeyWithValue("additionalProperties", false))
+		deep := opts["properties"].(map[string]any)["deep"].(map[string]any)
+		Expect(deep["type"]).To(ConsistOf("boolean", "null"))
+
+		item := props["list"].(map[string]any)["items"].(map[string]any)
+		Expect(item).To(HaveKeyWithValue("additionalProperties", false))
+		Expect(item["required"]).To(ConsistOf("k"))
+	})
+
+	It("wraps an optional property without a type in anyOf with null", func() {
+		m := parse(`{"type":"object","properties":{"v":{"anyOf":[{"type":"string"},{"type":"integer"}]}}}`)
+		Expect(strictSchema(m)).To(BeTrue())
+		v := m["properties"].(map[string]any)["v"].(map[string]any)
+		Expect(v["anyOf"]).To(HaveLen(2))
+		Expect(v["anyOf"].([]any)[1]).To(Equal(map[string]any{"type": "null"}))
+	})
+
+	DescribeTable("refuses what strict mode cannot express",
+		func(schema string) {
+			Expect(strictSchema(parse(schema))).To(BeFalse())
+		},
+		Entry("a map-typed object", `{"type":"object","additionalProperties":{"type":"string"}}`),
+		Entry("open objects", `{"type":"object","additionalProperties":true}`),
+		Entry("oneOf", `{"type":"object","properties":{"x":{"oneOf":[{"type":"string"}]}}}`),
+	)
+
+	It("leaves such a tool unchanged and not strict", func() {
+		inner := &fakeParamsTool{params: map[string]any{"type": "object", "additionalProperties": map[string]any{"type": "string"}}}
+		def := strictTool{inner: inner}.Tool()
+		Expect(def.Function.Strict).To(BeFalse())
+		Expect(paramsOf(def)["additionalProperties"]).To(Equal(map[string]any{"type": "string"}))
+	})
+})
+
+type fakeParamsTool struct{ params any }
+
+func (f *fakeParamsTool) Tool() openai.Tool {
+	return openai.Tool{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{Name: "f", Parameters: f.params}}
+}
+func (f *fakeParamsTool) Execute(map[string]any) (string, any, error) { return "", nil, nil }