StateSnapshot:每个超步结束后的状态快照

源码版本1.2.9

职责

StateSnapshot 是 LangGraph 在每个超步 (superstep) 边界上对外暴露的「这一步刚跑完时整张图长什么样」的不可变快照。它把五件事打包成一个 NamedTuple:当前各通道 (channel) 的值 (values)、下一步要跑哪些节点 (next)、本次用的 config、对应的 checkpoint 元数据 (metadata + created_at)、以及父检查点的 config (StateSnapshot 定义:643-661)。你调 graph.get_state(config)、graph.aget_state(config),或者遍历 graph.get_state_history(config),拿回来的都是这个类型。

它在架构里处在「检查点 (checkpoint) 之上、用户视图之下」。检查点 (CheckpointTuple) 是引擎内部的序列化格式,直接存的是 channel_values / channel_versions / pending_writes 这些低层字段,看不懂业务状态长什么样;StateSnapshot 则把这层翻译成「业务能直接读的 dict + 下一步要跑什么 + 任务列表」(_prepare_state_snapshot:1145-1162)。子图状态也走它——task.state 字段在递归 subgraphs=True 时会塞一个子图的 StateSnapshot(PregelTask.state:200-201)。

StateSnapshot 本身是 NamedTuple,不可变;要「改状态」必须通过 graph.update_state(config, values) 写新值,引擎会把它变成 NULL_TASK_ID 的 pending write 落到下一个检查点里,然后再生成一份新的 StateSnapshot。

设计动机

为什么不直接把 CheckpointTuple 暴露给用户?

• 解耦内部表示:检查点格式为了高效序列化,存的是 channel_versions 这种带版本号的扁平 dict;而用户关心的是「我的业务 state dict」,中间需要 read_channels 把通道值读出来 (read_channels:1258)。多一层 StateSnapshot 让检查点格式可以自由演进,不破坏对外 API。
• 带上「下一步」语义:引擎知道下一步要跑哪些节点,但 CheckpointTuple 本身只存写过的值;StateSnapshot 在生成时调 prepare_next_tasks(prepare_next_tasks:1178-1195),算出 next=(node_name, ...),直接回答「接下来会发生什么」。
• 任务列表可见:每个 StateSnapshot.tasks 是当前步要执行的 PregelTask 元组(tasks:658),包含 writes / interrupts / 子图状态——这是人机协作里判断「现在卡在哪个节点、等什么输入」的关键。
• 不可变 = 可对比:历史快照用 parent_config 串成链表(parent_config:656),get_state_history 拿到的就是按时间倒序的快照列表,任意一份都可以用 config 重新指向它做时间旅行。
• pending writes 透明应用:get_state 默认会把还没落检查点的 pending_writes 临时应用进通道再读 (apply pending writes:1239-1249),所以你在节点刚写完但还没落 checkpoint 时调 get_state,看到的就是最新值,而不是上一步的旧值。

关键文件

• StateSnapshot NamedTuple:643-661 — 7 个字段:values / next / config / metadata / created_at / parent_config / tasks / interrupts。
• PregelTask:200-217 — 任务对象,state 字段在递归时存子图的 StateSnapshot 或定位它的 RunnableConfig。
• _prepare_state_snapshot:1145-1162 — 空 saved 时的兜底:返回 values={} / next=() / tasks=() 的空快照。
• step 计算 + prepare_next_tasks:1167-1195 — 从 saved.metadata["step"] 推当前步号,调 prepare_next_tasks 算出本步要跑的任务字典。
• task_states 递归:1199-1227 — 对每个 task.name in subgraphs 的任务,组装 checkpoint_ns 后递归调子图的 get_state。
• apply pending writes:1229-1255 — 把 NULL_TASK_ID 的写和常规 pending writes 临时应用到通道,再算 tasks_with_writes。
• 组装 StateSnapshot:1257-1266 — 把通道值、next 节点元组、config、metadata、ts、parent_config、tasks、interrupts 拼成最终 StateSnapshot。
• get_state:1392-1434 — 同步入口:解析 checkpointer、处理 checkpoint_ns 路由到子图、决定是否 apply_pending_writes。
• get_state_history:1480 — 遍历同一 thread_id 下所有历史检查点,每个都生成 StateSnapshot。
• _aprepare_state_snapshot:1268 — 异步版,逻辑跟同步版对称,只是 channels_from_checkpoint 换成 achannels_from_checkpoint。

数据流

下面是 _prepare_state_snapshot 把检查点翻译成 StateSnapshot 的关键段:从 saved.metadata 推步号、调 prepare_next_tasks 拿到本步任务,再处理子图递归。

step = saved.metadata.get("step", -1) + 1
stop = step + 2
channels, managed = channels_from_checkpoint(
    self.channels,
    saved.checkpoint,
    saver=self.checkpointer
    if isinstance(self.checkpointer, BaseCheckpointSaver) else None,
    config=saved.config,
)
next_tasks = prepare_next_tasks(
    saved.checkpoint,
    saved.pending_writes or [],
    self.nodes,
    channels,
    managed,
    saved.config,
    step,
    stop,
    for_execution=True,
    store=self.store,
    checkpointer=(self.checkpointer
                  if isinstance(self.checkpointer, BaseCheckpointSaver) else None),
    manager=None,
)

(prepare_next_tasks 调用:1167-1195)

最后用读出来的通道值、过滤掉已写了 pending 的任务名,组装成不可变快照:

return StateSnapshot(
    read_channels(channels, self.stream_channels_asis),
    tuple(t.name for t in next_tasks.values() if not t.writes),
    patch_checkpoint_map(saved.config, saved.metadata),
    saved.metadata,
    saved.checkpoint["ts"],
    patch_checkpoint_map(saved.parent_config, saved.metadata),
    tasks_with_writes,
    tuple([i for task in tasks_with_writes for i in task.interrupts]),
)

(组装:1257-1266)

边界与失败

• 没有检查点时返回空快照而非报错(空 saved:1152-1162),values={} / next=(),调用方需要自己判断 snapshot.next 是不是空 tuple 来区分「图还没跑」和「图跑完了」。
• 没设 checkpointer 直接调 get_state 抛 ValueError(no checkpointer:1401-1402),No checkpointer set——纯内存跑的图没有历史快照。
• checkpoint_ns 找不到匹配子图时报错(subgraph not found:1415-1416),Subgraph {recast} not found,常见于 namespace 拼接错了 NS_SEP(|) 或 NS_END(:)。
• apply_pending_writes 仅在没显式给 checkpoint_id 时启用(apply 条件:1433),也就是说你拿历史快照看的就是当时落盘的值,不会被后续 pending write 污染;拿「最新」快照才会临时应用 pending。
• next 字段过滤掉已经有 writes 的任务(next 过滤:1259),t.writes 非空的任务不算「next」——因为它们的输入已经被本步消费了,所以 next 反映的是「还没跑的任务」。
• tasks 包含 interrupts(interrupts 聚合:1265),把所有任务的 Interrupt 对象扁平化拼成元组,用户调 Command(resume=...) 时按这个顺序回填。

小结

StateSnapshot 是 LangGraph 把内部检查点格式翻译成业务可读视图的边界类型——它不存数据,只把检查点 + pending writes 加上「下一步」语义组装成不可变快照。所有断点续跑、时间旅行、人机协作的 API 都建立在它之上。

它跟 RunnableConfig 是一对:config 是「怎么找到这个快照」的坐标,StateSnapshot 是「找到之后看到什么」。继续看 StateGraph 怎么定义出被快照的状态形状,或看 Pregel 引擎 里 _prepare_state_snapshot 是怎么被 get_state 调到的。

对照官方资料:LangGraph 文档 · README。