diff --git a/EvoScientist/EvoScientist.py b/EvoScientist/EvoScientist.py index 0b1614d..e4b0796 100644 --- a/EvoScientist/EvoScientist.py +++ b/EvoScientist/EvoScientist.py @@ -24,7 +24,7 @@ from .backends import CustomSandboxBackend, MergedReadOnlyBackend from .config import get_effective_config, apply_config_to_env from .llm import get_chat_model from .mcp import load_mcp_tools -from .middleware import create_memory_middleware +from .middleware import create_memory_middleware, ToolErrorHandlerMiddleware from .prompts import RESEARCHER_INSTRUCTIONS, get_system_prompt from .utils import load_subagents from .tools import tavily_search, think_tool, skill_manager @@ -219,6 +219,7 @@ prompt_refs = { } base_middleware = [ + ToolErrorHandlerMiddleware(), create_memory_middleware(MEMORY_DIR, extraction_model=chat_model), ] @@ -306,6 +307,7 @@ def create_cli_agent(workspace_dir: str | None = None, checkpointer=None): ) mw = [ + ToolErrorHandlerMiddleware(), create_memory_middleware(_mem_dir, extraction_model=chat_model), ] diff --git a/EvoScientist/middleware/__init__.py b/EvoScientist/middleware/__init__.py index 45c07e2..93e4dbb 100644 --- a/EvoScientist/middleware/__init__.py +++ b/EvoScientist/middleware/__init__.py @@ -10,10 +10,12 @@ from .memory import ( ExtractedMemory, create_memory_middleware, ) +from .tool_error_handler import ToolErrorHandlerMiddleware __all__ = [ "EvoMemoryMiddleware", "EvoMemoryState", "ExtractedMemory", + "ToolErrorHandlerMiddleware", "create_memory_middleware", ] diff --git a/EvoScientist/middleware/tool_error_handler.py b/EvoScientist/middleware/tool_error_handler.py new file mode 100644 index 0000000..63d90df --- /dev/null +++ b/EvoScientist/middleware/tool_error_handler.py @@ -0,0 +1,70 @@ +"""Middleware that catches tool execution exceptions and converts them to error ToolMessages. + +Without this, an MCP tool (or any tool) that raises an exception at runtime +crashes the entire agent loop because LangGraph's default ToolNode error handler +only catches argument-validation errors (ToolInvocationError), not execution +errors. + +With this middleware, the exception is caught and surfaced to the agent as a +ToolMessage with ``status="error"`` containing the traceback. The agent can +then decide to retry, use a different tool, or yield to the user. +""" + +from __future__ import annotations + +import logging +import traceback +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any + +from langchain.agents.middleware.types import AgentMiddleware +from langchain_core.messages import ToolMessage +from langgraph.types import Command + +if TYPE_CHECKING: + from langchain.agents.middleware.types import ToolCallRequest + +logger = logging.getLogger(__name__) + + +class ToolErrorHandlerMiddleware(AgentMiddleware): + """Catch tool execution exceptions and return them as error ToolMessages.""" + + name = "tool_error_handler" + + def wrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]], + ) -> ToolMessage | Command[Any]: + try: + return handler(request) + except Exception: + return _build_error_message(request) + + async def awrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]], + ) -> ToolMessage | Command[Any]: + try: + return await handler(request) + except Exception: + return _build_error_message(request) + + +def _build_error_message(request: ToolCallRequest) -> ToolMessage: + tb = traceback.format_exc() + tool_name = request.tool_call.get("name", "unknown_tool") + logger.error("Tool %r raised an exception:\n%s", tool_name, tb) + content = ( + f"[TOOL ERROR] Tool '{tool_name}' failed with an exception:\n\n{tb}\n" + "You may retry the tool call, try an alternative approach, " + "or inform the user about the failure." + ) + return ToolMessage( + content=content, + tool_call_id=request.tool_call["id"], + name=tool_name, + status="error", + ) diff --git a/tests/test_tool_error_handler.py b/tests/test_tool_error_handler.py new file mode 100644 index 0000000..c90b0bb --- /dev/null +++ b/tests/test_tool_error_handler.py @@ -0,0 +1,219 @@ +"""Tests for ToolErrorHandlerMiddleware.""" + +import asyncio +from unittest.mock import MagicMock + +import pytest +from langchain_core.messages import ToolMessage +from langgraph.types import Command + +from EvoScientist.middleware.tool_error_handler import ( + ToolErrorHandlerMiddleware, + _build_error_message, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_request(tool_name: str = "my_mcp_tool", call_id: str = "tc_001"): + """Create a minimal ToolCallRequest-like object.""" + req = MagicMock() + req.tool_call = {"id": call_id, "name": tool_name, "args": {"query": "test"}} + return req + + +# --------------------------------------------------------------------------- +# _build_error_message +# --------------------------------------------------------------------------- + + +class TestBuildErrorMessage: + def test_returns_tool_message(self): + req = _make_request("broken_tool", "tc_123") + try: + raise RuntimeError("connection refused") + except RuntimeError: + msg = _build_error_message(req) + + assert isinstance(msg, ToolMessage) + assert msg.tool_call_id == "tc_123" + assert msg.name == "broken_tool" + assert msg.status == "error" + + def test_content_includes_tool_name(self): + req = _make_request("broken_tool") + try: + raise ValueError("bad value") + except ValueError: + msg = _build_error_message(req) + + assert "broken_tool" in msg.content + + def test_content_includes_traceback(self): + req = _make_request() + try: + raise RuntimeError("something went wrong") + except RuntimeError: + msg = _build_error_message(req) + + assert "RuntimeError" in msg.content + assert "something went wrong" in msg.content + assert "Traceback" in msg.content + + def test_content_includes_retry_guidance(self): + req = _make_request() + try: + raise Exception("fail") + except Exception: + msg = _build_error_message(req) + + assert "retry" in msg.content.lower() + + +# --------------------------------------------------------------------------- +# Sync: wrap_tool_call +# --------------------------------------------------------------------------- + + +class TestWrapToolCallSync: + def setup_method(self): + self.mw = ToolErrorHandlerMiddleware() + + def test_success_passes_through(self): + expected = ToolMessage(content="ok", tool_call_id="tc_001", name="t") + handler = MagicMock(return_value=expected) + req = _make_request() + + result = self.mw.wrap_tool_call(req, handler) + + assert result is expected + handler.assert_called_once_with(req) + + def test_command_passes_through(self): + cmd = Command(update={"messages": []}) + handler = MagicMock(return_value=cmd) + req = _make_request() + + result = self.mw.wrap_tool_call(req, handler) + + assert result is cmd + + def test_exception_returns_error_tool_message(self): + handler = MagicMock(side_effect=RuntimeError("MCP server crashed")) + req = _make_request("flaky_tool", "tc_999") + + result = self.mw.wrap_tool_call(req, handler) + + assert isinstance(result, ToolMessage) + assert result.status == "error" + assert result.tool_call_id == "tc_999" + assert result.name == "flaky_tool" + assert "MCP server crashed" in result.content + + def test_exception_does_not_propagate(self): + handler = MagicMock(side_effect=ConnectionError("refused")) + req = _make_request() + + # Should NOT raise + result = self.mw.wrap_tool_call(req, handler) + assert isinstance(result, ToolMessage) + + def test_keyboard_interrupt_propagates(self): + """KeyboardInterrupt should NOT be caught (it's BaseException, not Exception).""" + handler = MagicMock(side_effect=KeyboardInterrupt()) + req = _make_request() + + with pytest.raises(KeyboardInterrupt): + self.mw.wrap_tool_call(req, handler) + + def test_various_exception_types(self): + """Different exception types are all caught and reported.""" + for exc_cls in (ValueError, TypeError, OSError, TimeoutError, ConnectionError): + handler = MagicMock(side_effect=exc_cls(f"{exc_cls.__name__} happened")) + req = _make_request() + + result = self.mw.wrap_tool_call(req, handler) + + assert isinstance(result, ToolMessage) + assert result.status == "error" + assert exc_cls.__name__ in result.content + + +# --------------------------------------------------------------------------- +# Async: awrap_tool_call +# --------------------------------------------------------------------------- + + +class TestWrapToolCallAsync: + def setup_method(self): + self.mw = ToolErrorHandlerMiddleware() + + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_success_passes_through(self): + expected = ToolMessage(content="ok", tool_call_id="tc_001", name="t") + + async def handler(req): + return expected + + req = _make_request() + result = self._run(self.mw.awrap_tool_call(req, handler)) + + assert result is expected + + def test_command_passes_through(self): + cmd = Command(update={"messages": []}) + + async def handler(req): + return cmd + + req = _make_request() + result = self._run(self.mw.awrap_tool_call(req, handler)) + + assert result is cmd + + def test_exception_returns_error_tool_message(self): + async def handler(req): + raise RuntimeError("MCP server timed out") + + req = _make_request("slow_tool", "tc_async") + result = self._run(self.mw.awrap_tool_call(req, handler)) + + assert isinstance(result, ToolMessage) + assert result.status == "error" + assert result.tool_call_id == "tc_async" + assert result.name == "slow_tool" + assert "MCP server timed out" in result.content + + def test_exception_does_not_propagate(self): + async def handler(req): + raise ConnectionError("connection lost") + + req = _make_request() + result = self._run(self.mw.awrap_tool_call(req, handler)) + assert isinstance(result, ToolMessage) + + def test_keyboard_interrupt_propagates(self): + async def handler(req): + raise KeyboardInterrupt() + + req = _make_request() + with pytest.raises(KeyboardInterrupt): + self._run(self.mw.awrap_tool_call(req, handler)) + + +# --------------------------------------------------------------------------- +# Middleware metadata +# --------------------------------------------------------------------------- + + +class TestMiddlewareMeta: + def test_name(self): + assert ToolErrorHandlerMiddleware.name == "tool_error_handler" + + def test_instantiation(self): + mw = ToolErrorHandlerMiddleware() + assert mw.name == "tool_error_handler"