Skip to content

Commit 4497e7c

Browse files
authored
fix(mcp): alias MCP tool properties that pydantic cannot use as field names (#2442)
Closes #2441 `_model_from_mcp_schema` used the property names of an MCP tool schema as Pydantic field names. Pydantic rejects names with a leading underscore and the names of its own `model_` members, so a tool taking `_id` failed with `NameError: Fields must not use names with leading underscores` and one taking `model_config` failed with a `TypeError`. The input models are built inside `MCPBaseClient.get_tools()`, so one such tool made the whole `mcp_client` / `per_user_mcp_client` function group fail and every tool on the server was lost. Such a property now gets a field named after it with the leading underscores stripped and a trailing underscore added (`_id` becomes `id_`, `model_config` becomes `model_config_`, with more underscores if that name is already a property), and the original name is the field alias. Those models set `validate_by_name`, because the input arrives keyed by the schema's names from the LLM and keyed by field name when `FunctionInfo` passes the validated model on as keyword arguments. The JSON schema the LLM sees still lists `_id`. `_dump_tool_args` now dumps with `by_alias=True`, and `_drop_optional_none_values` maps each dumped key back to its field, so the server receives `_id` again, also in nested objects. Properties with ordinary names get no alias and their models get the same config as before, so nothing changes for them. The `model_` rule also covers names that built before but broke at call time. On develop a `model_post_init` property made every validation fail with a TypeError, and a `model_extra` property was silently left out of the arguments sent to the server. Names such as `json` or `copy`, which only shadow a deprecated `BaseModel` method with a warning, keep working as they did. This is the client side counterpart of #2079, which handled property names that are not Python identifiers in the server's tool converter. I kept the same idea of a safe name plus a mapping back to the original, but here the mapping lives in the model as a Pydantic alias, which also covers nested objects. The new tests are in `packages/nvidia_nat_mcp/tests/client/test_mcp_schema.py` (models for `_id` and `model_config`, validation by alias and by field name, a renamed field that would collide with another property, a nested `_id`, and that ordinary names get no alias) and `packages/nvidia_nat_mcp/tests/client/test_mcp_client_impl.py` (the arguments a per-user and a session tool send keep `_id` and `model_config` as keys at both levels, and unset renamed optional arguments are left out). On develop 8 of them fail, and the ordinary-names check passes as it should. With the change all 102 tests in those two files pass. The whole `packages/nvidia_nat_mcp/tests` run gives 450 passed and 8 skipped, plus `test_mcp_plugin_discovered` and `test_mcp_command_registration`, which fail the same way on develop in my environment because the package is not installed and the `mcp` CLI entry point is not found. I also built the `mcp_client` group against a small stdio server with an `echo` tool and a `get_document(_id)` tool. Both tools are registered now, and calling `get_document` through `ainvoke`, `acall_invoke` and the LangChain tool wrapper delivered `{'_id': ...}` to the server each time. yapf 0.43.0 and ruff 0.15.0 check are clean on the four files. ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing/index.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit * **Bug Fixes** * Fixed handling of tool schemas with underscore-prefixed names or names that conflict with model attributes. These names are now preserved in server arguments, including in nested fields. * Optional renamed fields that are unset are omitted from arguments, while required fields continue to be validated. * **Compatibility** * Improved validation and schema output for fields with special names without changing how ordinary field names are handled. Authors: - Dmitry Voropaev (https://github.com/v0ropaev) Approvers: - Will Killian (https://github.com/willkill07) URL: #2442
1 parent 4c219e2 commit 4497e7c

4 files changed

Lines changed: 138 additions & 10 deletions

File tree

‎packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/client_impl.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,14 @@ def _drop_optional_none_values(value: object, dumped: object) -> object:
8181
if isinstance(value, BaseModel):
8282
fields = type(value).model_fields
8383
schema_required = getattr(type(value), SCHEMA_REQUIRED_FIELDS_ATTR, frozenset())
84+
# The dump is keyed by alias, which is the schema's name for a field Pydantic could not name after it
85+
field_names = {field.alias or name: name for name, field in fields.items()}
8486
filtered = {}
85-
# The dump is keyed by field name (no aliases), so each key maps back to its attribute.
8687
for key, item_dump in dumped.items():
87-
item = getattr(value, key, item_dump)
88+
name = field_names.get(key, key)
89+
item = getattr(value, name, item_dump)
8890
if item is None:
89-
if key in schema_required or (key in fields and fields[key].is_required()):
91+
if key in schema_required or (name in fields and fields[name].is_required()):
9092
filtered[key] = None
9193
continue
9294
filtered[key] = _drop_optional_none_values(item, item_dump)
@@ -108,7 +110,7 @@ def _dump_tool_args(model: BaseModel) -> dict:
108110
reject those extra nulls. A required field keeps an explicit None: JSON Schema requires the key to be
109111
present, and null can be a valid value for it.
110112
"""
111-
return _drop_optional_none_values(model, model.model_dump(mode='json'))
113+
return _drop_optional_none_values(model, model.model_dump(mode='json', by_alias=True))
112114

113115

114116
class PerUserMCPFunctionGroup(FunctionGroup):

‎packages/nvidia_nat_mcp/src/nat/plugins/mcp/utils.py‎

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,29 @@ def truncate_session_id(session_id: str, max_length: int = 10) -> str:
178178
SCHEMA_REQUIRED_FIELDS_ATTR = "__mcp_schema_required_fields__"
179179

180180

181+
def _field_names_for_properties(property_names: list[str]) -> dict[str, str]:
182+
"""Map each schema property name to the name of its Pydantic field.
183+
184+
JSON Schema allows any property name, but Pydantic rejects field names with a leading underscore (e.g. `_id`)
185+
and names of its own `model_` members (e.g. `model_config`). Such a property gets a field named after it with
186+
the leading underscores stripped and a trailing underscore added, and keeps its own name as the field alias.
187+
Every other property keeps its name.
188+
"""
189+
taken = set(property_names)
190+
field_names = {}
191+
for property_name in property_names:
192+
if not property_name.startswith("_") and not (property_name.startswith("model_")
193+
and hasattr(BaseModel, property_name)):
194+
field_names[property_name] = property_name
195+
continue
196+
field_name = (property_name.lstrip("_") or "field") + "_"
197+
while field_name in taken:
198+
field_name += "_"
199+
taken.add(field_name)
200+
field_names[property_name] = field_name
201+
return field_names
202+
203+
181204
def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
182205
"""
183206
Create a pydantic model from the input schema of the MCP tool
@@ -403,7 +426,7 @@ def _has_null_in_type(field_properties: dict[str, Any]) -> bool:
403426

404427
return False
405428

406-
def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
429+
def _generate_field(field_name: str, field_properties: dict[str, Any], aliased: bool = False) -> tuple:
407430
"""
408431
Generate a Pydantic field from JSON schema properties.
409432
Uses _resolve_schema_type for type resolution and handles field-specific logic.
@@ -444,16 +467,28 @@ def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
444467
field_type = field_type | None
445468

446469
description = field_properties.get("description", "")
470+
# A field Pydantic cannot name after its property keeps the property's name as its alias
471+
alias_kwargs = {"alias": field_name} if aliased else {}
447472

448-
return field_type, Field(default=default_value, description=description, **_constraint_kwargs(field_properties))
473+
return field_type, Field(default=default_value,
474+
description=description,
475+
**alias_kwargs,
476+
**_constraint_kwargs(field_properties))
449477

478+
field_names = _field_names_for_properties(list(properties))
450479
for field_name, field_props in properties.items():
451-
schema_dict[field_name] = _generate_field(field_name=field_name, field_properties=field_props)
452-
model_config = None
480+
schema_dict[field_names[field_name]] = _generate_field(field_name=field_name,
481+
field_properties=field_props,
482+
aliased=field_names[field_name] != field_name)
483+
model_config = ConfigDict()
453484
if mcp_input_schema.get("additionalProperties") is False:
454-
model_config = ConfigDict(extra="forbid")
485+
model_config["extra"] = "forbid"
455486
elif mcp_input_schema.get("additionalProperties") is True:
456-
model_config = ConfigDict(extra="allow")
487+
model_config["extra"] = "allow"
488+
if any(field_names[property_name] != property_name for property_name in properties):
489+
# The input arrives keyed by the schema's names from the model, and by field name when NAT passes the
490+
# validated input on as keyword arguments, so accept both
491+
model_config["validate_by_name"] = True
457492

458493
model = create_model(f"{_generate_valid_classname(name)}InputSchema", __config__=model_config, **schema_dict)
459494
# A field can declare a default and still be in `required`, so Pydantic alone cannot tell which keys the

‎packages/nvidia_nat_mcp/tests/client/test_mcp_client_impl.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,6 +480,30 @@ async def tracking_ctx(session_id):
480480
"required": ["pet"],
481481
}
482482

483+
_UNDERSCORE_SCHEMA = {
484+
"type": "object",
485+
"properties": {
486+
"_id": {
487+
"type": "string"
488+
},
489+
"model_config": {
490+
"type": "string"
491+
},
492+
"filter": {
493+
"type": "object",
494+
"properties": {
495+
"_id": {
496+
"type": "string"
497+
}, "_rev": {
498+
"type": "string"
499+
}
500+
},
501+
"required": ["_id"],
502+
},
503+
},
504+
"required": ["_id"],
505+
}
506+
483507

484508
class _RecordingTool:
485509
"""Fake MCP tool that records the arguments sent to the server."""
@@ -582,3 +606,15 @@ async def test_discriminated_union_argument_is_sent_as_the_matching_variant(
582606
"""The `const` tag selects the variant, so the call carries that variant's fields and defaults."""
583607
sent = await send_tool_args({"pet": {"pet_type": "dog", "name": "Rex"}}, _DISCRIMINATED_UNION_SCHEMA)
584608
assert sent == {"pet": {"pet_type": "dog", "name": "Rex", "good_boy": True, "lives": 1}}
609+
610+
611+
async def test_arguments_renamed_for_pydantic_are_sent_under_their_schema_names(send_tool_args):
612+
"""`_id` is a field named `id_` in the input schema, but the server expects the key `_id`."""
613+
tool_args = {"_id": "doc-1", "model_config": "strict", "filter": {"_id": "doc-2"}}
614+
sent = await send_tool_args(tool_args, _UNDERSCORE_SCHEMA)
615+
assert sent == {"_id": "doc-1", "model_config": "strict", "filter": {"_id": "doc-2"}}
616+
617+
618+
async def test_renamed_optional_arguments_left_unset_are_not_sent(send_tool_args):
619+
sent = await send_tool_args({"_id": "doc-1", "filter": {"_id": "doc-2"}}, _UNDERSCORE_SCHEMA)
620+
assert sent == {"_id": "doc-1", "filter": {"_id": "doc-2"}}

‎packages/nvidia_nat_mcp/tests/client/test_mcp_schema.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1692,3 +1692,58 @@ def test_const_string_null_is_a_string_value():
16921692
assert _model.model_validate({"mode": "null"}).mode == "null"
16931693
with pytest.raises(ValidationError):
16941694
_model.model_validate({"mode": "other"})
1695+
1696+
1697+
@pytest.mark.parametrize("property_name, field_name", [("_id", "id_"), ("model_config", "model_config_")])
1698+
def test_property_pydantic_cannot_name_a_field_after_keeps_its_name_as_alias(property_name: str, field_name: str):
1699+
"""Pydantic rejects `_id` (common on MongoDB-style servers) and `model_config` as field names."""
1700+
schema = {"type": "object", "properties": {property_name: {"type": "string"}}, "required": [property_name]}
1701+
1702+
_model = model_from_mcp_schema(f"test_{property_name}_property", schema)
1703+
1704+
assert _model.model_fields[field_name].alias == property_name
1705+
assert list(_model.model_json_schema()["properties"]) == [property_name]
1706+
assert getattr(_model.model_validate({property_name: "doc-1"}), field_name) == "doc-1"
1707+
assert getattr(_model.model_validate({field_name: "doc-1"}), field_name) == "doc-1"
1708+
with pytest.raises(ValidationError):
1709+
_model.model_validate({})
1710+
1711+
1712+
def test_renamed_field_does_not_take_the_name_of_another_property():
1713+
schema = {"type": "object", "properties": {"_id": {"type": "string"}, "id_": {"type": "integer"}}}
1714+
1715+
_model = model_from_mcp_schema("test_renamed_field_collision", schema)
1716+
1717+
assert {name: field.alias for name, field in _model.model_fields.items()} == {"id__": "_id", "id_": None}
1718+
assert _model.model_validate({"_id": "doc-1", "id_": 1}).model_dump(by_alias=True) == {"_id": "doc-1", "id_": 1}
1719+
1720+
1721+
def test_nested_property_with_leading_underscore_keeps_its_name_as_alias():
1722+
schema = {
1723+
"type": "object",
1724+
"properties": {
1725+
"filter": {
1726+
"type": "object", "properties": {
1727+
"_id": {
1728+
"type": "string"
1729+
}
1730+
}, "required": ["_id"]
1731+
}
1732+
},
1733+
"required": ["filter"],
1734+
}
1735+
1736+
_model = model_from_mcp_schema("test_nested_underscore_property", schema)
1737+
1738+
instance = _model.model_validate({"filter": {"_id": "doc-1"}})
1739+
assert instance.filter.id_ == "doc-1"
1740+
assert instance.model_dump(by_alias=True) == {"filter": {"_id": "doc-1"}}
1741+
1742+
1743+
def test_ordinary_property_names_get_no_alias():
1744+
schema = {"type": "object", "properties": {"id": {"type": "string"}, "model_name": {"type": "string"}}}
1745+
1746+
_model = model_from_mcp_schema("test_ordinary_property_names", schema)
1747+
1748+
assert {name: field.alias for name, field in _model.model_fields.items()} == {"id": None, "model_name": None}
1749+
assert "validate_by_name" not in _model.model_config

0 commit comments

Comments
 (0)