diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index da353c2b4..17829b7a0 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -123,13 +123,40 @@ class InMemorySaver( def _load_blobs( self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions ) -> dict[str, Any]: + from langgraph.checkpoint.base import DiffChainValue + channel_values: dict[str, Any] = {} + diff_channels: dict[str, Any] = {} + for k, v in versions.items(): kk = (thread_id, checkpoint_ns, k, v) - if kk in self.blobs: + if kk not in self.blobs: + continue + vv = self.blobs[kk] + if vv[0] == "diff": + diff_channels[k] = v + elif vv[0] != "empty": + channel_values[k] = self.serde.loads_typed(vv) + + for k, current_version in diff_channels.items(): + chain_deltas: list[list[Any]] = [] + base: list[Any] | None = None + version: str | None = current_version + while version is not None: + kk = (thread_id, checkpoint_ns, k, version) + if kk not in self.blobs: + break vv = self.blobs[kk] - if vv[0] != "empty": - channel_values[k] = self.serde.loads_typed(vv) + if vv[0] == "diff": + payload = self.serde.loads_typed(vv) # {"d": [...], "p": version|None} + chain_deltas.append(payload["d"]) + version = payload["p"] + else: + base = self.serde.loads_typed(vv) + break + chain_deltas.reverse() + channel_values[k] = DiffChainValue(base=base, deltas=chain_deltas) + return channel_values def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index a68a23d90..a0cb3407b 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -308,3 +308,58 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None: assert direct is not None expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict() assert direct.checkpoint["channel_values"]["foo"] == expected + + +class TestInMemorySaverDiffChannel: + def test_diff_channel_chain_reconstruction(self) -> None: + """_load_blobs follows the diff chain and returns DiffChainValue.""" + from langgraph.checkpoint.base import DiffChainValue, DiffDelta + + saver = InMemorySaver() + serde = JsonPlusSerializer() + + thread_id = "t1" + ns = "" + + # Simulate two steps: v1 (root) and v2 (chained to v1) + v1 = "00000000000000000000000000000001.1234567890000000" + v2 = "00000000000000000000000000000002.1234567890000000" + + delta1 = DiffDelta(delta=["msg1"], prev_version=None) + delta2 = DiffDelta(delta=["msg2"], prev_version=v1) + + saver.blobs[(thread_id, ns, "messages", v1)] = serde.dumps_typed(delta1) + saver.blobs[(thread_id, ns, "messages", v2)] = serde.dumps_typed(delta2) + + channel_values = saver._load_blobs(thread_id, ns, {"messages": v2}) + + assert "messages" in channel_values + result = channel_values["messages"] + assert isinstance(result, DiffChainValue) + assert result.base is None + assert result.deltas == [["msg1"], ["msg2"]] + + def test_diff_channel_mixed_old_and_new_blobs(self) -> None: + """When chain hits an old non-diff blob, it becomes base.""" + from langgraph.checkpoint.base import DiffChainValue, DiffDelta + + saver = InMemorySaver() + serde = JsonPlusSerializer() + + thread_id = "t2" + ns = "" + + v_old = "00000000000000000000000000000001.0000000000000000" + v_new = "00000000000000000000000000000002.0000000000000000" + + # Old-style full-list blob + saver.blobs[(thread_id, ns, "messages", v_old)] = serde.dumps_typed(["old_msg"]) + # New diff blob chained to old + delta = DiffDelta(delta=["new_msg"], prev_version=v_old) + saver.blobs[(thread_id, ns, "messages", v_new)] = serde.dumps_typed(delta) + + channel_values = saver._load_blobs(thread_id, ns, {"messages": v_new}) + result = channel_values["messages"] + assert isinstance(result, DiffChainValue) + assert result.base == ["old_msg"] + assert result.deltas == [["new_msg"]]