hanzo-mcp 0.7.3__py3-none-any.whl → 0.7.6__py3-none-any.whl

hanzo_mcp/tools/agent/agent_tool.py

@@ -1,124 +1,283 @@
-"""Agent tool implementation for Hanzo AI.
+"""Agent tool implementation using hanzo-agents SDK.
 
-This module implements the AgentTool that allows Claude to delegate tasks to sub-agents,
-enabling concurrent execution of multiple operations and specialized processing.
+This module implements the AgentTool that leverages the hanzo-agents SDK
+for sophisticated agent orchestration and execution.
 """
 
 import asyncio
 import json
+import os
 import re
 import time
-from collections.abc import Iterable
-from typing import Annotated, TypedDict, Unpack, final, override
+from typing import Annotated, TypedDict, Unpack, final, override, Optional, Dict, Any, List
 
-# Import litellm with warnings suppressed
-import warnings
-with warnings.catch_warnings():
-    warnings.simplefilter("ignore", DeprecationWarning)
-    import litellm
 from mcp.server.fastmcp import Context as MCPContext
 from mcp.server import FastMCP
-from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
 from pydantic import Field
 
-from hanzo_mcp.tools.agent.prompt import (
-    get_allowed_agent_tools,
-    get_default_model,
-    get_model_parameters,
-    get_system_prompt,
-)
-from hanzo_mcp.tools.agent.tool_adapter import (
-    convert_tools_to_openai_functions,
-)
-from hanzo_mcp.tools.agent.clarification_protocol import (
-    AgentClarificationMixin,
-    ClarificationType,
-)
-from hanzo_mcp.tools.agent.clarification_tool import ClarificationTool
-from hanzo_mcp.tools.agent.critic_tool import CriticTool, CriticProtocol
-from hanzo_mcp.tools.agent.review_tool import ReviewTool, ReviewProtocol
-from hanzo_mcp.tools.agent.iching_tool import IChingTool
+# Import hanzo-agents SDK
+try:
+    from hanzo_agents import (
+        Agent, State, Network, Tool, History,
+        ModelRegistry, InferenceResult, ToolCall,
+        create_memory_kv, create_memory_vector,
+        sequential_router, state_based_router,
+    )
+    from hanzo_agents.core.cli_agent import (
+        ClaudeCodeAgent, OpenAICodexAgent, 
+        GeminiAgent, GrokAgent
+    )
+    HANZO_AGENTS_AVAILABLE = True
+except ImportError:
+    HANZO_AGENTS_AVAILABLE = False
+    # Define stub classes when hanzo-agents is not available
+    class State:
+        """Stub State class when hanzo-agents is not available."""
+        def __init__(self):
+            pass
+        def to_dict(self):
+            return {}
+        @classmethod
+        def from_dict(cls, data):
+            return cls()
+    
+    class Tool:
+        """Stub Tool class when hanzo-agents is not available."""
+        pass
+    
+    class Agent:
+        """Stub Agent class when hanzo-agents is not available."""
+        pass
+    
+    class Network:
+        """Stub Network class when hanzo-agents is not available."""
+        pass
+    
+    class History:
+        """Stub History class when hanzo-agents is not available."""
+        pass
+    
+    class InferenceResult:
+        """Stub InferenceResult class when hanzo-agents is not available."""
+        def __init__(self, agent=None, content=None, metadata=None):
+            self.agent = agent
+            self.content = content
+            self.metadata = metadata or {}
+
 from hanzo_mcp.tools.common.base import BaseTool
-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.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, Edit, MultiEdit
 from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
-
-Prompt = Annotated[
-    str,
-    Field(
-        description="Task for the agent to perform (must include absolute paths starting with /)",
-        min_length=1,
-    ),
-]
+from hanzo_mcp.tools.common.batch_tool import BatchTool
+from hanzo_mcp.tools.agent.clarification_protocol import AgentClarificationMixin, ClarificationType
+from hanzo_mcp.tools.agent.clarification_tool import ClarificationTool
+from hanzo_mcp.tools.agent.critic_tool import CriticTool
+from hanzo_mcp.tools.agent.review_tool import ReviewTool
+from hanzo_mcp.tools.agent.iching_tool import IChingTool
 
 
 class AgentToolParams(TypedDict, total=False):
-    """Parameters for the AgentTool.
+    """Parameters for the AgentTool."""
+    prompts: str | list[str]
+    model: Optional[str]
+    use_memory: Optional[bool]
+    memory_backend: Optional[str]
 
-    Attributes:
-        prompts: Task(s) for the agent to perform (must include absolute paths starting with /)
-    """
 
-    prompts: str | list[str]
+class MCPAgentState(State):
+    """State for MCP agents."""
+    
+    def __init__(self, prompts: List[str], context: Dict[str, Any]):
+        """Initialize agent state."""
+        super().__init__()
+        self.prompts = prompts
+        self.context = context
+        self.current_prompt_index = 0
+        self.results = []
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        base_dict = super().to_dict()
+        base_dict.update({
+            "prompts": self.prompts,
+            "context": self.context,
+            "current_prompt_index": self.current_prompt_index,
+            "results": self.results
+        })
+        return base_dict
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "MCPAgentState":
+        """Create from dictionary."""
+        state = cls(
+            prompts=data.get("prompts", []),
+            context=data.get("context", {})
+        )
+        state.current_prompt_index = data.get("current_prompt_index", 0)
+        state.results = data.get("results", [])
+        for k, v in data.items():
+            if k not in ["prompts", "context", "current_prompt_index", "results"]:
+                state[k] = v
+        return state
 
 
-@final
-class AgentTool(AgentClarificationMixin, BaseTool):
-    """Tool for delegating tasks to sub-agents.
+class MCPToolAdapter(Tool):
+    """Adapter to wrap MCP tools for hanzo-agents."""
+    
+    def __init__(self, mcp_tool: BaseTool, ctx: MCPContext):
+        """Initialize adapter."""
+        self.mcp_tool = mcp_tool
+        self.ctx = ctx
+        
+    @property
+    def name(self) -> str:
+        """Get tool name."""
+        return self.mcp_tool.name
+    
+    @property
+    def description(self) -> str:
+        """Get tool description."""
+        return self.mcp_tool.description
+    
+    async def execute(self, **kwargs) -> str:
+        """Execute the MCP tool."""
+        return await self.mcp_tool.call(self.ctx, **kwargs)
+
 
-    The AgentTool allows Claude to create and manage sub-agents for performing
-    specialized tasks concurrently, such as code search, analysis, and more.
+class MCPAgent(Agent):
+    """Agent that executes MCP tasks."""
+    
+    name = "mcp_agent"
+    description = "Agent for executing MCP tasks"
+    
+    def __init__(self, 
+                 available_tools: List[BaseTool],
+                 permission_manager: PermissionManager,
+                 ctx: MCPContext,
+                 model: str = "model://anthropic/claude-3-5-sonnet-20241022",
+                 **kwargs):
+        """Initialize MCP agent."""
+        super().__init__(model=model, **kwargs)
+        
+        self.available_tools = available_tools
+        self.permission_manager = permission_manager
+        self.ctx = ctx
+        
+        # Register MCP tools as agent tools
+        for mcp_tool in available_tools:
+            adapter = MCPToolAdapter(mcp_tool, ctx)
+            self.register_tool(adapter)
+    
+    async def run(self, state: MCPAgentState, history: History, network: Network) -> InferenceResult:
+        """Execute the agent."""
+        # Get current prompt
+        if state.current_prompt_index >= len(state.prompts):
+            return InferenceResult(
+                agent=self.name,
+                content="All prompts completed",
+                metadata={"completed": True}
+            )
+        
+        prompt = state.prompts[state.current_prompt_index]
+        
+        # Execute with tools
+        messages = [
+            {"role": "system", "content": self._get_system_prompt()},
+            {"role": "user", "content": prompt}
+        ]
+        
+        # Add history context
+        for entry in history[-10:]:
+            if entry.role == "assistant":
+                messages.append({
+                    "role": "assistant",
+                    "content": entry.content
+                })
+            elif entry.role == "user":
+                messages.append({
+                    "role": "user", 
+                    "content": entry.content
+                })
+        
+        # Call model
+        from hanzo_agents import ModelRegistry
+        adapter = ModelRegistry.get_adapter(self.model)
+        response = await adapter.chat(messages)
+        
+        # Update state
+        state.current_prompt_index += 1
+        state.results.append(response)
+        
+        # Return result
+        return InferenceResult(
+            agent=self.name,
+            content=response,
+            metadata={
+                "prompt_index": state.current_prompt_index - 1,
+                "total_prompts": len(state.prompts)
+            }
+        )
     
-    Agents can request clarification from the main loop up to once per task.
-    """
+    def _get_system_prompt(self) -> str:
+        """Get system prompt for the agent."""
+        tool_descriptions = []
+        for tool in self.tools.values():
+            tool_descriptions.append(f"- {tool.name}: {tool.description}")
+        
+        return f"""You are an AI assistant with access to the following tools:
+
+{chr(10).join(tool_descriptions)}
+
+When you need to use a tool, respond with:
+TOOL: tool_name(arg1="value1", arg2="value2")
+
+Important guidelines:
+- Always include absolute paths starting with / when working with files
+- Be thorough in your searches and analysis
+- Provide clear, actionable results
+- Edit files when requested to make changes
+"""
+
 
+@final
+class AgentTool(AgentClarificationMixin, BaseTool):
+    """Tool for delegating tasks to sub-agents using hanzo-agents SDK."""
+    
     @property
     @override
     def name(self) -> str:
-        """Get the tool name.
-
-        Returns:
-            Tool name
-        """
+        """Get the tool name."""
         return "agent"
-
+    
     @property
     @override
     def description(self) -> str:
-        """Get the tool description.
-
-        Returns:
-            Tool description
-        """
-        # TODO: Add glob when it is implemented
+        """Get the tool description."""
+        if not HANZO_AGENTS_AVAILABLE:
+            return "Agent tool (hanzo-agents SDK not available - using fallback)"
+        
         at = [t.name for t in self.available_tools]
-
-        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.
+        return f"""Launch a new agent that has access to the following tools: {at}.
 
 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
-- When you need to perform edits across multiple files based on search results
+- If you are searching for a keyword like "config" or "logger"
+- When you need to perform edits across multiple files
 - When you need to delegate complex file modification tasks
 
 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
-- Writing code and running bash commands (use other tools for that)
-- Other tasks that are not related to searching for a keyword or file
+- If you want to read a specific file path
+- If you are searching for a specific class definition
+- Writing code and running bash commands
+- Other tasks that are not related to searching
 
 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.
+1. Launch multiple agents concurrently whenever possible
+2. Agent results are not visible to the user - summarize them
+3. Each agent invocation is stateless
 4. The agent's outputs should generally be trusted
-5. Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, web fetches, etc.), since it is not aware of the user's intent"""
-
+5. Clearly tell the agent whether you expect it to write code or just do research"""
+    
     def __init__(
         self,
         permission_manager: PermissionManager,
@@ -129,26 +288,16 @@ Usage notes:
         max_iterations: int = 10,
         max_tool_uses: int = 30,
     ) -> None:
-        """Initialize the agent tool.
-
-        Args:
-
-            permission_manager: Permission manager for access control
-            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)
-        """
-
+        """Initialize the agent tool."""
         self.permission_manager = permission_manager
         self.model_override = model
         self.api_key_override = api_key
-        self.base_url_override = base_url
+        self.base_url_override = base_url  
         self.max_tokens_override = max_tokens
         self.max_iterations = max_iterations
         self.max_tool_uses = max_tool_uses
+        
+        # Set up available tools
         self.available_tools: list[BaseTool] = []
         self.available_tools.extend(
             get_read_only_filesystem_tools(self.permission_manager)
@@ -157,485 +306,176 @@ Usage notes:
             get_read_only_jupyter_tools(self.permission_manager)
         )
         
-        # Always add edit tools - agents should have edit access
+        # Add edit tools
         self.available_tools.append(Edit(self.permission_manager))
         self.available_tools.append(MultiEdit(self.permission_manager))
         
-        # Add clarification tool for agents
+        # Add special tools
         self.available_tools.append(ClarificationTool())
-        
-        # Add critic tool for agents (devil's advocate)
         self.available_tools.append(CriticTool())
-        
-        # Add review tool for agents (balanced review)
         self.available_tools.append(ReviewTool())
-        
-        # Add I Ching tool for creative guidance
         self.available_tools.append(IChingTool())
         
         self.available_tools.append(
             BatchTool({t.name: t for t in self.available_tools})
         )
-        
-        # Initialize protocols
-        self.critic_protocol = CriticProtocol()
-        self.review_protocol = ReviewProtocol()
-
+    
     @override
     async def call(
         self,
         ctx: MCPContext,
         **params: Unpack[AgentToolParams],
     ) -> str:
-        """Execute the tool with the given parameters.
-
-        Args:
-            ctx: MCP context
-            **params: Tool parameters
-
-        Returns:
-            Tool execution result
-        """
+        """Execute the tool with the given parameters."""
         start_time = time.time()
-
+        
         # Create tool context
         tool_ctx = create_tool_context(ctx)
         await tool_ctx.set_tool_info(self.name)
-
-        # Extract and validate parameters
-        prompts = params.get("prompts")
         
+        # Extract parameters
+        prompts = params.get("prompts")
         if prompts is None:
             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
-
-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"]"""
-
+            return "Error: At least one prompt must be provided."
+        
         # Handle both string and list inputs
         if isinstance(prompts, str):
             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."
+                return "Error: At least one prompt must be provided."
             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
+            return "Error: Parameter 'prompts' must be a string or list of strings."
+        
+        # Validate absolute paths
         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)
+                await tool_ctx.error(f"Prompt missing absolute path: {prompt[:50]}...")
+                return "Error: All prompts must contain at least one absolute path."
+        
+        # Require hanzo-agents SDK
+        if not HANZO_AGENTS_AVAILABLE:
+            await tool_ctx.error("hanzo-agents SDK is required but not available")
+            return "Error: hanzo-agents SDK is required for agent tool functionality. Please install it with: pip install hanzo-agents"
+        
+        # Use hanzo-agents SDK
+        await tool_ctx.info(f"Launching {len(prompt_list)} agent(s) using hanzo-agents SDK")
+        
+        # Determine model and agent type
+        model = params.get("model", self.model_override)
+        use_memory = params.get("use_memory", False)
+        memory_backend = params.get("memory_backend", "sqlite")
+        
+        # Get appropriate agent class
+        agent_class = self._get_agent_class(model)
+        
+        # Create state
+        state = MCPAgentState(
+            prompts=prompt_list,
+            context={
+                "permission_manager": self.permission_manager,
+                "api_key": self.api_key_override,
+                "base_url": self.base_url_override,
+                "max_tokens": self.max_tokens_override,
+            }
+        )
+        
+        # Create memory if requested
+        memory_kv = None
+        memory_vector = None
+        if use_memory:
+            memory_kv = create_memory_kv(memory_backend)
+            memory_vector = create_memory_vector("simple")
+        
+        # Create network
+        network = Network(
+            state=state,
+            agents=[agent_class],
+            router=sequential_router([agent_class] * len(prompt_list)),
+            memory_kv=memory_kv,
+            memory_vector=memory_vector,
+            max_steps=self.max_iterations * len(prompt_list),
+        )
+        
+        # Execute
+        try:
+            final_state = await network.run()
             execution_time = time.time() - start_time
-            formatted_result = f"""Agent execution completed in {execution_time:.2f} seconds.
+            
+            # Format results
+            results = final_state.results
+            if len(results) == 1:
+                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).
+{results[0]}"""
+            else:
+                formatted_results = []
+                for i, result in enumerate(results):
+                    formatted_results.append(f"Agent {i+1} Result:\n{result}")
+                
+                formatted_result = f"""Multi-agent execution completed in {execution_time:.2f} seconds ({len(results)} agents).
 
 AGENT RESPONSES:
-{result}"""
-            await tool_ctx.info(f"Multi-agent execution completed in {execution_time:.2f}s")
+{chr(10).join(formatted_results)}"""
+            
+            await tool_ctx.info(f"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.
-
-        Args:
-            prompts: List of prompts for the agents
-            tool_ctx: Tool context for logging
-
-        Returns:
-            Combined results from all agents
-        """
-        # Get available tools for the agents
-        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)
-
-        # Create system prompt for the agents
-        system_prompt = get_system_prompt(
-            agent_tools,
-            self.permission_manager,
-        )
-
-        # Create tasks for parallel execution
-        tasks = []
-        for i, prompt in enumerate(prompts):
-            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,
-        system_prompt: str,
-        user_prompt: str,
-        available_tools: list[BaseTool],
-        openai_tools: list[ChatCompletionToolParam],
-        tool_ctx: ToolContext,
-    ) -> str:
-        """Execute agent with tool handling.
-
-        Args:
-            system_prompt: System prompt for the agent
-            user_prompt: User prompt for the agent
-            available_tools: List of available tools
-            openai_tools: List of tools in OpenAI format
-            tool_ctx: Tool context for logging
-
-        Returns:
-            Agent execution result
-        """
-        # 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.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
-
-        # 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 = {
-                    "model": model,
-                    "messages": messages,
-                    "tools": openai_tools,
-                    "tool_choice": "auto",
-                    "temperature": params["temperature"],
-                    "timeout": params["timeout"],
-                }
-
-                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
-                )
-
-                if len(response.choices) == 0:  # pyright: ignore
-                    raise ValueError("No response choices returned")
-
-                message = response.choices[0].message  # pyright: ignore
-
-                # Add message to conversation history
-                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
-                    )
-                    if not tool:
-                        tool_result = f"Error: Tool '{function_name}' not found"
-                    # Special handling for clarification requests
-                    elif function_name == "request_clarification":
-                        try:
-                            # Extract clarification parameters
-                            request_type = function_args.get("type", "ADDITIONAL_INFO")
-                            question = function_args.get("question", "")
-                            context = function_args.get("context", {})
-                            options = function_args.get("options", None)
-                            
-                            # Convert string type to enum
-                            clarification_type = ClarificationType[request_type]
-                            
-                            # Request clarification
-                            answer = await self.request_clarification(
-                                request_type=clarification_type,
-                                question=question,
-                                context=context,
-                                options=options
-                            )
-                            
-                            tool_result = self.format_clarification_in_output(question, answer)
-                        except Exception as e:
-                            tool_result = f"Error processing clarification: {str(e)}"
-                    # Special handling for critic requests
-                    elif function_name == "critic":
-                        try:
-                            # Extract critic parameters
-                            review_type = function_args.get("review_type", "GENERAL")
-                            work_description = function_args.get("work_description", "")
-                            code_snippets = function_args.get("code_snippets", None)
-                            file_paths = function_args.get("file_paths", None)
-                            specific_concerns = function_args.get("specific_concerns", None)
-                            
-                            # Request critical review
-                            tool_result = self.critic_protocol.request_review(
-                                review_type=review_type,
-                                work_description=work_description,
-                                code_snippets=code_snippets,
-                                file_paths=file_paths,
-                                specific_concerns=specific_concerns
-                            )
-                        except Exception as e:
-                            tool_result = f"Error processing critic review: {str(e)}"
-                    # Special handling for review requests
-                    elif function_name == "review":
-                        try:
-                            # Extract review parameters
-                            focus = function_args.get("focus", "GENERAL")
-                            work_description = function_args.get("work_description", "")
-                            code_snippets = function_args.get("code_snippets", None)
-                            file_paths = function_args.get("file_paths", None)
-                            context = function_args.get("context", None)
-                            
-                            # Request balanced review
-                            tool_result = self.review_protocol.request_review(
-                                focus=focus,
-                                work_description=work_description,
-                                code_snippets=code_snippets,
-                                file_paths=file_paths,
-                                context=context
-                            )
-                        except Exception as e:
-                            tool_result = f"Error processing review: {str(e)}"
-                    else:
-                        try:
-                            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}"
-                )
-
-            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:
-            messages.append(
-                {
-                    "role": "system",
-                    "content": "You have reached the maximum iteration. Please provide your final response.",
-                }
+            await tool_ctx.error(f"Agent execution failed: {str(e)}")
+            return f"Error: {str(e)}"
+    
+    def _get_agent_class(self, model: Optional[str]) -> type[Agent]:
+        """Get appropriate agent class based on model."""
+        if not model:
+            model = "model://anthropic/claude-3-5-sonnet-20241022"
+        
+        # Check for CLI agents
+        cli_agents = {
+            "claude_cli": ClaudeCodeAgent,
+            "codex_cli": OpenAICodexAgent,
+            "gemini_cli": GeminiAgent,
+            "grok_cli": GrokAgent,
+        }
+        
+        if model in cli_agents:
+            return cli_agents[model]
+        
+        # Return generic MCP agent
+        return type("DynamicMCPAgent", (MCPAgent,), {
+            "model": model,
+            "__init__": lambda self: MCPAgent.__init__(
+                self,
+                available_tools=self.available_tools,
+                permission_manager=self.permission_manager,
+                ctx=self.ctx,
+                model=model
             )
-
-            try:
-                # Make a final call to get the result
-                final_response = litellm.completion(
-                    model=model,
-                    messages=messages,
-                    temperature=params["temperature"],
-                    timeout=params["timeout"],
-                    max_tokens=params.get("max_tokens"),
-                )
-
-                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) -> str:
-        """Format agent result with metrics.
-
-        Args:
-            result: Raw result from agent
-            execution_time: Execution time in seconds
-
-        Returns:
-            Formatted result with metrics
-        """
-        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
-
+        """Register this agent tool with the MCP server."""
+        tool_self = self
+        
         @mcp_server.tool(name=self.name, description=self.description)
         async def dispatch_agent(
             prompts: str | list[str],
-            ctx: MCPContext
+            ctx: MCPContext,
+            model: Optional[str] = None,
+            use_memory: bool = False,
+            memory_backend: str = "sqlite"
         ) -> str:
-            return await tool_self.call(ctx, prompts=prompts)
+            return await tool_self.call(
+                ctx,
+                prompts=prompts,
+                model=model,
+                use_memory=use_memory,
+                memory_backend=memory_backend
+            )