From b2d151abe2494cb049b7b2f479f55613b57259ba Mon Sep 17 00:00:00 2001
From: Brooklyn Nicholson <brooklyn.bb.nicholson@gmail.com>
Date: Fri, 12 Jun 2026 00:30:51 -0500
Subject: [PATCH] fix(tools): strip default from $ref nodes in tool schemas

Fireworks-hosted Kimi rejects tool requests when nullable MCP/Pydantic
schemas collapse to {"$ref": "...", "default": null}. Strip that sibling
during global schema sanitization so gateway and CLI calls succeed again.
---
 tests/tools/test_schema_sanitizer.py | 66 ++++++++++++++++++++++++++++
 tools/schema_sanitizer.py            | 38 ++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/tests/tools/test_schema_sanitizer.py b/tests/tools/test_schema_sanitizer.py
index b856440ef40..360457de2b2 100644
--- a/tests/tools/test_schema_sanitizer.py
+++ b/tests/tools/test_schema_sanitizer.py
@@ -201,6 +201,72 @@ def test_items_sanitized_in_array_schema():
     assert items == {"type": "object", "properties": {}}
 
 
+def test_ref_with_default_sibling_stripped():
+    """Strict backends reject ``default`` alongside ``$ref``."""
+    tools = [_tool("t", {
+        "type": "object",
+        "properties": {
+            "payload": {"$ref": "#/$defs/Payload", "default": None},
+        },
+        "$defs": {
+            "Payload": {
+                "type": "object",
+                "properties": {"q": {"type": "string"}},
+            },
+        },
+    })]
+    out = sanitize_tool_schemas(tools)
+    payload = out[0]["function"]["parameters"]["properties"]["payload"]
+    assert payload == {"$ref": "#/$defs/Payload"}
+
+
+def test_nullable_union_collapse_does_not_leave_default_on_ref():
+    """Nullable anyOf collapse must not attach ``default`` to a ``$ref`` branch."""
+    tools = [_tool("t", {
+        "type": "object",
+        "properties": {
+            "input": {
+                "anyOf": [
+                    {"$ref": "#/$defs/Payload"},
+                    {"type": "null"},
+                ],
+                "default": None,
+            },
+        },
+        "$defs": {
+            "Payload": {
+                "type": "object",
+                "properties": {"q": {"type": "string"}},
+            },
+        },
+    })]
+    out = sanitize_tool_schemas(tools)
+    prop = out[0]["function"]["parameters"]["properties"]["input"]
+    assert prop["$ref"] == "#/$defs/Payload"
+    assert "default" not in prop
+    assert prop.get("nullable") is True
+
+
+def test_ref_description_preserved():
+    """Annotation siblings that strict backends allow should survive."""
+    tools = [_tool("t", {
+        "type": "object",
+        "properties": {
+            "payload": {
+                "$ref": "#/$defs/Payload",
+                "description": "The payload",
+            },
+        },
+        "$defs": {
+            "Payload": {"type": "object", "properties": {}},
+        },
+    })]
+    out = sanitize_tool_schemas(tools)
+    payload = out[0]["function"]["parameters"]["properties"]["payload"]
+    assert payload["description"] == "The payload"
+    assert payload["$ref"] == "#/$defs/Payload"
+
+
 def test_empty_tools_list_returns_empty():
     assert sanitize_tool_schemas([]) == []
 
diff --git a/tools/schema_sanitizer.py b/tools/schema_sanitizer.py
index e9677ac4a1b..a2a5892d927 100644
--- a/tools/schema_sanitizer.py
+++ b/tools/schema_sanitizer.py
@@ -21,6 +21,12 @@ The failure modes we've seen in the wild:
   optional fields (common Pydantic/MCP shape). Anthropic rejects these at
   the top of ``input_schema``; collapse them to the non-null branch.
 * Unconstrained ``additionalProperties`` on objects with empty properties.
+* ``default`` (and other annotation keywords) alongside ``$ref`` — strict
+  backends (Fireworks-hosted Kimi, JSON Schema draft-07 validators) reject
+  sibling keywords at the same level as ``$ref``.  Common MCP/Pydantic shape
+  after nullable-union collapse::
+
+      {"$ref": "#/$defs/Foo", "default": null}
 
 This module walks the final tool schema tree (after MCP-level normalization
 and any per-tool dynamic rebuilds) and fixes the known-hostile constructs
@@ -90,6 +96,35 @@ def _sanitize_single_tool(tool: dict) -> dict:
     fn["parameters"] = _strip_top_level_combinators(
         fn["parameters"], path=fn.get("name", "<tool>")
     )
+    fn["parameters"] = _strip_ref_siblings(fn["parameters"])
+    return out
+
+
+# Sibling keywords strict JSON Schema validators reject alongside ``$ref``.
+_REF_FORBIDDEN_SIBLINGS = frozenset({"default"})
+
+
+def _strip_ref_siblings(node: Any) -> Any:
+    """Drop forbidden sibling keywords from nodes that carry ``$ref``.
+
+    Fireworks (and other draft-07-strict backends) fail tool requests with::
+
+        JSON Schema not supported: keyword(s) ['default'] not allowed at
+        the same level as $ref.
+
+    Nullable-union collapse and MCP ingestion can leave ``default`` on a
+    ``$ref`` node; strip it recursively.
+    """
+    if isinstance(node, list):
+        return [_strip_ref_siblings(item) for item in node]
+    if not isinstance(node, dict):
+        return node
+
+    out = {key: _strip_ref_siblings(value) for key, value in node.items()}
+    if "$ref" in out:
+        for key in _REF_FORBIDDEN_SIBLINGS:
+            if key in out:
+                out.pop(key, None)
     return out
 
 
@@ -185,6 +220,9 @@ def strip_nullable_unions(
                 replacement.setdefault("nullable", True)
             for meta_key in ("title", "description", "default", "examples"):
                 if meta_key in stripped and meta_key not in replacement:
+                    # ``default`` is illegal alongside ``$ref`` on strict backends.
+                    if meta_key == "default" and "$ref" in replacement:
+                        continue
                     replacement[meta_key] = stripped[meta_key]
             return strip_nullable_unions(replacement, keep_nullable_hint=keep_nullable_hint)
     return stripped