hanzo-mcp 0.3.8__py3-none-any.whl → 0.5.1__py3-none-any.whl

hanzo_mcp/tools/agent/agent_tool.py

@@ -4,15 +4,19 @@ This module implements the AgentTool that allows Claude to delegate tasks to sub
 enabling concurrent execution of multiple operations and specialized processing.
 """
 
+import asyncio
 import json
+import re
 import time
 from collections.abc import Iterable
-from typing import Any, final, override
+from typing import Annotated, TypedDict, Unpack, final, override
 
 import litellm
-from mcp.server.fastmcp import Context as MCPContext
-from mcp.server.fastmcp import FastMCP
+from fastmcp import Context as MCPContext
+from fastmcp import FastMCP
+from fastmcp.server.dependencies import get_context
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
+from pydantic import Field
 
 from hanzo_mcp.tools.agent.prompt import (
     get_allowed_agent_tools,
@@ -24,12 +28,32 @@ from hanzo_mcp.tools.agent.tool_adapter import (
     convert_tools_to_openai_functions,
 )
 from hanzo_mcp.tools.common.base import BaseTool
-from hanzo_mcp.tools.common.context import DocumentContext, ToolContext, create_tool_context
+from hanzo_mcp.tools.common.batch_tool import BatchTool
+from hanzo_mcp.tools.common.context import (
+    ToolContext,
+    create_tool_context,
+)
 from hanzo_mcp.tools.common.permissions import PermissionManager
 from hanzo_mcp.tools.filesystem import get_read_only_filesystem_tools
 from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
-from hanzo_mcp.tools.project import get_project_tools
-from hanzo_mcp.tools.shell.command_executor import CommandExecutor
+
+Prompt = Annotated[
+    str,
+    Field(
+        description="Task for the agent to perform (must include absolute paths starting with /)",
+        min_length=1,
+    ),
+]
+
+
+class AgentToolParams(TypedDict, total=False):
+    """Parameters for the AgentTool.
+
+    Attributes:
+        prompts: Task(s) for the agent to perform (must include absolute paths starting with /)
+    """
+
+    prompts: str | list[str]
 
 
 @final
@@ -58,96 +82,73 @@ class AgentTool(BaseTool):
         Returns:
             Tool description
         """
-        return """Launch one or more agents that can perform tasks using read-only tools.
-
-This tool creates agents for delegation of tasks such as multi-step searches, complex analyses, 
-or other operations that benefit from focused processing. Multiple agents can work concurrently 
-on independent tasks, improving performance for complex operations.
+        # TODO: Add glob when it is implemented
+        at = [t.name for t in self.available_tools]
 
-Each agent works with its own context and provides a response containing the results of its work.
-Results from all agents are combined in the final response.
+        return f"""Launch a new agent that has access to the following tools: {at} When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries, use the Agent tool to perform the search for you.
 
-Args:
-    prompts: A list of task descriptions, where each item launches an independent agent.
+When to use the Agent tool:
+- If you are searching for a keyword like \"config\" or \"logger\", or for questions like \"which file does X?\", the Agent tool is strongly recommended
 
-Returns:
-    Combined results from all agent executions
-"""
+When NOT to use the Agent tool:
+- If you want to read a specific file path, use the read or glob tool instead of the Agent tool, to find the match more quickly
+- If you are searching for a specific class definition like \"class Foo\", use the glob tool instead, to find the match more quickly
+- If you are searching for code within a specific file or set of 2-3 files, use the read tool instead of the Agent tool, to find the match more quickly
 
-    @property
-    @override
-    def parameters(self) -> dict[str, Any]:
-        """Get the parameter specifications for the tool.
-
-        Returns:
-            Parameter specifications
-        """
-        return {
-            "properties": {
-                "prompts": {
-                    "anyOf": [
-                        {
-                            "type": "array",
-                            "items": {
-                                "type": "string"
-                            },
-                            "description": "List of tasks for agents to perform concurrently"
-                        },
-                        {
-                            "type": "string",
-                            "description": "Single task for the agent to perform"
-                        }
-                    ],
-                    "description": "Task(s) for agent(s) to perform"
-                }
-            },
-            "required": ["prompts"],
-            "type": "object",
-        }
-
-    @property
-    @override
-    def required(self) -> list[str]:
-        """Get the list of required parameter names.
-
-        Returns:
-            List of required parameter names
-        """
-        return ["prompts"]
+Usage notes:
+1. Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses
+2. When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result.
+3. Each agent invocation is stateless. You will not be able to send additional messages to the agent, nor will the agent be able to communicate with you outside of its final report. Therefore, your prompt should contain a highly detailed task description for the agent to perform autonomously and you should specify exactly what information the agent should return back to you in its final and only message to you.
+4. The agent's outputs should generally be trusted
+5.IMPORTANT: The Agent has no awareness of your context, so you must explicitly specify absolute project/file/directory paths and detailed background information about the current task. """
 
     def __init__(
-            self, document_context: DocumentContext, permission_manager: PermissionManager, command_executor: CommandExecutor,
-            model: str | None = None, api_key: str | None = None, max_tokens: int | None = None,
-            max_iterations: int = 10, max_tool_uses: int = 30
+        self,
+        permission_manager: PermissionManager,
+        model: str | None = None,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        max_tokens: int | None = None,
+        max_iterations: int = 10,
+        max_tool_uses: int = 30,
     ) -> None:
         """Initialize the agent tool.
 
         Args:
-            document_context: Document context for tracking file contents
+
             permission_manager: Permission manager for access control
-            command_executor: Command executor for running shell commands
             model: Optional model name override in LiteLLM format (e.g., "openai/gpt-4o")
             api_key: Optional API key for the model provider
+            base_url: Optional base URL for the model provider API endpoint
             max_tokens: Optional maximum tokens for model responses
             max_iterations: Maximum number of iterations for agent (default: 10)
             max_tool_uses: Maximum number of total tool uses for agent (default: 30)
         """
-        self.document_context = document_context
+
         self.permission_manager = permission_manager
-        self.command_executor = command_executor
         self.model_override = model
         self.api_key_override = api_key
+        self.base_url_override = base_url
         self.max_tokens_override = max_tokens
         self.max_iterations = max_iterations
         self.max_tool_uses = max_tool_uses
-        self.available_tools :list[BaseTool] = []
-        self.available_tools.extend(get_read_only_filesystem_tools(self.document_context, self.permission_manager))
-        self.available_tools.extend(get_read_only_jupyter_tools(self.document_context, self.permission_manager))
-        self.available_tools.extend(get_project_tools(self.document_context, self.permission_manager,self.command_executor))
-        
+        self.available_tools: list[BaseTool] = []
+        self.available_tools.extend(
+            get_read_only_filesystem_tools(self.permission_manager)
+        )
+        self.available_tools.extend(
+            get_read_only_jupyter_tools(self.permission_manager)
+        )
+        self.available_tools.append(
+            BatchTool({t.name: t for t in self.available_tools})
+        )
 
     @override
-    async def call(self, ctx: MCPContext, **params: Any) -> str:
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[AgentToolParams],
+    ) -> str:
         """Execute the tool with the given parameters.
 
         Args:
@@ -158,132 +159,180 @@ Returns:
             Tool execution result
         """
         start_time = time.time()
-        
+
         # Create tool context
         tool_ctx = create_tool_context(ctx)
-        tool_ctx.set_tool_info(self.name)
+        await tool_ctx.set_tool_info(self.name)
 
-        # Extract parameters
+        # Extract and validate parameters
         prompts = params.get("prompts")
+        
         if prompts is None:
-            await tool_ctx.error("Parameter 'prompts' is required but was not provided")
-            return "Error: Parameter 'prompts' is required but was not provided"
+            await tool_ctx.error("No prompts provided")
+            return """Error: At least one prompt must be provided.
+
+IMPORTANT REMINDER FOR CLAUDE:
+The dispatch_agent tool requires prompts parameter. Please provide either:
+- A single prompt as a string
+- Multiple prompts as an array of strings
 
-        if not isinstance(prompts, list) and not isinstance(prompts, str):
-            await tool_ctx.error("Parameter 'prompts' must be a string or an array of strings")
-            return "Error: Parameter 'prompts' must be a string or an array of strings"
+Each prompt must contain absolute paths starting with /.
+Example of correct usage:
+- prompts: "Search for all instances of the 'config' variable in /Users/bytedance/project/hanzo-mcp"
+- prompts: ["Find files in /path/to/project", "Search code in /path/to/src"]"""
 
+        # Handle both string and list inputs
         if isinstance(prompts, str):
-            prompts = [prompts]
-
-        if not prompts:  # Empty list
-            await tool_ctx.error("At least one prompt must be provided in the array")
-            return "Error: At least one prompt must be provided in the array"
-
-        # Check for empty strings in the list
-        if any(not isinstance(p, str) or not p.strip() for p in prompts):
-            await tool_ctx.error("All prompts must be non-empty strings")
-            return "Error: All prompts must be non-empty strings"
-                
-        # Always use _execute_multiple_agents, treating single agent as a special case
-        await tool_ctx.info(f"Launching {len(prompts)} agent{'s' if len(prompts) > 1 else ''}")
-        result = await self._execute_multiple_agents(prompts, tool_ctx)
-            
-        # Calculate execution time
-        execution_time = time.time() - start_time
-        
-        # Format the result
-        formatted_result = self._format_result(result, execution_time, len(prompts))
-        
-        # Log completion
-        await tool_ctx.info(f"Agent execution completed in {execution_time:.2f}s")
-        
-        return formatted_result
+            prompt_list = [prompts]
+        elif isinstance(prompts, list):
+            if not prompts:
+                await tool_ctx.error("Empty prompts list provided")
+                return "Error: At least one prompt must be provided when using a list."
+            if not all(isinstance(p, str) for p in prompts):
+                await tool_ctx.error("All prompts must be strings")
+                return "Error: All prompts in the list must be strings."
+            prompt_list = prompts
+        else:
+            await tool_ctx.error("Invalid prompts parameter type")
+            return "Error: Parameter 'prompts' must be a string or an array of strings."
+
+        # Validate absolute paths in all prompts
+        absolute_path_pattern = r"/(?:[^/\s]+/)*[^/\s]+"
+        for prompt in prompt_list:
+            if not re.search(absolute_path_pattern, prompt):
+                await tool_ctx.error(f"Prompt does not contain absolute path: {prompt[:50]}...")
+                return """Error: All prompts must contain at least one absolute path.
+
+IMPORTANT REMINDER FOR CLAUDE:
+When using the dispatch_agent tool, always include absolute paths in your prompts.
+Example of correct usage:
+- "Search for all instances of the 'config' variable in /Users/bytedance/project/hanzo-mcp"
+- "Find files that import the database module in /Users/bytedance/project/hanzo-mcp/src"
+
+The agent cannot access files without knowing their absolute locations."""
+
+        # Execute agent(s) - always use _execute_multiple_agents for list inputs
+        if len(prompt_list) == 1:
+            await tool_ctx.info("Launching agent")
+            result = await self._execute_multiple_agents(prompt_list, tool_ctx)
+            execution_time = time.time() - start_time
+            formatted_result = f"""Agent execution completed in {execution_time:.2f} seconds.
+
+AGENT RESPONSE:
+{result}"""
+            await tool_ctx.info(f"Agent execution completed in {execution_time:.2f}s")
+            return formatted_result
+        else:
+            await tool_ctx.info(f"Launching {len(prompt_list)} agents in parallel")
+            result = await self._execute_multiple_agents(prompt_list, tool_ctx)
+            execution_time = time.time() - start_time
+            formatted_result = f"""Multi-agent execution completed in {execution_time:.2f} seconds ({len(prompt_list)} agents).
+
+AGENT RESPONSES:
+{result}"""
+            await tool_ctx.info(f"Multi-agent execution completed in {execution_time:.2f}s")
+            return formatted_result
+
+    async def _execute_agent(self, prompt: str, tool_ctx: ToolContext) -> str:
+        """Execute a single agent with the given prompt.
+
+        Args:
+            prompt: The task prompt for the agent
+            tool_ctx: Tool context for logging
+
+        Returns:
+            Agent execution result
+        """
+        # Get available tools for the agent
+        agent_tools = get_allowed_agent_tools(
+            self.available_tools,
+            self.permission_manager,
+        )
+
+        # Convert tools to OpenAI format
+        openai_tools = convert_tools_to_openai_functions(agent_tools)
+
+        # Log execution start
+        await tool_ctx.info("Starting agent execution")
+
+        # Create a result container
+        result = ""
+
+        try:
+            # Create system prompt for this agent
+            system_prompt = get_system_prompt(
+                agent_tools,
+                self.permission_manager,
+            )
+
+            # Execute agent
+            await tool_ctx.info(f"Executing agent task: {prompt[:50]}...")
+            result = await self._execute_agent_with_tools(
+                system_prompt, prompt, agent_tools, openai_tools, tool_ctx
+            )
+        except Exception as e:
+            # Log and return error result
+            error_message = f"Error executing agent: {str(e)}"
+            await tool_ctx.error(error_message)
+            return f"Error: {error_message}"
+
+        return result if result else "No results returned from agent"
 
     async def _execute_multiple_agents(self, prompts: list[str], tool_ctx: ToolContext) -> str:
-        """Execute multiple agents concurrently with the given prompts.
+        """Execute multiple agents concurrently.
 
         Args:
             prompts: List of prompts for the agents
             tool_ctx: Tool context for logging
 
         Returns:
-            Combined agent execution results
+            Combined results from all agents
         """
-        # Get available tools for the agents (do this once to avoid redundant work)
+        # Get available tools for the agents
         agent_tools = get_allowed_agent_tools(
-            self.available_tools, 
+            self.available_tools,
             self.permission_manager,
         )
-        
-        # Convert tools to OpenAI format (do this once to avoid redundant work)
+
+        # Convert tools to OpenAI format
         openai_tools = convert_tools_to_openai_functions(agent_tools)
-        
-        # Log execution start
-        await tool_ctx.info(f"Starting execution of {len(prompts)} agent{'s' if len(prompts) > 1 else ''}")
-        
-        # Create a list to store the tasks
+
+        # Create system prompt for the agents
+        system_prompt = get_system_prompt(
+            agent_tools,
+            self.permission_manager,
+        )
+
+        # Create tasks for parallel execution
         tasks = []
-        results = []
-        
-        # Handle exceptions for individual agent executions
         for i, prompt in enumerate(prompts):
-            try:
-                # Create system prompt for this agent
-                system_prompt = get_system_prompt(
-                    agent_tools,
-                    self.permission_manager,
-                )
-                
-                # Execute agent and collect the task
-                await tool_ctx.info(f"Launching agent {i+1}/{len(prompts)}: {prompt[:50]}...")
-                task = self._execute_agent_with_tools(
-                    system_prompt, 
-                    prompt,
-                    agent_tools, 
-                    openai_tools, 
-                    tool_ctx
-                )
-                tasks.append(task)
-            except Exception as e:
-                # Log and add error result
-                error_message = f"Error preparing agent {i+1}: {str(e)}"
-                await tool_ctx.error(error_message)
-                results.append(f"Agent {i+1} Error: {error_message}")
-        
-        # Execute all pending tasks concurrently
-        if tasks:
-            import asyncio
-            try:
-                # Wait for all tasks to complete
-                completed_results = await asyncio.gather(*tasks, return_exceptions=True)
-                
-                # Process results, handling any exceptions
-                for i, result in enumerate(completed_results):
-                    if isinstance(result, Exception):
-                        results.append(f"Agent {i+1} Error: {str(result)}")
-                    else:
-                        # For multi-agent case, add agent number prefix
-                        if len(prompts) > 1:
-                            results.append(f"Agent {i+1} Result:\n{result}")
-                        else:
-                            # For single agent case, just add the result
-                            results.append(result)
-            except Exception as e:
-                # Handle any unexpected exceptions during gathering
-                error_message = f"Error executing agents concurrently: {str(e)}"
-                await tool_ctx.error(error_message)
-                results.append(f"Error: {error_message}")
-        
-        # Combine results - different handling for single vs multi-agent
-        if len(prompts) > 1:
-            # Multi-agent: add separator between results
-            combined_result = "\n\n" + "\n\n---\n\n".join(results)
-        else:
-            # Single agent: just return the result
-            combined_result = results[0] if results else "No results returned from agent"
-            
-        return combined_result
+            await tool_ctx.info(f"Creating agent task {i+1}: {prompt[:50]}...")
+            task = self._execute_agent_with_tools(
+                system_prompt, prompt, agent_tools, openai_tools, tool_ctx
+            )
+            tasks.append(task)
+
+        # Execute all agents concurrently
+        await tool_ctx.info(f"Executing {len(tasks)} agents in parallel")
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Handle single agent case
+        if len(results) == 1:
+            if isinstance(results[0], Exception):
+                await tool_ctx.error(f"Agent execution failed: {str(results[0])}")
+                return f"Error: {str(results[0])}"
+            return results[0]
+
+        # Format results for multiple agents
+        formatted_results = []
+        for i, result in enumerate(results):
+            if isinstance(result, Exception):
+                formatted_results.append(f"Agent {i+1} Error:\n{str(result)}")
+                await tool_ctx.error(f"Agent {i+1} failed: {str(result)}")
+            else:
+                formatted_results.append(f"Agent {i+1} Result:\n{result}")
+
+        return "\n\n---\n\n".join(formatted_results)
 
     async def _execute_agent_with_tools(
         self,
@@ -308,24 +357,26 @@ Returns:
         # Get model parameters and name
         model = get_default_model(self.model_override)
         params = get_model_parameters(max_tokens=self.max_tokens_override)
-        
+
         # Initialize messages
-        messages:Iterable[ChatCompletionMessageParam] = []
+        messages: Iterable[ChatCompletionMessageParam] = []
         messages.append({"role": "system", "content": system_prompt})
         messages.append({"role": "user", "content": user_prompt})
-        
+
         # Track tool usage for metrics
         tool_usage = {}
         total_tool_use_count = 0
         iteration_count = 0
         max_tool_uses = self.max_tool_uses  # Safety limit to prevent infinite loops
-        max_iterations = self.max_iterations  # Add a maximum number of iterations for safety
+        max_iterations = (
+            self.max_iterations
+        )  # Add a maximum number of iterations for safety
 
         # Execute until the agent completes or reaches the limit
         while total_tool_use_count < max_tool_uses and iteration_count < max_iterations:
             iteration_count += 1
             await tool_ctx.info(f"Calling model (iteration {iteration_count})...")
-            
+
             try:
                 # Configure model parameters based on capabilities
                 completion_params = {
@@ -339,88 +390,99 @@ Returns:
 
                 if self.api_key_override:
                     completion_params["api_key"] = self.api_key_override
-                
+
                 # Add max_tokens if provided
                 if params.get("max_tokens"):
                     completion_params["max_tokens"] = params.get("max_tokens")
-                
+
+                # Add base_url if provided
+                if self.base_url_override:
+                    completion_params["base_url"] = self.base_url_override
+
                 # Make the model call
                 response = litellm.completion(
-                        **completion_params #pyright: ignore
+                    **completion_params  # pyright: ignore
                 )
 
-                if len(response.choices) == 0: #pyright: ignore
+                if len(response.choices) == 0:  # pyright: ignore
                     raise ValueError("No response choices returned")
 
-                message = response.choices[0].message #pyright: ignore
+                message = response.choices[0].message  # pyright: ignore
 
                 # Add message to conversation history
-                messages.append(message) #pyright: ignore
+                messages.append(message)  # pyright: ignore
 
                 # If no tool calls, we're done
                 if not message.tool_calls:
                     return message.content or "Agent completed with no response."
-                    
+
                 # Process tool calls
                 tool_call_count = len(message.tool_calls)
                 await tool_ctx.info(f"Processing {tool_call_count} tool calls")
-                
+
                 for tool_call in message.tool_calls:
                     total_tool_use_count += 1
                     function_name = tool_call.function.name
-                    
+
                     # Track usage
                     tool_usage[function_name] = tool_usage.get(function_name, 0) + 1
-                    
+
                     # Log tool usage
                     await tool_ctx.info(f"Agent using tool: {function_name}")
-                    
+
                     # Parse the arguments
                     try:
                         function_args = json.loads(tool_call.function.arguments)
                     except json.JSONDecodeError:
                         function_args = {}
-                        
+
                     # Find the matching tool
-                    tool = next((t for t in available_tools if t.name == function_name), None)
+                    tool = next(
+                        (t for t in available_tools if t.name == function_name), None
+                    )
                     if not tool:
                         tool_result = f"Error: Tool '{function_name}' not found"
                     else:
                         try:
-                            tool_result = await tool.call(ctx=tool_ctx.mcp_context, **function_args)
+                            tool_result = await tool.call(
+                                ctx=tool_ctx.mcp_context, **function_args
+                            )
                         except Exception as e:
                             tool_result = f"Error executing {function_name}: {str(e)}"
-                            
+
+                    await tool_ctx.info(
+                        f"tool {function_name} run with args {function_args} and return {tool_result[: min(100, len(tool_result))]}"
+                    )
                     # Add the tool result to messages
                     messages.append(
                         {
                             "role": "tool",
                             "tool_call_id": tool_call.id,
+                            "name": function_name,
                             "content": tool_result,
                         }
                     )
-                
+
                 # Log progress
-                await tool_ctx.info(f"Processed {len(message.tool_calls)} tool calls. Total: {total_tool_use_count}")
-                
+                await tool_ctx.info(
+                    f"Processed {len(message.tool_calls)} tool calls. Total: {total_tool_use_count}"
+                )
+
             except Exception as e:
                 await tool_ctx.error(f"Error in model call: {str(e)}")
                 # Avoid trying to JSON serialize message objects
                 await tool_ctx.error(f"Message count: {len(messages)}")
                 return f"Error in agent execution: {str(e)}"
-                
+
         # If we've reached the limit, add a warning and get final response
         if total_tool_use_count >= max_tool_uses or iteration_count >= max_iterations:
-            limit_type = "tool usage" if total_tool_use_count >= max_tool_uses else "iterations"
-            await tool_ctx.info(f"Reached maximum {limit_type} limit. Getting final response.")
-            
             messages.append(
                 {
                     "role": "system",
-                    "content": f"You have reached the maximum number of {limit_type}. Please provide your final response.",
+                    "content": "You have reached the maximum iteration. Please provide your final response.",
                 }
             )
-            
+
             try:
                 # Make a final call to get the result
                 final_response = litellm.completion(
@@ -430,45 +492,49 @@ Returns:
                     timeout=params["timeout"],
                     max_tokens=params.get("max_tokens"),
                 )
-                
-                return final_response.choices[0].message.content or f"Agent reached {limit_type} limit without a response." #pyright: ignore
+
+                return (
+                    final_response.choices[0].message.content
+                    or "Agent reached max iteration limit without a response."
+                )  # pyright: ignore
             except Exception as e:
                 await tool_ctx.error(f"Error in final model call: {str(e)}")
                 return f"Error in final response: {str(e)}"
-        
+
         # Should not reach here but just in case
         return "Agent execution completed after maximum iterations."
 
-    def _format_result(self, result: str, execution_time: float, agent_count: int) -> str:
+    def _format_result(self, result: str, execution_time: float) -> str:
         """Format agent result with metrics.
 
         Args:
-            result: Raw result from agent(s)
+            result: Raw result from agent
             execution_time: Execution time in seconds
-            agent_count: Number of agents used
 
         Returns:
             Formatted result with metrics
         """
-        # Different format based on agent count
-        if agent_count > 1:
-            # Multi-agent response
-            return f"""Multi-agent execution completed in {execution_time:.2f} seconds ({agent_count} agents).
-
-{result}
-"""
-        else:
-            # Single agent response
-            return f"""Agent execution completed in {execution_time:.2f} seconds.
+        return f"""Agent execution completed in {execution_time:.2f} seconds.
 
 AGENT RESPONSE:
 {result}
 """
-    
+
     @override
     def register(self, mcp_server: FastMCP) -> None:
+        """Register this agent tool with the MCP server.
+
+        Creates a wrapper function with explicitly defined parameters that match
+        the tool's parameter schema and registers it with the MCP server.
+
+        Args:
+            mcp_server: The FastMCP server instance
+        """
         tool_self = self  # Create a reference to self for use in the closure
-        
-        @mcp_server.tool(name=self.name, description=self.mcp_description)
-        async def dispatch_agent(ctx: MCPContext, prompts: list[str] | str) -> str:
-             return await tool_self.call(ctx, prompts=prompts)
+
+        @mcp_server.tool(name=self.name, description=self.description)
+        async def dispatch_agent(
+            prompts: str | list[str],
+        ) -> str:
+            ctx = get_context()
+            return await tool_self.call(ctx, prompts=prompts)