Add timeout for tool calls

Signed-off-by: tamohannes <hovhannes.tamoyan@gmail.com>

Author: tamohannes

nemo_skills/inference/generate.py

@@ -202,6 +202,7 @@ class GenerationTaskConfig:
     schema_overrides: dict | None = field(default_factory=dict)
 
     max_tool_calls: int = -1  # If >= 0, will limit the number of tool calls executed during generation to this number
+    tool_call_timeout_s: float | None = 300.0  # Wall-clock timeout for each individual tool call
 
     # if True, will move full generation to _full_generation key and keep cfg.generation_key without thinking tokens
     # IMPORTANT: do not set this for non-reasoning models as it will make the generations empty!
@@ -482,6 +483,7 @@ def setup_llm(self):
                 tool_overrides=self.cfg.tool_overrides,
                 schema_overrides=self.cfg.schema_overrides,
                 max_tool_calls=self.cfg.max_tool_calls,
+                tool_call_timeout_s=self.cfg.tool_call_timeout_s,
                 tokenizer=self.tokenizer,
                 require_tokenizer=self.cfg.inference.tokens_to_generate is not None,
                 additional_config={"sandbox": self.cfg.sandbox},

nemo_skills/inference/model/__init__.py

@@ -136,6 +136,7 @@ def get_tool_calling_model(
     tool_overrides: dict | None = None,
     schema_overrides: dict | None = None,
     max_tool_calls: int = -1,
+    tool_call_timeout_s: float | None = 300.0,
     **kwargs,
 ):
     if isinstance(model, str):
@@ -147,6 +148,7 @@ def get_tool_calling_model(
         additional_config=additional_config,
         schema_overrides=schema_overrides,
         max_tool_calls=max_tool_calls,
+        tool_call_timeout_s=tool_call_timeout_s,
     )
 
 

nemo_skills/inference/model/tool_call.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import copy
 import json
 import logging
@@ -49,6 +50,7 @@ def __init__(
         additional_config: dict | None = None,
         schema_overrides: dict | None = None,
         max_tool_calls: int = -1,
+        tool_call_timeout_s: float | None = 300.0,
     ):
         self.model = model
         additional_config = additional_config or {}
@@ -63,6 +65,7 @@ def __init__(
         self.schema_overrides = load_schema_overrides(schema_overrides)
         self.schema_mappings = {}  # Built when tools are listed
         self.max_tool_calls = max_tool_calls
+        self.tool_call_timeout_s = tool_call_timeout_s
 
     async def _execute_tool_call(self, tool_call, request_id: str, endpoint_type: EndpointType):
         ## TODO(sanyamk): The correct key format needs to be cohesive with other formatters.
@@ -85,9 +88,16 @@ async def _execute_tool_call(self, tool_call, request_id: str, endpoint_type: En
 
         try:
             # Allow providers to specify extra_args behavior internally if needed in the future
-            result = await self.tool_manager.execute_tool(
+            tool_future = self.tool_manager.execute_tool(
                 original_tool_name, tool_args, extra_args={"request_id": request_id}
             )
+            if self.tool_call_timeout_s is None:
+                result = await tool_future
+            else:
+                result = await asyncio.wait_for(tool_future, timeout=self.tool_call_timeout_s)
+        except asyncio.TimeoutError:
+            LOG.error("Tool execution timed out after %s seconds: %s", self.tool_call_timeout_s, original_tool_name)
+            return {"error": f"Tool execution timed out after {self.tool_call_timeout_s} seconds."}
         except FatalToolError:
             # Fatal errors should propagate up and stop the process
             raise

tests/test_mcp_clients.py

@@ -12,11 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import types
 
 import pytest
 
 # Dummy client to exercise MCPClientMeta behavior without real I/O
+from nemo_skills.inference.model.base import EndpointType
+from nemo_skills.inference.model.tool_call import ToolCallingWrapper
 from nemo_skills.mcp.clients import MCPClient, MCPStdioClient, MCPStreamableHttpClient
 from nemo_skills.mcp.tool_manager import Tool, ToolManager
 
@@ -219,6 +222,27 @@ async def execute(self, tool_name: str, arguments: dict, extra_args: dict | None
         return {"unknown": tool_name, "args": arguments}
 
 
+class SlowTool(Tool):
+    def default_config(self):
+        return {}
+
+    def configure(self, overrides=None, context=None):
+        return None
+
+    async def list_tools(self):
+        return [
+            {
+                "name": "sleep",
+                "description": "Sleep long enough to exercise timeout handling",
+                "input_schema": {"type": "object", "properties": {}},
+            }
+        ]
+
+    async def execute(self, tool_name: str, arguments: dict, extra_args: dict | None = None):
+        await asyncio.sleep(10)
+        return "too late"
+
+
 # Helper class for test_tool_manager_cache_and_duplicate_detection
 # Defined at module level so it can be imported via locate()
 class CountingTool(DummyTool):
@@ -255,6 +279,24 @@ async def test_tool_manager_list_and_execute_with_class_locator():
     assert result == {"ran": True, "code": "x=1"}
 
 
+@pytest.mark.asyncio
+async def test_tool_calling_wrapper_times_out_slow_tool_call():
+    wrapper = ToolCallingWrapper(
+        model=object(),
+        tool_modules=[f"{__name__}::SlowTool"],
+        tool_call_timeout_s=0.01,
+    )
+    await wrapper.tool_manager.list_all_tools(use_cache=False)
+
+    result = await wrapper._execute_tool_call(
+        {"id": "call-1", "function": {"name": "sleep", "arguments": "{}"}},
+        request_id="req-1",
+        endpoint_type=EndpointType.chat,
+    )
+
+    assert result["error"].startswith("Tool execution timed out")
+
+
 @pytest.mark.asyncio
 async def test_tool_manager_cache_and_duplicate_detection():
     import sys