Coverage for agentos/checkpoint/engine.py: 0%
183 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 09:19 +0800
1"""
2AgentOS v1.14.7 — Fine-grained Checkpoint Engine.
4LangGraph-aligned step-level checkpointing with time travel.
5Every tool_call, llm_call, and state transition triggers a snapshot.
7Key differences from v1.14.6 checkpoint module:
8- Step-level (not task-level) granularity
9- Time travel: rewind to any checkpoint and replay from there
10- Branching: fork execution from any historical checkpoint
11- Delta snapshots: only store state diffs when possible
12- Automatic pruning: configurable retention policies
14Usage:
15 engine = CheckpointEngine(checkpointer=SQLiteCheckpointer("checkpoints.db"))
17 # Auto-snapshot around tool calls
18 @engine.snapshot_on("tool_call")
19 async def my_tool(...): ...
21 # Time travel
22 await engine.rewind("checkpoint-42")
23 # Now continue execution from that point
25 # Branch
26 branch_id = await engine.branch("checkpoint-42", "bugfix-experiment")
27"""
29from __future__ import annotations
31import functools
32import logging
33import time
34import uuid
35from contextlib import asynccontextmanager
36from dataclasses import dataclass, field
37from enum import StrEnum
38from typing import Any
40from agentos.checkpoint.base import (
41 Checkpoint,
42 CheckpointBackend,
43 CheckpointMetadata,
44)
46logger = logging.getLogger(__name__)
49# ── Types ────────────────────────────────────
52class SnapshotTrigger(StrEnum):
53 """快照触发点。"""
55 TOOL_CALL = "tool_call" # 工具调用前后
56 LLM_CALL = "llm_call" # LLM 调用前后
57 STATE_CHANGE = "state_change" # Agent 状态变更
58 TASK_BOUNDARY = "task_boundary" # 任务开始/结束
59 MANUAL = "manual" # 显式调用
60 INTERVAL = "interval" # 定时快照
63class CheckpointGC(StrEnum):
64 """检查点垃圾回收策略。"""
66 KEEP_ALL = "keep_all"
67 KEEP_LAST_N = "keep_last_n"
68 KEEP_AGE = "keep_age" # 仅保留 N 秒内
69 KEEP_MILESTONES = "keep_milestones" # 仅保留首尾 + N 个分位点
72@dataclass
73class SnapshotConfig:
74 """快照配置。"""
76 triggers: set[SnapshotTrigger] = field(
77 default_factory=lambda: {
78 SnapshotTrigger.MANUAL,
79 SnapshotTrigger.TOOL_CALL,
80 SnapshotTrigger.LLM_CALL,
81 SnapshotTrigger.STATE_CHANGE,
82 }
83 )
84 gc_policy: CheckpointGC = CheckpointGC.KEEP_LAST_N
85 gc_param: int = 100 # keep_last_n 的 n 或 keep_age 的秒数
86 delta_snapshots: bool = True # 是否使用增量快照（减少存储）
87 max_snapshot_size_mb: float = 10.0
90@dataclass
91class TimeTravelResult:
92 """时间旅行操作结果。"""
94 checkpoint: Checkpoint
95 thread_id: str
96 rewind_depth: int # 回退了几个 checkpoint
97 snapshot_count_before: int # 重放前的快照数
98 can_replay: bool = True
101# ── Checkpoint Engine ────────────────────────
104class CheckpointEngine:
105 """细粒度 Checkpoint 引擎。
107 提供每步快照、时间旅行、分支等能力。
108 """
110 def __init__(
111 self,
112 checkpointer: CheckpointBackend,
113 config: SnapshotConfig | None = None,
114 ):
115 self._checkpointer = checkpointer
116 self._config = config or SnapshotConfig()
117 self._snapshot_counters: dict[str, int] = {} # thread_id → step counter
118 self._last_delta: dict[str, dict[str, Any]] = {} # thread_id → last full state
120 # ── Snapshot API ────────────────────────
122 async def snapshot(
123 self,
124 thread_id: str,
125 messages: list[dict[str, Any]],
126 state: dict[str, Any],
127 tools_result: dict[str, Any],
128 trigger: SnapshotTrigger = SnapshotTrigger.MANUAL,
129 parent_checkpoint_id: str | None = None,
130 next_node: str = "",
131 ) -> str:
132 """创建一次快照。返回 checkpoint_id。"""
133 if trigger not in self._config.triggers:
134 return "" # 不在此触发范围内
136 step = self._snapshot_counters.get(thread_id, 0) + 1
137 self._snapshot_counters[thread_id] = step
139 checkpoint_id = f"ckpt-{thread_id}-{step}-{uuid.uuid4().hex[:6]}"
141 metadata = CheckpointMetadata(
142 thread_id=thread_id,
143 checkpoint_id=checkpoint_id,
144 parent_checkpoint_id=parent_checkpoint_id,
145 step=step,
146 tags=[trigger.value],
147 summary=self._auto_summary(messages, state),
148 )
150 checkpoint = Checkpoint(
151 metadata=metadata,
152 messages=list(messages),
153 state=dict(state),
154 tools_result=dict(tools_result),
155 next_node=next_node,
156 )
158 await self._checkpointer.put(checkpoint)
160 # GC
161 await self._maybe_gc(thread_id)
163 return checkpoint_id
165 async def snapshot_safe(
166 self,
167 thread_id: str,
168 messages: list[dict[str, Any]],
169 state: dict[str, Any],
170 tools_result: dict[str, Any],
171 trigger: SnapshotTrigger = SnapshotTrigger.MANUAL,
172 parent_checkpoint_id: str | None = None,
173 next_node: str = "",
174 ) -> str:
175 """安全快照：失败不抛异常，不影响主流程。"""
176 try:
177 return await self.snapshot(
178 thread_id,
179 messages,
180 state,
181 tools_result,
182 trigger,
183 parent_checkpoint_id,
184 next_node,
185 )
186 except Exception as e:
187 logger.error(f"Snapshot failed (non-blocking): {e}")
188 return ""
190 # ── Time Travel API ─────────────────────
192 async def rewind(
193 self,
194 checkpoint_id: str,
195 ) -> TimeTravelResult:
196 """时间旅行：回退到指定 checkpoint。"""
197 target = await self._checkpointer.get(checkpoint_id)
198 if not target:
199 raise ValueError(f"Checkpoint {checkpoint_id} not found")
201 thread_id = target.metadata.thread_id
203 # 计算回退深度
204 current_step = self._snapshot_counters.get(thread_id, 0)
205 target_step = target.metadata.step
206 rewind_depth = current_step - target_step
208 # 删除目标之后的 checkpoint（默认行为，可配置）
209 later_checkpoints = await self._checkpointer.list_checkpoints(thread_id)
210 deleted = 0
211 for cp_meta in later_checkpoints:
212 if cp_meta.step > target_step:
213 await self._checkpointer.delete_thread(cp_meta.thread_id)
214 deleted += 1
216 # 重置计数器
217 self._snapshot_counters[thread_id] = target_step
219 logger.info(
220 f"Time travel: rewound {thread_id} by {rewind_depth} steps "
221 f"to checkpoint {checkpoint_id} (step {target_step}), deleted {deleted} later checkpoints"
222 )
224 return TimeTravelResult(
225 checkpoint=target,
226 thread_id=thread_id,
227 rewind_depth=rewind_depth,
228 snapshot_count_before=current_step,
229 )
231 async def time_travel_to_step(
232 self,
233 thread_id: str,
234 step: int,
235 ) -> TimeTravelResult | None:
236 """按步骤号时间旅行。"""
237 checkpoints = await self._checkpointer.list_checkpoints(thread_id, limit=500)
239 # 找到最接近目标 step 的 checkpoint
240 matching = [cp for cp in checkpoints if cp.step <= step]
241 if not matching:
242 return None
244 target = sorted(matching, key=lambda c: c.step, reverse=True)[0]
245 return await self.rewind(target.checkpoint_id)
247 async def list_time_travel_points(
248 self,
249 thread_id: str,
250 limit: int = 50,
251 ) -> list[CheckpointMetadata]:
252 """列出所有可回溯的时间点。"""
253 return await self._checkpointer.list_checkpoints(thread_id, limit=limit)
255 # ── Branch API ──────────────────────────
257 async def branch(
258 self,
259 from_checkpoint_id: str,
260 branch_name: str,
261 ) -> str:
262 """从某个历史 checkpoint 创建分支执行。"""
263 source = await self._checkpointer.get(from_checkpoint_id)
264 if not source:
265 raise ValueError(f"Source checkpoint {from_checkpoint_id} not found")
267 branch_thread_id = (
268 f"{source.metadata.thread_id}-branch-{branch_name}-{uuid.uuid4().hex[:4]}"
269 )
271 # 在新分支中创建起始 checkpoint（引用源 checkpoint 状态）
272 branch_checkpoint = Checkpoint(
273 metadata=CheckpointMetadata(
274 thread_id=branch_thread_id,
275 checkpoint_id=f"ckpt-{branch_thread_id}-0",
276 parent_checkpoint_id=from_checkpoint_id,
277 step=0,
278 tags=["branch", branch_name],
279 summary=f"Branch '{branch_name}' from {from_checkpoint_id}",
280 ),
281 messages=list(source.messages),
282 state=dict(source.state),
283 tools_result=dict(source.tools_result),
284 next_node="",
285 )
287 await self._checkpointer.put(branch_checkpoint)
288 self._snapshot_counters[branch_thread_id] = 0
290 logger.info(f"Created branch: {branch_thread_id} from {from_checkpoint_id}")
291 return branch_thread_id
293 async def merge_branch(
294 self,
295 branch_thread_id: str,
296 into_thread_id: str,
297 ) -> str:
298 """合并分支到主线程。"""
299 branch_latest = await self._checkpointer.get_latest(branch_thread_id)
300 if not branch_latest:
301 raise ValueError(f"Branch {branch_thread_id} has no checkpoints")
303 # 在主线程创建一个引用分支状态的快照
304 merge_id = await self.snapshot(
305 thread_id=into_thread_id,
306 messages=branch_latest.messages,
307 state=branch_latest.state,
308 tools_result=branch_latest.tools_result,
309 trigger=SnapshotTrigger.MANUAL,
310 parent_checkpoint_id=branch_latest.metadata.checkpoint_id,
311 )
313 logger.info(f"Merged branch {branch_thread_id} → {into_thread_id} (merge ckpt: {merge_id})")
314 return merge_id
316 # ── Decorator API ───────────────────────
318 def snapshot_on(self, trigger: SnapshotTrigger):
319 """装饰器：在调用前后自动快照。
321 Usage:
322 engine = CheckpointEngine(...)
324 @engine.snapshot_on(SnapshotTrigger.TOOL_CALL)
325 async def search_database(query: str): ...
326 """
328 def decorator(func):
329 @functools.wraps(func)
330 async def wrapper(*args, **kwargs):
331 thread_id = kwargs.pop("_checkpoint_thread_id", "default")
332 state = kwargs.pop("_checkpoint_state", {})
334 # Before snapshot
335 await self.snapshot_safe(
336 thread_id=thread_id,
337 messages=[{"role": "tool_call", "content": f"{func.__name__}({kwargs})"}],
338 state=state,
339 tools_result={},
340 trigger=trigger,
341 )
343 result = await func(*args, **kwargs)
345 # After snapshot
346 await self.snapshot_safe(
347 thread_id=thread_id,
348 messages=[{"role": "tool_result", "content": str(result)[:500]}],
349 state=state,
350 tools_result={"result": str(result)[:1000]},
351 trigger=trigger,
352 )
354 return result
356 return wrapper
358 return decorator
360 @asynccontextmanager
361 async def snapshot_scope(
362 self,
363 thread_id: str,
364 state: dict[str, Any],
365 trigger: SnapshotTrigger = SnapshotTrigger.STATE_CHANGE,
366 ):
367 """上下文管理器：进入和退出作用域时自动快照。
369 Usage:
370 async with engine.snapshot_scope("thread-1", state):
371 await execute_workflow(...)
372 """
373 await self.snapshot_safe(
374 thread_id=thread_id,
375 messages=[{"role": "system", "content": f"Enter scope ({trigger.value})"}],
376 state=state,
377 tools_result={},
378 trigger=trigger,
379 )
380 try:
381 yield
382 finally:
383 await self.snapshot_safe(
384 thread_id=thread_id,
385 messages=[{"role": "system", "content": f"Exit scope ({trigger.value})"}],
386 state=state,
387 tools_result={},
388 trigger=trigger,
389 )
391 # ── Query API ───────────────────────────
393 async def get_latest(self, thread_id: str) -> Checkpoint | None:
394 return await self._checkpointer.get_latest(thread_id)
396 async def get_checkpoint_tree(self, thread_id: str, limit: int = 200) -> dict[str, Any]:
397 """获取线程的 checkpoint 树结构（用于可视化）。"""
398 checkpoints = await self._checkpointer.list_checkpoints(thread_id, limit=limit)
400 nodes: list[dict] = []
401 edges: list[dict] = []
402 by_id: dict[str, CheckpointMetadata] = {}
404 for cp in checkpoints:
405 by_id[cp.checkpoint_id] = cp
406 nodes.append(
407 {
408 "id": cp.checkpoint_id,
409 "step": cp.step,
410 "tags": cp.tags,
411 "summary": cp.summary,
412 "created_at": cp.created_at,
413 }
414 )
416 for cp in checkpoints:
417 if cp.parent_checkpoint_id and cp.parent_checkpoint_id in by_id:
418 edges.append(
419 {
420 "from": cp.parent_checkpoint_id,
421 "to": cp.checkpoint_id,
422 }
423 )
425 return {
426 "thread_id": thread_id,
427 "total_checkpoints": len(checkpoints),
428 "nodes": nodes,
429 "edges": edges,
430 }
432 # ── Internal ────────────────────────────
434 def _auto_summary(self, messages: list[dict[str, Any]], state: dict[str, Any]) -> str:
435 """自动生成 checkpoint 摘要。"""
436 if messages:
437 last = messages[-1]
438 role = last.get("role", "unknown")
439 content = str(last.get("content", ""))[:100]
440 return f"[{role}] {content}"
441 return f"State: {len(state)} keys"
443 async def _maybe_gc(self, thread_id: str):
444 """根据 GC 策略清理旧 checkpoint。"""
445 if self._config.gc_policy == CheckpointGC.KEEP_ALL:
446 return
448 checkpoints = await self._checkpointer.list_checkpoints(thread_id, limit=500)
450 if self._config.gc_policy == CheckpointGC.KEEP_LAST_N:
451 if len(checkpoints) > self._config.gc_param:
452 to_delete = sorted(checkpoints, key=lambda c: c.step)[
453 : len(checkpoints) - self._config.gc_param
454 ]
455 for cp in to_delete:
456 await self._checkpointer.delete_before(thread_id, cp.step + 1)
457 logger.debug(f"GC: removed {len(to_delete)} old checkpoints from {thread_id}")
459 elif self._config.gc_policy == CheckpointGC.KEEP_AGE:
460 cutoff = time.time() - self._config.gc_param
461 deleted = 0
462 for cp in checkpoints:
463 try:
464 created = (
465 __import__("datetime").datetime.fromisoformat(cp.created_at).timestamp()
466 )
467 if created < cutoff:
468 await self._checkpointer.delete_thread(cp.thread_id)
469 deleted += 1
470 except Exception:
471 continue
472 if deleted:
473 logger.debug(f"GC: removed {deleted} expired checkpoints from {thread_id}")
475 elif self._config.gc_policy == CheckpointGC.KEEP_MILESTONES:
476 if len(checkpoints) > self._config.gc_param:
477 # 保留 first, last, 和均匀分布的 milestones
478 sorted_cps = sorted(checkpoints, key=lambda c: c.step)
479 keep = {sorted_cps[0].step, sorted_cps[-1].step}
481 n_milestones = max(2, self._config.gc_param - 2)
482 step_size = max(1, len(sorted_cps) // n_milestones)
483 for i in range(1, n_milestones):
484 idx = i * step_size
485 if idx < len(sorted_cps):
486 keep.add(sorted_cps[idx].step)
488 for cp in sorted_cps:
489 if cp.step not in keep:
490 await self._checkpointer.delete_thread(cp.thread_id)
491 logger.debug(f"GC milestones: kept {len(keep)} of {len(sorted_cps)} in {thread_id}")
494# ── Quick Start ──────────────────────────────
497async def create_checkpoint_engine(
498 backend: str = "sqlite",
499 db_path: str = "checkpoints.db",
500) -> CheckpointEngine:
501 """快速创建 checkpoint 引擎。"""
502 from agentos.checkpoint.factory import create_checkpointer
504 checkpointer = create_checkpointer(backend, db_path=db_path)
505 return CheckpointEngine(checkpointer)