Add tests for sweeper (sync)

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

@@ -2,6 +2,7 @@ import asyncio
 import logging
 from collections.abc import AsyncIterator, Iterable, Sequence
 from contextlib import asynccontextmanager
+from types import TracebackType
 from typing import Any, Callable, Optional, Union, cast
 
 import orjson
@@ -25,6 +26,7 @@ from langgraph.store.postgres.base import (
     PoolConfig,
     PostgresIndexConfig,
     Row,
+    TTLConfig,
     _decode_ns_bytes,
     _ensure_index_config,
     _group_ops,
@@ -106,6 +108,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         Semantic search is disabled by default. You can enable it by providing an `index` configuration
         when creating the store. Without this configuration, all `index` arguments passed to
         `put` or `aput` will have no effect.
+
+    Note:
+        If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
+        the background task that removes expired items. Call `stop_ttl_sweeper()` to properly
+        clean up resources when you're done with the store.
     """
 
     __slots__ = (
@@ -115,7 +122,9 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         "supports_pipeline",
         "index_config",
         "embeddings",
-        "supports_ttl",
+        "ttl_config",
+        "_ttl_sweeper_task",
+        "_ttl_stop_event",
     )
     supports_ttl: bool = True
 
@@ -128,6 +137,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
         index: Optional[PostgresIndexConfig] = None,
+        ttl: Optional[TTLConfig] = None,
     ) -> None:
         if isinstance(conn, AsyncConnectionPool) and pipe is not None:
             raise ValueError(
@@ -143,10 +153,13 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         self.index_config = index
         if self.index_config:
             self.embeddings, self.index_config = _ensure_index_config(self.index_config)
-
         else:
             self.embeddings = None
 
+        self.ttl_config = ttl
+        self._ttl_sweeper_task: Optional[asyncio.Task[None]] = None
+        self._ttl_stop_event = asyncio.Event()
+
     async def abatch(self, ops: Iterable[Op]) -> list[Result]:
         grouped_ops, num_ops = _group_ops(ops)
         results: list[Result] = [None] * num_ops
@@ -169,6 +182,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         pipeline: bool = False,
         pool_config: Optional[PoolConfig] = None,
         index: Optional[PostgresIndexConfig] = None,
+        ttl: Optional[TTLConfig] = None,
     ) -> AsyncIterator["AsyncPostgresStore"]:
         """Create a new AsyncPostgresStore instance from a connection string.
 
@@ -200,16 +214,16 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
                     **cast(dict, pc),
                 ),
             ) as pool:
-                yield cls(conn=pool, index=index)
+                yield cls(conn=pool, index=index, ttl=ttl)
         else:
             async with await AsyncConnection.connect(
                 conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
             ) as conn:
                 if pipeline:
                     async with conn.pipeline() as pipe:
-                        yield cls(conn=conn, pipe=pipe, index=index)
+                        yield cls(conn=conn, pipe=pipe, index=index, ttl=ttl)
                 else:
-                    yield cls(conn=conn, index=index)
+                    yield cls(conn=conn, index=index, ttl=ttl)
 
     async def setup(self) -> None:
         """Set up the store database asynchronously.
@@ -276,30 +290,100 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
 
     async def start_ttl_sweeper(
         self, sweep_interval_minutes: Optional[int] = None
-    ) -> None:
-        """Periodically delete expired store items based on TTL."""
-        if not self.ttl_config:
-            return
-        sweep_interval_minutes_ = float(
-            cast(
-                float,
-                sweep_interval_minutes
-                or self.ttl_config.get("sweep_interval_minutes")
-                or 5,
-            )
-        )
-        logger.info(
-            f"Starting store TTL sweeper with interval {sweep_interval_minutes_} minutes",
-        )
+    ) -> asyncio.Task[None]:
+        """Periodically delete expired store items based on TTL.
 
-        while True:
-            await asyncio.sleep(sweep_interval_minutes_ * 60)
+        Returns:
+            Task that can be awaited or cancelled.
+        """
+        if not self.ttl_config:
+            return asyncio.create_task(asyncio.sleep(0))
+
+        if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done():
+            return self._ttl_sweeper_task
+
+        self._ttl_stop_event.clear()
+
+        interval = float(
+            sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
+        )
+        logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
+
+        async def _sweep_loop() -> None:
+            while not self._ttl_stop_event.is_set():
+                try:
+                    try:
+                        await asyncio.wait_for(
+                            self._ttl_stop_event.wait(),
+                            timeout=interval * 60,
+                        )
+                        break
+                    except asyncio.TimeoutError:
+                        pass
+
+                    expired_items = await self.sweep_ttl()
+                    if expired_items > 0:
+                        logger.info(f"Store swept {expired_items} expired items")
+                except asyncio.CancelledError:
+                    break
+                except Exception as exc:
+                    logger.exception("Store TTL sweep iteration failed", exc_info=exc)
+
+        task = asyncio.create_task(_sweep_loop())
+        task.set_name("ttl_sweeper")
+        self._ttl_sweeper_task = task
+        return task
+
+    async def stop_ttl_sweeper(self, timeout: Optional[float] = None) -> bool:
+        """Stop the TTL sweeper task if it's running.
+
+        Args:
+            timeout: Maximum time to wait for the task to stop, in seconds.
+                If None, wait indefinitely.
+
+        Returns:
+            bool: True if the task was successfully stopped or wasn't running,
+                False if the timeout was reached before the task stopped.
+        """
+        if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done():
+            return True
+
+        logger.info("Stopping TTL sweeper task")
+        self._ttl_stop_event.set()
+
+        if timeout is not None:
             try:
-                expired_items = await self.sweep_ttl()
-                if expired_items > 0:
-                    logger.info(f"Store swept {expired_items} expired items")
-            except Exception as exc:
-                logger.exception("Store TTL sweep iteration failed", exc_info=exc)
+                await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout)
+                success = True
+            except asyncio.TimeoutError:
+                success = False
+        else:
+            await self._ttl_sweeper_task
+            success = True
+
+        if success:
+            self._ttl_sweeper_task = None
+            logger.info("TTL sweeper task stopped")
+        else:
+            logger.warning("Timed out waiting for TTL sweeper task to stop")
+
+        return success
+
+    async def __aenter__(self) -> "AsyncPostgresStore":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional["TracebackType"],
+    ) -> None:
+        # Ensure the TTL sweeper task is stopped when exiting the context
+        if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None:
+            # Set the event to signal the task to stop
+            self._ttl_stop_event.set()
+            # We don't wait for the task to complete here to avoid blocking
+            # The task will clean up itself gracefully
 
     async def _execute_batch(
         self,

libs/checkpoint-postgres/langgraph/store/postgres/base.py

@@ -1,8 +1,8 @@
 import asyncio
+import concurrent.futures
 import json
 import logging
 import threading
-import time
 from collections import defaultdict
 from collections.abc import Iterable, Iterator, Sequence
 from contextlib import contextmanager
@@ -81,7 +81,8 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (p
 ALTER TABLE store
 ADD COLUMN expires_at TIMESTAMP WITH TIME ZONE,
 ADD COLUMN ttl_minutes INT;
-
+""",
+    """
 -- Add indexes for efficient TTL sweeping
 CREATE INDEX idx_store_expires_at ON store (expires_at)
 WHERE expires_at IS NOT NULL;
@@ -253,7 +254,7 @@ class BasePostgresStore(Generic[C]):
 
         results = []
         for namespace, items in namespace_groups.items():
-            _, keys = zip(*items, strict=True)
+            _, keys = zip(*items)
             this_refresh_ttls = refresh_ttls[namespace]
 
             query = """
@@ -292,7 +293,7 @@ class BasePostgresStore(Generic[C]):
         put_ops: Sequence[tuple[int, PutOp]],
     ) -> tuple[
         list[tuple[str, Sequence]],
-        tuple[str, Sequence[tuple[str, str, str, str]]] | None,
+        Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
     ]:
         dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
         for _, op in put_ops:
@@ -319,7 +320,9 @@ class BasePostgresStore(Generic[C]):
                 )
                 params = (_namespace_to_text(namespace), *keys)
                 queries.append((query, params))
-        embedding_request: tuple[str, Sequence[tuple[str, str, str, str]]] | None = None
+        embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = (
+            None
+        )
         if inserts:
             values = []
             insertion_params = []
@@ -400,7 +403,7 @@ class BasePostgresStore(Generic[C]):
         self,
         search_ops: Sequence[tuple[int, SearchOp]],
     ) -> tuple[
-        list[tuple[str, list[None | str | list[float]]]],  # queries, params
+        list[tuple[str, list[Union[None, str, list[float]]]]],  # queries, params
         list[tuple[int, str]],  # idx, query_text pairs to embed
     ]:
         """
@@ -412,7 +415,6 @@ class BasePostgresStore(Generic[C]):
 
         queries = []
         embedding_requests = []
-
         for idx, (_, op) in enumerate(search_ops):
             filter_params = []
             filter_clauses = []
@@ -430,7 +432,7 @@ class BasePostgresStore(Generic[C]):
                         filter_params.extend([key, orjson.dumps(value).decode("utf-8")])
 
             ns_condition = "TRUE"
-            ns_param: Sequence[str] | None = None
+            ns_param: Optional[Sequence[Union[str]]] = None
             if op.namespace_prefix:
                 ns_condition = "store.prefix LIKE %s"
                 ns_param = (f"{_namespace_to_text(op.namespace_prefix)}%",)
@@ -512,7 +514,7 @@ class BasePostgresStore(Generic[C]):
 
             else:
                 base_query = f"""
-                        SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at, 0 AS score
+                        SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at, NULL AS score
                         FROM store
                         WHERE {ns_condition} {extra_filters}
                         ORDER BY store.updated_at DESC
@@ -625,7 +627,7 @@ class BasePostgresStore(Generic[C]):
 
 
 class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
-    """Postgres-backed store with optional vector search using pgvector.
+    """Postgres-backed store with oktional vector search using pgvector.
 
     !!! example "Examples"
         Basic setup and usage:
@@ -694,6 +696,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         Make sure to call `setup()` before first use to create necessary tables and indexes.
         The pgvector extension must be available to use vector search.
 
+    Note:
+        If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
+        the background thread that removes expired items. Call `stop_ttl_sweeper()` to properly
+        clean up resources when you're done with the store.
+
     """
 
     __slots__ = (
@@ -703,7 +710,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         "supports_pipeline",
         "index_config",
         "embeddings",
-        "supports_ttl",
+        "_ttl_sweeper_thread",
+        "_ttl_stop_event",
     )
     supports_ttl: bool = True
 
@@ -730,6 +738,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         else:
             self.embeddings = None
         self.ttl_config = ttl
+        self._ttl_sweeper_thread: Optional[threading.Thread] = None
+        self._ttl_stop_event = threading.Event()
 
     @classmethod
     @contextmanager
@@ -740,6 +750,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         pipeline: bool = False,
         pool_config: Optional[PoolConfig] = None,
         index: Optional[PostgresIndexConfig] = None,
+        ttl: Optional[TTLConfig] = None,
     ) -> Iterator["PostgresStore"]:
         """Create a new PostgresStore instance from a connection string.
 
@@ -771,16 +782,16 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
                     **cast(dict, pc),
                 ),
             ) as pool:
-                yield cls(conn=pool, index=index)
+                yield cls(conn=pool, index=index, ttl=ttl)
         else:
             with Connection.connect(
                 conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
             ) as conn:
                 if pipeline:
                     with conn.pipeline() as pipe:
-                        yield cls(conn, pipe=pipe, index=index)
+                        yield cls(conn, pipe=pipe, index=index, ttl=ttl)
                 else:
-                    yield cls(conn, index=index)
+                    yield cls(conn, index=index, ttl=ttl)
 
     def sweep_ttl(self) -> int:
         """Delete expired store items based on TTL.
@@ -798,30 +809,96 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
             deleted_count = cur.rowcount
             return deleted_count
 
-    def start_ttl_sweeper(self, sweep_interval_minutes: Optional[int] = None) -> None:
-        """Periodically delete expired store items based on TTL."""
-        if not self.ttl_config:
-            return
-        sweep_interval_minutes_ = float(
-            cast(
-                float,
-                sweep_interval_minutes
-                or self.ttl_config.get("sweep_interval_minutes")
-                or 5,
-            )
-        )
-        logger.info(
-            f"Starting store TTL sweeper with interval {sweep_interval_minutes_} minutes",
-        )
+    def start_ttl_sweeper(
+        self, sweep_interval_minutes: Optional[int] = None
+    ) -> concurrent.futures.Future[None]:
+        """Periodically delete expired store items based on TTL.
 
-        while True:
-            time.sleep(sweep_interval_minutes_ * 60)
+        Returns:
+            Future that can be waited on or cancelled.
+        """
+        if not self.ttl_config:
+            future: concurrent.futures.Future[None] = concurrent.futures.Future()
+            future.set_result(None)
+            return future
+
+        if self._ttl_sweeper_thread and self._ttl_sweeper_thread.is_alive():
+            logger.info("TTL sweeper thread is already running")
+            # Return a future that can be used to cancel the existing thread
+            future = concurrent.futures.Future()
+            future.add_done_callback(
+                lambda f: self._ttl_stop_event.set() if f.cancelled() else None
+            )
+            return future
+
+        self._ttl_stop_event.clear()
+
+        interval = float(
+            sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
+        )
+        logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
+
+        future = concurrent.futures.Future()
+
+        def _sweep_loop() -> None:
             try:
-                expired_items = self.sweep_ttl()
-                if expired_items > 0:
-                    logger.info(f"Store swept {expired_items} expired items")
+                while not self._ttl_stop_event.is_set():
+                    if self._ttl_stop_event.wait(interval * 60):
+                        break
+
+                    try:
+                        expired_items = self.sweep_ttl()
+                        if expired_items > 0:
+                            logger.info(f"Store swept {expired_items} expired items")
+                    except Exception as exc:
+                        logger.exception(
+                            "Store TTL sweep iteration failed", exc_info=exc
+                        )
+                future.set_result(None)
             except Exception as exc:
-                logger.exception("Store TTL sweep iteration failed", exc_info=exc)
+                future.set_exception(exc)
+
+        thread = threading.Thread(target=_sweep_loop, daemon=True, name="ttl-sweeper")
+        self._ttl_sweeper_thread = thread
+        thread.start()
+
+        future.add_done_callback(
+            lambda f: self._ttl_stop_event.set() if f.cancelled() else None
+        )
+        return future
+
+    def stop_ttl_sweeper(self, timeout: Optional[float] = None) -> bool:
+        """Stop the TTL sweeper thread if it's running.
+
+        Args:
+            timeout: Maximum time to wait for the thread to stop, in seconds.
+                If None, wait indefinitely.
+
+        Returns:
+            bool: True if the thread was successfully stopped or wasn't running,
+                False if the timeout was reached before the thread stopped.
+        """
+        if not self._ttl_sweeper_thread or not self._ttl_sweeper_thread.is_alive():
+            return True
+
+        logger.info("Stopping TTL sweeper thread")
+        self._ttl_stop_event.set()
+
+        self._ttl_sweeper_thread.join(timeout)
+        success = not self._ttl_sweeper_thread.is_alive()
+
+        if success:
+            self._ttl_sweeper_thread = None
+            logger.info("TTL sweeper thread stopped")
+        else:
+            logger.warning("Timed out waiting for TTL sweeper thread to stop")
+
+        return success
+
+    def __del__(self) -> None:
+        """Ensure the TTL sweeper thread is stopped when the object is garbage collected."""
+        if hasattr(self, "_ttl_stop_event") and hasattr(self, "_ttl_sweeper_thread"):
+            self.stop_ttl_sweeper(timeout=0.1)
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -1020,8 +1097,14 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
         with self._cursor() as cur:
             version = _get_version(cur, table="store_migrations")
             for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
-                cur.execute(sql)
-                cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+                try:
+                    cur.execute(sql)
+                    cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+                except Exception as e:
+                    logger.error(
+                        f"Failed to apply migration {v}.\nSql={sql}\nError={e}"
+                    )
+                    raise
 
             if self.index_config:
                 version = _get_version(cur, table="vector_migrations")

libs/checkpoint-postgres/tests/test_async_store.py

@@ -26,6 +26,9 @@ from tests.conftest import (
     CharacterEmbeddings,
 )
 
+TTL_SECONDS = 6
+TTL_MINUTES = TTL_SECONDS / 60
+
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
 async def store(request) -> AsyncIterator[AsyncPostgresStore]:
@@ -42,28 +45,52 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
 
     conn_string = f"{uri_base}/{database}{query_params}"
     admin_conn_string = DEFAULT_URI
-
+    ttl_config = {
+        "default_ttl": TTL_MINUTES,
+        "refresh_on_read": True,
+        "sweep_interval_minutes": TTL_MINUTES / 2,
+    }
     async with await AsyncConnection.connect(
         admin_conn_string, autocommit=True
     ) as conn:
         await conn.execute(f"CREATE DATABASE {database}")
     try:
-        async with AsyncPostgresStore.from_conn_string(conn_string) as store:
+        async with AsyncPostgresStore.from_conn_string(
+            conn_string, ttl=ttl_config
+        ) as store:
+            store.MIGRATIONS = [
+                (
+                    mig.replace(
+                        "ADD COLUMN ttl_minutes INT;", "ADD COLUMN ttl_minutes FLOAT;"
+                    )
+                    if isinstance(mig, str)
+                    else mig
+                )
+                for mig in store.MIGRATIONS
+            ]
             await store.setup()
 
         if request.param == "pipe":
             async with AsyncPostgresStore.from_conn_string(
-                conn_string, pipeline=True
+                conn_string, pipeline=True, ttl=ttl_config
             ) as store:
+                await store.start_ttl_sweeper()
                 yield store
+                await store.stop_ttl_sweeper()
         elif request.param == "pool":
             async with AsyncPostgresStore.from_conn_string(
-                conn_string, pool_config={"min_size": 1, "max_size": 10}
+                conn_string, pool_config={"min_size": 1, "max_size": 10}, ttl=ttl_config
             ) as store:
+                await store.start_ttl_sweeper()
                 yield store
+                await store.stop_ttl_sweeper()
         else:  # default
-            async with AsyncPostgresStore.from_conn_string(conn_string) as store:
+            async with AsyncPostgresStore.from_conn_string(
+                conn_string, ttl=ttl_config
+            ) as store:
+                await store.start_ttl_sweeper()
                 yield store
+                await store.stop_ttl_sweeper()
     finally:
         async with await AsyncConnection.connect(
             admin_conn_string, autocommit=True
@@ -635,3 +662,28 @@ async def test_search_sorting(
         assert len(set(r.key for r in results)) == 10
         assert results[0].key == "M"
         assert results[0].score > results[1].score
+
+
+async def test_store_ttl(store):
+    # Assumes a TTL of 1 minute = 60 seconds
+    ns = ("foo",)
+    await store.start_ttl_sweeper()
+    await store.aput(
+        ns,
+        key="item1",
+        value={"foo": "bar"},
+        ttl=TTL_MINUTES,  # type: ignore
+    )
+    await asyncio.sleep(TTL_SECONDS - 2)
+    res = await store.aget(ns, key="item1", refresh_ttl=True)
+    assert res is not None
+    await asyncio.sleep(TTL_SECONDS - 2)
+    results = await store.asearch(ns, query="foo", refresh_ttl=True)
+    assert len(results) == 1
+    await asyncio.sleep(TTL_SECONDS - 2)
+    res = await store.aget(ns, key="item1", refresh_ttl=False)
+    assert res is not None
+    await asyncio.sleep(TTL_SECONDS - 1)
+    # Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
+    results = await store.asearch(ns, query="bar", refresh_ttl=False)
+    assert len(results) == 0

libs/checkpoint-postgres/tests/test_store.py

@@ -1,6 +1,7 @@
 # type: ignore
 
 import re
+import time
 from contextlib import contextmanager
 from typing import Any, Optional
 from uuid import uuid4
@@ -24,6 +25,9 @@ from tests.conftest import (
     CharacterEmbeddings,
 )
 
+TTL_SECONDS = 6
+TTL_MINUTES = TTL_SECONDS / 60
+
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
 def store(request) -> PostgresStore:
@@ -32,29 +36,58 @@ def store(request) -> PostgresStore:
     uri_base = "/".join(uri_parts[:-1])
     query_params = ""
     if "?" in uri_parts[-1]:
-        db_name, query_params = uri_parts[-1].split("?", 1)
+        _, query_params = uri_parts[-1].split("?", 1)
         query_params = "?" + query_params
 
     conn_string = f"{uri_base}/{database}{query_params}"
     admin_conn_string = DEFAULT_URI
-
+    ttl_config = {
+        "default_ttl": TTL_MINUTES,
+        "refresh_on_read": True,
+        "sweep_interval_minutes": TTL_MINUTES / 2,
+    }
     with Connection.connect(admin_conn_string, autocommit=True) as conn:
         conn.execute(f"CREATE DATABASE {database}")
     try:
-        with PostgresStore.from_conn_string(conn_string) as store:
+        with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
+            store.MIGRATIONS = [
+                (
+                    mig.replace(
+                        "ADD COLUMN ttl_minutes INT;", "ADD COLUMN ttl_minutes FLOAT;"
+                    )
+                    if isinstance(mig, str)
+                    else mig
+                )
+                for mig in store.MIGRATIONS
+            ]
             store.setup()
 
         if request.param == "pipe":
-            with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
+            with PostgresStore.from_conn_string(
+                conn_string,
+                pipeline=True,
+                ttl=ttl_config,
+            ) as store:
+                store.start_ttl_sweeper()
                 yield store
+
+                store.stop_ttl_sweeper()
         elif request.param == "pool":
             with PostgresStore.from_conn_string(
-                conn_string, pool_config={"min_size": 1, "max_size": 10}
+                conn_string,
+                pool_config={"min_size": 1, "max_size": 10},
+                ttl=ttl_config,
             ) as store:
+                store.start_ttl_sweeper()
                 yield store
+
+                store.stop_ttl_sweeper()
         else:  # default
-            with PostgresStore.from_conn_string(conn_string) as store:
+            with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
+                store.start_ttl_sweeper()
                 yield store
+
+                store.stop_ttl_sweeper()
     finally:
         with Connection.connect(admin_conn_string, autocommit=True) as conn:
             conn.execute(f"DROP DATABASE {database}")
@@ -220,134 +253,127 @@ def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
     assert all(ns[-1] == "public" for ns in results[2])
 
 
-class TestPostgresStore:
-    @pytest.fixture(autouse=True)
-    def setup(self) -> None:
-        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
-            store.setup()
+def test_basic_store_ops(store) -> None:
+    namespace = ("test", "documents")
+    item_id = "doc1"
+    item_value = {"title": "Test Document", "content": "Hello, World!"}
 
-    def test_basic_store_ops(self) -> None:
-        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
-            namespace = ("test", "documents")
-            item_id = "doc1"
-            item_value = {"title": "Test Document", "content": "Hello, World!"}
+    store.put(namespace, item_id, item_value)
+    item = store.get(namespace, item_id)
 
-            store.put(namespace, item_id, item_value)
-            item = store.get(namespace, item_id)
+    assert item
+    assert item.namespace == namespace
+    assert item.key == item_id
+    assert item.value == item_value
 
-            assert item
-            assert item.namespace == namespace
-            assert item.key == item_id
-            assert item.value == item_value
+    # Test update
+    updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
+    store.put(namespace, item_id, updated_value)
+    updated_item = store.get(namespace, item_id)
 
-            # Test update
-            updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
-            store.put(namespace, item_id, updated_value)
-            updated_item = store.get(namespace, item_id)
+    assert updated_item.value == updated_value
+    assert updated_item.updated_at > item.updated_at
 
-            assert updated_item.value == updated_value
-            assert updated_item.updated_at > item.updated_at
+    # Test get from non-existent namespace
+    different_namespace = ("test", "other_documents")
+    item_in_different_namespace = store.get(different_namespace, item_id)
+    assert item_in_different_namespace is None
 
-            # Test get from non-existent namespace
-            different_namespace = ("test", "other_documents")
-            item_in_different_namespace = store.get(different_namespace, item_id)
-            assert item_in_different_namespace is None
+    # Test delete
+    store.delete(namespace, item_id)
+    deleted_item = store.get(namespace, item_id)
+    assert deleted_item is None
 
-            # Test delete
-            store.delete(namespace, item_id)
-            deleted_item = store.get(namespace, item_id)
-            assert deleted_item is None
 
-    def test_list_namespaces(self) -> None:
-        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
-            # Create test data with various namespaces
-            test_namespaces = [
-                ("test", "documents", "public"),
-                ("test", "documents", "private"),
-                ("test", "images", "public"),
-                ("test", "images", "private"),
-                ("prod", "documents", "public"),
-                ("prod", "documents", "private"),
-            ]
+def test_list_namespaces(store) -> None:
+    # Create test data with various namespaces
+    test_namespaces = [
+        ("test", "documents", "public"),
+        ("test", "documents", "private"),
+        ("test", "images", "public"),
+        ("test", "images", "private"),
+        ("prod", "documents", "public"),
+        ("prod", "documents", "private"),
+    ]
 
-            # Insert test data
-            for namespace in test_namespaces:
-                store.put(namespace, "dummy", {"content": "dummy"})
+    # Insert test data
+    for namespace in test_namespaces:
+        store.put(namespace, "dummy", {"content": "dummy"})
 
-            # Test listing with various filters
-            all_namespaces = store.list_namespaces()
-            assert len(all_namespaces) == len(test_namespaces)
+    # Test listing with various filters
+    all_namespaces = store.list_namespaces()
+    assert len(all_namespaces) == len(test_namespaces)
 
-            # Test prefix filtering
-            test_prefix_namespaces = store.list_namespaces(prefix=["test"])
-            assert len(test_prefix_namespaces) == 4
-            assert all(ns[0] == "test" for ns in test_prefix_namespaces)
+    # Test prefix filtering
+    test_prefix_namespaces = store.list_namespaces(prefix=["test"])
+    assert len(test_prefix_namespaces) == 4
+    assert all(ns[0] == "test" for ns in test_prefix_namespaces)
 
-            # Test suffix filtering
-            public_namespaces = store.list_namespaces(suffix=["public"])
-            assert len(public_namespaces) == 3
-            assert all(ns[-1] == "public" for ns in public_namespaces)
+    # Test suffix filtering
+    public_namespaces = store.list_namespaces(suffix=["public"])
+    assert len(public_namespaces) == 3
+    assert all(ns[-1] == "public" for ns in public_namespaces)
 
-            # Test max depth
-            depth_2_namespaces = store.list_namespaces(max_depth=2)
-            assert all(len(ns) <= 2 for ns in depth_2_namespaces)
+    # Test max depth
+    depth_2_namespaces = store.list_namespaces(max_depth=2)
+    assert all(len(ns) <= 2 for ns in depth_2_namespaces)
 
-            # Test pagination
-            paginated_namespaces = store.list_namespaces(limit=3)
-            assert len(paginated_namespaces) == 3
+    # Test pagination
+    paginated_namespaces = store.list_namespaces(limit=3)
+    assert len(paginated_namespaces) == 3
 
-            # Cleanup
-            for namespace in test_namespaces:
-                store.delete(namespace, "dummy")
+    # Cleanup
+    for namespace in test_namespaces:
+        store.delete(namespace, "dummy")
 
-    def test_search(self) -> None:
-        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
-            # Create test data
-            test_data = [
-                (
-                    ("test", "docs"),
-                    "doc1",
-                    {"title": "First Doc", "author": "Alice", "tags": ["important"]},
-                ),
-                (
-                    ("test", "docs"),
-                    "doc2",
-                    {"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
-                ),
-                (
-                    ("test", "images"),
-                    "img1",
-                    {"title": "Image 1", "author": "Alice", "tags": ["final"]},
-                ),
-            ]
 
-            for namespace, key, value in test_data:
-                store.put(namespace, key, value)
+def test_search(store) -> None:
+    # Create test data
+    test_data = [
+        (
+            ("test", "docs"),
+            "doc1",
+            {"title": "First Doc", "author": "Alice", "tags": ["important"]},
+        ),
+        (
+            ("test", "docs"),
+            "doc2",
+            {"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
+        ),
+        (
+            ("test", "images"),
+            "img1",
+            {"title": "Image 1", "author": "Alice", "tags": ["final"]},
+        ),
+    ]
 
-            # Test basic search
-            all_items = store.search(["test"])
-            assert len(all_items) == 3
+    for namespace, key, value in test_data:
+        store.put(namespace, key, value)
 
-            # Test namespace filtering
-            docs_items = store.search(["test", "docs"])
-            assert len(docs_items) == 2
-            assert all(item.namespace == ("test", "docs") for item in docs_items)
+    # Test basic search
+    all_items = store.search(["test"])
+    assert len(all_items) == 3
 
-            # Test value filtering
-            alice_items = store.search(["test"], filter={"author": "Alice"})
-            assert len(alice_items) == 2
-            assert all(item.value["author"] == "Alice" for item in alice_items)
+    # Test namespace filtering
+    docs_items = store.search(["test", "docs"])
+    assert len(docs_items) == 2
+    assert all(item.namespace == ("test", "docs") for item in docs_items)
 
-            # Test pagination
-            paginated_items = store.search(["test"], limit=2)
-            assert len(paginated_items) == 2
+    # Test value filtering
+    alice_items = store.search(["test"], filter={"author": "Alice"})
+    assert len(alice_items) == 2
+    assert all(item.value["author"] == "Alice" for item in alice_items)
 
-            offset_items = store.search(["test"], offset=2)
-            assert len(offset_items) == 1
+    # Test pagination
+    paginated_items = store.search(["test"], limit=2)
+    assert len(paginated_items) == 2
 
-            # Cleanup
-            for namespace, key, _ in test_data:
-                store.delete(namespace, key)
+    offset_items = store.search(["test"], offset=2)
+    assert len(offset_items) == 1
+
+    # Cleanup
+    for namespace, key, _ in test_data:
+        store.delete(namespace, key)
 
 
 @contextmanager
@@ -356,6 +382,7 @@ def _create_vector_store(
     distance_type: str,
     fake_embeddings: Embeddings,
     text_fields: Optional[list[str]] = None,
+    enable_ttl: bool = True,
 ) -> PostgresStore:
     """Create a store with vector search enabled."""
     database = f"test_{uuid4().hex[:16]}"
@@ -385,6 +412,7 @@ def _create_vector_store(
         with PostgresStore.from_conn_string(
             conn_string,
             index=index_config,
+            ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
         ) as store:
             store.setup()
             yield store
@@ -393,15 +421,19 @@ def _create_vector_store(
             conn.execute(f"DROP DATABASE {database}")
 
 
+_vector_params = [
+    (vector_type, distance_type, True)
+    for vector_type in VECTOR_TYPES
+    for distance_type in (
+        ["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
+    )
+]
+_vector_params += [(*_vector_params[-1][:2], False)]
+
+
 @pytest.fixture(
     scope="function",
-    params=[
-        (vector_type, distance_type)
-        for vector_type in VECTOR_TYPES
-        for distance_type in (
-            ["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
-        )
-    ],
+    params=_vector_params,
     ids=lambda p: f"{p[0]}_{p[1]}",
 )
 def vector_store(
@@ -409,8 +441,10 @@ def vector_store(
     fake_embeddings: Embeddings,
 ) -> PostgresStore:
     """Create a store with vector search enabled."""
-    vector_type, distance_type = request.param
-    with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
+    vector_type, distance_type, enable_ttl = request.param
+    with _create_vector_store(
+        vector_type, distance_type, fake_embeddings, enable_ttl=enable_ttl
+    ) as store:
         yield store
 
 
@@ -474,7 +508,10 @@ def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
     assert not any(r.key == "doc4" for r in results_new)
 
 
-def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
+@pytest.mark.parametrize("refresh_ttl", [True, False])
+def test_vector_search_with_filters(
+    vector_store: PostgresStore, refresh_ttl: bool
+) -> None:
     """Test combining vector search with filters."""
     # Insert test documents
     docs = [
@@ -487,16 +524,23 @@ def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
     for key, value in docs:
         vector_store.put(("test",), key, value)
 
-    results = vector_store.search(("test",), query="apple", filter={"color": "red"})
+    results = vector_store.search(
+        ("test",), query="apple", filter={"color": "red"}, refresh_ttl=refresh_ttl
+    )
     assert len(results) == 2
     assert results[0].key == "doc1"
 
-    results = vector_store.search(("test",), query="car", filter={"color": "red"})
+    results = vector_store.search(
+        ("test",), query="car", filter={"color": "red"}, refresh_ttl=refresh_ttl
+    )
     assert len(results) == 2
     assert results[0].key == "doc2"
 
     results = vector_store.search(
-        ("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
+        ("test",),
+        query="bbbbluuu",
+        filter={"score": {"$gt": 3.2}},
+        refresh_ttl=refresh_ttl,
     )
     assert len(results) == 3
     assert results[0].key == "doc4"
@@ -688,7 +732,7 @@ def test_embed_with_path_operation_config(
         store.put(("test",), "doc5", doc5, index=False)
         results = store.search(("test",))
         assert len(results) == 3
-        assert all(r.score is None for r in results)
+        assert all(r.score is None for r in results), f"{results}"
         assert any(r.key == "doc5" for r in results)
 
         results = store.search(("test",), query="hhh")
@@ -790,3 +834,27 @@ def test_nonnull_migrations() -> None:
     for migration in PostgresStore.MIGRATIONS:
         statement = _leading_comment_remover.sub("", migration).split()[0]
         assert statement.strip()
+
+
+def test_store_ttl(store):
+    # Assumes a TTL of 1 minute = 60 seconds
+    ns = ("foo",)
+    store.put(
+        ns,
+        key="item1",
+        value={"foo": "bar"},
+        ttl=TTL_MINUTES,  # type: ignore
+    )
+    time.sleep(TTL_SECONDS - 2)
+    res = store.get(ns, key="item1", refresh_ttl=True)
+    assert res is not None
+    time.sleep(TTL_SECONDS - 2)
+    results = store.search(ns, query="foo", refresh_ttl=True)
+    assert len(results) == 1
+    time.sleep(TTL_SECONDS - 2)
+    res = store.get(ns, key="item1", refresh_ttl=False)
+    assert res is not None
+    time.sleep(TTL_SECONDS - 1)
+    # Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
+    res = store.search(ns, query="bar", refresh_ttl=False)
+    assert len(res) == 0

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

@@ -537,6 +537,12 @@ class TTLConfig(TypedDict, total=False):
     The expiration timer refreshes on both read and write operations.
     Defaults to None (no expiration).
     """
+    sweep_interval_minutes: Optional[int]
+    """Interval in minutes between TTL sweep operations.
+    
+    If provided, the store will periodically delete expired items based on TTL.
+    Defaults to None (no sweeping).
+    """
 
 
 class IndexConfig(TypedDict, total=False):

libs/cli/langgraph_cli/config.py

@@ -27,6 +27,12 @@ class TTLConfig(TypedDict, total=False):
     If provided, all new items will have this TTL unless explicitly overridden.
     If omitted, items will have no TTL by default.
     """
+    sweep_interval_minutes: Optional[int]
+    """Optional. Interval in minutes between TTL sweep iterations.
+    
+    If provided, the store will periodically delete expired items based on the TTL.
+    If omitted, no automatic sweeping will occur.
+    """
 
 
 class IndexConfig(TypedDict, total=False):