Skip to content

Commit 549f183

Browse files
ZsailerZsailer
authored
Merge pull request #4 from Zsailer/json-arg-conversion
Add automatic JSON argument conversion for MCP tools
2 parents 3ab6029 + e1d8301 commit 549f183

File tree

2 files changed

+481
-20
lines changed

2 files changed

+481
-20
lines changed

‎jupyter_server_mcp/mcp_server.py‎

Lines changed: 249 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,246 @@
11
"""Simple MCP server for registering Python functions as tools."""
22

3+
import inspect
4+
import json
35
import logging
46
from collections.abc import Callable
5-
from inspect import iscoroutinefunction
6-
from typing import Any
7+
from functools import wraps
8+
from inspect import iscoroutinefunction, signature
9+
from typing import Any, Union, get_args, get_origin
710

811
from fastmcp import FastMCP
9-
from traitlets import Bool, Int, Unicode
12+
from traitlets import Int, Unicode
1013
from traitlets.config.configurable import LoggingConfigurable
1114

1215
logger = logging.getLogger(__name__)
1316

1417

18+
def _is_dict_compatible_annotation(annotation) -> bool:
19+
"""Check if an annotation expects dict values that can be JSON-converted."""
20+
# Direct dict annotation
21+
if annotation is dict:
22+
return True
23+
24+
# Union types: Optional[dict], Union[dict, None], dict | None
25+
origin = get_origin(annotation)
26+
if origin is Union or (
27+
hasattr(annotation, "__class__")
28+
and annotation.__class__.__name__ == "UnionType"
29+
):
30+
args = get_args(annotation)
31+
return dict in args
32+
33+
# Typed dict annotations: Dict[K, V], dict[str, Any]
34+
return bool(hasattr(annotation, "__origin__") and annotation.__origin__ is dict)
35+
36+
37+
def _wrap_with_json_conversion(func: Callable) -> Callable:
38+
"""
39+
Wrapper that automatically converts JSON string arguments to dictionaries.
40+
41+
This addresses the common issue where MCP clients pass dictionary arguments
42+
as JSON strings instead of structured objects. The wrapper inspects the
43+
function signature and attempts JSON parsing for parameters annotated as
44+
dict types when they are received as strings.
45+
46+
Additionally, this function modifies the type annotations to accept Union[dict, str]
47+
for dict parameters to allow Pydantic validation to pass.
48+
49+
This conversion is always applied to all registered tools to ensure compatibility
50+
with various MCP clients that may serialize dict parameters differently.
51+
52+
Args:
53+
func: The function to wrap
54+
55+
Returns:
56+
Wrapped function that handles JSON string conversion with modified annotations
57+
"""
58+
sig = signature(func)
59+
60+
def _should_convert_to_dict(annotation, value):
61+
"""Check if a parameter should be converted from JSON string to dict."""
62+
return isinstance(value, str) and _is_dict_compatible_annotation(annotation)
63+
64+
def _add_string_to_annotation(annotation):
65+
"""Modify annotation to also accept strings for dict types."""
66+
# Direct dict annotation -> dict | str
67+
if annotation is dict:
68+
return dict | str
69+
70+
# Union types: add str to existing union
71+
origin = get_origin(annotation)
72+
if origin is Union:
73+
args = get_args(annotation)
74+
if dict in args and str not in args:
75+
return Union[(*tuple(args), str)]
76+
return annotation
77+
78+
# New Python 3.10+ union syntax: dict | None
79+
if (
80+
hasattr(annotation, "__class__")
81+
and annotation.__class__.__name__ == "UnionType"
82+
):
83+
args = get_args(annotation)
84+
if dict in args and str not in args:
85+
# Reconstruct the union with str added
86+
new_args = (*tuple(args), str)
87+
# Create new union type
88+
result = new_args[0]
89+
for arg in new_args[1:]:
90+
result = result | arg
91+
return result
92+
return annotation
93+
94+
# Typed dict annotations -> annotation | str
95+
if hasattr(annotation, "__origin__") and annotation.__origin__ is dict:
96+
return annotation | str
97+
98+
return annotation
99+
100+
# Create new annotations that accept strings for dict parameters
101+
new_annotations = {}
102+
for param_name, param in sig.parameters.items():
103+
if param.annotation != inspect.Parameter.empty:
104+
new_annotations[param_name] = _add_string_to_annotation(param.annotation)
105+
else:
106+
new_annotations[param_name] = param.annotation
107+
108+
# Keep the return annotation unchanged
109+
if hasattr(func, "__annotations__") and "return" in func.__annotations__:
110+
new_annotations["return"] = func.__annotations__["return"]
111+
112+
if iscoroutinefunction(func):
113+
114+
@wraps(func)
115+
async def async_wrapper(*args, **kwargs):
116+
# Convert keyword arguments that should be dicts but are strings
117+
converted_kwargs = {}
118+
for param_name, param_value in kwargs.items():
119+
if param_name in sig.parameters:
120+
param = sig.parameters[param_name]
121+
if _should_convert_to_dict(param.annotation, param_value):
122+
try:
123+
converted_kwargs[param_name] = json.loads(param_value)
124+
logger.debug(
125+
f"Converted JSON string to dict for parameter '{param_name}': {param_value}"
126+
)
127+
except json.JSONDecodeError:
128+
# If it's not valid JSON, pass the string as-is
129+
converted_kwargs[param_name] = param_value
130+
else:
131+
converted_kwargs[param_name] = param_value
132+
else:
133+
converted_kwargs[param_name] = param_value
134+
135+
return await func(*args, **converted_kwargs)
136+
137+
# Set the modified annotations on the wrapper
138+
async_wrapper.__annotations__ = new_annotations
139+
return async_wrapper
140+
141+
@wraps(func)
142+
def sync_wrapper(*args, **kwargs):
143+
# Convert keyword arguments that should be dicts but are strings
144+
converted_kwargs = {}
145+
for param_name, param_value in kwargs.items():
146+
if param_name in sig.parameters:
147+
param = sig.parameters[param_name]
148+
if _should_convert_to_dict(param.annotation, param_value):
149+
try:
150+
converted_kwargs[param_name] = json.loads(param_value)
151+
logger.debug(
152+
f"Converted JSON string to dict for parameter '{param_name}': {param_value}"
153+
)
154+
except json.JSONDecodeError:
155+
# If it's not valid JSON, pass the string as-is
156+
converted_kwargs[param_name] = param_value
157+
else:
158+
converted_kwargs[param_name] = param_value
159+
else:
160+
converted_kwargs[param_name] = param_value
161+
162+
return func(*args, **converted_kwargs)
163+
164+
# Set the modified annotations on the wrapper
165+
sync_wrapper.__annotations__ = new_annotations
166+
return sync_wrapper
167+
168+
169+
def _update_schema_for_json_args(func: Callable, tool) -> None:
170+
"""
171+
Modify the tool's JSON schema to accept strings for dict parameters.
172+
173+
This function updates the input schema to allow JSON strings in addition to objects for
174+
parameters that are annotated as dict types, enabling MCP clients to pass JSON strings
175+
that will be automatically converted to dicts.
176+
177+
This modification is always applied to ensure compatibility with various MCP clients.
178+
179+
Args:
180+
func: The original function
181+
tool: The FastMCP tool object
182+
"""
183+
try:
184+
sig = signature(func)
185+
186+
# Get the MCP tool representation to modify its schema
187+
mcp_tool_dict = tool.to_mcp_tool().model_dump()
188+
input_schema = mcp_tool_dict.get("inputSchema", {})
189+
properties = input_schema.get("properties", {})
190+
191+
# Check each parameter in the function signature
192+
for param_name, param in sig.parameters.items():
193+
if param_name in properties:
194+
param_schema = properties[param_name]
195+
196+
# Check if this parameter should support JSON string conversion
197+
annotation = param.annotation
198+
should_support_string = _is_dict_compatible_annotation(annotation)
199+
200+
if should_support_string:
201+
# Modify the schema to also accept strings
202+
if "anyOf" in param_schema:
203+
# For Optional[dict] - add string to the anyOf list
204+
existing_schemas = param_schema["anyOf"]
205+
# Check if string is already in the schema
206+
has_string = any(
207+
s.get("type") == "string" for s in existing_schemas
208+
)
209+
if not has_string:
210+
existing_schemas.append(
211+
{
212+
"type": "string",
213+
"description": "JSON string that will be parsed to object",
214+
}
215+
)
216+
elif param_schema.get("type") == "object":
217+
# For dict - convert to anyOf with object and string
218+
original_schema = param_schema.copy()
219+
properties[param_name] = {
220+
"anyOf": [
221+
original_schema,
222+
{
223+
"type": "string",
224+
"description": "JSON string that will be parsed to object",
225+
},
226+
],
227+
"title": param_schema.get("title", param_name.title()),
228+
}
229+
# Preserve default if it exists
230+
if "default" in param_schema:
231+
properties[param_name]["default"] = param_schema["default"]
232+
233+
# Update the tool's parameters with the modified schema
234+
tool.parameters = input_schema
235+
236+
logger.debug(
237+
f"Modified schema for tool '{tool.name}' to support JSON strings for dict parameters"
238+
)
239+
240+
except Exception as e:
241+
logger.warning(f"Could not modify schema for JSON string support: {e}")
242+
243+
15244
class MCPServer(LoggingConfigurable):
16245
"""Simple MCP server that allows registering Python functions as tools."""
17246

@@ -28,10 +257,6 @@ class MCPServer(LoggingConfigurable):
28257
default_value="localhost", help="Host for the MCP server to listen on"
29258
).tag(config=True)
30259

31-
enable_debug_logging = Bool(
32-
default_value=False, help="Enable debug logging for MCP operations"
33-
).tag(config=True)
34-
35260
def __init__(self, **kwargs):
36261
"""Initialize the MCP server.
37262
@@ -65,14 +290,24 @@ def register_tool(
65290
tool_description = description or func.__doc__ or f"Tool: {tool_name}"
66291

67292
self.log.info(f"Registering tool: {tool_name}")
68-
if self.enable_debug_logging:
69-
self.log.debug(
70-
f"Tool details - Name: {tool_name}, "
71-
f"Description: {tool_description}, Async: {iscoroutinefunction(func)}"
72-
)
293+
self.log.debug(
294+
f"Tool details - Name: {tool_name}, "
295+
f"Description: {tool_description}, Async: {iscoroutinefunction(func)}"
296+
)
297+
298+
# Apply auto-conversion wrapper (always enabled)
299+
registered_func = _wrap_with_json_conversion(func)
300+
self.log.debug(f"Applied JSON argument auto-conversion wrapper to {tool_name}")
73301

74302
# Register with FastMCP
75-
self.mcp.tool(func)
303+
tool = self.mcp.tool(registered_func)
304+
305+
# Modify schema to support JSON strings for dict parameters
306+
if tool:
307+
_update_schema_for_json_args(func, tool)
308+
self.log.debug(
309+
f"Modified schema for tool '{tool_name}' to accept JSON strings for dict parameters"
310+
)
76311

77312
# Keep track for listing
78313
self._registered_tools[tool_name] = {
@@ -115,11 +350,7 @@ async def start_server(self, host: str | None = None):
115350

116351
self.log.info(f"Starting MCP server '{self.name}' on {server_host}:{self.port}")
117352
self.log.info(f"Registered tools: {list(self._registered_tools.keys())}")
118-
119-
if self.enable_debug_logging:
120-
self.log.debug(
121-
f"Server configuration - Host: {server_host}, Port: {self.port}"
122-
)
353+
self.log.debug(f"Server configuration - Host: {server_host}, Port: {self.port}")
123354

124355
# Start FastMCP server with HTTP transport
125356
await self.mcp.run_http_async(host=server_host, port=self.port)

0 commit comments

Comments
 (0)