From 4d969844c88c1893b263478ce5ef627c69068272 Mon Sep 17 00:00:00 2001
From: openhands <openhands@all-hands.dev>
Date: Tue, 22 Apr 2025 15:13:20 +0000
Subject: [PATCH] Update iter_tool_classes to only iterate over tools in the
 same module by default

---
 src/serena/agent.py     | 14 ++++++--
 test/serena/test_mcp.py |  2 +-
 verify_fix.py           | 75 -----------------------------------------
 3 files changed, 13 insertions(+), 78 deletions(-)
 delete mode 100644 verify_fix.py

diff --git a/src/serena/agent.py b/src/serena/agent.py
index 80e80908..86e93a92 100644
--- a/src/serena/agent.py
+++ b/src/serena/agent.py
@@ -1370,8 +1370,18 @@ class InitialInstructionsTool(Tool):
         return self.agent.prompt_factory.create_system_prompt()
 
 
-def iter_tool_classes() -> Generator[type[Tool], None, None]:
-    return iter_subclasses(Tool)
+def iter_tool_classes(same_module_only: bool = True) -> Generator[type[Tool], None, None]:
+    """
+    Iterate over all Tool subclasses.
+    
+    Args:
+        same_module_only: If True, only iterate over tools defined in the same module as the Tool class.
+                         If False, iterate over all Tool subclasses.
+    """
+    for tool_class in iter_subclasses(Tool):
+        if same_module_only and tool_class.__module__ != Tool.__module__:
+            continue
+        yield tool_class
 
 
 def print_tool_overview() -> None:
diff --git a/test/serena/test_mcp.py b/test/serena/test_mcp.py
index 0a940d33..bac38da2 100644
--- a/test/serena/test_mcp.py
+++ b/test/serena/test_mcp.py
@@ -273,7 +273,7 @@ def get_real_tool_classes():
     """Get all non-test, non-abstract tool classes."""
     from serena.agent import iter_tool_classes
 
-    for tool_class in iter_tool_classes():
+    for tool_class in iter_tool_classes(same_module_only=False):
         # Skip abstract base classes that can't be instantiated
         if tool_class.__name__ == "Tool" or getattr(tool_class, "__abstractmethods__", set()):
             continue
diff --git a/verify_fix.py b/verify_fix.py
deleted file mode 100644
index d0a2833d..00000000
--- a/verify_fix.py
+++ /dev/null
@@ -1,75 +0,0 @@
-"""Simple script to verify our fix for the make_tool function."""
-
-import docstring_parser
-from typing import Any, Optional
-
-
-def make_tool_simplified(func_doc: str) -> str:
-    """Simplified version of make_tool function to test our fix."""
-    docstring = docstring_parser.parse(func_doc)
-
-    # Mount the tool description as a combination of the docstring description and
-    # the return value description, if it exists.
-    if docstring.description:
-        func_doc = f"{docstring.description.strip().strip('.')}."
-    else:
-        func_doc = ""
-    if (docstring.returns) and (docstring_returns := docstring.returns.description):
-        # Only add a space before "Returns" if func_doc is not empty
-        prefix = " " if func_doc else ""
-        func_doc = f"{func_doc}{prefix}Returns {docstring_returns.strip().strip('.')}."
-    
-    return func_doc
-
-
-def test_with_description():
-    """Test with a normal docstring that has a description."""
-    doc = """This is a test function.
-
-    :param name: The person's name
-    :param age: The person's age
-    :return: A greeting message
-    """
-    result = make_tool_simplified(doc)
-    assert result == "This is a test function. Returns A greeting message."
-    print("✅ Test with description passed")
-
-
-def test_without_description():
-    """Test with a docstring that has no description."""
-    doc = """
-    :param name: The person's name
-    :param age: The person's age
-    :return: A greeting message
-    """
-    result = make_tool_simplified(doc)
-    assert result == "Returns A greeting message."
-    print("✅ Test without description passed")
-
-
-def test_empty_docstring():
-    """Test with an empty docstring."""
-    doc = ""
-    result = make_tool_simplified(doc)
-    assert result == ""
-    print("✅ Test with empty docstring passed")
-
-
-def test_no_return():
-    """Test with a docstring that has no return description."""
-    doc = """This is a test function.
-
-    :param name: The person's name
-    :param age: The person's age
-    """
-    result = make_tool_simplified(doc)
-    assert result == "This is a test function."
-    print("✅ Test with no return passed")
-
-
-if __name__ == "__main__":
-    test_with_description()
-    test_without_description()
-    test_empty_docstring()
-    test_no_return()
-    print("All tests passed!")
\ No newline at end of file