Coverage for agentos/checkpoint/engine.py: 0%
183 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +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 Enum
38from typing import Any, Dict, List, Optional, Set
40from agentos.checkpoint.base import (
41 Checkpoint,
42 CheckpointMetadata,
43 CheckpointBackend,
44)
46logger = logging.getLogger(__name__)
49# ── Types ────────────────────────────────────
52class SnapshotTrigger(str, Enum):
53 """快照触发点。"""
54 TOOL_CALL = "tool_call" # 工具调用前后
55 LLM_CALL = "llm_call" # LLM 调用前后
56 STATE_CHANGE = "state_change" # Agent 状态变更
57 TASK_BOUNDARY = "task_boundary" # 任务开始/结束
58 MANUAL = "manual" # 显式调用
59 INTERVAL = "interval" # 定时快照
62class CheckpointGC(str, Enum):
63 """检查点垃圾回收策略。"""
64 KEEP_ALL = "keep_all"
65 KEEP_LAST_N = "keep_last_n"
66 KEEP_AGE = "keep_age" # 仅保留 N 秒内
67 KEEP_MILESTONES = "keep_milestones" # 仅保留首尾 + N 个分位点
70@dataclass
71class SnapshotConfig:
72 """快照配置。"""
73 triggers: Set[SnapshotTrigger] = field(default_factory=lambda: {
74 SnapshotTrigger.MANUAL,
75 SnapshotTrigger.TOOL_CALL,
76 SnapshotTrigger.LLM_CALL,
77 SnapshotTrigger.STATE_CHANGE,
78 })
79 gc_policy: CheckpointGC = CheckpointGC.KEEP_LAST_N
80 gc_param: int = 100 # keep_last_n 的 n 或 keep_age 的秒数
81 delta_snapshots: bool = True # 是否使用增量快照（减少存储）
82 max_snapshot_size_mb: float = 10.0
85@dataclass
86class TimeTravelResult:
87 """时间旅行操作结果。"""
88 checkpoint: Checkpoint
89 thread_id: str
90 rewind_depth: int # 回退了几个 checkpoint
91 snapshot_count_before: int # 重放前的快照数
92 can_replay: bool = True
95# ── Checkpoint Engine ────────────────────────
98class CheckpointEngine:
99 """细粒度 Checkpoint 引擎。
101 提供每步快照、时间旅行、分支等能力。
102 """
104 def __init__(
105 self,
106 checkpointer: CheckpointBackend,
107 config: Optional[SnapshotConfig] = None,
108 ):
109 self._checkpointer = checkpointer
110 self._config = config or SnapshotConfig()
111 self._snapshot_counters: Dict[str, int] = {} # thread_id → step counter
112 self._last_delta: Dict[str, Dict[str, Any]] = {} # thread_id → last full state
114 # ── Snapshot API ────────────────────────
116 async def snapshot(
117 self,
118 thread_id: str,
119 messages: List[Dict[str, Any]],
120 state: Dict[str, Any],
121 tools_result: Dict[str, Any],
122 trigger: SnapshotTrigger = SnapshotTrigger.MANUAL,
123 parent_checkpoint_id: Optional[str] = None,
124 next_node: str = "",
125 ) -> str:
126 """创建一次快照。返回 checkpoint_id。"""
127 if trigger not in self._config.triggers:
128 return "" # 不在此触发范围内
130 step = self._snapshot_counters.get(thread_id, 0) + 1
131 self._snapshot_counters[thread_id] = step
133 checkpoint_id = f"ckpt-{thread_id}-{step}-{uuid.uuid4().hex[:6]}"
135 metadata = CheckpointMetadata(
136 thread_id=thread_id,
137 checkpoint_id=checkpoint_id,
138 parent_checkpoint_id=parent_checkpoint_id,
139 step=step,
140 tags=[trigger.value],
141 summary=self._auto_summary(messages, state),
142 )
144 checkpoint = Checkpoint(
145 metadata=metadata,
146 messages=list(messages),
147 state=dict(state),
148 tools_result=dict(tools_result),
149 next_node=next_node,
150 )
152 await self._checkpointer.put(checkpoint)
154 # GC
155 await self._maybe_gc(thread_id)
157 return checkpoint_id
159 async def snapshot_safe(
160 self,
161 thread_id: str,
162 messages: List[Dict[str, Any]],
163 state: Dict[str, Any],
164 tools_result: Dict[str, Any],
165 trigger: SnapshotTrigger = SnapshotTrigger.MANUAL,
166 parent_checkpoint_id: Optional[str] = None,
167 next_node: str = "",
168 ) -> str:
169 """安全快照：失败不抛异常，不影响主流程。"""
170 try:
171 return await self.snapshot(
172 thread_id, messages, state, tools_result,
173 trigger, parent_checkpoint_id, next_node,
174 )
175 except Exception as e:
176 logger.error(f"Snapshot failed (non-blocking): {e}")
177 return ""
179 # ── Time Travel API ─────────────────────
181 async def rewind(
182 self,
183 checkpoint_id: str,
184 ) -> TimeTravelResult:
185 """时间旅行：回退到指定 checkpoint。"""
186 target = await self._checkpointer.get(checkpoint_id)
187 if not target:
188 raise ValueError(f"Checkpoint {checkpoint_id} not found")
190 thread_id = target.metadata.thread_id
192 # 计算回退深度
193 current_step = self._snapshot_counters.get(thread_id, 0)
194 target_step = target.metadata.step
195 rewind_depth = current_step - target_step
197 # 删除目标之后的 checkpoint（默认行为，可配置）
198 later_checkpoints = await self._checkpointer.list_checkpoints(thread_id)
199 deleted = 0
200 for cp_meta in later_checkpoints:
201 if cp_meta.step > target_step:
202 await self._checkpointer.delete_thread(cp_meta.thread_id)
203 deleted += 1
205 # 重置计数器
206 self._snapshot_counters[thread_id] = target_step
208 logger.info(
209 f"Time travel: rewound {thread_id} by {rewind_depth} steps "
210 f"to checkpoint {checkpoint_id} (step {target_step}), deleted {deleted} later checkpoints"
211 )
213 return TimeTravelResult(
214 checkpoint=target,
215 thread_id=thread_id,
216 rewind_depth=rewind_depth,
217 snapshot_count_before=current_step,
218 )
220 async def time_travel_to_step(
221 self,
222 thread_id: str,
223 step: int,
224 ) -> Optional[TimeTravelResult]:
225 """按步骤号时间旅行。"""
226 checkpoints = await self._checkpointer.list_checkpoints(thread_id, limit=500)
228 # 找到最接近目标 step 的 checkpoint
229 matching = [cp for cp in checkpoints if cp.step <= step]
230 if not matching:
231 return None
233 target = sorted(matching, key=lambda c: c.step, reverse=True)[0]
234 return await self.rewind(target.checkpoint_id)
236 async def list_time_travel_points(
237 self,
238 thread_id: str,
239 limit: int = 50,
240 ) -> List[CheckpointMetadata]:
241 """列出所有可回溯的时间点。"""
242 return await self._checkpointer.list_checkpoints(thread_id, limit=limit)
244 # ── Branch API ──────────────────────────
246 async def branch(
247 self,
248 from_checkpoint_id: str,
249 branch_name: str,
250 ) -> str:
251 """从某个历史 checkpoint 创建分支执行。"""
252 source = await self._checkpointer.get(from_checkpoint_id)
253 if not source:
254 raise ValueError(f"Source checkpoint {from_checkpoint_id} not found")
256 branch_thread_id = f"{source.metadata.thread_id}-branch-{branch_name}-{uuid.uuid4().hex[:4]}"
258 # 在新分支中创建起始 checkpoint（引用源 checkpoint 状态）
259 branch_checkpoint = Checkpoint(
260 metadata=CheckpointMetadata(
261 thread_id=branch_thread_id,
262 checkpoint_id=f"ckpt-{branch_thread_id}-0",
263 parent_checkpoint_id=from_checkpoint_id,
264 step=0,
265 tags=["branch", branch_name],
266 summary=f"Branch '{branch_name}' from {from_checkpoint_id}",
267 ),
268 messages=list(source.messages),
269 state=dict(source.state),
270 tools_result=dict(source.tools_result),
271 next_node="",
272 )
274 await self._checkpointer.put(branch_checkpoint)
275 self._snapshot_counters[branch_thread_id] = 0
277 logger.info(f"Created branch: {branch_thread_id} from {from_checkpoint_id}")
278 return branch_thread_id
280 async def merge_branch(
281 self,
282 branch_thread_id: str,
283 into_thread_id: str,
284 ) -> str:
285 """合并分支到主线程。"""
286 branch_latest = await self._checkpointer.get_latest(branch_thread_id)
287 if not branch_latest:
288 raise ValueError(f"Branch {branch_thread_id} has no checkpoints")
290 # 在主线程创建一个引用分支状态的快照
291 merge_id = await self.snapshot(
292 thread_id=into_thread_id,
293 messages=branch_latest.messages,
294 state=branch_latest.state,
295 tools_result=branch_latest.tools_result,
296 trigger=SnapshotTrigger.MANUAL,
297 parent_checkpoint_id=branch_latest.metadata.checkpoint_id,
298 )
300 logger.info(f"Merged branch {branch_thread_id} → {into_thread_id} (merge ckpt: {merge_id})")
301 return merge_id
303 # ── Decorator API ───────────────────────
305 def snapshot_on(self, trigger: SnapshotTrigger):
306 """装饰器：在调用前后自动快照。
308 Usage:
309 engine = CheckpointEngine(...)
311 @engine.snapshot_on(SnapshotTrigger.TOOL_CALL)
312 async def search_database(query: str): ...
313 """
314 def decorator(func):
315 @functools.wraps(func)
316 async def wrapper(*args, **kwargs):
317 thread_id = kwargs.pop("_checkpoint_thread_id", "default")
318 state = kwargs.pop("_checkpoint_state", {})
320 # Before snapshot
321 await self.snapshot_safe(
322 thread_id=thread_id,
323 messages=[{"role": "tool_call", "content": f"{func.__name__}({kwargs})"}],
324 state=state,
325 tools_result={},
326 trigger=trigger,
327 )
329 result = await func(*args, **kwargs)
331 # After snapshot
332 await self.snapshot_safe(
333 thread_id=thread_id,
334 messages=[{"role": "tool_result", "content": str(result)[:500]}],
335 state=state,
336 tools_result={"result": str(result)[:1000]},
337 trigger=trigger,
338 )
340 return result
341 return wrapper
342 return decorator
344 @asynccontextmanager
345 async def snapshot_scope(
346 self,
347 thread_id: str,
348 state: Dict[str, Any],
349 trigger: SnapshotTrigger = SnapshotTrigger.STATE_CHANGE,
350 ):
351 """上下文管理器：进入和退出作用域时自动快照。
353 Usage:
354 async with engine.snapshot_scope("thread-1", state):
355 await execute_workflow(...)
356 """
357 await self.snapshot_safe(
358 thread_id=thread_id,
359 messages=[{"role": "system", "content": f"Enter scope ({trigger.value})"}],
360 state=state,
361 tools_result={},
362 trigger=trigger,
363 )
364 try:
365 yield
366 finally:
367 await self.snapshot_safe(
368 thread_id=thread_id,
369 messages=[{"role": "system", "content": f"Exit scope ({trigger.value})"}],
370 state=state,
371 tools_result={},
372 trigger=trigger,
373 )
375 # ── Query API ───────────────────────────
377 async def get_latest(self, thread_id: str) -> Optional[Checkpoint]:
378 return await self._checkpointer.get_latest(thread_id)
380 async def get_checkpoint_tree(
381 self, thread_id: str, limit: int = 200
382 ) -> Dict[str, Any]:
383 """获取线程的 checkpoint 树结构（用于可视化）。"""
384 checkpoints = await self._checkpointer.list_checkpoints(thread_id, limit=limit)
386 nodes: List[Dict] = []
387 edges: List[Dict] = []
388 by_id: Dict[str, CheckpointMetadata] = {}
390 for cp in checkpoints:
391 by_id[cp.checkpoint_id] = cp
392 nodes.append({
393 "id": cp.checkpoint_id,
394 "step": cp.step,
395 "tags": cp.tags,
396 "summary": cp.summary,
397 "created_at": cp.created_at,
398 })
400 for cp in checkpoints:
401 if cp.parent_checkpoint_id and cp.parent_checkpoint_id in by_id:
402 edges.append({
403 "from": cp.parent_checkpoint_id,
404 "to": cp.checkpoint_id,
405 })
407 return {
408 "thread_id": thread_id,
409 "total_checkpoints": len(checkpoints),
410 "nodes": nodes,
411 "edges": edges,
412 }
414 # ── Internal ────────────────────────────
416 def _auto_summary(
417 self, messages: List[Dict[str, Any]], state: Dict[str, Any]
418 ) -> str:
419 """自动生成 checkpoint 摘要。"""
420 if messages:
421 last = messages[-1]
422 role = last.get("role", "unknown")
423 content = str(last.get("content", ""))[:100]
424 return f"[{role}] {content}"
425 return f"State: {len(state)} keys"
427 async def _maybe_gc(self, thread_id: str):
428 """根据 GC 策略清理旧 checkpoint。"""
429 if self._config.gc_policy == CheckpointGC.KEEP_ALL:
430 return
432 checkpoints = await self._checkpointer.list_checkpoints(thread_id, limit=500)
434 if self._config.gc_policy == CheckpointGC.KEEP_LAST_N:
435 if len(checkpoints) > self._config.gc_param:
436 to_delete = sorted(checkpoints, key=lambda c: c.step)[
437 :len(checkpoints) - self._config.gc_param
438 ]
439 for cp in to_delete:
440 await self._checkpointer.delete_before(thread_id, cp.step + 1)
441 logger.debug(f"GC: removed {len(to_delete)} old checkpoints from {thread_id}")
443 elif self._config.gc_policy == CheckpointGC.KEEP_AGE:
444 cutoff = time.time() - self._config.gc_param
445 deleted = 0
446 for cp in checkpoints:
447 try:
448 created = __import__('datetime').datetime.fromisoformat(cp.created_at).timestamp()
449 if created < cutoff:
450 await self._checkpointer.delete_thread(cp.thread_id)
451 deleted += 1
452 except Exception:
453 continue
454 if deleted:
455 logger.debug(f"GC: removed {deleted} expired checkpoints from {thread_id}")
457 elif self._config.gc_policy == CheckpointGC.KEEP_MILESTONES:
458 if len(checkpoints) > self._config.gc_param:
459 # 保留 first, last, 和均匀分布的 milestones
460 sorted_cps = sorted(checkpoints, key=lambda c: c.step)
461 keep = {sorted_cps[0].step, sorted_cps[-1].step}
463 n_milestones = max(2, self._config.gc_param - 2)
464 step_size = max(1, len(sorted_cps) // n_milestones)
465 for i in range(1, n_milestones):
466 idx = i * step_size
467 if idx < len(sorted_cps):
468 keep.add(sorted_cps[idx].step)
470 for cp in sorted_cps:
471 if cp.step not in keep:
472 await self._checkpointer.delete_thread(cp.thread_id)
473 logger.debug(f"GC milestones: kept {len(keep)} of {len(sorted_cps)} in {thread_id}")
476# ── Quick Start ──────────────────────────────
479async def create_checkpoint_engine(
480 backend: str = "sqlite",
481 db_path: str = "checkpoints.db",
482) -> CheckpointEngine:
483 """快速创建 checkpoint 引擎。"""
484 from agentos.checkpoint.factory import create_checkpointer
486 checkpointer = create_checkpointer(backend, db_path=db_path)
487 return CheckpointEngine(checkpointer)