"""Agent runner using PydanticAI.

Runs the LLM agent loop: takes PersonaSpec + task, calls LLM, manages tools.
Returns Result[T, E].
"""

# @invar:allow shell_result: Agent runner returns Result types
# @shell_complexity: Agent loop requires state machine branching
# @invar:allow forbidden_import: Shell module needs os for PID, pathlib for paths
# @invar:allow file_size: Agent runner is cohesive orchestrator; splitting would increase coupling
# @invar:allow function_size: AgentRunner._run_agent_loop keeps turn accounting and lifecycle completion decisions in one control loop

from __future__ import annotations

import asyncio
import os
from typing import TYPE_CHECKING, Any

from pydantic_ai import Agent
from pydantic_ai.exceptions import ModelHTTPError, UserError as PAgentUserError
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import UsageLimits
from returns.result import Failure, Result, Success

from anima.core.budget import grace_system_message, resolve_budget
from anima.core.lifecycle import LifecycleEvent
from anima.core.prune import assemble_tool_output_history
from anima.core.schemas import (
    AgentResult,
    AgentState,
    AnimaError,
    CompactionConfig,
    CapabilityToken,
    CompletionPolicy,
    InternalError,
    NotFoundError,
    JobSpec,
    LifecycleType,
    PersonaSpec,
    TokenUsage,
    AnimaTimeoutError,
)
from anima.core.tracker import BudgetTracker
from anima.shell.agent_compaction import run_compaction, should_compact
from anima.shell.agent_budget import (
    check_budget_exhaustion,
    handle_model_request_turn,
    update_state as budget_update_state,
)
from anima.shell.agent_factory import create_agent
from anima.shell.agent_helpers import (
    collect_mcp_server_context,
    log_ctl_tool_registration,
)
from anima.shell.agent_interact import (
    consume_stop_turn_request,
    drain_phase1_pending_messages,
    poll_pending_message,
    process_interact_message,
    read_current_state,
)
from anima.shell.agent_lifecycle import (
    configure_run_environment,
    finalize_run,
    on_timeout,
    setup_signal_handlers,
    setup_timeout_watchdog,
)
from anima.shell.agent_response import (
    check_auto_complete,
    extract_last_response_text,
    record_tool_usage,
)
from anima.shell.ctl_tools import create_ctl_tools
from anima.shell.model_config import (
    DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS,
    DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES,
    DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES,
    evaluate_runtime_model_context_limit,
    load_compaction_config,
    load_global_config,
    resolve_tool_output_budget_profile,
    resolve_model_context_limit,
)
from anima.shell.mcp_toolset_policy import (
    begin_tool_output_budget_scope,
    end_tool_output_budget_scope,
)
from anima.shell.mcp_agent import AgentMCPServer
from anima.shell.state import (
    append_log,
    read_contract_snapshot,
    read_state,
    transition_state,
    update_agent_state,
    write_agent_pid,
)
from anima.shell.timeout_watchdog import AsyncTimeoutWatchdog

if TYPE_CHECKING:
    from anima.core.sandbox import SandboxPolicy


class AgentRunner:
    """Agent runner that manages the LLM loop and MCP server.

    Each agent process creates one AgentRunner to:
    1. Run the agent transition loop
    2. Start AgentMCPServer for external control
    3. Handle graceful shutdown
    4. Track turns and token usage

    Attributes:
        instance_id: The agent's unique identifier (directory key for state).
        task: The task string for the agent.
        persona: The PersonaSpec for the agent.
        mcp_server: The AgentMCPServer instance (created during run).
        turn_count: Current turn number in the agent loop.
        token_usage: Cumulative token usage across all turns.
        completion_policy: How completion is triggered (explicit/auto).
        max_turns: Maximum turns allowed (None = unlimited).
        timeout: Wall-clock timeout string (e.g., '10m', '2h') or None.
        no_tool_streak: Consecutive turns without tool calls (for auto-completion).
        grace_message_injected: Whether grace system message has been injected.
        timeout_watchdog: AsyncTimeoutWatchdog instance for timeout enforcement.
    """

    def __init__(
        self,
        instance_id: str,
        task: str,
        persona: PersonaSpec | None = None,
        job: JobSpec | None = None,
        completion_policy: CompletionPolicy = CompletionPolicy.explicit,
        max_turns: int | None = None,
        model_override: str | None = None,
        model_instance: Any | None = None,
        base_url: str | None = None,
        memory_provider_name: str | None = None,
        timeout: str | None = None,
        mcp_servers: list[Any] | None = None,
        capability_token: CapabilityToken | None = None,
        sandbox_policy: "SandboxPolicy | None" = None,
        **legacy_kwargs: Any,
    ) -> None:
        """Initialize the agent runner.

        Args:
            instance_id: Unique identifier for this agent (directory key for state).
            task: The task string for the agent to execute.
            persona: Optional PersonaSpec (loaded from file if not provided).
            job: Optional JobSpec (for budget and constraints).
            completion_policy: How completion is triggered.
            max_turns: Maximum turns allowed (None = unlimited).
            model_override: Optional model override from CLI ``--model``.
            model_instance: Optional pre-constructed model instance (e.g., TestModel).
                If provided, bypasses model resolution in create_agent.
            base_url: Optional provider base URL override for model construction.
            memory_provider_name: Optional provider selector ('none'|'local'|'mcp').
            timeout: Optional wall-clock timeout string (e.g., '10m', '2h').
            mcp_servers: Optional list of MCP server toolsets to attach.
            capability_token: Optional CapabilityToken loaded from --token-file.
            sandbox_policy: Optional sandbox policy for tool execution.
        """
        if legacy_kwargs:
            unknown_arg = next(iter(legacy_kwargs))
            raise TypeError(f"Unexpected keyword argument: {unknown_arg}")

        self.instance_id = instance_id
        self.task = task
        self.persona = persona
        self.job = job
        self.completion_policy = completion_policy
        self.max_turns = max_turns
        self.model_override = model_override
        self.model_instance = model_instance
        self.base_url = base_url
        self.memory_provider_name = memory_provider_name
        self.timeout = timeout
        # Sanitize MCP tool names for LLM provider compatibility (e.g., dots → underscores)
        from anima.shell.agent_factory import _sanitize_mcp_tool_names

        self.mcp_servers = (
            _sanitize_mcp_tool_names(mcp_servers) if mcp_servers else mcp_servers
        )
        self.capability_token = capability_token
        self.sandbox_policy = sandbox_policy
        self.mcp_server: AgentMCPServer | None = None
        self.timeout_watchdog: AsyncTimeoutWatchdog | None = None
        self._interrupted: bool = False
        self._timeout_fired: bool = False
        # Spec source: event-mechanisms (plan.yaml) instant wake channels.
        self._approval_event: asyncio.Event = asyncio.Event()
        self._interactive_approval_mode: bool = False
        self._interactive_approval_future: (
            asyncio.Future[dict[str, bool | str | None]] | None
        ) = None
        self._resume_event: asyncio.Event = asyncio.Event()
        self.turn_count: int = 0
        self.token_usage: TokenUsage = TokenUsage()
        self.no_tool_streak: int = 0
        self._inject_followup_hint_next_message: bool = False

        # Budget tracking
        job_for_budget = job or JobSpec(
            persona_ref="default@v1",
            instance_id=instance_id,
        )
        budget_limit = resolve_budget(
            cli_max_tokens=None,
            job=job_for_budget,
        )
        grace_budget = job_for_budget.finalize_grace_budget if job else 500
        self.budget_tracker = BudgetTracker(budget_limit, grace_budget)

        # Grace window state
        self.in_grace_window: bool = False
        self.grace_exhausted_without_finish: bool = False
        self.grace_message_injected: bool = False
        self.max_turns_exceeded: bool = False

        # Previous RunUsage snapshot for delta computation.
        # PydanticAI's run.usage() returns cumulative values within an iter() call.
        # We must compute deltas to avoid double-counting tokens.
        self._prev_run_input_tokens: int = 0
        self._prev_run_output_tokens: int = 0

        # Last response text from the LLM, updated after each model response.
        # Used by ctl_finish wrapper to populate result.json with response_text
        # on the first write, eliminating the race where result.json lacks it.
        self.last_response_text: str | None = None

        # LLM call timeout (loaded from config at run time)
        # Spec source: talk-ux-fixes.llm-call-timeout
        self._llm_timeout: int | None = None

        # Spec source: context-compaction.wire-compaction-to-agent-loop
        self._compaction_config: CompactionConfig = CompactionConfig()
        # Spec source: context-compaction-remediation.spec-clarify-context-limit-source
        self._model_context_limit: int = DEFAULT_MODEL_CONTEXT_LIMIT_TOKENS
        self._compaction_pending: bool = False
        self._tool_output_per_call_cap_bytes: int = (
            DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES
        )
        self._tool_output_per_turn_cap_bytes: int = (
            DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES
        )

    _SYSTEM_REMINDER_TEMPLATE = (
        "<system-reminder>\n"
        "The user sent a new message while you were working:\n"
        "{original_message}\n\n"
        "IMPORTANT: After completing your current task, you MUST address "
        "the user's message above. Do not ignore it.\n"
        "</system-reminder>"
    )

    def _wrap_followup_message_with_system_reminder(self, message: str) -> str:
        """Wrap queued follow-up message in OpenCode system-reminder tags."""
        return self._SYSTEM_REMINDER_TEMPLATE.format(original_message=message)

    def _queue_followup_hint_if_pending_messages(self) -> None:
        """Queue one follow-up context hint when additional messages are pending."""
        if self.mcp_server is None:
            return
        get_pending_message = getattr(self.mcp_server, "get_pending_message", None)
        if not callable(get_pending_message):
            return
        queue_obj = get_pending_message()
        if not isinstance(queue_obj, asyncio.Queue):
            return
        if queue_obj.qsize() > 0:
            self._inject_followup_hint_next_message = True

    def _inject_followup_hint_if_needed(self, message: str) -> str:
        """Wrap the next dequeued message in a system-reminder block when flagged."""
        if not self._inject_followup_hint_next_message:
            return message
        self._inject_followup_hint_next_message = False
        return self._wrap_followup_message_with_system_reminder(message)

    def _get_model_settings_with_timeout(self) -> ModelSettings | None:
        """Get ModelSettings with per-turn LLM timeout if configured."""
        if self._llm_timeout is None:
            return None
        return ModelSettings(timeout=float(self._llm_timeout))

    def _setup_signal_handlers(self) -> None:
        """Set up signal handlers for graceful shutdown."""
        setup_signal_handlers(
            interrupted_setter=lambda: setattr(self, "_interrupted", True),
            mcp_server_getter=lambda: self.mcp_server,
        )

    def _on_timeout(self, instance_id: str) -> None:
        """Handle timeout by transitioning agent to failed state."""
        on_timeout(
            instance_id,
            timeout_fired_setter=lambda: setattr(self, "_timeout_fired", True),
            interrupted_setter=lambda: setattr(self, "_interrupted", True),
            mcp_server=self.mcp_server,
        )

    async def _initialize_run_state(self) -> Result[Any, AnimaError]:
        """Initialize persisted run state, PID, and watchdog."""
        state_result = read_state(self.instance_id)
        if isinstance(state_result, Failure):
            return Failure(
                InternalError(f"Failed to read agent state: {state_result.failure()}")
            )
        state = state_result.unwrap()

        configure_run_environment(self.instance_id, state, self.capability_token)

        pid_result = write_agent_pid(self.instance_id, os.getpid())
        if isinstance(pid_result, Failure):
            return Failure(
                InternalError(f"Failed to write PID: {pid_result.failure()}")
            )
        if state.state != AgentState.running:
            trans_result = transition_state(self.instance_id, LifecycleEvent.start)
            if isinstance(trans_result, Failure):
                return Failure(
                    InternalError(
                        f"Failed to transition to running: {trans_result.failure()}"
                    )
                )

        self._setup_signal_handlers()
        timeout_result = await self._setup_timeout_watchdog()
        if isinstance(timeout_result, Failure):
            return timeout_result

        return Success(state)

    async def _setup_timeout_watchdog(self) -> Result[None, AnimaError]:
        """Initialize and start timeout watchdog when timeout is configured."""
        result, watchdog = await setup_timeout_watchdog(
            self.timeout,
            self.job,
            self.instance_id,
            self._on_timeout,
        )
        if watchdog is not None:
            self.timeout_watchdog = watchdog
        return result

    async def _finalize_run(
        self,
        result: Result[AgentResult, AnimaError],
    ) -> Result[AgentResult, AnimaError]:
        """Finalize run cleanup and return the original result."""
        return await finalize_run(result, self.timeout_watchdog, self.mcp_server)

    def _load_runtime_config(self) -> Any:
        """Load optional runtime config and compaction settings."""
        global_config = None
        config_result = load_global_config()
        if isinstance(config_result, Success):
            global_config = config_result.unwrap()
            self._llm_timeout = global_config.defaults.llm_timeout

        compaction_result = load_compaction_config()
        if isinstance(compaction_result, Success):
            self._compaction_config = compaction_result.unwrap()
        else:
            append_log(
                self.instance_id,
                "compaction_config_load_failed",
                {"error": str(compaction_result.failure())},
            )
        return global_config

    def _resolve_persona(self) -> PersonaSpec:
        """Resolve persona with default fallback."""
        return self.persona or PersonaSpec(
            id="default",
            prompt="You are a helpful assistant.",
            model="gpt-5.2",
            description="Default runtime persona.",
        )

    async def _create_runtime_agent(
        self,
        persona: PersonaSpec,
        global_config: Any,
    ) -> Result[Agent[Any], AnimaError]:
        """Create configured runtime agent and update model context limits."""
        side_effect_policy = (
            self.job.side_effect_policy if self.job is not None else None
        )

        if self.base_url is not None:
            os.environ.setdefault("OPENAI_BASE_URL", self.base_url)
        agent_result = await create_agent(
            persona,
            mcp_servers=self.mcp_servers,
            agent_runner=self,
            side_effect_policy=side_effect_policy,
            model_override=self.model_override,
            model_instance=self.model_instance,
        )
        if isinstance(agent_result, Failure):
            transition_state(
                self.instance_id,
                LifecycleEvent.fail,
                error=str(agent_result.failure()),
            )
            return agent_result

        agent = agent_result.unwrap()
        runtime_model_context_limit = evaluate_runtime_model_context_limit(
            getattr(agent, "model", None)
        )
        self._model_context_limit = resolve_model_context_limit(
            global_config,
            runtime_model_context_limit,
            model_id=self.persona.model if self.persona else None,
        )

        snapshot_result = read_contract_snapshot(self.instance_id)
        if isinstance(snapshot_result, Success):
            snapshot_policy = snapshot_result.unwrap().output_policy
            per_call = snapshot_policy.per_call_tool_output_chars
            per_turn = snapshot_policy.per_turn_tool_output_chars
            self._tool_output_per_call_cap_bytes = (
                per_call
                if isinstance(per_call, int) and per_call > 0
                else DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES
            )
            self._tool_output_per_turn_cap_bytes = (
                per_turn
                if isinstance(per_turn, int) and per_turn > 0
                else DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES
            )
            append_log(
                self.instance_id,
                "tool_output_budget_resolved",
                {
                    "source": "contract_snapshot",
                    "per_call_cap_bytes": self._tool_output_per_call_cap_bytes,
                    "per_turn_cap_bytes": self._tool_output_per_turn_cap_bytes,
                    "provider_profile_name": snapshot_policy.provider_profile,
                    "truncation_strategy": snapshot_policy.truncation_strategy,
                },
            )
        else:
            snapshot_error = snapshot_result.failure()
            if not isinstance(snapshot_error, NotFoundError):
                return Failure(
                    InternalError(
                        f"Failed to load contract snapshot for tool-output budget: {snapshot_error}"
                    )
                )

            tool_output_config: dict[str, object] | None = None
            if global_config is not None:
                raw_tool_output = getattr(global_config, "tool_output", None)
                model_dump = getattr(raw_tool_output, "model_dump", None)
                if callable(model_dump):
                    dumped = model_dump()
                    if isinstance(dumped, dict):
                        tool_output_config = dumped

            resolved_tool_output = resolve_tool_output_budget_profile(
                cli=None,
                env={
                    "profile": os.environ.get("ANIMA_TOOL_OUTPUT_PROFILE"),
                    "per_call_cap_bytes": os.environ.get(
                        "ANIMA_TOOL_OUTPUT_PER_CALL_CAP_BYTES"
                    ),
                    "per_turn_cap_bytes": os.environ.get(
                        "ANIMA_TOOL_OUTPUT_PER_TURN_CAP_BYTES"
                    ),
                },
                config=tool_output_config,
                model_hint=str(getattr(agent, "model", "")),
                provider_hint=None,
                context_window_tokens=self._model_context_limit,
            )
            resolved_per_call = resolved_tool_output.get("per_call_cap_bytes")
            resolved_per_turn = resolved_tool_output.get("per_turn_cap_bytes")
            self._tool_output_per_call_cap_bytes = (
                resolved_per_call
                if isinstance(resolved_per_call, int) and resolved_per_call > 0
                else DEFAULT_TOOL_OUTPUT_PER_CALL_CAP_BYTES
            )
            self._tool_output_per_turn_cap_bytes = (
                resolved_per_turn
                if isinstance(resolved_per_turn, int) and resolved_per_turn > 0
                else DEFAULT_TOOL_OUTPUT_PER_TURN_CAP_BYTES
            )
            append_log(
                self.instance_id,
                "tool_output_budget_resolved",
                {
                    "source": "legacy_runtime_fallback",
                    "per_call_cap_bytes": self._tool_output_per_call_cap_bytes,
                    "per_turn_cap_bytes": self._tool_output_per_turn_cap_bytes,
                    "provider_profile_name": resolved_tool_output.get(
                        "provider_profile_name"
                    ),
                    "model_profile_name": resolved_tool_output.get(
                        "model_profile_name"
                    ),
                },
            )

        # Persist resolved model to state at startup
        # This ensures state.json has the effective model (after resolution chain)
        resolved_model = getattr(agent, "model", None)
        if resolved_model is not None:
            update_agent_state(self.instance_id, {"model": str(resolved_model)})

        log_ctl_tool_registration(self.instance_id, mcp_ctl_tools=[])
        return Success(agent)

    async def _start_runtime_mcp_server(self) -> Result[None, AnimaError]:
        """Start MCP server and persist socket path when available."""
        self.mcp_server = AgentMCPServer(self.instance_id, self)
        start_result = await self.mcp_server.start()
        if isinstance(start_result, Failure):
            transition_state(
                self.instance_id,
                LifecycleEvent.fail,
                error=f"MCP server start failed: {start_result.failure().message}",
            )
            return Failure(
                InternalError(
                    f"MCP server start failed: {start_result.failure().message}"
                )
            )

        socket_path = self.mcp_server.socket_path
        if socket_path is not None:
            socket_update_result = update_agent_state(
                self.instance_id, {"mcp_socket": str(socket_path)}
            )
            if isinstance(socket_update_result, Failure):
                append_log(
                    self.instance_id,
                    "mcp_socket_update_failed",
                    {"error": str(socket_update_result.failure())},
                )
        return Success(None)

    async def run(self) -> Result[AgentResult, AnimaError]:
        """Run the agent loop with MCP server.

        Returns:
            Success with AgentResult containing final state and metrics,
            or Failure with AnimaError on failure.
        """
        init_result = await self._initialize_run_state()
        if isinstance(init_result, Failure):
            return init_result
        state = init_result.unwrap()

        global_config = self._load_runtime_config()
        persona = self._resolve_persona()
        agent_result = await self._create_runtime_agent(persona, global_config)
        if isinstance(agent_result, Failure):
            return agent_result

        agent = agent_result.unwrap()
        start_result = await self._start_runtime_mcp_server()
        if isinstance(start_result, Failure):
            return start_result

        # Deferred persona resolution: if persona is a bootstrap placeholder,
        # resolve the real persona via larva MCP and rebuild the agent.
        if persona is not None and persona.id.startswith("bootstrap-"):
            resolved = await self._resolve_deferred_persona(persona, global_config)
            if isinstance(resolved, Success):
                agent = resolved.unwrap()
                persona = self.persona  # Updated by _apply_persona_switch

        run_result: Result[AgentResult, AnimaError]
        try:
            run_result = await self._run_agent_loop(agent, state)
        except Exception as e:
            transition_state(
                self.instance_id,
                LifecycleEvent.fail,
                error=f"Agent loop error ({type(e).__name__}): {e}",
            )
            run_result = Failure(InternalError(f"Agent loop failed: {e}"))

        return await self._finalize_run(run_result)

    def _reset_tracking(self) -> None:
        """Reset turn tracking counters."""
        self.turn_count = 0
        self.token_usage = TokenUsage()
        self.no_tool_streak = 0

    def _check_interrupt(self) -> Result[AgentResult, AnimaError] | None:
        """Check for interrupt signal or timeout."""
        if self._timeout_fired:
            return Failure(
                AnimaTimeoutError(
                    f"Agent timed out after {self.timeout} wall-clock duration",
                    details={"timeout": self.timeout},
                )
            )
        if self._interrupted:
            transition_state(self.instance_id, LifecycleEvent.kill)
            return Success(
                AgentResult(final_state=AgentState.killed, token_usage=self.token_usage)
            )
        return None

    def _handle_model_request(self) -> Result[AgentResult, AnimaError] | None:
        """Handle a model request turn via extracted budget logic."""
        self.turn_count += 1
        self._update_state()
        result, self.in_grace_window, self.max_turns_exceeded = (
            handle_model_request_turn(
                self.turn_count,
                self.max_turns,
                self.in_grace_window,
                self.budget_tracker,
                self.instance_id,
            )
        )
        if result is not None and isinstance(result, Failure):
            self.grace_exhausted_without_finish = True
        return result

    def _record_tool_usage(self, result: Any) -> None:
        """Record tool usage state for auto completion."""
        self.no_tool_streak = record_tool_usage(
            result, self.no_tool_streak, self.completion_policy
        )

    async def _run_pending_compaction(
        self,
        agent: Agent[Any],
        message_history: list[Any],
    ) -> list[Any]:
        """Run pending compaction at safe turn boundaries."""
        if not self._compaction_pending:
            return message_history
        if not message_history:
            self._compaction_pending = False
            return message_history

        self._compaction_pending = False

        try:
            compaction_result = await run_compaction(
                agent=agent,
                message_history=message_history,
                compaction_config=self._compaction_config,
                persona_compaction_prompt=(
                    self.persona.compaction_prompt if self.persona is not None else None
                ),
                instance_id=self.instance_id,
            )
        except Exception as exc:
            append_log(
                self.instance_id,
                "compaction_failed",
                {"error": str(exc), "error_type": type(exc).__name__},
            )
            return message_history

        if isinstance(compaction_result, Failure):
            append_log(
                self.instance_id,
                "compaction_failed",
                {"error": str(compaction_result.failure())},
            )
            return message_history

        return compaction_result.unwrap()

    def _process_model_result(
        self,
        result: Any,
        run: Any,
        model_request_node_type: type[Any],
        call_tools_node_type: type[Any],
    ) -> tuple[Result[AgentResult, AnimaError] | None, bool]:
        """Process one model-loop node and return terminal state if any."""
        if isinstance(result, model_request_node_type):
            exceeded = self._handle_model_request()
            if exceeded:
                return exceeded, False
        if isinstance(result, call_tools_node_type):
            self._record_tool_usage(result)

        usage = run.usage()
        if usage:
            self._track_token_usage(usage)
            if should_compact(
                self.token_usage,
                self._compaction_config,
                self._model_context_limit,
            ):
                self._compaction_pending = True
            budget_result = self._check_budget_exhaustion()
            if budget_result is not None:
                return budget_result, True

        # Cache latest response_text so ctl_finish wrapper can include it
        # in result.json without a race window.
        self.last_response_text = extract_last_response_text(run.all_messages())

        completed = self._check_auto_complete()
        if completed:
            response_text = self.last_response_text
            if response_text is not None:
                append_log(
                    self.instance_id,
                    "response",
                    {
                        "response_text": response_text,
                        "turn_number": self.turn_count,
                    },
                )
            if isinstance(completed, Success):
                val = completed.unwrap()
                completed = Success(
                    AgentResult(
                        final_state=val.final_state,
                        token_usage=val.token_usage,
                        result_ref=val.result_ref,
                        error=val.error,
                        response_text=response_text,
                    )
                )
            return completed, usage is not None
        return None, usage is not None

    @staticmethod
    def _extract_last_response_text(messages: list[Any]) -> str | None:
        """Extract text content from the last model response in message history."""
        return extract_last_response_text(messages)

    def _check_auto_complete(self) -> Result[AgentResult, AnimaError] | None:
        """Check for auto-completion."""
        return check_auto_complete(
            self.completion_policy,
            self.no_tool_streak,
            self.instance_id,
            self.token_usage,
        )

    async def _poll_pending_message(
        self, timeout_seconds: float | None = 30.0
    ) -> str | None:
        """Poll for pending messages from interact() MCP tool."""
        return await poll_pending_message(
            self.mcp_server,
            timeout_seconds,
            self._inject_followup_hint_if_needed,
        )

    def _read_current_state(self) -> AgentState | None:
        """Read persisted state for interactive-loop exit checks."""
        return read_current_state(self.instance_id)

    def _drain_phase1_pending_messages(self) -> list[str]:
        """Drain messages queued during initial task execution (phase 1)."""
        return drain_phase1_pending_messages(self.mcp_server)

    def _consume_stop_turn_request(self) -> bool:
        """Consume a stop-turn control request if present."""
        return consume_stop_turn_request(self.mcp_server)

    async def _resolve_deferred_persona(
        self,
        bootstrap_persona: PersonaSpec,
        global_config: Any,
    ) -> Result[Agent[Any], AnimaError]:
        """Resolve a deferred bootstrap persona via larva MCP and rebuild the agent.

        Called when persona.id starts with 'bootstrap-', indicating the real persona
        should be resolved from larva at runtime. Extracts spec_id from the bootstrap
        persona id (bootstrap-{spec_id}), calls larva_resolve via MCP, and applies
        a persona switch.

        Args:
            bootstrap_persona: The bootstrap placeholder persona.
            global_config: Global config for agent rebuild.

        Returns:
            Success(new_agent) if resolved and switched, Failure otherwise.
        """
        spec_id = bootstrap_persona.id.removeprefix("bootstrap-")
        append_log(
            self.instance_id,
            "deferred_persona_resolve_start",
            {"spec_id": spec_id},
        )

        from anima.shell.larva_persona_resolver import resolve_persona_via_larva

        resolve_result = await resolve_persona_via_larva(spec_id, self.mcp_servers)
        if isinstance(resolve_result, Failure):
            append_log(
                self.instance_id,
                "deferred_persona_resolve_failed",
                {"spec_id": spec_id, "error": str(resolve_result.failure())},
            )
            return resolve_result

        new_persona = resolve_result.unwrap()
        new_agent = await self._apply_persona_switch(
            None,  # current agent not needed, will be replaced
            new_persona,
            new_persona.id,
        )

        append_log(
            self.instance_id,
            "deferred_persona_resolved",
            {"spec_id": spec_id, "resolved_persona": new_persona.id},
        )
        return Success(new_agent)

    def _consume_pending_persona_switch(
        self,
    ) -> tuple[PersonaSpec, str] | None:
        """Check and consume a pending persona switch from the MCP server.

        Spec source: docs/INTERFACES.md D.5 — turn-boundary cooperative switch.

        Returns:
            Tuple of (PersonaSpec, persona_ref) if pending, None otherwise.
        """
        if self.mcp_server is None:
            return None
        return self.mcp_server.consume_pending_persona_switch()

    async def _apply_persona_switch(
        self,
        current_agent: Agent[Any],
        new_persona: PersonaSpec,
        new_persona_ref: str,
    ) -> Agent[Any]:
        """Rebuild agent with new persona, preserving message history.

        Spec source: docs/INTERFACES.md D.5 + docs/DESIGN.md Section 8.
        Preserves: message_history, turn counters, token accounting, MCP state.
        Rebuilds: PydanticAI Agent via existing create_agent path.

        Args:
            current_agent: The currently running agent (fallback on failure).
            new_persona: The new PersonaSpec to switch to.
            new_persona_ref: File path reference for the new persona.

        Returns:
            The new agent on success, or current_agent on failure.
        """
        self.persona = new_persona

        global_config = self._load_runtime_config()
        agent_result = await self._create_runtime_agent(new_persona, global_config)
        if isinstance(agent_result, Failure):
            append_log(
                self.instance_id,
                "persona_switch_failed",
                {"error": str(agent_result.failure())},
            )
            return current_agent

        new_agent = agent_result.unwrap()

        update_agent_state(
            self.instance_id,
            {
                "persona_name": new_persona.id,
                "persona_ref": new_persona_ref,
                "spec_id": new_persona.id,
                "spec_digest": new_persona.spec_digest,
                "model": new_persona.model,
            },
        )

        append_log(
            self.instance_id,
            "persona_switched",
            {
                "new_persona": new_persona.id,
                "new_model": new_persona.model,
                "persona_ref": new_persona_ref,
            },
        )

        return new_agent

    # @invar:allow function_size: Conversation continuation with message handling
    async def _process_interact_message(
        self,
        agent: Agent[Any],
        message: str,
        message_history: list[Any],
    ) -> Result[list[Any], AnimaError]:
        """Process an interact message as a continuation of the conversation."""
        (
            result,
            turn_count_delta,
            new_no_tool_streak,
            new_in_grace,
        ) = await process_interact_message(
            agent,
            message,
            message_history,
            instance_id=self.instance_id,
            mcp_server=self.mcp_server,
            model_settings=self._get_model_settings_with_timeout(),
            turn_count=self.turn_count,
            max_turns=self.max_turns,
            token_usage=self.token_usage,
            budget_tracker=self.budget_tracker,
            no_tool_streak=self.no_tool_streak,
            completion_policy=self.completion_policy,
            in_grace_window=self.in_grace_window,
            check_interrupt_fn=self._check_interrupt,
            queue_followup_hint_fn=self._queue_followup_hint_if_pending_messages,
            tool_output_per_call_cap_bytes=self._tool_output_per_call_cap_bytes,
            tool_output_per_turn_cap_bytes=self._tool_output_per_turn_cap_bytes,
        )
        self.turn_count += turn_count_delta
        self.no_tool_streak = new_no_tool_streak
        self.in_grace_window = new_in_grace
        if isinstance(result, Success) and should_compact(
            self.token_usage,
            self._compaction_config,
            self._model_context_limit,
        ):
            self._compaction_pending = True
        return result

    # @invar:allow function_size: AgentRunner._run_agent_loop keeps turn accounting and lifecycle completion decisions in one control loop
    async def _run_agent_loop(
        self,
        agent: Agent[Any],
        state: Any,
    ) -> Result[AgentResult, AnimaError]:
        """Run the main agent loop with turn tracking."""
        from pydantic_ai._agent_graph import CallToolsNode, ModelRequestNode

        self._reset_tracking()

        killed = self._check_interrupt()
        if killed:
            return killed

        mcp_instructions: str | None = None
        self.grace_message_injected = False

        async def get_dynamic_instructions(ctx: Any) -> str | None:
            nonlocal mcp_instructions

            # Collect MCP context only once (async with cache)
            if mcp_instructions is None:
                mcp_instructions = await collect_mcp_server_context(
                    self.mcp_servers, ctx
                )

            parts: list[str] = []

            # MCP instructions first
            if mcp_instructions:
                parts.append(mcp_instructions)

            # Grace window message (preserving exact behavior)
            if self.in_grace_window and not self.grace_message_injected:
                self.grace_message_injected = True
                grace_msg = grace_system_message(self.budget_tracker.grace_budget)
                append_log(
                    self.instance_id,
                    "grace_message_injected",
                    {
                        "message_length": len(grace_msg),
                        "grace_budget": self.budget_tracker.grace_budget,
                    },
                )
                parts.append(grace_msg)

            if not parts:
                return None

            return "\n\n".join(parts)

        model_settings = self._get_model_settings_with_timeout()

        # Load context bundle history if context_mode='shared' and bundle_ref is set
        # Runtime enforcement: context_bundle_ref is ONLY loaded when BOTH conditions are true:
        #   1. context_bundle_ref is not None
        #   2. context_mode.value == "shared"
        # This is defense-in-depth:
        #   - Schema validation rejects isolated + bundle_ref at construction time
        #   - Runtime check prevents loading even if schema is bypassed
        message_history: list[Any] = []
        if (
            self.job is not None
            and self.job.context_bundle_ref is not None
            and self.job.context_mode.value == "shared"
        ):
            from anima.shell.context_bundle import (
                load_message_history_from_context_bundle,
            )

            history_result = load_message_history_from_context_bundle(
                self.job.context_bundle_ref
            )
            if isinstance(history_result, Failure):
                return Failure(history_result.failure())
            message_history = history_result.unwrap()
        if message_history:
            message_history = assemble_tool_output_history(
                message_history,
                preserve_after_index=len(message_history),
                per_call_cap_bytes=self._tool_output_per_call_cap_bytes,
                per_turn_cap_bytes=self._tool_output_per_turn_cap_bytes,
            )

        try:
            # Skip initial LLM call when task is empty (interactive-only session).
            # Some models reject empty prompts, and there's nothing to process.
            if self.task.strip():
                # Spec source: plan step talk-ux-bugs-v3.fix-spawn-history-replay.
                # Initial spawn task must be logged as user_message so talk history
                # replay includes the first user turn.
                append_log(
                    self.instance_id,
                    "user_message",
                    {
                        "message_text": self.task,
                        "turn_number": self.turn_count + 1,
                    },
                )

                turn_budget_token = begin_tool_output_budget_scope(
                    per_call_cap_bytes=self._tool_output_per_call_cap_bytes,
                    per_turn_cap_bytes=self._tool_output_per_turn_cap_bytes,
                )
                try:
                    self._reset_run_usage_tracking()
                    async with agent.iter(
                        user_prompt=self.task,
                        instructions=get_dynamic_instructions,
                        model_settings=model_settings,
                        message_history=message_history,
                        usage_limits=UsageLimits(request_limit=None),
                    ) as run:
                        async for node in run:
                            killed = self._check_interrupt()
                            if killed:
                                return killed
                            node_result, _ = self._process_model_result(
                                node,
                                run,
                                ModelRequestNode,
                                CallToolsNode,
                            )
                            if node_result:
                                return node_result

                        message_history = run.all_messages()
                        message_history = await self._run_pending_compaction(
                            agent, message_history
                        )

                        for pending_message in self._drain_phase1_pending_messages():
                            message_for_agent = (
                                self._wrap_followup_message_with_system_reminder(
                                    pending_message
                                )
                            )
                            response_result = await self._process_interact_message(
                                agent,
                                message_for_agent,
                                message_history,
                            )
                            if isinstance(response_result, Failure):
                                return response_result
                            message_history = response_result.unwrap()
                            message_history = await self._run_pending_compaction(
                                agent, message_history
                            )
                finally:
                    end_tool_output_budget_scope(turn_budget_token)

            # Log initial task response for all lifecycles so talk history
            # replay shows the agent's first response when attaching later.
            response_text = extract_last_response_text(message_history)
            if response_text is not None:
                append_log(
                    self.instance_id,
                    "response",
                    {
                        "response_text": response_text,
                        "turn_number": self.turn_count,
                    },
                )

            if state.lifecycle == LifecycleType.ephemeral:
                transition_state(
                    self.instance_id, LifecycleEvent.finish, result_ref=None
                )
                return Success(
                    AgentResult(
                        final_state=AgentState.completed,
                        token_usage=self.token_usage,
                        response_text=response_text,
                    )
                )

            while True:
                killed = self._check_interrupt()
                if killed:
                    return killed

                current_state = self._read_current_state()
                if current_state == AgentState.completed:
                    return Success(
                        AgentResult(
                            final_state=AgentState.completed,
                            token_usage=self.token_usage,
                        )
                    )

                # Check for pending persona switch at turn boundary.
                # Spec source: docs/INTERFACES.md D.5 — turn-boundary cooperative switch.
                pending_switch = self._consume_pending_persona_switch()
                if pending_switch is not None:
                    new_persona, new_persona_ref = pending_switch
                    agent = await self._apply_persona_switch(
                        agent, new_persona, new_persona_ref
                    )

                pending_message = await self._poll_pending_message(timeout_seconds=30.0)
                if pending_message is not None:
                    # Proactive compaction: check if message history is approaching
                    # context limit BEFORE sending to the model. Without this,
                    # the model request fails with the full history and the agent
                    # crashes instead of compacting.
                    if should_compact(
                        self.token_usage,
                        self._compaction_config,
                        self._model_context_limit,
                    ):
                        message_history = await self._run_pending_compaction(
                            agent, message_history
                        )
                        self._compaction_pending = False

                    response_result = await self._process_interact_message(
                        agent, pending_message, message_history
                    )
                    if isinstance(response_result, Failure):
                        return response_result
                    message_history = response_result.unwrap()
                    message_history = await self._run_pending_compaction(  # @invar:allow dead_assign: compaction rewrites history for next turn
                        agent, message_history
                    )
                    continue

                await asyncio.sleep(0)

        except (
            ModelHTTPError,
            PAgentUserError,
            RuntimeError,
            OSError,
            ValueError,
            TypeError,
        ) as e:
            transition_state(
                self.instance_id,
                LifecycleEvent.fail,
                error=f"Agent execution error ({type(e).__name__}): {e}",
            )
            return Failure(InternalError(f"Agent execution error: {e}"))
        except Exception as e:
            # Catch-all for unexpected exceptions (e.g., openai.APIConnectionError,
            # httpx errors, LLM provider errors not in the explicit list above).
            # Without this, the agent process crashes silently and the talk
            # session sees "peer closed connection" with no error message.
            transition_state(
                self.instance_id,
                LifecycleEvent.fail,
                error=f"Agent crash ({type(e).__name__}): {e}",
            )
            return Failure(InternalError(f"Agent crash ({type(e).__name__}): {e}"))

    def _update_state(self) -> None:
        """Update state file with current turn count and token usage."""
        budget_update_state(self.instance_id, self.turn_count, self.token_usage)

    def _track_token_usage(self, usage: Any) -> None:
        """Track token usage delta from pydantic-ai run.

        PydanticAI's run.usage() returns cumulative values within an iter() call.
        We compute the delta since the last call to avoid double-counting.
        """
        if not usage:
            return
        curr_input = usage.input_tokens or 0
        curr_output = usage.output_tokens or 0
        delta_input = curr_input - self._prev_run_input_tokens
        delta_output = curr_output - self._prev_run_output_tokens
        self._prev_run_input_tokens = curr_input
        self._prev_run_output_tokens = curr_output

        if delta_input <= 0 and delta_output <= 0:
            return

        self.token_usage.prompt_tokens += delta_input
        self.token_usage.completion_tokens += delta_output
        # last_request_prompt_tokens = this request's actual input size
        # (the full cumulative input for this iter call = context window fill)
        self.token_usage.last_request_prompt_tokens = curr_input
        self.budget_tracker.record_usage(delta_input + delta_output)
        budget_update_state(self.instance_id, self.turn_count, self.token_usage)

    def _reset_run_usage_tracking(self) -> None:
        """Reset cumulative usage snapshot for a new iter() call."""
        self._prev_run_input_tokens = 0
        self._prev_run_output_tokens = 0

    def set_interactive_approval_mode(self, enabled: bool) -> None:
        """Set whether approvals should use interactive in-memory channel."""
        self._interactive_approval_mode = enabled
        if not enabled:
            self._interactive_approval_future = None

    def prepare_interactive_approval_future(
        self,
    ) -> asyncio.Future[dict[str, bool | str | None]] | None:
        """Create per-request approval future when interactive mode is active."""
        if not self._interactive_approval_mode:
            return None
        future: asyncio.Future[dict[str, bool | str | None]] = (
            asyncio.get_running_loop().create_future()
        )
        self._interactive_approval_future = future
        return future

    def resolve_interactive_approval_future(
        self,
        decision: dict[str, bool | str | None],
    ) -> None:
        """Resolve active interactive approval future after state transition writes."""
        future = self._interactive_approval_future
        if future is None or future.done():
            return
        future.set_result(decision)

    def _check_budget_exhaustion(self) -> Result[AgentResult, AnimaError] | None:
        """Check for budget exhaustion and grace window entry."""
        result, self.in_grace_window = check_budget_exhaustion(
            self.budget_tracker,
            self.in_grace_window,
            self.instance_id,
        )
        if result is not None and isinstance(result, Failure):
            self.grace_exhausted_without_finish = True
        return result


# =============================================================================
# Backward compatibility re-exports
# =============================================================================

__all__ = [
    "create_agent",
    "AgentRunner",
    "create_ctl_tools",
]
