Python: added more complete parsing for mcp tool arguments (#2756)

* added more complete parsing for mcp tool arguments

* fixed mypy

* added nonlocal model counter, and some fixes

* fixes in naming logic

* extracted json parsing function, added parametrized test and checked coverage

python/packages/core/agent_framework/_mcp.py

@@ -1,6 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import json
 import logging
 import re
 import sys
@@ -19,9 +18,9 @@ from mcp.client.websocket import websocket_client
 from mcp.shared.context import RequestContext
 from mcp.shared.exceptions import McpError
 from mcp.shared.session import RequestResponder
-from pydantic import BaseModel, Field, create_model
+from pydantic import BaseModel, create_model
 
-from ._tools import AIFunction, HostedMCPSpecificApproval
+from ._tools import AIFunction, HostedMCPSpecificApproval, _build_pydantic_model_from_json_schema
 from ._types import (
     ChatMessage,
     Contents,
@@ -274,95 +273,26 @@ def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> type[BaseModel]:
     if not prompt.arguments:
         return create_model(f"{prompt.name}_input")
 
-    field_definitions: dict[str, Any] = {}
+    # Convert prompt arguments to JSON schema format
+    properties: dict[str, Any] = {}
+    required: list[str] = []
+
     for prompt_argument in prompt.arguments:
-        # For prompts, all arguments are typically required and string type
-        # unless specified otherwise in the prompt argument
-        python_type = str  # Default type for prompt arguments
-
-        # Create field definition for create_model
+        # For prompts, all arguments are typically string type unless specified otherwise
+        properties[prompt_argument.name] = {
+            "type": "string",
+            "description": prompt_argument.description if hasattr(prompt_argument, "description") else "",
+        }
         if prompt_argument.required:
-            field_definitions[prompt_argument.name] = (python_type, ...)
-        else:
-            field_definitions[prompt_argument.name] = (python_type, None)
+            required.append(prompt_argument.name)
 
-    return create_model(f"{prompt.name}_input", **field_definitions)
+    schema = {"properties": properties, "required": required}
+    return _build_pydantic_model_from_json_schema(prompt.name, schema)
 
 
 def _get_input_model_from_mcp_tool(tool: types.Tool) -> type[BaseModel]:
     """Creates a Pydantic model from a tools parameters."""
-    properties = tool.inputSchema.get("properties", None)
-    required = tool.inputSchema.get("required", [])
-    definitions = tool.inputSchema.get("$defs", {})
-
-    # Check if 'properties' is missing or not a dictionary
-    if not properties:
-        return create_model(f"{tool.name}_input")
-
-    def resolve_type(prop_details: dict[str, Any]) -> type:
-        """Resolve JSON Schema type to Python type, handling $ref."""
-        # Handle $ref by resolving the reference
-        if "$ref" in prop_details:
-            ref = prop_details["$ref"]
-            # Extract the reference path (e.g., "#/$defs/CustomerIdParam" -> "CustomerIdParam")
-            if ref.startswith("#/$defs/"):
-                def_name = ref.split("/")[-1]
-                if def_name in definitions:
-                    # Resolve the reference and use its type
-                    resolved = definitions[def_name]
-                    return resolve_type(resolved)
-            # If we can't resolve the ref, default to dict for safety
-            return dict
-
-        # Map JSON Schema types to Python types
-        json_type = prop_details.get("type", "string")
-        match json_type:
-            case "integer":
-                return int
-            case "number":
-                return float
-            case "boolean":
-                return bool
-            case "array":
-                return list
-            case "object":
-                return dict
-            case _:
-                return str  # default
-
-    field_definitions: dict[str, Any] = {}
-    for prop_name, prop_details in properties.items():
-        prop_details = json.loads(prop_details) if isinstance(prop_details, str) else prop_details
-
-        python_type = resolve_type(prop_details)
-        description = prop_details.get("description", "")
-
-        # Build field kwargs (description, array items schema, etc.)
-        field_kwargs: dict[str, Any] = {}
-        if description:
-            field_kwargs["description"] = description
-
-        # Preserve array items schema if present
-        if prop_details.get("type") == "array" and "items" in prop_details:
-            items_schema = prop_details["items"]
-            if items_schema and items_schema != {}:
-                field_kwargs["json_schema_extra"] = {"items": items_schema}
-
-        # Create field definition for create_model
-        if prop_name in required:
-            if field_kwargs:
-                field_definitions[prop_name] = (python_type, Field(**field_kwargs))
-            else:
-                field_definitions[prop_name] = (python_type, ...)
-        else:
-            default_value = prop_details.get("default", None)
-            field_kwargs["default"] = default_value
-            if field_kwargs and any(k != "default" for k in field_kwargs):
-                field_definitions[prop_name] = (python_type, Field(**field_kwargs))
-            else:
-                field_definitions[prop_name] = (python_type, default_value)
-
-    return create_model(f"{tool.name}_input", **field_definitions)
+    return _build_pydantic_model_from_json_schema(tool.name, tool.inputSchema)
 
 
 def _normalize_mcp_name(name: str) -> str:

python/packages/core/agent_framework/_tools.py

@@ -25,7 +25,6 @@ from typing import (
 
 from opentelemetry.metrics import Histogram
 from pydantic import AnyUrl, BaseModel, Field, ValidationError, create_model
-from pydantic.fields import FieldInfo
 
 from ._logging import get_logger
 from ._serialization import SerializationMixin
@@ -932,6 +931,151 @@ TYPE_MAPPING = {
 }
 
 
+def _build_pydantic_model_from_json_schema(
+    model_name: str,
+    schema: Mapping[str, Any],
+) -> type[BaseModel]:
+    """Creates a Pydantic model from JSON Schema with support for $refs, nested objects, and typed arrays.
+
+    Args:
+        model_name: The name of the model to be created.
+        schema: The JSON Schema definition (should contain 'properties', 'required', '$defs', etc.).
+
+    Returns:
+        The dynamically created Pydantic model class.
+    """
+    properties = schema.get("properties")
+    required = schema.get("required", [])
+    definitions = schema.get("$defs", {})
+
+    # Check if 'properties' is missing or not a dictionary
+    if not properties:
+        return create_model(f"{model_name}_input")
+
+    def _resolve_type(prop_details: dict[str, Any], parent_name: str = "") -> type:
+        """Resolve JSON Schema type to Python type, handling $ref, nested objects, and typed arrays.
+
+        Args:
+            prop_details: The JSON Schema property details
+            parent_name: Name to use for creating nested models (for uniqueness)
+
+        Returns:
+            Python type annotation (could be int, str, list[str], or a nested Pydantic model)
+        """
+        # Handle $ref by resolving the reference
+        if "$ref" in prop_details:
+            ref = prop_details["$ref"]
+            # Extract the reference path (e.g., "#/$defs/CustomerIdParam" -> "CustomerIdParam")
+            if ref.startswith("#/$defs/"):
+                def_name = ref.split("/")[-1]
+                if def_name in definitions:
+                    # Resolve the reference and use its type
+                    resolved = definitions[def_name]
+                    return _resolve_type(resolved, def_name)
+            # If we can't resolve the ref, default to dict for safety
+            return dict
+
+        # Map JSON Schema types to Python types
+        json_type = prop_details.get("type", "string")
+        match json_type:
+            case "integer":
+                return int
+            case "number":
+                return float
+            case "boolean":
+                return bool
+            case "array":
+                # Handle typed arrays
+                items_schema = prop_details.get("items")
+                if items_schema and isinstance(items_schema, dict):
+                    # Recursively resolve the item type
+                    item_type = _resolve_type(items_schema, f"{parent_name}_item")
+                    # Return list[ItemType] instead of bare list
+                    return list[item_type]  # type: ignore
+                # If no items schema or invalid, return bare list
+                return list
+            case "object":
+                # Handle nested objects by creating a nested Pydantic model
+                nested_properties = prop_details.get("properties")
+                nested_required = prop_details.get("required", [])
+
+                if nested_properties and isinstance(nested_properties, dict):
+                    # Create the name for the nested model
+                    nested_model_name = f"{parent_name}_nested" if parent_name else "NestedModel"
+
+                    # Recursively build field definitions for the nested model
+                    nested_field_definitions: dict[str, Any] = {}
+                    for nested_prop_name, nested_prop_details in nested_properties.items():
+                        nested_prop_details = (
+                            json.loads(nested_prop_details)
+                            if isinstance(nested_prop_details, str)
+                            else nested_prop_details
+                        )
+
+                        nested_python_type = _resolve_type(
+                            nested_prop_details, f"{nested_model_name}_{nested_prop_name}"
+                        )
+                        nested_description = nested_prop_details.get("description", "")
+
+                        # Build field kwargs for nested property
+                        nested_field_kwargs: dict[str, Any] = {}
+                        if nested_description:
+                            nested_field_kwargs["description"] = nested_description
+
+                        # Create field definition
+                        if nested_prop_name in nested_required:
+                            nested_field_definitions[nested_prop_name] = (
+                                (
+                                    nested_python_type,
+                                    Field(**nested_field_kwargs),
+                                )
+                                if nested_field_kwargs
+                                else (nested_python_type, ...)
+                            )
+                        else:
+                            nested_field_kwargs["default"] = nested_prop_details.get("default", None)
+                            nested_field_definitions[nested_prop_name] = (
+                                nested_python_type,
+                                Field(**nested_field_kwargs),
+                            )
+
+                    # Create and return the nested Pydantic model
+                    return create_model(nested_model_name, **nested_field_definitions)  # type: ignore
+
+                # If no properties defined, return bare dict
+                return dict
+            case _:
+                return str  # default
+
+    field_definitions: dict[str, Any] = {}
+    for prop_name, prop_details in properties.items():
+        prop_details = json.loads(prop_details) if isinstance(prop_details, str) else prop_details
+
+        python_type = _resolve_type(prop_details, f"{model_name}_{prop_name}")
+        description = prop_details.get("description", "")
+
+        # Build field kwargs (description, etc.)
+        field_kwargs: dict[str, Any] = {}
+        if description:
+            field_kwargs["description"] = description
+
+        # Create field definition for create_model
+        if prop_name in required:
+            if field_kwargs:
+                field_definitions[prop_name] = (python_type, Field(**field_kwargs))
+            else:
+                field_definitions[prop_name] = (python_type, ...)
+        else:
+            default_value = prop_details.get("default", None)
+            field_kwargs["default"] = default_value
+            if field_kwargs and any(k != "default" for k in field_kwargs):
+                field_definitions[prop_name] = (python_type, Field(**field_kwargs))
+            else:
+                field_definitions[prop_name] = (python_type, default_value)
+
+    return create_model(f"{model_name}_input", **field_definitions)
+
+
 def _create_model_from_json_schema(tool_name: str, schema_json: Mapping[str, Any]) -> type[BaseModel]:
     """Creates a Pydantic model from a given JSON Schema.
 
@@ -948,29 +1092,8 @@ def _create_model_from_json_schema(tool_name: str, schema_json: Mapping[str, Any
             f"JSON schema for tool '{tool_name}' must contain a 'properties' key of type dict. "
             f"Got: {schema_json.get('properties', None)}"
         )
-    # Extract field definitions with type annotations
-    field_definitions: dict[str, tuple[type, FieldInfo]] = {}
-    for field_name, field_schema in schema_json["properties"].items():
-        field_args: dict[str, Any] = {}
-        if (field_description := field_schema.get("description", None)) is not None:
-            field_args["description"] = field_description
-        if (field_default := field_schema.get("default", None)) is not None:
-            field_args["default"] = field_default
-        field_type = field_schema.get("type", None)
-        if field_type is None:
-            raise ValueError(
-                f"Missing 'type' for field '{field_name}' in JSON schema. "
-                f"Got: {field_schema}, Supported types: {list(TYPE_MAPPING.keys())}"
-            )
-        python_type = TYPE_MAPPING.get(field_type)
-        if python_type is None:
-            raise ValueError(
-                f"Unsupported type '{field_type}' for field '{field_name}' in JSON schema. "
-                f"Got: {field_schema}, Supported types: {list(TYPE_MAPPING.keys())}"
-            )
-        field_definitions[field_name] = (python_type, Field(**field_args))
 
-    return create_model(f"{tool_name}_input", **field_definitions)  # type: ignore[call-overload, no-any-return]
+    return _build_pydantic_model_from_json_schema(tool_name, schema_json)
 
 
 @overload

python/packages/core/pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     # connectors and functions
     "openai>=1.99.0",
     "azure-identity>=1,<2",
-    "mcp[ws]>=1.13",
+    "mcp[ws]>=1.23",
     "packaging>=24.1",
 ]
 

python/packages/core/tests/core/test_mcp.py

@@ -9,7 +9,7 @@ import pytest
 from mcp import types
 from mcp.client.session import ClientSession
 from mcp.shared.exceptions import McpError
-from pydantic import AnyUrl, ValidationError
+from pydantic import AnyUrl, BaseModel, ValidationError
 
 from agent_framework import (
     ChatMessage,
@@ -357,122 +357,360 @@ def test_chat_message_to_mcp_types():
     assert isinstance(mcp_contents[1], types.ImageContent)
 
 
-def test_get_input_model_from_mcp_tool():
-    """Test creation of input model from MCP tool."""
-    tool = types.Tool(
-        name="test_tool",
-        description="A test tool",
-        inputSchema={
-            "type": "object",
-            "properties": {"param1": {"type": "string"}, "param2": {"type": "number"}},
-            "required": ["param1"],
-        },
-    )
-    model = _get_input_model_from_mcp_tool(tool)
-
-    # Create an instance to verify the model works
-    instance = model(param1="test", param2=42)
-    assert instance.param1 == "test"
-    assert instance.param2 == 42
-
-    # Test validation
-    with pytest.raises(ValidationError):  # Missing required param1
-        model(param2=42)
-
-
-def test_get_input_model_from_mcp_tool_with_nested_object():
-    """Test creation of input model from MCP tool with nested object property."""
-    tool = types.Tool(
-        name="get_customer_detail",
-        description="Get customer details",
-        inputSchema={
-            "type": "object",
-            "properties": {
-                "params": {
-                    "type": "object",
-                    "properties": {"customer_id": {"type": "integer"}},
-                    "required": ["customer_id"],
+@pytest.mark.parametrize(
+    "test_id,input_schema,valid_data,expected_values,invalid_data,validation_check",
+    [
+        # Basic types with required/optional fields
+        (
+            "basic_types",
+            {
+                "type": "object",
+                "properties": {"param1": {"type": "string"}, "param2": {"type": "number"}},
+                "required": ["param1"],
+            },
+            {"param1": "test", "param2": 42},
+            {"param1": "test", "param2": 42},
+            {"param2": 42},  # Missing required param1
+            None,
+        ),
+        # Nested object
+        (
+            "nested_object",
+            {
+                "type": "object",
+                "properties": {
+                    "params": {
+                        "type": "object",
+                        "properties": {"customer_id": {"type": "integer"}},
+                        "required": ["customer_id"],
+                    }
+                },
+                "required": ["params"],
+            },
+            {"params": {"customer_id": 251}},
+            {"params.customer_id": 251},
+            {"params": {}},  # Missing required customer_id
+            lambda instance: isinstance(instance.params, BaseModel),
+        ),
+        # $ref resolution
+        (
+            "ref_schema",
+            {
+                "type": "object",
+                "properties": {"params": {"$ref": "#/$defs/CustomerIdParam"}},
+                "required": ["params"],
+                "$defs": {
+                    "CustomerIdParam": {
+                        "type": "object",
+                        "properties": {"customer_id": {"type": "integer"}},
+                        "required": ["customer_id"],
+                    }
+                },
+            },
+            {"params": {"customer_id": 251}},
+            {"params.customer_id": 251},
+            {"params": {}},  # Missing required customer_id
+            lambda instance: isinstance(instance.params, BaseModel),
+        ),
+        # Array of strings (typed)
+        (
+            "array_of_strings",
+            {
+                "type": "object",
+                "properties": {
+                    "tags": {
+                        "type": "array",
+                        "description": "List of tags",
+                        "items": {"type": "string"},
+                    }
+                },
+                "required": ["tags"],
+            },
+            {"tags": ["tag1", "tag2", "tag3"]},
+            {"tags": ["tag1", "tag2", "tag3"]},
+            None,  # No validation error test for this case
+            None,
+        ),
+        # Array of integers (typed)
+        (
+            "array_of_integers",
+            {
+                "type": "object",
+                "properties": {
+                    "numbers": {
+                        "type": "array",
+                        "description": "List of integers",
+                        "items": {"type": "integer"},
+                    }
+                },
+                "required": ["numbers"],
+            },
+            {"numbers": [1, 2, 3]},
+            {"numbers": [1, 2, 3]},
+            None,
+            None,
+        ),
+        # Array of objects (complex nested)
+        (
+            "array_of_objects",
+            {
+                "type": "object",
+                "properties": {
+                    "users": {
+                        "type": "array",
+                        "description": "List of users",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "id": {"type": "integer", "description": "User ID"},
+                                "name": {"type": "string", "description": "User name"},
+                            },
+                            "required": ["id", "name"],
+                        },
+                    }
+                },
+                "required": ["users"],
+            },
+            {"users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]},
+            {"users[0].id": 1, "users[0].name": "Alice", "users[1].id": 2, "users[1].name": "Bob"},
+            {"users": [{"id": 1}]},  # Missing required 'name'
+            lambda instance: all(isinstance(user, BaseModel) for user in instance.users),
+        ),
+        # Deeply nested objects (3+ levels)
+        (
+            "deeply_nested",
+            {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "object",
+                        "properties": {
+                            "filters": {
+                                "type": "object",
+                                "properties": {
+                                    "date_range": {
+                                        "type": "object",
+                                        "properties": {
+                                            "start": {"type": "string"},
+                                            "end": {"type": "string"},
+                                        },
+                                        "required": ["start", "end"],
+                                    },
+                                    "categories": {"type": "array", "items": {"type": "string"}},
+                                },
+                                "required": ["date_range"],
+                            }
+                        },
+                        "required": ["filters"],
+                    }
+                },
+                "required": ["query"],
+            },
+            {
+                "query": {
+                    "filters": {
+                        "date_range": {"start": "2024-01-01", "end": "2024-12-31"},
+                        "categories": ["tech", "science"],
+                    }
                 }
             },
-            "required": ["params"],
-        },
-    )
-    model = _get_input_model_from_mcp_tool(tool)
+            {
+                "query.filters.date_range.start": "2024-01-01",
+                "query.filters.date_range.end": "2024-12-31",
+                "query.filters.categories": ["tech", "science"],
+            },
+            {"query": {"filters": {"date_range": {}}}},  # Missing required start and end
+            None,
+        ),
+        # Complex $ref with nested structure
+        (
+            "ref_nested_structure",
+            {
+                "type": "object",
+                "properties": {"order": {"$ref": "#/$defs/OrderParams"}},
+                "required": ["order"],
+                "$defs": {
+                    "OrderParams": {
+                        "type": "object",
+                        "properties": {
+                            "customer": {"$ref": "#/$defs/Customer"},
+                            "items": {"type": "array", "items": {"$ref": "#/$defs/OrderItem"}},
+                        },
+                        "required": ["customer", "items"],
+                    },
+                    "Customer": {
+                        "type": "object",
+                        "properties": {"id": {"type": "integer"}, "email": {"type": "string"}},
+                        "required": ["id", "email"],
+                    },
+                    "OrderItem": {
+                        "type": "object",
+                        "properties": {"product_id": {"type": "string"}, "quantity": {"type": "integer"}},
+                        "required": ["product_id", "quantity"],
+                    },
+                },
+            },
+            {
+                "order": {
+                    "customer": {"id": 123, "email": "test@example.com"},
+                    "items": [{"product_id": "prod1", "quantity": 2}],
+                }
+            },
+            {
+                "order.customer.id": 123,
+                "order.customer.email": "test@example.com",
+                "order.items[0].product_id": "prod1",
+                "order.items[0].quantity": 2,
+            },
+            {"order": {"customer": {"id": 123}, "items": []}},  # Missing email
+            lambda instance: isinstance(instance.order.customer, BaseModel),
+        ),
+        # Mixed types (primitives, arrays, nested objects)
+        (
+            "mixed_types",
+            {
+                "type": "object",
+                "properties": {
+                    "simple_string": {"type": "string"},
+                    "simple_number": {"type": "integer"},
+                    "string_array": {"type": "array", "items": {"type": "string"}},
+                    "nested_config": {
+                        "type": "object",
+                        "properties": {
+                            "enabled": {"type": "boolean"},
+                            "options": {"type": "array", "items": {"type": "string"}},
+                        },
+                        "required": ["enabled"],
+                    },
+                },
+                "required": ["simple_string", "nested_config"],
+            },
+            {
+                "simple_string": "test",
+                "simple_number": 42,
+                "string_array": ["a", "b"],
+                "nested_config": {"enabled": True, "options": ["opt1", "opt2"]},
+            },
+            {
+                "simple_string": "test",
+                "simple_number": 42,
+                "string_array": ["a", "b"],
+                "nested_config.enabled": True,
+                "nested_config.options": ["opt1", "opt2"],
+            },
+            None,
+            None,
+        ),
+        # Empty schema (no properties)
+        (
+            "empty_schema",
+            {"type": "object", "properties": {}},
+            {},
+            {},
+            None,
+            None,
+        ),
+        # All primitive types
+        (
+            "all_primitives",
+            {
+                "type": "object",
+                "properties": {
+                    "string_field": {"type": "string"},
+                    "integer_field": {"type": "integer"},
+                    "number_field": {"type": "number"},
+                    "boolean_field": {"type": "boolean"},
+                },
+            },
+            {"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True},
+            {"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True},
+            None,
+            None,
+        ),
+        # Edge case: unresolvable $ref (fallback to dict)
+        (
+            "unresolvable_ref",
+            {
+                "type": "object",
+                "properties": {"data": {"$ref": "#/$defs/NonExistent"}},
+                "$defs": {},
+            },
+            {"data": {"key": "value"}},
+            {"data": {"key": "value"}},
+            None,
+            None,
+        ),
+        # Edge case: array without items schema (fallback to bare list)
+        (
+            "array_no_items",
+            {
+                "type": "object",
+                "properties": {"items": {"type": "array"}},
+            },
+            {"items": [1, "two", 3.0]},
+            {"items": [1, "two", 3.0]},
+            None,
+            None,
+        ),
+        # Edge case: object without properties (fallback to dict)
+        (
+            "object_no_properties",
+            {
+                "type": "object",
+                "properties": {"config": {"type": "object"}},
+            },
+            {"config": {"arbitrary": "data", "nested": {"key": "value"}}},
+            {"config": {"arbitrary": "data", "nested": {"key": "value"}}},
+            None,
+            None,
+        ),
+    ],
+)
+def test_get_input_model_from_mcp_tool_parametrized(
+    test_id, input_schema, valid_data, expected_values, invalid_data, validation_check
+):
+    """Parametrized test for JSON schema to Pydantic model conversion.
 
-    # Create an instance to verify the model works with nested objects
-    instance = model(params={"customer_id": 251})
-    assert instance.params == {"customer_id": 251}
-    assert isinstance(instance.params, dict)
+    This test covers various edge cases including:
+    - Basic types with required/optional fields
+    - Nested objects
+    - $ref resolution
+    - Typed arrays (strings, integers, objects)
+    - Deeply nested structures
+    - Complex $ref with nested structures
+    - Mixed types
 
-    # Verify model_dump produces the correct nested structure
-    dumped = instance.model_dump()
-    assert dumped == {"params": {"customer_id": 251}}
-
-
-def test_get_input_model_from_mcp_tool_with_ref_schema():
-    """Test creation of input model from MCP tool with $ref schema.
-
-    This simulates a FastMCP tool that uses Pydantic models with $ref in the schema.
-    The schema should be resolved and nested objects should be preserved.
+    To add a new test case, add a tuple to the parametrize decorator with:
+    - test_id: A descriptive name for the test case
+    - input_schema: The JSON schema (inputSchema dict)
+    - valid_data: Valid data to instantiate the model
+    - expected_values: Dict of expected values (supports dot notation for nested access)
+    - invalid_data: Invalid data to test validation errors (None to skip)
+    - validation_check: Optional callable to perform additional validation checks
     """
-    # This is similar to what FastMCP generates when you have:
-    # async def get_customer_detail(params: CustomerIdParam) -> CustomerDetail
-    tool = types.Tool(
-        name="get_customer_detail",
-        description="Get customer details",
-        inputSchema={
-            "type": "object",
-            "properties": {"params": {"$ref": "#/$defs/CustomerIdParam"}},
-            "required": ["params"],
-            "$defs": {
-                "CustomerIdParam": {
-                    "type": "object",
-                    "properties": {"customer_id": {"type": "integer"}},
-                    "required": ["customer_id"],
-                }
-            },
-        },
-    )
+    tool = types.Tool(name="test_tool", description="A test tool", inputSchema=input_schema)
     model = _get_input_model_from_mcp_tool(tool)
 
-    # Create an instance to verify the model works with $ref schemas
-    instance = model(params={"customer_id": 251})
-    assert instance.params == {"customer_id": 251}
-    assert isinstance(instance.params, dict)
+    # Test valid data
+    instance = model(**valid_data)
 
-    # Verify model_dump produces the correct nested structure
-    dumped = instance.model_dump()
-    assert dumped == {"params": {"customer_id": 251}}
+    # Check expected values
+    for field_path, expected_value in expected_values.items():
+        # Support dot notation and array indexing for nested access
+        current = instance
+        parts = field_path.replace("]", "").replace("[", ".").split(".")
+        for part in parts:
+            current = current[int(part)] if part.isdigit() else getattr(current, part)
+        assert current == expected_value, f"Field {field_path} = {current}, expected {expected_value}"
 
+    # Run additional validation checks if provided
+    if validation_check:
+        assert validation_check(instance), f"Validation check failed for {test_id}"
 
-def test_get_input_model_from_mcp_tool_with_simple_array():
-    """Test array with simple items schema (items schema should be preserved in json_schema_extra)."""
-    tool = types.Tool(
-        name="simple_array_tool",
-        description="Tool with simple array",
-        inputSchema={
-            "type": "object",
-            "properties": {
-                "tags": {
-                    "type": "array",
-                    "description": "List of tags",
-                    "items": {"type": "string"},  # Simple string array
-                }
-            },
-            "required": ["tags"],
-        },
-    )
-    model = _get_input_model_from_mcp_tool(tool)
-
-    # Create an instance
-    instance = model(tags=["tag1", "tag2", "tag3"])
-    assert instance.tags == ["tag1", "tag2", "tag3"]
-
-    # Verify JSON schema still preserves items for simple types
-    json_schema = model.model_json_schema()
-    tags_property = json_schema["properties"]["tags"]
-    assert "items" in tags_property
-    assert tags_property["items"]["type"] == "string"
+    # Test invalid data if provided
+    if invalid_data is not None:
+        with pytest.raises(ValidationError):
+            model(**invalid_data)
 
 
 def test_get_input_model_from_mcp_prompt():