From a3cb9c1a94aa18ac24aeba8a77c0b2fe2df5e5d7 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 1 Oct 2024 18:43:45 -0700
Subject: [PATCH] Validate no empty namespace is added (#1961)

Also check the root label isn't "langgraph"
---
 .../langgraph/store/base/__init__.py          |  8 +++
 libs/checkpoint/tests/test_store.py           | 55 ++++++++++++++++++-
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py
index d1ad914a6..2e1dbf604 100644
--- a/libs/checkpoint/langgraph/store/base/__init__.py
+++ b/libs/checkpoint/langgraph/store/base/__init__.py
@@ -159,11 +159,19 @@ class InvalidNamespaceError(ValueError):
 
 
 def _validate_namespace(namespace: tuple[str, ...]) -> None:
+    if not namespace:
+        raise InvalidNamespaceError("Namespace cannot be empty.")
     for label in namespace:
         if "." in label:
             raise InvalidNamespaceError(
                 f"Invalid namespace label '{label}'. Namespace labels cannot contain periods ('.')."
             )
+        elif not label:
+            raise InvalidNamespaceError("Namespace labels cannot be empty strings.")
+    if namespace[0] == "langgraph":
+        raise InvalidNamespaceError(
+            f'Root label for namespace cannot be "langgraph". Got: {namespace}'
+        )
 
 
 class BaseStore(ABC):
diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py
index 915b91e0a..e1d3720a7 100644
--- a/libs/checkpoint/tests/test_store.py
+++ b/libs/checkpoint/tests/test_store.py
@@ -2,9 +2,10 @@ import asyncio
 from datetime import datetime
 from typing import Iterable
 
+import pytest
 from pytest_mock import MockerFixture
 
-from langgraph.store.base import GetOp, Item, Op, Result
+from langgraph.store.base import GetOp, InvalidNamespaceError, Item, Op, PutOp, Result
 from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.memory import InMemoryStore
 
@@ -259,3 +260,55 @@ def test_list_namespaces_empty_store() -> None:
 
     result = store.list_namespaces()
     assert result == []
+
+
+async def test_cannot_put_empty_namespace() -> None:
+    store = InMemoryStore()
+    doc = {"foo": "bar"}
+
+    with pytest.raises(InvalidNamespaceError):
+        store.put((), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await store.aput((), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        store.put(("the", "thing.about"), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await store.aput(("the", "thing.about"), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        store.put(("some", "fun", ""), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await store.aput(("some", "fun", ""), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await store.aput(("langgraph", "foo"), "bar", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        store.put(("langgraph", "foo"), "bar", doc)
+
+    await store.aput(("foo", "langgraph", "foo"), "bar", doc)
+    assert (await store.aget(("foo", "langgraph", "foo"), "bar")).value == doc  # type: ignore[union-attr]
+    assert (await store.asearch(("foo", "langgraph", "foo")))[0].value == doc
+    await store.adelete(("foo", "langgraph", "foo"), "bar")
+    assert (await store.aget(("foo", "langgraph", "foo"), "bar")) is None
+    store.put(("foo", "langgraph", "foo"), "bar", doc)
+    assert store.get(("foo", "langgraph", "foo"), "bar").value == doc  # type: ignore[union-attr]
+    assert store.search(("foo", "langgraph", "foo"))[0].value == doc
+    store.delete(("foo", "langgraph", "foo"), "bar")
+    assert store.get(("foo", "langgraph", "foo"), "bar") is None
+
+    # Do the same but go past the public put api
+    await store.abatch([PutOp(("langgraph", "foo"), "bar", doc)])
+    assert (await store.aget(("langgraph", "foo"), "bar")).value == doc  # type: ignore[union-attr]
+    assert (await store.asearch(("langgraph", "foo")))[0].value == doc
+    await store.adelete(("langgraph", "foo"), "bar")
+    assert (await store.aget(("langgraph", "foo"), "bar")) is None
+    store.batch([PutOp(("langgraph", "foo"), "bar", doc)])
+    assert store.get(("langgraph", "foo"), "bar").value == doc  # type: ignore[union-attr]
+    assert store.search(("langgraph", "foo"))[0].value == doc
+    store.delete(("langgraph", "foo"), "bar")
+    assert store.get(("langgraph", "foo"), "bar") is None