diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index 3cf1147c..7902e66b 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -315,13 +315,76 @@ def _extract_nullable_branch(options: Any) -> tuple[dict[str, Any], bool] | None return None -def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]: - """Normalize only nullable JSON Schema patterns for tool definitions.""" - if not isinstance(schema, dict): - return {"type": "object", "properties": {}} +def _resolve_local_schema_ref(root: dict[str, Any], ref: str) -> Any: + """Resolve a local JSON Pointer without accepting remote references.""" + if ref == "#": + return root + if not ref.startswith("#/"): + raise ValueError("not a local JSON Pointer") + current: Any = root + for raw_part in ref[2:].split("/"): + part = raw_part.replace("~1", "/").replace("~0", "~") + if isinstance(current, dict): + current = current[part] + elif isinstance(current, list): + current = current[int(part)] + else: + raise KeyError(part) + return current + + +def _rewrite_local_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: + """Hoist arbitrary local JSON-Pointer refs into provider-compatible ``$defs``.""" + rewritten_refs: dict[str, str] = {} + generated_defs: dict[str, Any] = {} + + def rewrite(value: Any) -> Any: + if isinstance(value, list): + return [rewrite(item) for item in value] + if not isinstance(value, dict): + return value + + rewritten = dict(value) + ref = rewritten.get("$ref") + is_rewritable_ref = isinstance(ref, str) and ( + ref == "#" or (ref.startswith("#/") and not ref.startswith("#/$defs/")) + ) + if is_rewritable_ref: + name = rewritten_refs.get(ref) + if name is None: + try: + target = _resolve_local_schema_ref(schema, ref) + except (KeyError, IndexError, TypeError, ValueError): + logger.warning("MCP tool schema contains an unresolved local $ref: {}", ref) + else: + assert isinstance(ref, str) + name = f"ref_{hashlib.sha256(ref.encode()).hexdigest()[:12]}" + existing_defs = schema.get("$defs") + while isinstance(existing_defs, dict) and name in existing_defs: + name += "_" + rewritten_refs[ref] = name + # Reserve the name before descending so recursive refs terminate. + generated_defs[name] = {} + generated_defs[name] = rewrite(target) + if name is not None: + rewritten["$ref"] = f"#/$defs/{name}" + + return {key: rewrite(item) for key, item in rewritten.items()} + + result = rewrite(schema) + if generated_defs: + existing_defs = result.get("$defs") + result["$defs"] = { + **(existing_defs if isinstance(existing_defs, dict) else {}), + **generated_defs, + } + return result + + +def _normalize_nullable_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Normalize nullable forms in structural subschemas only.""" normalized = dict(schema) - raw_type = normalized.get("type") if isinstance(raw_type, list): non_null = [item for item in raw_type if item != "null"] @@ -339,23 +402,34 @@ def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]: normalized["nullable"] = True break - if "properties" in normalized and isinstance(normalized["properties"], dict): + if isinstance(normalized.get("properties"), dict): normalized["properties"] = { - name: _normalize_schema_for_openai(prop) if isinstance(prop, dict) else prop + name: _normalize_nullable_schema(prop) if isinstance(prop, dict) else prop for name, prop in normalized["properties"].items() } + if isinstance(normalized.get("items"), dict): + normalized["items"] = _normalize_nullable_schema(normalized["items"]) + if isinstance(normalized.get("$defs"), dict): + normalized["$defs"] = { + name: _normalize_nullable_schema(definition) + if isinstance(definition, dict) + else definition + for name, definition in normalized["$defs"].items() + } - if "items" in normalized and isinstance(normalized["items"], dict): - normalized["items"] = _normalize_schema_for_openai(normalized["items"]) - - if normalized.get("type") != "object": - return normalized - - normalized.setdefault("properties", {}) - normalized.setdefault("required", []) + if normalized.get("type") == "object": + normalized.setdefault("properties", {}) + normalized.setdefault("required", []) return normalized +def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]: + """Normalize MCP JSON Schema patterns for tool definitions.""" + if not isinstance(schema, dict): + return {"type": "object", "properties": {}} + return _normalize_nullable_schema(_rewrite_local_schema_refs(schema)) + + class _MCPWrapperBase(Tool): """Common reconnect handling for wrappers bound to one MCP server session.""" diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 8bf932a7..233c51fa 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -236,6 +236,79 @@ def test_wrapper_normalizes_nullable_property_anyof() -> None: } +def test_wrapper_hoists_recursive_local_refs_into_defs() -> None: + recursive_items_ref = "#/properties/filter/properties/items" + tool_def = SimpleNamespace( + name="search_dataset", + description="search tool", + inputSchema={ + "type": "object", + "properties": { + "filter": { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": {"$ref": recursive_items_ref}, + } + }, + "required": ["items"], + } + }, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + generated_ref = wrapper.parameters["properties"]["filter"]["properties"]["items"][ + "items" + ]["$ref"] + assert generated_ref.startswith("#/$defs/ref_") + generated_name = generated_ref.removeprefix("#/$defs/") + generated_schema = wrapper.parameters["$defs"][generated_name] + assert generated_schema["type"] == "array" + assert generated_schema["items"]["$ref"] == generated_ref + + +def test_wrapper_hoists_root_self_ref_into_defs() -> None: + tool_def = SimpleNamespace( + name="tree", + description="tree tool", + inputSchema={ + "type": "object", + "properties": { + "children": {"type": "array", "items": {"$ref": "#"}}, + }, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + generated_ref = wrapper.parameters["properties"]["children"]["items"]["$ref"] + assert generated_ref.startswith("#/$defs/ref_") + generated_name = generated_ref.removeprefix("#/$defs/") + assert wrapper.parameters["$defs"][generated_name]["properties"]["children"]["items"] == { + "$ref": generated_ref + } + + +def test_wrapper_preserves_existing_defs_refs() -> None: + tool_def = SimpleNamespace( + name="demo", + description="demo tool", + inputSchema={ + "type": "object", + "$defs": {"value": {"type": "string"}}, + "properties": {"value": {"$ref": "#/$defs/value"}}, + }, + ) + + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def) + + assert wrapper.parameters["properties"]["value"]["$ref"] == "#/$defs/value" + assert wrapper.parameters["$defs"]["value"]["type"] == "string" + + def test_normalize_windows_stdio_command_is_noop_off_windows( monkeypatch: pytest.MonkeyPatch, ) -> None: