From 44926587028e71d37b2e65f72d3051ce7ae2dddf Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Thu, 13 Nov 2025 02:25:32 +0000
Subject: [PATCH] Optimize ToolManager._validate_backend_tool_arguments

The optimization achieves a **1206% speedup** by eliminating an expensive operation in the `_get_tool` method's fallback path when no source is specified.

**Key optimization:** In the original code, when `source=None`, the method called `self._get_all_tools()` which gathered all backend tools AND all MCP tools from an external client. The line profiler shows this single call consumed **99.8% of execution time** (21.5ms out of 21.5ms total).

The optimized version replaces this expensive all-tools scan with a **two-stage lookup**:
1. First, directly check the backend tools dictionary (`self._tools.get(name)`)
2. Only if not found, then query MCP tools (`self._list_mcp_tools()`)

**Why this works:** Most tool lookups likely target backend tools stored in the local dictionary. The optimization leverages this by checking the fast O(1) dictionary lookup first, avoiding the expensive MCP client call in the common case.

**Performance impact:** The test results show dramatic improvements across all scenarios:
- Simple successful lookups: **25,000-31,000% faster** (sub-microsecond vs hundreds of microseconds)
- Missing tool scenarios: **168-768% faster**
- Large-scale tests with 500+ parameters: **2,197-10,641% faster**

This optimization is particularly valuable if the function is called frequently during tool validation workflows, as it transforms an expensive operation into a fast dictionary lookup for the majority of cases.
---
 marimo/_server/ai/tools/tool_manager.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/marimo/_server/ai/tools/tool_manager.py b/marimo/_server/ai/tools/tool_manager.py
index cc616057e96..dabf9414167 100644
--- a/marimo/_server/ai/tools/tool_manager.py
+++ b/marimo/_server/ai/tools/tool_manager.py
@@ -108,10 +108,14 @@ def _get_tool(
             return None
         else:
             # No source specified, check all sources
-            all_tools = self._get_all_tools()
-            for tool in all_tools:
-                if tool.name == name:
-                    return tool
+            tool = self._tools.get(name)
+            if tool is not None:
+                return tool
+
+            mcp_tools = self._list_mcp_tools()
+            for mcp_tool in mcp_tools:
+                if mcp_tool.name == name:
+                    return mcp_tool
             return None
 
     def _validate_backend_tool_arguments(