emdash-core 0.1.25__py3-none-any.whl → 0.1.37__py3-none-any.whl

emdash_core/agent/runner/agent_runner.py

@@ -0,0 +1,765 @@
+"""Agent runner for LLM-powered exploration.
+
+This module contains the main AgentRunner class that orchestrates
+the agent loop, tool execution, and conversation management.
+"""
+
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Optional
+
+from ...utils.logger import log
+from ...core.exceptions import ContextLengthError
+from ..toolkit import AgentToolkit
+from ..events import AgentEventEmitter, NullEmitter
+from ..providers import get_provider
+from ..providers.factory import DEFAULT_MODEL
+from ..context_manager import (
+    truncate_tool_output,
+    reduce_context_for_retry,
+    is_context_overflow_error,
+)
+from ..prompts import build_system_prompt
+from ..tools.tasks import TaskState
+from ...checkpoint import CheckpointManager
+
+from .utils import SafeJSONEncoder, summarize_tool_result
+from .context import (
+    estimate_context_tokens,
+    get_context_breakdown,
+    maybe_compact_context,
+    emit_context_frame,
+)
+from .plan import PlanMixin
+
+
+class AgentRunner(PlanMixin):
+    """Runs an LLM agent with tool access for code exploration.
+
+    Example:
+        runner = AgentRunner()
+        response = runner.run("How does authentication work in this codebase?")
+        print(response)
+    """
+
+    def __init__(
+        self,
+        toolkit: Optional[AgentToolkit] = None,
+        model: str = DEFAULT_MODEL,
+        system_prompt: Optional[str] = None,
+        emitter: Optional[AgentEventEmitter] = None,
+        max_iterations: int = int(os.getenv("EMDASH_MAX_ITERATIONS", "100")),
+        verbose: bool = False,
+        show_tool_results: bool = False,
+        enable_thinking: Optional[bool] = None,
+        checkpoint_manager: Optional[CheckpointManager] = None,
+    ):
+        """Initialize the agent runner.
+
+        Args:
+            toolkit: AgentToolkit instance. If None, creates default.
+            model: LLM model to use.
+            system_prompt: Custom system prompt. If None, uses default.
+            emitter: Event emitter for streaming output.
+            max_iterations: Maximum tool call iterations.
+            verbose: Whether to print verbose output.
+            show_tool_results: Whether to show detailed tool results.
+            enable_thinking: Enable extended thinking. If None, auto-detect from model.
+            checkpoint_manager: Optional checkpoint manager for git-based checkpoints.
+        """
+        self.toolkit = toolkit or AgentToolkit()
+        self.provider = get_provider(model)
+        self.model = model
+        # Build system prompt
+        if system_prompt:
+            self.system_prompt = system_prompt
+        else:
+            self.system_prompt = build_system_prompt(self.toolkit)
+        self.emitter = emitter or NullEmitter()
+        # Inject emitter into tools that need it (e.g., TaskTool for sub-agent streaming)
+        self.toolkit.set_emitter(self.emitter)
+        self.max_iterations = max_iterations
+        self.verbose = verbose
+        self.show_tool_results = show_tool_results
+        # Extended thinking support
+        if enable_thinking is None:
+            # Auto-detect from provider capabilities
+            self.enable_thinking = (
+                hasattr(self.provider, "supports_thinking")
+                and self.provider.supports_thinking()
+            )
+        else:
+            self.enable_thinking = enable_thinking
+        # Conversation history for multi-turn support
+        self._messages: list[dict] = []
+        # Token usage tracking
+        self._total_input_tokens: int = 0
+        self._total_output_tokens: int = 0
+        self._total_thinking_tokens: int = 0
+        # Store query for reranking
+        self._current_query: str = ""
+        # Todo state tracking for injection
+        self._last_todo_snapshot: str = ""
+        # Checkpoint manager for git-based checkpoints
+        self._checkpoint_manager = checkpoint_manager
+        # Track tools used during current run (for checkpoint metadata)
+        self._tools_used_this_run: set[str] = set()
+        # Plan approval state (from PlanMixin)
+        self._pending_plan: Optional[dict] = None
+        # Callback for autosave after each iteration (set by API layer)
+        self._on_iteration_callback: Optional[callable] = None
+
+    def _get_default_plan_file_path(self) -> str:
+        """Get the default plan file path based on repo root.
+
+        Returns:
+            Path to the plan file (e.g., .emdash/plan.md)
+        """
+        repo_root = self.toolkit._repo_root
+        return str(repo_root / ".emdash" / "plan.md")
+
+    def _get_todo_snapshot(self) -> str:
+        """Get current todo state as string for comparison."""
+        state = TaskState.get_instance()
+        return json.dumps(state.get_all_tasks(), sort_keys=True)
+
+    def _format_todo_reminder(self) -> str:
+        """Format current todos as XML reminder for injection into context."""
+        state = TaskState.get_instance()
+        tasks = state.get_all_tasks()
+        if not tasks:
+            return ""
+
+        counts = {"pending": 0, "in_progress": 0, "completed": 0}
+        lines = []
+        for t in tasks:
+            status = t.get("status", "pending")
+            counts[status] = counts.get(status, 0) + 1
+            status_icon = {"pending": "⬚", "in_progress": "🔄", "completed": "✅"}.get(status, "?")
+            lines.append(f'  {t["id"]}. {status_icon} {t["title"]}')
+
+        header = f'Tasks: {counts["completed"]} completed, {counts["in_progress"]} in progress, {counts["pending"]} pending'
+        task_list = "\n".join(lines)
+        return f"<todo-state>\n{header}\n{task_list}\n</todo-state>"
+
+    def _execute_tools_parallel(self, parsed_calls: list) -> list:
+        """Execute multiple tool calls in parallel using a thread pool.
+
+        Args:
+            parsed_calls: List of (tool_call, args) tuples
+
+        Returns:
+            List of (tool_call, args, result) tuples in original order
+        """
+        # Emit tool start events for all calls
+        for tool_call, args in parsed_calls:
+            self.emitter.emit_tool_start(tool_call.name, args, tool_id=tool_call.id)
+
+        def execute_one(item):
+            tool_call, args = item
+            try:
+                result = self.toolkit.execute(tool_call.name, **args)
+                return (tool_call, args, result)
+            except Exception as e:
+                log.exception(f"Tool {tool_call.name} failed")
+                from ..tools.base import ToolResult
+                return (tool_call, args, ToolResult.error_result(str(e)))
+
+        # Execute in parallel with up to 3 workers
+        results: list = [None] * len(parsed_calls)
+        with ThreadPoolExecutor(max_workers=3) as executor:
+            futures = {executor.submit(execute_one, item): i for i, item in enumerate(parsed_calls)}
+            # Collect results maintaining order
+            for future in as_completed(futures):
+                idx = futures[future]
+                results[idx] = future.result()
+
+        # Emit tool result events for all calls
+        for tool_call, args, result in results:
+            self.emitter.emit_tool_result(
+                tool_call.name,
+                result.success,
+                summarize_tool_result(result),
+                tool_id=tool_call.id,
+            )
+
+        return results
+
+    def run(
+        self,
+        query: str,
+        context: Optional[str] = None,
+        images: Optional[list] = None,
+    ) -> str:
+        """Run the agent to answer a query.
+
+        Args:
+            query: User's question or request
+            context: Optional additional context
+            images: Optional list of images to include
+
+        Returns:
+            Agent's final response
+        """
+        # Store query for reranking context frame
+        self._current_query = query
+
+        # Reset per-cycle mode state (allows exit_plan to be called again)
+        from ..tools.modes import ModeState
+        ModeState.get_instance().reset_cycle()
+
+        # Build user message
+        if context:
+            user_message = {
+                "role": "user",
+                "content": f"Context:\n{context}\n\nQuestion: {query}",
+            }
+        else:
+            user_message = {
+                "role": "user",
+                "content": query,
+            }
+
+        # Save user message to history BEFORE running (so it's preserved even if interrupted)
+        self._messages.append(user_message)
+        messages = list(self._messages)  # Copy for the loop
+
+        # TODO: Handle images if provided
+
+        # Get tool schemas
+        tools = self.toolkit.get_all_schemas()
+
+        try:
+            response, final_messages = self._run_loop(messages, tools)
+            # Update conversation history with full exchange
+            self._messages = final_messages
+            self.emitter.emit_end(success=True)
+            # Create checkpoint if manager is configured
+            self._create_checkpoint()
+            return response
+
+        except Exception as e:
+            log.exception("Agent run failed")
+            self.emitter.emit_error(str(e))
+            # Keep user message in history even on error (already appended above)
+            return f"Error: {str(e)}"
+
+    def _run_loop(
+        self,
+        messages: list[dict],
+        tools: list[dict],
+    ) -> tuple[str, list[dict]]:
+        """Run the agent loop until completion.
+
+        Args:
+            messages: Initial messages
+            tools: Tool schemas
+
+        Returns:
+            Tuple of (final response text, conversation messages)
+        """
+        max_retries = 3
+
+        for iteration in range(self.max_iterations):
+            # When approaching max iterations, ask agent to wrap up
+            if iteration == self.max_iterations - 2:
+                messages.append({
+                    "role": "user",
+                    "content": "[SYSTEM: You are approaching your iteration limit. Please provide your findings and conclusions now, even if incomplete. Summarize what you've learned and any recommendations.]",
+                })
+
+            # Try API call with retry on context overflow
+            retry_count = 0
+            response = None
+
+            while retry_count < max_retries:
+                try:
+                    # Proactively compact context if approaching limit
+                    messages = maybe_compact_context(
+                        messages, self.provider, self.emitter, self.system_prompt
+                    )
+
+                    log.debug(
+                        "Calling LLM iteration={} messages={} tools={}",
+                        iteration,
+                        len(messages),
+                        len(tools) if tools else 0,
+                    )
+
+                    response = self.provider.chat(
+                        messages=messages,
+                        system=self.system_prompt,
+                        tools=tools,
+                        thinking=self.enable_thinking,
+                    )
+
+                    log.debug("LLM response received iteration={}", iteration)
+                    break  # Success
+
+                except Exception as exc:
+                    if is_context_overflow_error(exc):
+                        retry_count += 1
+                        log.warning(
+                            "Context overflow on attempt {}/{}, reducing context...",
+                            retry_count,
+                            max_retries,
+                        )
+
+                        if retry_count >= max_retries:
+                            raise ContextLengthError(
+                                f"Failed to reduce context after {max_retries} attempts: {exc}",
+                            )
+
+                        # Reduce context by removing old messages
+                        messages = reduce_context_for_retry(
+                            messages,
+                            keep_recent=max(2, 6 - retry_count * 2),  # Fewer messages each retry
+                        )
+                    else:
+                        raise  # Re-raise non-context errors
+
+            if response is None:
+                raise RuntimeError("Failed to get response from provider")
+
+            # Accumulate token usage
+            self._total_input_tokens += response.input_tokens
+            self._total_output_tokens += response.output_tokens
+            self._total_thinking_tokens += getattr(response, "thinking_tokens", 0)
+
+            # Emit thinking if present
+            if response.thinking:
+                self.emitter.emit_thinking(response.thinking)
+
+            # Check for tool calls
+            if response.tool_calls:
+                # Emit assistant text if present (shown as bullets between tool calls)
+                if response.content and response.content.strip():
+                    self.emitter.emit_assistant_text(response.content)
+
+                # Track if we need to pause for user input
+                needs_user_input = False
+
+                # Parse all tool call arguments first
+                parsed_calls = []
+                for tool_call in response.tool_calls:
+                    args = tool_call.arguments
+                    if isinstance(args, str):
+                        args = json.loads(args)
+                    parsed_calls.append((tool_call, args))
+
+                # CRITICAL: Check if exit_plan is in the batch - if so, execute it FIRST
+                # and skip all other tools. This prevents the agent from continuing to
+                # work after submitting a plan.
+                exit_plan_idx = None
+                for i, (tc, _) in enumerate(parsed_calls):
+                    if tc.name == "exit_plan":
+                        exit_plan_idx = i
+                        break
+
+                if exit_plan_idx is not None:
+                    # Execute ONLY exit_plan, skip everything else
+                    exit_call, exit_args = parsed_calls[exit_plan_idx]
+                    self.emitter.emit_tool_start(exit_call.name, exit_args, tool_id=exit_call.id)
+                    exit_result = self.toolkit.execute(exit_call.name, **exit_args)
+                    self.emitter.emit_tool_result(
+                        exit_call.name,
+                        exit_result.success,
+                        summarize_tool_result(exit_result),
+                        tool_id=exit_call.id,
+                    )
+
+                    # Build results list with exit_plan result and skipped placeholders
+                    results = []
+                    from ..tools.base import ToolResult
+                    for i, (tc, args) in enumerate(parsed_calls):
+                        if i == exit_plan_idx:
+                            results.append((tc, args, exit_result))
+                        else:
+                            # Skip this tool - don't execute it
+                            log.warning(f"Skipping tool {tc.name} - exit_plan takes priority")
+                            skip_result = ToolResult.error_result(
+                                "Tool skipped: exit_plan was called. Agent must stop and wait for user approval."
+                            )
+                            results.append((tc, args, skip_result))
+
+                elif len(parsed_calls) > 1:
+                    # No exit_plan - execute tools in parallel
+                    results = self._execute_tools_parallel(parsed_calls)
+                else:
+                    # Single tool - execute directly
+                    tool_call, args = parsed_calls[0]
+                    self.emitter.emit_tool_start(tool_call.name, args, tool_id=tool_call.id)
+                    result = self.toolkit.execute(tool_call.name, **args)
+                    self.emitter.emit_tool_result(
+                        tool_call.name,
+                        result.success,
+                        summarize_tool_result(result),
+                        tool_id=tool_call.id,
+                    )
+                    results = [(tool_call, args, result)]
+
+                # Track if we need to rebuild toolkit for mode change
+                mode_changed = False
+
+                # CRITICAL FIX: Add ONE assistant message with ALL tool calls
+                # This prevents the LLM from seeing multiple assistant messages
+                # which causes it to loop repeating the same tools
+                all_tool_calls = []
+                for tool_call, args in parsed_calls:
+                    all_tool_calls.append({
+                        "id": tool_call.id,
+                        "type": "function",
+                        "function": {
+                            "name": tool_call.name,
+                            "arguments": json.dumps(args),
+                        },
+                    })
+
+                messages.append({
+                    "role": "assistant",
+                    "content": response.content or "",
+                    "tool_calls": all_tool_calls,
+                })
+
+                # Now process results and add tool result messages
+                for tool_call, args, result in results:
+                    # Track tool for checkpoint metadata
+                    self._tools_used_this_run.add(tool_call.name)
+
+                    # Check if tool is asking a clarification question
+                    if (result.success and
+                        result.data and
+                        result.data.get("status") == "awaiting_response" and
+                        "question" in result.data):
+                        self.emitter.emit_clarification(
+                            question=result.data["question"],
+                            context="",
+                            options=result.data.get("options", []),
+                        )
+                        needs_user_input = True
+
+                    # Check if agent entered plan mode
+                    if (result.success and
+                        result.data and
+                        result.data.get("status") == "entered_plan_mode"):
+                        mode_changed = True
+                        # Get plan file path
+                        plan_file_path = self._get_default_plan_file_path()
+                        from ..tools.modes import ModeState
+                        ModeState.get_instance().set_plan_file_path(plan_file_path)
+                        # Rebuild toolkit with plan_mode=True
+                        self.toolkit = AgentToolkit(
+                            connection=self.toolkit.connection,
+                            repo_root=self.toolkit._repo_root,
+                            plan_mode=True,
+                            plan_file_path=plan_file_path,
+                        )
+                        self.toolkit.set_emitter(self.emitter)
+                        # Main agent uses normal prompt - delegates to Plan subagent
+                        self.system_prompt = build_system_prompt(self.toolkit)
+                        # Update tools for LLM
+                        tools = self.toolkit.get_all_schemas()
+
+                    # Check if agent requested to enter plan mode (enter_plan_mode)
+                    if (result.success and
+                        result.data and
+                        result.data.get("status") == "plan_mode_requested"):
+                        # Emit event for UI to show approval dialog
+                        self.emitter.emit_plan_mode_requested(
+                            reason=result.data.get("reason", ""),
+                        )
+                        # Pause and wait for user approval
+                        needs_user_input = True
+
+                    # Check if tool is submitting a plan for approval (exit_plan)
+                    if (result.success and
+                        result.data and
+                        result.data.get("status") == "plan_submitted"):
+                        # Store the pending plan (simple string)
+                        self._pending_plan = {
+                            "plan": result.data.get("plan", ""),
+                        }
+                        self.emitter.emit_plan_submitted(
+                            plan=self._pending_plan["plan"],
+                        )
+                        # Pause and wait for approval (similar to clarification flow)
+                        needs_user_input = True
+
+                    # Serialize and truncate tool result to prevent context overflow
+                    result_json = json.dumps(result.to_dict(), cls=SafeJSONEncoder)
+                    result_json = truncate_tool_output(result_json)
+
+                    # Check if todos changed and inject reminder
+                    if tool_call.name in ("write_todo", "update_todo_list"):
+                        new_snapshot = self._get_todo_snapshot()
+                        if new_snapshot != self._last_todo_snapshot:
+                            self._last_todo_snapshot = new_snapshot
+                            reminder = self._format_todo_reminder()
+                            if reminder:
+                                result_json += f"\n\n{reminder}"
+
+                    # Add tool result
+                    messages.append({
+                        "role": "tool",
+                        "tool_call_id": tool_call.id,
+                        "content": result_json,
+                    })
+
+                # Emit context frame after each iteration (for autosave and UI updates)
+                self._emit_context_frame(messages)
+
+                # If a clarification question was asked, pause and wait for user input
+                if needs_user_input:
+                    log.debug("Pausing agent loop - waiting for user input")
+                    return "", messages
+
+            else:
+                # No tool calls - check if response was truncated
+                if response.stop_reason in ("max_tokens", "length"):
+                    # Response was truncated, request continuation
+                    log.debug("Response truncated ({}), requesting continuation", response.stop_reason)
+                    if response.content:
+                        messages.append({
+                            "role": "assistant",
+                            "content": response.content,
+                        })
+                    messages.append({
+                        "role": "user",
+                        "content": "Your response was cut off. Please continue.",
+                    })
+                    continue
+
+                # PLAN MODE ENFORCEMENT: In plan mode, reject text-only responses
+                # Force the model to use tools (task, ask_followup_question, exit_plan)
+                if self.toolkit.plan_mode:
+                    log.warning("Plan mode: Agent output text without tools, forcing tool usage")
+                    if response.content:
+                        messages.append({
+                            "role": "assistant",
+                            "content": response.content,
+                        })
+                    messages.append({
+                        "role": "user",
+                        "content": """[SYSTEM ERROR] You are in plan mode but did not use any tools.
+
+In plan mode you MUST use tools - text-only responses are not allowed.
+
+YOUR REQUIRED ACTION NOW:
+Use your exploration tools directly to investigate the codebase:
+
+- glob(pattern="**/*.py") - Find files by pattern
+- grep(pattern="class User", path="src/") - Search file contents
+- read_file(path="path/to/file.py") - Read specific files
+- semantic_search(query="authentication") - Find conceptually related code
+
+You ARE the planner. Use tools to explore, then write your plan and call exit_plan.
+
+DO NOT output more text. Use a tool NOW.""",
+                    })
+                    continue  # Force another iteration with tool usage
+
+                # Agent is done - emit final response
+                if response.content:
+                    self.emitter.emit_message_start()
+                    self.emitter.emit_message_delta(response.content)
+                    self.emitter.emit_message_end()
+                    # Add final assistant message to history
+                    messages.append({
+                        "role": "assistant",
+                        "content": response.content,
+                    })
+                    # Emit final context frame summary
+                    self._emit_context_frame(messages)
+                    return response.content, messages
+
+                # Agent finished without providing a summary - request one
+                log.debug("Agent finished without summary, requesting completion message")
+                try:
+                    summary_response = self.provider.chat(
+                        messages=messages + [{
+                            "role": "user",
+                            "content": "[SYSTEM: You have completed your task. Please provide a brief summary of what you accomplished.]",
+                        }],
+                        system=self.system_prompt,
+                        tools=None,  # No tools - force text response
+                        thinking=self.enable_thinking,
+                    )
+                    if summary_response.content:
+                        self.emitter.emit_message_start()
+                        self.emitter.emit_message_delta(summary_response.content)
+                        self.emitter.emit_message_end()
+                        self._emit_context_frame(messages)
+                        return summary_response.content, messages
+                except Exception as e:
+                    log.warning(f"Failed to get completion summary: {e}")
+
+                # Fallback if summary request fails
+                fallback_message = "Task completed."
+                self.emitter.emit_message_start()
+                self.emitter.emit_message_delta(fallback_message)
+                self.emitter.emit_message_end()
+                self._emit_context_frame(messages)
+                return fallback_message, messages
+
+        # Hit max iterations - try one final request without tools to force a response
+        try:
+            final_response = self.provider.chat(
+                messages=messages + [{
+                    "role": "user",
+                    "content": "[SYSTEM: Maximum iterations reached. Provide your final response now with whatever information you have gathered. Do not use any tools.]",
+                }],
+                system=self.system_prompt,
+                tools=None,  # No tools - force text response
+                thinking=self.enable_thinking,
+            )
+            # Emit thinking if present
+            if final_response.thinking:
+                self.emitter.emit_thinking(final_response.thinking)
+            if final_response.content:
+                self.emitter.emit_message_start()
+                self.emitter.emit_message_delta(final_response.content)
+                self.emitter.emit_message_end()
+                self._emit_context_frame(messages)
+                return final_response.content, messages
+        except Exception as e:
+            log.warning(f"Failed to get final response: {e}")
+
+        # Fallback message if final response fails
+        final_message = "Reached maximum iterations. The agent was unable to complete the task within the allowed iterations."
+        self.emitter.emit_message_start()
+        self.emitter.emit_message_delta(final_message)
+        self.emitter.emit_message_end()
+        self._emit_context_frame(messages)
+        return final_message, messages
+
+    def _emit_context_frame(self, messages: list[dict] | None = None) -> None:
+        """Emit a context frame event with current exploration state.
+
+        Args:
+            messages: Current conversation messages to estimate context size
+        """
+        emit_context_frame(
+            toolkit=self.toolkit,
+            emitter=self.emitter,
+            messages=messages or [],
+            system_prompt=self.system_prompt,
+            current_query=self._current_query,
+            total_input_tokens=self._total_input_tokens,
+            total_output_tokens=self._total_output_tokens,
+        )
+
+        # Call iteration callback for autosave if set
+        if self._on_iteration_callback and messages:
+            try:
+                self._on_iteration_callback(messages)
+            except Exception as e:
+                log.debug(f"Iteration callback failed: {e}")
+
+    def chat(self, message: str, images: Optional[list] = None) -> str:
+        """Continue a conversation with a new message.
+
+        This method maintains conversation history for multi-turn interactions.
+        Call run() first to start a conversation, then chat() for follow-ups.
+
+        Args:
+            message: User's follow-up message
+            images: Optional list of images to include
+
+        Returns:
+            Agent's response
+        """
+        if not self._messages:
+            # No history, just run fresh
+            return self.run(message, images=images)
+
+        # Store query for reranking context frame
+        self._current_query = message
+
+        # Add new user message to history
+        self._messages.append({
+            "role": "user",
+            "content": message,
+        })
+
+        # Get tool schemas
+        tools = self.toolkit.get_all_schemas()
+
+        try:
+            response, final_messages = self._run_loop(self._messages, tools)
+            # Update conversation history
+            self._messages = final_messages
+            self.emitter.emit_end(success=True)
+            # Create checkpoint if manager is configured
+            self._create_checkpoint()
+            return response
+
+        except Exception as e:
+            log.exception("Agent chat failed")
+            self.emitter.emit_error(str(e))
+            return f"Error: {str(e)}"
+
+    def _create_checkpoint(self) -> None:
+        """Create a git checkpoint after successful run.
+
+        Only creates a checkpoint if:
+        - A checkpoint manager is configured
+        - There are file changes to commit
+        """
+        if not self._checkpoint_manager:
+            return
+
+        try:
+            self._checkpoint_manager.create_checkpoint(
+                messages=self._messages,
+                model=self.model,
+                system_prompt=self.system_prompt,
+                tools_used=list(self._tools_used_this_run),
+                token_usage={
+                    "input": self._total_input_tokens,
+                    "output": self._total_output_tokens,
+                    "thinking": self._total_thinking_tokens,
+                },
+            )
+        except Exception as e:
+            log.warning(f"Failed to create checkpoint: {e}")
+        finally:
+            # Clear tools for next run
+            self._tools_used_this_run.clear()
+
+    def reset(self) -> None:
+        """Reset the agent state."""
+        self.toolkit.reset_session()
+        self._total_input_tokens = 0
+        self._total_output_tokens = 0
+        self._current_query = ""
+
+    def answer_clarification(self, answer: str) -> str:
+        """Answer a pending clarification question and resume the agent.
+
+        This method is called when the user responds to a clarification question
+        asked via ask_followup_question tool. It clears the pending question state
+        and resumes the agent loop with the user's answer.
+
+        Args:
+            answer: The user's answer to the clarification question
+
+        Returns:
+            Agent's response after processing the answer
+        """
+        # Get current task state and clear the pending question
+        state = TaskState.get_instance()
+        pending_question = state.pending_question
+
+        # Clear the pending question state
+        state.pending_question = None
+        state.user_response = answer
+
+        # Build a context message that indicates this is the answer to the question
+        if pending_question:
+            context_message = f"[User answered the clarification question]\nQuestion: {pending_question}\nAnswer: {answer}"
+        else:
+            context_message = f"[User response]: {answer}"
+
+        # Continue the conversation with the answer
+        return self.chat(context_message)