fix(store): keep quotes and escape whitespace in store key queries

langgraph/store/redis/__init__.py

@@ -27,7 +27,6 @@
 from redisvl.index import SearchIndex
 from redisvl.query import FilterQuery, VectorQuery
 from redisvl.redis.connection import RedisConnectionFactory
-from redisvl.utils.token_escaper import TokenEscaper
 from ulid import ULID
 
 from langgraph.checkpoint.redis.util import unwrap_connection_errors
@@ -43,11 +42,11 @@
     _namespace_to_text,
     _row_to_item,
     _row_to_search_item,
+    _tag_escaper,
 )
 
 from .token_unescaper import TokenUnescaper
 
-_token_escaper = TokenEscaper()
 _token_unescaper = TokenUnescaper()
 
 logger = logging.getLogger(__name__)
@@ -324,7 +323,7 @@ def _batch_put_ops(
         # First delete any existing documents that are being updated/deleted
         for _, op in put_ops:
             namespace = _namespace_to_text(op.namespace)
-            query = f"@prefix:{namespace} @key:{{{_token_escaper.escape(op.key)}}}"
+            query = f"@prefix:{namespace} @key:{{{_tag_escaper.escape(op.key)}}}"
             results = self.store_index.search(query)
 
             if self.cluster_mode:

langgraph/store/redis/aio.py

@@ -26,7 +26,6 @@
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import FilterQuery, VectorQuery
 from redisvl.redis.connection import RedisConnectionFactory
-from redisvl.utils.token_escaper import TokenEscaper
 from ulid import ULID
 
 from langgraph.checkpoint.redis.util import unwrap_connection_errors
@@ -42,12 +41,12 @@
     _namespace_to_text,
     _row_to_item,
     _row_to_search_item,
+    _tag_escaper,
     logger,
 )
 
 from .token_unescaper import TokenUnescaper
 
-_token_escaper = TokenEscaper()
 _token_unescaper = TokenUnescaper()
 from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster
 
@@ -499,7 +498,7 @@ async def _aprepare_batch_PUT_queries(
             # Delete matching documents
             for op in deletes:
                 prefix = _namespace_to_text(op.namespace)
-                query = f"(@prefix:{prefix} @key:{{{_token_escaper.escape(op.key)}}})"
+                query = f"(@prefix:{prefix} @key:{{{_tag_escaper.escape(op.key)}}})"
                 results = await self.store_index.search(query)
                 for doc in results.docs:
                     await self._redis.delete(doc.id)
@@ -560,7 +559,7 @@ async def _batch_put_ops(
         # First delete any existing documents that are being updated/deleted
         for _, op in put_ops:
             namespace = _namespace_to_text(op.namespace)
-            query = f"@prefix:{namespace} @key:{{{_token_escaper.escape(op.key)}}}"
+            query = f"@prefix:{namespace} @key:{{{_tag_escaper.escape(op.key)}}}"
             results = await self.store_index.search(query)
 
             if self.cluster_mode:

langgraph/store/redis/base.py

@@ -5,6 +5,7 @@
 import copy
 import json
 import logging
+import re
 import threading
 from collections import defaultdict
 from datetime import datetime, timedelta, timezone
@@ -39,7 +40,6 @@
 from redis.asyncio import Redis as AsyncRedis
 from redis.exceptions import ResponseError
 from redisvl.index import SearchIndex
-from redisvl.query.filter import Tag, Text
 from redisvl.utils.token_escaper import TokenEscaper
 
 from langgraph.checkpoint.redis.jsonplus_redis import JsonPlusRedisSerializer
@@ -48,6 +48,11 @@
 from .types import IndexType, RedisClientType
 
 _token_escaper = TokenEscaper()
+# For tag values. The default set leaves "|", which separates tag alternatives,
+# and whitespace other than a space, which ends the value.
+_tag_escaper = TokenEscaper(
+    re.compile(TokenEscaper.DEFAULT_ESCAPED_CHARS[:-1] + r"|\n\r\t\v\f]")
+)
 _token_unescaper = TokenUnescaper()
 
 logger = logging.getLogger(__name__)
@@ -494,12 +499,12 @@ def _get_batch_GET_ops_queries(
         results: list[tuple[str, Sequence, tuple[str, ...], list]] = []
         for namespace, items in namespace_groups.items():
             _, keys = zip(*items)
-            # Use Tag helper to properly escape all special characters
-            prefix_filter = Text("prefix") == _namespace_to_text(namespace)
-            filter_str = f"({prefix_filter} "
+            # redisvl's Text filter turns a double quote into a space, so it
+            # can't match the escaped namespace that put() stored.
+            filter_str = f'(@prefix:("{_namespace_to_text(namespace)}") '
             if keys:
-                key_filter = Tag("key") == list(keys)
-                filter_str += f"{key_filter})"
+                key_filter = "|".join(_tag_escaper.escape(key) for key in keys)
+                filter_str += f"@key:{{{key_filter}}})"
             else:
                 filter_str += ")"
             results.append((filter_str, [], namespace, items))
@@ -532,7 +537,7 @@ def _prepare_batch_PUT_queries(
             # Delete matching documents
             for op in deletes:
                 prefix = _namespace_to_text(op.namespace)
-                query = f"(@prefix:{prefix} @key:{{{_token_escaper.escape(op.key)}}})"
+                query = f"(@prefix:{prefix} @key:{{{_tag_escaper.escape(op.key)}}})"
                 results = self.store_index.search(query)
                 for doc in results.docs:
                     self._redis.delete(doc.id)

tests/test_issue_237_store_get_escaping.py

@@ -0,0 +1,79 @@
+"""Test for issue #237 - get() must find every item that search() finds.
+
+``get()`` built its query with redisvl's ``Text`` and ``Tag`` filters. ``Text``
+turned a double quote in a namespace label into a space, and ``Tag`` left a
+newline in a key unescaped, so ``get()`` returned ``None`` for items that
+``put()`` had stored and ``search()`` could find.
+
+``delete()`` and the update path of ``put()`` had the same problem with keys,
+and with ``|`` too: ``delete()`` left the item in place, and ``put()`` kept the
+old document next to the new one.
+"""
+
+import pytest
+from langgraph.store.base import GetOp
+
+from langgraph.store.redis import AsyncRedisStore, RedisStore
+
+ITEMS = [
+    (("issue237", 'q"t'), "k"),
+    (("issue237", "newline"), "k\n1"),
+    (("issue237", "carriage-return"), "k\r1"),
+    (("issue237", "tab"), "k\t1"),
+    (("issue237", "vertical-tab"), "k\x0b1"),
+    (("issue237", "form-feed"), "k\x0c1"),
+    (("issue237", "pipe"), "k|1"),
+]
+
+
+@pytest.mark.parametrize("namespace,key", ITEMS)
+def test_sync_store_get_update_delete(redis_url, namespace, key) -> None:
+    with RedisStore.from_conn_string(redis_url) as store:
+        store.setup()
+        store.put(namespace, key, {"v": 1})
+        assert store.get(namespace, key).value == {"v": 1}
+
+        store.put(namespace, key, {"v": 2})
+        assert store.get(namespace, key).value == {"v": 2}
+        assert [item.value for item in store.search(namespace)] == [{"v": 2}]
+
+        store.delete(namespace, key)
+        assert store.get(namespace, key) is None
+        assert store.search(namespace) == []
+
+
+@pytest.mark.parametrize("namespace,key", ITEMS)
+async def test_async_store_get_update_delete(redis_url, namespace, key) -> None:
+    async with AsyncRedisStore.from_conn_string(redis_url) as store:
+        await store.setup()
+        await store.aput(namespace, key, {"v": 1})
+        assert (await store.aget(namespace, key)).value == {"v": 1}
+
+        await store.aput(namespace, key, {"v": 2})
+        assert (await store.aget(namespace, key)).value == {"v": 2}
+        assert [item.value for item in await store.asearch(namespace)] == [{"v": 2}]
+
+        await store.adelete(namespace, key)
+        assert await store.aget(namespace, key) is None
+        assert await store.asearch(namespace) == []
+
+
+def test_batched_get_with_several_keys(redis_url) -> None:
+    namespace = ("issue237", "batch")
+    keys = ["plain", "new\nline", "pi|pe"]
+    with RedisStore.from_conn_string(redis_url) as store:
+        store.setup()
+        for key in keys:
+            store.put(namespace, key, {"key": key})
+
+        items = store.batch([GetOp(namespace, key) for key in keys])
+
+        assert [item.value["key"] for item in items] == keys
+
+
+def test_get_stays_within_the_exact_namespace(redis_url) -> None:
+    with RedisStore.from_conn_string(redis_url) as store:
+        store.setup()
+        store.put(("issue237", "parent", "child"), "k", {"v": 1})
+
+        assert store.get(("issue237", "parent"), "k") is None