Ver código fonte

fix(trace-store): 加固文件账本并持久化审批批次

为每个 Trace 建立可重入异步锁,使消息、Goal、模型用量、认知日志、审批和事件的 read-modify-write 落在同一锁域。

JSON 与二进制文件统一使用同目录临时文件、flush、fsync 和原子 replace;events.jsonl 在锁内追加并 fsync,启动对账从真实消息与事件恢复计数。

新增严格路径组件和 resolve 越界校验,区分文件缺失与 JSON 损坏;附件迁入独立 attachments 目录,重复 message_id 在更新统计前直接失败。

引入 ToolApprovalBatchV1 持久化契约,保存调用身份、原始与有效参数、参数哈希、决策及执行状态,并扩展 TraceStore 协议支持审批和附件原子写入。
SamLee 14 horas atrás
pai
commit
b3c45a2a87

+ 121 - 0
cyber_agent/tools/approval.py

@@ -0,0 +1,121 @@
+"""Persistent contracts for human approval of tool-call batches."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from datetime import datetime
+from hashlib import sha256
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+def canonical_arguments(arguments: dict[str, Any]) -> str:
+    return json.dumps(
+        arguments,
+        ensure_ascii=False,
+        sort_keys=True,
+        separators=(",", ":"),
+        allow_nan=False,
+    )
+
+
+def tool_argument_hash(
+    *,
+    tool_call_id: str,
+    tool_name: str,
+    arguments: dict[str, Any],
+) -> str:
+    payload = f"{tool_call_id}\n{tool_name}\n{canonical_arguments(arguments)}"
+    return sha256(payload.encode("utf-8")).hexdigest()
+
+
+class _StrictModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class ToolApprovalCallV1(_StrictModel):
+    tool_call_id: str = Field(min_length=1)
+    tool_name: str = Field(min_length=1)
+    original_arguments: dict[str, Any]
+    effective_arguments: dict[str, Any]
+    argument_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
+    editable_params: list[str]
+    requires_confirmation: bool
+    decision: Literal["pending", "approved", "rejected", "auto_approved"]
+    execution_status: Literal[
+        "not_started",
+        "executing",
+        "executed",
+        "rejected",
+        "execution_unknown",
+    ] = "not_started"
+    result_message_id: str | None = None
+    decided_at: str | None = None
+
+    @model_validator(mode="after")
+    def validate_hash(self) -> "ToolApprovalCallV1":
+        expected = tool_argument_hash(
+            tool_call_id=self.tool_call_id,
+            tool_name=self.tool_name,
+            arguments=self.effective_arguments,
+        )
+        if self.argument_hash != expected:
+            raise ValueError("tool approval argument hash mismatch")
+        return self
+
+
+class ToolApprovalBatchV1(_StrictModel):
+    schema_version: Literal[1] = 1
+    batch_id: str = Field(min_length=1)
+    trace_id: str = Field(min_length=1)
+    assistant_message_id: str = Field(min_length=1)
+    assistant_sequence: int = Field(gt=0)
+    status: Literal[
+        "pending",
+        "decided",
+        "executing",
+        "completed",
+        "cancelled",
+        "execution_unknown",
+    ] = "pending"
+    calls: list[ToolApprovalCallV1] = Field(min_length=1)
+    created_at: str
+    updated_at: str
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        trace_id: str,
+        assistant_message_id: str,
+        assistant_sequence: int,
+        calls: list[ToolApprovalCallV1],
+    ) -> "ToolApprovalBatchV1":
+        now = datetime.now().isoformat()
+        return cls(
+            batch_id=f"approval_{uuid.uuid4().hex}",
+            trace_id=trace_id,
+            assistant_message_id=assistant_message_id,
+            assistant_sequence=assistant_sequence,
+            calls=calls,
+            created_at=now,
+            updated_at=now,
+        )
+
+    @property
+    def pending_calls(self) -> list[ToolApprovalCallV1]:
+        return [call for call in self.calls if call.decision == "pending"]
+
+
+def approval_grant(batch: ToolApprovalBatchV1, call: ToolApprovalCallV1) -> dict[str, Any]:
+    """Build the internal, hash-bound grant consumed by ToolRegistry.execute."""
+    return {
+        "batch_id": batch.batch_id,
+        "trace_id": batch.trace_id,
+        "tool_call_id": call.tool_call_id,
+        "tool_name": call.tool_name,
+        "argument_hash": call.argument_hash,
+        "decision": call.decision,
+    }

+ 9 - 1
cyber_agent/trace/models.py

@@ -51,7 +51,13 @@ class Trace:
     parent_goal_id: Optional[str] = None      # 哪个 Goal 启动的
 
     # 状态
-    status: Literal["running", "completed", "failed", "stopped"] = "running"
+    status: Literal[
+        "running",
+        "waiting_confirmation",
+        "completed",
+        "failed",
+        "stopped",
+    ] = "running"
 
     # 统计
     total_messages: int = 0      # 消息总数(改名自 total_steps)
@@ -442,6 +448,8 @@ class Message:
             "tool_call_id": self.tool_call_id,
             "content": self.content,
             "description": self.description,
+            "branch_type": self.branch_type,
+            "branch_id": self.branch_id,
             "tokens": self.tokens,  # 使用 @property 动态计算
             "prompt_tokens": self.prompt_tokens,
             "completion_tokens": self.completion_tokens,

+ 28 - 0
cyber_agent/trace/protocols.py

@@ -258,3 +258,31 @@ class TraceStore(Protocol):
             event_id: 新事件的 ID
         """
         ...
+
+    async def get_tool_approval_batch(self, trace_id: str):
+        """Load the persisted approval batch for a Trace, if any."""
+        ...
+
+    async def replace_tool_approval_batch(self, trace_id: str, batch) -> None:
+        """Atomically replace a Trace's tool approval batch."""
+        ...
+
+    async def mark_reflections_consumed(
+        self,
+        trace_id: str,
+        reflections: List[Dict[str, Any]],
+        consumed_at: str,
+    ) -> int:
+        """Atomically mark selected reflection cognition events consumed."""
+        ...
+
+    async def write_message_attachment(
+        self,
+        trace_id: str,
+        message_id: str,
+        *,
+        suffix: str,
+        data: bytes,
+    ):
+        """Atomically persist a Message-owned binary attachment."""
+        ...

+ 548 - 132
cyber_agent/trace/store.py

@@ -22,6 +22,7 @@ Sub-Trace 是完全独立的 Trace,有自己的目录:
 └── events.jsonl
 """
 
+import asyncio
 import json
 import logging
 import os
@@ -36,18 +37,150 @@ 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)
-        self.base_path.mkdir(exist_ok=True)
+        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 目录"""
-        return self.base_path / trace_id
+        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 文件路径"""
@@ -61,9 +194,16 @@ class FileSystemTraceStore:
         """获取 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 文件路径"""
-        return self._get_messages_dir(trace_id) / f"{message_id}.json"
+        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 文件路径"""
@@ -77,24 +217,30 @@ class FileSystemTraceStore:
         """获取根 Trace 的 resource_usage.json 文件路径。"""
         return self._get_trace_dir(root_trace_id) / "resource_usage.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"""
-        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
-        meta_file = self._get_meta_file(trace.trace_id)
-        meta_file.write_text(json.dumps(trace.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
+        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
-        events_file = self._get_events_file(trace.trace_id)
-        events_file.touch()
+            # 创建空的 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
 
@@ -104,7 +250,16 @@ class FileSystemTraceStore:
         if not meta_file.exists():
             return None
 
-        data = json.loads(meta_file.read_text(encoding="utf-8"))
+        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"):
@@ -116,18 +271,20 @@ class FileSystemTraceStore:
 
     async def update_trace(self, trace_id: str, **updates) -> None:
         """更新 Trace"""
-        trace = await self.get_trace(trace_id)
-        if not trace:
-            return
+        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)
+            # 更新字段
+            for key, value in updates.items():
+                if hasattr(trace, key):
+                    setattr(trace, key, value)
 
-        # 写回文件
-        meta_file = self._get_meta_file(trace_id)
-        meta_file.write_text(json.dumps(trace.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
+            self._atomic_write_json(
+                self._get_meta_file(trace_id),
+                trace.to_dict(),
+            )
 
     async def list_traces(
         self,
@@ -151,12 +308,21 @@ class FileSystemTraceStore:
             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:
@@ -179,8 +345,10 @@ class FileSystemTraceStore:
                     data["completed_at"] = datetime.fromisoformat(data["completed_at"])
 
                 traces.append(Trace.from_dict(data))
-            except Exception:
-                continue
+            except Exception as exc:
+                raise TraceStoreCorruptionError(
+                    f"corrupt trace metadata: {meta_file}"
+                ) from exc
 
         # 排序(最新的在前)
         traces.sort(key=lambda t: t.created_at, reverse=True)
@@ -199,10 +367,15 @@ class FileSystemTraceStore:
         usage_file = self._get_resource_usage_file(root_trace_id)
         if not usage_file.exists():
             return None
-        data = json.loads(usage_file.read_text(encoding="utf-8"))
-        if not isinstance(data, dict):
-            raise ValueError("resource_usage.json must contain an object")
-        return data
+        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,
@@ -212,31 +385,58 @@ class FileSystemTraceStore:
         """通过同目录临时文件 + os.replace 原子替换用量快照。
 
         ResourceBudgetController 在 Agent/LLM 预留和 Token/成本登记后调用,避免单进程留下半写入文件。"""
-        trace_dir = self._get_trace_dir(root_trace_id)
-        trace_dir.mkdir(exist_ok=True)
-        usage_file = self._get_resource_usage_file(root_trace_id)
-        temp_path: Optional[str] = None
+        async with self._lock_for(root_trace_id):
+            self._atomic_write_json(
+                self._get_resource_usage_file(root_trace_id),
+                usage,
+            )
+
+    # ===== 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:
-            with tempfile.NamedTemporaryFile(
-                mode="w",
-                encoding="utf-8",
-                dir=trace_dir,
-                prefix=".resource_usage.",
-                suffix=".tmp",
-                delete=False,
-            ) as temp_file:
-                temp_path = temp_file.name
-                json.dump(usage, temp_file, indent=2, ensure_ascii=False)
-                temp_file.flush()
-                os.fsync(temp_file.fileno())
-            os.replace(temp_path, usage_file)
-            temp_path = None
-        finally:
-            if temp_path is not None:
-                try:
-                    os.unlink(temp_path)
-                except FileNotFoundError:
-                    pass
+            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 操作 =====
 
@@ -248,30 +448,37 @@ class FileSystemTraceStore:
 
         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:
-            return None
+        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"""
-        goal_file = self._get_goal_file(trace_id)
-        goal_file.write_text(json.dumps(tree.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
+        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"""
-        tree = await self.get_goal_tree(trace_id)
-        if not tree:
-            return
-
-        tree.goals.append(goal)
-        await self.update_goal_tree(trace_id, tree)
+        async with self._lock_for(trace_id):
+            tree = await self.get_goal_tree(trace_id)
+            if not tree:
+                return
 
-        # 推送 goal_added 事件
-        event_data = {
-            "goal": goal.to_dict(),
-            "parent_id": goal.parent_id
-        }
-        await self.append_event(trace_id, "goal_added", event_data)
+            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
@@ -292,6 +499,22 @@ class FileSystemTraceStore:
         """更新 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
@@ -405,6 +628,10 @@ class FileSystemTraceStore:
     # ===== 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
 
@@ -413,9 +640,12 @@ class FileSystemTraceStore:
         trace_id = message.trace_id
 
         # 1. 写入 message 文件
-        messages_dir = self._get_messages_dir(trace_id)
-        message_file = messages_dir / f"{message.message_id}.json"
-        message_file.write_text(json.dumps(message.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
+        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)
@@ -566,22 +796,28 @@ class FileSystemTraceStore:
 
         return affected
 
-        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 目录
-            message_file = trace_dir / "messages" / f"{message_id}.json"
+            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:
-                    pass
+                except Exception as exc:
+                    raise TraceStoreCorruptionError(
+                        f"corrupt message file: {message_file}"
+                    ) from exc
 
         return None
 
@@ -601,8 +837,10 @@ class FileSystemTraceStore:
                 data = json.loads(message_file.read_text(encoding="utf-8"))
                 msg = Message.from_dict(data)
                 messages.append(msg)
-            except Exception:
-                continue
+            except Exception as exc:
+                raise TraceStoreCorruptionError(
+                    f"corrupt message file: {message_file}"
+                ) from exc
 
         # 按 sequence 排序
         messages.sort(key=lambda m: m.sequence)
@@ -652,42 +890,44 @@ class FileSystemTraceStore:
 
     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
-
-        # 更新字段
-        for key, value in updates.items():
-            if hasattr(message, key):
-                setattr(message, key, value)
-
-        # 确定文件路径
-        messages_dir = self._get_messages_dir(message.trace_id)
-        message_file = messages_dir / f"{message_id}.json"
-        message_file.write_text(json.dumps(message.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
+        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 列表。
         """
-        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
-                # 直接写回文件
-                message_file = self._get_messages_dir(trace_id) / f"{msg.message_id}.json"
-                message_file.write_text(
-                    json.dumps(msg.to_dict(), indent=2, ensure_ascii=False),
-                    encoding="utf-8"
-                )
-                abandoned_ids.append(msg.message_id)
-
-        return abandoned_ids
+        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
 
     # ===== 模型使用追踪 =====
 
@@ -715,11 +955,41 @@ class FileSystemTraceStore:
             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():
-            data = json.loads(usage_file.read_text(encoding="utf-8"))
+            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": {
@@ -785,7 +1055,7 @@ class FileSystemTraceStore:
         data["timeline"].append(timeline_entry)
 
         # 写回文件
-        usage_file.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
+        self._atomic_write_json(usage_file, data)
 
     # ===== 事件流操作(用于 WebSocket 断线续传)=====
 
@@ -800,14 +1070,23 @@ class FileSystemTraceStore:
             return []
 
         events = []
+        previous_id = 0
         with events_file.open('r', encoding='utf-8') as f:
-            for line in f:
+            for line_number, line in enumerate(f, start=1):
+                if not line.strip():
+                    continue
                 try:
                     event = json.loads(line.strip())
-                    if event.get("event_id", 0) > since_event_id:
+                    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:
-                    continue
+                except Exception as exc:
+                    raise TraceStoreCorruptionError(
+                        f"corrupt event journal {events_file}:{line_number}"
+                    ) from exc
 
         return events
 
@@ -818,6 +1097,15 @@ class FileSystemTraceStore:
         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:
@@ -841,9 +1129,51 @@ class FileSystemTraceStore:
         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:
@@ -863,12 +1193,28 @@ class FileSystemTraceStore:
     async def get_cognition_log(self, trace_id: str) -> Dict[str, Any]:
         """读取认知日志"""
         log_file = self._get_cognition_log_file(trace_id)
-        if log_file.exists():
-            return json.loads(log_file.read_text(encoding="utf-8"))
+        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"
-        if legacy_file.exists():
-            return json.loads(legacy_file.read_text(encoding="utf-8"))
+        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]:
@@ -915,13 +1261,51 @@ class FileSystemTraceStore:
 
         其他字段可按需附加,不做强校验(演进友好)。
         """
-        log = await self.get_cognition_log(trace_id)
-        if "events" not in log:
-            log["events"] = log.pop("entries", [])
-        event["timestamp"] = datetime.now().isoformat()
-        log["events"].append(event)
-        log_file = self._get_cognition_log_file(trace_id)
-        log_file.write_text(json.dumps(log, indent=2, ensure_ascii=False), encoding="utf-8")
+        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,
@@ -958,6 +1342,21 @@ class FileSystemTraceStore:
         旧格式:更新 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", []))
 
@@ -973,8 +1372,10 @@ class FileSystemTraceStore:
                 entry["eval_result"] = eval_result
                 entry["evaluated_at"] = datetime.now().isoformat()
                 entry["evaluated_at_trigger"] = trigger_event
-                log_file = self._get_knowledge_log_file(trace_id)
-                log_file.write_text(json.dumps(log, indent=2, ensure_ascii=False), encoding="utf-8")
+                self._atomic_write_json(
+                    self._get_knowledge_log_file(trace_id),
+                    log,
+                )
                 return
 
         # 新格式:追加 evaluation 事件
@@ -1036,6 +1437,19 @@ class FileSystemTraceStore:
         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)
 
         # 旧格式
@@ -1048,8 +1462,10 @@ class FileSystemTraceStore:
                 matching.sort(key=lambda x: x[1].get("injected_at_sequence", 0), reverse=True)
                 _, entry = matching[0]
                 entry["user_feedback"] = user_feedback
-            log_file = self._get_knowledge_log_file(trace_id)
-            log_file.write_text(json.dumps(log, indent=2, ensure_ascii=False), encoding="utf-8")
+            self._atomic_write_json(
+                self._get_knowledge_log_file(trace_id),
+                log,
+            )
             return
 
         # 新格式:追加 user_feedback 事件(或直接记录在 evaluation 上)