Add history tracking to in-memory, sqlite and aiosqlite checkpointers

- Add Pregel.get_state_history and .aget_state_history methods to get history iterator - Update checkpointer base class with new list and get_tuple methods - Rewrite in-memory checkpointer class to track history - Rewrite sqlite and aiosqlite checkpointers to track history - Add new tests for history tracking
2026-09-03 08:18:47 +02:00 · 2024-03-13 17:22:39 -07:00
parent 418267e4de
commit 0d13c6b159
8 changed files with 486 additions and 180 deletions
@@ -1,15 +1,16 @@
 import pickle
-from typing import Optional
+from contextlib import AbstractAsyncContextManager
+from types import TracebackType
+from typing import AsyncIterator, Optional, Self

 import aiosqlite
 from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.utils import ConfigurableFieldSpec

-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
+from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple


-class AsyncSqliteSaver(BaseCheckpointSaver):
+class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
    conn: aiosqlite.Connection

    is_setup: bool = Field(False, init=False, repr=False)
@@ -21,65 +22,92 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
    def from_conn_string(cls, conn_string: str) -> "AsyncSqliteSaver":
        return AsyncSqliteSaver(conn=aiosqlite.connect(conn_string))

-    @property
-    def config_specs(self) -> list[ConfigurableFieldSpec]:
-        return [
-            ConfigurableFieldSpec(
-                id="thread_id",
-                annotation=str,
-                name="Thread ID",
-                description=None,
-                default="",
-                is_shared=True,
-            ),
-        ]
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> bool | None:
+        return await self.conn.close()

    async def setup(self) -> None:
-        print("hello")
        if self.is_setup:
            return

-        try:
-            await self.conn
-            await self.conn.executescript(
-                """
-                CREATE TABLE IF NOT EXISTS checkpoints (
-                    thread_id TEXT PRIMARY KEY,
-                    checkpoint BLOB
-                );
-                """
-            )
+        await self.conn
+        async with self.conn.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS checkpoints (
+                thread_id TEXT NOT NULL,
+                thread_ts TEXT NOT NULL,
+                checkpoint BLOB,
+                PRIMARY KEY (thread_id, thread_ts)
+            );
+            """
+        ):
            await self.conn.commit()

-            print("good bye")
+        self.is_setup = True

-            self.is_setup = True
-        except BaseException as e:
-            print(e)
-            raise e
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        await self.setup()
+        if config["configurable"].get("thread_ts"):
+            async with self.conn.execute(
+                "SELECT checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
+                (
+                    config["configurable"]["thread_id"],
+                    config["configurable"]["thread_ts"],
+                ),
+            ) as cursor:
+                if value := await cursor.fetchone():
+                    return CheckpointTuple(config, pickle.loads(value[0]))
+        else:
+            async with self.conn.execute(
+                "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
+                (config["configurable"]["thread_id"],),
+            ) as cursor:
+                if value := await cursor.fetchone():
+                    return CheckpointTuple(
+                        {
+                            "configurable": {
+                                "thread_id": value[0],
+                                "thread_ts": value[1],
+                            }
+                        },
+                        pickle.loads(value[2]),
+                    )

-    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
-        raise NotImplementedError
-
-    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
-        raise NotImplementedError
-
-    async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
+    async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]:
        await self.setup()
        async with self.conn.execute(
-            "SELECT checkpoint FROM checkpoints WHERE thread_id = ?",
+            "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
            (config["configurable"]["thread_id"],),
        ) as cursor:
-            if value := await cursor.fetchone():
-                return pickle.loads(value[0])
+            async for thread_id, thread_ts, value in cursor:
+                yield CheckpointTuple(
+                    {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
+                    pickle.loads(value),
+                )

-    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+    async def aput(
+        self, config: RunnableConfig, checkpoint: Checkpoint
+    ) -> RunnableConfig:
        await self.setup()
-        await self.conn.execute(
-            "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)",
+        async with self.conn.execute(
+            "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, checkpoint) VALUES (?, ?, ?)",
            (
                config["configurable"]["thread_id"],
+                checkpoint["ts"],
                pickle.dumps(checkpoint),
            ),
-        )
-        await self.conn.commit()
+        ):
+            await self.conn.commit()
+        return {
+            "configurable": {
+                "thread_id": config["configurable"]["thread_id"],
+                "thread_ts": checkpoint["ts"],
+            }
+        }
@@ -1,9 +1,9 @@
 import asyncio
-from abc import ABC, abstractmethod
+from abc import ABC
 from collections import defaultdict
 from copy import deepcopy
 from datetime import datetime, timezone
-from typing import Any, Optional, TypedDict
+from typing import Any, AsyncIterator, Iterator, NamedTuple, Optional, TypedDict

 from langchain_core.load.serializable import Serializable
 from langchain_core.runnables import RunnableConfig
@@ -49,25 +49,71 @@ class CheckpointAt(StrEnum):
    END_OF_RUN = "end_of_run"


+class CheckpointTuple(NamedTuple):
+    config: RunnableConfig
+    checkpoint: Checkpoint
+
+
+CheckpointThreadId = ConfigurableFieldSpec(
+    id="thread_id",
+    annotation=str,
+    name="Thread ID",
+    description=None,
+    default="",
+    is_shared=True,
+)
+
+CheckpointThreadTs = ConfigurableFieldSpec(
+    id="thread_ts",
+    annotation=Optional[str],
+    name="Thread Timestamp",
+    description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.",
+    default=None,
+    is_shared=True,
+)
+
+
 class BaseCheckpointSaver(Serializable, ABC):
    at: CheckpointAt = CheckpointAt.END_OF_RUN

    @property
    def config_specs(self) -> list[ConfigurableFieldSpec]:
-        return []
+        return [CheckpointThreadId, CheckpointThreadTs]

-    @abstractmethod
    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
-        ...
+        if value := self.get_tuple(config):
+            return value.checkpoint

-    @abstractmethod
-    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
-        ...
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        raise NotImplementedError
+
+    def list(self, config: RunnableConfig) -> Iterator[CheckpointTuple]:
+        raise NotImplementedError
+
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
+        raise NotImplementedError

    async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
-        return await asyncio.get_running_loop().run_in_executor(None, self.get, config)
+        if value := await self.aget_tuple(config):
+            return value.checkpoint

-    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.get_tuple, config
+        )
+
+    async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]:
+        loop = asyncio.get_running_loop()
+        iter = loop.run_in_executor(None, self.list, config)
+        while True:
+            try:
+                yield await loop.run_in_executor(None, next, iter)
+            except StopIteration:
+                return
+
+    async def aput(
+        self, config: RunnableConfig, checkpoint: Checkpoint
+    ) -> RunnableConfig:
        return await asyncio.get_running_loop().run_in_executor(
            None, self.put, config, checkpoint
        )
@@ -1,30 +1,47 @@
+from collections import defaultdict
 from typing import Optional

 from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.utils import ConfigurableFieldSpec

-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
+from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple


 class MemorySaver(BaseCheckpointSaver):
-    storage: dict[str, Checkpoint] = Field(default_factory=dict)
-
-    @property
-    def config_specs(self) -> list[ConfigurableFieldSpec]:
-        return [
-            ConfigurableFieldSpec(
-                id="thread_id",
-                annotation=str,
-                name="Thread ID",
-                description=None,
-                default="",
-                is_shared=True,
-            ),
-        ]
+    storage: defaultdict[str, dict[str, Checkpoint]] = Field(
+        default_factory=lambda: defaultdict(dict)
+    )

    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
-        return self.storage.get(config["configurable"]["thread_id"], None)
+        if value := self.get_tuple(config):
+            return value.checkpoint

-    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
-        return self.storage.update({config["configurable"]["thread_id"]: checkpoint})
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        if config["configurable"].get("thread_ts"):
+            if checkpoint := self.storage[config["configurable"]["thread_id"]].get(
+                config["configurable"]["thread_ts"]
+            ):
+                return CheckpointTuple(config=config, checkpoint=checkpoint)
+        else:
+            if checkpoints := self.storage[config["configurable"]["thread_id"]]:
+                thread_ts = max(checkpoints.keys())
+                return CheckpointTuple(
+                    config={
+                        "configurable": {
+                            "thread_id": config["configurable"]["thread_id"],
+                            "thread_ts": thread_ts,
+                        }
+                    },
+                    checkpoint=checkpoints[thread_ts],
+                )
+
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
+        self.storage[config["configurable"]["thread_id"]].update(
+            {checkpoint["ts"]: checkpoint}
+        )
+        return {
+            "configurable": {
+                "thread_id": config["configurable"]["thread_id"],
+                "thread_ts": checkpoint["ts"],
+            }
+        }
@@ -1,16 +1,16 @@
 import pickle
 import sqlite3
-from contextlib import contextmanager
-from typing import Optional
+from contextlib import AbstractContextManager, contextmanager
+from types import TracebackType
+from typing import AsyncIterator, Iterator, Optional, Self

 from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.utils import ConfigurableFieldSpec

-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
+from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple


-class SqliteSaver(BaseCheckpointSaver):
+class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
    conn: sqlite3.Connection

    is_setup: bool = Field(False, init=False, repr=False)
@@ -22,18 +22,16 @@ class SqliteSaver(BaseCheckpointSaver):
    def from_conn_string(cls, conn_string: str) -> "SqliteSaver":
        return SqliteSaver(conn=sqlite3.connect(conn_string))

-    @property
-    def config_specs(self) -> list[ConfigurableFieldSpec]:
-        return [
-            ConfigurableFieldSpec(
-                id="thread_id",
-                annotation=str,
-                name="Thread ID",
-                description=None,
-                default="",
-                is_shared=True,
-            ),
-        ]
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> bool | None:
+        return self.conn.close()

    def setup(self) -> None:
        if self.is_setup:
@@ -42,8 +40,10 @@ class SqliteSaver(BaseCheckpointSaver):
        self.conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS checkpoints (
-                thread_id TEXT PRIMARY KEY,
-                checkpoint BLOB
+                thread_id TEXT NOT NULL,
+                thread_ts TEXT NOT NULL,
+                checkpoint BLOB,
+                PRIMARY KEY (thread_id, thread_ts)
            );
            """
        )
@@ -61,27 +61,72 @@ class SqliteSaver(BaseCheckpointSaver):
                self.conn.commit()
            cur.close()

-    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        with self.cursor(transaction=False) as cur:
+            if config["configurable"].get("thread_ts"):
+                cur.execute(
+                    "SELECT checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
+                    (
+                        config["configurable"]["thread_id"],
+                        config["configurable"]["thread_ts"],
+                    ),
+                )
+                if value := cur.fetchone():
+                    return CheckpointTuple(config, pickle.loads(value[0]))
+            else:
+                cur.execute(
+                    "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
+                    (config["configurable"]["thread_id"],),
+                )
+                if value := cur.fetchone():
+                    return CheckpointTuple(
+                        {
+                            "configurable": {
+                                "thread_id": value[0],
+                                "thread_ts": value[1],
+                            }
+                        },
+                        pickle.loads(value[2]),
+                    )
+
+    def list(self, config: RunnableConfig) -> Iterator[CheckpointTuple]:
        with self.cursor(transaction=False) as cur:
            cur.execute(
-                "SELECT checkpoint FROM checkpoints WHERE thread_id = ?",
+                "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
                (config["configurable"]["thread_id"],),
            )
-            if value := cur.fetchone():
-                return pickle.loads(value[0])
+            for thread_id, thread_ts, value in cur:
+                yield CheckpointTuple(
+                    {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
+                    pickle.loads(value),
+                )

-    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
        with self.cursor() as cur:
            cur.execute(
-                "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)",
+                "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, checkpoint) VALUES (?, ?, ?)",
                (
                    config["configurable"]["thread_id"],
+                    checkpoint["ts"],
                    pickle.dumps(checkpoint),
                ),
            )
+        return {
+            "configurable": {
+                "thread_id": config["configurable"]["thread_id"],
+                "thread_ts": checkpoint["ts"],
+            }
+        }

-    async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
        raise NotImplementedError

-    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+    async def alist(
+        self, config: RunnableConfig
+    ) -> AsyncIterator[tuple[RunnableConfig, Checkpoint]]:
+        raise NotImplementedError
+
+    async def aput(
+        self, config: RunnableConfig, checkpoint: Checkpoint
+    ) -> RunnableConfig:
        raise NotImplementedError
@@ -163,6 +163,8 @@ class StateSnapshot(NamedTuple):
    """Current values of channels"""
    next: tuple[str]
    """Nodes to execute in the next step, if any"""
+    config: RunnableConfig
+    """Config used to fetch this snapshot"""


 class Pregel(
@@ -278,8 +280,9 @@ class Pregel(
        if not self.checkpointer:
            raise ValueError("No checkpointer set")

-        checkpoint = self.checkpointer.get(config)
-        checkpoint = checkpoint or empty_checkpoint()
+        saved = self.checkpointer.get_tuple(config)
+        checkpoint = saved.checkpoint if saved else empty_checkpoint()
+        config = saved.config if saved else config
        with ChannelsManager(self.channels, checkpoint) as channels:
            _, next_tasks = _prepare_next_tasks(
                checkpoint, self.nodes, channels, update_seen=False
@@ -294,14 +297,16 @@ class Pregel(
                if isinstance(self.snapshot_channels, str)
                else values,
                tuple(name for _, _, name in next_tasks),
+                config,
            )

    async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
        if not self.checkpointer:
            raise ValueError("No checkpointer set")

-        checkpoint = await self.checkpointer.aget(config)
-        checkpoint = checkpoint or empty_checkpoint()
+        saved = await self.checkpointer.aget_tuple(config)
+        checkpoint = saved.checkpoint if saved else empty_checkpoint()
+        config = saved.config if saved else config
        async with AsyncChannelsManager(self.channels, checkpoint) as channels:
            _, next_tasks = _prepare_next_tasks(
                checkpoint, self.nodes, channels, update_seen=False
@@ -316,11 +321,58 @@ class Pregel(
                if isinstance(self.snapshot_channels, str)
                else values,
                tuple(name for _, _, name in next_tasks),
+                config,
            )

+    def get_state_history(self, config: RunnableConfig) -> Iterator[StateSnapshot]:
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        for config, checkpoint in self.checkpointer.list(config):
+            with ChannelsManager(self.channels, checkpoint) as channels:
+                _, next_tasks = _prepare_next_tasks(
+                    checkpoint, self.nodes, channels, update_seen=False
+                )
+                values = {
+                    k: _read_channel(channels, k)
+                    for k in channels
+                    if k in self.snapshot_channels_list
+                }
+                yield StateSnapshot(
+                    values[self.snapshot_channels]
+                    if isinstance(self.snapshot_channels, str)
+                    else values,
+                    tuple(name for _, _, name in next_tasks),
+                    config,
+                )
+
+    async def aget_state_history(
+        self, config: RunnableConfig
+    ) -> AsyncIterator[StateSnapshot]:
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        async for config, checkpoint in self.checkpointer.alist(config):
+            async with AsyncChannelsManager(self.channels, checkpoint) as channels:
+                _, next_tasks = _prepare_next_tasks(
+                    checkpoint, self.nodes, channels, update_seen=False
+                )
+                values = {
+                    k: _read_channel(channels, k)
+                    for k in channels
+                    if k in self.snapshot_channels_list
+                }
+                yield StateSnapshot(
+                    values[self.snapshot_channels]
+                    if isinstance(self.snapshot_channels, str)
+                    else values,
+                    tuple(name for _, _, name in next_tasks),
+                    config,
+                )
+
    def update_state(
        self, config: RunnableConfig, values: dict[str, Any] | Any
-    ) -> None:
+    ) -> RunnableConfig:
        if not self.checkpointer:
            raise ValueError("No checkpointer set")

@@ -338,11 +390,13 @@ class Pregel(
            for k in self.snapshot_channels_list:
                version = checkpoint["channel_versions"][k]
                checkpoint["versions_seen"][INTERRUPT][k] = version
-            self.checkpointer.put(config, create_checkpoint(checkpoint, channels))
+            return self.checkpointer.put(
+                config, create_checkpoint(checkpoint, channels)
+            )

    async def aupdate_state(
        self, config: RunnableConfig, values: dict[str, Any] | Any
-    ) -> None:
+    ) -> RunnableConfig:
        if not self.checkpointer:
            raise ValueError("No checkpointer set")

@@ -360,7 +414,7 @@ class Pregel(
            for k in self.snapshot_channels or self.channels:
                version = checkpoint["channel_versions"][k]
                checkpoint["versions_seen"][INTERRUPT][k] = version
-            await self.checkpointer.aput(
+            return await self.checkpointer.aput(
                config, create_checkpoint(checkpoint, channels)
            )

@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 from langchain_core.pydantic_v1 import Field

 from langgraph.checkpoint.base import Checkpoint, CheckpointAt, copy_checkpoint
@@ -5,15 +7,19 @@ from langgraph.checkpoint.memory import MemorySaver


 class MemorySaverAssertImmutable(MemorySaver):
-    storage_for_copies: dict[str, Checkpoint] = Field(default_factory=dict)
+    storage_for_copies: defaultdict[str, dict[str, Checkpoint]] = Field(
+        default_factory=lambda: defaultdict(dict)
+    )

    at = CheckpointAt.END_OF_STEP

-    def put(self, config: dict, checkpoint: dict) -> None:
+    def put(self, config: dict, checkpoint: Checkpoint) -> None:
        # assert checkpoint hasn't been modified since last written
        thread_id = config["configurable"]["thread_id"]
        if saved := super().get(config):
-            assert self.storage_for_copies[thread_id] == saved
-        self.storage_for_copies[thread_id] = copy_checkpoint(checkpoint)
+            assert self.storage_for_copies[thread_id][saved["ts"]] == saved
+        self.storage_for_copies[thread_id][checkpoint["ts"]] = copy_checkpoint(
+            checkpoint
+        )
        # call super to write checkpoint
        super().put(config, checkpoint)
@@ -503,39 +503,79 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
        | raise_if_above_10
    )

-    memory = SqliteSaver.from_conn_string(":memory:")
+    with SqliteSaver.from_conn_string(":memory:") as memory:
+        app = Pregel(
+            nodes={"one": one},
+            channels={"total": BinaryOperatorAggregate(int, operator.add)},
+            checkpointer=memory,
+        )

-    app = Pregel(
-        nodes={"one": one},
-        channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpointer=memory,
-    )
+        thread_1 = {"configurable": {"thread_id": "1"}}
+        # total starts out as 0, so output is 0+2=2
+        assert app.invoke(2, thread_1) == 2
+        state = app.get_state(thread_1)
+        assert state is not None
+        assert state.values.get("total") == 2
+        assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
+        # total is now 2, so output is 2+3=5
+        assert app.invoke(3, thread_1) == 5
+        state = app.get_state(thread_1)
+        assert state is not None
+        assert state.values.get("total") == 7
+        assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
+        # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+        with pytest.raises(ValueError):
+            app.invoke(4, thread_1)
+        # checkpoint is not updated
+        state = app.get_state(thread_1)
+        assert state is not None
+        assert state.values.get("total") == 7

-    # total starts out as 0, so output is 0+2=2
-    assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
-    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 2
-    # total is now 2, so output is 2+3=5
-    assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
-    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 7
-    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
-    with pytest.raises(ValueError):
-        app.invoke(4, {"configurable": {"thread_id": "1"}})
-    # checkpoint is not updated
-    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 7
-    # on a new thread, total starts out as 0, so output is 0+5=5
-    assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
-    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 7
-    checkpoint = memory.get({"configurable": {"thread_id": "2"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 5
+        thread_2 = {"configurable": {"thread_id": "2"}}
+        # on a new thread, total starts out as 0, so output is 0+5=5
+        assert app.invoke(5, thread_2) == 5
+        state = app.get_state({"configurable": {"thread_id": "1"}})
+        assert state is not None
+        assert state.values.get("total") == 7
+        state = app.get_state(thread_2)
+        assert state is not None
+        assert state.values.get("total") == 5
+
+        # list all checkpoints for thread 1
+        thread_1_history = [c for c in app.get_state_history(thread_1)]
+        # there are 2: one for each successful ainvoke()
+        assert len(thread_1_history) == 2
+        # sorted descending
+        assert (
+            thread_1_history[0].config["configurable"]["thread_ts"]
+            > thread_1_history[1].config["configurable"]["thread_ts"]
+        )
+        # the second checkpoint
+        assert thread_1_history[0].values["total"] == 7
+        # the first checkpoint
+        assert thread_1_history[1].values["total"] == 2
+        # can get each checkpoint using aget with config
+        assert (
+            memory.get(thread_1_history[0].config)["ts"]
+            == thread_1_history[0].config["configurable"]["thread_ts"]
+        )
+        assert (
+            memory.get(thread_1_history[1].config)["ts"]
+            == thread_1_history[1].config["configurable"]["thread_ts"]
+        )
+
+        thread_1_next_config = app.update_state(
+            thread_1_history[1].config, {"total": 10}
+        )
+        # update creates a new checkpoint
+        assert (
+            thread_1_next_config["configurable"]["thread_ts"]
+            > thread_1_history[0].config["configurable"]["thread_ts"]
+        )
+        # 1 more checkpoint in history
+        assert len(list(app.get_state_history(thread_1))) == 3
+        # the latest checkpoint is the updated one
+        assert app.get_state(thread_1) == app.get_state(thread_1_next_config)


 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
@@ -942,6 +982,13 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+    assert (
+        app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][
+            "thread_ts"
+        ]
+        is not None
    )

    app_w_interrupt.update_state(
@@ -971,6 +1018,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    assert [c for c in app_w_interrupt.stream(None, config)] == [
@@ -1088,6 +1136,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    app_w_interrupt.update_state(
@@ -1117,6 +1166,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    assert [c for c in app_w_interrupt.stream(None, config)] == [
@@ -1234,6 +1284,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    assert [c for c in app_w_interrupt.stream(None, config)] == [
@@ -1571,6 +1622,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    app_w_interrupt.update_state(
@@ -1595,6 +1647,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    assert [c for c in app_w_interrupt.stream(None, config)] == [
@@ -1686,6 +1739,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    app_w_interrupt.update_state(
@@ -1710,6 +1764,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    assert [c for c in app_w_interrupt.stream(None, config)] == [
@@ -2478,6 +2533,7 @@ def test_message_graph(snapshot: SnapshotAssertion) -> None:
            ),
        ],
        next=("agent:edges",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
    )

    # TODO use update_state once we have message ids
@@ -502,7 +502,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
    assert checkpoint["channel_values"].get("total") == 5


-async def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
+async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])

    def raise_if_above_10(input: int) -> int:
@@ -517,42 +517,86 @@ async def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
        | raise_if_above_10
    )

-    memory = AsyncSqliteSaver.from_conn_string(":memory:")
+    async with AsyncSqliteSaver.from_conn_string(":memory:") as memory:
+        app = Pregel(
+            nodes={"one": one},
+            channels={"total": BinaryOperatorAggregate(int, operator.add)},
+            checkpointer=memory,
+            debug=True,
+        )

-    app = Pregel(
-        nodes={"one": one},
-        channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpointer=memory,
-        debug=True,
-    )
+        thread_1 = {"configurable": {"thread_id": "1"}}
+        # total starts out as 0, so output is 0+2=2
+        assert await app.ainvoke(2, thread_1) == 2
+        state = await app.aget_state(thread_1)
+        assert state is not None
+        assert state.values.get("total") == 2
+        assert (
+            state.config["configurable"]["thread_ts"]
+            == (await memory.aget(thread_1))["ts"]
+        )
+        # total is now 2, so output is 2+3=5
+        assert await app.ainvoke(3, thread_1) == 5
+        state = await app.aget_state(thread_1)
+        assert state is not None
+        assert state.values.get("total") == 7
+        assert (
+            state.config["configurable"]["thread_ts"]
+            == (await memory.aget(thread_1))["ts"]
+        )
+        # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+        with pytest.raises(ValueError):
+            await app.ainvoke(4, thread_1)
+        # checkpoint is not updated
+        state = await app.aget_state(thread_1)
+        assert state is not None
+        assert state.values.get("total") == 7

-    # total starts out as 0, so output is 0+2=2
-    assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
-    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 2
-    # total is now 2, so output is 2+3=5
-    assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
-    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 7
-    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
-    with pytest.raises(ValueError):
-        await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
-    # checkpoint is not updated
-    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 7
-    # on a new thread, total starts out as 0, so output is 0+5=5
-    assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
-    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 7
-    checkpoint = await memory.aget({"configurable": {"thread_id": "2"}})
-    assert checkpoint is not None
-    assert checkpoint["channel_values"].get("total") == 5
+        thread_2 = {"configurable": {"thread_id": "2"}}
+        # on a new thread, total starts out as 0, so output is 0+5=5
+        assert await app.ainvoke(5, thread_2) == 5
+        state = await app.aget_state({"configurable": {"thread_id": "1"}})
+        assert state is not None
+        assert state.values.get("total") == 7
+        state = await app.aget_state(thread_2)
+        assert state is not None
+        assert state.values.get("total") == 5

-    await memory.conn.close()
+        # list all checkpoints for thread 1
+        thread_1_history = [c async for c in app.aget_state_history(thread_1)]
+        # there are 2: one for each successful ainvoke()
+        assert len(thread_1_history) == 2
+        # sorted descending
+        assert (
+            thread_1_history[0].config["configurable"]["thread_ts"]
+            > thread_1_history[1].config["configurable"]["thread_ts"]
+        )
+        # the second checkpoint
+        assert thread_1_history[0].values["total"] == 7
+        # the first checkpoint
+        assert thread_1_history[1].values["total"] == 2
+        # can get each checkpoint using aget with config
+        assert (await memory.aget(thread_1_history[0].config))[
+            "ts"
+        ] == thread_1_history[0].config["configurable"]["thread_ts"]
+        assert (await memory.aget(thread_1_history[1].config))[
+            "ts"
+        ] == thread_1_history[1].config["configurable"]["thread_ts"]
+
+        thread_1_next_config = await app.aupdate_state(
+            thread_1_history[1].config, {"total": 10}
+        )
+        # update creates a new checkpoint
+        assert (
+            thread_1_next_config["configurable"]["thread_ts"]
+            > thread_1_history[0].config["configurable"]["thread_ts"]
+        )
+        # 1 more checkpoint in history
+        assert len([h async for h in app.aget_state_history(thread_1)]) == 3
+        # the latest checkpoint is the updated one
+        assert await app.aget_state(thread_1) == await app.aget_state(
+            thread_1_next_config
+        )


 async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
@@ -988,6 +1032,7 @@ async def test_conditional_graph() -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    await app_w_interrupt.aupdate_state(
@@ -1017,6 +1062,7 @@ async def test_conditional_graph() -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1137,6 +1183,7 @@ async def test_conditional_graph() -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    await app_w_interrupt.aupdate_state(
@@ -1166,6 +1213,7 @@ async def test_conditional_graph() -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1286,6 +1334,7 @@ async def test_conditional_graph() -> None:
            "tools": None,
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1621,6 +1670,7 @@ async def test_conditional_graph_state() -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    await app_w_interrupt.aupdate_state(
@@ -1645,6 +1695,7 @@ async def test_conditional_graph_state() -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -1737,6 +1788,7 @@ async def test_conditional_graph_state() -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    await app_w_interrupt.aupdate_state(
@@ -1761,6 +1813,7 @@ async def test_conditional_graph_state() -> None:
            "intermediate_steps": [],
        },
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    assert [c async for c in app_w_interrupt.astream(None, config)] == [
@@ -2512,6 +2565,7 @@ async def test_message_graph() -> None:
            ),
        ],
        next=("agent:edges",),
+        config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
    )

    # TODO use update_state once we have message ids