Agent/cyber_agent/trace/store.py

"""
文件系统 Trace 存储实现。

Runner、Trace API 和 Sub-Agent 共享 .trace/ 目录中的血缘、消息与 Goal 数据；
Recursive 资源预算还在根 Trace 目录内维护独立的用量快照。

目录结构：
.trace/{trace_id}/
├── meta.json           # Trace 元数据
├── goal.json           # GoalTree（扁平 JSON，通过 parent_id 构建层级）
├── resource_usage.json # Recursive 根 Trace 的树级资源用量
├── messages/           # Messages（每条独立文件）
│   ├── {message_id}.json
│   └── ...
└── events.jsonl        # 事件流（WebSocket 续传）

Sub-Trace 是完全独立的 Trace，有自己的目录：
.trace/{parent_id}@{mode}-{timestamp}-{seq}/
├── meta.json           # parent_trace_id 指向父 Trace
├── goal.json
├── messages/
└── events.jsonl
"""

import asyncio
import json
import logging
import os
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any

from .models import Trace, Message
from .goal_models import GoalTree, Goal, GoalStats

logger = logging.getLogger(__name__)


class TraceStoreCorruptionError(RuntimeError):
    """A persisted Trace file exists but cannot be trusted or decoded."""


class _ReentrantAsyncLock:
    """Small task-reentrant lock used to serialize one trace's file updates."""

    def __init__(self) -> None:
        self._lock = asyncio.Lock()
        self._owner: Optional[asyncio.Task] = None
        self._depth = 0

    async def acquire(self) -> None:
        task = asyncio.current_task()
        if task is not None and self._owner is task:
            self._depth += 1
            return
        await self._lock.acquire()
        self._owner = task
        self._depth = 1

    def release(self) -> None:
        task = asyncio.current_task()
        if self._owner is not task or self._depth <= 0:
            raise RuntimeError("trace lock released by a non-owner task")
        self._depth -= 1
        if self._depth == 0:
            self._owner = None
            self._lock.release()

    async def __aenter__(self) -> "_ReentrantAsyncLock":
        await self.acquire()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        self.release()


class FileSystemTraceStore:
    """按 Trace 分目录持久化数据的 TraceStore 默认实现。

    它实现 protocols.TraceStore，由 Runner、API、Goal 工具和 ResourceBudgetController 共用。"""

    def __init__(self, base_path: str = ".trace"):
        self.base_path = Path(base_path).resolve()
        self.base_path.mkdir(parents=True, exist_ok=True)
        self._trace_locks: Dict[str, _ReentrantAsyncLock] = {}

    @staticmethod
    def _validate_component(value: str, *, field: str) -> str:
        """Validate a caller-controlled path component without narrowing valid IDs."""
        if not isinstance(value, str) or not value or value in {".", ".."}:
            raise ValueError(f"invalid {field}")
        if "\x00" in value or "/" in value or "\\" in value:
            raise ValueError(f"invalid {field}: path separators are not allowed")
        if Path(value).name != value or Path(value).is_absolute():
            raise ValueError(f"invalid {field}")
        return value

    def _lock_for(self, trace_id: str) -> _ReentrantAsyncLock:
        trace_id = self._validate_component(trace_id, field="trace_id")
        return self._trace_locks.setdefault(trace_id, _ReentrantAsyncLock())

    def _require_contained(self, path: Path) -> Path:
        resolved = path.resolve(strict=False)
        if not resolved.is_relative_to(self.base_path):
            raise ValueError(f"trace path escapes base directory: {path}")
        return path

    @staticmethod
    def _fsync_directory(directory: Path) -> None:
        try:
            fd = os.open(directory, os.O_RDONLY)
        except OSError:
            return
        try:
            os.fsync(fd)
        finally:
            os.close(fd)

    def _atomic_write_text(self, path: Path, text: str) -> None:
        """Atomically replace one UTF-8 file and durably publish its directory entry."""
        self._require_contained(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        temp_path: Optional[str] = None
        try:
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=path.parent,
                prefix=f".{path.name}.",
                suffix=".tmp",
                delete=False,
            ) as temp_file:
                temp_path = temp_file.name
                temp_file.write(text)
                temp_file.flush()
                os.fsync(temp_file.fileno())
            os.replace(temp_path, path)
            temp_path = None
            self._fsync_directory(path.parent)
        finally:
            if temp_path is not None:
                try:
                    os.unlink(temp_path)
                except FileNotFoundError:
                    pass

    def _atomic_write_json(self, path: Path, value: Any) -> None:
        self._atomic_write_text(
            path,
            json.dumps(value, indent=2, ensure_ascii=False),
        )

    def _atomic_write_bytes(self, path: Path, data: bytes) -> None:
        self._require_contained(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        temp_path: Optional[str] = None
        try:
            with tempfile.NamedTemporaryFile(
                mode="wb",
                dir=path.parent,
                prefix=f".{path.name}.",
                suffix=".tmp",
                delete=False,
            ) as temp_file:
                temp_path = temp_file.name
                temp_file.write(data)
                temp_file.flush()
                os.fsync(temp_file.fileno())
            os.replace(temp_path, path)
            temp_path = None
            self._fsync_directory(path.parent)
        finally:
            if temp_path is not None:
                try:
                    os.unlink(temp_path)
                except FileNotFoundError:
                    pass

    def _get_trace_dir(self, trace_id: str) -> Path:
        """获取 trace 目录"""
        safe_id = self._validate_component(trace_id, field="trace_id")
        return self._require_contained(self.base_path / safe_id)

    def _get_meta_file(self, trace_id: str) -> Path:
        """获取 meta.json 文件路径"""
        return self._get_trace_dir(trace_id) / "meta.json"

    def _get_goal_file(self, trace_id: str) -> Path:
        """获取 goal.json 文件路径"""
        return self._get_trace_dir(trace_id) / "goal.json"

    def _get_messages_dir(self, trace_id: str) -> Path:
        """获取 messages 目录"""
        return self._get_trace_dir(trace_id) / "messages"

    def _get_attachments_dir(self, trace_id: str) -> Path:
        """Keep binary attachments outside the ``*.json`` message namespace."""
        return self._get_trace_dir(trace_id) / "attachments"

    def _get_message_file(self, trace_id: str, message_id: str) -> Path:
        """获取 message 文件路径"""
        safe_id = self._validate_component(message_id, field="message_id")
        return self._require_contained(
            self._get_messages_dir(trace_id) / f"{safe_id}.json"
        )

    def _get_events_file(self, trace_id: str) -> Path:
        """获取 events.jsonl 文件路径"""
        return self._get_trace_dir(trace_id) / "events.jsonl"

    def _get_model_usage_file(self, trace_id: str) -> Path:
        """获取 model_usage.json 文件路径"""
        return self._get_trace_dir(trace_id) / "model_usage.json"

    def _get_resource_usage_file(self, root_trace_id: str) -> Path:
        """获取根 Trace 的 resource_usage.json 文件路径。"""
        return self._get_trace_dir(root_trace_id) / "resource_usage.json"

    def _get_candidate_ledger_file(self, root_trace_id: str) -> Path:
        return self._get_trace_dir(root_trace_id) / "candidate_ledger.json"

    def _get_tool_approvals_file(self, trace_id: str) -> Path:
        return self._get_trace_dir(trace_id) / "tool_approvals.json"

    # ===== Trace 操作 =====

    async def create_trace(self, trace: Trace) -> str:
        """创建新的 Trace"""
        async with self._lock_for(trace.trace_id):
            trace_dir = self._get_trace_dir(trace.trace_id)
            trace_dir.mkdir(exist_ok=True)

            # 创建 messages 目录
            messages_dir = self._get_messages_dir(trace.trace_id)
            messages_dir.mkdir(exist_ok=True)

            # 写入 meta.json
            self._atomic_write_json(
                self._get_meta_file(trace.trace_id),
                trace.to_dict(),
            )

            # 创建空的 events.jsonl without truncating a pre-existing journal.
            events_file = self._get_events_file(trace.trace_id)
            events_file.touch(exist_ok=True)

        return trace.trace_id

    async def get_trace(self, trace_id: str) -> Optional[Trace]:
        """获取 Trace"""
        meta_file = self._get_meta_file(trace_id)
        if not meta_file.exists():
            return None

        try:
            data = json.loads(meta_file.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError, TypeError) as exc:
            raise TraceStoreCorruptionError(
                f"corrupt trace metadata: {meta_file}"
            ) from exc
        if not isinstance(data, dict):
            raise TraceStoreCorruptionError(
                f"trace metadata must be a JSON object: {meta_file}"
            )

        # 解析 datetime 字段
        if data.get("created_at"):
            data["created_at"] = datetime.fromisoformat(data["created_at"])
        if data.get("completed_at"):
            data["completed_at"] = datetime.fromisoformat(data["completed_at"])

        return Trace.from_dict(data)

    async def update_trace(self, trace_id: str, **updates) -> None:
        """更新 Trace"""
        async with self._lock_for(trace_id):
            trace = await self.get_trace(trace_id)
            if not trace:
                return

            # 更新字段
            for key, value in updates.items():
                if hasattr(trace, key):
                    setattr(trace, key, value)

            self._atomic_write_json(
                self._get_meta_file(trace_id),
                trace.to_dict(),
            )

    async def list_traces(
        self,
        mode: Optional[str] = None,
        agent_type: Optional[str] = None,
        uid: Optional[str] = None,
        status: Optional[str] = None,
        limit: int = 50,
        parent_trace_id: Optional[str] = None,
        created_by_tool: Optional[str] = None,
    ) -> List[Trace]:
        """扫描 meta.json 并在截断前完成所有过滤。

        Sub-Agent Spawn Guard 依赖父 Trace 和 created_by_tool 条件，避免 Validator 混入六孩子配额。"""
        traces = []

        if not self.base_path.exists():
            return []

        for trace_dir in self.base_path.iterdir():
            if not trace_dir.is_dir():
                continue

            try:
                self._require_contained(trace_dir)
            except ValueError as exc:
                raise TraceStoreCorruptionError(
                    f"trace directory escapes base path: {trace_dir}"
                ) from exc

            meta_file = trace_dir / "meta.json"
            if not meta_file.exists():
                continue

            try:
                data = json.loads(meta_file.read_text(encoding="utf-8"))
                if not isinstance(data, dict):
                    raise TypeError("meta.json must contain an object")

                # 过滤
                if mode and data.get("mode") != mode:
                    continue
                if agent_type and data.get("agent_type") != agent_type:
                    continue
                if uid and data.get("uid") != uid:
                    continue
                if status and data.get("status") != status:
                    continue
                if parent_trace_id and data.get("parent_trace_id") != parent_trace_id:
                    continue
                if created_by_tool and data.get("context", {}).get("created_by_tool") != created_by_tool:
                    continue

                # 解析 datetime
                if data.get("created_at"):
                    data["created_at"] = datetime.fromisoformat(data["created_at"])
                if data.get("completed_at"):
                    data["completed_at"] = datetime.fromisoformat(data["completed_at"])

                traces.append(Trace.from_dict(data))
            except Exception as exc:
                raise TraceStoreCorruptionError(
                    f"corrupt trace metadata: {meta_file}"
                ) from exc

        # 排序（最新的在前）
        traces.sort(key=lambda t: t.created_at, reverse=True)

        return traces[:limit]

    # ===== Recursive 树级资源用量 =====

    async def get_resource_usage(
        self,
        root_trace_id: str,
    ) -> Optional[Dict[str, Any]]:
        """读取根 Trace 的树级用量快照。

        ResourceBudgetController 在进程内锁中调用；文件缺失由上层按 fail closed 处理。"""
        usage_file = self._get_resource_usage_file(root_trace_id)
        if not usage_file.exists():
            return None
        try:
            data = json.loads(usage_file.read_text(encoding="utf-8"))
            if not isinstance(data, dict):
                raise TypeError("resource_usage.json must contain an object")
            return data
        except Exception as exc:
            raise TraceStoreCorruptionError(
                f"corrupt resource usage: {usage_file}"
            ) from exc

    async def replace_resource_usage(
        self,
        root_trace_id: str,
        usage: Dict[str, Any],
    ) -> None:
        """通过同目录临时文件 + os.replace 原子替换用量快照。

        ResourceBudgetController 在 Agent/LLM 预留和 Token/成本登记后调用，避免单进程留下半写入文件。"""
        async with self._lock_for(root_trace_id):
            self._atomic_write_json(
                self._get_resource_usage_file(root_trace_id),
                usage,
            )

    # ===== Application candidate ledger =====

    async def get_candidate_ledger(
        self,
        root_trace_id: str,
    ) -> Optional[Dict[str, Any]]:
        path = self._get_candidate_ledger_file(root_trace_id)
        if not path.exists():
            return None
        try:
            raw = json.loads(path.read_text(encoding="utf-8"))
            if not isinstance(raw, dict):
                raise TypeError("candidate_ledger.json must contain an object")
            return raw
        except Exception as exc:
            raise TraceStoreCorruptionError(
                f"corrupt candidate ledger: {path}"
            ) from exc

    async def replace_candidate_ledger(
        self,
        root_trace_id: str,
        ledger: Dict[str, Any],
    ) -> None:
        async with self._lock_for(root_trace_id):
            if await self.get_trace(root_trace_id) is None:
                raise ValueError(f"Root Trace not found: {root_trace_id}")
            self._atomic_write_json(
                self._get_candidate_ledger_file(root_trace_id),
                ledger,
            )

    # ===== Tool approval state =====

    async def get_tool_approval_batch(self, trace_id: str):
        from cyber_agent.tools.approval import ToolApprovalBatchV1

        path = self._get_tool_approvals_file(trace_id)
        if not path.exists():
            return None
        try:
            raw = json.loads(path.read_text(encoding="utf-8"))
            return ToolApprovalBatchV1.model_validate(raw)
        except Exception as exc:
            raise TraceStoreCorruptionError(
                f"corrupt tool approval state: {path}"
            ) from exc

    async def replace_tool_approval_batch(self, trace_id: str, batch) -> None:
        async with self._lock_for(trace_id):
            self._atomic_write_json(
                self._get_tool_approvals_file(trace_id),
                batch.model_dump(mode="json"),
            )

    async def write_message_attachment(
        self,
        trace_id: str,
        message_id: str,
        *,
        suffix: str,
        data: bytes,
    ) -> Path:
        """Atomically persist an attachment owned by one stored Message."""
        self._validate_component(message_id, field="message_id")
        if (
            not suffix.startswith(".")
            or len(suffix) < 2
            or len(suffix) > 17
            or not suffix[1:].isalnum()
        ):
            raise ValueError("invalid attachment suffix")
        async with self._lock_for(trace_id):
            if not self._get_message_file(trace_id, message_id).exists():
                raise ValueError(f"message not found for attachment: {message_id}")
            path = self._get_attachments_dir(trace_id) / f"{message_id}{suffix}"
            self._atomic_write_bytes(path, data)
            return path

    # ===== GoalTree 操作 =====

    async def get_goal_tree(self, trace_id: str) -> Optional[GoalTree]:
        """获取 GoalTree"""
        goal_file = self._get_goal_file(trace_id)
        if not goal_file.exists():
            return None

        try:
            data = json.loads(goal_file.read_text(encoding="utf-8"))
            if not isinstance(data, dict):
                raise TypeError("goal.json must contain an object")
            return GoalTree.from_dict(data)
        except Exception as exc:
            raise TraceStoreCorruptionError(
                f"corrupt goal tree: {goal_file}"
            ) from exc

    async def update_goal_tree(self, trace_id: str, tree: GoalTree) -> None:
        """更新完整 GoalTree"""
        async with self._lock_for(trace_id):
            self._atomic_write_json(
                self._get_goal_file(trace_id),
                tree.to_dict(),
            )

    async def add_goal(self, trace_id: str, goal: Goal) -> None:
        """添加 Goal 到 GoalTree"""
        async with self._lock_for(trace_id):
            tree = await self.get_goal_tree(trace_id)
            if not tree:
                return

            tree.goals.append(goal)
            await self.update_goal_tree(trace_id, tree)
            # 推送 goal_added 事件
            event_data = {
                "goal": goal.to_dict(),
                "parent_id": goal.parent_id,
            }
            await self.append_event(trace_id, "goal_added", event_data)

        # 打印详细的 goal 信息
        desc_preview = goal.description[:80] + "..." if len(goal.description) > 80 else goal.description
        print(f"[Goal Added] ID={goal.id}, Parent={goal.parent_id or 'root'}")
        print(f"  📝 {desc_preview}")
        if goal.reason:
            reason_preview = goal.reason[:60] + "..." if len(goal.reason) > 60 else goal.reason
            print(f"  💡 {reason_preview}")

    async def update_goal(
        self,
        trace_id: str,
        goal_id: str,
        *,
        cascade_completion: bool = True,
        **updates,
    ) -> None:
        """更新 Goal 并推送 goal_updated 事件。

        goal_tool、Sub-Agent 和审核工具调用它；Recursive 传入 cascade_completion=False，防止存储层代替父级完成。"""
        async with self._lock_for(trace_id):
            await self._update_goal_locked(
                trace_id,
                goal_id,
                cascade_completion=cascade_completion,
                **updates,
            )

    async def _update_goal_locked(
        self,
        trace_id: str,
        goal_id: str,
        *,
        cascade_completion: bool = True,
        **updates,
    ) -> None:
        tree = await self.get_goal_tree(trace_id)
        if not tree:
            return

        goal = tree.find(goal_id)
        if not goal:
            return

        # 更新字段
        for key, value in updates.items():
            if hasattr(goal, key):
                # 特殊处理 stats 字段（可能是 dict）
                if key in ["self_stats", "cumulative_stats"] and isinstance(value, dict):
                    value = GoalStats.from_dict(value)
                setattr(goal, key, value)

        await self.update_goal_tree(trace_id, tree)

        # 推送 goal_updated 事件
        # 如果状态变为 completed，检查是否需要级联完成父 Goal
        affected_goals = [{"goal_id": goal_id, "updates": updates}]

        if cascade_completion and updates.get("status") == "completed":
            # 检查级联完成：如果所有兄弟 Goal 都完成，父 Goal 也完成
            cascade_completed = await self._check_cascade_completion(trace_id, goal)
            affected_goals.extend(cascade_completed)

        await self.append_event(trace_id, "goal_updated", {
            "goal_id": goal_id,
            "updates": updates,
            "affected_goals": affected_goals
        })
        print(f"[DEBUG] Pushed goal_updated event: goal_id={goal_id}, updates={updates}, affected={len(affected_goals)}")

        # Goal 完成时触发知识评估
        if updates.get("status") in ["completed", "abandoned"]:
            pending = await self.get_pending_knowledge_entries(trace_id)
            if pending:
                # 在trace.context中设置标志，由runner主循环检查
                trace = await self.get_trace(trace_id)
                if trace:
                    if not trace.context:
                        trace.context = {}
                    trace.context["pending_knowledge_eval"] = True
                    trace.context["knowledge_eval_trigger"] = "goal_completion"
                    await self.update_trace(trace_id, context=trace.context)
                    logger.info(f"[Knowledge Eval] Goal {goal_id} 完成，设置评估标志，待评估知识: {len(pending)} 条")

    async def _check_cascade_completion(
        self,
        trace_id: str,
        completed_goal: Goal
    ) -> List[Dict[str, Any]]:
        """
        Legacy 的 Goal 级联完成：所有子 Goal 都 completed 或 abandoned 时自动完成父 Goal。

        update_goal() 只在 cascade_completion=True 时调用；Recursive 审核流明确禁用该路径。

        Args:
            trace_id: Trace ID
            completed_goal: 刚完成的 Goal

        Returns:
            受影响的父 Goals 列表（自动完成的）
        """
        if not completed_goal.parent_id:
            return []

        tree = await self.get_goal_tree(trace_id)
        if not tree:
            return []

        affected = []
        parent = tree.find(completed_goal.parent_id)

        if not parent:
            return []

        # 获取父 Goal 的所有子 Goal
        children = tree.get_children(parent.id)

        # 检查是否所有子 Goal 都已完成（排除 abandoned）
        all_completed = all(
            child.status in ["completed", "abandoned"]
            for child in children
        )

        if all_completed and parent.status != "completed":
            # 自动完成父 Goal
            parent.status = "completed"
            if not parent.summary:
                # 生成自动摘要
                completed_count = sum(1 for c in children if c.status == "completed")
                parent.summary = f"所有子目标已完成 ({completed_count}/{len(children)})"

            await self.update_goal_tree(trace_id, tree)

            affected.append({
                "goal_id": parent.id,
                "status": "completed",
                "summary": parent.summary,
                "cumulative_stats": parent.cumulative_stats.to_dict()
            })

            # 递归检查祖父 Goal
            grandparent_affected = await self._check_cascade_completion(trace_id, parent)
            affected.extend(grandparent_affected)

        return affected

    # ===== Message 操作 =====

    async def add_message(self, message: Message) -> str:
        async with self._lock_for(message.trace_id):
            return await self._add_message_locked(message)

    async def _add_message_locked(self, message: Message) -> str:
        """
        添加 Message

        自动更新关联 Goal 的 stats（self_stats 和祖先的 cumulative_stats）
        """
        trace_id = message.trace_id

        # 1. 写入 message 文件
        message_file = self._get_message_file(trace_id, message.message_id)
        if message_file.exists():
            raise ValueError(
                f"message_id already exists in Trace {trace_id}: {message.message_id}"
            )
        self._atomic_write_json(message_file, message.to_dict())

        # 2. 更新 trace 统计
        trace = await self.get_trace(trace_id)
        if trace:
            trace.total_messages += 1
            trace.last_sequence = max(trace.last_sequence, message.sequence)

            # 累计 tokens（完整版）
            if message.prompt_tokens:
                trace.total_prompt_tokens += message.prompt_tokens
            if message.completion_tokens:
                trace.total_completion_tokens += message.completion_tokens
            if message.reasoning_tokens:
                trace.total_reasoning_tokens += message.reasoning_tokens
            if message.cache_creation_tokens:
                trace.total_cache_creation_tokens += message.cache_creation_tokens
            if message.cache_read_tokens:
                trace.total_cache_read_tokens += message.cache_read_tokens

            # 向后兼容：也更新 total_tokens
            if message.tokens:
                trace.total_tokens += message.tokens
            elif message.prompt_tokens or message.completion_tokens:
                trace.total_tokens += (message.prompt_tokens or 0) + (message.completion_tokens or 0)

            if message.cost:
                trace.total_cost += message.cost
            if message.duration_ms:
                trace.total_duration_ms += message.duration_ms

            # 更新 Trace
            await self.update_trace(
                trace_id,
                total_messages=trace.total_messages,
                last_sequence=trace.last_sequence,
                total_tokens=trace.total_tokens,
                total_prompt_tokens=trace.total_prompt_tokens,
                total_completion_tokens=trace.total_completion_tokens,
                total_reasoning_tokens=trace.total_reasoning_tokens,
                total_cache_creation_tokens=trace.total_cache_creation_tokens,
                total_cache_read_tokens=trace.total_cache_read_tokens,
                total_cost=trace.total_cost,
                total_duration_ms=trace.total_duration_ms
            )

        # 3. 更新 Goal stats
        await self._update_goal_stats(trace_id, message)

        # 4. 追加 message_added 事件
        affected_goals = await self._get_affected_goals(trace_id, message)
        event_id = await self.append_event(trace_id, "message_added", {
            "message": message.to_dict(),
            "affected_goals": affected_goals
        })
        if event_id:
            try:
                from . import websocket as trace_ws

                await trace_ws.broadcast_message_added(
                    trace_id=trace_id,
                    event_id=event_id,
                    message_dict=message.to_dict(),
                    affected_goals=affected_goals,
                )
            except Exception:
                logger.exception("Failed to broadcast message_added (trace_id=%s, event_id=%s)", trace_id, event_id)

        return message.message_id

    async def _update_goal_stats(self, trace_id: str, message: Message) -> None:
        """更新 Goal 的 self_stats 和祖先的 cumulative_stats"""
        tree = await self.get_goal_tree(trace_id)

        if not tree:
            return

        # 找到关联的 Goal
        goal = tree.find(message.goal_id)
        if not goal:
            return

        # 更新自身 self_stats
        goal.self_stats.message_count += 1
        if message.tokens:
            goal.self_stats.total_tokens += message.tokens
        if message.cost:
            goal.self_stats.total_cost += message.cost
        # TODO: 更新 preview（工具调用摘要）

        # 更新自身 cumulative_stats
        goal.cumulative_stats.message_count += 1
        if message.tokens:
            goal.cumulative_stats.total_tokens += message.tokens
        if message.cost:
            goal.cumulative_stats.total_cost += message.cost

        # 沿祖先链向上更新 cumulative_stats
        current_goal = goal
        while current_goal.parent_id:
            parent = tree.find(current_goal.parent_id)
            if not parent:
                break

            parent.cumulative_stats.message_count += 1
            if message.tokens:
                parent.cumulative_stats.total_tokens += message.tokens
            if message.cost:
                parent.cumulative_stats.total_cost += message.cost

            current_goal = parent

        # 保存更新后的 tree
        await self.update_goal_tree(trace_id, tree)

    async def _get_affected_goals(self, trace_id: str, message: Message) -> List[Dict[str, Any]]:
        """获取受影响的 Goals（自身 + 所有祖先）"""
        tree = await self.get_goal_tree(trace_id)

        if not tree:
            return []

        goal = tree.find(message.goal_id)
        if not goal:
            return []

        affected = []

        # 添加自身（包含 self_stats 和 cumulative_stats）
        affected.append({
            "goal_id": goal.id,
            "self_stats": goal.self_stats.to_dict(),
            "cumulative_stats": goal.cumulative_stats.to_dict()
        })

        # 添加所有祖先（仅 cumulative_stats）
        current_goal = goal
        while current_goal.parent_id:
            parent = tree.find(current_goal.parent_id)
            if not parent:
                break

            affected.append({
                "goal_id": parent.id,
                "cumulative_stats": parent.cumulative_stats.to_dict()
            })

            current_goal = parent

        return affected

    async def get_message(self, message_id: str) -> Optional[Message]:
        """获取 Message（扫描所有 trace）"""
        self._validate_component(message_id, field="message_id")
        for trace_dir in self.base_path.iterdir():
            if not trace_dir.is_dir():
                continue

            # 检查 messages 目录
            try:
                message_file = self._get_message_file(trace_dir.name, message_id)
            except ValueError as exc:
                raise TraceStoreCorruptionError(
                    f"invalid trace directory in store: {trace_dir}"
                ) from exc
            if message_file.exists():
                try:
                    data = json.loads(message_file.read_text(encoding="utf-8"))
                    return Message.from_dict(data)
                except Exception as exc:
                    raise TraceStoreCorruptionError(
                        f"corrupt message file: {message_file}"
                    ) from exc

        return None

    async def get_trace_messages(
        self,
        trace_id: str,
    ) -> List[Message]:
        """获取 Trace 的所有 Messages（包含所有分支，按 sequence 排序）"""
        messages_dir = self._get_messages_dir(trace_id)

        if not messages_dir.exists():
            return []

        messages = []
        for message_file in messages_dir.glob("*.json"):
            try:
                data = json.loads(message_file.read_text(encoding="utf-8"))
                msg = Message.from_dict(data)
                messages.append(msg)
            except Exception as exc:
                raise TraceStoreCorruptionError(
                    f"corrupt message file: {message_file}"
                ) from exc

        # 按 sequence 排序
        messages.sort(key=lambda m: m.sequence)
        return messages

    async def get_main_path_messages(
        self,
        trace_id: str,
        head_sequence: int
    ) -> List[Message]:
        """
        获取从 head_sequence 沿 parent_sequence 链回溯到 root 的完整路径

        此函数是通用的路径追溯函数，返回从指定 head 到 root 的完整消息链。
        只要 trace.head_sequence 管理正确（指向主路径），此函数自然返回主路径消息。
        侧分支消息通过 parent_sequence 链自然被跳过（因为主路径的 parent 不指向侧分支）。

        Returns:
            按 sequence 正序排列的路径 Message 列表
        """
        # 加载所有消息，建立 sequence -> Message 索引
        all_messages = await self.get_trace_messages(trace_id)
        messages_by_seq = {m.sequence: m for m in all_messages}

        # 从 head 沿 parent chain 回溯
        path = []
        seq = head_sequence
        while seq is not None:
            msg = messages_by_seq.get(seq)
            if not msg:
                break
            path.append(msg)
            seq = msg.parent_sequence

        # 反转为正序（root → head）
        path.reverse()
        return path

    async def get_messages_by_goal(
        self,
        trace_id: str,
        goal_id: str
    ) -> List[Message]:
        """获取指定 Goal 关联的所有 Messages"""
        all_messages = await self.get_trace_messages(trace_id)
        return [m for m in all_messages if m.goal_id == goal_id]

    async def update_message(self, message_id: str, **updates) -> None:
        """更新 Message 字段"""
        self._validate_component(message_id, field="message_id")
        message = await self.get_message(message_id)
        if not message:
            return
        async with self._lock_for(message.trace_id):
            # Re-read after taking the trace lock so updates cannot overwrite each other.
            message = await self.get_message(message_id)
            if not message:
                return
            for key, value in updates.items():
                if hasattr(message, key):
                    setattr(message, key, value)
            self._atomic_write_json(
                self._get_message_file(message.trace_id, message_id),
                message.to_dict(),
            )

    async def abandon_messages_after(self, trace_id: str, cutoff_sequence: int) -> List[str]:
        """
        将 sequence > cutoff_sequence 的 active messages 标记为 abandoned。
        返回被 abandon 的 message_id 列表。
        """
        async with self._lock_for(trace_id):
            all_messages = await self.get_trace_messages(trace_id)
            abandoned_ids = []
            now = datetime.now()

            for msg in all_messages:
                if msg.sequence > cutoff_sequence and msg.status == "active":
                    msg.status = "abandoned"
                    msg.abandoned_at = now
                    self._atomic_write_json(
                        self._get_message_file(trace_id, msg.message_id),
                        msg.to_dict(),
                    )
                    abandoned_ids.append(msg.message_id)

            return abandoned_ids

    # ===== 模型使用追踪 =====

    async def record_model_usage(
        self,
        trace_id: str,
        sequence: int,
        role: str,
        model: str,
        prompt_tokens: int,
        completion_tokens: int,
        cache_read_tokens: int = 0,
        tool_name: Optional[str] = None,
    ) -> None:
        """
        记录模型使用情况到 model_usage.json

        Args:
            trace_id: Trace ID
            sequence: 消息序号
            role: 角色（assistant/tool）
            model: 模型名称
            prompt_tokens: 输入tokens
            completion_tokens: 输出tokens
            cache_read_tokens: 缓存读取tokens
            tool_name: 工具名称（role=tool时）
        """
        async with self._lock_for(trace_id):
            await self._record_model_usage_locked(
                trace_id=trace_id,
                sequence=sequence,
                role=role,
                model=model,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                cache_read_tokens=cache_read_tokens,
                tool_name=tool_name,
            )

    async def _record_model_usage_locked(
        self,
        trace_id: str,
        sequence: int,
        role: str,
        model: str,
        prompt_tokens: int,
        completion_tokens: int,
        cache_read_tokens: int = 0,
        tool_name: Optional[str] = None,
    ) -> None:
        usage_file = self._get_model_usage_file(trace_id)

        # 读取现有数据
        if usage_file.exists():
            try:
                data = json.loads(usage_file.read_text(encoding="utf-8"))
                if not isinstance(data, dict):
                    raise TypeError("model_usage.json must contain an object")
            except Exception as exc:
                raise TraceStoreCorruptionError(
                    f"corrupt model usage: {usage_file}"
                ) from exc
        else:
            data = {
                "summary": {
                    "total_models": 0,
                    "total_tokens": 0,
                    "total_cache_read_tokens": 0,
                    "agent_tokens": 0,
                    "tool_tokens": 0,
                },
                "models": [],
                "timeline": [],
            }

        # 更新summary
        total_tokens = prompt_tokens + completion_tokens
        data["summary"]["total_tokens"] += total_tokens
        data["summary"]["total_cache_read_tokens"] += cache_read_tokens

        if role == "assistant":
            data["summary"]["agent_tokens"] += total_tokens
            source = "agent"
        else:
            data["summary"]["tool_tokens"] += total_tokens
            source = f"tool:{tool_name}" if tool_name else "tool"

        # 更新models列表
        model_entry = None
        for m in data["models"]:
            if m["model"] == model and m["source"] == source:
                model_entry = m
                break

        if model_entry:
            model_entry["prompt_tokens"] += prompt_tokens
            model_entry["completion_tokens"] += completion_tokens
            model_entry["total_tokens"] += total_tokens
            model_entry["cache_read_tokens"] += cache_read_tokens
            model_entry["call_count"] += 1
        else:
            data["models"].append({
                "model": model,
                "source": source,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": total_tokens,
                "cache_read_tokens": cache_read_tokens,
                "call_count": 1,
            })
            data["summary"]["total_models"] = len(data["models"])

        # 添加到timeline
        timeline_entry = {
            "sequence": sequence,
            "role": role,
            "model": model,
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
        }
        if cache_read_tokens > 0:
            timeline_entry["cache_read_tokens"] = cache_read_tokens
        if tool_name:
            timeline_entry["tool_name"] = tool_name
        data["timeline"].append(timeline_entry)

        # 写回文件
        self._atomic_write_json(usage_file, data)

    # ===== 事件流操作（用于 WebSocket 断线续传）=====

    async def get_events(
        self,
        trace_id: str,
        since_event_id: int = 0
    ) -> List[Dict[str, Any]]:
        """获取事件流"""
        events_file = self._get_events_file(trace_id)
        if not events_file.exists():
            return []

        events = []
        previous_id = 0
        with events_file.open('r', encoding='utf-8') as f:
            for line_number, line in enumerate(f, start=1):
                if not line.strip():
                    continue
                try:
                    event = json.loads(line.strip())
                    event_id = int(event.get("event_id", 0))
                    if event_id <= previous_id:
                        raise ValueError("event IDs are not strictly increasing")
                    previous_id = event_id
                    if event_id > since_event_id:
                        events.append(event)
                except Exception as exc:
                    raise TraceStoreCorruptionError(
                        f"corrupt event journal {events_file}:{line_number}"
                    ) from exc

        return events

    async def append_event(
        self,
        trace_id: str,
        event_type: str,
        payload: Dict[str, Any]
    ) -> int:
        """追加事件，返回 event_id"""
        async with self._lock_for(trace_id):
            return await self._append_event_locked(trace_id, event_type, payload)

    async def _append_event_locked(
        self,
        trace_id: str,
        event_type: str,
        payload: Dict[str, Any],
    ) -> int:
        # 获取 trace 并递增 event_id
        trace = await self.get_trace(trace_id)
        if not trace:
            return 0

        trace.last_event_id += 1
        event_id = trace.last_event_id

        # 更新 trace 的 last_event_id
        await self.update_trace(trace_id, last_event_id=event_id)

        # 创建事件
        event = {
            "event_id": event_id,
            "event": event_type,
            "ts": datetime.now().isoformat(),
            **payload
        }

        # 追加到 events.jsonl
        events_file = self._get_events_file(trace_id)
        with events_file.open('a', encoding='utf-8') as f:
            f.write(json.dumps(event, ensure_ascii=False) + '\n')
            f.flush()
            os.fsync(f.fileno())

        return event_id

    async def reconcile_trace_integrity(self, trace_id: str) -> Dict[str, int]:
        """Rebuild recoverable Trace counters from durable message/event files."""
        async with self._lock_for(trace_id):
            trace = await self.get_trace(trace_id)
            if trace is None:
                raise ValueError(f"Trace not found: {trace_id}")
            messages = await self.get_trace_messages(trace_id)
            sequences = [message.sequence for message in messages]
            if len(sequences) != len(set(sequences)):
                raise TraceStoreCorruptionError(
                    f"duplicate message sequence values in Trace {trace_id}"
                )
            events = await self.get_events(trace_id)
            totals = {
                "total_messages": len(messages),
                "last_sequence": max(sequences, default=0),
                "last_event_id": max(
                    (int(event.get("event_id", 0)) for event in events),
                    default=0,
                ),
                "total_prompt_tokens": sum(m.prompt_tokens or 0 for m in messages),
                "total_completion_tokens": sum(m.completion_tokens or 0 for m in messages),
                "total_reasoning_tokens": sum(m.reasoning_tokens or 0 for m in messages),
                "total_cache_creation_tokens": sum(
                    m.cache_creation_tokens or 0 for m in messages
                ),
                "total_cache_read_tokens": sum(
                    m.cache_read_tokens or 0 for m in messages
                ),
                "total_tokens": sum(m.tokens for m in messages),
                "total_cost": sum(m.cost or 0.0 for m in messages),
                "total_duration_ms": sum(m.duration_ms or 0 for m in messages),
            }
            await self.update_trace(trace_id, **totals)
            return {
                "messages": totals["total_messages"],
                "last_sequence": totals["last_sequence"],
                "last_event_id": totals["last_event_id"],
            }

    # ===== Cognition Log 管理 =====

    def _get_cognition_log_file(self, trace_id: str) -> Path:
        """获取 cognition_log.json 文件路径"""
        return self._get_trace_dir(trace_id) / "cognition_log.json"

    def _get_knowledge_log_file(self, trace_id: str) -> Path:
        """兼容旧接口：优先使用 cognition_log，回退到 knowledge_log"""
        cognition_file = self._get_cognition_log_file(trace_id)
        if cognition_file.exists():
            return cognition_file
        legacy_file = self._get_trace_dir(trace_id) / "knowledge_log.json"
        if legacy_file.exists():
            return legacy_file
        return cognition_file  # 新建时用 cognition_log

    async def get_cognition_log(self, trace_id: str) -> Dict[str, Any]:
        """读取认知日志"""
        log_file = self._get_cognition_log_file(trace_id)
        try:
            if log_file.exists():
                data = json.loads(log_file.read_text(encoding="utf-8"))
                if not isinstance(data, dict):
                    raise TypeError("cognition log must be an object")
                return data
        except Exception as exc:
            raise TraceStoreCorruptionError(
                f"corrupt cognition log: {log_file}"
            ) from exc
        # 兼容旧格式：如果只有 knowledge_log.json，读取并转换
        legacy_file = self._get_trace_dir(trace_id) / "knowledge_log.json"
        try:
            if legacy_file.exists():
                data = json.loads(legacy_file.read_text(encoding="utf-8"))
                if not isinstance(data, dict):
                    raise TypeError("legacy knowledge log must be an object")
                return data
        except Exception as exc:
            raise TraceStoreCorruptionError(
                f"corrupt knowledge log: {legacy_file}"
            ) from exc
        return {"trace_id": trace_id, "events": []}

    async def get_knowledge_log(self, trace_id: str) -> Dict[str, Any]:
        """兼容旧接口"""
        log = await self.get_cognition_log(trace_id)
        # 旧格式用 entries，新格式用 events
        if "entries" not in log and "events" in log:
            log["entries"] = log["events"]
        return log

    async def append_cognition_event(
        self,
        trace_id: str,
        event: Dict[str, Any],
    ) -> None:
        """追加认知事件到 cognition_log.json。

        所有事件共有字段：
            type: str         事件类型（见下表）
            timestamp: str    ISO 格式时间戳（框架自动写入）

        已定义的事件类型及典型字段：

            type="query" — 知识注入查询（goal focus 时触发）
                sequence, goal_id, query, response, source_ids, sources

            type="evaluation" — 知识评估（Goal 完成/压缩前/任务结束触发）
                knowledge_id, eval_result{relevance, utility, notes}, trigger_event

            type="extraction_pending" — 反思侧分支暂存的待审核提取（Phase 1.2+）
                extraction_id, sequence, goal_id, branch_id, payload
                (payload 字段与 knowledge_save 参数一一对应)

            type="extraction_reviewed" — 人工审核决策（CLI / HTTP API 写入）
                extraction_id, decision("approve"/"edit"/"discard"), edited_payload?

            type="extraction_committed" — 已上传到 KnowHub
                extraction_id, knowledge_id

            type="reflection" — Dream 的 per-trace 反思摘要（Phase 2.4 / 3.1）
                sequence_range: [start, end]    本次反思覆盖的消息区间
                summary: str                    LLM 生成的反思摘要
                consumed_at: 可选, ISO 时间戳   当跨 trace 整合已消化此反思时写入

        其他字段可按需附加，不做强校验（演进友好）。
        """
        async with self._lock_for(trace_id):
            log = await self.get_cognition_log(trace_id)
            if "events" not in log:
                log["events"] = log.pop("entries", [])
            event = dict(event)
            event["timestamp"] = datetime.now().isoformat()
            log["events"].append(event)
            self._atomic_write_json(self._get_cognition_log_file(trace_id), log)

    async def mark_reflections_consumed(
        self,
        trace_id: str,
        reflections: List[Dict[str, Any]],
        consumed_at: str,
    ) -> int:
        """Atomically mark the specified still-pending reflection events consumed."""
        wanted = {
            (
                tuple(item.get("sequence_range") or []),
                str(item.get("summary") or ""),
            )
            for item in reflections
        }
        async with self._lock_for(trace_id):
            log = await self.get_cognition_log(trace_id)
            events = log.get("events", log.get("entries", []))
            changed = 0
            for event in events:
                key = (
                    tuple(event.get("sequence_range") or []),
                    str(event.get("summary") or ""),
                )
                if (
                    event.get("type") == "reflection"
                    and not event.get("consumed_at")
                    and key in wanted
                ):
                    event["consumed_at"] = consumed_at
                    changed += 1
            if changed:
                self._atomic_write_json(
                    self._get_cognition_log_file(trace_id),
                    log,
                )
            return changed

    async def append_knowledge_entry(
        self,
        trace_id: str,
        knowledge_id: str,
        goal_id: str,
        injected_at_sequence: int,
        task: str,
        content: str
    ) -> None:
        """兼容旧接口：追加知识注入记录（转换为 query 事件）"""
        await self.append_cognition_event(
            trace_id=trace_id,
            event={
                "type": "query",
                "sequence": injected_at_sequence,
                "goal_id": goal_id,
                "query": task,
                "response": "",
                "source_ids": [knowledge_id],
                "sources": [{"id": knowledge_id, "task": task, "content": content[:500]}],
            }
        )

    async def update_knowledge_evaluation(
        self,
        trace_id: str,
        knowledge_id: str,
        eval_result: Dict[str, Any],
        trigger_event: str
    ) -> None:
        """更新知识评估结果（兼容旧格式 + 新 cognition_log 格式）

        旧格式：更新 entries[] 中匹配 knowledge_id 的条目的 eval_result
        新格式：追加 evaluation 事件到 events[]
        """
        async with self._lock_for(trace_id):
            await self._update_knowledge_evaluation_locked(
                trace_id,
                knowledge_id,
                eval_result,
                trigger_event,
            )

    async def _update_knowledge_evaluation_locked(
        self,
        trace_id: str,
        knowledge_id: str,
        eval_result: Dict[str, Any],
        trigger_event: str,
    ) -> None:
        log = await self.get_cognition_log(trace_id)
        events = log.get("events", log.get("entries", []))

        # 旧格式兼容：直接更新 entries 中的 eval_result 字段
        if "entries" in log:
            matching = [
                (i, e) for i, e in enumerate(log["entries"])
                if e.get("knowledge_id") == knowledge_id and e.get("eval_result") is None
            ]
            if matching:
                matching.sort(key=lambda x: x[1].get("injected_at_sequence", 0), reverse=True)
                _, entry = matching[0]
                entry["eval_result"] = eval_result
                entry["evaluated_at"] = datetime.now().isoformat()
                entry["evaluated_at_trigger"] = trigger_event
                self._atomic_write_json(
                    self._get_knowledge_log_file(trace_id),
                    log,
                )
                return

        # 新格式：追加 evaluation 事件
        # 找到包含该 knowledge_id 的最近 query 事件
        query_events = [
            e for e in events
            if e.get("type") == "query" and knowledge_id in e.get("source_ids", [])
        ]
        query_sequence = query_events[-1]["sequence"] if query_events else None

        await self.append_cognition_event(
            trace_id=trace_id,
            event={
                "type": "evaluation",
                "sequence": max((e.get("sequence", 0) for e in events), default=0) + 1,
                "query_sequence": query_sequence,
                "trigger": trigger_event,
                "assessments": [
                    {"source_id": knowledge_id, "status": eval_result.get("eval_status", ""), "reason": eval_result.get("reason", "")}
                ],
            }
        )

    async def get_pending_knowledge_entries(self, trace_id: str) -> List[Dict[str, Any]]:
        """获取所有待评估的知识条目（兼容旧格式 + 新格式）"""
        log = await self.get_cognition_log(trace_id)

        # 旧格式
        if "entries" in log:
            return [e for e in log["entries"] if e.get("eval_result") is None]

        # 新格式：找没有对应 evaluation 事件的 query 事件
        events = log.get("events", [])
        query_events = [e for e in events if e.get("type") == "query"]
        eval_events = [e for e in events if e.get("type") == "evaluation"]

        # 已评估的 query sequences
        evaluated_sequences = {e.get("query_sequence") for e in eval_events}

        pending = []
        for qe in query_events:
            if qe.get("sequence") not in evaluated_sequences:
                # 转为旧格式兼容（runner 中的评估逻辑期望此格式）
                for source in qe.get("sources", []):
                    pending.append({
                        "knowledge_id": source.get("id", ""),
                        "goal_id": qe.get("goal_id", ""),
                        "injected_at_sequence": qe.get("sequence", 0),
                        "task": source.get("task", ""),
                        "content": source.get("content", ""),
                        "query_sequence": qe.get("sequence"),
                    })
        return pending

    async def update_user_feedback(
        self,
        trace_id: str,
        knowledge_id: str,
        user_feedback: Dict[str, Any]
    ) -> None:
        """记录用户对知识的反馈（confirm/override）"""
        async with self._lock_for(trace_id):
            await self._update_user_feedback_locked(
                trace_id,
                knowledge_id,
                user_feedback,
            )

    async def _update_user_feedback_locked(
        self,
        trace_id: str,
        knowledge_id: str,
        user_feedback: Dict[str, Any],
    ) -> None:
        log = await self.get_cognition_log(trace_id)

        # 旧格式
        if "entries" in log:
            matching = [
                (i, e) for i, e in enumerate(log["entries"])
                if e.get("knowledge_id") == knowledge_id
            ]
            if matching:
                matching.sort(key=lambda x: x[1].get("injected_at_sequence", 0), reverse=True)
                _, entry = matching[0]
                entry["user_feedback"] = user_feedback
            self._atomic_write_json(
                self._get_knowledge_log_file(trace_id),
                log,
            )
            return

        # 新格式：追加 user_feedback 事件（或直接记录在 evaluation 上）
        await self.append_cognition_event(
            trace_id=trace_id,
            event={
                "type": "user_feedback",
                "knowledge_id": knowledge_id,
                "feedback": user_feedback,
            }
        )