From 92c1b0f76862cd54ae864f9b8f9ee2896da527c0 Mon Sep 17 00:00:00 2001
From: PM-pinou <2504420230@qq.com>
Date: Fri, 24 Jul 2026 12:45:29 +0800
Subject: [PATCH] feat(tools): LLM-mandated _timeout parameter on every MCP
 tool, enforced in dispatch layer

---
 analysis/tools/_dispatch.py | 24 ++++++++++++++++++++++--
 analysis/tools/_registry.py | 29 ++++++++++++++++++++++++++---
 server_stderr.txt           |  6 ++++++
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/analysis/tools/_dispatch.py b/analysis/tools/_dispatch.py
index 5b04410..e1b5082 100644
--- a/analysis/tools/_dispatch.py
+++ b/analysis/tools/_dispatch.py
@@ -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
diff --git a/analysis/tools/_registry.py b/analysis/tools/_registry.py
index 9644eec..1bab357 100644
--- a/analysis/tools/_registry.py
+++ b/analysis/tools/_registry.py
@@ -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
diff --git a/server_stderr.txt b/server_stderr.txt
index 3582d34..1c0f409 100644
--- a/server_stderr.txt
+++ b/server_stderr.txt
@@ -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