From 3502b4cec5957f79e5719e3d8e1ddb878d677e6f Mon Sep 17 00:00:00 2001
From: Vishal Habib <vishalhabib99@gmail.com>
Date: Wed, 14 Oct 2026 12:31:49 -0400
Subject: [PATCH] fix(tools): let optional MCP tool arguments accept null

read_file, search_code and read_document_segments declared optional
arguments as `int = None` / `str = None` / `List[str] = None`. FastMCP
advertises those as {"type": "integer", "default": null}, so a client
that sends the advertised default back explicitly gets a validation error
before the tool runs. Annotate the five arguments as `T | None`.
---
 tests/test_bundled_mcp_optional_args.py | 66 +++++++++++++++++++++++++
 tools/code_implementation_server.py     |  4 +-
 tools/document_segmentation_server.py   |  4 +-
 3 files changed, 70 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_bundled_mcp_optional_args.py

diff --git a/tests/test_bundled_mcp_optional_args.py b/tests/test_bundled_mcp_optional_args.py
new file mode 100644
index 000000000..59c13978a
--- /dev/null
+++ b/tests/test_bundled_mcp_optional_args.py
@@ -0,0 +1,66 @@
+"""Regression cover for optional tool arguments in the bundled MCP servers.
+
+A parameter written ``start_line: int = None`` makes FastMCP advertise
+``{"type": "integer", "default": null}``: the default is null, but the type
+excludes null. A client that sends the advertised default back explicitly
+(``"start_line": null``) is rejected with a validation error before the tool
+runs. Every argument whose schema default is null must also accept null.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from pathlib import Path
+
+import pytest
+from pydantic import ValidationError
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+import tools.code_implementation_server as code_implementation_server
+import tools.document_segmentation_server as document_segmentation_server
+
+SERVERS = {
+    "code-implementation": code_implementation_server.mcp,
+    "document-segmentation": document_segmentation_server.mcp,
+}
+
+
+def _accepts_null(prop: dict) -> bool:
+    if prop.get("type") == "null":
+        return True
+    if isinstance(prop.get("type"), list) and "null" in prop["type"]:
+        return True
+    return any(_accepts_null(option) for option in prop.get("anyOf", []))
+
+
+def _null_default_args(server):
+    tools = asyncio.run(server.list_tools())
+    for tool in tools:
+        for name, prop in tool.inputSchema.get("properties", {}).items():
+            if "default" in prop and prop["default"] is None:
+                yield tool.name, name, prop
+
+
+@pytest.mark.parametrize("server_name", sorted(SERVERS))
+def test_null_default_arguments_accept_null(server_name):
+    server = SERVERS[server_name]
+    found = list(_null_default_args(server))
+    assert found, f"{server_name}: expected optional arguments with a null default"
+
+    rejected = []
+    for tool_name, arg_name, prop in found:
+        tool = server._tool_manager.get_tool(tool_name)
+        required = {req: "x" for req in tool.parameters.get("required", [])}
+        if not _accepts_null(prop):
+            rejected.append(f"{tool_name}.{arg_name}: schema {prop}")
+            continue
+        try:
+            tool.fn_metadata.arg_model.model_validate({**required, arg_name: None})
+        except ValidationError as exc:
+            rejected.append(f"{tool_name}.{arg_name}: {exc}")
+
+    assert not rejected, "null default but null rejected:\n" + "\n".join(rejected)
diff --git a/tools/code_implementation_server.py b/tools/code_implementation_server.py
index d16a9cd58..ff475cd92 100644
--- a/tools/code_implementation_server.py
+++ b/tools/code_implementation_server.py
@@ -109,7 +109,7 @@ def log_operation(action: str, details: Dict[str, Any]):
 
 @mcp.tool()
 async def read_file(
-    file_path: str, start_line: int = None, end_line: int = None
+    file_path: str, start_line: int | None = None, end_line: int | None = None
 ) -> str:
     """
     Read file content, supports specifying line number range
@@ -1184,7 +1184,7 @@ async def search_code(
     pattern: str,
     file_pattern: str = "*.json",
     use_regex: bool = False,
-    search_directory: str = None,
+    search_directory: str | None = None,
 ) -> str:
     """
     Search patterns in code files
diff --git a/tools/document_segmentation_server.py b/tools/document_segmentation_server.py
index 91e12a783..454d64866 100644
--- a/tools/document_segmentation_server.py
+++ b/tools/document_segmentation_server.py
@@ -1627,9 +1627,9 @@ async def analyze_and_segment_document(
 async def read_document_segments(
     paper_dir: str,
     query_type: str,
-    keywords: List[str] = None,
+    keywords: List[str] | None = None,
     max_segments: int = 3,
-    max_total_chars: int = None,
+    max_total_chars: int | None = None,
 ) -> str:
     """
     Intelligently retrieve relevant document segments based on query type