"""Shared execution loop for tool-using agents.

Architecture Note (Runner vs Loop):
-----------------------------------
This module contains **AgentRunner**, the pure logic layer for LLM execution.
It knows nothing about messaging, sessions, or product-layer concerns.

**AgentRunner** (this file) responsibilities:
  - Execute the LLM iteration loop (request → tool calls → response)
  - Call the LLM provider and handle streaming/non-streaming responses
  - Execute tool calls and collect results
  - Manage context budget and history truncation (_snip_history)
  - Handle finalization retries and error states
  - Emit checkpoints and context usage callbacks

**AgentLoop** (loop.py) responsibilities:
  - Integrate with the message bus (receive InboundMessage, send OutboundMessage)
  - Build context (system prompt, session history, skills, memory)
  - Manage sessions and conversation history
  - Handle streaming output to channels (CLI, DingTalk, Feishu, etc.)
  - Coordinate components: TaskManager, VectorStore, Consolidator, Dream
  - Process MCP server connections
  - Route slash commands

The separation allows Runner to be reused in different contexts (CLI, SDK, tests)
while Loop handles all product-specific integration.
"""

from __future__ import annotations

import asyncio
from datetime import datetime
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from loguru import logger

from tinybot.agent.hook import AgentHook, AgentHookContext
from tinybot.agent.tools.base import AwaitingUserInputResult
from tinybot.agent.tools.registry import ToolRegistry
from tinybot.providers.base import LLMProvider, ToolCallRequest
from tinybot.security.approval import ApprovalAction, ApprovalManager, format_approval_required
from tinybot.utils.fs import find_legal_message_start
from tinybot.utils.helper import (
    build_assistant_message,
    maybe_persist_tool_result,
    truncate_text,
)
from tinybot.utils.prompt_templates import render_template
from tinybot.utils.runtime import (
    EMPTY_FINAL_RESPONSE_MESSAGE,
    build_finalization_retry_message,
    ensure_nonempty_tool_result,
    is_blank_text,
    repeated_external_lookup_error,
)
from tinybot.utils.tokens import estimate_message_tokens, estimate_prompt_tokens_chain

_DEFAULT_ERROR_MESSAGE = "Sorry, I encountered an error calling the AI model."
_SNIP_SAFETY_BUFFER = 1024
@dataclass(slots=True)
class AgentRunSpec:
    """Configuration for a single agent execution."""

    initial_messages: list[dict[str, Any]]
    tools: ToolRegistry
    model: str
    max_iterations: int
    max_tool_result_chars: int
    temperature: float | None = None
    max_tokens: int | None = None
    reasoning_effort: str | None = None
    hook: AgentHook | None = None
    error_message: str | None = _DEFAULT_ERROR_MESSAGE
    max_iterations_message: str | None = None
    concurrent_tools: bool = False
    fail_on_tool_error: bool = False
    experience_analyzer: Any | None = None  # ErrorAnalyzer for auto error diagnosis
    experience_store: Any | None = None
    workspace: Path | None = None
    session_key: str | None = None
    session: Any | None = None
    context_window_tokens: int | None = None
    context_block_limit: int | None = None
    provider_retry_mode: str = "standard"
    progress_callback: Any | None = None
    checkpoint_callback: Any | None = None
    context_usage_callback: Any | None = None



@dataclass(slots=True)
class AgentRunResult:
    """Outcome of a shared agent execution."""

    final_content: str | None
    messages: list[dict[str, Any]]
    tools_used: list[str] = field(default_factory=list)
    usage: dict[str, int] = field(default_factory=dict)
    stop_reason: str = "completed"
    error: str | None = None
    tool_events: list[dict[str, str]] = field(default_factory=list)


class AgentRunner:
    """Run a tool-capable LLM loop without product-layer concerns."""

    def __init__(self, provider: LLMProvider):
        self.provider = provider

    async def run(self, spec: AgentRunSpec) -> AgentRunResult:
        hook = spec.hook or AgentHook()
        messages = list(spec.initial_messages)
        final_content: str | None = None
        tools_used: list[str] = []
        usage: dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0}
        error: str | None = None
        stop_reason = "completed"
        tool_events: list[dict[str, str]] = []
        external_lookup_counts: dict[str, int] = {}
        recovery_hints: dict[str, list[str]] = {}
        tool_attempts: dict[str, int] = {}

        for iteration in range(spec.max_iterations):
            try:
                messages = self._apply_tool_result_budget(spec, messages)
                messages_for_model = self._snip_history(spec, messages)
            except Exception as exc:
                logger.warning(
                    "Context governance failed on turn {} for {}: {}; using raw messages",
                    iteration,
                    spec.session_key or "default",
                    exc,
                )
                messages_for_model = messages
            context = AgentHookContext(iteration=iteration, messages=messages)
            await hook.before_iteration(context)
            await self._emit_context_usage(
                spec,
                self._build_context_usage_payload(
                    spec,
                    messages_for_model,
                    iteration=iteration,
                    phase="before_request",
                ),
            )
            response = await self._request_model(spec, messages_for_model, hook, context)
            raw_usage = self._usage_dict(response.usage)
            context.response = response
            context.usage = dict(raw_usage)
            context.tool_calls = list(response.tool_calls)
            self._accumulate_usage(usage, raw_usage)
            await self._emit_context_usage(
                spec,
                self._build_actual_usage_payload(
                    raw_usage,
                    iteration=iteration,
                    phase="after_response",
                ),
            )


            if response.has_tool_calls:
                if hook.wants_streaming():
                    await hook.on_stream_end(context, resuming=True)

                assistant_message = build_assistant_message(
                    response.content or "",
                    tool_calls=[tc.to_openai_tool_call() for tc in response.tool_calls],
                    reasoning_content=response.reasoning_content,
                    thinking_blocks=response.thinking_blocks,
                )
                messages.append(assistant_message)
                tools_used.extend(tc.name for tc in response.tool_calls)
                await self._emit_checkpoint(
                    spec,
                    {
                        "phase": "awaiting_tools",
                        "iteration": iteration,
                        "model": spec.model,
                        "assistant_message": assistant_message,
                        "completed_tool_results": [],
                        "pending_tool_calls": [tc.to_openai_tool_call() for tc in response.tool_calls],
                    },
                )

                await hook.before_execute_tools(context)

                results, new_events, fatal_error = await self._execute_tools(
                    spec,
                    response.tool_calls,
                    external_lookup_counts,
                    recovery_hints,
                    tool_attempts,
                    hook,
                    context,
                )
                tool_events.extend(new_events)
                context.tool_results = list(results)
                context.tool_events = list(new_events)
                await hook.after_execute_tools(context)
                if any(event.get("status") == "approval_required" for event in new_events):
                    stop_reason = "awaiting_approval"
                    context.stop_reason = stop_reason
                    await hook.after_iteration(context)
                    break
                if fatal_error is not None:
                    error = f"Error: {type(fatal_error).__name__}: {fatal_error}"
                    final_content = error
                    stop_reason = "tool_error"
                    self._append_final_message(messages, final_content)
                    context.final_content = final_content
                    context.error = error
                    context.stop_reason = stop_reason
                    await hook.after_iteration(context)
                    break
                completed_tool_results: list[dict[str, Any]] = []
                for tool_call, result in zip(response.tool_calls, results, strict=False):
                    tool_message = {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": tool_call.name,
                        "content": self._normalize_tool_result(
                            spec,
                            tool_call.id,
                            tool_call.name,
                            result,
                        ),
                    }
                    if isinstance(result, AwaitingUserInputResult):
                        tool_message["_awaiting_user_input"] = True
                        tool_message["_agent_ui_internal"] = True
                    messages.append(tool_message)
                    completed_tool_results.append(tool_message)
                await self._emit_checkpoint(
                    spec,
                    {
                        "phase": "tools_completed",
                        "iteration": iteration,
                        "model": spec.model,
                        "assistant_message": assistant_message,
                        "completed_tool_results": completed_tool_results,
                        "pending_tool_calls": [],
                    },
                )
                awaiting_user_input = next(
                    (result for result in results if isinstance(result, AwaitingUserInputResult)),
                    None,
                )
                if awaiting_user_input is not None:
                    if hook.wants_streaming():
                        await hook.on_stream_end(context, resuming=False)
                    final_content = None
                    stop_reason = awaiting_user_input.stop_reason
                    context.final_content = None
                    context.stop_reason = stop_reason
                    await hook.after_iteration(context)
                    break
                await hook.after_iteration(context)
                continue

            clean = hook.finalize_content(context, response.content)
            did_finalization_retry = False
            if response.finish_reason != "error" and is_blank_text(clean):
                logger.warning(
                    "Empty final response on turn {} for {}; retrying with explicit finalization prompt",
                    iteration,
                    spec.session_key or "default",
                )
                # Stream end before retry: this clears the current buffers
                # (reasoning was streamed, but content is empty)
                if hook.wants_streaming():
                    await hook.on_stream_end(context, resuming=False)
                await self._emit_context_usage(
                    spec,
                    self._build_context_usage_payload(
                        spec,
                        messages_for_model + [build_finalization_retry_message()],
                        iteration=iteration,
                        phase="before_finalization_retry",
                    ),
                )
                # Retry uses non-stream API, so no stream deltas will follow
                response = await self._request_finalization_retry(spec, messages_for_model)
                did_finalization_retry = True
                retry_usage = self._usage_dict(response.usage)
                self._accumulate_usage(usage, retry_usage)
                raw_usage = self._merge_usage(raw_usage, retry_usage)
                context.response = response
                context.usage = dict(raw_usage)
                context.tool_calls = list(response.tool_calls)
                await self._emit_context_usage(
                    spec,
                    self._build_actual_usage_payload(
                        retry_usage,
                        iteration=iteration,
                        phase="after_finalization_retry",
                    ),
                )
                clean = hook.finalize_content(context, response.content)


            # Only call stream_end if we didn't do a finalization_retry.
            # After retry, content comes from non-stream API and will be
            # delivered via OutboundMessage (not stream deltas).
            if hook.wants_streaming() and not did_finalization_retry:
                await hook.on_stream_end(context, resuming=False)

            if response.finish_reason == "error":
                final_content = clean or spec.error_message or _DEFAULT_ERROR_MESSAGE
                stop_reason = "error"
                error = final_content
                self._append_final_message(messages, final_content)
                context.final_content = final_content
                context.error = error
                context.stop_reason = stop_reason
                await hook.after_iteration(context)
                break
            if is_blank_text(clean):
                final_content = EMPTY_FINAL_RESPONSE_MESSAGE
                stop_reason = "empty_final_response"
                error = final_content
                self._append_final_message(messages, final_content)
                context.final_content = final_content
                context.error = error
                context.stop_reason = stop_reason
                await hook.after_iteration(context)
                break

            messages.append(build_assistant_message(
                clean,
                reasoning_content=response.reasoning_content,
                thinking_blocks=response.thinking_blocks,
            ))
            await self._emit_checkpoint(
                spec,
                {
                    "phase": "final_response",
                    "iteration": iteration,
                    "model": spec.model,
                    "assistant_message": messages[-1],
                    "completed_tool_results": [],
                    "pending_tool_calls": [],
                },
            )
            final_content = clean
            context.final_content = final_content
            context.stop_reason = stop_reason
            await hook.after_iteration(context)
            break
        else:
            stop_reason = "max_iterations"
            if spec.max_iterations_message:
                final_content = spec.max_iterations_message.format(
                    max_iterations=spec.max_iterations,
                )
            else:
                final_content = render_template(
                    "agent/max_iterations_message.md",
                    strip=True,
                    max_iterations=spec.max_iterations,
                )
            self._append_final_message(messages, final_content)

        return AgentRunResult(
            final_content=final_content,
            messages=messages,
            tools_used=tools_used,
            usage=usage,
            stop_reason=stop_reason,
            error=error,
            tool_events=tool_events,
        )

    def _build_request_kwargs(
        self,
        spec: AgentRunSpec,
        messages: list[dict[str, Any]],
        *,
        tools: list[dict[str, Any]] | None,
    ) -> dict[str, Any]:
        kwargs: dict[str, Any] = {
            "messages": messages,
            "tools": tools,
            "model": spec.model,
            "retry_mode": spec.provider_retry_mode,
            "on_retry_wait": spec.progress_callback,
        }
        if spec.temperature is not None:
            kwargs["temperature"] = spec.temperature
        if spec.max_tokens is not None:
            kwargs["max_tokens"] = spec.max_tokens
        if spec.reasoning_effort is not None:
            kwargs["reasoning_effort"] = spec.reasoning_effort
        return kwargs

    async def _request_model(
        self,
        spec: AgentRunSpec,
        messages: list[dict[str, Any]],
        hook: AgentHook,
        context: AgentHookContext,
    ):
        kwargs = self._build_request_kwargs(
            spec,
            messages,
            tools=spec.tools.get_definitions(),
        )
        if hook.wants_streaming():
            async def _stream(delta: str) -> None:
                await hook.on_stream(context, delta)

            async def _reasoning_stream(delta: str) -> None:
                await hook.on_reasoning_stream(context, delta)

            return await self.provider.chat_stream_with_retry(
                **kwargs,
                on_content_delta=_stream,
                on_reasoning_delta=_reasoning_stream,
            )
        return await self.provider.chat_with_retry(**kwargs)

    async def _request_finalization_retry(
        self,
        spec: AgentRunSpec,
        messages: list[dict[str, Any]],
    ):
        retry_messages = list(messages)
        retry_messages.append(build_finalization_retry_message())
        kwargs = self._build_request_kwargs(spec, retry_messages, tools=None)
        return await self.provider.chat_with_retry(**kwargs)

    @staticmethod
    def _usage_dict(usage: dict[str, Any] | None) -> dict[str, int]:
        if not usage:
            return {}
        result: dict[str, int] = {}
        for key, value in usage.items():
            try:
                result[key] = int(value or 0)
            except (TypeError, ValueError):
                continue
        return result

    @staticmethod
    def _accumulate_usage(target: dict[str, int], addition: dict[str, int]) -> None:
        for key, value in addition.items():
            target[key] = target.get(key, 0) + value

    @staticmethod
    def _merge_usage(left: dict[str, int], right: dict[str, int]) -> dict[str, int]:
        merged = dict(left)
        for key, value in right.items():
            merged[key] = merged.get(key, 0) + value
        return merged

    async def _execute_tools(
        self,
        spec: AgentRunSpec,
        tool_calls: list[ToolCallRequest],
        external_lookup_counts: dict[str, int],
        recovery_hints: dict[str, list[str]],
        tool_attempts: dict[str, int],
        hook: AgentHook | None = None,
        context: AgentHookContext | None = None,
    ) -> tuple[list[Any], list[dict[str, str]], BaseException | None]:
        batches = self._partition_tool_batches(spec, tool_calls)
        tool_results: list[tuple[Any, dict[str, str], BaseException | None]] = []
        for batch in batches:
            if spec.concurrent_tools and len(batch) > 1:
                tool_results.extend(await asyncio.gather(*(
                    self._run_tool(
                        spec,
                        tool_call,
                        external_lookup_counts,
                        recovery_hints,
                        tool_attempts,
                        hook,
                        context,
                    )
                    for tool_call in batch
                )))
            else:
                for tool_call in batch:
                    tool_results.append(
                        await self._run_tool(
                            spec,
                            tool_call,
                            external_lookup_counts,
                            recovery_hints,
                            tool_attempts,
                            hook,
                            context,
                        )
                    )

        results: list[Any] = []
        events: list[dict[str, str]] = []
        fatal_error: BaseException | None = None
        for result, event, error in tool_results:
            results.append(result)
            events.append(event)
            if error is not None and fatal_error is None:
                fatal_error = error
        return results, events, fatal_error

    async def _run_tool(
        self,
        spec: AgentRunSpec,
        tool_call: ToolCallRequest,
        external_lookup_counts: dict[str, int],
        recovery_hints: dict[str, list[str]],
        tool_attempts: dict[str, int],
        hook: AgentHook | None = None,
        context: AgentHookContext | None = None,
    ) -> tuple[Any, dict[str, str], BaseException | None]:
        _HINT = "\n\n[Analyze the error above and try a different approach. Consider using `query_experience` to search for past solutions to similar problems.]"
        lookup_error = repeated_external_lookup_error(
            tool_call.name,
            tool_call.arguments,
            external_lookup_counts,
        )
        attempt_no = tool_attempts.get(tool_call.name, 0) + 1
        tool_attempts[tool_call.name] = attempt_no
        if lookup_error:
            event = {
                "name": tool_call.name,
                "status": "error",
                "detail": "repeated external lookup blocked",
            }
            if spec.fail_on_tool_error:
                return lookup_error + _HINT, event, RuntimeError(lookup_error)
            return lookup_error + _HINT, event, None

        # Auto-analyze error with experience analyzer if available
        def _get_error_suggestions(tool_name: str, error: Exception | str) -> str:
            if spec.experience_analyzer:
                strategy = spec.experience_analyzer.build_recovery_strategy(tool_name, error)
                if strategy:
                    recovery_hints[tool_name] = [exp.id for exp in strategy["experiences"]]
                    primary_block = ""
                    retry_policy = strategy.get("retry_policy")
                    if retry_policy:
                        primary_block = (
                            "\n\n"
                            + spec.experience_analyzer.format_retry_policy(retry_policy)
                            + "\n"
                        )
                    elif strategy.get("primary_action", "").strip():
                        primary_block = (
                            "\n\n[PRIMARY RECOVERY ACTION]\n"
                            f"{strategy['primary_action'].strip()}\n"
                        )
                    suggestions = spec.experience_analyzer._format_suggestions(  # type: ignore[attr-defined]
                        strategy["experiences"],
                        strategy["error_type"],
                    )
                    return primary_block + "\n\n" + suggestions
            return _HINT
        prepare_call = getattr(spec.tools, "prepare_call", None)
        tool, params, prep_error = None, tool_call.arguments, None
        if callable(prepare_call):
            try:
                prepared = prepare_call(tool_call.name, tool_call.arguments)
                if isinstance(prepared, tuple) and len(prepared) == 3:
                    tool, params, prep_error = prepared
            except Exception:
                pass
        if prep_error:
            event = {
                "name": tool_call.name,
                "status": "error",
                "detail": prep_error.split(": ", 1)[-1][:120],
            }
            suggestions = _get_error_suggestions(tool_call.name, prep_error)
            if spec.experience_store and spec.session_key:
                related_id = (recovery_hints.get(tool_call.name) or [""])[0]
                spec.experience_store.record_tool_event(
                    tool_name=tool_call.name,
                    params=params,
                    status="error",
                    detail=prep_error,
                    session_key=spec.session_key,
                    attempt_no=attempt_no,
                    related_experience_id=related_id,
                )
            return prep_error + suggestions, event, RuntimeError(prep_error) if spec.fail_on_tool_error else None

        approval = ApprovalManager.evaluate(
            session=spec.session,
            tool=tool,
            tool_name=tool_call.name,
            params=params,
        )
        if approval.action == ApprovalAction.REQUIRE_APPROVAL and approval.request is not None:
            event = {
                "name": tool_call.name,
                "status": "approval_required",
                "detail": f"approval required: {approval.request.id}",
            }
            if spec.experience_store and spec.session_key:
                spec.experience_store.record_tool_event(
                    tool_name=tool_call.name,
                    params=params,
                    status="error",
                    detail=format_approval_required(approval.request),
                    session_key=spec.session_key,
                    attempt_no=attempt_no,
                    related_experience_id="",
                )
            if spec.fail_on_tool_error:
                result = format_approval_required(approval.request)
                return result, event, PermissionError(result)
            return "", event, None

        # Trigger on_tool_start hook before execution
        if hook and context:
            context.current_tool = tool_call.name
            await hook.on_tool_start(context, tool_call.name, params)

        try:
            if tool is not None:
                result = await tool.execute(**params)
            else:
                result = await spec.tools.execute(tool_call.name, params)

            # Trigger on_tool_end hook after successful execution
            if hook and context:
                await hook.on_tool_end(context, tool_call.name, result)
                context.current_tool = None

        except asyncio.CancelledError:
            raise
        except BaseException as exc:
            # Trigger on_error hook on exception
            if hook and context:
                await hook.on_error(context, exc)
                context.current_tool = None

            event = {
                "name": tool_call.name,
                "status": "error",
                "detail": str(exc),
            }
            suggestions = _get_error_suggestions(tool_call.name, exc)
            if spec.experience_store and spec.session_key:
                related_id = (recovery_hints.get(tool_call.name) or [""])[0]
                spec.experience_store.record_tool_event(
                    tool_name=tool_call.name,
                    params=params,
                    status="error",
                    detail=f"{type(exc).__name__}: {exc}",
                    session_key=spec.session_key,
                    attempt_no=attempt_no,
                    related_experience_id=related_id,
                )
            if spec.fail_on_tool_error:
                return f"Error: {type(exc).__name__}: {exc}" + suggestions, event, exc
            return f"Error: {type(exc).__name__}: {exc}" + suggestions, event, None

        if isinstance(result, str) and result.startswith("Error"):
            event = {
                "name": tool_call.name,
                "status": "error",
                "detail": result.replace("\n", " ").strip()[:120],
            }
            suggestions = _get_error_suggestions(tool_call.name, result)
            if spec.experience_store and spec.session_key:
                related_id = (recovery_hints.get(tool_call.name) or [""])[0]
                spec.experience_store.record_tool_event(
                    tool_name=tool_call.name,
                    params=params,
                    status="error",
                    detail=result,
                    session_key=spec.session_key,
                    attempt_no=attempt_no,
                    related_experience_id=related_id,
                )
            if spec.fail_on_tool_error:
                return result + suggestions, event, RuntimeError(result)
            return result + suggestions, event, None

        if isinstance(result, AwaitingUserInputResult):
            return result, {"name": tool_call.name, "status": result.stop_reason, "detail": ""}, None

        detail = "" if result is None else str(result)
        detail = detail.replace("\n", " ").strip()
        if not detail:
            detail = "(empty)"
        elif len(detail) > 120:
            detail = detail[:120] + "..."
        if spec.experience_store and spec.session_key:
            related_id = (recovery_hints.get(tool_call.name) or [""])[0]
            spec.experience_store.record_tool_event(
                tool_name=tool_call.name,
                params=params,
                status="ok",
                detail=detail,
                session_key=spec.session_key,
                attempt_no=attempt_no,
                related_experience_id=related_id,
            )
            if related_id:
                recovery_hints.pop(tool_call.name, None)
        return result, {"name": tool_call.name, "status": "ok", "detail": detail}, None

    def _build_context_usage_payload(

        self,
        spec: AgentRunSpec,
        messages: list[dict[str, Any]],
        *,
        iteration: int,
        phase: str,
    ) -> dict[str, Any]:
        provider_max_tokens = getattr(getattr(self.provider, "generation", None), "max_tokens", 4096)
        max_output = spec.max_tokens if isinstance(spec.max_tokens, int) else (
            provider_max_tokens if isinstance(provider_max_tokens, int) else 4096
        )
        budget = spec.context_block_limit or (
            spec.context_window_tokens - max_output - _SNIP_SAFETY_BUFFER
            if spec.context_window_tokens else None
        )
        estimated, source = estimate_prompt_tokens_chain(
            self.provider,
            spec.model,
            messages,
            spec.tools.get_definitions(),
        )
        return {
            "phase": phase,
            "iteration": iteration,
            "tokens": estimated,
            "source": source,
            "budget": budget,
            "message_count": len(messages),
            "estimated": True,
        }

    @staticmethod
    def _build_actual_usage_payload(
        usage: dict[str, int],
        *,
        iteration: int,
        phase: str,
    ) -> dict[str, Any]:
        return {
            "phase": phase,
            "iteration": iteration,
            "tokens": int(usage.get("prompt_tokens") or 0),
            "source": "provider_usage",
            "estimated": False,
        }

    async def _emit_checkpoint(
        self,
        spec: AgentRunSpec,
        payload: dict[str, Any],
    ) -> None:
        callback = spec.checkpoint_callback
        if callback is not None:
            await callback(payload)

    async def _emit_context_usage(
        self,
        spec: AgentRunSpec,
        payload: dict[str, Any],
    ) -> None:
        callback = spec.context_usage_callback
        if callback is not None and int(payload.get("tokens") or 0) > 0:
            await callback(payload)

    @staticmethod
    def _append_final_message(messages: list[dict[str, Any]], content: str | None) -> None:

        if not content:
            return
        if (
            messages
            and messages[-1].get("role") == "assistant"
            and not messages[-1].get("tool_calls")
        ):
            if messages[-1].get("content") == content:
                return
            messages[-1] = build_assistant_message(content)
            return
        messages.append(build_assistant_message(content))

    def _normalize_tool_result(
        self,
        spec: AgentRunSpec,
        tool_call_id: str,
        tool_name: str,
        result: Any,
    ) -> Any:
        result = ensure_nonempty_tool_result(tool_name, result)
        try:
            content = maybe_persist_tool_result(
                spec.workspace,
                spec.session_key,
                tool_call_id,
                result,
                max_chars=spec.max_tool_result_chars,
            )
        except Exception as exc:
            logger.warning(
                "Tool result persist failed for {} in {}: {}; using raw result",
                tool_call_id,
                spec.session_key or "default",
                exc,
            )
            content = result
        if isinstance(content, str) and len(content) > spec.max_tool_result_chars:
            return truncate_text(content, spec.max_tool_result_chars)
        return content

    def _apply_tool_result_budget(
        self,
        spec: AgentRunSpec,
        messages: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        updated = messages
        for idx, message in enumerate(messages):
            if message.get("role") != "tool":
                continue
            normalized = self._normalize_tool_result(
                spec,
                str(message.get("tool_call_id") or f"tool_{idx}"),
                str(message.get("name") or "tool"),
                message.get("content"),
            )
            if normalized != message.get("content"):
                if updated is messages:
                    updated = [dict(m) for m in messages]
                updated[idx]["content"] = normalized
        return updated

    def _snip_history(
        self,
        spec: AgentRunSpec,
        messages: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        if not messages or not spec.context_window_tokens:
            return messages

        provider_max_tokens = getattr(getattr(self.provider, "generation", None), "max_tokens", 4096)
        max_output = spec.max_tokens if isinstance(spec.max_tokens, int) else (
            provider_max_tokens if isinstance(provider_max_tokens, int) else 4096
        )
        budget = spec.context_block_limit or (
            spec.context_window_tokens - max_output - _SNIP_SAFETY_BUFFER
        )
        if budget <= 0:
            return messages

        estimate, _ = estimate_prompt_tokens_chain(
            self.provider,
            spec.model,
            messages,
            spec.tools.get_definitions(),
        )
        if estimate <= budget:
            return messages

        # Only the first system message is the durable agent contract. Later
        # system messages are dynamic retrieval/context blocks, so they must be
        # allowed to fall out of the prompt when the turn is over budget.
        core_system: list[dict[str, Any]] = []
        candidates: list[dict[str, Any]] = []
        seen_core_system = False
        for msg in messages:
            entry = dict(msg)
            if entry.get("role") == "system" and not seen_core_system:
                core_system.append(entry)
                seen_core_system = True
            else:
                candidates.append(entry)

        if not candidates:
            return messages

        core_system_tokens = sum(
            estimate_message_tokens(msg, model=spec.model) for msg in core_system
        )
        remaining_budget = max(128, budget - core_system_tokens)
        kept: list[dict[str, Any]] = []
        kept_tokens = 0
        for message in reversed(candidates):
            msg_tokens = estimate_message_tokens(message, model=spec.model)

            if kept and kept_tokens + msg_tokens > remaining_budget:
                break
            kept.append(message)
            kept_tokens += msg_tokens
        kept.reverse()

        if kept:
            kept_dynamic_system = [msg for msg in kept if msg.get("role") == "system"]
            kept_turn_messages = [msg for msg in kept if msg.get("role") != "system"]
            for i, message in enumerate(kept_turn_messages):
                if message.get("role") == "user":
                    kept_turn_messages = kept_turn_messages[i:]
                    break
            start = find_legal_message_start(kept_turn_messages)
            if start:
                kept_turn_messages = kept_turn_messages[start:]
            kept = kept_dynamic_system + kept_turn_messages
        if not kept:
            fallback = [msg for msg in candidates if msg.get("role") != "system"]
            kept = fallback[-min(len(fallback), 4) :]
            start = find_legal_message_start(kept)
            if start:
                kept = kept[start:]
        return core_system + kept

    def _partition_tool_batches(
        self,
        spec: AgentRunSpec,
        tool_calls: list[ToolCallRequest],
    ) -> list[list[ToolCallRequest]]:
        if not spec.concurrent_tools:
            return [[tool_call] for tool_call in tool_calls]

        batches: list[list[ToolCallRequest]] = []
        current: list[ToolCallRequest] = []
        for tool_call in tool_calls:
            get_tool = getattr(spec.tools, "get", None)
            tool = get_tool(tool_call.name) if callable(get_tool) else None
            can_batch = bool(tool and tool.concurrency_safe)
            if can_batch:
                current.append(tool_call)
                continue
            if current:
                batches.append(current)
                current = []
            batches.append([tool_call])
        if current:
            batches.append(current)
        return batches
