Inline the tool schema generator into func_metadata

func_metadata already defined StrictJsonSchema; keep the SEP-2106 external-$ref
rejection there instead of a separate _schema_generator module.

Author: Kludex

src/mcp/server/mcpserver/tools/base.py

@@ -7,9 +7,8 @@
 from pydantic import BaseModel, Field
 
 from mcp.server.mcpserver.exceptions import ToolError
-from mcp.server.mcpserver.utilities._schema_generator import StrictJsonSchema
 from mcp.server.mcpserver.utilities.context_injection import find_context_parameter
-from mcp.server.mcpserver.utilities.func_metadata import FuncMetadata, func_metadata
+from mcp.server.mcpserver.utilities.func_metadata import FuncMetadata, StrictJsonSchema, func_metadata
 from mcp.shared._callable_inspection import is_async_callable
 from mcp.shared.exceptions import UrlElicitationRequiredError
 from mcp.shared.tool_name_validation import validate_and_warn_tool_name

src/mcp/server/mcpserver/utilities/_schema_generator.py

@@ -11,6 +11,8 @@
 import pydantic_core
 from pydantic import BaseModel, ConfigDict, Field, PydanticUserError, WithJsonSchema, create_model
 from pydantic.fields import FieldInfo
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue, JsonSchemaWarningKind
+from pydantic_core import CoreSchema
 from typing_extensions import is_typeddict
 from typing_inspection.introspection import (
     UNKNOWN,
@@ -21,14 +23,58 @@
 )
 
 from mcp.server.mcpserver.exceptions import InvalidSignature
-from mcp.server.mcpserver.utilities._schema_generator import ExternalSchemaRefError, StrictJsonSchema
 from mcp.server.mcpserver.utilities.logging import get_logger
 from mcp.server.mcpserver.utilities.types import Audio, Image
 from mcp.types import CallToolResult, ContentBlock, TextContent
 
 logger = get_logger(__name__)
 
 
+class ExternalSchemaRefError(ValueError):
+    """A tool schema contains a `$ref` that is not a same-document reference."""
+
+
+class StrictJsonSchema(GenerateJsonSchema):
+    """Render tool schemas, raising on pydantic warnings and external `$ref`s.
+
+    Warnings (e.g. a non-serializable type) become errors so they surface at tool
+    registration instead of silently producing a degenerate schema. External
+    `$ref`s -- which pydantic never emits itself, but a user can inject via
+    `Field(json_schema_extra=...)` -- are an SSRF / fetch-DoS vector and are
+    rejected for the same reason (SEP-2106).
+
+    See: https://modelcontextprotocol.io/seps/2106-json-schema-2020-12#security-implications
+    """
+
+    def emit_warning(self, kind: JsonSchemaWarningKind, detail: str) -> None:
+        raise ValueError(f"JSON schema warning: {kind} - {detail}")
+
+    def generate(self, schema: CoreSchema, mode: Any = "validation") -> JsonSchemaValue:
+        json_schema = super().generate(schema, mode)
+        external = sorted(_find_external_refs(json_schema))
+        if external:
+            raise ExternalSchemaRefError(
+                f"Tool schema contains external $ref(s) that MUST NOT be dereferenced (SEP-2106): "
+                f"{', '.join(external)}. Only same-document references (e.g. '#/$defs/Foo') are allowed."
+            )
+        return json_schema
+
+
+def _find_external_refs(node: Any) -> set[str]:
+    external: set[str] = set()
+    if isinstance(node, dict):
+        mapping = cast("dict[str, Any]", node)
+        ref = mapping.get("$ref")
+        if isinstance(ref, str) and not ref.startswith("#"):
+            external.add(ref)
+        for value in mapping.values():
+            external |= _find_external_refs(value)
+    elif isinstance(node, list):
+        for item in cast("list[Any]", node):
+            external |= _find_external_refs(item)
+    return external
+
+
 class ArgModelBase(BaseModel):
     """A model representing the arguments to a function."""
 

src/mcp/server/mcpserver/utilities/func_metadata.py

@@ -13,7 +13,7 @@
 from pydantic import BaseModel, Field
 
 from mcp.server.mcpserver.exceptions import InvalidSignature
-from mcp.server.mcpserver.utilities.func_metadata import func_metadata
+from mcp.server.mcpserver.utilities.func_metadata import ExternalSchemaRefError, StrictJsonSchema, func_metadata
 from mcp.types import CallToolResult
 
 
@@ -1191,3 +1191,34 @@ def func_with_metadata() -> Annotated[int, Field(gt=1)]: ...  # pragma: no branc
 
     assert meta.output_schema is not None
     assert meta.output_schema["properties"]["result"] == {"exclusiveMinimum": 1, "title": "Result", "type": "integer"}
+
+
+def test_strict_json_schema_allows_same_document_refs():
+    class Inner(BaseModel):
+        x: int
+
+    class Model(BaseModel):
+        inner: Inner
+
+    schema = Model.model_json_schema(schema_generator=StrictJsonSchema)
+    assert "$defs" in schema
+    assert schema["properties"]["inner"]["$ref"] == "#/$defs/Inner"
+
+
+def test_strict_json_schema_rejects_external_ref_in_property():
+    class Model(BaseModel):
+        profile: Annotated[dict[str, Any], Field(json_schema_extra={"$ref": "https://evil.example/s.json"})]
+
+    with pytest.raises(ExternalSchemaRefError, match="https://evil.example/s.json"):
+        Model.model_json_schema(schema_generator=StrictJsonSchema)
+
+
+def test_strict_json_schema_rejects_external_ref_nested_in_list():
+    class Model(BaseModel):
+        items: Annotated[
+            list[str],
+            Field(json_schema_extra={"prefixItems": [{"$ref": "https://evil.example/a.json"}]}),
+        ]
+
+    with pytest.raises(ExternalSchemaRefError, match="https://evil.example/a.json"):
+        Model.model_json_schema(schema_generator=StrictJsonSchema)

tests/server/mcpserver/test_func_metadata.py

@@ -10,8 +10,7 @@
 from mcp.server.mcpserver import Context, MCPServer
 from mcp.server.mcpserver.exceptions import ToolError
 from mcp.server.mcpserver.tools import Tool, ToolManager
-from mcp.server.mcpserver.utilities._schema_generator import ExternalSchemaRefError
-from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, FuncMetadata
+from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, ExternalSchemaRefError, FuncMetadata
 from mcp.types import CallToolResult, TextContent, ToolAnnotations