From d4871e9d8826a0a64607fd5325d12cbea27e4a6f Mon Sep 17 00:00:00 2001
From: Dominik Jain <dominik.jain@outlook.com>
Date: Sun, 30 Mar 2025 00:25:04 +0100
Subject: [PATCH] Add Tool.apply_ex to handle logging and exceptions

---
 src/serena/agent.py | 23 +++++++++++++++++++++++
 src/serena/mcp.py   |  9 ++-------
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/serena/agent.py b/src/serena/agent.py
index b36eb12d..a6e44c70 100644
--- a/src/serena/agent.py
+++ b/src/serena/agent.py
@@ -7,6 +7,7 @@ import json
 import os
 import platform
 import sys
+import traceback
 from abc import ABC
 from collections import defaultdict
 from collections.abc import Iterator, Sequence
@@ -167,6 +168,28 @@ class Tool(Component):
             )
         return result
 
+    def apply_ex(self, *args, log_call: bool = True, catch_exceptions: bool = True, **kwargs) -> str:
+        """
+        Applies the tool with the given arguments
+        """
+        apply_fn = getattr(self, "apply")
+        if apply_fn is None:
+            raise ValueError(f"apply not defined in {self}")
+
+        if log_call:
+            self._log_tool_application(inspect.currentframe())
+        try:
+            result = apply_fn(*args, **kwargs)
+        except Exception as e:
+            if not catch_exceptions:
+                raise
+            msg = f"Error executing tool: {e}\n{traceback.format_exc()}"
+            log.error(f"Error executing tool: {e}", exc_info=e)
+            result = msg
+        if log_call:
+            log.info("Result: {}")
+        return result
+
 
 class ReadFileTool(Tool):
     def apply(
diff --git a/src/serena/mcp.py b/src/serena/mcp.py
index c2f7fb29..f8936a23 100644
--- a/src/serena/mcp.py
+++ b/src/serena/mcp.py
@@ -3,7 +3,6 @@ The Serena Model Context Protocol (MCP) Server
 """
 
 import sys
-import traceback
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -56,11 +55,7 @@ def make_tool(
 
     def execute_fn(ctx: Context, *args, **kwargs) -> str:  # type: ignore
         mark_used(ctx)
-        try:
-            return apply_fn(*args, **kwargs)
-        except Exception as e:
-            msg = f"Error executing tool: {e}\n{traceback.format_exc()}"
-            return msg
+        tool.apply_ex(*args, log_call=True, catch_exceptions=True, **kwargs)
 
     return MCPTool(
         fn=execute_fn,
@@ -96,5 +91,5 @@ def create_mcp_server() -> FastMCP:
     return mcp
 
 
-def start_mcp_server():
+def start_mcp_server() -> None:
     create_mcp_server().run()