Agent/agent/core/runner.py

"""
Agent Runner - Agent 执行引擎

核心职责：
1. 执行 Agent 任务（循环调用 LLM + 工具）
2. 记录执行轨迹（Trace + Messages + GoalTree）
3. 检索和注入记忆（Experience + Skill）
4. 管理执行计划（GoalTree）
5. 支持续跑（continue）和回溯重跑（rewind）

参数分层：
- Infrastructure: AgentRunner 构造时设置（trace_store, llm_call 等）
- RunConfig: 每次 run 时指定（model, trace_id, after_sequence 等）
- Messages: OpenAI SDK 格式的任务消息
"""

import asyncio
import json
import logging
import os
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import AsyncIterator, Optional, Dict, Any, List, Callable, Literal, Tuple, Union

from agent.trace.models import Trace, Message
from agent.trace.protocols import TraceStore
from agent.trace.goal_models import GoalTree
from agent.memory.models import Skill
from agent.memory.protocols import MemoryStore, StateStore
from agent.memory.skill_loader import load_skills_from_dir
from agent.tools import ToolRegistry, get_tool_registry

logger = logging.getLogger(__name__)


# ===== 运行配置 =====

@dataclass
class RunConfig:
    """
    运行参数 — 控制 Agent 如何执行

    分为模型层参数（由上游 agent 或用户决定）和框架层参数（由系统注入）。
    """
    # --- 模型层参数 ---
    model: str = "gpt-4o"
    temperature: float = 0.3
    max_iterations: int = 200
    tools: Optional[List[str]] = None          # None = 全部已注册工具

    # --- 框架层参数 ---
    agent_type: str = "default"
    uid: Optional[str] = None
    system_prompt: Optional[str] = None        # None = 从 skills 自动构建
    enable_memory: bool = True
    auto_execute_tools: bool = True
    name: Optional[str] = None                 # 显示名称（空则由 utility_llm 自动生成）

    # --- Trace 控制 ---
    trace_id: Optional[str] = None             # None = 新建
    parent_trace_id: Optional[str] = None      # 子 Agent 专用
    parent_goal_id: Optional[str] = None

    # --- 续跑控制 ---
    after_sequence: Optional[int] = None       # 从哪条消息后续跑（message sequence）

    # --- 额外 LLM 参数（传给 llm_call 的 **kwargs）---
    extra_llm_params: Dict[str, Any] = field(default_factory=dict)


# 内置工具列表（始终自动加载）
BUILTIN_TOOLS = [
    # 文件操作工具
    "read_file",
    "edit_file",
    "write_file",
    "glob_files",
    "grep_content",

    # 系统工具
    "bash_command",

    # 技能和目标管理
    "skill",
    "list_skills",
    "goal",
    "agent",
    "evaluate",

    # 搜索工具
    "search_posts",
    "get_search_suggestions",

    # 沙箱工具
    "sandbox_create_environment",
    "sandbox_run_shell",
    "sandbox_rebuild_with_ports",
    "sandbox_destroy_environment",

    # 浏览器工具
    "browser_navigate_to_url",
    "browser_search_web",
    "browser_go_back",
    "browser_wait",
    "browser_click_element",
    "browser_input_text",
    "browser_send_keys",
    "browser_upload_file",
    "browser_scroll_page",
    "browser_find_text",
    "browser_screenshot",
    "browser_switch_tab",
    "browser_close_tab",
    "browser_get_dropdown_options",
    "browser_select_dropdown_option",
    "browser_extract_content",
    "browser_read_long_content",
    "browser_download_direct_url",
    "browser_get_page_html",
    "browser_get_visual_selector_map",
    "browser_evaluate",
    "browser_ensure_login_with_cookies",
    "browser_wait_for_user_action",
    "browser_done",
    "browser_export_cookies",
    "browser_load_cookies"
]


# ===== 向后兼容 =====

@dataclass
class AgentConfig:
    """[向后兼容] Agent 配置，新代码请使用 RunConfig"""
    agent_type: str = "default"
    max_iterations: int = 200
    enable_memory: bool = True
    auto_execute_tools: bool = True


@dataclass
class CallResult:
    """单次调用结果"""
    reply: str
    tool_calls: Optional[List[Dict]] = None
    trace_id: Optional[str] = None
    step_id: Optional[str] = None
    tokens: Optional[Dict[str, int]] = None
    cost: float = 0.0


# ===== 执行引擎 =====

CONTEXT_INJECTION_INTERVAL = 10  # 每 N 轮注入一次 GoalTree + Collaborators


class AgentRunner:
    """
    Agent 执行引擎

    支持三种运行模式（通过 RunConfig 区分）：
    1. 新建：trace_id=None
    2. 续跑：trace_id=已有ID, after_sequence=None 或 == head
    3. 回溯：trace_id=已有ID, after_sequence=N（N < head_sequence）
    """

    def __init__(
        self,
        trace_store: Optional[TraceStore] = None,
        memory_store: Optional[MemoryStore] = None,
        state_store: Optional[StateStore] = None,
        tool_registry: Optional[ToolRegistry] = None,
        llm_call: Optional[Callable] = None,
        utility_llm_call: Optional[Callable] = None,
        config: Optional[AgentConfig] = None,
        skills_dir: Optional[str] = None,
        experiences_path: Optional[str] = "./cache/experiences.md",
        goal_tree: Optional[GoalTree] = None,
        debug: bool = False,
    ):
        """
        初始化 AgentRunner

        Args:
            trace_store: Trace 存储
            memory_store: Memory 存储（可选）
            state_store: State 存储（可选）
            tool_registry: 工具注册表（默认使用全局注册表）
            llm_call: 主 LLM 调用函数
            utility_llm_call: 轻量 LLM（用于生成任务标题等），可选
            config: [向后兼容] AgentConfig
            skills_dir: Skills 目录路径
            experiences_path: 经验文件路径（默认 ./cache/experiences.md）
            goal_tree: 初始 GoalTree（可选）
            debug: 保留参数（已废弃）
        """
        self.trace_store = trace_store
        self.memory_store = memory_store
        self.state_store = state_store
        self.tools = tool_registry or get_tool_registry()
        self.llm_call = llm_call
        self.utility_llm_call = utility_llm_call
        self.config = config or AgentConfig()
        self.skills_dir = skills_dir
        self.experiences_path = experiences_path
        self.goal_tree = goal_tree
        self.debug = debug
        self._cancel_events: Dict[str, asyncio.Event] = {}  # trace_id → cancel event

    # ===== 核心公开方法 =====

    async def run(
        self,
        messages: List[Dict],
        config: Optional[RunConfig] = None,
    ) -> AsyncIterator[Union[Trace, Message]]:
        """
        Agent 模式执行（核心方法）

        Args:
            messages: OpenAI SDK 格式的输入消息
                新建: 初始任务消息 [{"role": "user", "content": "..."}]
                续跑: 追加的新消息
                回溯: 在插入点之后追加的消息
            config: 运行配置

        Yields:
            Union[Trace, Message]: Trace 对象（状态变化）或 Message 对象（执行过程）
        """
        if not self.llm_call:
            raise ValueError("llm_call function not provided")

        config = config or RunConfig()
        trace = None

        try:
            # Phase 1: PREPARE TRACE
            trace, goal_tree, sequence = await self._prepare_trace(messages, config)
            # 注册取消事件
            self._cancel_events[trace.trace_id] = asyncio.Event()
            yield trace

            # Phase 2: BUILD HISTORY
            history, sequence, created_messages, head_seq = await self._build_history(
                trace.trace_id, messages, goal_tree, config, sequence
            )
            # Update trace's head_sequence in memory
            trace.head_sequence = head_seq
            for msg in created_messages:
                yield msg

            # Phase 3: AGENT LOOP
            async for event in self._agent_loop(trace, history, goal_tree, config, sequence):
                yield event

        except Exception as e:
            logger.error(f"Agent run failed: {e}")
            tid = config.trace_id or (trace.trace_id if trace else None)
            if self.trace_store and tid:
                await self.trace_store.update_trace(
                    tid,
                    status="failed",
                    error_message=str(e),
                    completed_at=datetime.now()
                )
                trace_obj = await self.trace_store.get_trace(tid)
                if trace_obj:
                    yield trace_obj
            raise
        finally:
            # 清理取消事件
            if trace:
                self._cancel_events.pop(trace.trace_id, None)

    async def run_result(
        self,
        messages: List[Dict],
        config: Optional[RunConfig] = None,
    ) -> Dict[str, Any]:
        """
        结果模式 — 消费 run()，返回结构化结果。

        主要用于 agent/evaluate 工具内部。
        """
        last_assistant_text = ""
        final_trace: Optional[Trace] = None

        async for item in self.run(messages=messages, config=config):
            if isinstance(item, Message) and item.role == "assistant":
                content = item.content
                text = ""
                if isinstance(content, dict):
                    text = content.get("text", "") or ""
                elif isinstance(content, str):
                    text = content
                if text and text.strip():
                    last_assistant_text = text
            elif isinstance(item, Trace):
                final_trace = item

        config = config or RunConfig()
        if not final_trace and config.trace_id and self.trace_store:
            final_trace = await self.trace_store.get_trace(config.trace_id)

        status = final_trace.status if final_trace else "unknown"
        error = final_trace.error_message if final_trace else None
        summary = last_assistant_text

        if not summary:
            status = "failed"
            error = error or "Agent 没有产生 assistant 文本结果"

        return {
            "status": status,
            "summary": summary,
            "trace_id": final_trace.trace_id if final_trace else config.trace_id,
            "error": error,
            "stats": {
                "total_messages": final_trace.total_messages if final_trace else 0,
                "total_tokens": final_trace.total_tokens if final_trace else 0,
                "total_cost": final_trace.total_cost if final_trace else 0.0,
            },
        }

    async def stop(self, trace_id: str) -> bool:
        """
        停止运行中的 Trace

        设置取消信号，agent loop 在下一个 LLM 调用前检查并退出。
        Trace 状态置为 "stopped"。

        Returns:
            True 如果成功发送停止信号，False 如果该 trace 不在运行中
        """
        cancel_event = self._cancel_events.get(trace_id)
        if cancel_event is None:
            return False
        cancel_event.set()
        return True

    # ===== 单次调用（保留）=====

    async def call(
        self,
        messages: List[Dict],
        model: str = "gpt-4o",
        tools: Optional[List[str]] = None,
        uid: Optional[str] = None,
        trace: bool = True,
        **kwargs
    ) -> CallResult:
        """
        单次 LLM 调用（无 Agent Loop）
        """
        if not self.llm_call:
            raise ValueError("llm_call function not provided")

        trace_id = None
        message_id = None

        tool_schemas = self._get_tool_schemas(tools)

        if trace and self.trace_store:
            trace_obj = Trace.create(mode="call", uid=uid, model=model, tools=tool_schemas, llm_params=kwargs)
            trace_id = await self.trace_store.create_trace(trace_obj)

        result = await self.llm_call(messages=messages, model=model, tools=tool_schemas, **kwargs)

        if trace and self.trace_store and trace_id:
            msg = Message.create(
                trace_id=trace_id, role="assistant", sequence=1, goal_id=None,
                content={"text": result.get("content", ""), "tool_calls": result.get("tool_calls")},
                prompt_tokens=result.get("prompt_tokens", 0),
                completion_tokens=result.get("completion_tokens", 0),
                finish_reason=result.get("finish_reason"),
                cost=result.get("cost", 0),
            )
            message_id = await self.trace_store.add_message(msg)
            await self.trace_store.update_trace(trace_id, status="completed", completed_at=datetime.now())

        return CallResult(
            reply=result.get("content", ""),
            tool_calls=result.get("tool_calls"),
            trace_id=trace_id,
            step_id=message_id,
            tokens={"prompt": result.get("prompt_tokens", 0), "completion": result.get("completion_tokens", 0)},
            cost=result.get("cost", 0)
        )

    # ===== Phase 1: PREPARE TRACE =====

    async def _prepare_trace(
        self,
        messages: List[Dict],
        config: RunConfig,
    ) -> Tuple[Trace, Optional[GoalTree], int]:
        """
        准备 Trace：创建新的或加载已有的

        Returns:
            (trace, goal_tree, next_sequence)
        """
        if config.trace_id:
            return await self._prepare_existing_trace(config)
        else:
            return await self._prepare_new_trace(messages, config)

    async def _prepare_new_trace(
        self,
        messages: List[Dict],
        config: RunConfig,
    ) -> Tuple[Trace, Optional[GoalTree], int]:
        """创建新 Trace"""
        trace_id = str(uuid.uuid4())

        # 生成任务名称
        task_name = config.name or await self._generate_task_name(messages)

        # 准备工具 Schema
        tool_schemas = self._get_tool_schemas(config.tools)

        trace_obj = Trace(
            trace_id=trace_id,
            mode="agent",
            task=task_name,
            agent_type=config.agent_type,
            parent_trace_id=config.parent_trace_id,
            parent_goal_id=config.parent_goal_id,
            uid=config.uid,
            model=config.model,
            tools=tool_schemas,
            llm_params={"temperature": config.temperature, **config.extra_llm_params},
            status="running",
        )

        goal_tree = self.goal_tree or GoalTree(mission=task_name)

        if self.trace_store:
            await self.trace_store.create_trace(trace_obj)
            await self.trace_store.update_goal_tree(trace_id, goal_tree)

        return trace_obj, goal_tree, 1

    async def _prepare_existing_trace(
        self,
        config: RunConfig,
    ) -> Tuple[Trace, Optional[GoalTree], int]:
        """加载已有 Trace（续跑或回溯）"""
        if not self.trace_store:
            raise ValueError("trace_store required for continue/rewind")

        trace_obj = await self.trace_store.get_trace(config.trace_id)
        if not trace_obj:
            raise ValueError(f"Trace not found: {config.trace_id}")

        goal_tree = await self.trace_store.get_goal_tree(config.trace_id)

        # 自动判断行为：after_sequence 为 None 或 == head → 续跑；< head → 回溯
        after_seq = config.after_sequence
        if after_seq is not None and after_seq < trace_obj.head_sequence:
            # 回溯模式
            sequence = await self._rewind(config.trace_id, after_seq, goal_tree)
        else:
            # 续跑模式：从 last_sequence + 1 开始
            sequence = trace_obj.last_sequence + 1

        # 状态置为 running
        await self.trace_store.update_trace(
            config.trace_id,
            status="running",
            completed_at=None,
        )
        trace_obj.status = "running"

        return trace_obj, goal_tree, sequence

    # ===== Phase 2: BUILD HISTORY =====

    async def _build_history(
        self,
        trace_id: str,
        new_messages: List[Dict],
        goal_tree: Optional[GoalTree],
        config: RunConfig,
        sequence: int,
    ) -> Tuple[List[Dict], int, List[Message]]:
        """
        构建完整的 LLM 消息历史

        1. 从 head_sequence 沿 parent chain 加载主路径消息（续跑/回溯场景）
        2. 构建 system prompt（新建时注入 skills）
        3. 新建时：在第一条 user message 末尾注入当前经验
        4. 追加 input messages（设置 parent_sequence 链接到当前 head）

        Returns:
            (history, next_sequence, created_messages, head_sequence)
            created_messages: 本次新创建并持久化的 Message 列表，供 run() yield 给调用方
            head_sequence: 当前主路径头节点的 sequence
        """
        history: List[Dict] = []
        created_messages: List[Message] = []
        head_seq: Optional[int] = None  # 当前主路径的头节点 sequence

        # 1. 加载已有 messages（通过主路径遍历）
        if config.trace_id and self.trace_store:
            trace_obj = await self.trace_store.get_trace(trace_id)
            if trace_obj and trace_obj.head_sequence > 0:
                main_path = await self.trace_store.get_main_path_messages(
                    trace_id, trace_obj.head_sequence
                )
                history = [msg.to_llm_dict() for msg in main_path]
                if main_path:
                    head_seq = main_path[-1].sequence

        # 2. 构建 system prompt（如果历史中没有 system message）
        has_system = any(m.get("role") == "system" for m in history)
        has_system_in_new = any(m.get("role") == "system" for m in new_messages)

        if not has_system and not has_system_in_new:
            system_prompt = await self._build_system_prompt(config)
            if system_prompt:
                history = [{"role": "system", "content": system_prompt}] + history

                if self.trace_store:
                    system_msg = Message.create(
                        trace_id=trace_id, role="system", sequence=sequence,
                        goal_id=None, content=system_prompt,
                        parent_sequence=None,  # system message 是 root
                    )
                    await self.trace_store.add_message(system_msg)
                    created_messages.append(system_msg)
                    head_seq = sequence
                    sequence += 1

        # 3. 新建时：在第一条 user message 末尾注入当前经验
        if not config.trace_id:  # 新建模式
            experiences_text = self._load_experiences()
            if experiences_text:
                for msg in new_messages:
                    if msg.get("role") == "user" and isinstance(msg.get("content"), str):
                        msg["content"] += f"\n\n## 参考经验\n\n{experiences_text}"
                        break

        # 4. 追加新 messages（设置 parent_sequence 链接到当前 head）
        for msg_dict in new_messages:
            history.append(msg_dict)

            if self.trace_store:
                stored_msg = Message.from_llm_dict(
                    msg_dict, trace_id=trace_id, sequence=sequence,
                    goal_id=None, parent_sequence=head_seq,
                )
                await self.trace_store.add_message(stored_msg)
                created_messages.append(stored_msg)
                head_seq = sequence
                sequence += 1

        # 5. 更新 trace 的 head_sequence
        if self.trace_store and head_seq is not None:
            await self.trace_store.update_trace(trace_id, head_sequence=head_seq)

        return history, sequence, created_messages, head_seq or 0

    # ===== Phase 3: AGENT LOOP =====

    async def _agent_loop(
        self,
        trace: Trace,
        history: List[Dict],
        goal_tree: Optional[GoalTree],
        config: RunConfig,
        sequence: int,
    ) -> AsyncIterator[Union[Trace, Message]]:
        """ReAct 循环"""
        trace_id = trace.trace_id
        tool_schemas = self._get_tool_schemas(config.tools)

        # 当前主路径头节点的 sequence（用于设置 parent_sequence）
        head_seq = trace.head_sequence

        # 设置 goal_tree 到 goal 工具
        if goal_tree and self.trace_store:
            from agent.trace.goal_tool import set_goal_tree
            set_goal_tree(goal_tree)

        for iteration in range(config.max_iterations):
            # 检查取消信号
            cancel_event = self._cancel_events.get(trace_id)
            if cancel_event and cancel_event.is_set():
                logger.info(f"Trace {trace_id} stopped by user")
                if self.trace_store:
                    await self.trace_store.update_trace(
                        trace_id,
                        status="stopped",
                        completed_at=datetime.now(),
                    )
                    trace_obj = await self.trace_store.get_trace(trace_id)
                    if trace_obj:
                        yield trace_obj
                return

            # 构建 LLM messages（注入上下文）
            llm_messages = list(history)

            # 周期性注入 GoalTree + Collaborators
            if iteration % CONTEXT_INJECTION_INTERVAL == 0:
                context_injection = self._build_context_injection(trace, goal_tree)
                if context_injection:
                    llm_messages.append({"role": "system", "content": context_injection})

            # 调用 LLM
            result = await self.llm_call(
                messages=llm_messages,
                model=config.model,
                tools=tool_schemas,
                temperature=config.temperature,
                **config.extra_llm_params,
            )

            response_content = result.get("content", "")
            tool_calls = result.get("tool_calls")
            finish_reason = result.get("finish_reason")
            prompt_tokens = result.get("prompt_tokens", 0)
            completion_tokens = result.get("completion_tokens", 0)
            step_cost = result.get("cost", 0)

            # 按需自动创建 root goal
            if goal_tree and not goal_tree.goals and tool_calls:
                has_goal_call = any(
                    tc.get("function", {}).get("name") == "goal"
                    for tc in tool_calls
                )
                if not has_goal_call:
                    mission = goal_tree.mission
                    root_desc = mission[:200] if len(mission) > 200 else mission
                    goal_tree.add_goals(
                        descriptions=[root_desc],
                        reasons=["系统自动创建：Agent 未显式创建目标"],
                        parent_id=None
                    )
                    goal_tree.focus(goal_tree.goals[0].id)
                    if self.trace_store:
                        await self.trace_store.update_goal_tree(trace_id, goal_tree)
                        await self.trace_store.add_goal(trace_id, goal_tree.goals[0])
                    logger.info(f"自动创建 root goal: {goal_tree.goals[0].id}")

            # 获取当前 goal_id
            current_goal_id = goal_tree.current_id if (goal_tree and goal_tree.current_id) else None

            # 记录 assistant Message（parent_sequence 指向当前 head）
            assistant_msg = Message.create(
                trace_id=trace_id,
                role="assistant",
                sequence=sequence,
                goal_id=current_goal_id,
                parent_sequence=head_seq if head_seq > 0 else None,
                content={"text": response_content, "tool_calls": tool_calls},
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                finish_reason=finish_reason,
                cost=step_cost,
            )

            if self.trace_store:
                await self.trace_store.add_message(assistant_msg)

            yield assistant_msg
            head_seq = sequence
            sequence += 1

            # 处理工具调用
            if tool_calls and config.auto_execute_tools:
                history.append({
                    "role": "assistant",
                    "content": response_content,
                    "tool_calls": tool_calls,
                })

                for tc in tool_calls:
                    current_goal_id = goal_tree.current_id if (goal_tree and goal_tree.current_id) else None

                    tool_name = tc["function"]["name"]
                    tool_args = tc["function"]["arguments"]

                    if isinstance(tool_args, str):
                        tool_args = json.loads(tool_args) if tool_args.strip() else {}
                    elif tool_args is None:
                        tool_args = {}

                    tool_result = await self.tools.execute(
                        tool_name,
                        tool_args,
                        uid=config.uid or "",
                        context={
                            "store": self.trace_store,
                            "trace_id": trace_id,
                            "goal_id": current_goal_id,
                            "runner": self,
                        }
                    )

                    # --- 支持多模态工具反馈 ---
                    # execute() 返回 dict{"text","images"} 或 str
                    if isinstance(tool_result, dict) and tool_result.get("images"):
                        tool_result_text = tool_result["text"]
                        # 构建多模态消息格式
                        tool_content_for_llm = [{"type": "text", "text": tool_result_text}]
                        for img in tool_result["images"]:
                            if img.get("type") == "base64" and img.get("data"):
                                media_type = img.get("media_type", "image/png")
                                tool_content_for_llm.append({
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:{media_type};base64,{img['data']}"
                                    }
                                })
                        img_count = len(tool_content_for_llm) - 1  # 减去 text 块
                        print(f"[Runner] 多模态工具反馈: tool={tool_name}, images={img_count}, text_len={len(tool_result_text)}")
                    else:
                        tool_result_text = str(tool_result)
                        tool_content_for_llm = tool_result_text

                    tool_msg = Message.create(
                        trace_id=trace_id,
                        role="tool",
                        sequence=sequence,
                        goal_id=current_goal_id,
                        parent_sequence=head_seq,
                        tool_call_id=tc["id"],
                        content={"tool_name": tool_name, "result": tool_result_text},
                    )

                    if self.trace_store:
                        await self.trace_store.add_message(tool_msg)
                        # 截图单独存为同名 PNG 文件
                        if isinstance(tool_result, dict) and tool_result.get("images"):
                            import base64 as b64mod
                            for img in tool_result["images"]:
                                if img.get("data"):
                                    png_path = self.trace_store._get_messages_dir(trace_id) / f"{tool_msg.message_id}.png"
                                    png_path.write_bytes(b64mod.b64decode(img["data"]))
                                    print(f"[Runner] 截图已保存: {png_path.name}")
                                    break  # 只存第一张

                    yield tool_msg
                    head_seq = sequence
                    sequence += 1

                    history.append({
                        "role": "tool",
                        "tool_call_id": tc["id"],
                        "name": tool_name,
                        "content": tool_content_for_llm, # 这里传入 list 即可触发模型的视觉能力
                    })
                    # ------------------------------------------

                continue  # 继续循环

            # 无工具调用，任务完成
            break

        # 更新 head_sequence 并完成 Trace
        if self.trace_store:
            await self.trace_store.update_trace(
                trace_id,
                status="completed",
                head_sequence=head_seq,
                completed_at=datetime.now(),
            )
            trace_obj = await self.trace_store.get_trace(trace_id)
            if trace_obj:
                yield trace_obj

    # ===== 回溯（Rewind）=====

    async def _rewind(
        self,
        trace_id: str,
        after_sequence: int,
        goal_tree: Optional[GoalTree],
    ) -> int:
        """
        执行回溯：快照 GoalTree，重建干净树，设置 head_sequence

        新消息的 parent_sequence 将指向 rewind 点，旧消息通过树结构自然脱离主路径。

        Returns:
            下一个可用的 sequence 号
        """
        if not self.trace_store:
            raise ValueError("trace_store required for rewind")

        # 1. 加载所有 messages（用于 safe cutoff 和 max sequence）
        all_messages = await self.trace_store.get_trace_messages(trace_id)

        if not all_messages:
            return 1

        # 2. 找到安全截断点（确保不截断在 tool_call 和 tool response 之间）
        cutoff = self._find_safe_cutoff(all_messages, after_sequence)

        # 3. 快照并重建 GoalTree
        if goal_tree:
            # 通过主路径消息来判断：从 cutoff 沿 parent_sequence 回溯，只检查实际在主路径上的消息
            main_path_before = await self.trace_store.get_main_path_messages(
                trace_id, cutoff
            )
            completed_goal_ids = set()
            for goal in goal_tree.goals:
                if goal.status == "completed":
                    # 检查该 goal 是否在主路径上有关联消息（即确实在 rewind 点之前就存在）
                    goal_msgs = [m for m in main_path_before if m.goal_id == goal.id]
                    if goal_msgs:
                        completed_goal_ids.add(goal.id)

            # 快照到 events（含 head_sequence 供前端感知分支切换）
            await self.trace_store.append_event(trace_id, "rewind", {
                "after_sequence": cutoff,
                "head_sequence": cutoff,
                "goal_tree_snapshot": goal_tree.to_dict(),
            })

            # 重建干净的 GoalTree
            new_tree = goal_tree.rebuild_for_rewind(completed_goal_ids)
            await self.trace_store.update_goal_tree(trace_id, new_tree)

            # 更新内存中的引用
            goal_tree.goals = new_tree.goals
            goal_tree.current_id = new_tree.current_id

        # 4. 更新 head_sequence 到 rewind 点
        await self.trace_store.update_trace(trace_id, head_sequence=cutoff)

        # 5. 返回 next sequence（全局递增，不复用）
        max_seq = max((m.sequence for m in all_messages), default=0)
        return max_seq + 1

    def _find_safe_cutoff(self, messages: List[Message], after_sequence: int) -> int:
        """
        找到安全的截断点。

        如果 after_sequence 指向一条带 tool_calls 的 assistant message，
        则自动扩展到其所有对应的 tool response 之后。
        """
        cutoff = after_sequence

        # 找到 after_sequence 对应的 message
        target_msg = None
        for msg in messages:
            if msg.sequence == after_sequence:
                target_msg = msg
                break

        if not target_msg:
            return cutoff

        # 如果是 assistant 且有 tool_calls，找到所有对应的 tool responses
        if target_msg.role == "assistant":
            content = target_msg.content
            if isinstance(content, dict) and content.get("tool_calls"):
                tool_call_ids = set()
                for tc in content["tool_calls"]:
                    if isinstance(tc, dict) and tc.get("id"):
                        tool_call_ids.add(tc["id"])

                # 找到这些 tool_call 对应的 tool messages
                for msg in messages:
                    if (msg.role == "tool" and msg.tool_call_id
                            and msg.tool_call_id in tool_call_ids):
                        cutoff = max(cutoff, msg.sequence)

        return cutoff

    # ===== 上下文注入 =====

    def _build_context_injection(
        self,
        trace: Trace,
        goal_tree: Optional[GoalTree],
    ) -> str:
        """构建周期性注入的上下文（GoalTree + Active Collaborators）"""
        parts = []

        # GoalTree
        if goal_tree and goal_tree.goals:
            parts.append(f"## Current Plan\n\n{goal_tree.to_prompt()}")

        # Active Collaborators
        collaborators = trace.context.get("collaborators", [])
        if collaborators:
            lines = ["## Active Collaborators"]
            for c in collaborators:
                status_str = c.get("status", "unknown")
                ctype = c.get("type", "agent")
                summary = c.get("summary", "")
                name = c.get("name", "unnamed")
                lines.append(f"- {name} [{ctype}, {status_str}]: {summary}")
            parts.append("\n".join(lines))

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

    # ===== 辅助方法 =====

    def _get_tool_schemas(self, tools: Optional[List[str]]) -> List[Dict]:
        """
        获取工具 Schema

        - tools=None: 使用 registry 中全部已注册工具（含内置 + 外部注册的）
        - tools=["a", "b"]: 在 BUILTIN_TOOLS 基础上追加指定工具
        """
        if tools is None:
            # 全部已注册工具
            tool_names = self.tools.get_tool_names()
        else:
            # BUILTIN_TOOLS + 显式指定的额外工具
            tool_names = BUILTIN_TOOLS.copy()
            for t in tools:
                if t not in tool_names:
                    tool_names.append(t)
        return self.tools.get_schemas(tool_names)

    # 默认 system prompt 前缀（当 config.system_prompt 和前端都未提供 system message 时使用）
    DEFAULT_SYSTEM_PREFIX = "你是最顶尖的AI助手，可以拆分并调用工具逐步解决复杂问题。"

    async def _build_system_prompt(self, config: RunConfig) -> Optional[str]:
        """构建 system prompt（注入 skills）"""
        system_prompt = config.system_prompt

        # 加载 Skills
        skills_text = ""
        skills = load_skills_from_dir(self.skills_dir)
        if skills:
            skills_text = self._format_skills(skills)

        # 拼装：有自定义 system_prompt 则用它，否则用默认前缀
        if system_prompt:
            if skills_text:
                system_prompt += f"\n\n## Skills\n{skills_text}"
        else:
            system_prompt = self.DEFAULT_SYSTEM_PREFIX
            if skills_text:
                system_prompt += f"\n\n## Skills\n{skills_text}"

        return system_prompt

    async def _generate_task_name(self, messages: List[Dict]) -> str:
        """生成任务名称：优先使用 utility_llm，fallback 到文本截取"""
        # 提取 messages 中的文本内容
        text_parts = []
        for msg in messages:
            content = msg.get("content", "")
            if isinstance(content, str):
                text_parts.append(content)
            elif isinstance(content, list):
                for part in content:
                    if isinstance(part, dict) and part.get("type") == "text":
                        text_parts.append(part.get("text", ""))
        raw_text = " ".join(text_parts).strip()

        if not raw_text:
            return "未命名任务"

        # 尝试使用 utility_llm 生成标题
        if self.utility_llm_call:
            try:
                result = await self.utility_llm_call(
                    messages=[
                        {"role": "system", "content": "用中文为以下任务生成一个简短标题（10-30字），只输出标题本身："},
                        {"role": "user", "content": raw_text[:2000]},
                    ],
                    model="gpt-4o-mini",  # 使用便宜模型
                )
                title = result.get("content", "").strip()
                if title and len(title) < 100:
                    return title
            except Exception:
                pass

        # Fallback: 截取前 50 字符
        return raw_text[:50] + ("..." if len(raw_text) > 50 else "")

    def _format_skills(self, skills: List[Skill]) -> str:
        if not skills:
            return ""
        return "\n\n".join(s.to_prompt_text() for s in skills)

    def _load_experiences(self) -> str:
        """从文件加载经验（./cache/experiences.md）"""
        if not self.experiences_path:
            return ""
        try:
            if os.path.exists(self.experiences_path):
                with open(self.experiences_path, "r", encoding="utf-8") as f:
                    return f.read().strip()
        except Exception as e:
            logger.warning(f"Failed to load experiences from {self.experiences_path}: {e}")
        return ""