Add delete_thread method to Checkpointer class (#4328)

- Deletes all data associated with a thread_id
- Implemented in InMemory, Sqlite and Postgres checkpointers

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -357,6 +357,29 @@ class PostgresSaver(BasePostgresSaver):
                 ),
             )
 
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        with self._cursor(pipeline=True) as cur:
+            cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            cur.execute(
+                "DELETE FROM checkpoint_blobs WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            cur.execute(
+                "DELETE FROM checkpoint_writes WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
         """Create a database cursor as a context manager.

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -314,6 +314,29 @@ class AsyncPostgresSaver(BasePostgresSaver):
         async with self._cursor(pipeline=True) as cur:
             await cur.executemany(query, params)
 
+    async def adelete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        async with self._cursor(pipeline=True) as cur:
+            await cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM checkpoint_blobs WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM checkpoint_writes WHERE thread_id = %s",
+                (str(thread_id),),
+            )
+
     @asynccontextmanager
     async def _cursor(
         self, *, pipeline: bool = False
@@ -481,5 +504,30 @@ class AsyncPostgresSaver(BasePostgresSaver):
             self.aput_writes(config, writes, task_id, task_path), self.loop
         ).result()
 
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `await checkpointer.aget_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.adelete_thread(thread_id), self.loop
+        ).result()
+
 
 __all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py

@@ -464,6 +464,25 @@ class SqliteSaver(BaseCheckpointSaver[str]):
                 ],
             )
 
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        with self.cursor() as cur:
+            cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = ?",
+                (str(thread_id),),
+            )
+            cur.execute(
+                "DELETE FROM writes WHERE thread_id = ?",
+                (str(thread_id),),
+            )
+
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database asynchronously.
 

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py

@@ -244,6 +244,31 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
             self.aput_writes(config, writes, task_id, task_path), self.loop
         ).result()
 
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncSqliteSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface. "
+                    "For example, use `checkpointer.alist(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
+        return asyncio.run_coroutine_threadsafe(
+            self.adelete_thread(thread_id), self.loop
+        ).result()
+
     async def setup(self) -> None:
         """Set up the checkpoint database asynchronously.
 
@@ -535,6 +560,26 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
             )
             await self.conn.commit()
 
+    async def adelete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        async with self.lock, self.conn.cursor() as cur:
+            await cur.execute(
+                "DELETE FROM checkpoints WHERE thread_id = ?",
+                (str(thread_id),),
+            )
+            await cur.execute(
+                "DELETE FROM writes WHERE thread_id = ?",
+                (str(thread_id),),
+            )
+            await self.conn.commit()
+
     def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
         """Generate the next version ID for a channel.
 

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -321,6 +321,17 @@ class BaseCheckpointSaver(Generic[V]):
         """
         raise NotImplementedError
 
+    def delete_thread(
+        self,
+        thread_id: str,
+    ) -> None:
+        """Delete all checkpoints and writes associated with a specific thread ID.
+
+        Args:
+            thread_id (str): The thread ID whose checkpoints should be deleted.
+        """
+        raise NotImplementedError
+
     async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
         """Asynchronously fetch a checkpoint using the given configuration.
 
@@ -415,6 +426,17 @@ class BaseCheckpointSaver(Generic[V]):
         """
         raise NotImplementedError
 
+    async def adelete_thread(
+        self,
+        thread_id: str,
+    ) -> None:
+        """Delete all checkpoints and writes associated with a specific thread ID.
+
+        Args:
+            thread_id (str): The thread ID whose checkpoints should be deleted.
+        """
+        raise NotImplementedError
+
     def get_next_version(self, current: Optional[V], channel: ChannelProtocol) -> V:
         """Generate the next version ID for a channel.
 

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -69,7 +69,7 @@ class InMemorySaver(
         ],
     ]
     writes: defaultdict[
-        tuple[str, str, str],
+        tuple[str, str, str],  # thread ID, checkpoint NS, checkpoint ID
         dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]],
     ]
     blobs: dict[
@@ -451,6 +451,24 @@ class InMemorySaver(
                 task_path,
             )
 
+    def delete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        if thread_id in self.storage:
+            del self.storage[thread_id]
+        for k in list(self.writes.keys()):
+            if k[0] == thread_id:
+                del self.writes[k]
+        for k in list(self.blobs.keys()):
+            if k[0] == thread_id:
+                del self.blobs[k]
+
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Asynchronous version of get_tuple.
 
@@ -530,6 +548,17 @@ class InMemorySaver(
         """
         return self.put_writes(config, writes, task_id, task_path)
 
+    async def adelete_thread(self, thread_id: str) -> None:
+        """Delete all checkpoints and writes associated with a thread ID.
+
+        Args:
+            thread_id (str): The thread ID to delete.
+
+        Returns:
+            None
+        """
+        return self.delete_thread(thread_id)
+
     def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
         if current is None:
             current_v = 0
@@ -615,4 +644,4 @@ class PersistentDict(defaultdict):
                 except Exception:
                     logging.error(f"Failed to load file: {fileobj.name}")
                     raise
-            raise ValueError("File not in a supported f ormat")
+            raise ValueError("File not in a supported format")

libs/langgraph/tests/test_pregel.py

@@ -4601,7 +4601,7 @@ def test_checkpoint_metadata() -> None:
         assert chkpnt_tuple.metadata["test_config_4"] == "bar"
 
 
-@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
 def test_remove_message_via_state_update(
     request: pytest.FixtureRequest, checkpointer_name: str
 ) -> None:
@@ -4631,6 +4631,12 @@ def test_remove_message_via_state_update(
     assert len(updated_state.values) == 1
     assert updated_state.values[-1].content == "Hi"
 
+    app.checkpointer.delete_thread(config["configurable"]["thread_id"])
+
+    # Verify that the message was removed from the checkpointer
+    assert app.checkpointer.get_tuple(config) is None
+    assert [*app.get_state_history(config)] == []
+
 
 def test_remove_message_from_node():
     from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
@@ -6034,7 +6040,7 @@ def test_concurrent_execution_thread_safety():
         assert result["counter"] == 1
 
 
-@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
 def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: str):
     """Test recovery from checkpoints after failures."""
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
@@ -6087,6 +6093,11 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name:
     failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
     assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error
 
+    # Verify delete leaves it empty
+    graph.checkpointer.delete_thread(config["configurable"]["thread_id"])
+    assert graph.checkpointer.get_tuple(config) is None
+    assert [*graph.get_state_history(config)] == []
+
 
 def test_multiple_updates_root() -> None:
     def node_a(state):