more tests passing

Author: naji247

src/mcpcat/modules/tools.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from mcp import Tool
+from mcp import ServerResult, Tool
 from mcp.server import FastMCP
 from mcp.types import CallToolRequest, CallToolResult, ListToolsRequest, TextContent
 
@@ -25,7 +25,9 @@ async def wrapped_list_tools_handler(request: ListToolsRequest) -> ListToolsResu
         """Intercept list_tools requests to add MCPCat tools and modify existing ones."""
         # Call the original handler to get the tools
         original_result = await original_list_tools_handler(request)
-        tools_list = original_result.tools if hasattr(original_result, 'tools') else []
+        if not original_result or not hasattr(original_result, 'root') or not hasattr(original_result.root, 'tools'):
+            return original_result
+        tools_list = original_result.root.tools
 
         # Add report_missing tool if enabled
         if data.options.enableReportMissing:
@@ -77,7 +79,7 @@ async def wrapped_list_tools_handler(request: ListToolsRequest) -> ListToolsResu
                         else:
                             tool.inputSchema["required"] = ["context"]
 
-        return ListToolsResult(tools=tools_list)
+        return ServerResult(ListToolsResult(tools=tools_list))
 
     async def wrapped_call_tool_handler(request: CallToolRequest) -> CallToolResult:
         """Intercept call_tool requests to add MCPCat tracking and handle special tools."""

tests/test_preserve_existing_tools.py

@@ -0,0 +1,137 @@
+"""Test that mcpcat.track preserves existing tools and only adds report_missing."""
+
+import pytest
+
+from mcpcat import track
+from mcpcat.types import MCPCatOptions
+
+from .test_utils.client import create_test_client
+from .test_utils.todo_server import create_todo_server
+
+
+class TestPreserveExistingTools:
+    """Test that existing tools are preserved when tracking."""
+
+    async def test_track_preserves_all_existing_tools(self):
+        """Should preserve all existing tools and only add report_missing."""
+        # Create server with existing tools
+        server = create_todo_server()
+        
+        # Get original tools before tracking
+        async with create_test_client(server) as client:
+            result = await client.list_tools()
+            original_tools = result.tools
+            original_tool_names = {tool.name for tool in original_tools}
+        
+        # Track the server
+        tracked_server = track(server)
+        
+        # Get tools after tracking
+        async with create_test_client(tracked_server) as client:
+            result = await client.list_tools()
+            tracked_tools = result.tools
+            tracked_tool_names = {tool.name for tool in tracked_tools}
+        
+        # Verify all original tools are preserved
+        assert original_tool_names.issubset(tracked_tool_names), "Original tools should be preserved"
+        
+        # Verify only report_missing was added
+        added_tools = tracked_tool_names - original_tool_names
+        assert added_tools == {"report_missing"}, "Only report_missing should be added"
+        
+        # Verify the exact count
+        assert len(tracked_tools) == len(original_tools) + 1, "Should have exactly one more tool"
+        
+        # Verify original tools are still present
+        assert "add_todo" in tracked_tool_names
+        assert "list_todos" in tracked_tool_names
+        assert "complete_todo" in tracked_tool_names
+
+    async def test_track_without_report_missing_preserves_exact_tools(self):
+        """Should preserve exact tool list when report_missing is disabled."""
+        # Create server with existing tools
+        server = create_todo_server()
+        
+        # Get original tools before tracking
+        async with create_test_client(server) as client:
+            result = await client.list_tools()
+            original_tools = result.tools
+            original_tool_names = {tool.name for tool in original_tools}
+        
+        # Track the server with report_missing disabled
+        tracked_server = track(server, MCPCatOptions(enableReportMissing=False))
+        
+        # Get tools after tracking
+        async with create_test_client(tracked_server) as client:
+            result = await client.list_tools()
+            tracked_tools = result.tools
+            tracked_tool_names = {tool.name for tool in tracked_tools}
+        
+        # Verify tool lists are identical
+        assert tracked_tool_names == original_tool_names, "Tool names should be identical"
+        assert len(tracked_tools) == len(original_tools), "Tool count should be identical"
+
+    async def test_track_with_context_modifies_existing_tools_but_preserves_them(self):
+        """Should modify existing tools to add context but preserve all of them."""
+        # Create server with existing tools
+        server = create_todo_server()
+        
+        # Get original tools before tracking
+        async with create_test_client(server) as client:
+            result = await client.list_tools()
+            original_tools = result.tools
+            original_tool_names = {tool.name for tool in original_tools}
+        
+        # Track the server with context enabled
+        tracked_server = track(server, MCPCatOptions(enableToolCallContext=True))
+        
+        # Get tools after tracking
+        async with create_test_client(tracked_server) as client:
+            result = await client.list_tools()
+            tracked_tools = result.tools
+        
+        # Find modified tools (excluding report_missing)
+        for tool in tracked_tools:
+            if tool.name in original_tool_names:
+                # Verify context was added to schema
+                assert "context" in tool.inputSchema["properties"], f"Context should be added to {tool.name}"
+                assert "context" in tool.inputSchema["required"], f"Context should be required for {tool.name}"
+            elif tool.name == "report_missing":
+                # Verify report_missing doesn't have context parameter
+                assert "context" not in tool.inputSchema["properties"], "report_missing should not have context"
+        
+        # Verify all original tools are still present
+        tracked_tool_names = {tool.name for tool in tracked_tools}
+        assert original_tool_names.issubset(tracked_tool_names), "All original tools should be preserved"
+
+    async def test_multiple_track_calls_do_not_duplicate_tools(self):
+        """Should not duplicate tools when track is called multiple times."""
+        # Create server with existing tools
+        server = create_todo_server()
+        
+        # Track the server multiple times
+        track(server)
+        track(server)  # Should be a no-op
+        track(server)  # Should be a no-op
+        
+        # Get tools after multiple track calls
+        async with create_test_client(server) as client:
+            result = await client.list_tools()
+            tools = result.tools
+            tool_names = [tool.name for tool in tools]
+        
+        # Count occurrences of each tool
+        tool_counts = {}
+        for name in tool_names:
+            tool_counts[name] = tool_counts.get(name, 0) + 1
+        
+        # Verify no duplicates
+        for name, count in tool_counts.items():
+            assert count == 1, f"Tool {name} should appear exactly once, but appears {count} times"
+        
+        # Verify expected tools are present
+        assert "add_todo" in tool_names
+        assert "list_todos" in tool_names
+        assert "complete_todo" in tool_names
+        assert "report_missing" in tool_names
+        assert len(tools) == 4, "Should have exactly 4 tools"

tests/test_utils/todo_server.py

@@ -22,7 +22,7 @@ def create_todo_server() -> FastMCP:
     next_id = 1
 
     @server.tool()
-    def add_todo(text: str, context: str | None = None) -> str:
+    def add_todo(text: str) -> str:
         """Add a new todo item."""
         nonlocal next_id
         todo = Todo(str(next_id), text)
@@ -31,7 +31,7 @@ def add_todo(text: str, context: str | None = None) -> str:
         return f'Added todo: "{text}" with ID {todo.id}'
 
     @server.tool()
-    def list_todos(context: str | None = None) -> str:
+    def list_todos() -> str:
         """List all todo items."""
         if not todos:
             return "No todos found"
@@ -44,7 +44,7 @@ def list_todos(context: str | None = None) -> str:
         return "\n".join(todo_list)
 
     @server.tool()
-    def complete_todo(id: str, context: str | None = None) -> str:
+    def complete_todo(id: str) -> str:
         """Mark a todo item as completed."""
         for todo in todos:
             if todo.id == id: