fix(checkpoint): replay migrated delta writes through a plain seed

InMemorySaver.get_delta_channel_history skipped on-path writes when the
terminating ancestor's blob was a plain (pre-delta migration) value
instead of a _DeltaSnapshot, dropping post-migration writes on reload.
Collect writes regardless of seed type, matching the base implementation.
This commit is contained in:
Christian Bromann
2026-06-01 12:54:11 -07:00
parent 83dd61feac
commit b60ef4adb0
2 changed files with 56 additions and 12 deletions
@@ -148,16 +148,16 @@ class InMemorySaver(
whose stored blob is non-empty. Other channels keep walking until
they find their own terminator or hit the root.
Pre-delta plain-value blobs subsume their ancestor's pending
writes (the value already includes them); `_DeltaSnapshot` blobs
do not (snapshot is the value AT that ancestor, prior to its own
pending writes that produce the child).
The seed value (whether a `_DeltaSnapshot` or a plain pre-delta
migration blob) is the value AT that ancestor, prior to its own
pending writes that produce the child. Those on-path writes —
including the ones stored on the terminating ancestor — are always
collected and replayed on top of the seed, so a thread migrated from
a pre-delta channel does not drop the writes saved under the
migration boundary checkpoint.
"""
if not channels:
return {}
# Imported lazily to avoid a hard checkpoint→serde-types coupling at
# module import; only this override needs the runtime check.
from langgraph.checkpoint.serde.types import _DeltaSnapshot
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -205,11 +205,12 @@ class InMemorySaver(
):
if ch not in remaining:
continue
blob_value = blob_value_by_ch.get(ch)
if blob_value is not None and not isinstance(
blob_value, _DeltaSnapshot
):
continue
# Collect on-path writes regardless of seed type. A plain
# (pre-delta migration) blob is the settled value AT that
# ancestor; its own pending writes produce the child and must
# still be replayed, just like a `_DeltaSnapshot` seed.
# Skipping them would drop post-migration writes saved under
# the migration boundary checkpoint.
collected_by_ch[ch].append(
(tid, ch, self.serde.loads_typed(serialized))
)
@@ -616,3 +616,46 @@ async def test_add_messages_to_delta_migration_preserves_message_history_async()
assert [m.id for m in snap.values["messages"]] == ["h1", "a1"], (
f"async tip hydration mismatch: got {[m.id for m in snap.values['messages']]}"
)
def test_post_migration_write_survives_reload_through_plain_seed() -> None:
"""A write made on the first post-migration super-step must be preserved
when the checkpoint is reloaded (reconstructed via the plain seed)."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "post-migration-reload"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
delta = _delta_graph(checkpointer)
live_result = delta.invoke({"items": ["POST"]}, config)
reloaded = delta.get_state(config)
assert list(live_result.get("items", [])) == ["u0", "u1", "u2", "POST"], (
f"sanity: live invoke should include POST, got {live_result.get('items')}"
)
assert list(reloaded.values.get("items", [])) == ["u0", "u1", "u2", "POST"], (
"post-migration write dropped on reload through plain seed: "
f"got {reloaded.values.get('items')}"
)
def test_post_migration_reload_base_matches_optimized_override() -> None:
"""The reference `BaseCheckpointSaver` path and the optimized
`InMemorySaver` override must agree once a post-migration write is
reconstructed through a plain seed (the scenario that triggers the
write-collection guard)."""
def _run(saver: Any) -> list:
config = {"configurable": {"thread_id": "parity"}}
_drive(_binop_graph(saver), config, "u", 2)
delta = _delta_graph(saver)
delta.invoke({"items": ["POST"]}, config)
return list(delta.get_state(config).values.get("items", []))
fast = _run(InMemorySaver())
slow = _run(_ThirdPartyStyleSaver())
assert fast == ["u0", "u1", "POST"], f"optimized override wrong: {fast}"
assert slow == fast, f"base fallback diverged from override: {slow} != {fast}"