Skip to content

Commit 4d50565

Browse files
naji247naji247
committed
working with fastmcp and server (maybe)
1 parent 8ecc7c5 commit 4d50565

File tree

9 files changed

+256
-198
lines changed

9 files changed

+256
-198
lines changed

‎src/mcpcat/__init__.py‎

Lines changed: 12 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,17 @@
22

33
from typing import Any, Optional
44

5-
from .modules.compatibility import is_compatible_server
5+
from mcpcat.modules.overrides.fastmcp import override_fastmcp
6+
from mcpcat.modules.overrides.mcp_server import override_lowlevel_mcp_server
7+
8+
from .modules.compatibility import is_compatible_server, is_fastmcp_server
69
from .modules.internal import get_mcpcat_data, has_mcpcat_data, set_mcpcat_data
710
from .modules.logging import log_error, log_info
8-
from .modules.tools import setup_tool_handlers
11+
from .modules.session import get_unknown_or_stdio_session
12+
from .modules.tools import handle_report_missing as handleReportMissing
913
from .types import MCPCatData, MCPCatOptions, UserData
1014

11-
__version__ = "0.1.0"
15+
__version__ = "1.0.0"
1216

1317

1418
def track(server: Any, options: MCPCatOptions | None = None) -> Any:
@@ -32,7 +36,7 @@ def track(server: Any, options: MCPCatOptions | None = None) -> Any:
3236
# Validate server compatibility
3337
if not is_compatible_server(server):
3438
raise TypeError(
35-
"Server must be a FastMCP instance or implement the MCP server protocol"
39+
"Server must be a FastMCP instance or MCP Low-level Server instance"
3640
)
3741

3842
# Check if already tracked
@@ -50,7 +54,10 @@ def track(server: Any, options: MCPCatOptions | None = None) -> Any:
5054

5155
try:
5256
# Set up tool handlers
53-
setup_tool_handlers(server, data)
57+
if is_fastmcp_server(server):
58+
override_fastmcp(server, data)
59+
else:
60+
override_lowlevel_mcp_server(server, data)
5461

5562
# Log initialization
5663
log_info(
@@ -84,37 +91,10 @@ def track(server: Any, options: MCPCatOptions | None = None) -> Any:
8491
return server
8592

8693

87-
def get_unknown_or_stdio_session(server: Any) -> dict:
88-
"""Get or create an unknown/STDIO session for the server."""
89-
from .modules.session import get_or_create_unknown_session
90-
91-
data = get_mcpcat_data(server)
92-
if not data:
93-
raise Exception("Server tracking data not found")
94-
95-
session_id = get_or_create_unknown_session(data)
96-
return {
97-
"sessionId": session_id,
98-
"expiresIn": data.unknown_session.last_used if data.unknown_session else None
99-
}
100-
101-
10294
def _getServerTrackingData(server: Any) -> MCPCatData | None:
10395
"""Get server tracking data (for testing)."""
10496
return get_mcpcat_data(server)
10597

106-
107-
async def handleReportMissing(arguments: dict) -> None:
108-
"""Handle report missing tool call (for testing)."""
109-
from .modules.logging import log_info
110-
from .types import MCPCatOptions
111-
112-
# Create minimal options for logging
113-
options = MCPCatOptions()
114-
115-
log_info("Missing tool reported", arguments, options)
116-
117-
11898
__all__ = [
11999
"track",
120100
"get_unknown_or_stdio_session",

‎src/mcpcat/modules/__init__.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from .internal import get_mcpcat_data, has_mcpcat_data, set_mcpcat_data
99
from .logging import log_error, log_info, log_trace, log_warning
1010
from .session import capture_session_info, get_unknown_or_stdio_session
11-
from .tools import handle_report_missing, setup_tool_handlers
11+
from .tools import handle_report_missing
1212
from .tracing import record_trace
1313

1414
__all__ = [
@@ -33,7 +33,6 @@
3333
"capture_session_info",
3434
# Tools
3535
"handle_report_missing",
36-
"setup_tool_handlers",
3736
# Tracing
3837
"record_trace",
3938
]

‎src/mcpcat/modules/compatibility.py‎

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,45 @@ def call_tool(self, name: str, arguments: dict) -> Any:
1919
def is_fastmcp_server(server: Any) -> bool:
2020
"""Check if the server is a FastMCP instance."""
2121
# Check for FastMCP class name or specific attributes
22-
return (
23-
server.__class__.__name__ == "FastMCP" or
24-
(hasattr(server, "list_tools") and
25-
hasattr(server, "call_tool") and
26-
hasattr(server, "_tool_manager"))
27-
)
22+
return hasattr(server, "_mcp_server")
23+
24+
def has_neccessary_attributes(server: Any) -> bool:
25+
"""Check if the server has necessary attributes for compatibility."""
26+
required_methods = ["list_tools", "call_tool"]
27+
28+
# Check for core methods that both FastMCP and Server implementations have
29+
for method in required_methods:
30+
if not hasattr(server, method):
31+
return False
32+
33+
# For FastMCP servers, verify internal MCP server exists
34+
if hasattr(server, "_mcp_server"):
35+
# FastMCP server - check that internal MCP server has request_context
36+
# Use dir() to avoid triggering property getters that might raise exceptions
37+
if "request_context" not in dir(server._mcp_server):
38+
return False
39+
# Check for get_context method which is FastMCP specific
40+
if not hasattr(server, "get_context"):
41+
return False
42+
# Check for request_handlers dictionary on internal server
43+
if not hasattr(server._mcp_server, "request_handlers"):
44+
return False
45+
if not isinstance(server._mcp_server.request_handlers, dict):
46+
return False
47+
else:
48+
# Regular Server implementation - check for request_context directly
49+
# Use dir() to avoid triggering property getters that might raise exceptions
50+
if "request_context" not in dir(server):
51+
return False
52+
# Check for request_handlers dictionary
53+
if not hasattr(server, "request_handlers"):
54+
return False
55+
if not isinstance(server.request_handlers, dict):
56+
return False
57+
58+
return True
2859

2960

3061
def is_compatible_server(server: Any) -> bool:
3162
"""Check if the server is compatible with MCPCat."""
32-
return is_fastmcp_server(server) or isinstance(server, MCPServerProtocol)
63+
return has_neccessary_attributes(server)

‎src/mcpcat/modules/internal.py‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,37 @@
44
from typing import Any
55

66
from ..types import MCPCatData
7+
from .compatibility import is_fastmcp_server
78

89
# WeakKeyDictionary to store data associated with server instances
910
_server_data_map: weakref.WeakKeyDictionary[Any, MCPCatData] = weakref.WeakKeyDictionary()
1011

1112

1213
def set_mcpcat_data(server: Any, data: MCPCatData) -> None:
1314
"""Store MCPCat data for a server instance."""
14-
_server_data_map[server] = data
15+
# Always use low-level server as key
16+
if is_fastmcp_server(server):
17+
key = server._mcp_server
18+
else:
19+
key = server
20+
_server_data_map[key] = data
1521

1622

1723
def get_mcpcat_data(server: Any) -> MCPCatData | None:
1824
"""Retrieve MCPCat data for a server instance."""
19-
return _server_data_map.get(server)
25+
# Always use low-level server as key
26+
if is_fastmcp_server(server):
27+
key = server._mcp_server
28+
else:
29+
key = server
30+
return _server_data_map.get(key)
2031

2132

2233
def has_mcpcat_data(server: Any) -> bool:
2334
"""Check if a server instance has MCPCat data."""
24-
return server in _server_data_map
35+
# Always use low-level server as key
36+
if is_fastmcp_server(server):
37+
key = server._mcp_server
38+
else:
39+
key = server
40+
return key in _server_data_map

‎src/mcpcat/modules/overrides/fastmcp.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,21 @@
1+
from typing import Any
12

2-
"""Tool management and interception for MCPCat."""
3+
from mcp import ServerResult, Tool
4+
from mcp.server import FastMCP
5+
from mcp.types import CallToolRequest, CallToolResult, ListToolsRequest, TextContent
6+
7+
from mcpcat.modules.overrides.mcp_server import override_lowlevel_mcp_server
8+
from mcpcat.modules.tools import handle_report_missing
9+
from mcpcat.modules.tracing import record_trace
310

11+
from ...types import MCPCatData
12+
from ..logging import log_info, log_warning
13+
from ..session import capture_session_info
414

15+
"""Tool management and interception for MCPCat."""
516

617

18+
def override_fastmcp(server: FastMCP, data: MCPCatData) -> None:
19+
"""Set up tool list and call handlers for FastMCP."""
20+
from mcp.types import CallToolResult, ListToolsResult
21+
override_lowlevel_mcp_server(server._mcp_server, data)

‎src/mcpcat/modules/overrides/mcp_server.py‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
from typing import Any
2+
3+
from mcp import ServerResult, Tool
4+
from mcp.server.lowlevel import Server
5+
from mcp.types import CallToolRequest, CallToolResult, ListToolsRequest, TextContent
6+
7+
from mcpcat.modules.tools import handle_report_missing
8+
from mcpcat.modules.tracing import record_trace
9+
10+
from ...types import MCPCatData
11+
from ..logging import log_info, log_warning
12+
from ..session import capture_session_info
13+
14+
"""Tool management and interception for MCPCat."""
15+
16+
17+
def override_lowlevel_mcp_server(server: Server, data: MCPCatData) -> None:
18+
"""Set up tool list and call handlers for FastMCP."""
19+
from mcp.types import CallToolResult, ListToolsResult
20+
21+
# Store original request handlers - we only need to intercept at the low-level
22+
original_call_tool_handler = server.request_handlers.get(CallToolRequest)
23+
original_list_tools_handler = server.request_handlers.get(ListToolsRequest)
24+
25+
async def wrapped_list_tools_handler(request: ListToolsRequest) -> ServerResult:
26+
"""Intercept list_tools requests to add MCPCat tools and modify existing ones."""
27+
# Call the original handler to get the tools
28+
original_result = await original_list_tools_handler(request)
29+
if not original_result or not hasattr(original_result, 'root') or not hasattr(original_result.root, 'tools'):
30+
return original_result
31+
tools_list = original_result.root.tools
32+
33+
# Add report_missing tool if enabled
34+
if data.options.enableReportMissing:
35+
report_missing_tool = Tool(
36+
name="report_missing",
37+
description="Report when a tool you need is missing from this server",
38+
inputSchema={
39+
"type": "object",
40+
"properties": {
41+
"missing_tool": {
42+
"type": "string",
43+
"description": "Name of the missing tool"
44+
},
45+
"description": {
46+
"type": "string",
47+
"description": "Description of what the tool should do"
48+
}
49+
},
50+
"required": ["missing_tool", "description"]
51+
}
52+
)
53+
tools_list.append(report_missing_tool)
54+
55+
# Add context parameters to existing tools if enabled
56+
if data.options.enableToolCallContext:
57+
for tool in tools_list:
58+
if tool.name != "report_missing": # Don't modify our own tool
59+
if not tool.inputSchema:
60+
tool.inputSchema = {
61+
"type": "object",
62+
"properties": {},
63+
"required": []
64+
}
65+
66+
# Add context property if it doesn't exist
67+
if "context" not in tool.inputSchema.get("properties", {}):
68+
if "properties" not in tool.inputSchema:
69+
tool.inputSchema["properties"] = {}
70+
71+
tool.inputSchema["properties"]["context"] = {
72+
"type": "string",
73+
"description": "Describe why you are calling this tool and how it fits into your overall task"
74+
}
75+
76+
# Add context to required array if it exists
77+
if isinstance(tool.inputSchema.get("required"), list):
78+
if "context" not in tool.inputSchema["required"]:
79+
tool.inputSchema["required"].append("context")
80+
else:
81+
tool.inputSchema["required"] = ["context"]
82+
83+
return ServerResult(ListToolsResult(tools=tools_list))
84+
85+
async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult:
86+
"""Intercept call_tool requests to add MCPCat tracking and handle special tools."""
87+
tool_name = request.params.name
88+
arguments = request.params.arguments or {}
89+
90+
# Handle report_missing tool directly
91+
if tool_name == "report_missing":
92+
return await handle_report_missing(arguments, data)
93+
94+
# Extract MCPCat context if enabled
95+
mcpcat_user_context = None
96+
if data.options.enableToolCallContext:
97+
mcpcat_user_context = arguments.pop("context", None)
98+
# Log warning if context is missing and tool is not report_missing
99+
if mcpcat_user_context is None and tool_name != "report_missing":
100+
log_warning("Missing context parameter", {"tool_name": tool_name}, data.options)
101+
102+
# Get session info for tracking
103+
try:
104+
request_context = server.request_context
105+
session_id, user_id = await capture_session_info(server, arguments=arguments, request_context=request_context)
106+
except:
107+
request_context = None
108+
session_id, user_id = None, None
109+
110+
# If tracing is enabled, wrap the call with timing and logging
111+
if data.options.enableTracing:
112+
import time
113+
start_time = time.time()
114+
115+
try:
116+
# Call the original handler
117+
result = await original_call_tool_handler(request)
118+
duration = time.time() - start_time
119+
120+
# Record the trace using existing infrastructure
121+
await record_trace(
122+
server=server,
123+
tool_name=tool_name,
124+
arguments=arguments,
125+
request_context=request_context,
126+
session_id=session_id,
127+
user_id=user_id,
128+
tool_result=result.model_dump() if result else None,
129+
duration=duration,
130+
mcpcat_context=mcpcat_user_context
131+
)
132+
133+
return result
134+
135+
except Exception as e:
136+
duration = time.time() - start_time
137+
138+
# Record the error trace
139+
await record_trace(
140+
server=server,
141+
tool_name=tool_name,
142+
arguments=arguments,
143+
request_context=request_context,
144+
session_id=session_id,
145+
user_id=user_id,
146+
tool_result=None,
147+
duration=duration,
148+
mcpcat_context=mcpcat_user_context,
149+
error=str(e)
150+
)
151+
# Re-raise the exception
152+
raise
153+
else:
154+
# No tracing, just call the original handler
155+
return await original_call_tool_handler(request)
156+
157+
# Replace only the low-level request handlers
158+
server.request_handlers[CallToolRequest] = wrapped_call_tool_handler
159+
server.request_handlers[ListToolsRequest] = wrapped_list_tools_handler
160+

0 commit comments

Comments
 (0)