Add comprehensive unit tests for JSON argument conversion

- Add TestJSONArgumentConversion class with 9 comprehensive tests
- Test basic dict conversion, Optional/Union types, typed dicts
- Test error handling for invalid JSON strings
- Test async and sync function wrappers
- Test complex nested JSON structures
- Test annotation modification behavior
- Test that non-dict parameters remain unchanged

- Add TestJSONSchemaModification class with schema tests
- Test tool registration with schema modification
- Test multiple dict parameters in single function

- Fix union type handling for Python 3.10+ syntax (dict | None)
- Support both typing.Union and new UnionType for compatibility
- Update schema modification to handle both union syntaxes

All 44 tests passing, including 11 new JSON conversion tests

Author: Zsailer

jupyter_server_mcp/mcp_server.py

@@ -47,12 +47,18 @@ def _should_convert_to_dict(annotation, value):
         if annotation is dict:
             return True
 
-        # Optional[dict] or Union[dict, None] etc.
+        # Optional[dict] or Union[dict, None] etc. (old typing.Union)
         origin = get_origin(annotation)
         if origin is Union:
             args = get_args(annotation)
             return dict in args
 
+        # New Python 3.10+ union syntax: dict | None
+        if hasattr(annotation, '__class__') and annotation.__class__.__name__ == 'UnionType':
+            # For dict | None style unions
+            args = get_args(annotation)
+            return dict in args
+            
         # Dict[K, V] style annotations
         return bool(hasattr(annotation, '__origin__') and annotation.__origin__ is dict)
 
@@ -62,7 +68,7 @@ def _modify_annotation_for_string_support(annotation):
         if annotation is dict:
             return dict | str
 
-        # Optional[dict] or Union[dict, None] etc.
+        # Optional[dict] or Union[dict, None] etc. (old typing.Union)
         origin = get_origin(annotation)
         if origin is Union:
             args = get_args(annotation)
@@ -71,6 +77,21 @@ def _modify_annotation_for_string_support(annotation):
                 if str not in args:
                     return Union[(*tuple(args), str)]
                 return annotation
+        
+        # New Python 3.10+ union syntax: dict | None
+        if hasattr(annotation, '__class__') and annotation.__class__.__name__ == 'UnionType':
+            args = get_args(annotation)
+            if dict in args:
+                # Add str to the union if it's not already there
+                if str not in args:
+                    # Reconstruct the union with str added
+                    new_args = (*tuple(args), str)
+                    # Create new union type
+                    result = new_args[0]
+                    for arg in new_args[1:]:
+                        result = result | arg
+                    return result
+                return annotation
 
         # Dict[K, V] style annotations -> annotation | str
         if hasattr(annotation, '__origin__') and annotation.__origin__ is dict:
@@ -175,8 +196,8 @@ def _modify_schema_for_json_string_support(func: Callable, tool) -> None:
                 # Direct dict annotation
                 if annotation is dict:
                     should_support_string = True
-                # Optional[dict] or Union[dict, None] etc.
-                elif get_origin(annotation) is Union:
+                # Optional[dict] or Union[dict, None] etc. (old typing.Union)
+                elif get_origin(annotation) is Union or (hasattr(annotation, '__class__') and annotation.__class__.__name__ == 'UnionType'):
                     args = get_args(annotation)
                     if dict in args:
                         should_support_string = True

tests/test_mcp_server.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from jupyter_server_mcp.mcp_server import MCPServer
+from jupyter_server_mcp.mcp_server import MCPServer, _auto_convert_json_args
 
 
 def simple_function(x: int, y: int) -> int:
@@ -212,3 +212,240 @@ def test_server_with_multiple_tools(self):
         assert server._registered_tools["simple_function"]["is_async"] is False
         assert server._registered_tools["async_function"]["is_async"] is True
         assert server._registered_tools["printer"]["is_async"] is False
+
+
+class TestJSONArgumentConversion:
+    """Test JSON argument conversion functionality."""
+
+    def test_simple_dict_conversion(self):
+        """Test basic JSON string to dict conversion."""
+
+        def func_with_dict(data: dict) -> dict:
+            """Function that expects a dict."""
+            return {"received": data, "type": type(data).__name__}
+
+        wrapped_func = _auto_convert_json_args(func_with_dict)
+
+        # Test with actual dict (should pass through)
+        result = wrapped_func(data={"key": "value"})
+        assert result["received"] == {"key": "value"}
+        assert result["type"] == "dict"
+
+        # Test with JSON string (should be converted)
+        result = wrapped_func(data='{"key": "value"}')
+        assert result["received"] == {"key": "value"}
+        assert result["type"] == "dict"
+
+    def test_optional_dict_conversion(self):
+        """Test JSON conversion with Optional[dict] annotation."""
+
+        def func_with_optional_dict(data: dict | None = None) -> dict:
+            """Function with optional dict parameter."""
+            return {"received": data, "type": type(data).__name__ if data else "NoneType"}
+
+        wrapped_func = _auto_convert_json_args(func_with_optional_dict)
+
+        # Test with None (should pass through)
+        result = wrapped_func(data=None)
+        assert result["received"] is None
+        assert result["type"] == "NoneType"
+
+        # Test with JSON string (should be converted)
+        result = wrapped_func(data='{"optional": true}')
+        assert result["received"] == {"optional": True}
+        assert result["type"] == "dict"
+
+    def test_union_dict_conversion(self):
+        """Test JSON conversion with Union type annotations."""
+
+        def func_with_union_dict(data: dict | None) -> dict:
+            """Function with Union[dict, None] parameter."""
+            return {"received": data, "type": type(data).__name__ if data else "NoneType"}
+
+        wrapped_func = _auto_convert_json_args(func_with_union_dict)
+
+        # Test with JSON string (should be converted)
+        result = wrapped_func(data='{"union": "test"}')
+        assert result["received"] == {"union": "test"}
+        assert result["type"] == "dict"
+
+    def test_typed_dict_conversion(self):
+        """Test JSON conversion with typed dict annotations."""
+
+        def func_with_typed_dict(config: dict[str, str]) -> dict:
+            """Function with Dict[str, str] annotation."""
+            return {"received": config, "type": type(config).__name__}
+
+        wrapped_func = _auto_convert_json_args(func_with_typed_dict)
+
+        # Test with JSON string (should be converted)
+        result = wrapped_func(config='{"name": "test", "value": "data"}')
+        assert result["received"] == {"name": "test", "value": "data"}
+        assert result["type"] == "dict"
+
+    def test_invalid_json_handling(self):
+        """Test handling of invalid JSON strings."""
+
+        def func_with_dict(data: dict) -> dict:
+            """Function that expects a dict."""
+            return {"received": data, "type": type(data).__name__}
+
+        wrapped_func = _auto_convert_json_args(func_with_dict)
+
+        # Test with invalid JSON (should pass string as-is)
+        result = wrapped_func(data="invalid json {")
+        assert result["received"] == "invalid json {"
+        assert result["type"] == "str"
+
+        # Test with empty string (should pass as-is)
+        result = wrapped_func(data="")
+        assert result["received"] == ""
+        assert result["type"] == "str"
+
+    def test_non_dict_parameters_unchanged(self):
+        """Test that non-dict parameters are not affected."""
+
+        def mixed_func(name: str, count: int, data: dict) -> dict:
+            """Function with mixed parameter types."""
+            return {
+                "name": name,
+                "name_type": type(name).__name__,
+                "count": count,
+                "count_type": type(count).__name__,
+                "data": data,
+                "data_type": type(data).__name__
+            }
+
+        wrapped_func = _auto_convert_json_args(mixed_func)
+
+        # Only the dict parameter should be converted
+        result = wrapped_func(
+            name="test",
+            count=42,
+            data='{"converted": true}'
+        )
+
+        assert result["name"] == "test"
+        assert result["name_type"] == "str"
+        assert result["count"] == 42
+        assert result["count_type"] == "int"
+        assert result["data"] == {"converted": True}
+        assert result["data_type"] == "dict"
+
+    @pytest.mark.asyncio
+    async def test_async_function_conversion(self):
+        """Test JSON conversion with async functions."""
+
+        async def async_func_with_dict(config: dict) -> dict:
+            """Async function that expects a dict."""
+            await asyncio.sleep(0.001)  # Small delay
+            return {"async_result": config, "type": type(config).__name__}
+
+        wrapped_func = _auto_convert_json_args(async_func_with_dict)
+
+        # Test with JSON string (should be converted)
+        result = await wrapped_func(config='{"async": true, "value": 123}')
+        assert result["async_result"] == {"async": True, "value": 123}
+        assert result["type"] == "dict"
+
+    def test_complex_nested_json(self):
+        """Test conversion of complex nested JSON structures."""
+
+        def func_with_nested_dict(data: dict) -> dict:
+            """Function that processes nested dict data."""
+            return {"processed": data}
+
+        wrapped_func = _auto_convert_json_args(func_with_nested_dict)
+
+        complex_json = '''{
+            "users": [
+                {"name": "Alice", "age": 30},
+                {"name": "Bob", "age": 25}
+            ],
+            "metadata": {
+                "version": "1.0",
+                "created": "2024-01-01"
+            }
+        }'''
+
+        result = wrapped_func(data=complex_json)
+        expected = {
+            "users": [
+                {"name": "Alice", "age": 30},
+                {"name": "Bob", "age": 25}
+            ],
+            "metadata": {
+                "version": "1.0",
+                "created": "2024-01-01"
+            }
+        }
+        assert result["processed"] == expected
+
+    def test_annotation_modification(self):
+        """Test that function annotations are properly modified."""
+
+        def original_func(data: dict) -> dict:
+            """Original function with dict annotation."""
+            return data
+
+        wrapped_func = _auto_convert_json_args(original_func)
+
+        # Check that annotations were modified to accept strings
+        annotations = wrapped_func.__annotations__
+        assert 'data' in annotations
+        
+        # The annotation should now be dict | str (or Union equivalent)
+        data_annotation = annotations['data']
+        # We can check this works by ensuring both dict and str are acceptable
+        assert hasattr(data_annotation, '__args__') or data_annotation == (dict | str)
+
+
+class TestJSONSchemaModification:
+    """Test JSON schema modification for MCP tools."""
+
+    def test_schema_modification_applied(self):
+        """Test that schema modification is applied during tool registration."""
+        server = MCPServer()
+
+        def func_with_dict_param(config: dict) -> str:
+            """Function with dict parameter."""
+            return f"Received config: {config}"
+
+        # Register the function - schema should be automatically modified
+        server.register_tool(func_with_dict_param)
+
+        # Verify the tool was registered
+        assert "func_with_dict_param" in server._registered_tools
+        tool_info = server._registered_tools["func_with_dict_param"]
+        assert tool_info["name"] == "func_with_dict_param"
+
+    def test_multiple_dict_parameters(self):
+        """Test conversion with multiple dict parameters."""
+
+        def func_multiple_dicts(config: dict, metadata: dict, name: str) -> dict:
+            """Function with multiple dict parameters."""
+            return {
+                "config": config,
+                "metadata": metadata,
+                "name": name,
+                "types": {
+                    "config": type(config).__name__,
+                    "metadata": type(metadata).__name__,
+                    "name": type(name).__name__
+                }
+            }
+
+        wrapped_func = _auto_convert_json_args(func_multiple_dicts)
+
+        result = wrapped_func(
+            config='{"key1": "value1"}',
+            metadata='{"version": 2}',
+            name="test_function"
+        )
+
+        assert result["config"] == {"key1": "value1"}
+        assert result["metadata"] == {"version": 2}
+        assert result["name"] == "test_function"
+        assert result["types"]["config"] == "dict"
+        assert result["types"]["metadata"] == "dict"
+        assert result["types"]["name"] == "str"