tunacode-cli 0.0.29__py3-none-any.whl → 0.0.31__py3-none-any.whl

tunacode/core/agents/main.py

@@ -4,16 +4,21 @@ Main agent functionality and coordination for the TunaCode CLI.
 Handles agent creation, configuration, and request processing.
 """
 
+import asyncio
 import json
+import os
 import re
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Optional
+from typing import Any, Iterator, List, Optional, Tuple
 
+from tunacode.constants import READ_ONLY_TOOLS
 from tunacode.core.state import StateManager
 from tunacode.services.mcp import get_mcp_servers
 from tunacode.tools.bash import bash
+from tunacode.tools.glob import glob
 from tunacode.tools.grep import grep
+from tunacode.tools.list_dir import list_dir
 from tunacode.tools.read_file import read_file
 from tunacode.tools.run_command import run_command
 from tunacode.tools.update_file import update_file
@@ -22,6 +27,27 @@ from tunacode.types import (AgentRun, ErrorMessage, FallbackResponse, ModelName,
                             ResponseState, SimpleResult, ToolCallback, ToolCallId, ToolName)
 
 
+class ToolBuffer:
+    """Buffer for collecting read-only tool calls to execute in parallel."""
+
+    def __init__(self):
+        self.read_only_tasks: List[Tuple[Any, Any]] = []
+
+    def add(self, part: Any, node: Any) -> None:
+        """Add a read-only tool call to the buffer."""
+        self.read_only_tasks.append((part, node))
+
+    def flush(self) -> List[Tuple[Any, Any]]:
+        """Return buffered tasks and clear the buffer."""
+        tasks = self.read_only_tasks
+        self.read_only_tasks = []
+        return tasks
+
+    def has_tasks(self) -> bool:
+        """Check if there are buffered tasks."""
+        return len(self.read_only_tasks) > 0
+
+
 # Lazy import for Agent and Tool
 def get_agent_tool():
     import importlib
@@ -37,7 +63,137 @@ def get_model_messages():
     return messages.ModelRequest, messages.ToolReturnPart
 
 
-async def _process_node(node, tool_callback: Optional[ToolCallback], state_manager: StateManager):
+async def execute_tools_parallel(
+    tool_calls: List[Tuple[Any, Any]], callback: ToolCallback, return_exceptions: bool = True
+) -> List[Any]:
+    """
+    Execute multiple tool calls in parallel using asyncio.
+
+    Args:
+        tool_calls: List of (part, node) tuples
+        callback: The tool callback function to execute
+        return_exceptions: Whether to return exceptions or raise them
+
+    Returns:
+        List of results in the same order as input, with exceptions for failed calls
+    """
+    # Get max parallel from environment or default to CPU count
+    max_parallel = int(os.environ.get("TUNACODE_MAX_PARALLEL", os.cpu_count() or 4))
+
+    async def execute_with_error_handling(part, node):
+        try:
+            return await callback(part, node)
+        except Exception as e:
+            return e
+
+    # If we have more tools than max_parallel, execute in batches
+    if len(tool_calls) > max_parallel:
+        results = []
+        for i in range(0, len(tool_calls), max_parallel):
+            batch = tool_calls[i : i + max_parallel]
+            batch_tasks = [execute_with_error_handling(part, node) for part, node in batch]
+            batch_results = await asyncio.gather(*batch_tasks, return_exceptions=return_exceptions)
+            results.extend(batch_results)
+        return results
+    else:
+        tasks = [execute_with_error_handling(part, node) for part, node in tool_calls]
+        return await asyncio.gather(*tasks, return_exceptions=return_exceptions)
+
+
+def batch_read_only_tools(tool_calls: List[Any]) -> Iterator[List[Any]]:
+    """
+    Batch tool calls so read-only tools can be executed in parallel.
+
+    Yields batches where:
+    - Read-only tools are grouped together
+    - Write/execute tools are in their own batch (single item)
+    - Order within each batch is preserved
+
+    Args:
+        tool_calls: List of tool call objects with 'tool' attribute
+
+    Yields:
+        Batches of tool calls
+    """
+    if not tool_calls:
+        return
+
+    current_batch = []
+
+    for tool_call in tool_calls:
+        tool_name = tool_call.tool_name if hasattr(tool_call, "tool_name") else None
+
+        if tool_name in READ_ONLY_TOOLS:
+            # Add to current batch
+            current_batch.append(tool_call)
+        else:
+            # Yield any pending read-only batch
+            if current_batch:
+                yield current_batch
+                current_batch = []
+
+            # Yield write/execute tool as single-item batch
+            yield [tool_call]
+
+    # Yield any remaining read-only tools
+    if current_batch:
+        yield current_batch
+
+
+async def create_buffering_callback(
+    original_callback: ToolCallback, buffer: ToolBuffer, state_manager: StateManager
+) -> ToolCallback:
+    """
+    Create a callback wrapper that buffers read-only tools for parallel execution.
+
+    Args:
+        original_callback: The original tool callback
+        buffer: ToolBuffer instance to store read-only tools
+        state_manager: StateManager for UI access
+
+    Returns:
+        A wrapped callback function
+    """
+
+    async def buffering_callback(part, node):
+        tool_name = getattr(part, "tool_name", None)
+
+        if tool_name in READ_ONLY_TOOLS:
+            # Buffer read-only tools
+            buffer.add(part, node)
+            # Don't execute yet - will be executed in parallel batch
+            return None
+
+        # Non-read-only tool encountered - flush buffer first
+        if buffer.has_tasks():
+            buffered_tasks = buffer.flush()
+
+            # Execute buffered read-only tools in parallel
+            if state_manager.session.show_thoughts:
+                from tunacode.ui import console as ui
+
+                await ui.muted(f"Executing {len(buffered_tasks)} read-only tools in parallel")
+
+            await execute_tools_parallel(buffered_tasks, original_callback)
+
+        # Execute the non-read-only tool
+        return await original_callback(part, node)
+
+    return buffering_callback
+
+
+async def _process_node(
+    node,
+    tool_callback: Optional[ToolCallback],
+    state_manager: StateManager,
+    tool_buffer: Optional[ToolBuffer] = None,
+):
+    from tunacode.ui import console as ui
+    from tunacode.utils.token_counter import estimate_tokens
+
+    # Use the original callback directly - parallel execution will be handled differently
+    buffering_callback = tool_callback
+
     if hasattr(node, "request"):
         state_manager.session.messages.append(node.request)
 
@@ -45,36 +201,81 @@ async def _process_node(node, tool_callback: Optional[ToolCallback], state_manag
         state_manager.session.messages.append({"thought": node.thought})
         # Display thought immediately if show_thoughts is enabled
         if state_manager.session.show_thoughts:
-            from tunacode.ui import console as ui
-
-            await ui.muted(f"💭 THOUGHT: {node.thought}")
+            await ui.muted(f"THOUGHT: {node.thought}")
 
     if hasattr(node, "model_response"):
         state_manager.session.messages.append(node.model_response)
 
-        # Enhanced ReAct thought processing
+        # Enhanced display when thoughts are enabled
         if state_manager.session.show_thoughts:
+            # Show raw API response data
             import json
             import re
 
-            from tunacode.ui import console as ui
+            # Display the raw model response parts
+            await ui.muted("\n" + "=" * 60)
+            await ui.muted(" RAW API RESPONSE DATA:")
+            await ui.muted("=" * 60)
+
+            for idx, part in enumerate(node.model_response.parts):
+                part_data = {"part_index": idx, "part_kind": getattr(part, "part_kind", "unknown")}
+
+                # Add part-specific data
+                if hasattr(part, "content"):
+                    part_data["content"] = (
+                        part.content[:200] + "..." if len(str(part.content)) > 200 else part.content
+                    )
+                if hasattr(part, "tool_name"):
+                    part_data["tool_name"] = part.tool_name
+                if hasattr(part, "args"):
+                    part_data["args"] = part.args
+                if hasattr(part, "tool_call_id"):
+                    part_data["tool_call_id"] = part.tool_call_id
+
+                await ui.muted(json.dumps(part_data, indent=2))
+
+            await ui.muted("=" * 60)
+
+            # Count how many tool calls are in this response
+            tool_count = sum(
+                1
+                for part in node.model_response.parts
+                if hasattr(part, "part_kind") and part.part_kind == "tool-call"
+            )
+            if tool_count > 0:
+                await ui.muted(f"\n MODEL RESPONSE: Contains {tool_count} tool call(s)")
 
+            # Display LLM response content
             for part in node.model_response.parts:
                 if hasattr(part, "content") and isinstance(part.content, str):
                     content = part.content.strip()
 
+                    # Skip empty content
+                    if not content:
+                        continue
+
+                    # Estimate tokens in this response
+                    token_count = estimate_tokens(content)
+
+                    # Display non-JSON content as LLM response
+                    if not content.startswith('{"thought"'):
+                        # Truncate very long responses for display
+                        display_content = content[:500] + "..." if len(content) > 500 else content
+                        await ui.muted(f"\nRESPONSE: {display_content}")
+                        await ui.muted(f"TOKENS: ~{token_count}")
+
                     # Pattern 1: Inline JSON thoughts {"thought": "..."}
                     thought_pattern = r'\{"thought":\s*"([^"]+)"\}'
                     matches = re.findall(thought_pattern, content)
                     for thought in matches:
-                        await ui.muted(f"💭 REASONING: {thought}")
+                        await ui.muted(f"REASONING: {thought}")
 
                     # Pattern 2: Standalone thought JSON objects
                     try:
                         if content.startswith('{"thought"'):
                             thought_obj = json.loads(content)
                             if "thought" in thought_obj:
-                                await ui.muted(f"💭 REASONING: {thought_obj['thought']}")
+                                await ui.muted(f"REASONING: {thought_obj['thought']}")
                     except (json.JSONDecodeError, KeyError):
                         pass
 
@@ -85,41 +286,148 @@ async def _process_node(node, tool_callback: Optional[ToolCallback], state_manag
                         if thought not in [m for m in matches]:  # Avoid duplicates
                             # Clean up escaped characters
                             cleaned_thought = thought.replace('\\"', '"').replace("\\n", " ")
-                            await ui.muted(f"💭 REASONING: {cleaned_thought}")
-
-                    # Pattern 4: Text-based reasoning indicators
-                    reasoning_indicators = [
-                        (r"I need to (.+?)\.", "PLANNING"),
-                        (r"Let me (.+?)\.", "ACTION"),
-                        (r"The output shows (.+?)\.", "OBSERVATION"),
-                        (r"Based on (.+?), I should (.+?)\.", "DECISION"),
-                    ]
-
-                    for pattern, label in reasoning_indicators:
-                        indicator_matches = re.findall(pattern, content, re.IGNORECASE)
-                        for match in indicator_matches:
-                            if isinstance(match, tuple):
-                                match_text = " ".join(match)
-                            else:
-                                match_text = match
-                            await ui.muted(f"🎯 {label}: {match_text}")
-                            break  # Only show first match per pattern
-
-        # Check for tool calls and fallback to JSON parsing if needed
+                            await ui.muted(f"REASONING: {cleaned_thought}")
+
+        # Check for tool calls and collect them for potential parallel execution
         has_tool_calls = False
+        tool_parts = []  # Collect all tool calls from this node
+
         for part in node.model_response.parts:
             if part.part_kind == "tool-call" and tool_callback:
                 has_tool_calls = True
-                await tool_callback(part, node)
-            elif part.part_kind == "tool-return":
+                tool_parts.append(part)
+
+                # Display tool call details when thoughts are enabled
+                if state_manager.session.show_thoughts:
+                    # Show each tool as it's collected
+                    tool_desc = f" COLLECTED: {part.tool_name}"
+                    if hasattr(part, "args") and isinstance(part.args, dict):
+                        if part.tool_name == "read_file" and "file_path" in part.args:
+                            tool_desc += f" → {part.args['file_path']}"
+                        elif part.tool_name == "grep" and "pattern" in part.args:
+                            tool_desc += f" → pattern: '{part.args['pattern']}'"
+                        elif part.tool_name == "list_dir" and "directory" in part.args:
+                            tool_desc += f" → {part.args['directory']}"
+                        elif part.tool_name == "run_command" and "command" in part.args:
+                            tool_desc += f" → {part.args['command']}"
+                    await ui.muted(tool_desc)
+
+                # Track this tool call (moved outside thoughts block)
+                state_manager.session.tool_calls.append(
+                    {
+                        "tool": part.tool_name,
+                        "args": part.args if hasattr(part, "args") else {},
+                        "iteration": state_manager.session.current_iteration,
+                    }
+                )
+
+                # Track files if this is read_file (moved outside thoughts block)
+                if (
+                    part.tool_name == "read_file"
+                    and hasattr(part, "args")
+                    and isinstance(part.args, dict)
+                    and "file_path" in part.args
+                ):
+                    state_manager.session.files_in_context.add(part.args["file_path"])
+                    # Show files in context when thoughts are enabled
+                    if state_manager.session.show_thoughts:
+                        await ui.muted(
+                            f"\nFILES IN CONTEXT: {list(state_manager.session.files_in_context)}"
+                        )
+
+        # Execute tool calls - with ACTUAL parallel execution for read-only batches
+        if tool_parts:
+            if state_manager.session.show_thoughts:
+                await ui.muted(
+                    f"\n NODE SUMMARY: {len(tool_parts)} tool(s) collected in this response"
+                )
+
+            # Check if ALL tools in this node are read-only
+            all_read_only = all(part.tool_name in READ_ONLY_TOOLS for part in tool_parts)
+
+            if all_read_only and len(tool_parts) > 1 and buffering_callback:
+                # Execute read-only tools in parallel!
+                import time
+
+                start_time = time.time()
+
+                if state_manager.session.show_thoughts:
+                    await ui.muted("\n" + "=" * 60)
+                    await ui.muted(
+                        f" PARALLEL BATCH: Executing {len(tool_parts)} read-only tools concurrently"
+                    )
+                    await ui.muted("=" * 60)
+
+                    for idx, part in enumerate(tool_parts, 1):
+                        tool_desc = f"  [{idx}] {part.tool_name}"
+                        if hasattr(part, "args") and isinstance(part.args, dict):
+                            if part.tool_name == "read_file" and "file_path" in part.args:
+                                tool_desc += f" → {part.args['file_path']}"
+                            elif part.tool_name == "grep" and "pattern" in part.args:
+                                tool_desc += f" → pattern: '{part.args['pattern']}'"
+                            elif part.tool_name == "list_dir" and "directory" in part.args:
+                                tool_desc += f" → {part.args['directory']}"
+                            elif part.tool_name == "glob" and "pattern" in part.args:
+                                tool_desc += f" → pattern: '{part.args['pattern']}'"
+                        await ui.muted(tool_desc)
+                    await ui.muted("=" * 60)
+
+                # Execute in parallel
+                tool_tuples = [(part, node) for part in tool_parts]
+                await execute_tools_parallel(tool_tuples, buffering_callback)
+
+                if state_manager.session.show_thoughts:
+                    elapsed_time = (time.time() - start_time) * 1000
+                    sequential_estimate = len(tool_parts) * 100
+                    speedup = sequential_estimate / elapsed_time if elapsed_time > 0 else 1.0
+                    await ui.muted(
+                        f" Parallel batch completed in {elapsed_time:.0f}ms ({speedup:.1f}x faster than sequential)"
+                    )
+
+            else:
+                # Sequential execution for mixed or write/execute tools
+                for part in tool_parts:
+                    if (
+                        state_manager.session.show_thoughts
+                        and part.tool_name not in READ_ONLY_TOOLS
+                    ):
+                        await ui.muted(f"\n SEQUENTIAL: {part.tool_name} (write/execute tool)")
+
+                    # Execute the tool
+                    if buffering_callback:
+                        await buffering_callback(part, node)
+
+        # Handle tool returns
+        for part in node.model_response.parts:
+            if part.part_kind == "tool-return":
                 obs_msg = f"OBSERVATION[{part.tool_name}]: {part.content[:2_000]}"
                 state_manager.session.messages.append(obs_msg)
 
+                # Display tool return when thoughts are enabled
+                if state_manager.session.show_thoughts:
+                    # Truncate for display
+                    display_content = (
+                        part.content[:200] + "..." if len(part.content) > 200 else part.content
+                    )
+                    await ui.muted(f"TOOL RESULT: {display_content}")
+
         # If no structured tool calls found, try parsing JSON from text content
-        if not has_tool_calls and tool_callback:
+        if not has_tool_calls and buffering_callback:
             for part in node.model_response.parts:
                 if hasattr(part, "content") and isinstance(part.content, str):
-                    await extract_and_execute_tool_calls(part.content, tool_callback, state_manager)
+                    await extract_and_execute_tool_calls(
+                        part.content, buffering_callback, state_manager
+                    )
+
+    # Final flush: disabled temporarily while fixing the parallel execution design
+    # The buffer is not being used in the current implementation
+    # if tool_callback and buffer.has_tasks():
+    #     buffered_tasks = buffer.flush()
+    #     if state_manager.session.show_thoughts:
+    #         await ui.muted(
+    #             f"Final flush: Executing {len(buffered_tasks)} remaining read-only tools in parallel"
+    #         )
+    #     await execute_tools_parallel(buffered_tasks, tool_callback)
 
 
 def get_or_create_agent(model: ModelName, state_manager: StateManager) -> PydanticAgent:
@@ -149,7 +457,9 @@ def get_or_create_agent(model: ModelName, state_manager: StateManager) -> Pydant
             system_prompt=system_prompt,
             tools=[
                 Tool(bash, max_retries=max_retries),
+                Tool(glob, max_retries=max_retries),
                 Tool(grep, max_retries=max_retries),
+                Tool(list_dir, max_retries=max_retries),
                 Tool(read_file, max_retries=max_retries),
                 Tool(run_command, max_retries=max_retries),
                 Tool(update_file, max_retries=max_retries),
@@ -276,13 +586,13 @@ async def parse_json_tool_calls(
             if state_manager.session.show_thoughts:
                 from tunacode.ui import console as ui
 
-                await ui.muted(f"🔧 FALLBACK: Executed {tool_name} via JSON parsing")
+                await ui.muted(f"FALLBACK: Executed {tool_name} via JSON parsing")
 
         except Exception as e:
             if state_manager.session.show_thoughts:
                 from tunacode.ui import console as ui
 
-                await ui.error(f"❌ Error executing fallback tool {tool_name}: {str(e)}")
+                await ui.error(f"Error executing fallback tool {tool_name}: {str(e)}")
 
 
 async def extract_and_execute_tool_calls(
@@ -324,13 +634,13 @@ async def extract_and_execute_tool_calls(
                 if state_manager.session.show_thoughts:
                     from tunacode.ui import console as ui
 
-                    await ui.muted(f"🔧 FALLBACK: Executed {tool_data['tool']} from code block")
+                    await ui.muted(f"FALLBACK: Executed {tool_data['tool']} from code block")
 
         except (json.JSONDecodeError, KeyError, Exception) as e:
             if state_manager.session.show_thoughts:
                 from tunacode.ui import console as ui
 
-                await ui.error(f"❌ Error parsing code block tool call: {str(e)}")
+                await ui.error(f"Error parsing code block tool call: {str(e)}")
 
 
 async def process_request(
@@ -349,28 +659,101 @@ async def process_request(
 
     response_state = ResponseState()
 
+    # Reset iteration tracking for this request
+    state_manager.session.iteration_count = 0
+
+    # Create a request-level buffer for batching read-only tools across nodes
+    tool_buffer = ToolBuffer()
+
+    # Show what we're sending to the API when thoughts are enabled
+    if state_manager.session.show_thoughts:
+        from tunacode.ui import console as ui
+
+        await ui.muted("\n" + "=" * 60)
+        await ui.muted("📤 SENDING TO API:")
+        await ui.muted(f"Message: {message}")
+        await ui.muted(f"Model: {model}")
+        await ui.muted(f"Message History Length: {len(mh)}")
+        await ui.muted("=" * 60)
+
     async with agent.iter(message, message_history=mh) as agent_run:
         i = 0
         async for node in agent_run:
-            await _process_node(node, tool_callback, state_manager)
+            state_manager.session.current_iteration = i + 1
+            await _process_node(node, tool_callback, state_manager, tool_buffer)
             if hasattr(node, "result") and node.result and hasattr(node.result, "output"):
                 if node.result.output:
                     response_state.has_user_response = True
             i += 1
+            state_manager.session.iteration_count = i
 
             # Display iteration progress if thoughts are enabled
-            if state_manager.session.show_thoughts and i > 1:
+            if state_manager.session.show_thoughts:
                 from tunacode.ui import console as ui
 
-                await ui.muted(f"🔄 Iteration {i}/{max_iterations}")
+                await ui.muted(f"\nITERATION: {i}/{max_iterations}")
+
+                # Show summary of tools used so far
+                if state_manager.session.tool_calls:
+                    tool_summary = {}
+                    for tc in state_manager.session.tool_calls:
+                        tool_name = tc.get("tool", "unknown")
+                        tool_summary[tool_name] = tool_summary.get(tool_name, 0) + 1
+
+                    summary_str = ", ".join(
+                        [f"{name}: {count}" for name, count in tool_summary.items()]
+                    )
+                    await ui.muted(f"TOOLS USED: {summary_str}")
 
             if i >= max_iterations:
                 if state_manager.session.show_thoughts:
                     from tunacode.ui import console as ui
 
-                    await ui.warning(f"⚠️ Reached maximum iterations ({max_iterations})")
+                    await ui.warning(f"Reached maximum iterations ({max_iterations})")
                 break
 
+        # Final flush: execute any remaining buffered read-only tools
+        if tool_callback and tool_buffer.has_tasks():
+            import time
+
+            from tunacode.ui import console as ui
+
+            buffered_tasks = tool_buffer.flush()
+            start_time = time.time()
+
+            await ui.muted("\n" + "=" * 60)
+            await ui.muted(
+                f"🚀 FINAL BATCH: Executing {len(buffered_tasks)} buffered read-only tools"
+            )
+            await ui.muted("=" * 60)
+
+            for idx, (part, node) in enumerate(buffered_tasks, 1):
+                tool_desc = f"  [{idx}] {part.tool_name}"
+                if hasattr(part, "args") and isinstance(part.args, dict):
+                    if part.tool_name == "read_file" and "file_path" in part.args:
+                        tool_desc += f" → {part.args['file_path']}"
+                    elif part.tool_name == "grep" and "pattern" in part.args:
+                        tool_desc += f" → pattern: '{part.args['pattern']}'"
+                        if "include_files" in part.args:
+                            tool_desc += f", files: '{part.args['include_files']}'"
+                    elif part.tool_name == "list_dir" and "directory" in part.args:
+                        tool_desc += f" → {part.args['directory']}"
+                    elif part.tool_name == "glob" and "pattern" in part.args:
+                        tool_desc += f" → pattern: '{part.args['pattern']}'"
+                await ui.muted(tool_desc)
+            await ui.muted("=" * 60)
+
+            await execute_tools_parallel(buffered_tasks, tool_callback)
+
+            elapsed_time = (time.time() - start_time) * 1000
+            sequential_estimate = len(buffered_tasks) * 100
+            speedup = sequential_estimate / elapsed_time if elapsed_time > 0 else 1.0
+
+            await ui.muted(
+                f"✅ Final batch completed in {elapsed_time:.0f}ms "
+                f"(~{speedup:.1f}x faster than sequential)\n"
+            )
+
         # If we need to add a fallback response, create a wrapper
         if not response_state.has_user_response and i >= max_iterations and fallback_enabled:
             patch_tool_messages("Task incomplete", state_manager=state_manager)
@@ -391,10 +774,10 @@ async def process_request(
 
                             # Track specific operations
                             if tool_name in ["write_file", "update_file"] and hasattr(part, "args"):
-                                if "file_path" in part.args:
+                                if isinstance(part.args, dict) and "file_path" in part.args:
                                     files_modified.add(part.args["file_path"])
                             elif tool_name in ["run_command", "bash"] and hasattr(part, "args"):
-                                if "command" in part.args:
+                                if isinstance(part.args, dict) and "command" in part.args:
                                     commands_run.append(part.args["command"])
 
             # Build fallback response with context
@@ -465,12 +848,25 @@ async def process_request(
             class AgentRunWrapper:
                 def __init__(self, wrapped_run, fallback_result):
                     self._wrapped = wrapped_run
-                    self.result = fallback_result
+                    self._result = fallback_result
                     self.response_state = response_state
 
-                def __getattr__(self, name):
+                def __getattribute__(self, name):
+                    # Handle special attributes first to avoid conflicts
+                    if name in ["_wrapped", "_result", "response_state"]:
+                        return object.__getattribute__(self, name)
+
+                    # Explicitly handle 'result' to return our fallback result
+                    if name == "result":
+                        return object.__getattribute__(self, "_result")
+
                     # Delegate all other attributes to the wrapped object
-                    return getattr(self._wrapped, name)
+                    try:
+                        return getattr(object.__getattribute__(self, "_wrapped"), name)
+                    except AttributeError:
+                        raise AttributeError(
+                            f"'{type(self).__name__}' object has no attribute '{name}'"
+                        )
 
             return AgentRunWrapper(agent_run, SimpleResult(comprehensive_output))
 
@@ -481,8 +877,17 @@ async def process_request(
                 self._wrapped = wrapped_run
                 self.response_state = response_state
 
-            def __getattr__(self, name):
+            def __getattribute__(self, name):
+                # Handle special attributes first
+                if name in ["_wrapped", "response_state"]:
+                    return object.__getattribute__(self, name)
+
                 # Delegate all other attributes to the wrapped object
-                return getattr(self._wrapped, name)
+                try:
+                    return getattr(object.__getattribute__(self, "_wrapped"), name)
+                except AttributeError:
+                    raise AttributeError(
+                        f"'{type(self).__name__}' object has no attribute '{name}'"
+                    )
 
         return AgentRunWithState(agent_run)