fix tests

Author: naji247

.claude/settings.local.json

@@ -5,7 +5,11 @@
       "Bash(find:*)",
       "Bash(grep:*)",
       "Bash(uv run pytest tests/test_get_unknown_or_stdio_session.py -v)",
-      "Bash(uv run:*)"
+      "Bash(uv run:*)",
+      "Bash(rg:*)",
+      "Bash(/Users/naseemalnaji/.nvm/versions/node/v23.5.0/lib/node_modules/@anthropic-ai/claude-code/vendor/ripgrep/arm64-darwin/rg \"CallToolResult\" --type py -A 3 -B 3)",
+      "Bash(/Users/naseemalnaji/.nvm/versions/node/v23.5.0/lib/node_modules/@anthropic-ai/claude-code/vendor/ripgrep/arm64-darwin/rg \"from mcp.types import.*CallToolResult\" --type py)",
+      "Bash(pip show:*)"
     ],
     "deny": []
   }

examples/basic_usage.py

@@ -1,8 +1,8 @@
 """Basic usage example for MCPCat with FastMCP."""
 
 from fastmcp import FastMCP
-from mcpcat import track, MCPCatOptions
 
+from mcpcat import MCPCatOptions, track
 
 # Create a FastMCP server
 mcp = FastMCP("example-server")
@@ -37,4 +37,4 @@ def get_weather(city: str) -> str:
 
 # Run the server
 if __name__ == "__main__":
-    mcp.run()
+    mcp.run()

pyproject.toml

@@ -19,7 +19,9 @@ classifiers = [
 ]
 
 dependencies = [
+    "datafog>=4.1.1",
     "mcp>=1.0.0",
+    "pydantic>=2.0.0",
 ]
 
 [project.optional-dependencies]
@@ -82,4 +84,4 @@ exclude = [
 ]
 
 [tool.ruff.isort]
-known-first-party = ["mcpcat"]
+known-first-party = ["mcpcat"]

src/mcpcat/__init__.py

@@ -11,7 +11,7 @@
 __version__ = "0.1.0"
 
 
-def track(server: Any, options: Optional[MCPCatOptions] = None) -> Any:
+def track(server: Any, options: MCPCatOptions | None = None) -> Any:
     """
     Enable analytics tracking for an MCP server.
     
@@ -28,13 +28,13 @@ def track(server: Any, options: Optional[MCPCatOptions] = None) -> Any:
     # Use default options if not provided
     if options is None:
         options = MCPCatOptions()
-    
+
     # Validate server compatibility
     if not is_compatible_server(server):
         raise TypeError(
             "Server must be a FastMCP instance or implement the MCP server protocol"
         )
-    
+
     # Check if already tracked
     if has_mcpcat_data(server):
         log_info(
@@ -43,15 +43,15 @@ def track(server: Any, options: Optional[MCPCatOptions] = None) -> Any:
             options
         )
         return server
-    
+
     # Create and store tracking data
     data = MCPCatData(options=options)
     set_mcpcat_data(server, data)
-    
+
     try:
         # Set up tool handlers
         setup_tool_handlers(server, data)
-        
+
         # Log initialization
         log_info(
             "MCPCat tracking initialized",
@@ -70,36 +70,36 @@ def track(server: Any, options: Optional[MCPCatOptions] = None) -> Any:
         if has_mcpcat_data(server):
             # Remove from tracking
             get_mcpcat_data(server)  # This will remove the weak reference
-        
+
         log_error(
             "Failed to initialize MCPCat",
             e,
             {"server": server.__class__.__name__},
             options
         )
-        
+
         raise
-    
+
     # Return the server (like TypeScript version)
     return server
 
 
 def get_unknown_or_stdio_session(server: Any) -> dict:
     """Get or create an unknown/STDIO session for the server."""
     from .modules.session import get_or_create_unknown_session
-    
+
     data = get_mcpcat_data(server)
     if not data:
         raise Exception("Server tracking data not found")
-    
+
     session_id = get_or_create_unknown_session(data)
     return {
         "sessionId": session_id,
         "expiresIn": data.unknown_session.last_used if data.unknown_session else None
     }
 
 
-def _getServerTrackingData(server: Any) -> Optional[MCPCatData]:
+def _getServerTrackingData(server: Any) -> MCPCatData | None:
     """Get server tracking data (for testing)."""
     return get_mcpcat_data(server)
 
@@ -108,19 +108,19 @@ async def handleReportMissing(arguments: dict) -> None:
     """Handle report missing tool call (for testing)."""
     from .modules.logging import log_info
     from .types import MCPCatOptions
-    
+
     # Create minimal options for logging
     options = MCPCatOptions()
-    
+
     log_info("Missing tool reported", arguments, options)
 
 
 __all__ = [
     "track",
     "get_unknown_or_stdio_session",
-    "_getServerTrackingData", 
+    "_getServerTrackingData",
     "handleReportMissing",
-    "MCPCatOptions", 
+    "MCPCatOptions",
     "UserData",
     "__version__",
-]
+]

src/mcpcat/modules/__init__.py

@@ -1,20 +1,22 @@
 """MCPCat modules."""
 
 from .compatibility import is_compatible_server, is_fastmcp_server
-from .context_parameters import add_context_parameter_to_schema, add_context_parameter_to_tools
+from .context_parameters import (
+    add_context_parameter_to_schema,
+    add_context_parameter_to_tools,
+)
 from .internal import get_mcpcat_data, has_mcpcat_data, set_mcpcat_data
-from .logging import log_error, log_info, log_tool_call, log_warning
-from .redaction import redact_sensitive_info
-from .session import get_unknown_or_stdio_session, capture_session_info
+from .logging import log_error, log_info, log_trace, log_warning
+from .session import capture_session_info, get_unknown_or_stdio_session
 from .tools import handle_report_missing, setup_tool_handlers
-from .tracing import trace_tool_call
+from .tracing import record_trace
 
 __all__ = [
     # Compatibility
     "is_compatible_server",
     "is_fastmcp_server",
     # Context parameters
-    "add_context_parameter_to_schema", 
+    "add_context_parameter_to_schema",
     "add_context_parameter_to_tools",
     # Internal
     "get_mcpcat_data",
@@ -23,16 +25,15 @@
     # Logging
     "log_error",
     "log_info",
-    "log_tool_call",
+    "log_trace",
     "log_warning",
     # Redaction
-    "redact_sensitive_info",
     # Session
     "get_unknown_or_stdio_session",
     "capture_session_info",
     # Tools
     "handle_report_missing",
     "setup_tool_handlers",
     # Tracing
-    "trace_tool_call",
-]
+    "record_trace",
+]

src/mcpcat/modules/compatibility.py

@@ -6,11 +6,11 @@
 @runtime_checkable
 class MCPServerProtocol(Protocol):
     """Protocol for MCP server compatibility."""
-    
+
     def list_tools(self) -> Any:
         """List available tools."""
         ...
-    
+
     def call_tool(self, name: str, arguments: dict) -> Any:
         """Call a tool by name."""
         ...
@@ -21,12 +21,12 @@ def is_fastmcp_server(server: Any) -> bool:
     # Check for FastMCP class name or specific attributes
     return (
         server.__class__.__name__ == "FastMCP" or
-        (hasattr(server, "list_tools") and 
+        (hasattr(server, "list_tools") and
          hasattr(server, "call_tool") and
          hasattr(server, "_tool_manager"))
     )
 
 
 def is_compatible_server(server: Any) -> bool:
     """Check if the server is compatible with MCPCat."""
-    return is_fastmcp_server(server) or isinstance(server, MCPServerProtocol)
+    return is_fastmcp_server(server) or isinstance(server, MCPServerProtocol)

src/mcpcat/modules/context_parameters.py

@@ -1,52 +1,52 @@
 """Context parameter injection for MCP tools."""
 
-from typing import Any, Dict, List
+from typing import Any
 
 
-def add_context_parameter_to_tools(tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def add_context_parameter_to_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
     """Add context parameter to tool schemas."""
     modified_tools = []
-    
+
     for tool in tools:
         # Create a copy to avoid modifying original
         modified_tool = tool.copy()
-        
+
         if "inputSchema" in modified_tool:
             modified_tool["inputSchema"] = add_context_parameter_to_schema(
                 modified_tool["inputSchema"]
             )
-        
+
         modified_tools.append(modified_tool)
-    
+
     return modified_tools
 
 
-def add_context_parameter_to_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
+def add_context_parameter_to_schema(schema: dict[str, Any]) -> dict[str, Any]:
     """Add context parameter to a JSON schema."""
     # Create a copy to avoid modifying original
     modified_schema = schema.copy()
-    
+
     # Ensure properties exists
     if "properties" not in modified_schema:
         modified_schema["properties"] = {}
     else:
         # Deep copy properties
         modified_schema["properties"] = modified_schema["properties"].copy()
-    
+
     # Add context parameter
     modified_schema["properties"]["context"] = {
         "type": "string",
         "description": "Describe why you are calling this tool and how it fits into your overall task"
     }
-    
+
     # Add to required fields
     if "required" not in modified_schema:
         modified_schema["required"] = []
     else:
         # Copy required list
         modified_schema["required"] = list(modified_schema["required"])
-    
+
     if "context" not in modified_schema["required"]:
         modified_schema["required"].append("context")
-    
-    return modified_schema
+
+    return modified_schema

src/mcpcat/modules/internal.py

@@ -1,7 +1,7 @@
 """Internal data storage for MCPCat."""
 
 import weakref
-from typing import Any, Optional
+from typing import Any
 
 from ..types import MCPCatData
 
@@ -14,11 +14,11 @@ def set_mcpcat_data(server: Any, data: MCPCatData) -> None:
     _server_data_map[server] = data
 
 
-def get_mcpcat_data(server: Any) -> Optional[MCPCatData]:
+def get_mcpcat_data(server: Any) -> MCPCatData | None:
     """Retrieve MCPCat data for a server instance."""
     return _server_data_map.get(server)
 
 
 def has_mcpcat_data(server: Any) -> bool:
     """Check if a server instance has MCPCat data."""
-    return server in _server_data_map
+    return server in _server_data_map