|
|
@@ -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 上)
|