From bfcfea554ed5c7f7be562cebf8825e911b493ab1 Mon Sep 17 00:00:00 2001 From: Baihao Wang <25713361+byhow@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:10:55 -0700 Subject: [PATCH] fix(checkpoint,checkpoint-postgres,checkpoint-sqlite): validate namespace labels before store dispatch (#9217) The SQLite and Postgres stores save a namespace as its labels joined with `.`, so a label containing `.` could resolve to a different namespace: `("foo.bar",)` and `("foo", "bar")` share the same row. This rejects those labels at every entry point. - `BaseStore` and `AsyncBatchedBaseStore`: `get`, `search`, `delete`, and `list_namespaces` (sync and async) reject labels that are empty, not strings, or contain `.`. The check runs before an op is queued, so one caller's bad label can't fail other callers' concurrent requests. - New `langgraph.store.base.validate_op_namespace(op)` applies the same label checks to a single op. `SqliteStore`/`AsyncSqliteStore` and `PostgresStore`/`AsyncPostgresStore` call it from `_group_ops` before any SQL runs, which covers direct `batch`/`abatch` calls. Other stores that serialize namespaces as delimited text can call it too. - Unchanged: hierarchical prefix search, empty-prefix `search(())`, `*` wildcards in `list_namespaces`, and the existing `put` rules. ## Release note `get`, `search`, `delete`, and `list_namespaces` now raise `InvalidNamespaceError` for labels that are empty, not strings, or contain `.`, where they previously returned `None` or `[]` (or, on the SQLite and Postgres stores, matched a different namespace). `put` already rejected these labels, so nothing written through `put` becomes unreachable. ## Release order Release `langgraph-checkpoint` first. `langgraph-checkpoint-sqlite` and `langgraph-checkpoint-postgres` import `validate_op_namespace`, so they need the `langgraph-checkpoint` release that ships it. `main` already has `libs/checkpoint` at 4.3.0 and requires `langgraph-checkpoint>=4.3.0` in both stores (from #8544), so this must merge before 4.3.0 is released; if 4.3.0 ships first, bump to the next version and raise both requirements to it. That also guarantees the pre-queue checks are installed: with an older `langgraph-checkpoint`, an invalid op would fail inside the shared `abatch` and reject the calls batched with it. --------- Co-authored-by: Baihao Wang Co-authored-by: open-swe[bot] --- .../langgraph/store/postgres/base.py | 2 + .../tests/test_async_store.py | 64 +++++++++ libs/checkpoint-postgres/tests/test_store.py | 49 +++++++ .../langgraph/store/sqlite/base.py | 2 + .../tests/test_async_store.py | 64 +++++++++ libs/checkpoint-sqlite/tests/test_store.py | 49 +++++++ .../langgraph/store/base/__init__.py | 86 +++++++++++- libs/checkpoint/langgraph/store/base/batch.py | 6 + libs/checkpoint/tests/test_store.py | 125 ++++++++++++++++++ 9 files changed, 443 insertions(+), 4 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index c407f26f1..b36671b92 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -36,6 +36,7 @@ from langgraph.store.base import ( ensure_embeddings, get_text_at_path, tokenize_path, + validate_op_namespace, ) from psycopg import Capabilities, Connection, Cursor, Pipeline from psycopg.rows import DictRow, dict_row @@ -1386,6 +1387,7 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list) tot = 0 for idx, op in enumerate(ops): + validate_op_namespace(op) grouped_ops[type(op)].append((idx, op)) tot += 1 return grouped_ops, tot diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index 080a78852..0f0fd5ce6 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -13,8 +13,10 @@ import pytest from langchain_core.embeddings import Embeddings from langgraph.store.base import ( GetOp, + InvalidNamespaceError, Item, ListNamespacesOp, + MatchCondition, PutOp, SearchOp, ) @@ -871,3 +873,65 @@ async def test_omit_expired_search_pagination(store: AsyncPostgresStore) -> None page2 = await store.asearch(ns, limit=2, offset=2) assert [i.key for i in page1] == ["a", "b"] assert [i.key for i in page2] == ["c"] + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +async def test_abatch_rejects_invalid_namespace_labels( + store: AsyncPostgresStore, namespace: tuple +) -> None: + await store.aput(("foo", "bar"), "key", {"original": True}) + + for op in ( + GetOp(namespace, "key"), + GetOp(namespace, "key", refresh_ttl=True), + PutOp(namespace, "key", {"changed": True}), + PutOp(namespace, "key", None), + SearchOp(namespace), + ListNamespacesOp((MatchCondition("prefix", namespace),)), + ListNamespacesOp((MatchCondition("suffix", namespace),)), + ): + with pytest.raises(InvalidNamespaceError): + await store.abatch([op]) + + item = await store.aget(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + + +async def test_invalid_namespace_only_fails_its_own_call( + store: AsyncPostgresStore, +) -> None: + """Concurrent calls share one `abatch`, which fails every op if it raises. + + Labels are checked before an op is queued, so one caller's bad label cannot + fail another caller's request. + """ + await store.aput(("foo", "bar"), "key", {"original": True}) + + valid, invalid = await asyncio.gather( + store.aget(("foo", "bar"), "key"), + store.aget(("foo.bar",), "key"), + return_exceptions=True, + ) + + assert isinstance(valid, Item) and valid.value == {"original": True} + assert isinstance(invalid, InvalidNamespaceError) + + +async def test_sync_methods_reject_invalid_namespace_labels( + store: AsyncPostgresStore, +) -> None: + """The sync wrappers run off the event loop thread and must validate too.""" + await store.aput(("foo", "bar"), "key", {"original": True}) + + for call in ( + lambda: store.get(("foo.bar",), "key"), + lambda: store.search(("foo.bar",)), + lambda: store.delete(("foo.bar",), "key"), + lambda: store.list_namespaces(prefix=("foo.bar",)), + lambda: store.batch([GetOp(("foo.bar",), "key")]), + ): + with pytest.raises(InvalidNamespaceError): + await asyncio.to_thread(call) + + item = await store.aget(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index af5655f3e..abe6475e6 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -11,6 +11,7 @@ import pytest from langchain_core.embeddings import Embeddings from langgraph.store.base import ( GetOp, + InvalidNamespaceError, Item, ListNamespacesOp, MatchCondition, @@ -1164,3 +1165,51 @@ def test_namespace_labels_with_trailing_newline(store) -> None: assert set(store.list_namespaces(prefix=["users", "alice"], limit=100)) == { ("users", "alice"), } + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +@pytest.mark.parametrize( + "kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"] +) +def test_batch_rejects_invalid_namespace_labels( + store, namespace: tuple, kind: str +) -> None: + """Ops passed straight to `batch` must not reach another namespace. + + Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same + text as `("foo", "bar")`. `BaseStore` methods validate labels themselves, + but `batch` takes ops as given. + """ + op = { + "get": GetOp(namespace, "key"), + "put": PutOp(namespace, "key", {"changed": True}), + "delete": PutOp(namespace, "key", None), + "search": SearchOp(namespace), + "list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)), + "list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)), + }[kind] + store.put(("foo", "bar"), "key", {"original": True}) + + with pytest.raises(InvalidNamespaceError): + store.batch([PutOp(("valid",), "key", {}), op]) + + item = store.get(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + # The whole batch is rejected before any SQL runs. + assert store.get(("valid",), "key") is None + + +def test_batch_allows_empty_search_prefix_and_listing_wildcards( + store, +) -> None: + store.put(("foo", "bar"), "key", {"v": 1}) + + found, listed = store.batch( + [ + SearchOp(()), + ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)), + ] + ) + + assert [item.namespace for item in found] == [("foo", "bar")] + assert listed == [("foo", "bar")] diff --git a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py index fbc23beda..c4f8f6d18 100644 --- a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py +++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py @@ -28,6 +28,7 @@ from langgraph.store.base import ( ensure_embeddings, get_text_at_path, tokenize_path, + validate_op_namespace, ) _AIO_ERROR_MSG = ( @@ -257,6 +258,7 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list) tot = 0 for idx, op in enumerate(ops): + validate_op_namespace(op) grouped_ops[type(op)].append((idx, op)) tot += 1 return grouped_ops, tot diff --git a/libs/checkpoint-sqlite/tests/test_async_store.py b/libs/checkpoint-sqlite/tests/test_async_store.py index 428c4e90c..4af125c9f 100644 --- a/libs/checkpoint-sqlite/tests/test_async_store.py +++ b/libs/checkpoint-sqlite/tests/test_async_store.py @@ -9,8 +9,10 @@ from typing import cast import pytest from langgraph.store.base import ( GetOp, + InvalidNamespaceError, Item, ListNamespacesOp, + MatchCondition, PutOp, SearchOp, ) @@ -745,3 +747,65 @@ async def test_async_namespace_segment_boundary(store: AsyncSqliteStore) -> None assert set(await store.alist_namespaces(suffix=["alice"], limit=100)) == { ("uid", "users", "alice"), } + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +async def test_abatch_rejects_invalid_namespace_labels( + store: AsyncSqliteStore, namespace: tuple +) -> None: + await store.aput(("foo", "bar"), "key", {"original": True}) + + for op in ( + GetOp(namespace, "key"), + GetOp(namespace, "key", refresh_ttl=True), + PutOp(namespace, "key", {"changed": True}), + PutOp(namespace, "key", None), + SearchOp(namespace), + ListNamespacesOp((MatchCondition("prefix", namespace),)), + ListNamespacesOp((MatchCondition("suffix", namespace),)), + ): + with pytest.raises(InvalidNamespaceError): + await store.abatch([op]) + + item = await store.aget(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + + +async def test_invalid_namespace_only_fails_its_own_call( + store: AsyncSqliteStore, +) -> None: + """Concurrent calls share one `abatch`, which fails every op if it raises. + + Labels are checked before an op is queued, so one caller's bad label cannot + fail another caller's request. + """ + await store.aput(("foo", "bar"), "key", {"original": True}) + + valid, invalid = await asyncio.gather( + store.aget(("foo", "bar"), "key"), + store.aget(("foo.bar",), "key"), + return_exceptions=True, + ) + + assert isinstance(valid, Item) and valid.value == {"original": True} + assert isinstance(invalid, InvalidNamespaceError) + + +async def test_sync_methods_reject_invalid_namespace_labels( + store: AsyncSqliteStore, +) -> None: + """The sync wrappers run off the event loop thread and must validate too.""" + await store.aput(("foo", "bar"), "key", {"original": True}) + + for call in ( + lambda: store.get(("foo.bar",), "key"), + lambda: store.search(("foo.bar",)), + lambda: store.delete(("foo.bar",), "key"), + lambda: store.list_namespaces(prefix=("foo.bar",)), + lambda: store.batch([GetOp(("foo.bar",), "key")]), + ): + with pytest.raises(InvalidNamespaceError): + await asyncio.to_thread(call) + + item = await store.aget(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} diff --git a/libs/checkpoint-sqlite/tests/test_store.py b/libs/checkpoint-sqlite/tests/test_store.py index d4b06847e..91835b512 100644 --- a/libs/checkpoint-sqlite/tests/test_store.py +++ b/libs/checkpoint-sqlite/tests/test_store.py @@ -14,6 +14,7 @@ import pytest from langchain_core.embeddings import Embeddings from langgraph.store.base import ( GetOp, + InvalidNamespaceError, Item, ListNamespacesOp, MatchCondition, @@ -1435,3 +1436,51 @@ def test_list_namespaces_metacharacter_labels(store: SqliteStore) -> None: assert set(store.list_namespaces(prefix=[label, "child"], limit=100)) == { (label, "child"), } + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +@pytest.mark.parametrize( + "kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"] +) +def test_batch_rejects_invalid_namespace_labels( + store: SqliteStore, namespace: tuple, kind: str +) -> None: + """Ops passed straight to `batch` must not reach another namespace. + + Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same + text as `("foo", "bar")`. `BaseStore` methods validate labels themselves, + but `batch` takes ops as given. + """ + op = { + "get": GetOp(namespace, "key"), + "put": PutOp(namespace, "key", {"changed": True}), + "delete": PutOp(namespace, "key", None), + "search": SearchOp(namespace), + "list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)), + "list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)), + }[kind] + store.put(("foo", "bar"), "key", {"original": True}) + + with pytest.raises(InvalidNamespaceError): + store.batch([PutOp(("valid",), "key", {}), op]) + + item = store.get(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + # The whole batch is rejected before any SQL runs. + assert store.get(("valid",), "key") is None + + +def test_batch_allows_empty_search_prefix_and_listing_wildcards( + store: SqliteStore, +) -> None: + store.put(("foo", "bar"), "key", {"v": 1}) + + found, listed = store.batch( + [ + SearchOp(()), + ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)), + ] + ) + + assert [item.namespace for item in found] == [("foo", "bar")] + assert listed == [("foo", "bar")] diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index 727da8987..6e5d9f4d9 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -771,7 +771,12 @@ class BaseStore(ABC): Returns: The retrieved item or `None` if not found. + + Raises: + InvalidNamespaceError: If a namespace label is empty, is not a string, + or contains a period (`.`). """ + _validate_namespace_labels(namespace) return self.batch( [GetOp(namespace, str(key), _ensure_refresh(self.ttl_config, refresh_ttl))] )[0] @@ -801,6 +806,10 @@ class BaseStore(ABC): Returns: List of items matching the search criteria. + Raises: + InvalidNamespaceError: If a `namespace_prefix` label is empty, is not a + string, or contains a period (`.`). + ???+ example "Examples" Basic filtering: @@ -840,6 +849,7 @@ class BaseStore(ABC): Natural language search support depends on your store implementation and requires proper embedding configuration. """ + _validate_namespace_labels(namespace_prefix) return self.batch( [ SearchOp( @@ -887,6 +897,11 @@ class BaseStore(ABC): By default, the expiration timer refreshes on both read operations (get/search) and write operations (put/update), whenever the item is included in the operation. + Raises: + InvalidNamespaceError: If the namespace is empty, its root label is + `"langgraph"`, or a label is empty, is not a string, or contains a + period (`.`). + Note: Indexing support depends on your store implementation. If you do not initialize the store with indexing capabilities, @@ -940,7 +955,12 @@ class BaseStore(ABC): Args: namespace: Hierarchical path for the item. key: Unique identifier within the namespace. + + Raises: + InvalidNamespaceError: If a namespace label is empty, is not a string, + or contains a period (`.`). """ + _validate_namespace_labels(namespace) self.batch([PutOp(namespace, str(key), None, ttl=None)]) def list_namespaces( @@ -969,6 +989,10 @@ class BaseStore(ABC): A list of namespace tuples that match the criteria. Each tuple represents a full namespace path up to `max_depth`. + Raises: + InvalidNamespaceError: If a `prefix` or `suffix` label is empty, is not a + string, or contains a period (`.`). + ???+ example "Examples": Setting `max_depth=3`. Given the namespaces: @@ -984,6 +1008,8 @@ class BaseStore(ABC): # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] ``` """ + _validate_namespace_labels(prefix or ()) + _validate_namespace_labels(suffix or ()) match_conditions = [] if prefix: match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) @@ -1013,7 +1039,12 @@ class BaseStore(ABC): Returns: The retrieved item or `None` if not found. + + Raises: + InvalidNamespaceError: If a namespace label is empty, is not a string, + or contains a period (`.`). """ + _validate_namespace_labels(namespace) return ( await self.abatch( [ @@ -1052,6 +1083,10 @@ class BaseStore(ABC): Returns: List of items matching the search criteria. + Raises: + InvalidNamespaceError: If a `namespace_prefix` label is empty, is not a + string, or contains a period (`.`). + ???+ example "Examples" Basic filtering: @@ -1091,6 +1126,7 @@ class BaseStore(ABC): Natural language search support depends on your store implementation and requires proper embedding configuration. """ + _validate_namespace_labels(namespace_prefix) return ( await self.abatch( [ @@ -1140,6 +1176,11 @@ class BaseStore(ABC): By default, the expiration timer refreshes on both read operations (get/search) and write operations (put/update), whenever the item is included in the operation. + Raises: + InvalidNamespaceError: If the namespace is empty, its root label is + `"langgraph"`, or a label is empty, is not a string, or contains a + period (`.`). + Note: Indexing support depends on your store implementation. If you do not initialize the store with indexing capabilities, @@ -1201,7 +1242,12 @@ class BaseStore(ABC): Args: namespace: Hierarchical path for the item. key: Unique identifier within the namespace. + + Raises: + InvalidNamespaceError: If a namespace label is empty, is not a string, + or contains a period (`.`). """ + _validate_namespace_labels(namespace) await self.abatch([PutOp(namespace, str(key), None)]) async def alist_namespaces( @@ -1230,6 +1276,10 @@ class BaseStore(ABC): A list of namespace tuples that match the criteria. Each tuple represents a full namespace path up to `max_depth`. + Raises: + InvalidNamespaceError: If a `prefix` or `suffix` label is empty, is not a + string, or contains a period (`.`). + ???+ example "Examples" Setting `max_depth=3` with existing namespaces: @@ -1245,6 +1295,8 @@ class BaseStore(ABC): # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] ``` """ + _validate_namespace_labels(prefix or ()) + _validate_namespace_labels(suffix or ()) match_conditions = [] if prefix: match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) @@ -1263,6 +1315,14 @@ class BaseStore(ABC): def _validate_namespace(namespace: tuple[str, ...]) -> None: if not namespace: raise InvalidNamespaceError("Namespace cannot be empty.") + _validate_namespace_labels(namespace) + if namespace[0] == "langgraph": + raise InvalidNamespaceError( + f'Root label for namespace cannot be "langgraph". Got: {namespace}' + ) + + +def _validate_namespace_labels(namespace: tuple[str, ...]) -> None: for label in namespace: if not isinstance(label, str): raise InvalidNamespaceError( @@ -1277,10 +1337,27 @@ def _validate_namespace(namespace: tuple[str, ...]) -> None: raise InvalidNamespaceError( f"Namespace labels cannot be empty strings. Got {label} in {namespace}" ) - if namespace[0] == "langgraph": - raise InvalidNamespaceError( - f'Root label for namespace cannot be "langgraph". Got: {namespace}' - ) + + +def validate_op_namespace(op: Op) -> None: + """Validate the namespace labels an op carries before a store executes it. + + `BaseStore` methods check labels before batching, but ops passed directly to + `batch`/`abatch` skip those methods. Stores that serialize namespaces as + delimited text should call this for every op they execute, so a label such + as `"foo.bar"` cannot address the namespace `("foo", "bar")`. + + Raises: + InvalidNamespaceError: If a label is empty, is not a string, or contains + a period (`.`). + """ + if isinstance(op, (GetOp, PutOp)): + _validate_namespace_labels(op.namespace) + elif isinstance(op, SearchOp): + _validate_namespace_labels(op.namespace_prefix) + elif isinstance(op, ListNamespacesOp): + for condition in op.match_conditions or (): + _validate_namespace_labels(condition.path) def _ensure_refresh( @@ -1319,4 +1396,5 @@ __all__ = [ "ensure_embeddings", "tokenize_path", "get_text_at_path", + "validate_op_namespace", ] diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 9deb02e2a..8b591df76 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -25,6 +25,7 @@ from langgraph.store.base import ( _ensure_refresh, _ensure_ttl, _validate_namespace, + _validate_namespace_labels, ) F = TypeVar("F", bound=Callable) @@ -86,6 +87,7 @@ class AsyncBatchedBaseStore(BaseStore): *, refresh_ttl: bool | None = None, ) -> Item | None: + _validate_namespace_labels(namespace) self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait( @@ -111,6 +113,7 @@ class AsyncBatchedBaseStore(BaseStore): offset: int = 0, refresh_ttl: bool | None = None, ) -> list[SearchItem]: + _validate_namespace_labels(namespace_prefix) self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait( @@ -155,6 +158,7 @@ class AsyncBatchedBaseStore(BaseStore): namespace: tuple[str, ...], key: str, ) -> None: + _validate_namespace_labels(namespace) self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait((fut, PutOp(namespace, key, None))) @@ -169,6 +173,8 @@ class AsyncBatchedBaseStore(BaseStore): limit: int = 100, offset: int = 0, ) -> list[tuple[str, ...]]: + _validate_namespace_labels(prefix or ()) + _validate_namespace_labels(suffix or ()) self._ensure_task() fut = self._loop.create_future() match_conditions = [] diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index 9079aeecc..bee5c68d2 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -13,10 +13,14 @@ from langgraph.store.base import ( GetOp, InvalidNamespaceError, Item, + ListNamespacesOp, + MatchCondition, Op, PutOp, Result, + SearchOp, get_text_at_path, + validate_op_namespace, ) from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.memory import InMemoryStore @@ -528,6 +532,127 @@ async def test_cannot_put_empty_namespace() -> None: assert (await async_store.aget(("valid", "namespace"), "key")) is None +INVALID_NAMESPACES = [("foo.bar",), ("foo", ""), (123,)] +NAMESPACE_METHODS = ["get", "delete", "search", "prefix", "suffix"] + + +@pytest.mark.parametrize("namespace", INVALID_NAMESPACES) +@pytest.mark.parametrize("method", NAMESPACE_METHODS) +def test_rejects_invalid_namespace_labels( + mocker: MockerFixture, namespace: tuple, method: str +) -> None: + store = InMemoryStore() + batch = mocker.spy(InMemoryStore, "batch") + call = { + "get": lambda: store.get(namespace, "key"), + "delete": lambda: store.delete(namespace, "key"), + "search": lambda: store.search(namespace), + "prefix": lambda: store.list_namespaces(prefix=namespace), + "suffix": lambda: store.list_namespaces(suffix=namespace), + }[method] + + with pytest.raises(InvalidNamespaceError): + call() + + batch.assert_not_called() + + +@pytest.mark.parametrize("batched", [False, True]) +@pytest.mark.parametrize("namespace", INVALID_NAMESPACES) +@pytest.mark.parametrize("method", NAMESPACE_METHODS) +async def test_async_rejects_invalid_namespace_labels( + mocker: MockerFixture, batched: bool, namespace: tuple, method: str +) -> None: + # The batched store must reject before queueing: a failure inside the + # shared `abatch` would fail every op queued alongside this one. + store = MockAsyncBatchedStore() if batched else InMemoryStore() + # `MockAsyncBatchedStore` dispatches through `InMemoryStore.batch`. + batch = mocker.spy(InMemoryStore, "batch") + abatch = mocker.spy(InMemoryStore, "abatch") + call = { + "get": lambda: store.aget(namespace, "key"), + "delete": lambda: store.adelete(namespace, "key"), + "search": lambda: store.asearch(namespace), + "prefix": lambda: store.alist_namespaces(prefix=namespace), + "suffix": lambda: store.alist_namespaces(suffix=namespace), + }[method] + + with pytest.raises(InvalidNamespaceError): + await call() + + batch.assert_not_called() + abatch.assert_not_called() + + +def test_search_and_listing_keep_empty_prefixes_and_wildcards() -> None: + store = InMemoryStore() + store.put(("tenant", "a_%"), "key", {"v": 1}) + store.put(("tenant", "b", "child"), "key", {"v": 1}) + + assert len(store.search(())) == 2 + assert [item.namespace for item in store.search(("tenant", "a_%"))] == [ + ("tenant", "a_%") + ] + assert sorted(store.list_namespaces(prefix=("tenant", "*"), suffix=("*",))) == [ + ("tenant", "a_%"), + ("tenant", "b", "child"), + ] + + +@pytest.mark.parametrize("batched", [False, True]) +async def test_async_search_and_listing_keep_empty_prefixes_and_wildcards( + batched: bool, +) -> None: + store = MockAsyncBatchedStore() if batched else InMemoryStore() + await store.aput(("tenant", "a_%"), "key", {"v": 1}) + await store.aput(("tenant", "b", "child"), "key", {"v": 1}) + + assert len(await store.asearch(())) == 2 + assert [item.namespace for item in await store.asearch(("tenant", "a_%"))] == [ + ("tenant", "a_%") + ] + assert sorted( + await store.alist_namespaces(prefix=("tenant", "*"), suffix=("*",)) + ) == [("tenant", "a_%"), ("tenant", "b", "child")] + + +@pytest.mark.parametrize("namespace", INVALID_NAMESPACES) +@pytest.mark.parametrize( + "kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"] +) +def test_validate_op_namespace_rejects_invalid_labels( + namespace: tuple, kind: str +) -> None: + op = { + "get": GetOp(namespace, "key"), + "put": PutOp(namespace, "key", {"v": 1}), + "delete": PutOp(namespace, "key", None), + "search": SearchOp(namespace), + "list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)), + "list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)), + }[kind] + + with pytest.raises(InvalidNamespaceError): + validate_op_namespace(op) + + +def test_validate_op_namespace_allows_empty_prefix_and_wildcards() -> None: + for op in ( + SearchOp(()), + ListNamespacesOp(), + ListNamespacesOp( + ( + MatchCondition("prefix", ("tenant", "*")), + MatchCondition("suffix", ("*",)), + ) + ), + GetOp(("tenant", "a_%"), "key"), + # Write-only rules belong to `put`, not to op validation. + PutOp(("langgraph", "x"), "key", {"v": 1}), + ): + validate_op_namespace(op) + + async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None: abatch = mocker.spy(InMemoryStore, "batch") store = MockAsyncBatchedStore()