diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index c407f26f1..9f7906ab8 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -25,6 +25,7 @@ from langgraph.store.base import ( BaseStore, GetOp, IndexConfig, + InvalidNamespaceError, Item, ListNamespacesOp, Op, @@ -1386,11 +1387,36 @@ 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 +def _validate_op_namespace(op: Op) -> None: + """Reject namespace labels that would alias another namespace once dot-joined. + + `("foo.bar",)` and `("foo", "bar")` flatten to the same `prefix` text. + `BaseStore` methods check labels before batching, but ops passed directly to + `batch`/`abatch` reach here unchecked. + """ + if isinstance(op, (GetOp, PutOp)): + paths = (op.namespace,) + elif isinstance(op, SearchOp): + paths = (op.namespace_prefix,) + elif isinstance(op, ListNamespacesOp): + paths = tuple(condition.path for condition in op.match_conditions or ()) + else: + return + for path in paths: + for label in path: + if not isinstance(label, str) or not label or "." in label: + raise InvalidNamespaceError( + f"Invalid namespace label {label!r} found in {path}. Namespace " + "labels must be non-empty strings without periods ('.')." + ) + + def _json_loads(content: bytes | orjson.Fragment) -> Any: if isinstance(content, orjson.Fragment): if hasattr(content, "buf"): diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index 080a78852..998f85574 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,62 @@ 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"] + + +async def test_abatch_rejects_ambiguous_namespace_labels( + store: AsyncPostgresStore, +) -> None: + await store.aput(("foo", "bar"), "key", {"original": True}) + + for op in ( + GetOp(("foo.bar",), "key"), + PutOp(("foo.bar",), "key", {"changed": True}), + PutOp(("foo.bar",), "key", None), + SearchOp(("foo.bar",)), + ListNamespacesOp((MatchCondition("prefix", ("foo.bar",)),)), + ): + 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_ambiguous_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..495b4395d 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_ambiguous_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..5e901c49b 100644 --- a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py +++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py @@ -17,6 +17,7 @@ from langgraph.store.base import ( BaseStore, GetOp, IndexConfig, + InvalidNamespaceError, Item, ListNamespacesOp, Op, @@ -257,11 +258,36 @@ 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 +def _validate_op_namespace(op: Op) -> None: + """Reject namespace labels that would alias another namespace once dot-joined. + + `("foo.bar",)` and `("foo", "bar")` flatten to the same `prefix` text. + `BaseStore` methods check labels before batching, but ops passed directly to + `batch`/`abatch` reach here unchecked. + """ + if isinstance(op, (GetOp, PutOp)): + paths = (op.namespace,) + elif isinstance(op, SearchOp): + paths = (op.namespace_prefix,) + elif isinstance(op, ListNamespacesOp): + paths = tuple(condition.path for condition in op.match_conditions or ()) + else: + return + for path in paths: + for label in path: + if not isinstance(label, str) or not label or "." in label: + raise InvalidNamespaceError( + f"Invalid namespace label {label!r} found in {path}. Namespace " + "labels must be non-empty strings without periods ('.')." + ) + + class PreparedGetQuery(NamedTuple): query: str # Main query to execute params: tuple # Parameters for the main query diff --git a/libs/checkpoint-sqlite/tests/test_async_store.py b/libs/checkpoint-sqlite/tests/test_async_store.py index 428c4e90c..767ceabf0 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,62 @@ async def test_async_namespace_segment_boundary(store: AsyncSqliteStore) -> None assert set(await store.alist_namespaces(suffix=["alice"], limit=100)) == { ("uid", "users", "alice"), } + + +async def test_abatch_rejects_ambiguous_namespace_labels( + store: AsyncSqliteStore, +) -> None: + await store.aput(("foo", "bar"), "key", {"original": True}) + + for op in ( + GetOp(("foo.bar",), "key"), + PutOp(("foo.bar",), "key", {"changed": True}), + PutOp(("foo.bar",), "key", None), + SearchOp(("foo.bar",)), + ListNamespacesOp((MatchCondition("prefix", ("foo.bar",)),)), + ): + 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_ambiguous_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..0d9b16b12 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_ambiguous_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 fdfdf7299..e6996675c 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -771,6 +771,10 @@ 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( @@ -802,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: @@ -942,6 +950,10 @@ 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)]) @@ -972,6 +984,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: @@ -1018,6 +1034,10 @@ 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 ( @@ -1058,6 +1078,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: @@ -1208,6 +1232,10 @@ 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)]) @@ -1238,6 +1266,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: diff --git a/libs/checkpoint/tests/test_store_namespace_validation.py b/libs/checkpoint/tests/test_store_namespace_validation.py index 2ab0c86c8..7a193b71f 100644 --- a/libs/checkpoint/tests/test_store_namespace_validation.py +++ b/libs/checkpoint/tests/test_store_namespace_validation.py @@ -3,53 +3,64 @@ import pytest from langgraph.store.base import BaseStore, InvalidNamespaceError from langgraph.store.base.batch import AsyncBatchedBaseStore +INVALID_NAMESPACES = [("foo.bar",), ("foo", ""), (123,)] +METHODS = ["get", "delete", "search", "prefix", "suffix"] + class RecordingStore(BaseStore): def __init__(self): self.operations = [] def batch(self, ops): + ops = list(ops) self.operations.extend(ops) - return [None for _ in self.operations] + return [None for _ in ops] async def abatch(self, ops): return self.batch(ops) class RecordingBatchedStore(AsyncBatchedBaseStore): + def __init__(self): + super().__init__() + self.operations = [] + async def abatch(self, ops): + ops = list(ops) + self.operations.extend(ops) return [None for _ in ops] -@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), (123,)]) -@pytest.mark.parametrize("method", ["get", "delete", "search", "list_namespaces"]) +@pytest.mark.parametrize("namespace", INVALID_NAMESPACES) +@pytest.mark.parametrize("method", METHODS) def test_base_rejects_invalid_labels(namespace, method): store = RecordingStore() - args = (namespace, "key") if method in ("get", "delete") else (namespace,) with pytest.raises(InvalidNamespaceError): - if method == "list_namespaces": - store.list_namespaces(prefix=namespace) + if method in ("prefix", "suffix"): + store.list_namespaces(**{method: namespace}) + elif method == "search": + store.search(namespace) else: - getattr(store, method)(*args) + getattr(store, method)(namespace, "key") assert store.operations == [] @pytest.mark.parametrize("store_type", [RecordingStore, RecordingBatchedStore]) -@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), (123,)]) -@pytest.mark.parametrize("method", ["aget", "adelete", "asearch", "alist_namespaces"]) +@pytest.mark.parametrize("namespace", INVALID_NAMESPACES) +@pytest.mark.parametrize("method", METHODS) async def test_async_rejects_invalid_labels(store_type, namespace, method): store = store_type() - args = (namespace, "key") if method in ("aget", "adelete") else (namespace,) try: with pytest.raises(InvalidNamespaceError): - if method == "alist_namespaces": - await store.alist_namespaces(suffix=namespace) + if method in ("prefix", "suffix"): + await store.alist_namespaces(**{method: namespace}) + elif method == "search": + await store.asearch(namespace) else: - await getattr(store, method)(*args) + await getattr(store, f"a{method}")(namespace, "key") if isinstance(store, RecordingBatchedStore): assert store._aqueue.empty() - else: - assert store.operations == [] + assert store.operations == [] finally: if isinstance(store, RecordingBatchedStore) and store._task: store._task.cancel() @@ -62,7 +73,10 @@ def test_search_and_listing_preserve_prefixes(): store.list_namespaces(prefix=("tenant", "*"), suffix=("*",)) assert store.operations[0].namespace_prefix == () assert store.operations[1].namespace_prefix == ("tenant", "a_%") - assert store.operations[2].match_conditions[0].path == ("tenant", "*") + assert [c.path for c in store.operations[2].match_conditions] == [ + ("tenant", "*"), + ("*",), + ] @pytest.mark.parametrize("store_type", [RecordingStore, RecordingBatchedStore]) @@ -70,7 +84,14 @@ async def test_async_empty_search_and_listing_wildcards(store_type): store = store_type() try: await store.asearch(()) - await store.alist_namespaces(prefix=("tenant", "*")) + await store.asearch(("tenant", "a_%")) + await store.alist_namespaces(prefix=("tenant", "*"), suffix=("*",)) finally: if isinstance(store, RecordingBatchedStore) and store._task: store._task.cancel() + assert store.operations[0].namespace_prefix == () + assert store.operations[1].namespace_prefix == ("tenant", "a_%") + assert [c.path for c in store.operations[2].match_conditions] == [ + ("tenant", "*"), + ("*",), + ]