import logging from typing import Any import pytest from langchain_core.runnables import RunnableConfig from pydantic import BaseModel from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, create_checkpoint, empty_checkpoint, ) from langgraph.checkpoint.memory import InMemorySaver from langgraph.checkpoint.serde.jsonplus import ( JsonPlusSerializer, _warned_blocked_types, _warned_unregistered_types, ) class MemoryPydantic(BaseModel): foo: str @pytest.fixture(autouse=True) def _reset_warned_types() -> None: # Warning dedup state is process-global; reset per-test so each case sees # a fresh slate and assertions about warning emission are stable. _warned_unregistered_types.clear() _warned_blocked_types.clear() class TestMemorySaver: @pytest.fixture(autouse=True) def setup(self) -> None: self.memory_saver = InMemorySaver() # objects for test setup self.config_1: RunnableConfig = { "configurable": { "thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "1", } } self.config_2: RunnableConfig = { "configurable": { "thread_id": "thread-2", "checkpoint_ns": "", "checkpoint_id": "2", } } self.config_3: RunnableConfig = { "configurable": { "thread_id": "thread-2", "checkpoint_id": "2-inner", "checkpoint_ns": "inner", } } self.chkpnt_1: Checkpoint = empty_checkpoint() self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) self.chkpnt_3: Checkpoint = empty_checkpoint() self.metadata_1: CheckpointMetadata = { "source": "input", "step": 2, "writes": {}, "score": 1, } self.metadata_2: CheckpointMetadata = { "source": "loop", "step": 1, "writes": {"foo": "bar"}, "score": None, } self.metadata_3: CheckpointMetadata = {} def test_combined_metadata(self) -> None: config: RunnableConfig = { "configurable": { "thread_id": "thread-2", "checkpoint_ns": "", "__super_private_key": "super_private_value", }, "metadata": {"run_id": "my_run_id"}, } self.memory_saver.put( config, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"] ) checkpoint = self.memory_saver.get_tuple(config) assert checkpoint is not None assert checkpoint.metadata == { **self.metadata_2, "run_id": "my_run_id", } async def test_search(self) -> None: # set up test # save checkpoints self.memory_saver.put( self.config_1, self.chkpnt_1, self.metadata_1, self.chkpnt_1["channel_versions"], ) self.memory_saver.put( self.config_2, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"], ) self.memory_saver.put( self.config_3, self.chkpnt_3, self.metadata_3, self.chkpnt_3["channel_versions"], ) # call method / assertions query_1 = {"source": "input"} # search by 1 key query_2 = { "step": 1, "writes": {"foo": "bar"}, } # search by multiple keys query_3: dict[str, Any] = {} # search by no keys, return all checkpoints query_4 = {"source": "update", "step": 1} # no match search_results_1 = list(self.memory_saver.list(None, filter=query_1)) assert len(search_results_1) == 1 assert search_results_1[0].metadata == self.metadata_1 search_results_2 = list(self.memory_saver.list(None, filter=query_2)) assert len(search_results_2) == 1 assert search_results_2[0].metadata == self.metadata_2 search_results_3 = list(self.memory_saver.list(None, filter=query_3)) assert len(search_results_3) == 3 search_results_4 = list(self.memory_saver.list(None, filter=query_4)) assert len(search_results_4) == 0 # search by config (defaults to checkpoints across all namespaces) search_results_5 = list( self.memory_saver.list({"configurable": {"thread_id": "thread-2"}}) ) assert len(search_results_5) == 2 assert { search_results_5[0].config["configurable"]["checkpoint_ns"], search_results_5[1].config["configurable"]["checkpoint_ns"], } == {"", "inner"} # TODO: test before and limit params async def test_asearch(self) -> None: # set up test # save checkpoints self.memory_saver.put( self.config_1, self.chkpnt_1, self.metadata_1, self.chkpnt_1["channel_versions"], ) self.memory_saver.put( self.config_2, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"], ) self.memory_saver.put( self.config_3, self.chkpnt_3, self.metadata_3, self.chkpnt_3["channel_versions"], ) # call method / assertions query_1 = {"source": "input"} # search by 1 key query_2 = { "step": 1, "writes": {"foo": "bar"}, } # search by multiple keys query_3: dict[str, Any] = {} # search by no keys, return all checkpoints query_4 = {"source": "update", "step": 1} # no match search_results_1 = [ c async for c in self.memory_saver.alist(None, filter=query_1) ] assert len(search_results_1) == 1 assert search_results_1[0].metadata == self.metadata_1 search_results_2 = [ c async for c in self.memory_saver.alist(None, filter=query_2) ] assert len(search_results_2) == 1 assert search_results_2[0].metadata == self.metadata_2 search_results_3 = [ c async for c in self.memory_saver.alist(None, filter=query_3) ] assert len(search_results_3) == 3 search_results_4 = [ c async for c in self.memory_saver.alist(None, filter=query_4) ] assert len(search_results_4) == 0 async def test_memory_saver() -> None: from langgraph.checkpoint.memory import InMemorySaver memory_saver = InMemorySaver() assert isinstance(memory_saver, InMemorySaver) async with memory_saver as async_memory_saver: assert async_memory_saver is memory_saver with memory_saver as sync_memory_saver: assert sync_memory_saver is memory_saver def test_memory_saver_warns_on_unregistered_msgpack( caplog: pytest.LogCaptureFixture, ) -> None: serde = JsonPlusSerializer() memory_saver = InMemorySaver(serde=serde) obj = MemoryPydantic(foo="bar") checkpoint = empty_checkpoint() checkpoint["channel_values"] = {"foo": obj} checkpoint["channel_versions"] = {"foo": 1} config: RunnableConfig = { "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} } caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1}) result = memory_saver.get_tuple(new_config) assert result is not None assert "unregistered type" in caplog.text.lower() assert result.checkpoint["channel_values"]["foo"] == obj def test_memory_saver_allowlist_silences_warning( caplog: pytest.LogCaptureFixture, ) -> None: serde = JsonPlusSerializer( allowed_msgpack_modules=[("tests.test_memory", "MemoryPydantic")] ) memory_saver = InMemorySaver(serde=serde) obj = MemoryPydantic(foo="bar") checkpoint = empty_checkpoint() checkpoint["channel_values"] = {"foo": obj} checkpoint["channel_versions"] = {"foo": 1} config: RunnableConfig = { "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} } caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1}) result = memory_saver.get_tuple(new_config) assert result is not None assert "unregistered type" not in caplog.text.lower() assert result.checkpoint["channel_values"]["foo"] == obj def test_memory_saver_strict_blocks_unregistered( caplog: pytest.LogCaptureFixture, ) -> None: serde = JsonPlusSerializer(allowed_msgpack_modules=None) memory_saver = InMemorySaver(serde=serde) obj = MemoryPydantic(foo="bar") checkpoint = empty_checkpoint() checkpoint["channel_values"] = {"foo": obj} checkpoint["channel_versions"] = {"foo": 1} config: RunnableConfig = { "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} } caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1}) result = memory_saver.get_tuple(new_config) assert result is not None assert "blocked" in caplog.text.lower() expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict() assert result.checkpoint["channel_values"]["foo"] == expected def test_memory_saver_with_allowlist_proxy_isolated() -> None: serde = JsonPlusSerializer(allowed_msgpack_modules=None) memory_saver = InMemorySaver(serde=serde) proxy = memory_saver.with_allowlist([("tests.test_memory", "MemoryPydantic")]) obj = MemoryPydantic(foo="bar") checkpoint = empty_checkpoint() checkpoint["channel_values"] = {"foo": obj} checkpoint["channel_versions"] = {"foo": 1} config: RunnableConfig = { "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} } new_config = proxy.put(config, checkpoint, {}, {"foo": 1}) proxied = proxy.get_tuple(new_config) assert proxied is not None assert proxied.checkpoint["channel_values"]["foo"] == obj direct = memory_saver.get_tuple(new_config) 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 TestInMemorySaverDeltaChannel: def test_delta_channel_chain_reconstruction(self) -> None: """_load_blobs follows the diff chain and returns DeltaChainValue.""" from langgraph.checkpoint.base import DeltaChainValue, DeltaValue 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 = DeltaValue(delta=["msg1"], prev_version=None) delta2 = DeltaValue(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, DeltaChainValue) assert result.base is None assert result.deltas == [["msg1"], ["msg2"]] def test_delta_channel_mixed_old_and_new_blobs(self) -> None: """When chain hits an old non-diff blob, it becomes base.""" from langgraph.checkpoint.base import DeltaChainValue, DeltaValue 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 = DeltaValue(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, DeltaChainValue) assert result.base == ["old_msg"] assert result.deltas == [["new_msg"]]