feat(tools): LLM-mandated _timeout parameter on every MCP tool, enforced in dispatch layer
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
"""Tool call dispatcher — maps tool names to handler functions."""
|
||||
import asyncio
|
||||
from ._helpers import _truncate_response
|
||||
from .load_data import _handle_load_data
|
||||
from .profile import _handle_profile_data
|
||||
@@ -27,9 +28,12 @@ from .distance_matrix import _handle_compute_distance_matrix
|
||||
async def handle_call(name: str, arguments: dict) -> dict:
|
||||
"""Dispatch a tool call to the appropriate handler.
|
||||
|
||||
Extracts the mandatory ``_timeout`` parameter (set by the LLM) from
|
||||
*arguments* and enforces it as an ``asyncio.wait_for`` deadline.
|
||||
|
||||
Args:
|
||||
name: Tool name (must match one of the 30 tools).
|
||||
arguments: Tool-specific parameters.
|
||||
arguments: Tool-specific parameters, must include ``_timeout``.
|
||||
|
||||
Returns:
|
||||
A JSON-serialisable dict.
|
||||
@@ -72,7 +76,23 @@ async def handle_call(name: str, arguments: dict) -> dict:
|
||||
handler = _handlers.get(name)
|
||||
if handler is None:
|
||||
raise ValueError(f"Unknown tool: '{name}'")
|
||||
result = await handler(**arguments)
|
||||
|
||||
# Extract mandatory _timeout from the LLM-provided arguments
|
||||
timeout = arguments.pop('_timeout', 0)
|
||||
cleaned_args = arguments # _timeout already removed by pop
|
||||
|
||||
# Execute the handler with optional timeout
|
||||
if timeout and timeout > 0:
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
handler(**cleaned_args),
|
||||
timeout=timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
result = {'error': f'工具 {name} 执行超时 ({timeout}秒)'}
|
||||
else:
|
||||
result = await handler(**cleaned_args)
|
||||
|
||||
# Always ensure truncated key exists
|
||||
if 'truncated' not in result:
|
||||
result['truncated'] = False
|
||||
|
||||
@@ -8,9 +8,10 @@ from mcp.types import Tool
|
||||
def get_tools_meta() -> list[Tool]:
|
||||
"""Return the list of Tool metadata objects for MCP server registration.
|
||||
|
||||
Call this from :meth:`TLSAnalyzerMCPServer._setup_tools`.
|
||||
Each tool gets a mandatory ``_timeout`` parameter so the LLM must specify
|
||||
how many seconds to allow for each tool invocation (0 = no limit).
|
||||
"""
|
||||
return [
|
||||
raw_tools = [
|
||||
Tool(
|
||||
name="load_data",
|
||||
description=(
|
||||
@@ -751,4 +752,26 @@ def get_tools_meta() -> list[Tool]:
|
||||
"required": ["dataset_id", "python_function"],
|
||||
},
|
||||
),
|
||||
]
|
||||
] # end of raw_tools
|
||||
|
||||
# Inject mandatory _timeout parameter into every tool
|
||||
result = []
|
||||
for tool in raw_tools:
|
||||
schema = dict(tool.inputSchema)
|
||||
props = dict(schema.get("properties", {}))
|
||||
req = list(schema.get("required", []))
|
||||
# Add _timeout as the first required property
|
||||
props["_timeout"] = {
|
||||
"type": "integer",
|
||||
"description": "工具调用超时秒数(必填,LLM决定每次调用的最长等待时间;0表示不限制)",
|
||||
"default": 120,
|
||||
}
|
||||
req.insert(0, "_timeout")
|
||||
schema["properties"] = props
|
||||
schema["required"] = req
|
||||
result.append(Tool(
|
||||
name=tool.name,
|
||||
description=tool.description,
|
||||
inputSchema=schema,
|
||||
))
|
||||
return result
|
||||
|
||||
@@ -72,3 +72,9 @@ Restored 1 datasets from disk
|
||||
[24/Jul/2026 12:37:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:37:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:38:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:38:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:39:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:39:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:40:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:40:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:41:07] "HEAD / HTTP/1.1" 200 0
|
||||
|
||||
Reference in New Issue
Block a user