From 73ba261b395d611ca713543194dbe1525608f92b Mon Sep 17 00:00:00 2001
From: gaoflow <gaobing1230@gmail.com>
Date: Tue, 2 Jun 2026 19:47:31 +0200
Subject: [PATCH] fix(tools): dereference draft-07 `definitions` in MCP tool
 schemas
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`_dereference_schema` only resolved `$ref`s against `$defs` (JSON Schema
draft 2019-09+/2020-12) and only stripped the `$defs` block afterwards. A
tool whose `inputSchema` uses JSON Schema draft-07 — `definitions` plus
`$ref: "#/definitions/..."`, which the MCP specification explicitly
permits — was therefore left with its refs unresolved and the
`definitions` block intact, so `_to_gemini_schema` raised
`KeyError: 'definitions'` while building the Gemini schema.

Resolve refs against both `definitions` and `$defs` (the latter wins on a
key collision) and strip both blocks once resolved.

Fixes #5940.
---
 src/google/adk/tools/_gemini_schema_util.py   | 15 +++--
 .../tools/test_gemini_schema_util.py          | 60 +++++++++++++++++++
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/src/google/adk/tools/_gemini_schema_util.py b/src/google/adk/tools/_gemini_schema_util.py
index 08e8d4e6c1..6935a118b7 100644
--- a/src/google/adk/tools/_gemini_schema_util.py
+++ b/src/google/adk/tools/_gemini_schema_util.py
@@ -106,7 +106,12 @@ def _sanitize_schema_type(
 def _dereference_schema(schema: dict[str, Any]) -> dict[str, Any]:
   """Resolves $ref pointers in a JSON schema."""
 
-  defs = schema.get("$defs", {})
+  # Support both the draft 2019-09+/2020-12 keyword (`$defs`) and the
+  # draft-07 keyword (`definitions`). The MCP specification allows tool
+  # `inputSchema`s to use either, so a server sending draft-07 schemas with
+  # `definitions` + `$ref: "#/definitions/..."` must dereference correctly.
+  # `$defs` takes precedence on the (pathological) key collision.
+  defs = {**schema.get("definitions", {}), **schema.get("$defs", {})}
 
   def _resolve_refs(sub_schema: Any, path_refs: frozenset[str]) -> Any:
     if isinstance(sub_schema, dict):
@@ -148,9 +153,11 @@ def _resolve_refs(sub_schema: Any, path_refs: frozenset[str]) -> Any:
       return sub_schema
 
   dereferenced_schema = _resolve_refs(schema, frozenset())
-  # Remove the definitions block after resolving.
-  if "$defs" in dereferenced_schema:
-    del dereferenced_schema["$defs"]
+  # Remove the definition blocks after resolving so the leftover keywords do
+  # not leak into the Gemini schema (which would otherwise raise a KeyError).
+  for defs_keyword in ("$defs", "definitions"):
+    if defs_keyword in dereferenced_schema:
+      del dereferenced_schema[defs_keyword]
   return dereferenced_schema
 
 
diff --git a/tests/unittests/tools/test_gemini_schema_util.py b/tests/unittests/tools/test_gemini_schema_util.py
index d919172527..6aaa4ddea0 100644
--- a/tests/unittests/tools/test_gemini_schema_util.py
+++ b/tests/unittests/tools/test_gemini_schema_util.py
@@ -337,6 +337,66 @@ def test_to_gemini_schema_nested_dict_with_defs_and_ref(self):
     ]
     assert gemini_schema.properties["payload"].required == ["adDomain"]
 
+  def test_to_gemini_schema_draft_07_definitions_and_ref(self):
+    """Draft-07 schemas use `definitions`/`#/definitions/...` instead of `$defs`.
+
+    The MCP spec allows tool `inputSchema`s to use JSON Schema draft-07, so a
+    server sending `definitions` + `$ref: "#/definitions/..."` must dereference
+    correctly instead of raising `KeyError: 'definitions'`.
+    """
+    openapi_schema = {
+        "$schema": "http://json-schema.org/draft-07/schema#",
+        "definitions": {
+            "DeviceEnum": {
+                "enum": ["GLOBAL", "desktop", "mobile"],
+                "title": "DeviceEnum",
+                "type": "string",
+            },
+            "DomainPayload": {
+                "properties": {
+                    "adDomain": {
+                        "description": "List of one or many domains.",
+                        "items": {"type": "string"},
+                        "title": "Addomain",
+                        "type": "array",
+                    },
+                    "device": {
+                        "$ref": "#/definitions/DeviceEnum",
+                        "default": "GLOBAL",
+                    },
+                },
+                "required": ["adDomain"],
+                "title": "DomainPayload",
+                "type": "object",
+            },
+        },
+        "properties": {"payload": {"$ref": "#/definitions/DomainPayload"}},
+        "required": ["payload"],
+        "title": "query_domainsArguments",
+        "type": "object",
+    }
+    gemini_schema = _to_gemini_schema(openapi_schema)
+    assert gemini_schema.type == Type.OBJECT
+    assert gemini_schema.properties["payload"].type == Type.OBJECT
+    assert (
+        gemini_schema.properties["payload"].properties["adDomain"].type
+        == Type.ARRAY
+    )
+    assert (
+        gemini_schema.properties["payload"].properties["adDomain"].items.type
+        == Type.STRING
+    )
+    assert (
+        gemini_schema.properties["payload"].properties["device"].type
+        == Type.STRING
+    )
+    assert gemini_schema.properties["payload"].properties["device"].enum == [
+        "GLOBAL",
+        "desktop",
+        "mobile",
+    ]
+    assert gemini_schema.properties["payload"].required == ["adDomain"]
+
   def test_sanitize_integer_formats(self):
     """Test that int32 and int64 formats are preserved for integer types"""
     openapi_schema = {