import asyncio
import copy
import threading
import time
import uuid
from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Literal, Optional, Union, TYPE_CHECKING


from upsonic.utils.logging_config import sentry_sdk, get_env_bool_optional, get_logger

_pl_logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# Persistent event loop for sync wrappers (do, stream, print_do, etc.)
#
# asyncio.run() creates a *new* event loop each call and closes it afterward.
# Cached httpx.AsyncClient connections inside the OpenAI SDK stay bound to
# the first loop.  When asyncio.run() destroys that loop, subsequent calls
# fail with "RuntimeError: Event loop is closed".
#
# Solution: keep a single background loop alive for the process lifetime so
# async clients and their connection pools remain valid across calls.
# ---------------------------------------------------------------------------
_bg_loop: Optional[asyncio.AbstractEventLoop] = None
_bg_loop_lock = threading.Lock()


def _get_bg_loop() -> asyncio.AbstractEventLoop:
    """Return a persistent event loop running in a daemon thread."""
    global _bg_loop
    if _bg_loop is not None and not _bg_loop.is_closed():
        return _bg_loop
    with _bg_loop_lock:
        # Double-check after acquiring lock
        if _bg_loop is not None and not _bg_loop.is_closed():
            return _bg_loop
        loop = asyncio.new_event_loop()
        thread = threading.Thread(target=loop.run_forever, daemon=True, name="upsonic-bg-loop")
        thread.start()
        _bg_loop = loop
        return _bg_loop


def _run_in_bg_loop(coro):
    """Submit *coro* to the persistent background loop and block until done.
    
    If already running on the background loop thread, runs the coroutine
    directly via a nested event loop to avoid deadlock.
    """
    loop = _get_bg_loop()
    try:
        running_loop = asyncio.get_running_loop()
    except RuntimeError:
        running_loop = None
    
    if running_loop is loop:
        import nest_asyncio as _nest_asyncio
        _nest_asyncio.apply(loop)
        return loop.run_until_complete(coro)
    
    future = asyncio.run_coroutine_threadsafe(coro, loop)
    return future.result()
from upsonic.agent.base import BaseAgent
from upsonic.run.agent.output import AgentRunOutput
from upsonic.run.agent.input import AgentRunInput
from upsonic.run.base import RunStatus

from upsonic._utils import now_utc
from upsonic.utils.retry import retryable
from upsonic.tools.processor import ExternalExecutionPause, ConfirmationPause, UserInputPause
from upsonic.run.cancel import register_run, cleanup_run, raise_if_cancelled, cancel_run as cancel_run_func, is_cancelled
from upsonic.session.base import SessionType
from upsonic.output import DEFAULT_OUTPUT_TOOL_NAME

if TYPE_CHECKING:
    from upsonic.models import Model, ModelRequest, ModelRequestParameters, ModelResponse
    from upsonic.messages import ToolCallPart, ToolReturnPart
    from upsonic.tasks.tasks import Task
    from upsonic.storage.memory.memory import Memory
    from upsonic.canvas.canvas import Canvas
    from upsonic.models.settings import ModelSettings
    from upsonic.profiles import ModelProfile
    from upsonic.reflection import ReflectionConfig
    from upsonic.safety_engine.base import Policy
    from upsonic.tools import ToolDefinition, ToolManager
    from upsonic.skills import Skills
    from upsonic.run.tools.tools import ToolExecution
    from upsonic.run.requirements import RunRequirement
    from upsonic.usage import RequestUsage, TaskUsage, AgentUsage
    from upsonic.agent.context_managers import (
        MemoryManager
    )
    from upsonic.graph.graph import State
    from upsonic.run.events.events import AgentStreamEvent
    from upsonic.db.database import DatabaseBase
    from upsonic.models.model_selector import ModelRecommendation
    from upsonic.culture.culture import Culture
    from upsonic.culture.manager import CultureManager
    from upsonic.session.agent import RunData
    from upsonic.models.instrumented import InstrumentationSettings
    from upsonic.integrations.tracing import TracingProvider
    from upsonic.integrations.promptlayer import PromptLayer
    from upsonic.agent.otel_manager import AgentOTelManager
    from fastmcp import FastMCP
else:
    Model = "Model"
    ModelRequest = "ModelRequest"
    ModelRequestParameters = "ModelRequestParameters"
    ModelResponse = "ModelResponse"
    Task = "Task"
    Memory = "Memory"
    Canvas = "Canvas"
    ModelSettings = "ModelSettings"
    ModelProfile = "ModelProfile"
    ReflectionConfig = "ReflectionConfig"
    Policy = "Policy"
    ToolDefinition = "ToolDefinition"
    RequestUsage = "RequestUsage"
    MemoryManager = "MemoryManager"
    CulturalKnowledge = "CulturalKnowledge"
    CultureManager = "CultureManager"
    State = "State"
    ModelRecommendation = "ModelRecommendation"
    DatabaseBase = "DatabaseBase"
    InstrumentationSettings = "InstrumentationSettings"
    TracingProvider = "TracingProvider"
    RunData = "RunData"


PromptCompressor = None

RetryMode = Literal["raise", "return_false"]


class Agent(BaseAgent):
    """
    A comprehensive, high-level AI Agent that integrates all framework components.

    This Agent class provides:
    - Complete model abstraction through Model/Provider/Profile system
    - Advanced tool handling with ToolManager and Orchestrator
    - Streaming and non-streaming execution modes
    - Memory management and conversation history
    - Context management and prompt engineering
    - Caching capabilities
    - Safety policies and guardrails
    - Reliability layers
    - Canvas integration
    - External tool execution support
    - OpenTelemetry instrumentation for full observability
    
    Usage:
        Basic usage:
        ```python
        from upsonic import Agent, Task
        
        agent = Agent("openai/gpt-4o")
        task = Task("What is 1 + 1?")
        result = agent.do(task)
        ```
        
        Advanced usage:
        ```python
        agent = Agent(
            model="openai/gpt-4o",
            name="Math Teacher",
            memory=memory,
            enable_thinking_tool=True,
            user_policy=safety_policy
        )
        result = agent.stream(task)
        ```
        
        With OpenTelemetry:
        ```python
        agent = Agent("openai/gpt-4o", instrument=True)
        result = agent.do("What is 2 + 2?")
        
        # Or instrument all agents globally:
        Agent.instrument_all()
        ```
    """

    _global_tracing_provider: Optional["TracingProvider"] = None

    @classmethod
    def instrument_all(
        cls,
        instrument: Union[bool, "TracingProvider", "InstrumentationSettings"] = True,
    ) -> None:
        """Enable OpenTelemetry instrumentation globally for all Agent instances.

        Args:
            instrument: If True, creates a DefaultTracingProvider from env vars.
                If a TracingProvider subclass, uses it directly.
                If an InstrumentationSettings instance, wraps it.
                If False, disables global instrumentation.
        """
        if instrument is False:
            cls._global_tracing_provider = None
            return

        from upsonic.integrations.tracing import TracingProvider as _TP

        if instrument is True:
            from upsonic.integrations.tracing import DefaultTracingProvider
            cls._global_tracing_provider = DefaultTracingProvider()
        elif isinstance(instrument, _TP):
            cls._global_tracing_provider = instrument
        else:
            cls._global_tracing_provider = instrument  # type: ignore[assignment]

    def __init__(
        self,
        model: Union[str, "Model"] = "openai/gpt-4o",
        *,
        name: Optional[str] = None,
        memory: Optional["Memory"] = None,
        db: Optional["DatabaseBase"] = None,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        debug: bool = False,
        debug_level: int = 1,
        print: Optional[bool] = None,
        company_url: Optional[str] = None,
        company_objective: Optional[str] = None,
        company_description: Optional[str] = None,
        company_name: Optional[str] = None,
        system_prompt: Optional[str] = None,
        reflection: bool = False,
        context_management: bool = False,
        context_management_keep_recent: int = 5,
        context_management_model: Optional[str] = None,
        reliability_layer: Optional[Any] = None,
        agent_id_: Optional[str] = None,
        canvas: Optional["Canvas"] = None,
        retry: int = 1,
        mode: RetryMode = "raise",
        role: Optional[str] = None,
        goal: Optional[str] = None,
        instructions: Optional[str] = None,
        education: Optional[str] = None,
        work_experience: Optional[str] = None,
        feed_tool_call_results: Optional[bool] = None,
        show_tool_calls: bool = True,
        tool_call_limit: int = 100,
        enable_thinking_tool: bool = False,
        enable_reasoning_tool: bool = False,
        tools: Optional[list] = None,
        skills: Optional["Skills"] = None,
        user_policy: Optional[Union["Policy", List["Policy"]]] = None,
        agent_policy: Optional[Union["Policy", List["Policy"]]] = None,
        tool_policy_pre: Optional[Union["Policy", List["Policy"]]] = None,
        tool_policy_post: Optional[Union["Policy", List["Policy"]]] = None,
        # Policy feedback loop settings
        user_policy_feedback: bool = False,
        agent_policy_feedback: bool = False,
        user_policy_feedback_loop: int = 1,
        agent_policy_feedback_loop: int = 1,
        settings: Optional["ModelSettings"] = None,
        profile: Optional["ModelProfile"] = None,
        reflection_config: Optional["ReflectionConfig"] = None,
        model_selection_criteria: Optional[Dict[str, Any]] = None,
        use_llm_for_selection: bool = False,
        # Common reasoning/thinking attributes
        reasoning_effort: Optional[Literal["low", "medium", "high"]] = None,
        reasoning_summary: Optional[Literal["concise", "detailed"]] = None,
        thinking_enabled: Optional[bool] = None,
        thinking_budget: Optional[int] = None,
        thinking_include_thoughts: Optional[bool] = None,
        reasoning_format: Optional[Literal["hidden", "raw", "parsed"]] = None,
        culture: Optional["Culture"] = None,
        # Agent metadata (passed to prompt)
        metadata: Optional[Dict[str, Any]] = None,
        # Workspace settings
        workspace: Optional[str] = None,
        # OpenTelemetry instrumentation
        instrument: Union[bool, "TracingProvider", "InstrumentationSettings", None] = None,
        # PromptLayer integration
        promptlayer: Optional["PromptLayer"] = None,
    ):
        """
        Initialize the Agent with comprehensive configuration options.
        
        Args:
            model: Model identifier or Model instance
            name: Agent name for identification
            memory: Memory instance for conversation history
            db: Database instance (overrides memory if provided)
            debug: Enable debug logging
            debug_level: Debug level (1 = standard, 2 = detailed). Only used when debug=True
            print: Enable printing for do() (and allow print_do() when not False). If None, do() does not print unless UPSONIC_AGENT_PRINT=true. If False, print_do() also does not print. UPSONIC_AGENT_PRINT=false overrides everything.
            company_url: Company URL for context
            company_objective: Company objective for context
            company_description: Company description for context
            system_prompt: Custom system prompt
            reflection: Reflection capabilities (default is False)
            context_management: Enable automatic context window management (default True).
                When enabled, the middleware automatically prunes tool call history and
                summarizes old messages when the context approaches the model's limit.
            context_management_keep_recent: Number of recent tool-call events /
                messages to preserve when the context management middleware prunes or
                summarizes the history (default 5).
            context_management_model: Optional model identifier (e.g. 'openai/gpt-4.1')
                with a larger context window to use specifically for context compression /
                summarization. If None, the agent's primary model is used.
            reliability_layer: Reliability layer for robustness
            agent_id_: Specific agent ID
            canvas: Canvas instance for visual interactions
            retry: Number of retry attempts
            mode: Retry mode behavior
            role: Agent role
            goal: Agent goal
            instructions: Specific instructions
            education: Agent education background
            work_experience: Agent work experience
            feed_tool_call_results: Include tool results in memory
            show_tool_calls: Display tool calls
            tool_call_limit: Maximum tool calls per execution
            enable_thinking_tool: Enable orchestrated thinking
            enable_reasoning_tool: Enable reasoning capabilities
            tools: List of tools to register with this agent (can be functions, ToolKits, or other agents)
            user_policy: User input safety policy (single policy or list of policies)
            agent_policy: Agent output safety policy (single policy or list of policies)
            settings: Model-specific settings
            profile: Model profile configuration
            reflection_config: Configuration for reflection and self-evaluation
            model_selection_criteria: Default criteria dictionary for recommend_model_for_task() (see SelectionCriteria)
            use_llm_for_selection: Default flag for whether to use LLM in recommend_model_for_task()
            
            # Common reasoning/thinking attributes (mapped to model-specific settings):
            reasoning_effort: Reasoning effort level for OpenAI models ("low", "medium", "high")
            reasoning_summary: Reasoning summary type for OpenAI models ("concise", "detailed")
            thinking_enabled: Enable thinking for Anthropic/Google models (True/False)
            thinking_budget: Token budget for thinking (Anthropic: budget_tokens, Google: thinking_budget)
            thinking_include_thoughts: Include thoughts in output (Google models)
            reasoning_format: Reasoning format for Groq models ("hidden", "raw", "parsed")
            tool_policy_pre: Tool safety policy for pre-execution validation (single policy or list of policies)
            tool_policy_post: Tool safety policy for post-execution validation (single policy or list of policies)
            user_policy_feedback: Enable feedback loop for user policy violations (returns helpful message instead of blocking)
            agent_policy_feedback: Enable feedback loop for agent policy violations (re-executes agent with feedback)
            user_policy_feedback_loop: Maximum retry count for user policy feedback (default 1)
            agent_policy_feedback_loop: Maximum retry count for agent policy feedback (default 1)
            
            culture: Culture instance defining agent behavior and communication guidelines.
                Includes description, add_system_prompt, repeat, and repeat_interval settings.
            workspace: Path to workspace folder containing AGENTS.md file with agent configuration.
                When set, the AGENTS.md content is included in system prompt and a greeting 
                message is generated before the first task/chat, integrated into message history.
            instrument: OpenTelemetry instrumentation configuration.
                If True, creates a DefaultTracingProvider from env vars.
                If a TracingProvider subclass (Langfuse, DefaultTracingProvider, …),
                    uses its InstrumentationSettings.
                If an InstrumentationSettings instance, uses it directly.
                If None/False, no instrumentation (default).
                When enabled, all LLM calls, pipeline steps, and tool executions
                are traced with OpenTelemetry spans following GenAI semantic conventions.
            promptlayer: PromptLayer integration instance.
                When set, every agent execution automatically logs the request
                (task description, output, model, cost) to PromptLayer for
                tracking, scoring, and prompt version correlation.
        """
        from upsonic.models import infer_model
        self.model = infer_model(model)
        self.model_name=model

        self._tracing_provider, self._instrument_settings, self._otel = self._resolve_instrumentation(instrument)
        self.promptlayer: Optional["PromptLayer"] = promptlayer
        self._suppress_promptlayer_logging: bool = False

        self.name = name
        self.agent_id_ = agent_id_
        
        # Session/user overrides
        self._override_session_id = session_id
        self._override_user_id = user_id
        
        # Common reasoning/thinking attributes
        self.reasoning_effort = reasoning_effort
        self.reasoning_summary = reasoning_summary
        self.thinking_enabled = thinking_enabled
        self.thinking_budget = thinking_budget
        self.thinking_include_thoughts = thinking_include_thoughts
        self.reasoning_format = reasoning_format
        
        self.role = role
        self.goal = goal
        self.instructions = instructions
        self.education = education
        self.work_experience = work_experience
        self._user_system_prompt: Optional[str] = system_prompt
        self._last_built_system_prompt: Optional[str] = None
        
        self.company_url = company_url
        self.company_objective = company_objective
        self.company_description = company_description
        self.company_name = company_name
        
        self.debug = debug
        self.debug_level = debug_level if debug else 1
        self.reflection = reflection
        self._print_env: Optional[bool] = get_env_bool_optional("UPSONIC_AGENT_PRINT")
        self._print_param: Optional[bool] = print
        self.print: bool = self._print_env if self._print_env is not None else (print if print is not None else False)

        self.db = db
        
        if db is not None:
            self.memory = db.memory
        else:
            self.memory = memory
        
        self.model_selection_criteria = model_selection_criteria
        self.use_llm_for_selection = use_llm_for_selection
        self._model_recommendation: Optional[Any] = None  # Store last recommendation

        self.context_management: bool = context_management
        self.context_management_keep_recent: int = context_management_keep_recent
        self.context_management_model: Optional[str] = context_management_model
        self._context_management_middleware: Optional[Any] = None

        if self.context_management:
            from upsonic.agent.context_managers import ContextManagementMiddleware
            from upsonic.models import infer_model as _infer_model

            compression_model: Optional[Any] = None
            if self.context_management_model is not None:
                compression_model = _infer_model(self.context_management_model)

            self._context_management_middleware = ContextManagementMiddleware(
                model=self.model,
                keep_recent_count=self.context_management_keep_recent,
                context_compression_model=compression_model,
            )

        self.reliability_layer = reliability_layer
        
        if retry < 1:
            raise ValueError("The 'retry' count must be at least 1.")
        if mode not in ("raise", "return_false"):
            raise ValueError(f"Invalid retry_mode '{mode}'. Must be 'raise' or 'return_false'.")
        
        self.retry = retry
        self.mode = mode
        
        self.show_tool_calls = show_tool_calls
        self.tool_call_limit = tool_call_limit
        self.enable_thinking_tool = enable_thinking_tool
        self.enable_reasoning_tool = enable_reasoning_tool
        
        self.tools = tools if tools is not None else []
        self.skills = skills

        # Register skill tools if skills are provided
        if self.skills is not None:
            self.tools.extend(self.skills.get_tools())

        if self.memory and feed_tool_call_results is not None:
            self.memory.feed_tool_call_results = feed_tool_call_results
        
        self.canvas = canvas
        
        # Agent metadata (injected into prompts)
        self.metadata = metadata or {}
        
        self._culture_input = culture
        self._culture_manager: Optional["CultureManager"] = None
        if culture is not None:
            from upsonic.culture.manager import CultureManager
            self._culture_manager = CultureManager(
                model=self.model_name,
                debug=self.debug,
                debug_level=self.debug_level,
                print=self.print,
            )
            self._culture_manager.set_culture(culture)
        
        # Initialize policy managers
        from upsonic.agent.policy_manager import PolicyManager
        self.user_policy_manager = PolicyManager(
            policies=user_policy,
            debug=self.debug,
            enable_feedback=user_policy_feedback,
            feedback_loop_count=user_policy_feedback_loop,
            policy_type="user_policy"
        )
        self.agent_policy_manager = PolicyManager(
            policies=agent_policy,
            debug=self.debug,
            enable_feedback=agent_policy_feedback,
            feedback_loop_count=agent_policy_feedback_loop,
            policy_type="agent_policy"
        )
        
        # Store feedback settings for reference
        self.user_policy_feedback = user_policy_feedback
        self.agent_policy_feedback = agent_policy_feedback
        self.user_policy_feedback_loop = user_policy_feedback_loop
        self.agent_policy_feedback_loop = agent_policy_feedback_loop
        
        # Keep backward compatibility - expose as single policy if only one
        self.user_policy = user_policy
        self.agent_policy = agent_policy
        
        # Initialize tool policy managers
        from upsonic.agent.tool_policy_manager import ToolPolicyManager
        self.tool_policy_pre_manager = ToolPolicyManager(policies=tool_policy_pre, debug=self.debug)
        self.tool_policy_post_manager = ToolPolicyManager(policies=tool_policy_post, debug=self.debug)
        
        # Keep references
        self.tool_policy_pre = tool_policy_pre
        self.tool_policy_post = tool_policy_post
        
        # Handle reflection configuration
        if reflection and not reflection_config:
            # Create default reflection config if reflection=True but no config provided
            from upsonic.reflection import ReflectionConfig
            reflection_config = ReflectionConfig()
        
        self.reflection_config = reflection_config
        if reflection_config:
            from upsonic.reflection import ReflectionProcessor
            self.reflection_processor = ReflectionProcessor(reflection_config)
        else:
            self.reflection_processor = None
        
        if settings:
            self.model._settings = settings
        if profile:
            self.model._profile = profile
            
        self._apply_reasoning_settings()
        
        from upsonic.cache import CacheManager
        from upsonic.tools import ToolManager
        
        self._cache_manager = CacheManager(session_id=f"agent_{self.agent_id}")
        self.tool_manager = ToolManager()
        
        # Track registered agent tools
        self.registered_agent_tools = {}
        
        # Track agent-level builtin tools
        self.agent_builtin_tools = []
        
        # Register agent-level tools immediately
        self._register_agent_tools()
        
        # Tool tracking (deprecated - now tracked in AgentRunOutput)
        # Kept for backwards compatibility
        self._tool_call_count = 0
        self._tool_limit_reached = False
        
        # Agent-level accumulated usage across all tasks
        self.usage: Optional["AgentUsage"] = None
        
        # Run cancellation tracking
        self.run_id: Optional[str] = None
        
        self._setup_policy_models()

        self.session_type = SessionType.AGENT
        
        # Workspace settings
        self.workspace: Optional[str] = workspace
        self._workspace_greeting_executed: bool = False
        self._workspace_agents_md_content: Optional[str] = None
        
        # Pre-load workspace AGENTS.md content if workspace is set
        if self.workspace:
            self._workspace_agents_md_content = self._read_workspace_agents_md()
    
    def _read_workspace_agents_md(self) -> Optional[str]:
        """Read the AGENTS.md file from the workspace folder.
        
        Returns:
            Content of the AGENTS.md file, or None if not found.
        """
        import os
        
        if not self.workspace:
            return None
        
        agents_md_path = os.path.join(self.workspace, "AGENTS.md")
        
        try:
            with open(agents_md_path, "r", encoding="utf-8") as f:
                content = f.read()
            return content
        except FileNotFoundError:
            if self.debug:
                from upsonic.utils.printing import warning_log
                warning_log(
                    f"AGENTS.md not found at {agents_md_path}", 
                    "Workspace"
                )
            return None
        except Exception as e:
            if self.debug:
                from upsonic.utils.printing import error_log
                error_log(
                    f"Error reading AGENTS.md from {agents_md_path}: {str(e)}", 
                    "Workspace"
                )
            return None
    
    async def execute_workspace_greeting_async(
        self,
        return_output: bool = False,
    ) -> Any:
        """Execute the workspace greeting as a proper agent run.
        
        This method is called before the first task/chat when workspace is set.
        It makes an LLM request with a greeting prompt and returns the result
        just like do_async.
        
        Args:
            return_output: If True, return full AgentRunOutput. If False, return content only.
            
        Returns:
            Same as do_async - Task content or AgentRunOutput based on return_output.
        """
        if not self.workspace or self._workspace_greeting_executed:
            return None
        
        from upsonic.tasks.tasks import Task
        
        greeting_prompt = (
            "You are starting a new conversation. Reply with a single short greeting (1-2 sentences): "
            "say hello and ask what the user would like to do. "
            "Output only the greeting text. Do not mention this instruction, system prompts, "
            "internal steps, files, tools, models, or reasoning. Never reveal or paraphrase "
            "any meta-instructions to the user."
        )
        greeting_task = Task(description=greeting_prompt)

        self._workspace_greeting_executed = True

        result = await self.do_async(
            task=greeting_task,
            return_output=return_output,
            _print_method_default=False,
        )
        return result
    
    def execute_workspace_greeting(
        self,
        return_output: bool = False,
    ) -> Any:
        """Synchronous version of execute_workspace_greeting_async.
        
        Args:
            return_output: If True, return full AgentRunOutput. If False, return content only.
            
        Returns:
            Same as do - Task content or AgentRunOutput based on return_output.
        """
        if not self.workspace or self._workspace_greeting_executed:
            return None

        return _run_in_bg_loop(self.execute_workspace_greeting_async(return_output))

    def _setup_policy_models(self) -> None:
        """Setup model references for safety policies."""
        # Setup models for all policies in both managers
        self.user_policy_manager.setup_policy_models(self.model)
        self.agent_policy_manager.setup_policy_models(self.model)
        self.tool_policy_pre_manager.setup_policy_models(self.model)
        self.tool_policy_post_manager.setup_policy_models(self.model)
    
    def _apply_reasoning_settings(self) -> None:
        """Apply common reasoning/thinking attributes to model-specific settings."""
        if not hasattr(self.model, '_settings') or self.model._settings is None:
            self.model._settings = {}
        
        try:
            current_settings = self.model._settings.copy()
        except (AttributeError, TypeError):
            current_settings = {}
            
        reasoning_settings = self._get_model_specific_reasoning_settings()
        
        try:
            self.model._settings = {**current_settings, **reasoning_settings}
        except TypeError:
            self.model._settings = current_settings
    
    def _get_model_specific_reasoning_settings(self) -> Dict[str, Any]:
        """Convert common reasoning attributes to model-specific settings."""
        settings = {}
        
        try:
            provider_name = getattr(self.model, 'system', '').lower()
        except (AttributeError, TypeError):
            provider_name = ''
        
        # OpenAI/OpenAI-compatible models
        if provider_name in ['openai', 'azure', 'deepseek', 'cerebras', 'fireworks', 'github', 'grok', 'heroku', 'moonshotai', 'openrouter', 'together', 'vercel', 'litellm']:
            # Apply reasoning_effort to all OpenAI models
            if self.reasoning_effort is not None:
                settings['openai_reasoning_effort'] = self.reasoning_effort
            
            # Only apply reasoning_summary to OpenAIResponsesModel
            if self.reasoning_summary is not None:
                from upsonic.models.openai import OpenAIResponsesModel
                if isinstance(self.model, OpenAIResponsesModel):
                    settings['openai_reasoning_summary'] = self.reasoning_summary
        
        # Anthropic models
        elif provider_name == 'anthropic':
            if self.thinking_enabled is not None or self.thinking_budget is not None:
                thinking_config = {}
                if self.thinking_enabled is not None:
                    thinking_config['type'] = 'enabled' if self.thinking_enabled else 'disabled'
                if self.thinking_budget is not None:
                    thinking_config['budget_tokens'] = self.thinking_budget
                settings['anthropic_thinking'] = thinking_config
        
        # Google models
        elif provider_name in ['google-gla', 'google-vertex']:
            if self.thinking_enabled is not None or self.thinking_budget is not None or self.thinking_include_thoughts is not None:
                thinking_config = {}
                if self.thinking_enabled is not None:
                    thinking_config['include_thoughts'] = self.thinking_include_thoughts if self.thinking_include_thoughts is not None else self.thinking_enabled
                if self.thinking_budget is not None:
                    thinking_config['thinking_budget'] = self.thinking_budget
                settings['google_thinking_config'] = thinking_config
        
        # Groq models
        elif provider_name == 'groq':
            if self.reasoning_format is not None:
                settings['groq_reasoning_format'] = self.reasoning_format
        
        return settings
    
    @property
    def system_prompt(self) -> Optional[str]:
        """Return the latest fully-built system prompt, falling back to the
        user-provided value when no build has occurred yet."""
        if self._last_built_system_prompt is not None:
            return self._last_built_system_prompt
        return self._user_system_prompt

    @system_prompt.setter
    def system_prompt(self, value: Optional[str]) -> None:
        self._user_system_prompt = value
        self._last_built_system_prompt = None

    @property
    def agent_id(self) -> str:
        """Get or generate agent ID."""
        if self.agent_id_ is None:
            self.agent_id_ = str(uuid.uuid4())
        return self.agent_id_
    
    @property
    def session_id(self) -> Optional[str]:
        """Get session_id from override, memory, or db."""
        if self._override_session_id:
            return self._override_session_id
        if self.memory and hasattr(self.memory, 'session_id'):
            return self.memory.session_id
        if self.db and hasattr(self.db, 'session_id'):
            return self.db.session_id
        if self.db and hasattr(self.db, 'memory') and hasattr(self.db.memory, 'session_id'):
            return self.db.memory.session_id
        return None
    
    @property
    def user_id(self) -> Optional[str]:
        """Get user_id from override, memory, or db."""
        if self._override_user_id:
            return self._override_user_id
        if self.memory and hasattr(self.memory, 'user_id'):
            return self.memory.user_id
        if self.db and hasattr(self.db, 'user_id'):
            return self.db.user_id
        if self.db and hasattr(self.db, 'memory') and hasattr(self.db.memory, 'user_id'):
            return self.db.memory.user_id
        return None
    
    def get_agent_id(self) -> str:
        """Get display-friendly agent ID."""
        if self.name:
            return self.name
        return f"Agent_{self.agent_id[:8]}"
    
    def get_entity_id(self) -> str:
        """Get entity ID for unified interface (Agent + Team)."""
        return self.get_agent_id()
    
    def get_cache_stats(self) -> Dict[str, Any]:
        """Get cache statistics for this agent's session."""
        return self._cache_manager.get_cache_stats()
    
    def clear_cache(self) -> None:
        """Clear the agent's session cache."""
        self._cache_manager.clear_cache()
    
    def get_run_output(self) -> Optional[AgentRunOutput]:
        """
        Get the AgentRunOutput from the last execution.
        
        Returns:
            AgentRunOutput: The complete run output, or None if no run has been executed
        """
        return getattr(self, '_agent_run_output', None)
    
    def get_session_usage(self) -> "TaskUsage":
        """
        Get the aggregated usage for the current session.

        Returns the session-level usage from the AgentSession stored in storage.
        If no memory is configured, returns an empty TaskUsage.

        Usage:
            ```python
            agent = Agent("openai/gpt-4o", memory=memory)
            result = agent.do(task1)
            result = agent.do(task2)

            session_usage = agent.get_session_usage()
            print(session_usage.to_dict())
            ```

        Returns:
            TaskUsage: Aggregated usage metrics for the session.
        """
        from upsonic.usage import TaskUsage as TaskUsageCls

        if not self.memory:
            return TaskUsageCls()

        session = self.memory.get_session()
        if session and hasattr(session, 'get_session_usage'):
            return session.get_session_usage()

        return TaskUsageCls()

    async def aget_session_usage(self) -> "TaskUsage":
        """
        Get the aggregated usage for the current session (async version).

        Returns the session-level usage from the AgentSession stored in storage.
        If no memory is configured, returns an empty TaskUsage.

        Returns:
            TaskUsage: Aggregated usage metrics for the session.
        """
        from upsonic.usage import TaskUsage as TaskUsageCls

        if not self.memory:
            return TaskUsageCls()

        session = await self.memory.get_session_async()
        if session and hasattr(session, 'get_session_usage'):
            return session.get_session_usage()

        return TaskUsageCls()
    
    def _create_agent_run_input(self, task: "Task") -> AgentRunInput:
        """
        Create AgentRunInput from Task, separating images and documents.
        
        Categorizes file attachments by mime type without processing:
        - Images: jpg, png, gif, webp, etc.
        - Documents: pdf, docx, txt, etc.
        
        Args:
            task: The task with attachments
            
        Returns:
            AgentRunInput with user_prompt, images (file paths), and documents (file paths)
        """
        import mimetypes
        
        images: List[str] = []
        documents: List[str] = []
        
        if task.attachments:
            for file_path in task.attachments:
                mime_type, _ = mimetypes.guess_type(file_path)
                if mime_type is None:
                    mime_type = "application/octet-stream"
                
                if mime_type.startswith('image/'):
                    images.append(file_path)
                else:
                    documents.append(file_path)
        
        return AgentRunInput(
            user_prompt=task.description,
            images=images if images else None,
            documents=documents if documents else None
        )
    
    def _convert_to_task(self, task: Union[str, "Task"]) -> "Task":
        """
        Convert a string to a Task object if needed.
        
        Args:
            task: Task object or string description
            
        Returns:
            Task: Task object (converted if string was provided)
        """
        from upsonic.tasks.tasks import Task as TaskClass
        
        if isinstance(task, str):
            return TaskClass(description=task)
        return task

    def _handle_task_list(
        self,
        task: Union[str, "Task", List[Union[str, "Task"]]],
        executor: Callable[..., Any],
        *args: Any,
        **kwargs: Any,
    ) -> tuple[bool, Any]:
        """
        Handle list-of-tasks input for synchronous execution methods.

        Returns:
            (handled, result):
                handled=True  → result is the final return value (empty list or list of results)
                handled=False → result is the unwrapped single task for normal processing
        """
        if not isinstance(task, list):
            return False, task
        if len(task) == 0:
            return True, []
        if len(task) == 1:
            return False, task[0]
        results: List[Any] = []
        for single_task in task:
            result: Any = executor(single_task, *args, **kwargs)
            results.append(result)
        return True, results

    async def _handle_task_list_async(
        self,
        task: Union[str, "Task", List[Union[str, "Task"]]],
        executor: Callable[..., Any],
        *args: Any,
        **kwargs: Any,
    ) -> tuple[bool, Any]:
        """
        Handle list-of-tasks input for asynchronous execution methods.

        Returns:
            (handled, result):
                handled=True  → result is the final return value (empty list or list of results)
                handled=False → result is the unwrapped single task for normal processing
        """
        if not isinstance(task, list):
            return False, task
        if len(task) == 0:
            return True, []
        if len(task) == 1:
            return False, task[0]
        results: List[Any] = []
        for single_task in task:
            result: Any = await executor(single_task, *args, **kwargs)
            results.append(result)
        return True, results

    def _resolve_print_flag(self, method_default: bool) -> bool:
        """
        Resolve the print flag based on hierarchy.
        
        Priority (highest to lowest):
        1. UPSONIC_AGENT_PRINT env variable (overrides everything)
        2. Agent constructor print parameter
        3. Method default (print_do=True, do=False)
        
        Args:
            method_default: Default value based on method name (True for print_do, False for do)
            
        Returns:
            bool: Whether to print output
        """
        # 1. ENV has highest priority - overrides everything
        if self._print_env is not None:
            return self._print_env
        
        # 2. Agent constructor parameter
        if self._print_param is not None:
            return self._print_param
        
        # 3. Method default (print_do=True, do=False)
        return method_default
    
    def _validate_task_for_new_run(
        self,
        task: "Task",
        is_resuming: bool = False,
        allow_problematic_for_retry: bool = False,
    ) -> Optional[AgentRunOutput]:
        """
        Validate task state before starting a new run.

        Checks if task is already completed or has problematic status.
        When allow_problematic_for_retry is True (internal retry attempt), problematic
        status is allowed so retry can override durable-execution requirement until retries are exhausted.

        Args:
            task: Task to validate
            is_resuming: Whether this is a resume operation (skips problematic check)
            allow_problematic_for_retry: If True, do not block on problematic status (used when retrying).

        Returns:
            AgentRunOutput if task cannot be run (error/warning case), None if OK to proceed
        """
        from upsonic.utils.printing import warning_log

        if task.is_completed:
            run_id = task.run_id or "unknown"
            warning_log(
                f"Task is already completed (run_id={run_id}). Cannot re-run a completed task.",
                "Agent"
            )
            return AgentRunOutput(
                run_id=run_id,
                agent_id=self.agent_id,
                agent_name=self.name,
                session_id=self.session_id,
                user_id=self.user_id,
                status=RunStatus.completed,
                output=f"Task is already completed (run_id={run_id}). Cannot re-run a completed task.",
            )

        if not is_resuming and not allow_problematic_for_retry and task.is_problematic:
            status_str = task.status.value
            run_id = task.run_id
            warning_log(
                f"Task has a problematic run (run_id={run_id}, status={status_str}). "
                f"Use continue_run_async() to continue this run.",
                "Agent"
            )
            return AgentRunOutput(
                run_id=run_id,
                agent_id=self.agent_id,
                agent_name=self.name,
                session_id=self.session_id,
                user_id=self.user_id,
                status=task.status,
                output=f"Cannot start new run: Task has a {status_str} run (run_id={run_id}). "
                       f"Call continue_run_async() to continue this run.",
            )
        
        return None  # OK to proceed
    
    def _create_agent_run_output(
        self,
        run_id: str,
        task: "Task",
        run_input: AgentRunInput,
        is_streaming: bool = False
    ) -> AgentRunOutput:
        """
        Create AgentRunOutput with all required fields.
        
        This is a centralized factory method for creating AgentRunOutput instances,
        ensuring consistency across do_async and astream methods.
        
        Args:
            run_id: Unique run identifier
            task: Task being executed
            run_input: Agent run input data
            is_streaming: Whether this is a streaming execution
            
        Returns:
            AgentRunOutput: Initialized output context
        """
        from upsonic.schemas.kb_filter import KBFilterExpr
        
        kb_filter = KBFilterExpr.from_task(task) if hasattr(KBFilterExpr, 'from_task') else None
        
        return AgentRunOutput(
            run_id=run_id,
            agent_id=self.agent_id,
            agent_name=self.name,
            session_id=self.session_id,
            user_id=self.user_id,
            task=task,
            input=run_input,
            output=None,
            output_schema=task.response_format if hasattr(task, 'response_format') else None,
            thinking_content=None,
            thinking_parts=None,
            model_name=self.model.model_name if self.model else None,
            model_provider=self.model.system if self.model else None,
            model_provider_profile=self.model.profile if self.model else None,
            chat_history=[],
            messages=None,  # Will be set by finalize_run_messages()
            response=None,
            usage=None,
            additional_input_message=None,
            memory_message_count=0,
            tools=[],
            tool_call_count=0,
            tool_limit_reached=False,
            images=None,
            files=None,
            status=RunStatus.running,
            requirements=[],
            step_results=[],
            execution_stats=None,
            events=[],
            agent_knowledge_base_filter=kb_filter,
            metadata=None,
            session_state=None,
            is_streaming=is_streaming,
            accumulated_text="",
            pause_reason=None,
            error_details=None,
            created_at=int(now_utc().timestamp()),
            updated_at=None
        )
    
    def _resolve_instrumentation(
        self,
        instrument: Union[bool, "TracingProvider", "InstrumentationSettings", None],
    ) -> tuple[Optional["TracingProvider"], Optional["InstrumentationSettings"], "AgentOTelManager"]:
        """Resolve the ``instrument`` parameter into a TracingProvider, InstrumentationSettings, and AgentOTelManager.

        Handles all supported input types (bool, TracingProvider, InstrumentationSettings, None)
        and wraps ``self.model`` with ``InstrumentedModel`` when instrumentation is active.
        """
        from upsonic.integrations.tracing import TracingProvider as _TP
        from upsonic.agent.otel_manager import AgentOTelManager

        resolved: Any = instrument if instrument is not None else self._global_tracing_provider
        tracing_provider: Optional["TracingProvider"] = None
        settings: Optional["InstrumentationSettings"] = None

        if resolved is True:
            from upsonic.integrations.tracing import DefaultTracingProvider
            tracing_provider = DefaultTracingProvider()
            settings = tracing_provider.settings
        elif isinstance(resolved, _TP):
            tracing_provider = resolved
            settings = resolved.settings
        elif resolved and resolved is not False:
            settings = resolved
        
        if settings is not None:
            from upsonic.models.instrumented import instrument_model
            self.model = instrument_model(self.model, settings)

        return tracing_provider, settings, AgentOTelManager(settings, tracing_provider)

    async def _log_to_promptlayer_unified(
        self,
        task: "Task",
        output: Any,
        start_time: float,
        end_time: float,
    ) -> None:
        if self.promptlayer is None or self._suppress_promptlayer_logging:
            return
        try:
            import json as _json
            from upsonic.eval._pl_helpers import extract_model_parameters, accumulate_agent_usage

            task_description: str = str(task.description) if task is not None else ""
            output_text: str = str(output) if output is not None else ""
            model_name: str = str(self.model_name) if self.model_name else "unknown"

            provider: str
            model_short: str
            provider, model_short = self.promptlayer._parse_provider_model(model_name)

            tags: List[str] = ["upsonic-agent"]
            if self.name:
                tags.append(f"agent:{self.name}")

            # ── Collect all metrics from AgentRunOutput ───────────────
            run_output: Any = getattr(self, "_agent_run_output", None)
            run_usage: Any = getattr(run_output, "usage", None) if run_output else None

            input_tokens: int
            output_tokens: int
            price: float
            input_tokens, output_tokens, price = accumulate_agent_usage(self)

            metadata_dict: Dict[str, Any] = {
                "agent_name": self.name or "",
                "model": model_name,
            }

            total_cost: Optional[float] = self._calculate_aggregated_cost()
            if total_cost is not None:
                metadata_dict["total_cost"] = total_cost

            # TaskUsage metrics
            if run_usage is not None:
                requests = getattr(run_usage, "requests", 0) or 0
                if requests:
                    metadata_dict["requests"] = requests

                cache_write = getattr(run_usage, "cache_write_tokens", 0) or 0
                if cache_write:
                    metadata_dict["cache_write_tokens"] = cache_write

                cache_read = getattr(run_usage, "cache_read_tokens", 0) or 0
                if cache_read:
                    metadata_dict["cache_read_tokens"] = cache_read

                reasoning = getattr(run_usage, "reasoning_tokens", 0) or 0
                if reasoning:
                    metadata_dict["reasoning_tokens"] = reasoning

                total_tok = getattr(run_usage, "total_tokens", 0) or 0
                if total_tok:
                    metadata_dict["total_tokens"] = total_tok

                input_audio = getattr(run_usage, "input_audio_tokens", 0) or 0
                if input_audio:
                    metadata_dict["input_audio_tokens"] = input_audio

                output_audio = getattr(run_usage, "output_audio_tokens", 0) or 0
                if output_audio:
                    metadata_dict["output_audio_tokens"] = output_audio

                model_exec_time = getattr(run_usage, "model_execution_time", None)
                if model_exec_time is not None and model_exec_time > 0:
                    metadata_dict["model_execution_time"] = round(model_exec_time, 3)

                tool_exec_time = getattr(run_usage, "tool_execution_time", None)
                if tool_exec_time is not None and tool_exec_time > 0:
                    metadata_dict["tool_execution_time"] = round(tool_exec_time, 3)

                framework_time = getattr(run_usage, "upsonic_execution_time", None)
                if framework_time is not None and framework_time > 0:
                    metadata_dict["framework_overhead_time"] = round(framework_time, 3)

                duration = getattr(run_usage, "duration", None)
                if duration is not None and duration > 0:
                    metadata_dict["duration"] = round(duration, 3)

                ttft = getattr(run_usage, "time_to_first_token", None)
                if ttft is not None and ttft > 0:
                    metadata_dict["time_to_first_token"] = round(ttft, 3)

            # AgentRunOutput metadata
            if run_output is not None:
                run_id = getattr(run_output, "run_id", None)
                if run_id:
                    metadata_dict["run_id"] = run_id

                status = getattr(run_output, "status", None)
                if status is not None:
                    metadata_dict["status"] = str(status.value) if hasattr(status, "value") else str(status)

                out_model_name = getattr(run_output, "model_name", None)
                if out_model_name:
                    metadata_dict["model_name"] = out_model_name

                out_model_provider = getattr(run_output, "model_provider", None)
                if out_model_provider:
                    metadata_dict["model_provider"] = out_model_provider

                tool_call_count = getattr(run_output, "tool_call_count", 0) or 0
                if tool_call_count:
                    metadata_dict["tool_call_count"] = tool_call_count

                tool_limit_reached = getattr(run_output, "tool_limit_reached", False)
                if tool_limit_reached:
                    metadata_dict["tool_limit_reached"] = True

                is_streaming = getattr(run_output, "is_streaming", False)
                if is_streaming:
                    metadata_dict["is_streaming"] = True

                # Pipeline execution stats
                exec_stats = getattr(run_output, "execution_stats", None)
                if exec_stats is not None:
                    executed_steps = getattr(exec_stats, "executed_steps", 0) or 0
                    if executed_steps:
                        metadata_dict["pipeline_executed_steps"] = executed_steps

                    total_steps = getattr(exec_stats, "total_steps", 0) or 0
                    if total_steps:
                        metadata_dict["pipeline_total_steps"] = total_steps

                    step_timing = getattr(exec_stats, "step_timing", None)
                    if step_timing:
                        metadata_dict["pipeline_step_timing"] = step_timing

                    step_statuses = getattr(exec_stats, "step_statuses", None)
                    if step_statuses:
                        metadata_dict["pipeline_step_statuses"] = step_statuses

            model_parameters: Optional[Dict[str, Any]] = extract_model_parameters(self)

            pl_tools: Optional[List[Dict[str, Any]]] = self._build_pl_tools()

            # ── Build tool calls and tool results from AgentRunOutput.tools ──
            pl_tool_calls: Optional[List[Dict[str, Any]]] = None
            pl_tool_results: Optional[List[Dict[str, Any]]] = None

            tool_executions = getattr(run_output, "tools", None) if run_output else None
            if tool_executions:
                pl_tool_calls = []
                pl_tool_results = []
                for idx, te in enumerate(tool_executions):
                    call_id = getattr(te, "tool_call_id", None) or f"call_{idx}"
                    tool_name = getattr(te, "tool_name", None) or "unknown"
                    tool_args = getattr(te, "tool_args", None) or {}
                    result = getattr(te, "result", None)

                    pl_tool_calls.append({
                        "type": "function",
                        "id": call_id,
                        "function": {
                            "name": tool_name,
                            "arguments": _json.dumps(tool_args, default=str),
                        },
                    })

                    pl_tool_results.append({
                        "tool_call_id": call_id,
                        "name": tool_name,
                        "content": str(result)[:5000] if result is not None else "",
                    })
            elif task is not None and task.tool_calls:
                # Fallback to legacy task.tool_calls
                pl_tool_calls = [
                    {
                        "type": "function",
                        "id": tc.get("tool_call_id", f"call_{idx}"),
                        "function": {
                            "name": tc.get("tool_name", ""),
                            "arguments": _json.dumps(tc.get("params", {}), default=str),
                        },
                    }
                    for idx, tc in enumerate(task.tool_calls)
                ]

            request_id: int = await self.promptlayer.alog(
                provider=provider,
                model=model_short,
                input_text=task_description,
                output_text=output_text,
                start_time=start_time,
                end_time=end_time,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                price=price,
                parameters=model_parameters,
                tags=tags,
                metadata=metadata_dict,
                function_name=model_name,
                prompt_name=self.promptlayer._last_prompt_name,
                prompt_id=self.promptlayer._last_prompt_id,
                prompt_version=self.promptlayer._last_prompt_version,
                system_prompt=self.system_prompt,
                tools=pl_tools,
                tool_calls=pl_tool_calls,
                tool_results=pl_tool_results,
            )

            if task is not None:
                task._promptlayer_request_id = request_id

            # ── Register agent as a workflow in PromptLayer ──────────
            await self._create_promptlayer_workflow(
                task_description=task_description,
                output_text=output_text,
                model_name=model_name,
                pl_tool_calls=pl_tool_calls,
                pl_tool_results=pl_tool_results,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                price=price,
                start_time=start_time,
                end_time=end_time,
                tags=tags,
                run_output=run_output,
            )
        except Exception as e:
            _pl_logger.warning("Error in _log_to_promptlayer_unified: %s", e)

    def _log_to_promptlayer_background(
        self,
        task: "Task",
        output: Any,
        start_time: float,
        end_time: float,
    ) -> None:
        """Fire-and-forget: launches PromptLayer logging in a background thread."""
        def _run():
            try:
                asyncio.run(self._log_to_promptlayer_unified(task, output, start_time, end_time))
            except Exception as e:
                _pl_logger.warning("Background PromptLayer logging failed: %s", e)

        thread = threading.Thread(target=_run, daemon=True)
        thread.start()

    def _build_pl_tools(self) -> Optional[List[Dict[str, Any]]]:
        """Build PromptLayer-compatible tool definitions from the agent's registered tools."""
        try:
            tool_defs = list(self.tool_manager.get_tool_definitions())
        except Exception:
            return None
        if not tool_defs:
            return None
        result: List[Dict[str, Any]] = []
        for td in tool_defs:
            result.append({
                "type": "function",
                "function": {
                    "name": td.name,
                    "description": td.description or "",
                    "parameters": td.parameters_json_schema,
                },
            })
        return result

    async def _create_promptlayer_workflow(
        self,
        *,
        task_description: str,
        output_text: str,
        model_name: str,
        pl_tool_calls: Optional[List[Dict[str, Any]]],
        pl_tool_results: Optional[List[Dict[str, Any]]],
        input_tokens: int,
        output_tokens: int,
        price: float,
        start_time: float,
        end_time: float,
        tags: List[str],
        run_output: Any,
    ) -> None:
        """Register or update the agent as a workflow in PromptLayer."""
        if self.promptlayer is None:
            return
        try:
            import json as _json

            agent_name: str = self.name or "upsonic-agent"
            workflow_display_name: str = f"upsonic-{agent_name}"

            # ── Collect pipeline data ────────────────────────────────
            step_results: List[Any] = []
            if run_output is not None:
                step_results = getattr(run_output, "step_results", None) or []

            is_streaming: bool = getattr(run_output, "is_streaming", False) if run_output else False

            # ── Agent-level usage from self.usage ────────────────────
            agent_usage: Any = self.usage
            agent_metrics: Dict[str, Any] = {}
            if agent_usage is not None:
                agent_metrics = {
                    "total_requests": getattr(agent_usage, "requests", 0),
                    "total_tool_calls": getattr(agent_usage, "tool_calls", 0),
                    "total_input_tokens": getattr(agent_usage, "input_tokens", 0),
                    "total_output_tokens": getattr(agent_usage, "output_tokens", 0),
                    "total_tokens": (getattr(agent_usage, "input_tokens", 0)
                                     + getattr(agent_usage, "output_tokens", 0)),
                    "reasoning_tokens": getattr(agent_usage, "reasoning_tokens", 0),
                    "cache_write_tokens": getattr(agent_usage, "cache_write_tokens", 0),
                    "cache_read_tokens": getattr(agent_usage, "cache_read_tokens", 0),
                    "total_cost_usd": round(getattr(agent_usage, "cost", 0) or 0, 6),
                    "total_duration_seconds": round(getattr(agent_usage, "duration", 0) or 0, 3),
                    "model_execution_time": round(getattr(agent_usage, "model_execution_time", 0) or 0, 3),
                    "tool_execution_time": round(getattr(agent_usage, "tool_execution_time", 0) or 0, 3),
                }
                fw_time = getattr(agent_usage, "upsonic_execution_time", None)
                if fw_time is not None:
                    agent_metrics["framework_overhead_time"] = round(fw_time, 3)

            # ── Build nodes ──────────────────────────────────────────
            nodes: List[Dict[str, Any]] = []
            last_step_node: Optional[str] = None

            # 1. Input node
            nodes.append({
                "name": "input_node",
                "node_type": "VARIABLE",
                "configuration": {
                    "value": {"type": "string", "value": task_description},
                },
                "dependencies": [],
                "is_output_node": False,
            })
            last_step_node = "input_node"

            # 2. System prompt node (if present) — connects to first step like input_node
            if self.system_prompt:
                nodes.append({
                    "name": "system_prompt_node",
                    "node_type": "VARIABLE",
                    "configuration": {
                        "value": {"type": "string", "value": self.system_prompt},
                    },
                    "dependencies": [],
                    "is_output_node": False,
                })

            # 3. Pipeline step nodes — one per executed step
            #    The first step depends on both input_node and system_prompt_node (if present)
            first_step: bool = True
            for step in step_results:
                step_name_raw: str = getattr(step, "name", "unknown")
                step_status_val: Any = getattr(step, "status", None)
                step_status_str: str = step_status_val.value if hasattr(step_status_val, "value") else str(step_status_val or "UNKNOWN")
                step_time: float = getattr(step, "execution_time", 0.0)
                step_msg: str = getattr(step, "message", "") or ""

                step_info: Dict[str, Any] = {
                    "step": step_name_raw,
                    "status": step_status_str,
                    "execution_time": f"{step_time:.4f}s",
                }
                if step_msg:
                    step_info["message"] = step_msg

                node_name: str = f"step_{step_name_raw}"
                step_deps: List[str] = [last_step_node] if last_step_node else []
                if first_step and self.system_prompt:
                    step_deps.append("system_prompt_node")
                    first_step = False
                elif first_step:
                    first_step = False
                nodes.append({
                    "name": node_name,
                    "node_type": "VARIABLE",
                    "configuration": {
                        "value": {"type": "json", "value": step_info},
                    },
                    "dependencies": step_deps,
                    "is_output_node": False,
                })
                last_step_node = node_name

            # 4. Tool call nodes — separate branches off step_model_execution
            #    Each tool node depends ONLY on step_model_execution.
            #    Nothing depends on tool nodes — they are dead-end branches,
            #    isolated from the main pipeline.
            model_exec_node: str = "step_model_execution"
            model_exec_exists: bool = any(
                n["name"] == model_exec_node for n in nodes
            )

            if pl_tool_calls and model_exec_exists:
                for idx, tc in enumerate(pl_tool_calls):
                    func_info = tc.get("function", {})
                    tool_name: str = func_info.get("name", f"tool_{idx}")
                    node_name: str = f"tool_{tool_name}_{idx}"

                    tool_data: Dict[str, Any] = {
                        "tool_name": tool_name,
                        "arguments": _json.loads(func_info.get("arguments", "{}")),
                    }
                    if pl_tool_results and idx < len(pl_tool_results):
                        tool_data["result"] = pl_tool_results[idx].get("content", "")

                    nodes.append({
                        "name": node_name,
                        "node_type": "VARIABLE",
                        "configuration": {
                            "value": {"type": "json", "value": tool_data},
                        },
                        "dependencies": [model_exec_node],
                        "is_output_node": False,
                    })

            # 5. Output node — shows the agent's last result
            nodes.append({
                "name": "output",
                "node_type": "VARIABLE",
                "configuration": {
                    "value": {"type": "string", "value": output_text},
                },
                "dependencies": [last_step_node] if last_step_node else [],
                "is_output_node": True,
            })

            # ── Required input variables ─────────────────────────────
            required_input_variables: Dict[str, str] = {
                "task_description": "string",
            }
            if self.system_prompt:
                required_input_variables["system_prompt"] = "string"

            # ── Build commit message with agent-level usage ──────────
            run_latency: float = round(end_time - start_time, 3)
            total_cost: float = agent_metrics.get("total_cost_usd", price)
            total_reqs: int = agent_metrics.get("total_requests", 1)
            total_dur: float = agent_metrics.get("total_duration_seconds", run_latency)
            commit_message: str = (
                f"Agent: {agent_name} | model={model_name} | "
                f"{'streaming' if is_streaming else 'sync'} | "
                f"requests={total_reqs} | "
                f"tokens={agent_metrics.get('total_input_tokens', input_tokens)}"
                f"+{agent_metrics.get('total_output_tokens', output_tokens)} | "
                f"cost=${total_cost:.4f} | duration={total_dur:.3f}s | "
                f"latency={run_latency}s"
            )

            # ── Release labels ───────────────────────────────────────
            release_labels: List[str] = list(tags)

            # ── Create or patch the workflow ──────────────────────────
            existing_id: Optional[int] = self.promptlayer._created_workflows.get(workflow_display_name)

            # If not cached locally, check if it already exists in PromptLayer
            if existing_id is None:
                try:
                    wf_list = await self.promptlayer.alist_workflows(per_page=100)
                    for item in wf_list.get("items", []):
                        if item.get("name") == workflow_display_name:
                            existing_id = item.get("id")
                            self.promptlayer._created_workflows[workflow_display_name] = existing_id
                            break
                except Exception:
                    pass

            if existing_id is not None:
                # Agent already registered — patch to update with latest run data
                patch_nodes: Dict[str, Optional[Dict[str, Any]]] = {}
                for node in nodes:
                    patch_nodes[node["name"]] = {
                        k: v for k, v in node.items() if k != "name"
                    }
                await self.promptlayer.apatch_workflow(
                    workflow_id_or_name=existing_id,
                    nodes=patch_nodes,
                    required_input_variables=required_input_variables,
                    commit_message=commit_message,
                    release_labels=release_labels,
                )
            else:
                # First time — create the workflow
                wf_result = await self.promptlayer.acreate_workflow(
                    name=workflow_display_name,
                    nodes=nodes,
                    required_input_variables=required_input_variables,
                    commit_message=commit_message,
                    release_labels=release_labels,
                    folder_id=None,
                )
                wf_id = wf_result.get("workflow_id")
                if wf_id:
                    self.promptlayer._created_workflows[workflow_display_name] = wf_id
        except Exception as e:
            _pl_logger.warning("Error in _create_promptlayer_workflow: %s", e)

    def _apply_model_override(self, model: Optional[Union[str, "Model"]]) -> Optional["Model"]:
        """Apply a per-call model override and return the original model for restoration.

        Re-wraps the new model with InstrumentedModel if instrumentation is active,
        so OTel LLM-level spans are preserved even when the model is overridden per-call.

        Args:
            model: Optional model override (string or Model instance)

        Returns:
            The original model instance if an override was applied (caller must restore
            it after the run), or None if no override was needed.
        """
        if model:
            original_model: "Model" = self.model
            from upsonic.models import infer_model
            self.model = infer_model(model)
            if self._instrument_settings is not None:
                from upsonic.models.instrumented import instrument_model
                self.model = instrument_model(self.model, self._instrument_settings)
            return original_model
        return None
    
    
    def get_run_id(self) -> Optional[str]:
        """
        Get the current run ID.
        
        Returns:
            str: The current run ID, or None if no run is active.
        """
        return self.run_id
    
    def cancel_run(self, run_id: Optional[str] = None) -> bool:
        """
        Cancel a run by its ID.
        
        If no run_id is provided, cancels the current run.
        
        Args:
            run_id: The ID of the run to cancel. If None, cancels the current run.
            
        Returns:
            bool: True if the run was found and cancelled, False otherwise.
        """
        target_run_id = run_id or self.run_id
        if not target_run_id:
            return False
        return cancel_run_func(target_run_id)
    
    def _validate_tools_with_policy_pre(
        self, 
        context_description: str = "Tool Validation",
        task: Optional["Task"] = None,
        registered_tools_dicts: Optional[List[Dict[str, Any]]] = None
    ) -> None:
        """
        Validate all currently registered tools with tool_policy_pre before use.
        
        Combines tool definitions from both the agent's ToolManager and the task's
        ToolManager (if provided) for comprehensive validation.
        
        Args:
            context_description: Description of where this validation is being called from
            task: Optional task whose ToolManager should also be validated
            registered_tools_dicts: List of registered tools dictionaries to check when removing tools.
                                   If None, defaults to [self.registered_agent_tools]
            
        Raises:
            DisallowedOperation: If any tool is blocked by the safety policy
        """
        if not hasattr(self, 'tool_policy_pre_manager') or not self.tool_policy_pre_manager.has_policies():
            return
        
        if registered_tools_dicts is None:
            registered_tools_dicts = [self.registered_agent_tools]
        
        import asyncio
        tool_definitions = list(self.tool_manager.get_tool_definitions())
        if task is not None and task.tool_manager is not None:
            tool_definitions.extend(task.tool_manager.get_tool_definitions())
        
        for tool_def in tool_definitions:
            if tool_def.name == 'plan_and_execute':
                continue
                
            tool_info = {
                "name": tool_def.name,
                "description": tool_def.description or "",
                "parameters": tool_def.parameters_json_schema or {},
                "metadata": tool_def.metadata or {}
            }
            
            validation_result = _run_in_bg_loop(
                self.tool_policy_pre_manager.execute_tool_validation_async(
                    tool_info=tool_info,
                    check_type=f"Pre-Execution Tool Validation ({context_description})"
                )
            )
            
            if validation_result.should_block():
                if validation_result.disallowed_exception:
                    raise validation_result.disallowed_exception
                
                if self.debug:
                    from upsonic.utils.printing import warning_log
                    warning_log(
                        f"Tool '{tool_def.name}' blocked by safety policy: {validation_result.get_final_message()}",
                        "Tool Safety"
                    )
                
                for registered_tools_dict in registered_tools_dicts:
                    if tool_def.name in registered_tools_dict:
                        target_manager = self.tool_manager
                        if task is not None and task.tool_manager is not None and tool_def.name in (task.tool_manager.wrapped_tools or {}):
                            target_manager = task.tool_manager
                        target_manager.remove_tools(
                            tools=[tool_def.name],
                            registered_tools=registered_tools_dict
                        )
                        del registered_tools_dict[tool_def.name]
    
    def _register_agent_tools(self) -> None:
        """
        Register agent-level tools with the ToolManager.
        
        This is called in __init__ to ensure agent tools are registered immediately.
        Automatically includes canvas tools if canvas is provided.
        """
        # Prepare tools list starting with user-provided tools
        final_tools = list(self.tools) if self.tools else []
        
        if self.canvas:
            canvas_functions = self.canvas.functions()
            for canvas_func in canvas_functions:
                if canvas_func not in final_tools:
                    final_tools.append(canvas_func)
            self.tools = final_tools
        
        if not final_tools:
            self.registered_agent_tools = {}
            self.agent_builtin_tools = []
            return
        
        # Add thinking tool if enabled
        if self.enable_thinking_tool:
            from upsonic.tools.orchestration import plan_and_execute
            if plan_and_execute not in final_tools:
                final_tools.append(plan_and_execute)
        
        # Separate builtin tools from regular tools
        from upsonic.tools.builtin_tools import AbstractBuiltinTool
        builtin_tools = []
        regular_tools = []
        
        for tool in final_tools:
            if tool is not None and isinstance(tool, AbstractBuiltinTool):
                builtin_tools.append(tool)
            else:
                regular_tools.append(tool)
        
        # Handle builtin tools separately - they don't need ToolManager/ToolProcessor
        self.agent_builtin_tools = builtin_tools
        
        # Register only regular tools with ToolManager
        if regular_tools:
            self.registered_agent_tools = self.tool_manager.register_tools(
                tools=regular_tools,
                task=None,  # Agent tools not task-specific
                agent_instance=self
            )
        else:
            self.registered_agent_tools = {}
        
        # PRE-EXECUTION TOOL VALIDATION
        # Validate all registered agent tools with tool_policy_pre
        self._validate_tools_with_policy_pre(
            context_description="Agent Tool Registration",
            registered_tools_dicts=[self.registered_agent_tools]
        )
    
    def add_tools(self, tools: Union[Any, List[Any]]) -> None:
        """
        Dynamically add tools to the agent and register them.
        
        This method:
        1. Separates builtin tools from regular tools
        2. For builtin tools: Updates self.tools and self.agent_builtin_tools directly
        3. For regular tools: Calls ToolManager to register them
        4. Updates self.registered_agent_tools with wrapped tools
        5. Validates tools with tool_policy_pre if configured
        
        Args:
            tools: A single tool or list of tools to add
            
        Raises:
            DisallowedOperation: If any tool is blocked by the safety policy
        """
        if not isinstance(tools, list):
            tools = [tools]
        
        # Prepare tools with plan_and_execute if needed
        tools_to_add = list(tools)
        
        # Add thinking tool if enabled and not already in the list
        if self.enable_thinking_tool:
            from upsonic.tools.orchestration import plan_and_execute
            if plan_and_execute not in tools_to_add and plan_and_execute not in self.tools:
                tools_to_add.append(plan_and_execute)
        
        # Separate builtin tools from regular tools
        from upsonic.tools.builtin_tools import AbstractBuiltinTool
        builtin_tools = []
        regular_tools = []
        
        for tool in tools_to_add:
            if tool is not None and isinstance(tool, AbstractBuiltinTool):
                builtin_tools.append(tool)
            else:
                regular_tools.append(tool)
        
        # Handle builtin tools separately - they don't need ToolManager/ToolProcessor
        if builtin_tools:
            if not hasattr(self, 'agent_builtin_tools'):
                self.agent_builtin_tools = []
            
            # Merge builtin tools (avoid duplicates based on unique_id)
            existing_ids = {tool.unique_id for tool in self.agent_builtin_tools}
            for tool in builtin_tools:
                if tool.unique_id not in existing_ids:
                    self.agent_builtin_tools.append(tool)
                    existing_ids.add(tool.unique_id)
        
        # Handle regular tools through ToolManager
        if regular_tools:
            # Call ToolManager to register new tools (filters already registered ones)
            newly_registered = self.tool_manager.register_tools(
                tools=regular_tools,
                task=None,  # Agent tools are not task-specific
                agent_instance=self
            )
            
            # Update self.registered_agent_tools with newly registered tools
            self.registered_agent_tools.update(newly_registered)
        
        # Update self.tools - add original tool objects (not plan_and_execute if auto-added)
        for tool in tools:
            if tool not in self.tools:
                self.tools.append(tool)
        
        # PRE-EXECUTION TOOL VALIDATION
        # Validate newly added tools with tool_policy_pre
        self._validate_tools_with_policy_pre(
            context_description="Dynamic Tool Addition (add_tools)",
            registered_tools_dicts=[self.registered_agent_tools]
        )
    
    def remove_tools(self, tools: Union[str, List[str], Any, List[Any]]) -> None:
        """
        Remove tools from the agent.
        
        Supports removing:
        - Tool names (strings)
        - Function objects
        - Agent objects
        - MCP handlers (and all their tools)
        - Class instances (ToolKit or regular classes, and all their tools)
        - Builtin tools (AbstractBuiltinTool instances)
        
        Args:
            tools: Single tool or list of tools to remove (any type)
        """
        if not isinstance(tools, list):
            tools = [tools]
        
        # Separate builtin tools from regular tools
        from upsonic.tools.builtin_tools import AbstractBuiltinTool
        builtin_tools_to_remove = []
        regular_tools_to_remove = []
        
        for tool in tools:
            if tool is not None and isinstance(tool, AbstractBuiltinTool):
                builtin_tools_to_remove.append(tool)
            else:
                regular_tools_to_remove.append(tool)
        
        # Handle regular tools through ToolManager
        removed_tool_names = []
        removed_objects = []
        
        if regular_tools_to_remove:
            # Call ToolManager to handle all removal logic for regular tools
            removed_tool_names, removed_objects = self.tool_manager.remove_tools(
                tools=regular_tools_to_remove,
                registered_tools=self.registered_agent_tools
            )
            
            # Update self.registered_agent_tools - remove the tool names
            for tool_name in removed_tool_names:
                if tool_name in self.registered_agent_tools:
                    del self.registered_agent_tools[tool_name]
        
        # Handle builtin tools separately - they don't use ToolManager/ToolProcessor
        if builtin_tools_to_remove and hasattr(self, 'agent_builtin_tools'):
            # Remove from agent_builtin_tools by unique_id
            builtin_ids_to_remove = {tool.unique_id for tool in builtin_tools_to_remove}
            self.agent_builtin_tools = [
                tool for tool in self.agent_builtin_tools 
                if tool.unique_id not in builtin_ids_to_remove
            ]
            # Add to removed_objects for self.tools cleanup
            removed_objects.extend(builtin_tools_to_remove)
        
        # Update self.tools - remove all removed objects (regular + builtin)
        if removed_objects:
            self.tools = [t for t in self.tools if t not in removed_objects]
    
    def get_tool_defs(self) -> List["ToolDefinition"]:
        """
        Get the tool definitions for all currently registered tools.
        
        Returns:
            List[ToolDefinition]: List of tool definitions from the ToolManager
        """
        return self.tool_manager.get_tool_definitions()
    
    def _setup_task_tools(self, task: "Task") -> None:
        """Setup tools with a dedicated ToolManager on the task (task tools only)."""
        self._tool_limit_reached = False
        
        from upsonic.tools import ToolMetrics
        self._tool_metrics = ToolMetrics(
            tool_call_count=self._tool_call_count,
            tool_call_limit=self.tool_call_limit
        )
        
        task_tools: list = task.tools if task.tools else []
        
        is_thinking_enabled: bool = self.enable_thinking_tool
        if task.enable_thinking_tool is not None:
            is_thinking_enabled = task.enable_thinking_tool
        
        is_reasoning_enabled: bool = self.enable_reasoning_tool
        if task.enable_reasoning_tool is not None:
            is_reasoning_enabled = task.enable_reasoning_tool

        if is_reasoning_enabled and not is_thinking_enabled:
            raise ValueError("Configuration error: 'enable_reasoning_tool' cannot be True if 'enable_thinking_tool' is False.")

        from upsonic.tools.orchestration import plan_and_execute
        
        tools_to_register: list = list(task_tools) if task_tools else []

        # Register task-level skill tools with prefix to avoid name collision
        # with agent-level skill tools (both generate get_skill_instructions, etc.)
        if hasattr(task, 'skills') and task.skills is not None:
            tools_to_register.extend(task.skills.get_tools(prefix="task_"))

        if is_thinking_enabled and plan_and_execute not in tools_to_register:
            tools_to_register.append(plan_and_execute)
        
        from upsonic.tools import ToolManager
        task_tool_manager: ToolManager = task._ensure_tool_manager()
        
        if not tools_to_register:
            return

        agent_for_this_run = copy.copy(self)
        agent_for_this_run.enable_thinking_tool = is_thinking_enabled
        agent_for_this_run.enable_reasoning_tool = is_reasoning_enabled

        from upsonic.tools.builtin_tools import AbstractBuiltinTool
        builtin_tools: list = []
        regular_tools: list = []
        
        for tool in tools_to_register:
            if tool is not None and isinstance(tool, AbstractBuiltinTool):
                builtin_tools.append(tool)
            else:
                regular_tools.append(tool)
        
        task.task_builtin_tools = builtin_tools
        
        if regular_tools:
            newly_registered = task_tool_manager.register_tools(
                tools=regular_tools,
                task=task,
                agent_instance=agent_for_this_run
            )
        else:
            newly_registered = {}
        
        task.registered_task_tools.update(newly_registered)
        
        self._validate_tools_with_policy_pre(
            context_description="Task Tool Setup",
            task=task,
            registered_tools_dicts=[self.registered_agent_tools, task.registered_task_tools]
        )
    
    def _get_combined_tool_definitions(self) -> List["ToolDefinition"]:
        """Combine tool definitions from the agent's ToolManager and the current task's ToolManager."""
        tool_definitions: list = list(self.tool_manager.get_tool_definitions())
        current_task = getattr(self, 'current_task', None)
        if current_task is not None and current_task.tool_manager is not None:
            tool_definitions.extend(current_task.tool_manager.get_tool_definitions())
        return tool_definitions

    def _resolve_tool_manager(self, tool_name: str) -> "ToolManager":
        """Determine which ToolManager owns the given tool by name.
        
        Checks the agent's ToolManager first, then the current task's ToolManager.
        
        Args:
            tool_name: Name of the tool to look up
            
        Returns:
            The ToolManager that owns the tool
            
        Raises:
            ValueError: If the tool is not found in any manager
        """
        if tool_name in self.tool_manager.wrapped_tools:
            return self.tool_manager
        current_task = getattr(self, 'current_task', None)
        if current_task is not None and current_task.tool_manager is not None:
            if tool_name in current_task.tool_manager.wrapped_tools:
                return current_task.tool_manager
        raise ValueError(f"Tool '{tool_name}' not found in any ToolManager")

    def get_skill_metrics(self) -> Dict[str, Any]:
        """Return skill metrics from agent-level skills."""
        if self.skills is not None:
            return {k: v.to_dict() for k, v in self.skills.get_metrics().items()}
        return {}

    async def _build_model_request_with_input(
        self, 
        task: "Task", 
        memory_handler: Optional["MemoryManager"], 
        current_input: Any, 
        temporary_message_history: List["ModelRequest"],
        state: Optional["State"] = None,
    ) -> tuple[List["ModelRequest"], Optional["ModelResponse"]]:
        """Build model request with custom input and message history for guardrail retries.

        Returns:
            A tuple of (messages, context_full_response).
            context_full_response is None when the context fits within the
            model's window, or a ModelResponse with a fixed message when
            the context is full after all reduction strategies.
        """
        from upsonic.agent.context_managers import SystemPromptManager, ContextManager
        from upsonic.messages import SystemPromptPart, UserPromptPart, ModelRequest
        
        messages: List["ModelRequest"] = list(temporary_message_history)
        
        system_prompt_manager = SystemPromptManager(self, task)
        context_manager = ContextManager(self, task, state)
        
        async with system_prompt_manager.manage_system_prompt(memory_handler) as sp_handler, \
                   context_manager.manage_context(memory_handler) as _ctx_handler:
            
            user_part = UserPromptPart(content=current_input)
            
            parts = []
            
            if not messages:
                system_prompt = sp_handler.get_system_prompt()
                if system_prompt:
                    self._last_built_system_prompt = system_prompt
                    system_part = SystemPromptPart(content=system_prompt)
                    parts.append(system_part)
            
            parts.append(user_part)
            
            current_request = ModelRequest(parts=parts)
            messages.append(current_request)

        # Apply context management middleware
        context_full_response: Optional["ModelResponse"] = None
        if self.context_management and self._context_management_middleware:
            managed_msgs, ctx_full = await self._context_management_middleware.apply(messages)
            messages = managed_msgs
            # Propagate summarization usage to the parent run context
            self._propagate_context_management_usage()
            if ctx_full:
                context_full_response = self._context_management_middleware._build_context_full_response(
                    model_name=self.model.model_name
                )

        return messages, context_full_response
    
    def _build_model_request_parameters(self, task: "Task") -> "ModelRequestParameters":
        """Build model request parameters including tools and structured output."""
        from pydantic import BaseModel
        from upsonic.output import OutputObjectDefinition
        from upsonic.models import ModelRequestParameters
        
        if hasattr(self, '_tool_limit_reached') and self._tool_limit_reached:
            tool_definitions = []
        elif self.tool_call_limit and self._tool_call_count >= self.tool_call_limit:
            tool_definitions = []
            self._tool_limit_reached = True
        else:
            tool_definitions = list(self.tool_manager.get_tool_definitions())
            if task is not None and task.tool_manager is not None:
                tool_definitions.extend(task.tool_manager.get_tool_definitions())
        
        # Combine agent-level and task-level builtin tools
        agent_builtin_tools = getattr(self, 'agent_builtin_tools', [])
        task_builtin_tools = getattr(task, 'task_builtin_tools', [])

        # Merge builtin tools, avoiding duplicates based on unique_id
        builtin_tools_dict = {}
        for tool in agent_builtin_tools:
            builtin_tools_dict[tool.unique_id] = tool
        for tool in task_builtin_tools:
            builtin_tools_dict[tool.unique_id] = tool
        builtin_tools = list(builtin_tools_dict.values())
        
        output_mode = 'text'
        output_object = None
        output_tools = []
        allow_text_output = True
        
        if task.response_format and task.response_format != str and task.response_format is not str:
            if isinstance(task.response_format, type) and issubclass(task.response_format, BaseModel):
                output_mode = 'auto'
                allow_text_output = False
                
                schema = task.response_format.model_json_schema()
                output_object = OutputObjectDefinition(
                    json_schema=schema,
                    name=task.response_format.__name__,
                    description=task.response_format.__doc__,
                    strict=True
                )
                
                # Create output tool for tool-based structured output
                output_tools = self._build_output_tools(task.response_format, schema)
        
        return ModelRequestParameters(
            function_tools=tool_definitions,
            builtin_tools=builtin_tools,
            output_mode=output_mode,
            output_object=output_object,
            output_tools=output_tools,
            allow_text_output=allow_text_output
        )
    
    def _build_output_tools(self, response_format: type, schema: dict) -> list:
        """Build output tools for tool-based structured output.
        
        Creates a ToolDefinition that the model can use to return structured data
        when native JSON schema output is not supported.
        
        Args:
            response_format: The Pydantic model class for the response
            schema: The JSON schema for the response format
            
        Returns:
            List containing a single ToolDefinition for structured output
        """
        from upsonic.tools import ToolDefinition
        
        return [ToolDefinition(
            name=DEFAULT_OUTPUT_TOOL_NAME,
            parameters_json_schema=schema,
            description=response_format.__doc__ or f"Return the final result as a {response_format.__name__}",
            kind='output',
            strict=True
        )]
    
    def _handle_output_tool_call(
        self,
        tool_call: "ToolCallPart",
        response_format: type | None,
        use_output_tool: bool,
    ) -> "ToolReturnPart | None":
        """Handle the structured-output tool (final_result) if applicable. Returns ToolReturnPart or None."""
        from upsonic.messages import ToolReturnPart
        
        if not use_output_tool or response_format is None or tool_call.tool_name != DEFAULT_OUTPUT_TOOL_NAME:
            return None
        try:
            validated = response_format.model_validate(tool_call.args_as_dict())
        except Exception as e:
            return ToolReturnPart(
                tool_name=tool_call.tool_name,
                content=f"Validation error: {str(e)}",
                tool_call_id=tool_call.tool_call_id,
                timestamp=now_utc(),
            )
        return ToolReturnPart(
            tool_name=tool_call.tool_name,
            content={"result": validated.model_dump()},
            tool_call_id=tool_call.tool_call_id,
            timestamp=now_utc(),
        )
    
    async def _execute_tool_calls(self, tool_calls: List["ToolCallPart"]) -> List["ToolReturnPart"]:
        """
        Execute tool calls and return results.
        
        Handles both sequential and parallel execution based on tool configuration.
        Tools marked as sequential will be executed one at a time.
        Other tools can be executed in parallel if multiple are called.
        """
        from upsonic.messages import ToolReturnPart
        
        if not tool_calls:
            return []
        
        # Check for cancellation before executing tools
        if self.run_id:
            raise_if_cancelled(self.run_id)
        
        if self.tool_call_limit and self._tool_call_count >= self.tool_call_limit:
            error_results = []
            for tool_call in tool_calls:
                error_results.append(ToolReturnPart(
                    tool_name=tool_call.tool_name,
                    content=f"Tool call limit of {self.tool_call_limit} reached. Cannot execute more tools.",
                    tool_call_id=tool_call.tool_call_id
                ))
            self._tool_limit_reached = True
            return error_results
        
        current_task = getattr(self, "current_task", None)
        response_format = getattr(current_task, "response_format", None) if current_task else None
        try:
            from pydantic import BaseModel
            use_output_tool = (
                isinstance(response_format, type)
                and issubclass(response_format, BaseModel)
                and response_format is not str
                and response_format != str
            )
        except TypeError:
            use_output_tool = False
        
        tool_defs = {td.name: td for td in self._get_combined_tool_definitions()}
        
        sequential_calls = []
        parallel_calls = []
        
        for tool_call in tool_calls:
            tool_def = tool_defs.get(tool_call.tool_name)
            if tool_def and tool_def.sequential:
                sequential_calls.append(tool_call)
            else:
                parallel_calls.append(tool_call)
        
        results = []
        
        for tool_call in sequential_calls:
            output_part = self._handle_output_tool_call(tool_call, response_format, use_output_tool)
            if output_part is not None:
                results.append(output_part)
                continue
            # POST-EXECUTION TOOL CALL VALIDATION
            if hasattr(self, 'tool_policy_post_manager') and self.tool_policy_post_manager.has_policies():
                tool_def = tool_defs.get(tool_call.tool_name)
                tool_call_info = {
                    "name": tool_call.tool_name,
                    "description": tool_def.description if tool_def else "",
                    "parameters": tool_def.parameters_json_schema if tool_def else {},
                    "arguments": tool_call.args_as_dict(),
                    "call_id": tool_call.tool_call_id
                }
                
                validation_result = await self.tool_policy_post_manager.execute_tool_call_validation_async(
                    tool_call_info=tool_call_info,
                    check_type="Post-Execution Tool Call Validation"
                )
                
                # Drain and aggregate sub-agent usage from tool policy LLM calls
                tool_policy_usage = self.tool_policy_post_manager.drain_accumulated_usage()
                if tool_policy_usage is not None and hasattr(self, '_agent_run_output') and self._agent_run_output:
                    usage = self._agent_run_output._ensure_usage()
                    usage.incr(tool_policy_usage)
                
                if validation_result.should_block():
                    # Handle blocking based on action type
                    # If DisallowedOperation was raised by a RAISE action policy, re-raise it
                    if validation_result.disallowed_exception:
                        raise validation_result.disallowed_exception
                    
                    # Otherwise it's a BLOCK action - return error message without raising
                    results.append(ToolReturnPart(
                        tool_name=tool_call.tool_name,
                        content=validation_result.get_final_message(),
                        tool_call_id=tool_call.tool_call_id,
                        timestamp=now_utc()
                    ))
                    continue  # Skip execution
            
            import time
            tool_start_time = time.time()
            _tool_args = tool_call.args_as_dict()
            with self._otel.tool_span(tool_call.tool_name, tool_call.tool_call_id, tool_args=_tool_args) as otel_tool_span:
                try:
                    target_manager = self._resolve_tool_manager(tool_call.tool_name)
                    result = await target_manager.execute_tool(
                        tool_name=tool_call.tool_name,
                        args=_tool_args,
                        metrics=self._tool_metrics,
                        tool_call_id=tool_call.tool_call_id
                    )
                    tool_execution_time = time.time() - tool_start_time

                    if hasattr(self, '_agent_run_output') and self._agent_run_output:
                        self._agent_run_output.add_tool_execution_time(tool_execution_time)

                    self._tool_call_count += 1
                    if hasattr(self, '_tool_metrics') and self._tool_metrics:
                        self._tool_metrics.tool_call_count = self._tool_call_count
                    if hasattr(self, '_agent_run_output') and self._agent_run_output is not None:
                        self._agent_run_output.tool_call_count = self._tool_call_count
                        self._agent_run_output.increment_tool_calls(1)

                    tool_return = ToolReturnPart(
                        tool_name=result.tool_name,
                        content=result.content,
                        tool_call_id=result.tool_call_id,
                        timestamp=now_utc()
                    )
                    results.append(tool_return)

                    self._otel.set_tool_result(otel_tool_span, tool_execution_time, success=True, output=result.content)
                    
                    if hasattr(self, '_agent_run_output') and self._agent_run_output:
                        from upsonic.run.tools.tools import ToolExecution
                        tool_exec = ToolExecution(
                            tool_call_id=tool_call.tool_call_id,
                            tool_name=tool_call.tool_name,
                            tool_args=tool_call.args_as_dict(),
                            result=str(result.content) if result.content else None,
                        )
                        if self._agent_run_output.tools is None:
                            self._agent_run_output.tools = []
                        self._agent_run_output.tools.append(tool_exec)
                    
                    if hasattr(self, '_agent_run_output') and self._agent_run_output:
                        self._drain_agent_tool_usage(tool_call.tool_name)
                    
                    if self.debug and self.debug_level >= 2:
                        from upsonic.utils.printing import debug_log_level2
                        tool_def = tool_defs.get(tool_call.tool_name)
                        debug_log_level2(
                            f"Tool executed: {tool_call.tool_name}",
                            "Agent",
                            debug=self.debug,
                            debug_level=self.debug_level,
                            tool_name=tool_call.tool_name,
                            tool_description=tool_def.description if tool_def else "Unknown",
                            tool_parameters=tool_call.args_as_dict(),
                            tool_result=str(result.content)[:1000] if result.content else None,
                            tool_execution_time=tool_execution_time,
                            tool_call_id=tool_call.tool_call_id,
                            total_tool_calls=self._tool_call_count,
                            tool_call_limit=self.tool_call_limit,
                            tool_sequential=tool_def.sequential if tool_def else False
                        )
                    
                except (ExternalExecutionPause, ConfirmationPause, UserInputPause) as e:
                    raise e
                except Exception as e:
                    tool_execution_time = time.time() - tool_start_time
                    self._otel.set_tool_result(otel_tool_span, tool_execution_time, success=False, error=e)

                    if hasattr(self, '_agent_run_output') and self._agent_run_output:
                        self._agent_run_output.add_tool_execution_time(tool_execution_time)

                    error_return = ToolReturnPart(
                        tool_name=tool_call.tool_name,
                        content=f"Error executing tool: {str(e)}",
                        tool_call_id=tool_call.tool_call_id,
                        timestamp=now_utc()
                    )
                    results.append(error_return)
                    
                    if self.debug and self.debug_level >= 2:
                        from upsonic.utils.printing import debug_log_level2
                        import traceback
                        error_traceback = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
                        debug_log_level2(
                            f"Tool execution error: {tool_call.tool_name}",
                            "Agent",
                            debug=self.debug,
                            debug_level=self.debug_level,
                            tool_name=tool_call.tool_name,
                            tool_parameters=tool_call.args_as_dict(),
                            error_type=type(e).__name__,
                            error_message=str(e),
                            error_traceback=error_traceback[-1500:],
                            tool_call_id=tool_call.tool_call_id
                        )
        
        if parallel_calls:
            async def execute_single_tool(tool_call: "ToolCallPart") -> "ToolReturnPart":
                """Execute a single tool call and return the result."""
                output_part = self._handle_output_tool_call(tool_call, response_format, use_output_tool)
                if output_part is not None:
                    return output_part
                # POST-EXECUTION TOOL CALL VALIDATION (for parallel execution)
                if hasattr(self, 'tool_policy_post_manager') and self.tool_policy_post_manager.has_policies():
                    tool_def = tool_defs.get(tool_call.tool_name)
                    tool_call_info = {
                        "name": tool_call.tool_name,
                        "description": tool_def.description if tool_def else "",
                        "parameters": tool_def.parameters_json_schema if tool_def else {},
                        "arguments": tool_call.args_as_dict(),
                        "call_id": tool_call.tool_call_id
                    }
                    
                    validation_result = await self.tool_policy_post_manager.execute_tool_call_validation_async(
                        tool_call_info=tool_call_info,
                        check_type="Post-Execution Tool Call Validation"
                    )
                    
                    # Drain and aggregate sub-agent usage from tool policy LLM calls (parallel path)
                    tool_policy_usage = self.tool_policy_post_manager.drain_accumulated_usage()
                    if tool_policy_usage is not None and hasattr(self, '_agent_run_output') and self._agent_run_output:
                        usage = self._agent_run_output._ensure_usage()
                        usage.incr(tool_policy_usage)
                    
                    if validation_result.should_block():
                        # Handle blocking based on action type
                        # If DisallowedOperation was raised by a RAISE action policy, re-raise it
                        if validation_result.disallowed_exception:
                            raise validation_result.disallowed_exception
                        
                        # Otherwise it's a BLOCK action - return error message without raising
                        return ToolReturnPart(
                            tool_name=tool_call.tool_name,
                            content=validation_result.get_final_message(),
                            tool_call_id=tool_call.tool_call_id,
                            timestamp=now_utc()
                        )
                
                import time as _time
                _tool_start = _time.time()
                _tool_args2 = tool_call.args_as_dict()
                with self._otel.tool_span(tool_call.tool_name, tool_call.tool_call_id, tool_args=_tool_args2) as otel_tool_span:
                    try:
                        target_manager = self._resolve_tool_manager(tool_call.tool_name)
                        result = await target_manager.execute_tool(
                            tool_name=tool_call.tool_name,
                            args=_tool_args2,
                            metrics=self._tool_metrics,
                            tool_call_id=tool_call.tool_call_id
                        )
                        _tool_elapsed = _time.time() - _tool_start

                        if hasattr(self, '_agent_run_output') and self._agent_run_output:
                            from upsonic.run.tools.tools import ToolExecution
                            tool_exec = ToolExecution(
                                tool_call_id=tool_call.tool_call_id,
                                tool_name=tool_call.tool_name,
                                tool_args=_tool_args2,
                                result=str(result.content) if result.content else None,
                            )
                            if self._agent_run_output.tools is None:
                                self._agent_run_output.tools = []
                            self._agent_run_output.tools.append(tool_exec)

                        if hasattr(self, '_agent_run_output') and self._agent_run_output:
                            self._drain_agent_tool_usage(tool_call.tool_name)

                        self._otel.set_tool_result(otel_tool_span, _tool_elapsed, success=True, output=result.content)
                        
                        return ToolReturnPart(
                            tool_name=result.tool_name,
                            content=result.content,
                            tool_call_id=result.tool_call_id,
                            timestamp=now_utc()
                        )
                    
                    except (ExternalExecutionPause, ConfirmationPause, UserInputPause):
                        raise
                    except Exception as e:
                        _tool_elapsed = _time.time() - _tool_start
                        self._otel.set_tool_result(otel_tool_span, _tool_elapsed, success=False, error=e)
                        return ToolReturnPart(
                            tool_name=tool_call.tool_name,
                            content=f"Error executing tool: {str(e)}",
                            tool_call_id=tool_call.tool_call_id,
                            timestamp=now_utc()
                        )
            
            import time as _time_mod
            _parallel_batch_start: float = _time_mod.time()
            parallel_results = await asyncio.gather(
                *[execute_single_tool(tc) for tc in parallel_calls],
                return_exceptions=True
            )
            _parallel_batch_elapsed: float = _time_mod.time() - _parallel_batch_start
            if hasattr(self, '_agent_run_output') and self._agent_run_output:
                self._agent_run_output.add_tool_execution_time(_parallel_batch_elapsed)
            
            # Separate successful results from HITL pauses
            external_pauses: List[ExternalExecutionPause] = []
            confirmation_pauses: List[ConfirmationPause] = []
            user_input_pauses: List[UserInputPause] = []
            successful_results: List["ToolReturnPart"] = []
            other_errors: List[Exception] = []
            
            for tc, result in zip(parallel_calls, parallel_results):
                if isinstance(result, ConfirmationPause):
                    confirmation_pauses.append(result)
                elif isinstance(result, UserInputPause):
                    user_input_pauses.append(result)
                elif isinstance(result, ExternalExecutionPause):
                    external_pauses.append(result)
                elif isinstance(result, Exception):
                    other_errors.append(result)
                    successful_results.append(ToolReturnPart(
                        tool_name=tc.tool_name,
                        content=f"Error executing tool: {str(result)}",
                        tool_call_id=tc.tool_call_id,
                        timestamp=now_utc()
                    ))
                else:
                    successful_results.append(result)
            
            if confirmation_pauses:
                all_calls = []
                for pause in confirmation_pauses:
                    if pause.paused_calls:
                        all_calls.extend(pause.paused_calls)
                raise ConfirmationPause(paused_calls=all_calls)

            if user_input_pauses:
                all_calls = []
                all_schema = []
                for pause in user_input_pauses:
                    if pause.paused_calls:
                        all_calls.extend(pause.paused_calls)
                    if pause.user_input_schema:
                        all_schema.extend(pause.user_input_schema)
                raise UserInputPause(paused_calls=all_calls, user_input_schema=all_schema)

            if external_pauses:
                all_paused_calls = []
                for pause in external_pauses:
                    if pause.paused_calls:
                        all_paused_calls.extend(pause.paused_calls)
                raise ExternalExecutionPause(paused_calls=all_paused_calls)
            
            self._tool_call_count += len(parallel_calls)
            if hasattr(self, '_tool_metrics') and self._tool_metrics:
                self._tool_metrics.tool_call_count = self._tool_call_count
            if hasattr(self, '_agent_run_output') and self._agent_run_output is not None:
                self._agent_run_output.tool_call_count = self._tool_call_count
                self._agent_run_output.increment_tool_calls(len(parallel_calls))
            
            results.extend(successful_results)
        
        return results
    
    async def _handle_model_response(
        self, 
        response: "ModelResponse", 
        messages: List["ModelRequest"]
    ) -> "ModelResponse":
        """Handle model response including tool calls."""
        from upsonic.messages import ToolCallPart, TextPart, UserPromptPart, ModelRequest, ModelResponse
        from upsonic._utils import now_utc
        
        if hasattr(self, '_tool_limit_reached') and self._tool_limit_reached:
            return response
        
        # Handle culture repeat logic
        if self._culture_manager and self._culture_manager.enabled:
            culture = self._culture_manager.culture
            if culture and culture.repeat:
                if self._culture_manager.should_repeat():
                    # Ensure culture is prepared
                    if not self._culture_manager.prepared:
                        await self._culture_manager.aprepare()
                        
                        # Drain culture extraction LLM usage into run output
                        if hasattr(self, '_agent_run_output') and self._agent_run_output:
                            culture_usage = self._culture_manager.drain_accumulated_usage()
                            if culture_usage:
                                self._agent_run_output.usage.incr(culture_usage)
                    
                    culture_formatted = self._culture_manager.format_for_system_prompt()
                    if culture_formatted:
                        # Create mock ModelRequest with culture guidelines
                        culture_system_part = UserPromptPart(content=culture_formatted)
                        culture_request = ModelRequest(parts=[culture_system_part])
                        
                        # Create mock ModelResponse acknowledging culture
                        culture_response = ModelResponse(
                            parts=[TextPart(content="Culture guidelines acknowledged.")],
                            model_name=response.model_name if response else None,
                            timestamp=now_utc(),
                            usage=response.usage if response else None,
                            provider_name=response.provider_name if response else None,
                        )
                        
                        # Insert culture into messages before the current response
                        messages.append(culture_request)
                        messages.append(culture_response)
        
        tool_calls = [
            part for part in response.parts 
            if isinstance(part, ToolCallPart)
        ]
        
        # Filter out output tool calls - these are used for structured output
        # and should not be executed as regular tools
        output_tool_names = {DEFAULT_OUTPUT_TOOL_NAME}
        regular_tool_calls = [tc for tc in tool_calls if tc.tool_name not in output_tool_names]
        
        # If all tool calls are output tools, return response directly (structured output)
        if tool_calls and not regular_tool_calls:
            return response
        
        # Handle max_tokens truncation: when finish_reason is 'length', tool call
        # arguments may be incomplete/truncated. Instead of executing broken calls,
        # inform the model and let it retry with smaller output.
        if response.finish_reason == 'length' and regular_tool_calls:
            from upsonic.messages import ToolReturnPart
            truncation_results: List[ToolReturnPart] = []
            for tc in regular_tool_calls:
                truncation_results.append(
                    ToolReturnPart(
                        tool_name=tc.tool_name,
                        content=(
                            "[ERROR] Your response was truncated (hit max_tokens limit) so this tool call "
                            "could not be executed — its arguments are likely incomplete. "
                            "Please retry with shorter content, or split the operation into smaller steps."
                        ),
                        tool_call_id=tc.tool_call_id,
                    )
                )
            
            truncation_request = ModelRequest(parts=truncation_results)
            messages.append(response)
            messages.append(truncation_request)
            
            # Apply context management middleware before retry
            if self.context_management and self._context_management_middleware:
                managed_msgs, ctx_full = await self._context_management_middleware.apply(messages)
                messages.clear()
                messages.extend(managed_msgs)
                self._propagate_context_management_usage()
                if ctx_full:
                    return self._context_management_middleware._build_context_full_response(
                        model_name=self.model.model_name
                    )
            
            model_params = self._build_model_request_parameters(getattr(self, 'current_task', None))
            model_params = self.model.customize_request_parameters(model_params)
            
            _retry_model_start: float = time.time()
            retry_response: "ModelResponse" = await self.model.request(
                messages=messages,
                model_settings=self.model.settings,
                model_request_parameters=model_params
            )
            _retry_model_elapsed: float = time.time() - _retry_model_start
            
            if hasattr(self, '_agent_run_output') and self._agent_run_output:
                self._agent_run_output.add_model_execution_time(_retry_model_elapsed)
                if hasattr(retry_response, 'usage') and retry_response.usage:
                    self._agent_run_output._ensure_usage().incr(retry_response.usage)
                    try:
                        from upsonic.utils.usage import calculate_cost_from_usage
                        cost_value: float = calculate_cost_from_usage(retry_response.usage, self.model)
                        self._agent_run_output.set_usage_cost(cost_value)
                    except Exception:
                        pass
            
            return await self._handle_model_response(retry_response, messages)
        
        if regular_tool_calls:
            tool_results = await self._execute_tool_calls(regular_tool_calls)
            
            if hasattr(self, '_tool_limit_reached') and self._tool_limit_reached:
                tool_request = ModelRequest(parts=tool_results)
                messages.append(response)
                messages.append(tool_request)
                
                limit_notification = UserPromptPart(
                    content=f"[SYSTEM] Tool call limit of {self.tool_call_limit} has been reached. "
                    f"No more tools are available. Please provide a final response based on the information you have."
                )
                limit_message = ModelRequest(parts=[limit_notification])
                messages.append(limit_message)
                
                # Apply context management middleware before model request
                if self.context_management and self._context_management_middleware:
                    managed_msgs, ctx_full = await self._context_management_middleware.apply(messages)
                    messages.clear()
                    messages.extend(managed_msgs)
                    self._propagate_context_management_usage()
                    if ctx_full:
                        return self._context_management_middleware._build_context_full_response(
                            model_name=self.model.model_name
                        )
                
                model_params = self._build_model_request_parameters(getattr(self, 'current_task', None))
                model_params = self.model.customize_request_parameters(model_params)
                
                _limit_model_start: float = time.time()
                final_response = await self.model.request(
                    messages=messages,
                    model_settings=self.model.settings,
                    model_request_parameters=model_params
                )
                _limit_model_elapsed: float = time.time() - _limit_model_start
                
                if hasattr(self, '_agent_run_output') and self._agent_run_output:
                    self._agent_run_output.add_model_execution_time(_limit_model_elapsed)
                    if hasattr(final_response, 'usage') and final_response.usage:
                        self._agent_run_output._ensure_usage().incr(final_response.usage)
                        try:
                            from upsonic.utils.usage import calculate_cost_from_usage
                            cost_value: float = calculate_cost_from_usage(final_response.usage, self.model)
                            self._agent_run_output.set_usage_cost(cost_value)
                        except Exception:
                            pass
                
                return final_response
            
            should_stop = False
            for tool_result in tool_results:
                if hasattr(tool_result, 'content') and isinstance(tool_result.content, dict):
                    if (
                        getattr(tool_result, "tool_name", None) == DEFAULT_OUTPUT_TOOL_NAME
                        and "result" in tool_result.content
                    ):
                        should_stop = True
                        break
                    if tool_result.content.get('_stop_execution'):
                        should_stop = True
                        tool_result.content.pop('_stop_execution', None)

            tool_request = ModelRequest(parts=tool_results)
            messages.append(response)
            messages.append(tool_request)
            
            if should_stop:
                final_text = ""
                for tool_result in tool_results:
                    if hasattr(tool_result, 'content'):
                        if isinstance(tool_result.content, dict):
                            final_text = str(tool_result.content.get('func', tool_result.content))
                        else:
                            final_text = str(tool_result.content)
                
                stop_response = ModelResponse(
                    parts=[TextPart(content=final_text)],
                    model_name=response.model_name,
                    timestamp=response.timestamp,
                    usage=response.usage,
                    provider_name=response.provider_name,
                    provider_response_id=response.provider_response_id,
                    provider_details=response.provider_details,
                    finish_reason="stop"
                )
                return stop_response
            
            # Apply context management middleware before follow-up model request
            if self.context_management and self._context_management_middleware:
                managed_msgs, ctx_full = await self._context_management_middleware.apply(messages)
                messages.clear()
                messages.extend(managed_msgs)
                self._propagate_context_management_usage()
                if ctx_full:
                    return self._context_management_middleware._build_context_full_response(
                        model_name=self.model.model_name
                    )
            
            model_params = self._build_model_request_parameters(getattr(self, 'current_task', None))
            model_params = self.model.customize_request_parameters(model_params)
            
            _followup_model_start: float = time.time()
            follow_up_response = await self.model.request(
                messages=messages,
                model_settings=self.model.settings,
                model_request_parameters=model_params
            )
            _followup_model_elapsed: float = time.time() - _followup_model_start
            
            if hasattr(self, '_agent_run_output') and self._agent_run_output:
                self._agent_run_output.add_model_execution_time(_followup_model_elapsed)
                if hasattr(follow_up_response, 'usage') and follow_up_response.usage:
                    self._agent_run_output._ensure_usage().incr(follow_up_response.usage)
                    try:
                        from upsonic.utils.usage import calculate_cost_from_usage
                        cost_value: float = calculate_cost_from_usage(follow_up_response.usage, self.model)
                        self._agent_run_output.set_usage_cost(cost_value)
                    except Exception:
                        pass
            
            return await self._handle_model_response(follow_up_response, messages)
        
        return response
    
    def _propagate_context_management_usage(self) -> None:
        """Propagate usage from ContextManagementMiddleware summarization LLM call
        to the parent AgentRunOutput context.
        """
        if not hasattr(self, '_agent_run_output') or not self._agent_run_output:
            return
        if not self._context_management_middleware:
            return
        summarization_usage = getattr(self._context_management_middleware, '_last_summarization_usage', None)
        if summarization_usage is not None:
            self._agent_run_output._ensure_usage().incr(summarization_usage)
            try:
                from upsonic.utils.usage import calculate_cost_from_usage
                summarization_model: "Model" = self._context_management_middleware._get_summarization_model()
                cost_value: float = calculate_cost_from_usage(summarization_usage, summarization_model)
                self._agent_run_output.set_usage_cost(cost_value)
            except Exception:
                pass
            # Reset to prevent double-counting on next apply() call
            self._context_management_middleware._last_summarization_usage = None
    
    def _drain_agent_tool_usage(self, tool_name: str) -> None:
        """Drain accumulated usage from AgentTool sub-agent executions into the parent AgentRunOutput.
        
        Args:
            tool_name: The name of the tool that was just executed
        """
        from upsonic.tools.wrappers import AgentTool
        
        registered_tool = self.tool_manager.processor.registered_tools.get(tool_name)
        if registered_tool is None:
            current_task = getattr(self, 'current_task', None)
            if current_task is not None and current_task.tool_manager is not None:
                registered_tool = current_task.tool_manager.processor.registered_tools.get(tool_name)
        if registered_tool and isinstance(registered_tool, AgentTool):
            agent_tool_usage = registered_tool.drain_accumulated_usage()
            if agent_tool_usage is not None:
                self._agent_run_output.usage.incr(agent_tool_usage)
    
    async def _handle_cache(self, task: "Task") -> Optional[Any]:
        """Handle cache operations for the task."""
        if not task.enable_cache:
            return None
        
        if self.debug:
            from upsonic.utils.printing import cache_configuration, debug_log_level2
            embedding_provider_name = None
            if task.cache_embedding_provider:
                embedding_provider_name = task.cache_embedding_provider.model_name
            
            cache_configuration(
                enable_cache=task.enable_cache,
                cache_method=task.cache_method,
                cache_threshold=task.cache_threshold if task.cache_method == "vector_search" else None,
                cache_duration_minutes=task.cache_duration_minutes,
                embedding_provider=embedding_provider_name
            )
            
            # Level 2: Detailed cache information
            if self.debug_level >= 2:
                debug_log_level2(
                    "Cache check details",
                    "Agent",
                    debug=self.debug,
                    debug_level=self.debug_level,
                    task_description=task.description[:100] if task.description else None,
                    cache_enabled=task.enable_cache,
                    cache_method=task.cache_method,
                    cache_threshold=task.cache_threshold,
                    cache_duration_minutes=task.cache_duration_minutes,
                    embedding_provider=embedding_provider_name,
                    model_name=self.model.model_name
                )
        
        input_text = task._original_input or task.description
        cached_response = await task.get_cached_response(input_text, self.model)
        
        if cached_response is not None:
            similarity = None
            if hasattr(task, '_last_cache_entry') and 'similarity' in task._last_cache_entry:
                similarity = task._last_cache_entry['similarity']
            
            from upsonic.utils.printing import cache_hit
            cache_hit(
                cache_method=task.cache_method,
                similarity=similarity,
                input_preview=(task._original_input or task.description)[:100] if (task._original_input or task.description) else None
            )
            
            return cached_response
        else:
            from upsonic.utils.printing import cache_miss
            cache_miss(
                cache_method=task.cache_method,
                input_preview=(task._original_input or task.description)[:100] if (task._original_input or task.description) else None
            )
            return None
    
    async def _apply_user_policy(
        self, 
        task: "Task", 
        context: Optional[AgentRunOutput] = None
    ) -> tuple[Optional["Task"], bool]:
        """
        Apply user policy to task input.
        
        This method now uses PolicyManager to handle multiple policies.
        When feedback is enabled, returns a helpful message to the user instead
        of hard blocking, explaining what was wrong and how to correct it.
        
        Args:
            task: The task to apply policy to
            context: Optional AgentRunOutput for event emission
        
        Returns:
            tuple: (task, should_continue)
                - task: The task (possibly modified with feedback response)
                - should_continue: False if task should stop (blocked or feedback given)
        """
        if not self.user_policy_manager.has_policies() or not task.description:
            # Emit ALLOW event if no policies
            if context and context.is_streaming:
                from upsonic.utils.agent.events import ayield_policy_check_event
                async for event in ayield_policy_check_event(
                    run_id=context.run_id or "",
                    policy_type='user_policy',
                    action='ALLOW',
                    policies_checked=0
                ):
                    context.events.append(event)
            return task, True
        
        from upsonic.safety_engine.models import PolicyInput
        
        policy_input = PolicyInput(input_texts=[task.description])
        result = await self.user_policy_manager.execute_policies_async(
            policy_input,
            check_type="User Input Check"
        )
        
        # Get policies checked count
        policies_checked = len(self.user_policy_manager.policies)
        
        # Map action_taken to event action
        action_mapping = {
            "ALLOW": "ALLOW",
            "BLOCK": "BLOCK",
            "REPLACE": "REPLACE",
            "ANONYMIZE": "ANONYMIZE",
            "DISALLOWED_EXCEPTION": "RAISE ERROR"
        }
        event_action = action_mapping.get(result.action_taken, "ALLOW")
        
        # Emit PolicyCheckEvent
        if context and context.is_streaming:
            from upsonic.utils.agent.events import ayield_policy_check_event
            content_modified = result.action_taken in ["REPLACE", "ANONYMIZE"]
            blocked_reason = result.message if result.action_taken == "BLOCK" else None
            
            async for event in ayield_policy_check_event(
                run_id=context.run_id or "",
                policy_type='user_policy',
                action=event_action,
                policies_checked=policies_checked,
                content_modified=content_modified,
                blocked_reason=blocked_reason
            ):
                context.events.append(event)
        
        # Emit PolicyFeedbackEvent if feedback was generated
        if context and context.is_streaming and result.feedback_message:
            from upsonic.utils.agent.events import ayield_policy_feedback_event
            async for event in ayield_policy_feedback_event(
                run_id=context.run_id or "",
                policy_type='user_policy',
                feedback_message=result.feedback_message,
                retry_count=self.user_policy_manager._current_retry_count,
                max_retries=self.user_policy_manager.feedback_loop_count,
                violated_policy=result.violated_policy_name
            ):
                context.events.append(event)
        
        if result.should_block():
            # Re-raise DisallowedOperation if it was caught by PolicyManager
            # (unless feedback was generated - then we want to return the feedback)
            if result.disallowed_exception and not result.feedback_message:
                raise result.disallowed_exception
            
            task.task_end()
            # Use feedback message if available (gives helpful guidance to user)
            task._response = result.get_final_message()
            
            # Print feedback info if debug mode and feedback was generated
            if self.debug and result.feedback_message:
                from upsonic.utils.printing import user_policy_feedback_returned, debug_log_level2
                user_policy_feedback_returned(
                    policy_name=result.violated_policy_name or "Unknown Policy",
                    feedback_message=result.feedback_message
                )
                # Level 2: Detailed policy feedback information
                if self.debug_level >= 2:
                    debug_log_level2(
                        "User policy feedback details",
                        "Agent",
                        debug=self.debug,
                        debug_level=self.debug_level,
                        policy_name=result.violated_policy_name or "Unknown Policy",
                        feedback_message=result.feedback_message,
                        action_taken=result.action_taken,
                        confidence=result.confidence if hasattr(result, 'confidence') else None,
                        original_input=task.description[:200] if task.description else None
                    )
            return task, False
        elif result.action_taken in ["REPLACE", "ANONYMIZE"]:
            task.description = result.final_output or task.description
            # Store transformation map for de-anonymization of LLM response
            if result.transformation_map:
                task._anonymization_map = result.transformation_map
                # Add anonymization notice to the prompt - prepend for visibility
                anonymization_notice = (
                    "[PRIVACY MODE ACTIVE: Personal data has been anonymized with random placeholders. "
                    "Answer the question directly using the placeholder values shown. "
                    "Do NOT comment on, question, or mention the format of any data.]\n\n"
                )
                task.description = anonymization_notice + task.description
                if self.debug:
                    from upsonic.utils.printing import debug_log, anonymization_debug_panel
                    debug_log(
                        f"Stored anonymization map with {len(result.transformation_map)} entries for de-anonymization",
                        "Agent"
                    )
                    # Show detailed anonymization info - what LLM will see vs original
                    anonymization_debug_panel(
                        original_values=[entry.get("original", "") for entry in result.transformation_map.values()],
                        anonymized_values=[entry.get("anonymous", "") for entry in result.transformation_map.values()],
                        llm_input=task.description
                    )
            return task, True
        
        return task, True
    
    async def _execute_with_guardrail(self, task: "Task", memory_handler: Optional["MemoryManager"], state: Optional["State"] = None) -> "ModelResponse":
        """
        Executes the agent's run method with a validation and retry loop based on a task guardrail.
        This method encapsulates the retry logic, hiding it from the main `do_async` pipeline.
        It returns a single, "clean" ModelResponse that represents the final, successful interaction.
        """
        from upsonic.messages import TextPart, ModelResponse
        retry_counter = 0
        validation_passed = False
        final_model_response = None
        last_error_message = ""
        
        temporary_message_history = copy.deepcopy(memory_handler.get_message_history())
        
        if hasattr(self, '_agent_run_output') and self._agent_run_output and self._agent_run_output.input:
            run_input = self._agent_run_output.input
            if run_input.input is None:
                run_input.build_input(context_formatted=task.context_formatted)
            current_input = run_input.input
            if task.context_formatted:
                task.context_formatted = None
        else:
            raise RuntimeError("AgentRunInput not available. This should not happen.")

        if task.guardrail_retries is not None and task.guardrail_retries > 0:
            max_retries = task.guardrail_retries + 1
        else:
            max_retries = 1

        while not validation_passed and retry_counter < max_retries:
            messages, context_full_response = await self._build_model_request_with_input(task, memory_handler, current_input, temporary_message_history, state)
            
            if context_full_response is not None:
                return context_full_response
            
            model_params = self._build_model_request_parameters(task)
            model_params = self.model.customize_request_parameters(model_params)
            
            _guardrail_model_start: float = time.time()
            response = await self.model.request(
                messages=messages,
                model_settings=self.model.settings,
                model_request_parameters=model_params
            )
            _guardrail_model_elapsed: float = time.time() - _guardrail_model_start
            
            if hasattr(self, '_agent_run_output') and self._agent_run_output:
                self._agent_run_output.add_model_execution_time(_guardrail_model_elapsed)
                if hasattr(response, 'usage') and response.usage:
                    self._agent_run_output._ensure_usage().incr(response.usage)
                    try:
                        from upsonic.utils.usage import calculate_cost_from_usage
                        cost_value: float = calculate_cost_from_usage(response.usage, self.model)
                        self._agent_run_output.set_usage_cost(cost_value)
                    except Exception:
                        pass
            
            current_model_response = await self._handle_model_response(response, messages)
            
            if task.guardrail is None:
                validation_passed = True
                final_model_response = current_model_response
                break

            final_text_output = ""
            text_parts = [part.content for part in current_model_response.parts if isinstance(part, TextPart)]
            final_text_output = "".join(text_parts)

            if not final_text_output:
                validation_passed = True
                final_model_response = current_model_response
                break

            try:
                # Parse structured output if response_format is a Pydantic model
                guardrail_input = final_text_output
                if task.response_format and task.response_format != str:
                    try:
                        import json
                        parsed = json.loads(final_text_output)
                        if hasattr(task.response_format, 'model_validate'):
                            guardrail_input = task.response_format.model_validate(parsed)
                    except:
                        # If parsing fails, use the text output
                        guardrail_input = final_text_output
                
                guardrail_result = task.guardrail(guardrail_input)
                
                if isinstance(guardrail_result, tuple) and len(guardrail_result) == 2:
                    is_valid, result = guardrail_result
                elif isinstance(guardrail_result, bool):
                    is_valid = guardrail_result
                    result = final_text_output if guardrail_result else "Guardrail validation failed"
                else:
                    is_valid = bool(guardrail_result)
                    result = guardrail_result if guardrail_result else "Guardrail validation failed"

                if is_valid:
                    validation_passed = True
                    
                    if result != final_text_output:
                        updated_parts = []
                        found_and_updated = False
                        for part in current_model_response.parts:
                            if isinstance(part, TextPart) and not found_and_updated:
                                updated_parts.append(TextPart(content=str(result)))
                                found_and_updated = True
                            elif isinstance(part, TextPart):
                                updated_parts.append(TextPart(content=""))
                            else:
                                updated_parts.append(part)
                        
                        final_model_response = ModelResponse(
                            parts=updated_parts,
                            model_name=current_model_response.model_name,
                            timestamp=current_model_response.timestamp,
                            usage=current_model_response.usage,
                            provider_name=current_model_response.provider_name,
                            provider_response_id=current_model_response.provider_response_id,
                            provider_details=current_model_response.provider_details,
                            finish_reason=current_model_response.finish_reason
                        )
                    else:
                        final_model_response = current_model_response
                    break
                else:
                    retry_counter += 1
                    last_error_message = str(result)
                    
                    temporary_message_history.append(current_model_response)
                    
                    correction_prompt = f"Your previous response failed a validation check. Please review the reason and provide a corrected response. Failure Reason: {last_error_message}"
                    current_input = correction_prompt
                    
            except Exception as e:
                retry_counter += 1
                last_error_message = f"Guardrail execution error: {str(e)}"
                
                temporary_message_history.append(current_model_response)
                
                correction_prompt = f"Your previous response failed a validation check. Please review the reason and provide a corrected response. Failure Reason: {last_error_message}"
                current_input = correction_prompt

        if not validation_passed:
            error_msg = f"Task failed after {max_retries-1} retry(s). Last error: {last_error_message}"
            if self.mode == "raise":
                from upsonic.utils.package.exception import GuardrailValidationError
                raise GuardrailValidationError(error_msg)
            else:
                error_response = ModelResponse(
                    parts=[TextPart(content="Guardrail validation failed after retries")],
                    model_name=self.model.model_name,
                    timestamp=now_utc(),
                    usage=RequestUsage()
                )
                return error_response
                
        return final_model_response
    
    async def recommend_model_for_task_async(
        self,
        task: Union["Task", str],
        criteria: Optional[Dict[str, Any]] = None,
        use_llm: Optional[bool] = None
    ) -> "ModelRecommendation":
        """
        Get a model recommendation for a specific task.
        
        This method analyzes the task and returns a recommendation for the best model to use.
        The user can then decide whether to use the recommended model or stick with the default.
        
        Args:
            task: Task object or task description string
            criteria: Optional criteria dictionary for model selection (overrides agent's default)
            use_llm: Optional flag to use LLM for selection (overrides agent's default)
        
        Returns:
            ModelRecommendation: Object containing:
                - model_name: Recommended model identifier
                - reason: Explanation for the recommendation
                - confidence_score: Confidence level (0.0 to 1.0)
                - selection_method: "rule_based" or "llm_based"
                - estimated_cost_tier: Cost estimate (1-10)
                - estimated_speed_tier: Speed estimate (1-10)
                - alternative_models: List of alternative model names
        
        Example:
            ```python
            # Get recommendation
            recommendation = await agent.recommend_model_for_task_async(task)
            print(f"Recommended: {recommendation.model_name}")
            print(f"Reason: {recommendation.reason}")
            print(f"Confidence: {recommendation.confidence_score}")
            
            # Use it if you have credentials
            if user_has_credentials(recommendation.model_name):
                result = await agent.do_async(task, model=recommendation.model_name)
            else:
                result = await agent.do_async(task)  # Use default
            ```
        """
        try:
            from upsonic.models.model_selector import select_model_async, SelectionCriteria
            
            task_description = task.description if hasattr(task, 'description') else str(task)
            
            selection_criteria = None
            if criteria:
                selection_criteria = SelectionCriteria(**criteria)
            elif self.model_selection_criteria:
                selection_criteria = SelectionCriteria(**self.model_selection_criteria)
            
            use_llm_selection = use_llm if use_llm is not None else self.use_llm_for_selection
            
            recommendation = await select_model_async(
                task_description=task_description,
                criteria=selection_criteria,
                use_llm=use_llm_selection,
                agent=self if use_llm_selection else None,
                default_model=self.model.model_name
            )
            
            self._model_recommendation = recommendation
            
            if self.debug:
                from upsonic.utils.printing import model_recommendation_summary
                model_recommendation_summary(recommendation)
            
            return recommendation
            
        except Exception as e:
            if self.debug:
                from upsonic.utils.printing import model_recommendation_error
                model_recommendation_error(str(e))
            raise
    
    def recommend_model_for_task(
        self,
        task: Union["Task", str],
        criteria: Optional[Dict[str, Any]] = None,
        use_llm: Optional[bool] = None
    ) -> "ModelRecommendation":
        """
        Synchronous version of recommend_model_for_task_async.
        
        Get a model recommendation for a specific task.
        
        Args:
            task: Task object or task description string
            criteria: Optional criteria dictionary for model selection
            use_llm: Optional flag to use LLM for selection
        
        Returns:
            ModelRecommendation: Object containing recommendation details
        
        Example:
            ```python
            recommendation = agent.recommend_model_for_task("Write a sorting algorithm")
            print(f"Use: {recommendation.model_name}")
            ```
        """
        return _run_in_bg_loop(self.recommend_model_for_task_async(task, criteria, use_llm))

    def get_last_model_recommendation(self) -> Optional[Any]:
        """
        Get the last model recommendation made by the agent.
        
        Returns:
            ModelRecommendation object or None if no recommendation was made
        """
        return self._model_recommendation
    

    async def _apply_agent_policy(
        self, 
        task: "Task", 
        context: Optional[AgentRunOutput] = None
    ) -> tuple["Task", Optional[str]]:
        """
        Apply agent policy to task output.
        
        This method uses PolicyManager to handle multiple policies.
        When feedback is enabled and a violation occurs, it returns the feedback
        message along with the task so the caller can decide to retry.
        
        Args:
            task: The task to apply policy to
            context: Optional AgentRunOutput for event emission
        
        Returns:
            tuple: (task, feedback_message_or_none)
                - task: The task (possibly modified with blocked response)
                - feedback_message: If not None, agent should retry with this feedback
        """
        if not self.agent_policy_manager.has_policies() or not task or not task.response:
            # Emit ALLOW event if no policies
            if context and context.is_streaming:
                from upsonic.utils.agent.events import ayield_policy_check_event
                async for event in ayield_policy_check_event(
                    run_id=context.run_id or "",
                    policy_type='agent_policy',
                    action='ALLOW',
                    policies_checked=0
                ):
                    context.events.append(event)
            return task, None
        
        from upsonic.safety_engine.models import PolicyInput
        
        # Convert response to text
        response_text = ""
        if isinstance(task.response, str):
            response_text = task.response
        elif hasattr(task.response, 'model_dump_json'):
            response_text = task.response.model_dump_json()
        else:
            response_text = str(task.response)
        
        if not response_text:
            return task, None
        
        agent_policy_input = PolicyInput(input_texts=[response_text])
        result = await self.agent_policy_manager.execute_policies_async(
            agent_policy_input,
            check_type="Agent Output Check"
        )
        
        # Get policies checked count
        policies_checked = len(self.agent_policy_manager.policies)
        original_response = task.response
        
        # Map action_taken to event action
        action_mapping = {
            "ALLOW": "ALLOW",
            "BLOCK": "BLOCK",
            "REPLACE": "REPLACE",
            "ANONYMIZE": "ANONYMIZE",
            "DISALLOWED_EXCEPTION": "RAISE ERROR"
        }
        event_action = action_mapping.get(result.action_taken, "ALLOW")
        
        # Emit PolicyCheckEvent
        if context and context.is_streaming:
            from upsonic.utils.agent.events import ayield_policy_check_event
            content_modified = result.action_taken in ["REPLACE", "ANONYMIZE"] or (
                result.final_output and str(result.final_output) != str(original_response)
            )
            blocked_reason = result.message if result.action_taken == "BLOCK" else None
            
            async for event in ayield_policy_check_event(
                run_id=context.run_id or "",
                policy_type='agent_policy',
                action=event_action,
                policies_checked=policies_checked,
                content_modified=content_modified,
                blocked_reason=blocked_reason
            ):
                context.events.append(event)
        
        # Check if retry with feedback should be attempted
        if result.should_retry_with_feedback() and self.agent_policy_manager.can_retry():
            # Emit PolicyFeedbackEvent
            if context and context.is_streaming:
                from upsonic.utils.agent.events import ayield_policy_feedback_event
                async for event in ayield_policy_feedback_event(
                    run_id=context.run_id or "",
                    policy_type='agent_policy',
                    feedback_message=result.feedback_message,
                    retry_count=self.agent_policy_manager._current_retry_count,
                    max_retries=self.agent_policy_manager.feedback_loop_count,
                    violated_policy=result.violated_policy_name
                ):
                    context.events.append(event)
            # Return feedback message for retry - don't modify task yet
            return task, result.feedback_message
        
        # Apply the result (no retry - either passed or exhausted retries)
        if result.should_block():
            # Re-raise DisallowedOperation if it was caught by PolicyManager
            if result.disallowed_exception and not result.feedback_message:
                raise result.disallowed_exception
            
            task._response = result.get_final_message()
        elif result.action_taken in ["REPLACE", "ANONYMIZE"]:
            task._response = result.final_output or "Response modified by agent policy."
        elif result.final_output:
            task._response = result.final_output
        
        return task, None

    
    
    @retryable(retries_from_param="retry")
    async def do_async(
        self,
        task: Union[str, "Task", List[Union[str, "Task"]]],
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        return_output: bool = False,
        state: Optional["State"] = None,
        *,
        timeout: Optional[float] = None,
        partial_on_timeout: bool = False,
        graph_execution_id: Optional[str] = None,
        _resume_output: Optional[AgentRunOutput] = None,
        _resume_step_index: Optional[int] = None,
        _print_method_default: bool = False,
    ) -> Union[Any, List[Any], List[AgentRunOutput]]:
        """
        Execute a task (or list of tasks) asynchronously using the pipeline architecture.

        When a list of tasks is provided with more than one element, each task is
        executed sequentially and a list of results is returned.  A single-element
        list is unwrapped and treated identically to a non-list input.

        Args:
            task: Task to execute. Accepts a single Task/str or a list of them.
            model: Override model for this execution
            debug: Enable debug mode
            retry: Number of retries
            return_output: If True, return full AgentRunOutput. If False (default), return content only.
            state: Graph execution state
            timeout: Maximum execution time in seconds. None means no timeout.
            partial_on_timeout: If True and timeout is set, return whatever text was
                generated so far instead of raising an error. Requires timeout to be set.
                Internally uses streaming to enable progressive text capture.
            graph_execution_id: Graph execution identifier
            _resume_output: Internal - output for HITL resumption
            _resume_step_index: Internal - step index to resume from
            _print_method_default: Internal - default print value based on method (do=False, print_do=True)

        Returns:
            Single task:
                Task content (str, BaseModel, etc.) if return_output=False
                Full AgentRunOutput if return_output=True
            List of tasks (len > 1):
                List of task contents if return_output=False
                List of AgentRunOutput if return_output=True

        Example:
            ```python
            # Single task
            result = await agent.do_async(task)

            # With timeout and partial results
            result = await agent.do_async(task, timeout=120, partial_on_timeout=True)
            # If timeout hits: returns whatever text was generated so far
            # If completed normally: returns full result
            ```
        """
        handled, task_or_results = await self._handle_task_list_async(
            task, self.do_async,
            model, debug, retry, return_output, state,
            timeout=timeout, partial_on_timeout=partial_on_timeout,
            graph_execution_id=graph_execution_id,
            _print_method_default=_print_method_default,
        )
        if handled:
            return task_or_results
        task = task_or_results

        # Validate timeout parameters
        if partial_on_timeout and timeout is None:
            raise ValueError("partial_on_timeout=True requires timeout to be set")
        if timeout is not None and timeout <= 0:
            raise ValueError("timeout must be a positive number")

        resolved_print_flag = self._resolve_print_flag(_print_method_default)

        # Convert string to Task if needed
        task = self._convert_to_task(task)
        
        start_step_index = _resume_step_index if _resume_step_index is not None else 0
        is_resuming = _resume_output is not None
        effective_retry = self.retry if getattr(self, "retry", None) is not None else retry

        validation_error = self._validate_task_for_new_run(
            task, is_resuming, allow_problematic_for_retry=(effective_retry > 1)
        )
        if validation_error is not None:
            if return_output:
                return validation_error
            return validation_error.output

        if not is_resuming and effective_retry > 1 and task.is_problematic:
            task.status = None
            task.run_id = None
            # Clear the stale output from the previous failed attempt so no code
            # can accidentally read its error status before the new output is created.
            self._agent_run_output = None

        # Only reset per-run agent state for fresh runs — HITL resume should
        # keep existing state so cost and tool counts aggregate correctly.
        if not is_resuming:
            import uuid as _uuid
            task.price_id_ = str(_uuid.uuid4())
            self._tool_call_count = 0
            self._tool_limit_reached = False
        self._last_built_system_prompt = None

        if debug or self.debug:
            self.user_policy_manager.debug = True
            self.agent_policy_manager.debug = True

        if is_resuming:
            run_id = _resume_output.run_id
            self.run_id = run_id
            self._agent_run_output = _resume_output
            self._agent_run_output.is_streaming = False
            self._agent_run_output.print_flag = resolved_print_flag
        else:
            run_id = str(uuid.uuid4())
            self.run_id = run_id
            register_run(run_id)

        original_model: Optional["Model"] = None
        try:
            if not is_resuming:
                run_input = self._create_agent_run_input(task)
                self._agent_run_output = self._create_agent_run_output(
                    run_id=run_id,
                    task=task,
                    run_input=run_input,
                    is_streaming=False
                )
                self._agent_run_output.print_flag = resolved_print_flag
                if task is not None:
                    task.run_id = run_id

            original_model = self._apply_model_override(model)

            _pl_start_time: float = time.time()
            task_desc: str = str(task.description) if task is not None and hasattr(task, "description") else ""
            with self._otel.agent_run_span(
                run_id,
                name=self.name or "",
                model=str(self.model_name) if self.model_name else "",
                task_description=task_desc,
                user_id=self.user_id,
                session_id=self.session_id,
            ) as otel_span:
                try:
                    result = await self._do_async_pipeline(
                        task=task,
                        run_id=run_id,
                        debug=debug,
                        return_output=return_output,
                        timeout=timeout,
                        partial_on_timeout=partial_on_timeout,
                        start_step_index=start_step_index,
                    )
                    self._otel.finalize_agent_run(
                        otel_span,
                        getattr(self, "_agent_run_output", None),
                        agent_user_id=self.user_id,
                        agent_session_id=self.session_id,
                        total_cost=self._calculate_aggregated_cost(),
                        tool_definitions=self._get_combined_tool_definitions(),
                    )
                    _trace_id = self._otel.extract_trace_id(otel_span)
                    if _trace_id and self._agent_run_output is not None:
                        self._agent_run_output.trace_id = _trace_id
                    self._log_to_promptlayer_background(
                        task=task,
                        output=self._agent_run_output.output if self._agent_run_output else result,
                        start_time=_pl_start_time,
                        end_time=time.time(),
                    )
                    return result
                except Exception as exc:
                    self._otel.record_error(otel_span, exc)
                    raise
        finally:
            if original_model is not None:
                self.model = original_model
            self._finalize_agent_usage(resolved_print_flag)
            # Keep run_id alive when task is paused (HITL) — the run is still active
            _output = getattr(self, '_agent_run_output', None)
            _is_paused = getattr(_output.task, 'is_paused', False) if _output and _output.task else False
            if not _is_paused:
                self.run_id = None

    def _calculate_aggregated_cost(self) -> Optional[float]:
        """Calculate the aggregated monetary cost across the agent run.

        Attempts to derive cost from the current task's pricing data first,
        then falls back to computing it from RunUsage and model.
        """
        task = getattr(self, "current_task", None)
        if task is not None:
            task_cost = getattr(task, "total_cost", None)
            if task_cost is not None:
                return float(task_cost)

        output = getattr(self, "_agent_run_output", None)
        run_usage = getattr(output, "usage", None) if output else None
        if run_usage is None:
            return None

        try:
            from upsonic.utils.usage import calculate_cost_from_usage
            model = getattr(self, "model", None) or getattr(self, "model_name", None)
            if model is not None:
                return calculate_cost_from_usage(run_usage, model)
        except Exception:
            pass
        return None

    async def _do_async_pipeline(
        self,
        task: "Task",
        run_id: str,
        debug: bool,
        return_output: bool,
        timeout: Optional[float],
        partial_on_timeout: bool,
        start_step_index: int,
    ) -> Union[Any, "AgentRunOutput"]:
        """Core pipeline execution logic extracted from do_async for OTel span wrapping."""
        from upsonic.agent.pipeline import PipelineManager

        if partial_on_timeout and timeout is not None:
            self._agent_run_output.is_streaming = True

            pipeline = PipelineManager(
                steps=self._create_streaming_pipeline_steps(),
                task=self._agent_run_output.task,
                agent=self,
                model=self.model,
                debug=debug or self.debug
            )

            async def _consume_stream():
                async for pipeline_event in pipeline.execute_stream(
                    context=self._agent_run_output, start_step_index=start_step_index
                ):
                    text_content = self._extract_text_from_stream_event(pipeline_event)
                    if text_content:
                        self._agent_run_output.accumulated_text += text_content

            try:
                await asyncio.wait_for(_consume_stream(), timeout=timeout)
            except (asyncio.TimeoutError, asyncio.CancelledError):
                cancel_run_func(run_id)

                partial_text = self._agent_run_output.accumulated_text or None

                self._agent_run_output.mark_cancelled()
                if self._agent_run_output.metadata is None:
                    self._agent_run_output.metadata = {}
                self._agent_run_output.metadata["timeout"] = True
                self._agent_run_output.metadata["partial_result"] = bool(partial_text)
                self._agent_run_output.metadata["timeout_seconds"] = timeout

                self._agent_run_output.output = partial_text
                if task is not None:
                    task._response = partial_text

                await self._run_call_management_step(task, debug)

                cleanup_run(run_id)
                sentry_sdk.flush()

                if return_output:
                    return self._agent_run_output
                return self._agent_run_output.output

            self._agent_run_output.is_streaming = False

            cleanup_run(run_id)
            sentry_sdk.flush()
            if return_output:
                return self._agent_run_output
            return self._agent_run_output.output

        elif timeout is not None:
            from upsonic.exceptions import ExecutionTimeoutError

            pipeline = PipelineManager(
                steps=self._create_direct_pipeline_steps(),
                task=self._agent_run_output.task,
                agent=self,
                model=self.model,
                debug=debug or self.debug
            )

            try:
                await asyncio.wait_for(
                    pipeline.execute(self._agent_run_output, start_step_index=start_step_index),
                    timeout=timeout
                )
            except (asyncio.TimeoutError, asyncio.CancelledError):
                cancel_run_func(run_id)
                cleanup_run(run_id)
                raise ExecutionTimeoutError(
                    f"Agent execution timed out after {timeout} seconds",
                    timeout=timeout
                )

            cleanup_run(run_id)
            sentry_sdk.flush()
            if return_output:
                return self._agent_run_output
            return self._agent_run_output.output

        else:
            pipeline = PipelineManager(
                steps=self._create_direct_pipeline_steps(),
                task=self._agent_run_output.task,
                agent=self,
                model=self.model,
                debug=debug or self.debug
            )

            await pipeline.execute(self._agent_run_output, start_step_index=start_step_index)
            cleanup_run(run_id)
            sentry_sdk.flush()
            if return_output:
                return self._agent_run_output
            return self._agent_run_output.output

    async def _run_call_management_step(self, task: "Task", debug: bool = False) -> None:
        """Run CallManagementStep to record usage for streaming-based execution.

        The streaming pipeline lacks CallManagementStep, so when do_async()
        uses streaming internally (partial_on_timeout), we run it manually
        so that task.total_input_token / total_output_token work.
        """
        try:
            from upsonic.agent.pipeline.steps import CallManagementStep
            step = CallManagementStep()
            await step.execute(
                self._agent_run_output, task, self, self.model,
                step_number=0, pipeline_manager=None,
            )
        except Exception:
            pass  # Best-effort — don't let tracking errors mask the result

    def _accumulate_run_usage(self, task_usage: Optional["TaskUsage"]) -> None:
        """Accumulate a task's usage metrics into the agent-level usage.

        Args:
            task_usage: The TaskUsage from a completed task. If None, no-op.
        """
        if task_usage is None:
            return

        from upsonic.usage import AgentUsage

        if self.usage is None:
            self.usage = AgentUsage()

        self.usage.incr(task_usage)

    def _finalize_agent_usage(self, print_flag: bool) -> None:
        """Accumulate current run usage into agent-level usage and optionally print agent metrics.
        
        Should be called at the end of every do_async / astream execution.
        Skips both accumulation AND printing when the task is paused (waiting
        for HITL) because the TaskUsage object persists across rounds and
        accumulates model_execution_time cumulatively.  Accumulating while
        paused would double-count the initial round's metrics once the
        continuation completes and accumulates the full cumulative TaskUsage.
        
        Args:
            print_flag: Whether printing is enabled for this execution.
        """
        output = getattr(self, '_agent_run_output', None)
        if output is None:
            return
        
        task = output.task
        is_paused: bool = getattr(task, 'is_paused', False) if task else False
        
        if not is_paused:
            self._accumulate_run_usage(output.usage)
        
        if print_flag and task and not getattr(task, 'not_main_task', False) and not is_paused:
            from upsonic.utils.printing import print_agent_metrics
            print_agent_metrics(self, print_output=print_flag)

    def _extract_output(self, task: "Task", response: "ModelResponse") -> Any:
        """Extract the output from a model response."""
        from upsonic.messages import TextPart, ToolCallPart
        
        # Check for image outputs first
        images = response.images
        if images:
            # If there are multiple images, return a list; if single, return the image data
            if len(images) == 1:
                return images[0].data
            else:
                return [img.data for img in images]
        
        # Check for tool call output from structured output tool
        if task.response_format and task.response_format != str and task.response_format is not str:
            tool_call_parts = [part for part in response.parts if isinstance(part, ToolCallPart)]
            for tool_call in tool_call_parts:
                # Look for the output tool
                if tool_call.tool_name == DEFAULT_OUTPUT_TOOL_NAME:
                    try:
                        args = tool_call.args_as_dict()
                        if hasattr(task.response_format, 'model_validate'):
                            return task.response_format.model_validate(args)
                        return args
                    except Exception:
                        pass
        
        # Extract text parts for non-image responses
        text_parts = [part.content for part in response.parts if isinstance(part, TextPart)]
        
        if task.response_format == str or task.response_format is str:
            return "".join(text_parts)
        
        text_content = "".join(text_parts)
        if task.response_format != str and text_content:
            try:
                import json
                parsed = json.loads(text_content)
                if hasattr(task.response_format, 'model_validate'):
                    return task.response_format.model_validate(parsed)
                return parsed
            except:
                pass
        
        return text_content
    
    def do(
        self,
        task: Union[str, "Task", List[Union[str, "Task"]]],
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        return_output: bool = False,
        timeout: Optional[float] = None,
        partial_on_timeout: bool = False,
    ) -> Union[Any, List[Any], List[AgentRunOutput]]:
        """
        Execute a task (or list of tasks) synchronously.

        When a list of tasks is provided with more than one element, each task is
        executed sequentially and a list of results is returned.  A single-element
        list is unwrapped and treated identically to a non-list input.

        Args:
            task: Task to execute. Accepts a single Task/str or a list of them.
            model: Override model for this execution
            debug: Enable debug mode
            retry: Number of retries
            return_output: If True, return full AgentRunOutput. If False (default), return content only.
            timeout: Maximum execution time in seconds. None means no timeout.
            partial_on_timeout: If True and timeout is set, return whatever text was
                generated so far instead of raising an error. Requires timeout to be set.
                Internally uses streaming to enable progressive text capture.

        Returns:
            Single task:
                Task content (str, BaseModel, etc.) if return_output=False
                Full AgentRunOutput if return_output=True
            List of tasks (len > 1):
                List of task contents if return_output=False
                List of AgentRunOutput if return_output=True
        """
        handled, task_or_results = self._handle_task_list(
            task, self.do, model, debug, retry, return_output,
            timeout, partial_on_timeout,
        )
        if handled:
            return task_or_results
        task = task_or_results

        # Auto-convert string to Task object if needed
        from upsonic.tasks.tasks import Task as TaskClass
        if isinstance(task, str):
            task = TaskClass(description=task)

        return _run_in_bg_loop(self.do_async(
            task, model, debug, retry, return_output,
            timeout=timeout, partial_on_timeout=partial_on_timeout,
            _print_method_default=False,
        ))

    def print_do(
        self,
        task: Union[str, "Task", List[Union[str, "Task"]]],
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        return_output: bool = False
    ) -> Union[Any, List[Any], List[AgentRunOutput]]:
        """
        Execute a task (or list of tasks) synchronously and print the result.
        
        When a list of tasks is provided with more than one element, each task is
        executed sequentially and a list of results is returned.  A single-element
        list is unwrapped and treated identically to a non-list input.
        
        Args:
            task: Task to execute. Accepts a single Task/str or a list of them.
            model: Override model for this execution
            debug: Enable debug mode
            retry: Number of retries
            return_output: If True, return full AgentRunOutput. If False (default), return content only.
            
        Returns:
            Single task:
                Task content (str, BaseModel, etc.) if return_output=False
                Full AgentRunOutput if return_output=True
            List of tasks (len > 1):
                List of task contents if return_output=False
                List of AgentRunOutput if return_output=True
        """
        handled, task_or_results = self._handle_task_list(
            task, self.print_do, model, debug, retry, return_output,
        )
        if handled:
            return task_or_results
        task = task_or_results

        # Auto-convert string to Task object if needed
        from upsonic.tasks.tasks import Task as TaskClass
        if isinstance(task, str):
            task = TaskClass(description=task)

        return _run_in_bg_loop(self.do_async(
            task, model, debug, retry, return_output,
            _print_method_default=True,
        ))

    async def print_do_async(
        self,
        task: Union[str, "Task", List[Union[str, "Task"]]],
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        return_output: bool = False
    ) -> Union[Any, List[Any], List[AgentRunOutput]]:
        """
        Execute a task (or list of tasks) asynchronously with print output enabled
        (unless overridden by ENV or Agent param).
        
        When a list of tasks is provided with more than one element, each task is
        executed sequentially and a list of results is returned.  A single-element
        list is unwrapped and treated identically to a non-list input.
        
        Print hierarchy (highest to lowest priority):
        1. UPSONIC_AGENT_PRINT env variable
        2. Agent constructor print parameter
        3. Method default (print_do_async=True)
        
        Args:
            task: Task to execute. Accepts a single Task/str or a list of them.
            model: Override model for this execution
            debug: Enable debug mode
            retry: Number of retries
            return_output: If True, return full AgentRunOutput. If False (default), return content only.
            
        Returns:
            Single task:
                Task content (str, BaseModel, etc.) if return_output=False
                Full AgentRunOutput if return_output=True
            List of tasks (len > 1):
                List of task contents if return_output=False
                List of AgentRunOutput if return_output=True
        """
        return await self.do_async(task, model, debug, retry, return_output, _print_method_default=True)

    def as_mcp(self, name: Optional[str] = None) -> "FastMCP":
        """
        Expose this agent as an MCP server.

        Creates a FastMCP server with a ``do`` tool that delegates task
        execution to this agent.  The returned server can be started with
        any transport (stdio, sse, streamable-http) via its ``.run()``
        method.

        Args:
            name: MCP server name. Defaults to the agent's name or
                  ``"Upsonic Agent"``.

        Returns:
            A :class:`fastmcp.FastMCP` server instance ready to ``.run()``.
        """
        from fastmcp import FastMCP as _FastMCP

        server_name: str = name or self.name or "Upsonic Agent"
        server: _FastMCP = _FastMCP(server_name)

        description_parts: List[str] = []
        if self.role:
            description_parts.append(f"Role: {self.role}")
        if self.goal:
            description_parts.append(f"Goal: {self.goal}")
        if self.instructions:
            description_parts.append(f"Instructions: {self.instructions}")

        tool_description: str = f"Execute a task using the {server_name} agent."
        if description_parts:
            tool_description += " " + " | ".join(description_parts)

        agent_ref: "Agent" = self

        @server.tool(description=tool_description)
        def do(task: str) -> str:
            """Give a task to this agent and get the result."""
            result: Any = agent_ref.print_do(task)
            if result is None:
                return ""
            return str(result)

        return server

    def stream(
        self,
        task: Union[str, "Task"],
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        events: bool = False,
        state: Optional["State"] = None,
        *,
        event: Optional[bool] = None,
    ) -> Iterator[Union[str, "AgentStreamEvent"]]:
        """
        Stream task execution synchronously - yields events/text as they arrive.
        
        For async streaming, use `astream()` instead.
        
        Args:
            task: Task to execute
            model: Override model for this execution
            debug: Enable debug mode
            retry: Number of retries
            events: If True, yield AgentEvent objects. If False (default), yield text chunks.
            state: Graph execution state
            event: Deprecated, use 'events' instead.
            
        Yields:
            AgentEvent if events=True, str if events=False
            
        Example:
            ```python
            # Stream text synchronously
            for text in agent.stream(task):
                print(text, end='', flush=True)
            
            # Stream events
            from upsonic.run.events import RunEvent
            for chunk in agent.stream(task, events=True):
                if chunk.event_kind == RunEvent.run_content:
                    print(chunk.content)
            ```
        """
        import queue
        import threading
        
        if event is not None:
            events = event
        
        result_queue: queue.Queue = queue.Queue()
        error_holder: List[Exception] = []
        
        async def stream_to_queue():
            try:
                async for item in self.astream(task, model, debug, retry, events, state):
                    result_queue.put(item)
            except Exception as e:
                error_holder.append(e)
            finally:
                result_queue.put(None)
        
        def run_async_stream():
            loop = _get_bg_loop()
            asyncio.run_coroutine_threadsafe(stream_to_queue(), loop).result()

        thread = threading.Thread(target=run_async_stream, daemon=True)
        thread.start()

        while True:
            item = result_queue.get()
            if item is None:
                if error_holder:
                    raise error_holder[0]
                break
            yield item
    
    async def astream(
        self,
        task: Union[str, "Task"],
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        events: bool = False,
        state: Optional["State"] = None,
        *,
        event: Optional[bool] = None,
    ) -> AsyncIterator[Union[str, "AgentStreamEvent"]]:
        """
        Stream task execution asynchronously - yields events or text as they arrive.
        
        Note: HITL (Human-in-the-Loop) features are not supported in streaming mode.
        Use do_async() for HITL functionality.
        
        Args:
            task: Task to execute
            model: Override model for this execution
            debug: Enable debug mode
            retry: Number of retries
            events: If True, yield AgentEvent objects. If False (default), yield text chunks.
            state: Graph execution state
            event: Deprecated, use 'events' instead.
            
        Yields:
            AgentEvent if events=True, str if events=False
            
        Example:
            ```python
            # Stream text
            async for text in agent.astream(task):
                print(text, end='', flush=True)
            
            # Stream events
            from upsonic.run.events import RunEvent
            async for evt in agent.astream(task, events=True):
                if evt.event_kind == RunEvent.run_content:
                    print(evt.content, end='')
            ```
        """
        if event is not None:
            events = event
        
        from upsonic.agent.pipeline import PipelineManager
        from upsonic.utils.printing import warning_log
        
        # Convert string to Task if needed
        task = self._convert_to_task(task)
        
        # Validate task state (completed/problematic checks)
        # For streaming, we just return early instead of yielding error output
        validation_error = self._validate_task_for_new_run(task, is_resuming=False)
        if validation_error is not None:
            # For streaming with problematic status, add specific warning
            if task.is_problematic:
                warning_log(
                    "Streaming does not support HITL continuation. Use continue_run_async() instead.",
                    "Agent"
                )
            return

        # Reset per-run state (same as do_async for fresh runs)
        import uuid as _uuid
        task.price_id_ = str(_uuid.uuid4())
        self._tool_call_count = 0
        self._tool_limit_reached = False
        self._last_built_system_prompt = None

        run_id = str(uuid.uuid4())
        self.run_id = run_id
        register_run(run_id)
        
        original_model: Optional["Model"] = None
        try:
            run_input = self._create_agent_run_input(task)
            
            self._agent_run_output = self._create_agent_run_output(
                run_id=run_id,
                task=task,
                run_input=run_input,
                is_streaming=True
            )
            
            self._agent_run_output.print_flag = self._resolve_print_flag(False)
            
            if task is not None:
                task.run_id = run_id
            
            original_model = self._apply_model_override(model)
            
            pipeline = PipelineManager(
                steps=self._create_streaming_pipeline_steps(),
                task=self._agent_run_output.task,
                agent=self,
                model=self.model,
                debug=debug or self.debug
            )
            
            _pl_start_time: float = time.time()
            _stream_succeeded: bool = False
            stream_task_desc: str = str(task.description) if task is not None and hasattr(task, "description") else ""
            with self._otel.agent_run_span(
                run_id,
                name=self.name or "",
                model=str(self.model_name) if self.model_name else "",
                task_description=stream_task_desc,
                user_id=self.user_id,
                session_id=self.session_id,
            ) as otel_span:
                try:
                    async for pipeline_event in pipeline.execute_stream(context=self._agent_run_output, start_step_index=0):
                        if events:
                            yield pipeline_event
                        else:
                            text_content = self._extract_text_from_stream_event(pipeline_event)
                            if text_content:
                                self._agent_run_output.accumulated_text += text_content
                                yield text_content
                    _stream_succeeded = True
                except Exception as stream_error:
                    self._otel.record_error(otel_span, stream_error)
                    raise stream_error
                finally:
                    self._otel.finalize_agent_run(
                        otel_span,
                        getattr(self, "_agent_run_output", None),
                        agent_user_id=self.user_id,
                        agent_session_id=self.session_id,
                        total_cost=self._calculate_aggregated_cost(),
                        tool_definitions=self._get_combined_tool_definitions(),
                    )
                    _trace_id = self._otel.extract_trace_id(otel_span)
                    if _trace_id and self._agent_run_output is not None:
                        self._agent_run_output.trace_id = _trace_id
                    if _stream_succeeded:
                        _stream_output: str = ""
                        if self._agent_run_output:
                            _stream_output = self._agent_run_output.accumulated_text or str(self._agent_run_output.output or "")
                        self._log_to_promptlayer_background(
                            task=task,
                            output=_stream_output,
                            start_time=_pl_start_time,
                            end_time=time.time(),
                        )

        finally:
            if original_model is not None:
                self.model = original_model
            stream_print_flag = self._resolve_print_flag(False)
            self._finalize_agent_usage(stream_print_flag)
            cleanup_run(run_id)
            self.run_id = None
    
    def _extract_text_from_stream_event(self, event: Any) -> Optional[str]:
        """Extract text content from a streaming event.
        
        Handles both Agent events (TextDeltaEvent) and raw LLM events
        (PartStartEvent, PartDeltaEvent).
        """
        from upsonic.messages import PartStartEvent, PartDeltaEvent, TextPart, TextPartDelta
        from upsonic.run.events.events import TextDeltaEvent
        
        # Handle Agent events (new event system)
        if isinstance(event, TextDeltaEvent):
            return event.content
        
        # Handle raw LLM events (legacy/internal)
        if isinstance(event, PartStartEvent) and isinstance(event.part, TextPart):
            return event.part.content
        elif isinstance(event, PartDeltaEvent):
            # Check if delta is a TextPartDelta specifically
            if isinstance(event.delta, TextPartDelta):
                return event.delta.content_delta
            # Fallback to hasattr check for compatibility
            elif hasattr(event.delta, 'content_delta'):
                return event.delta.content_delta
        return None
    
    
    # HITL checkpoint loading - delegates to Memory class
    
    async def _load_incomplete_run_data_from_storage(self, run_id: str) -> Optional["RunData"]:
        """
        Load a resumable RunData from storage. Delegates to Memory class.
        """
        if not self.memory:
            raise ValueError("No memory configured. Agent must have memory configured to load paused runs.")
        
        return await self.memory.load_resumable_run_async(run_id=run_id, session_type=self.session_type, agent_id=self.agent_id)
    
    async def _check_if_run_is_problematic(
        self,
        task: Optional["Task"],
        run_id: Optional[str]
    ) -> tuple[bool, Optional[RunStatus]]:
        """
        Check if a run is problematic (paused, cancelled, or error).
        
        Uses task.is_problematic if task provided, otherwise checks in-memory output,
        and only loads from storage if run_id is provided and no matching in-memory output.
        
        Args:
            task: Task instance (if provided, uses task.status directly - fast path)
            run_id: Run ID to check (if provided without task, may load from storage)
            
        Returns:
            Tuple of (is_problematic, run_status)
        """
        if task is not None:
            # Task provided - use task.status directly (fast path, no storage call)
            return task.is_problematic, task.status
        
        if run_id is not None:
            # Only run_id provided - check in-memory output first, then storage if needed
            _output = getattr(self, '_agent_run_output', None)
            if _output and _output.run_id == run_id:
                # In-memory output matches - use it
                return _output.is_problematic, _output.status
            else:
                # No matching in-memory output - must load from storage
                run_data = await self._load_incomplete_run_data_from_storage(run_id)
                if run_data and run_data.output:
                    return run_data.output.is_problematic, run_data.output.status
        
        return False, None
    
    async def _check_if_run_is_completed(
        self,
        task: Optional["Task"],
        run_id: Optional[str]
    ) -> tuple[bool, Optional[RunStatus]]:
        """
        Check if a run is already completed.
        
        Uses task.is_completed if task provided, otherwise checks in-memory output,
        and only loads from storage if run_id is provided and no matching in-memory output.
        
        Args:
            task: Task instance (if provided, uses task.status directly - fast path)
            run_id: Run ID to check (if provided without task, may load from storage)
            
        Returns:
            Tuple of (is_completed, run_status)
        """
        if task is not None:
            # Task provided - use task.is_completed directly (fast path, no storage call)
            return task.is_completed, task.status
        
        if run_id is not None:
            # Only run_id provided - check in-memory output first, then storage if needed
            _output = getattr(self, '_agent_run_output', None)
            if _output and _output.run_id == run_id:
                # In-memory output matches - use it
                return _output.is_complete, _output.status
            else:
                # No matching in-memory output - load ANY run from storage (not just resumable)
                if self.memory:
                    run_data = await self.memory.load_run_async(run_id=run_id, session_type=self.session_type, agent_id=self.agent_id)
                    if run_data and run_data.output:
                        return run_data.output.is_complete, run_data.output.status
        
        return False, None
    
    def _create_direct_pipeline_steps(self) -> List[Any]:
        """
        Create pipeline steps for direct call mode (do_async).
        
        Returns:
            List of all pipeline steps for direct execution
        """
        from upsonic.agent.pipeline import (
            InitializationStep, CacheCheckStep, UserPolicyStep,
            StorageConnectionStep, LLMManagerStep, ModelSelectionStep,
            ToolSetupStep,
            MemoryPrepareStep, SystemPromptBuildStep, ContextBuildStep,
            UserInputBuildStep, ChatHistoryStep, MessageAssemblyStep,
            CallManagerSetupStep,
            ModelExecutionStep, ResponseProcessingStep,
            ReflectionStep, CallManagementStep, TaskManagementStep,
            MemorySaveStep,
            ReliabilityStep, AgentPolicyStep,
            CacheStorageStep, FinalizationStep
        )

        return [
            InitializationStep(),          # 0
            StorageConnectionStep(),       # 1
            CacheCheckStep(),              # 2
            UserPolicyStep(),              # 3
            LLMManagerStep(),              # 4
            ModelSelectionStep(),          # 5
            ToolSetupStep(),               # 6
            MemoryPrepareStep(),           # 7  <-- Load memory (history, profile, metadata)
            SystemPromptBuildStep(),       # 8  <-- Build system prompt
            ContextBuildStep(),            # 9  <-- Build task context (KB/RAG)
            UserInputBuildStep(),          # 10 <-- Build user input (prompt + context + attachments)
            ChatHistoryStep(),             # 11 <-- Load chat history + mark run boundary
            MessageAssemblyStep(),         # 12 <-- Assemble ModelRequest + apply middleware
            CallManagerSetupStep(),        # 13 <-- Create and register CallManager
            ModelExecutionStep(),          # 14 <-- External tool resumes here
            ResponseProcessingStep(),      # 15 <-- Tracks messages to AgentRunOutput
            ReflectionStep(),              # 16
            TaskManagementStep(),          # 17
            ReliabilityStep(),             # 18
            AgentPolicyStep(),             # 19
            CacheStorageStep(),            # 20
            FinalizationStep(),            # 21
            MemorySaveStep(),              # 22
            CallManagementStep(),          # 23 <-- LAST: calls task_end() & prints metrics
        ]
    
    def _get_step_index_by_name(self, step_name: str, is_streaming: bool = False) -> int:
        """
        Get the index of a pipeline step by its name.
        
        This allows dynamic lookup of step indices instead of hardcoding,
        which is important for HITL resumption logic that needs to know
        which steps set _run_boundaries.
        
        Args:
            step_name: The name of the step (e.g., "chat_history", "model_execution")
            is_streaming: Whether to use streaming pipeline (default: direct)
            
        Returns:
            The index of the step in the pipeline
            
        Raises:
            ValueError: If the step is not found
        """
        steps = self._create_streaming_pipeline_steps() if is_streaming else self._create_direct_pipeline_steps()
        
        for i, step in enumerate(steps):
            if step.name == step_name:
                return i
        
        raise ValueError(f"Step '{step_name}' not found in {'streaming' if is_streaming else 'direct'} pipeline")
    
    def _create_streaming_pipeline_steps(self) -> List[Any]:
        """
        Create pipeline steps for streaming mode (stream).
        
        Returns:
            List of all pipeline steps for streaming execution
        """
        from upsonic.agent.pipeline import (
            InitializationStep, CacheCheckStep, UserPolicyStep,
            StorageConnectionStep, LLMManagerStep, ModelSelectionStep,
            ToolSetupStep,
            MemoryPrepareStep, SystemPromptBuildStep, ContextBuildStep,
            UserInputBuildStep, ChatHistoryStep, MessageAssemblyStep,
            CallManagerSetupStep,
            StreamModelExecutionStep,
            ReflectionStep, ReliabilityStep,
            AgentPolicyStep, CacheStorageStep,
            StreamFinalizationStep, CallManagementStep,
            StreamMemoryMessageTrackingStep
        )

        return [
            InitializationStep(),              # 0
            StorageConnectionStep(),           # 1
            CacheCheckStep(),                  # 2
            UserPolicyStep(),                  # 3
            LLMManagerStep(),                  # 4
            ModelSelectionStep(),              # 5
            ToolSetupStep(),                   # 6
            MemoryPrepareStep(),               # 7  <-- Load memory
            SystemPromptBuildStep(),           # 8  <-- Build system prompt
            ContextBuildStep(),                # 9  <-- Build task context (KB/RAG)
            UserInputBuildStep(),              # 10 <-- Build user input
            ChatHistoryStep(),                 # 11 <-- Load chat history + mark run boundary
            MessageAssemblyStep(),             # 12 <-- Assemble ModelRequest + apply middleware
            CallManagerSetupStep(),            # 13 <-- Create and register CallManager
            StreamModelExecutionStep(),        # 14 <-- Streaming model execution
            ReflectionStep(),                  # 15 <-- Improve output quality
            ReliabilityStep(),                 # 16 <-- Verify and clean output
            AgentPolicyStep(),                 # 17
            CacheStorageStep(),                # 18
            StreamFinalizationStep(),          # 19
            StreamMemoryMessageTrackingStep(), # 20 <-- Saves AgentSession + task_end()
            CallManagementStep(),              # 21 <-- LAST: Records usage to price_id_summary
        ]
    
    def _create_full_pipeline_steps(self, is_streaming: bool = False) -> List[Any]:
        """
        Create complete pipeline steps based on execution mode.
        
        Args:
            is_streaming: If True, return streaming pipeline steps.
                         If False, return direct call pipeline steps.
        
        Returns:
            List of all pipeline steps in order
        """
        if is_streaming:
            return self._create_streaming_pipeline_steps()
        return self._create_direct_pipeline_steps()
    
    async def _inject_external_tool_results(
        self, 
        output: AgentRunOutput, 
        requirements: list
    ) -> None:
        """Inject external tool results into output chat_history."""
        await self._inject_hitl_results(output, requirements)

    async def _inject_hitl_results(
        self,
        output: AgentRunOutput,
        requirements: list,
    ) -> None:
        """
        Inject resolved HITL requirement results into output chat_history.

        Handles external execution, confirmation, and user input requirements.
        For confirmation (rejected), injects a rejection message.
        For confirmation (approved) and user input, executes the tool and injects the result.
        """
        from upsonic.messages import ModelRequest, ToolReturnPart
        from upsonic._utils import now_utc

        if not requirements:
            return

        if output.response:
            if not output.chat_history or output.chat_history[-1] != output.response:
                output.chat_history.append(output.response)

        if output.tools is None:
            output.tools = []

        tool_return_parts: list = []
        for requirement in requirements:
            te = requirement.tool_execution
            if not te:
                continue
            if te.result_injected:
                continue

            result_content: Optional[str] = None

            # --- External execution: result already set by user ---
            if te.external_execution_required and te.result is not None:
                result_content = te.result

            # --- Confirmation ---
            elif te.requires_confirmation and requirement.confirmation is not None:
                if requirement.confirmation:
                    import time as _time
                    _tool_start = _time.time()
                    result_content = await self._execute_confirmed_tool(te)
                    _tool_elapsed = _time.time() - _tool_start
                    output.add_tool_execution_time(_tool_elapsed)
                else:
                    note = requirement.confirmation_note or "Tool execution rejected by user."
                    result_content = f"Tool execution rejected: {note}"

            # --- User input ---
            elif te.requires_user_input and te.answered:
                import time as _time
                _tool_start = _time.time()
                result_content = await self._execute_user_input_tool(te, requirement)
                _tool_elapsed = _time.time() - _tool_start
                output.add_tool_execution_time(_tool_elapsed)

            if result_content is not None:
                te.result = result_content
                te.result_injected = True
                tool_return_parts.append(ToolReturnPart(
                    tool_name=te.tool_name,
                    content=result_content,
                    tool_call_id=te.tool_call_id,
                    timestamp=now_utc(),
                ))
                output.tools.append(te)
                self._tool_call_count += 1
                output.increment_tool_calls(1)

        if tool_return_parts:
            output.chat_history.append(ModelRequest(parts=tool_return_parts))

    async def _execute_hitl_tool_directly(
        self,
        tool_name: str,
        tool_args: Dict[str, Any],
    ) -> str:
        """Execute a tool bypassing HITL checks (for confirmed/user-input-answered tools).
        
        Calls the underlying function directly to avoid re-triggering pause exceptions.
        HITL pause exceptions are NOT caught — they propagate to the caller.
        """
        import asyncio
        from upsonic.tools.processor import (
            ExternalExecutionPause as _hitl_ExternalExecutionPause,
            ConfirmationPause as _hitl_ConfirmationPause,
            UserInputPause as _hitl_UserInputPause,
        )

        try:
            manager = self._resolve_tool_manager(tool_name)
        except ValueError:
            return f"Error: Tool '{tool_name}' not found in any ToolManager"

        tool_obj = manager.processor.registered_tools.get(tool_name)
        if not tool_obj:
            return f"Error: Tool '{tool_name}' not registered"

        try:
            if hasattr(tool_obj, 'function'):
                func = tool_obj.function
                if asyncio.iscoroutinefunction(func):
                    result = await func(**tool_args)
                else:
                    loop = asyncio.get_running_loop()
                    result = await loop.run_in_executor(None, lambda: func(**tool_args))
            else:
                result = await tool_obj.execute(**tool_args)
            return str(result) if result is not None else ""
        except (
            _hitl_ExternalExecutionPause,
            _hitl_ConfirmationPause,
            _hitl_UserInputPause,
            TypeError,
        ):
            raise
        except Exception as exc:
            return f"Error executing tool {tool_name}: {exc}"

    async def _execute_confirmed_tool(
        self,
        te: "ToolExecution",
    ) -> str:
        """Execute a tool that was confirmed by the user."""
        if not te.tool_name or te.tool_args is None:
            return "Error: Missing tool name or arguments for confirmed tool"
        return await self._execute_hitl_tool_directly(te.tool_name, te.tool_args)

    async def _execute_user_input_tool(
        self,
        te: "ToolExecution",
        requirement: "RunRequirement",
    ) -> str:
        """Execute a tool after user has provided input values.

        For static user-input tools (``@tool(requires_user_input=True)``),
        the underlying function is re-executed with the merged arguments
        (agent-provided + user-provided).

        For dynamic user-input tools (e.g. ``get_user_input`` from
        ``UserControlFlowTools``), re-executing would raise ``UserInputPause``
        again.  In that case the user's answers are formatted and returned
        directly as the tool result so the agent can proceed.
        """
        if not te.tool_name:
            return "Error: Missing tool name for user input tool"

        merged_args: Dict[str, Any] = dict(te.tool_args or {})
        user_provided: Dict[str, str] = {}
        if requirement.user_input_schema:
            for field_dict in requirement.user_input_schema:
                if isinstance(field_dict, dict) and field_dict.get("value") is not None:
                    user_provided[field_dict["name"]] = field_dict["value"]
            merged_args.update(user_provided)

        from upsonic.tools.processor import UserInputPause
        try:
            return await self._execute_hitl_tool_directly(te.tool_name, merged_args)
        except (UserInputPause, TypeError):
            return self._format_user_input_result(user_provided)

    @staticmethod
    def _format_user_input_result(user_provided: Dict[str, str]) -> str:
        """Format user-provided field values into a tool-result string."""
        if not user_provided:
            return "User input received."
        parts: list = [f"{k}: {v}" for k, v in user_provided.items()]
        return "User provided input — " + ", ".join(parts)
    
    # HITL continuation support
    
    def continue_run(
        self,
        task: Optional["Task"] = None,
        run_id: Optional[str] = None,
        requirements: Optional[List["RunRequirement"]] = None,
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        return_output: bool = False,
        *,
        streaming: Optional[bool] = None,
        event: bool = False,
        hitl_handler: Optional[Callable[["RunRequirement"], None]] = None,
    ) -> Any:
        """
        Continue a paused agent run (synchronous wrapper).
        
        Automatically detects if the original run was streaming and continues
        in the same mode, or you can override with the streaming parameter.
        
        Supports all HITL continuation scenarios:
        1. External tool execution: Pass task object with external results filled
        2. User confirmation: Approve or reject tool calls
        3. User input: Provide field values for tool calls
        4. Durable execution (error recovery): Pass run_id to load from storage
        5. Cancel run resumption: Pass run_id to load from storage
        
        Args:
            task: Task object (for HITL continuation with results)
            run_id: Run ID to load from storage (for durable/cancel)
            requirements: Resolved requirements from a previous pause
            model: Override model
            debug: Enable debug mode
            retry: Number of retries
            return_output: If True, return full AgentRunOutput. If False (default), return content only.
            streaming: If True, return list of events/text. If False, return result. 
                      If None (default), auto-detect from original run.
            event: If True (with streaming), return list of AgentEvent objects.
                   If False (with streaming), return list of text chunks.
            hitl_handler: Unified HITL handler that resolves any paused requirement.
                Called for each active RunRequirement when the agent pauses again.
                The handler must mutate the requirement in-place:
                - External execution: set requirement.tool_execution.result
                - Confirmation: call requirement.confirm() or requirement.reject()
                - User input: fill requirement.user_input_schema field values
                Signature: (requirement: RunRequirement) -> None
            
        Returns:
            - For direct mode: Task content if return_output=False, AgentRunOutput if return_output=True
            - For streaming mode: List of events (if event=True) or text chunks (if event=False)
            
        Example:
            # Force direct mode
            result = agent.continue_run(run_id=result.run_id, streaming=False, return_output=True)
            
            # With unified HITL handler
            def my_handler(req):
                if req.needs_external_execution:
                    req.tool_execution.result = execute_my_tool(req.tool_execution.tool_args)
                elif req.needs_confirmation:
                    req.confirm()
                elif req.needs_user_input:
                    for field in req.user_input_schema:
                        field["value"] = get_value_for(field["name"])
            result = agent.continue_run(run_id=result.run_id, hitl_handler=my_handler)
        """
        if not task and not run_id:
            raise ValueError("Either 'task' or 'run_id' must be provided")
        
        use_streaming = streaming
        if use_streaming is None:
            _output = getattr(self, '_agent_run_output', None)
            if _output:
                use_streaming = _output.is_streaming
            else:
                use_streaming = False
        
        if use_streaming:
            async def collect_stream():
                results = []
                async_gen = await self.continue_run_async(
                    task, run_id, requirements, model, debug, retry, return_output,
                    streaming=True,
                    event=event,
                    hitl_handler=hitl_handler,
                )
                async for item in async_gen:
                    results.append(item)
                return results
            
            return _run_in_bg_loop(collect_stream())
        else:
            return _run_in_bg_loop(
                self.continue_run_async(
                    task, run_id, requirements, model, debug, retry, return_output,
                    streaming=False,
                    event=event,
                    hitl_handler=hitl_handler,
                )
            )
    
    async def _prepare_continuation_context(
        self,
        task: Optional["Task"],
        run_id: Optional[str],
        model: Optional[Union[str, "Model"]],
        debug: bool,
        requirements: Optional[List["RunRequirement"]] = None,
    ) -> tuple:
        """
        Prepare output context and determine resume point for continue_run_async.
        
        Status-based processing:
        - PAUSED status: External tool execution - inject tool results
        - ERROR/CANCELLED status: Resume from the problematic step
        
        All state needed for resumption is in AgentRunOutput (chat_history, step_results, etc.)
        
        Args:
            requirements: Optional list of RunRequirement with tool results set.
                         For new agents, pass the requirements from original output
                         (with results set) to copy results to the loaded state.
            
        Returns:
            Tuple of (output, task, resume_step_index)
        """
        if not task and not run_id:
            raise ValueError("Either 'task' or 'run_id' must be provided")
        
        output: Optional[AgentRunOutput] = None
        
        # Step 1: Determine if we should load from storage
        need_storage_load = False
        
        # Get in-memory output (may not exist for new agents)
        _in_memory_output = getattr(self, '_agent_run_output', None)
        
        if run_id:
            # run_id provided - check if we already have matching in-memory output
            if _in_memory_output:
                in_memory_run_id = _in_memory_output.run_id
                in_memory_agent_id = _in_memory_output.agent_id
                
                if in_memory_run_id != run_id:
                    need_storage_load = True
                elif in_memory_agent_id != self.agent_id:
                    need_storage_load = True
                else:
                    output = _in_memory_output
            else:
                need_storage_load = True
        else:
            # No run_id provided - use in-memory output if available
            if _in_memory_output:
                in_memory_agent_id = _in_memory_output.agent_id
                if in_memory_agent_id != self.agent_id:
                    # Different agent - try to load from storage using task.run_id
                    if task and hasattr(task, 'run_id') and task.run_id:
                        need_storage_load = True
                        run_id = task.run_id
                    else:
                        raise ValueError(
                            "In-memory output belongs to different agent. "
                            "Provide run_id to load from storage."
                        )
                else:
                    output = _in_memory_output
            else:
                # No in-memory checkpoint - try to load from storage using task.run_id
                if task and hasattr(task, 'run_id') and task.run_id:
                    need_storage_load = True
                    run_id = task.run_id
                else:
                    raise ValueError(
                        "No in-memory checkpoint found. Provide run_id to load from storage."
                    )
        
        # Step 2: Load from storage if needed
        if need_storage_load:
            run_data = await self._load_incomplete_run_data_from_storage(run_id)
            if not run_data:
                raise ValueError(f"No resumable run found with run_id: {run_id}")
            
            if not run_data.output:
                raise ValueError("Run has no output (checkpoint not found)")
            
            self._agent_run_output = run_data.output
            output = self._agent_run_output
            
            # Step 2b: Copy HITL state from passed requirements to loaded output
            # This is needed for new agents - they load state from storage (without results)
            # and the user passes requirements (with resolved state) to inject
            if requirements:
                passed_map: Dict[str, "RunRequirement"] = {}
                for req in requirements:
                    if req.tool_execution and req.tool_execution.tool_call_id:
                        passed_map[req.tool_execution.tool_call_id] = req
                
                for loaded_req in output.requirements or []:
                    if not loaded_req.tool_execution or not loaded_req.tool_execution.tool_call_id:
                        continue
                    tcid = loaded_req.tool_execution.tool_call_id
                    passed_req = passed_map.get(tcid)
                    if not passed_req:
                        continue

                    # External execution result
                    if passed_req.tool_execution and passed_req.tool_execution.result is not None:
                        loaded_req.tool_execution.result = passed_req.tool_execution.result

                    # Confirmation state
                    if passed_req.confirmation is not None:
                        loaded_req.confirmation = passed_req.confirmation
                        loaded_req.confirmation_note = passed_req.confirmation_note
                        if loaded_req.tool_execution:
                            loaded_req.tool_execution.confirmed = passed_req.confirmation

                    # User input state
                    if passed_req.user_input_schema is not None:
                        loaded_req.user_input_schema = passed_req.user_input_schema
                    if passed_req.tool_execution and passed_req.tool_execution.answered:
                        if loaded_req.tool_execution:
                            loaded_req.tool_execution.answered = True
            
        # Step 3: Validate we have what we need
        if output is None:
            raise ValueError(
                "No checkpoint found. Either provide run_id to load from storage, "
                "or call continue_run_async while agent still has in-memory output."
            )
        
        # Get task from output if not provided
        if task is None:
            task = output.task
        self.current_task = task
        if task is None:
            raise ValueError("Cannot extract task from checkpoint")

        from upsonic.usage import TaskUsage
        if getattr(task, "_usage", None) is None:
            # Restore usage from the stored output if available (cross-process resume),
            # otherwise create a fresh TaskUsage.
            stored_usage = getattr(output, 'usage', None)
            if stored_usage is not None:
                if isinstance(stored_usage, dict):
                    task._usage = TaskUsage.from_dict(stored_usage)
                elif isinstance(stored_usage, TaskUsage):
                    task._usage = stored_usage
                else:
                    task._usage = TaskUsage()
            else:
                task._usage = TaskUsage()
        output.usage = task._usage

        run_status = output.status
        
        # Centralized: All problematic runs use get_problematic_step() to find resume point
        if run_status not in (RunStatus.paused, RunStatus.error, RunStatus.cancelled):
            raise ValueError(
                f"Cannot continue run with status '{run_status}'. "
                "Only paused, error, or cancelled runs can be continued."
            )
        
        # Get the problematic step for resume point (works for paused, error, cancelled)
        problematic_step = output.get_problematic_step()
        if not problematic_step:
            raise ValueError(f"Run has {run_status.value} status but no problematic step found in output")
        
        resume_step_index = problematic_step.step_number
        
        # ChatHistoryStep sets _run_boundaries when loading history.
        # Only call start_new_run() if we're resuming AFTER ChatHistoryStep.
        # If we resume AT or BEFORE ChatHistoryStep, it will rebuild chat_history and
        # set the correct boundary itself.
        chat_history_step_index: int = self._get_step_index_by_name("chat_history")
        if resume_step_index > chat_history_step_index:
            # CRITICAL: Mark the start of this resumed run for message tracking BEFORE any modifications.
            # When resuming AFTER ChatHistoryStep, it's skipped (we resume from a later step),
            # so we need to record the current chat_history length here.
            # This ensures finalize_run_messages() includes ALL messages from THIS resumed run,
            # including injected tool results.
            output.start_new_run()
        
        # For paused runs, inject HITL results (external tool, confirmation, user input)
        if run_status == RunStatus.paused:
            resolved_reqs = [
                r for r in (output.requirements or [])
                if r.is_resolved and r.tool_execution is not None
            ]
            
            if not resolved_reqs:
                raise ValueError(
                    "Run is paused but no resolved requirements found. "
                    "For external tools: set result via requirement.tool_execution.result = ... "
                    "For confirmation: call requirement.confirm() or requirement.reject(). "
                    "For user input: fill requirement.user_input_schema field values and set "
                    "requirement.tool_execution.answered = True."
                )
            
            await self._inject_hitl_results(output, resolved_reqs)
        
        # Clear paused state and set up for continuation
        task.is_paused = False
        output.task = task

        # Restore agent-level tool call count from the output so that
        # tool_call_limit checks work correctly across HITL resume
        # (especially important for cross-process resume where agent.__init__
        # resets _tool_call_count to 0).
        self._tool_call_count = getattr(output, 'tool_call_count', 0)
        self._tool_limit_reached = False

        if task.enable_cache:
            task.set_cache_manager(self._cache_manager)

        return output, task, resume_step_index
    
    async def continue_run_async(
        self,
        task: Optional["Task"] = None,
        run_id: Optional[str] = None,
        requirements: Optional[List["RunRequirement"]] = None,
        model: Optional[Union[str, "Model"]] = None,
        debug: bool = False,
        retry: int = 1,
        return_output: bool = False,
        state: Optional["State"] = None,
        *,
        streaming: bool = False,
        event: bool = False,
        hitl_handler: Optional[Callable[["RunRequirement"], None]] = None,
        graph_execution_id: Optional[str] = None
    ) -> Any:
        """
        Continue a paused agent run using StepResult-based intelligent resumption.
        
        Note: HITL continuation is only supported in direct call mode (streaming=False).
        
        Supports HITL continuation for:
        - External tool execution: inject tool results and resume
        - User confirmation: inject confirmed/rejected status and resume
        - User input: inject user-provided field values and resume
        - ERROR/CANCELLED status: Resume from the problematic step
        
        Args:
            hitl_handler: Unified HITL handler that resolves any paused requirement.
                Called for each active RunRequirement when the agent pauses again.
                The handler must mutate the requirement in-place.
        """
        from upsonic.utils.printing import info_log, warning_log
        
        # HITL continuation is only supported in direct call mode
        if streaming:
            raise ValueError(
                "Streaming mode is not supported for HITL continuation. "
                "Use streaming=False (default) for continue_run_async."
            )
        
        # Check if task is already completed - cannot continue a completed task
        is_completed, run_status = await self._check_if_run_is_completed(task, run_id)
        if is_completed:
            resolved_run_id = (task.run_id if task else run_id) or "unknown"
            warning_log(
                f"Task is already completed (run_id={resolved_run_id}). Cannot continue a completed task.",
                "Agent"
            )
            completed_output = AgentRunOutput(
                run_id=resolved_run_id,
                agent_id=self.agent_id,
                agent_name=self.name,
                session_id=self.session_id,
                user_id=self.user_id,
                status=RunStatus.completed,
                output=f"Task is already completed (run_id={resolved_run_id}). Cannot continue a completed task.",
            )
            if return_output:
                return completed_output
            return completed_output.output
        
        # Check if run is problematic (paused, cancelled, error) before preparing continuation
        is_problematic, run_status = await self._check_if_run_is_problematic(task, run_id)
        
        if not is_problematic:
            # Run is not problematic - log and call do_async for a fresh run
            status_str = run_status.value if run_status else "unknown"
            info_log(
                f"Run status is '{status_str}' (not paused, cancelled, or error). "
                f"Starting fresh run via do_async.",
                "Agent"
            )
            return await self.do_async(
                task=task,
                model=model,
                debug=debug,
                retry=retry,
                return_output=return_output,
                state=state,
                graph_execution_id=graph_execution_id,
            )
        
        # Prepare context and determine resume point for problematic runs
        output, task, resume_step_index = await self._prepare_continuation_context(
            task, run_id, model, debug, requirements
        )
        
        return await self._continue_run_direct_impl(
            task, model, debug, retry, return_output, output, resume_step_index, hitl_handler
        )
    
    async def _continue_run_direct_impl(
        self,
        task: "Task",
        model: Optional[Union[str, "Model"]],
        debug: bool,
        retry: int,
        return_output: bool,
        output: AgentRunOutput,
        resume_step_index: int,
        hitl_handler: Optional[Callable[["RunRequirement"], None]] = None,
    ) -> Any:
        """
        Internal direct call implementation for continue_run_async.
        
        Handles the loop for sequential HITL interactions (external tools,
        confirmation, user input) automatically when a hitl_handler is provided.
        The handler is called for every active requirement regardless of type.
        """
        max_rounds = 10
        rounds = 0
        result = None
        original_print_flag: bool = getattr(output, 'print_flag', False)

        task._usage.start_timer()
        
        while rounds < max_rounds:
            rounds += 1
            
            result = await self.do_async(
                task,
                model=model,
                debug=debug,
                retry=retry,
                return_output=True,
                _resume_output=output,
                _resume_step_index=resume_step_index,
                _print_method_default=original_print_flag,
            )
            
            if result.is_complete:
                if return_output:
                    return result
                return result.output if hasattr(result, 'output') else result
            
            active_reqs = result.active_requirements
            if not active_reqs:
                if return_output:
                    return result
                return result.output if hasattr(result, 'output') else result

            if not hitl_handler:
                if return_output:
                    return result
                return result.output if hasattr(result, 'output') else result

            for requirement in active_reqs:
                hitl_handler(requirement)

            output = getattr(self, '_agent_run_output', None)
            problematic_step = output.get_problematic_step()
            if problematic_step:
                resume_step_index = problematic_step.step_number
            
            task.is_paused = False
            
            resolved_reqs = [
                r for r in (output.requirements or [])
                if r.tool_execution and r.tool_execution.result is not None
            ]
            confirmed_or_input_reqs = [
                r for r in (output.requirements or [])
                if (r.tool_execution and r.tool_execution.requires_confirmation and r.confirmation is not None)
                or (r.tool_execution and r.tool_execution.requires_user_input and r.tool_execution.answered)
            ]
            all_to_inject = resolved_reqs + [r for r in confirmed_or_input_reqs if r not in resolved_reqs]

            if all_to_inject:
                await self._inject_hitl_results(output, all_to_inject)
        
        if return_output:
            return result
        return result.output if hasattr(result, 'output') else result



Clanker = Agent
