fix(tools): dereference draft-07 definitions in MCP tool schemas

Merge google#5941

Co-authored-by: Xuan Yang <xygoogle@google.com>
COPYBARA_INTEGRATE_REVIEW=google#5941 from gaoflow:fix-5940-draft07-definitions 4e56be5
PiperOrigin-RevId: 929433469

Author: 2 people

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
 
 

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 = {