feat(tools): LLM-mandated _timeout parameter on every MCP tool, enforced in dispatch layer

This commit is contained in:
PM-pinou
2026-07-24 12:45:29 +08:00
parent d24507e541
commit 92c1b0f768
3 changed files with 54 additions and 5 deletions
+22 -2
View File
@@ -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
+26 -3
View File
@@ -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
+6
View File
@@ -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