|
1 | 1 | """Internal data storage for MCPCat.""" |
2 | 2 |
|
3 | 3 | import weakref |
4 | | -from typing import Any |
| 4 | +from datetime import datetime, timezone |
| 5 | +from typing import Any, Dict, List, Optional |
5 | 6 |
|
6 | | -from ..types import MCPCatData |
| 7 | +from ..types import MCPCatData, ToolRegistration |
7 | 8 | from .compatibility import is_fastmcp_server |
| 9 | +from .logging import write_to_log |
8 | 10 |
|
9 | 11 | # WeakKeyDictionary to store data associated with server instances |
10 | 12 | _server_data_map: weakref.WeakKeyDictionary[Any, MCPCatData] = weakref.WeakKeyDictionary() |
11 | 13 |
|
| 14 | +# Global storage for original unpatched methods (keyed by tool_manager or handler id) |
| 15 | +# This is global because tool managers might be shared between servers |
| 16 | +_original_methods: Dict[str, Any] = {} |
| 17 | + |
| 18 | + |
| 19 | +def _get_server_key(server: Any) -> Any: |
| 20 | + """Get the canonical key for a server (handles FastMCP vs low-level).""" |
| 21 | + if is_fastmcp_server(server): |
| 22 | + return server._mcp_server |
| 23 | + return server |
| 24 | + |
12 | 25 |
|
13 | 26 | def set_server_tracking_data(server: Any, data: MCPCatData) -> None: |
14 | 27 | """Store MCPCat data for a server instance.""" |
15 | | - # Always use low-level server as key |
16 | | - if is_fastmcp_server(server): |
17 | | - key = server._mcp_server |
18 | | - else: |
19 | | - key = server |
| 28 | + key = _get_server_key(server) |
20 | 29 | _server_data_map[key] = data |
21 | 30 |
|
22 | 31 |
|
23 | 32 | def get_server_tracking_data(server: Any) -> MCPCatData | None: |
24 | 33 | """Retrieve MCPCat data for a server instance.""" |
25 | | - # Always use low-level server as key |
26 | | - if is_fastmcp_server(server): |
27 | | - key = server._mcp_server |
28 | | - else: |
29 | | - key = server |
| 34 | + key = _get_server_key(server) |
30 | 35 | return _server_data_map.get(key, None) |
| 36 | + |
| 37 | + |
| 38 | +def reset_server_tracking_data(server: Any) -> None: |
| 39 | + """Reset tracking data for a specific server (mainly for testing).""" |
| 40 | + key = _get_server_key(server) |
| 41 | + if key in _server_data_map: |
| 42 | + del _server_data_map[key] |
| 43 | + write_to_log(f"Reset tracking data for server {id(key)}") |
| 44 | + |
| 45 | + |
| 46 | +def reset_all_tracking_data() -> None: |
| 47 | + """Reset all server tracking data (mainly for testing).""" |
| 48 | + _server_data_map.clear() |
| 49 | + _original_methods.clear() |
| 50 | + write_to_log("Reset all server tracking data") |
| 51 | + |
| 52 | + |
| 53 | +# Dynamic tracking helper methods |
| 54 | +def register_tool(server: Any, name: str) -> None: |
| 55 | + """Register a tool in the server's tracking system.""" |
| 56 | + data = get_server_tracking_data(server) |
| 57 | + if data and name not in data.tool_registry: |
| 58 | + data.tool_registry[name] = ToolRegistration( |
| 59 | + name=name, |
| 60 | + registered_at=datetime.now(timezone.utc) |
| 61 | + ) |
| 62 | + write_to_log(f"Registered tool '{name}'") |
| 63 | + |
| 64 | + |
| 65 | +def mark_tool_tracked(server: Any, name: str) -> None: |
| 66 | + """Mark a tool as being tracked by MCPCat for this server.""" |
| 67 | + data = get_server_tracking_data(server) |
| 68 | + if data and name in data.tool_registry: |
| 69 | + data.tool_registry[name].tracked = True |
| 70 | + data.tool_registry[name].wrapped = True |
| 71 | + data.wrapped_tools.add(name) |
| 72 | + |
| 73 | + |
| 74 | +def is_tool_tracked(server: Any, name: str) -> bool: |
| 75 | + """Check if a tool is already being tracked for this server.""" |
| 76 | + data = get_server_tracking_data(server) |
| 77 | + return data and name in data.wrapped_tools |
| 78 | + |
| 79 | + |
| 80 | +def get_untracked_tools(server: Any) -> List[str]: |
| 81 | + """Get list of tools that aren't tracked yet for this server.""" |
| 82 | + data = get_server_tracking_data(server) |
| 83 | + if not data: |
| 84 | + return [] |
| 85 | + return [ |
| 86 | + name for name, reg in data.tool_registry.items() |
| 87 | + if not reg.tracked |
| 88 | + ] |
| 89 | + |
| 90 | + |
| 91 | +def discover_new_tools(server: Any, tools: List[Any]) -> List[str]: |
| 92 | + """Discover tools that weren't previously known for this server.""" |
| 93 | + data = get_server_tracking_data(server) |
| 94 | + if not data: |
| 95 | + return [] |
| 96 | + |
| 97 | + new_tools = [] |
| 98 | + for tool in tools: |
| 99 | + if tool.name not in data.tool_registry: |
| 100 | + register_tool(server, tool.name) |
| 101 | + new_tools.append(tool.name) |
| 102 | + return new_tools |
| 103 | + |
| 104 | + |
| 105 | +# Original methods storage (global, not per-server) |
| 106 | +def store_original_method(key: str, method: Any) -> None: |
| 107 | + """Store an original unpatched method.""" |
| 108 | + if key not in _original_methods: |
| 109 | + _original_methods[key] = method |
| 110 | + |
| 111 | + |
| 112 | +def get_original_method(key: str) -> Optional[Any]: |
| 113 | + """Get an original unpatched method.""" |
| 114 | + return _original_methods.get(key) |
| 115 | + |
| 116 | + |
| 117 | +def get_original_methods() -> Dict[str, Any]: |
| 118 | + """Get the global original methods storage.""" |
| 119 | + return _original_methods |
| 120 | + |
| 121 | + |
| 122 | +def get_tool_timeline(server: Any) -> List[Dict[str, Any]]: |
| 123 | + """Get a timeline of tool registrations for debugging. |
| 124 | + |
| 125 | + Args: |
| 126 | + server: MCP server instance |
| 127 | + |
| 128 | + Returns: |
| 129 | + List of tool registration events sorted by time |
| 130 | + """ |
| 131 | + data = get_server_tracking_data(server) |
| 132 | + if not data: |
| 133 | + return [] |
| 134 | + timeline = [] |
| 135 | + |
| 136 | + for name, reg in data.tool_registry.items(): |
| 137 | + timeline.append({ |
| 138 | + "name": name, |
| 139 | + "registered_at": reg.registered_at.isoformat(), |
| 140 | + "tracked": reg.tracked, |
| 141 | + "wrapped": reg.wrapped |
| 142 | + }) |
| 143 | + |
| 144 | + # Sort by registration time |
| 145 | + timeline.sort(key=lambda x: x["registered_at"]) |
| 146 | + |
| 147 | + return timeline |
0 commit comments