From e57af6cc4a027910c31b45615f0a240e2391e74e Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 16 Oct 2026 17:41:35 +0200 Subject: [PATCH 1/6] Python: preserve Redis context scope identity --- .../workflows/python-integration-tests.yml | 6 +- .github/workflows/python-merge-tests.yml | 6 +- python/packages/redis/AGENTS.md | 4 + python/packages/redis/README.md | 31 +++ .../_context_provider.py | 141 ++++++++++-- .../tests/test_context_provider_edges.py | 147 +++++++++++- .../test_context_provider_integration.py | 214 ++++++++++++++++++ .../redis/tests/test_context_provider_sdk.py | 13 +- python/packages/redis/tests/test_providers.py | 7 +- 9 files changed, 526 insertions(+), 43 deletions(-) create mode 100644 python/packages/redis/tests/test_context_provider_integration.py diff --git a/.github/workflows/python-integration-tests.yml b/.github/workflows/python-integration-tests.yml index 97bcb8cb20f..b7dff8ca41c 100644 --- a/.github/workflows/python-integration-tests.yml +++ b/.github/workflows/python-integration-tests.yml @@ -654,9 +654,11 @@ jobs: with: python-version: ${{ env.UV_PYTHON }} os: ${{ runner.os }} - - name: Test Redis HASH and JSON vector collections + - name: Test Redis integrations run: > - uv run --no-sync pytest packages/redis/tests/test_vector_store.py + uv run --no-sync pytest + packages/redis/tests/test_vector_store.py + packages/redis/tests/test_context_provider_integration.py -m integration --timeout=120 --junitxml=pytest-redis.xml - name: Upload test results if: always() diff --git a/.github/workflows/python-merge-tests.yml b/.github/workflows/python-merge-tests.yml index 3b747555f66..6a7de1a7a4d 100644 --- a/.github/workflows/python-merge-tests.yml +++ b/.github/workflows/python-merge-tests.yml @@ -894,9 +894,11 @@ jobs: with: python-version: ${{ env.UV_PYTHON }} os: ${{ runner.os }} - - name: Test Redis HASH and JSON vector collections + - name: Test Redis integrations run: > - uv run --no-sync pytest packages/redis/tests/test_vector_store.py + uv run --no-sync pytest + packages/redis/tests/test_vector_store.py + packages/redis/tests/test_context_provider_integration.py -m integration --timeout=120 --junitxml=pytest-redis.xml - name: Upload test results if: always() diff --git a/python/packages/redis/AGENTS.md b/python/packages/redis/AGENTS.md index 00aff40a2f6..d1921183dd4 100644 --- a/python/packages/redis/AGENTS.md +++ b/python/packages/redis/AGENTS.md @@ -10,6 +10,10 @@ Redis-based storage for agent threads and context. an owned client from a URL or credential provider, or borrow a standalone async Redis client configured with `decode_responses=True`; borrowed clients remain caller-owned. - **`RedisContextProvider`** - Context provider with Redis-backed retrieval + Scope identifiers are stored as raw metadata but indexed through versioned, + domain-separated tokens in case-sensitive TAG fields. Existing legacy + indexes fail closed unless `overwrite_index=True` is explicitly used for a + coordinated in-place rebuild and token backfill. - **`RedisSettings`** - TypedDict connection settings for vector stores, resolved with core `load_settings` from explicit URL overrides, an optional .env file, or `REDIS_URL`. URLs use `SecretString` to mask credentials. - **`RedisCollection` / `RedisStore`** - Experimental generic vector storage and search over native HASH and JSON diff --git a/python/packages/redis/README.md b/python/packages/redis/README.md index 018fa4ecedf..c7b2ce8ffbc 100644 --- a/python/packages/redis/README.md +++ b/python/packages/redis/README.md @@ -148,6 +148,37 @@ and session into the corresponding scoped key using an application-owned migration process. After verifying the copied history, remove legacy keys separately according to the application's retention policy. +## Isolate retrieved memories + +`RedisContextProvider` stores exact application, agent, user, and session +identifiers as metadata, but indexes only deterministic versioned tokens for +those values. Scope comparisons therefore preserve case and treat punctuation +as ordinary identifier content. Values supplied inside documents cannot +override the provider's configured scopes. + +Indexes created by earlier releases use a different scope schema. A provider +using the current schema fails closed when it encounters such an index. During +a coordinated deployment, stop older writers and rebuild the index once while +preserving its hashes: + +```python +context_provider = RedisContextProvider( + redis_url="redis://localhost:6379", + index_name="context", + prefix="context", + application_id="support-app", + user_id="user-123", + overwrite_index=True, +) +``` + +The rebuild replaces the index definition first, then backfills the current +scope tokens from the exact identifiers stored in each matching Redis hash. +Existing memory content remains in place. After the migration completes, +remove `overwrite_index=True`; leaving it enabled rebuilds the index whenever a +new provider process initializes. Review or remove records whose historical +scope metadata is not authoritative before migration. + ## Store and search documents This example uses precomputed vectors, so no embedding service is needed. diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index 117a0dd19b7..16fbf3f68fb 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -8,6 +8,8 @@ from __future__ import annotations +import asyncio +import hashlib import json import sys from typing import TYPE_CHECKING, Any, ClassVar, Literal @@ -20,10 +22,9 @@ IntegrationInvalidRequestException, ) from redisvl.index import AsyncSearchIndex -from redisvl.query import AggregateHybridQuery, TextQuery -from redisvl.query.filter import FilterExpression, Tag +from redisvl.query import AggregateHybridQuery, FilterQuery, TextQuery +from redisvl.query.filter import FilterExpression from redisvl.redis.utils import array_to_buffer -from redisvl.utils.token_escaper import TokenEscaper from redisvl.utils.vectorize import BaseVectorizer from ._feature_usage import FeatureIndex @@ -41,6 +42,23 @@ if TYPE_CHECKING: from agent_framework._agents import SupportsAgentRun +_SCOPE_SCHEMA_VERSION = 1 +_SCOPE_FIELDS = ("application_id", "agent_id", "user_id", "thread_id", "conversation_id") +_SCOPE_TAG_SEPARATOR = "\x1f" +_SCOPE_MIGRATION_BATCH_SIZE = 100 +_SCOPE_MIGRATION_TIMEOUT_SECONDS = 10.0 + + +def _scope_token_field(name: str) -> str: + if name not in _SCOPE_FIELDS: + raise ValueError(f"Unsupported Redis context scope field: {name}") + return f"af_scope_v{_SCOPE_SCHEMA_VERSION}_{name}" + + +def _scope_token(name: str, value: str) -> str: + payload = f"agent-framework:redis-context:v{_SCOPE_SCHEMA_VERSION}:{name}\0{value}".encode() + return hashlib.sha256(payload).hexdigest() + class RedisContextProvider(ContextProvider): """Redis context provider using the new ContextProvider hooks pattern. @@ -93,7 +111,8 @@ def __init__( retrieval is not filtered by user. context_prompt: The context prompt to use for the provider. redis_index: The Redis index to use for the provider. - overwrite_index: Whether to overwrite the existing Redis index. + overwrite_index: Whether to rebuild an existing index in place and + migrate its stored scope identifiers to the current schema. """ super().__init__(source_id) self.redis_url = redis_url @@ -112,7 +131,6 @@ def __init__( self.user_id = user_id self.context_prompt = context_prompt or self.DEFAULT_CONTEXT_PROMPT self.overwrite_index = overwrite_index - self._token_escaper: TokenEscaper = TokenEscaper() self._index_initialized: bool = False self._schema_dict: dict[str, Any] | None = None # An index built here owns its connection pool and must release it on close; @@ -203,7 +221,11 @@ def schema_dict(self) -> dict[str, Any]: def _build_filter_from_dict(self, filters: dict[str, str | None]) -> Any | None: """Builds a combined filter expression from simple equality tags.""" - parts: list[FilterExpression] = [Tag(k) == v for k, v in filters.items() if v] + parts = [ + FilterExpression(f"@{_scope_token_field(name)}:{{{_scope_token(name, value)}}}") + for name, value in filters.items() + if value + ] if not parts: return None combined = parts[0] @@ -227,13 +249,16 @@ def _build_schema_dict( {"name": "role", "type": "tag"}, {"name": "mime_type", "type": "tag"}, {"name": "content", "type": "text"}, - {"name": "conversation_id", "type": "tag"}, {"name": "message_id", "type": "tag"}, {"name": "author_name", "type": "tag"}, - {"name": "application_id", "type": "tag"}, - {"name": "agent_id", "type": "tag"}, - {"name": "user_id", "type": "tag"}, - {"name": "thread_id", "type": "tag"}, + *[ + { + "name": _scope_token_field(name), + "type": "tag", + "attrs": {"case_sensitive": True, "separator": _SCOPE_TAG_SEPARATOR}, + } + for name in _SCOPE_FIELDS + ], ] if vector_field_name is not None and vector_dims is not None: fields.append({ @@ -256,9 +281,11 @@ async def _ensure_index(self) -> None: if self._index_initialized: return index_exists = await self.redis_index.exists() - if not self.overwrite_index and index_exists: + if index_exists and not self.overwrite_index: await self._validate_schema_compatibility() await self.redis_index.create(overwrite=self.overwrite_index, drop=False) + if index_exists and self.overwrite_index: + await self._migrate_scope_tokens() self._index_initialized = True async def _validate_schema_compatibility(self) -> None: @@ -297,8 +324,11 @@ def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]: sig["fields"][name] = {"type": ftype} return sig + existing_client = self.redis_index.client existing_index: Any = await AsyncSearchIndex.from_existing( # pyright: ignore[reportUnknownMemberType] - self.index_name, redis_url=self.redis_url + self.index_name, + redis_client=existing_client, + redis_url=None if existing_client is not None else self.redis_url, ) existing_schema = existing_index.schema.to_dict() current_schema = self.schema_dict @@ -309,9 +339,65 @@ def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]: "Existing Redis index schema is incompatible with the current configuration.\n" f"Existing (significant): {json.dumps(existing_sig, indent=2, sort_keys=True)}\n" f"Current (significant): {json.dumps(current_sig, indent=2, sort_keys=True)}\n" - "Set overwrite_index=True to rebuild if this change is intentional." + "Set overwrite_index=True during a coordinated migration to rebuild the index and populate " + "versioned scope tokens while preserving stored hashes." ) + async def _migrate_scope_tokens(self) -> None: + """Populate versioned scope tokens in hashes preserved during an index rebuild.""" + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index rebuild did not initialize a client for scope migration.") + + prefix = self.prefix.encode() + pending: list[tuple[Any, dict[str, str]]] = [] + migrated_documents = 0 + + async def flush() -> None: + nonlocal migrated_documents + if not pending: + return + async with client.pipeline(transaction=False) as pipeline: + for key, mapping in pending: + pipeline.hset(key, mapping=mapping) + await pipeline.execute() + migrated_documents += len(pending) + pending.clear() + + async for key in client.scan_iter(_type="hash"): + encoded_key = key if isinstance(key, bytes) else str(key).encode() + if not encoded_key.startswith(prefix): + continue + record = await client.hgetall(key) + tokens: dict[str, str] = {} + for name in _SCOPE_FIELDS: + raw_value = record.get(name) + if raw_value is None: + raw_value = record.get(name.encode()) + if raw_value is None: + continue + value = raw_value.decode() if isinstance(raw_value, bytes) else str(raw_value) + if value: + tokens[_scope_token_field(name)] = _scope_token(name, value) + if tokens: + pending.append((key, tokens)) + if len(pending) >= _SCOPE_MIGRATION_BATCH_SIZE: + await flush() + await flush() + + if not migrated_documents: + return + deadline = asyncio.get_running_loop().time() + _SCOPE_MIGRATION_TIMEOUT_SECONDS + while True: + info = await self.redis_index.info() + if int(info.get("num_docs", 0)) >= migrated_documents: + return + if asyncio.get_running_loop().time() >= deadline: + raise RuntimeError( + f"Redis context index rebuild timed out after migrating {migrated_documents} stored hashes." + ) + await asyncio.sleep(0.05) + async def _add( self, *, @@ -327,11 +413,16 @@ async def _add( prepared: list[dict[str, Any]] = [] for doc in docs: d = dict(doc) - d.setdefault("application_id", self.application_id) - d.setdefault("agent_id", self.agent_id) - d.setdefault("user_id", self.user_id) - d.setdefault("thread_id", session_id) - d.setdefault("conversation_id", session_id) + scope_values = { + "application_id": self.application_id, + "agent_id": self.agent_id, + "user_id": self.user_id, + "thread_id": session_id, + "conversation_id": session_id, + } + for name, value in scope_values.items(): + d[name] = value or None + d[_scope_token_field(name)] = _scope_token(name, value) if value else None if "content" not in d: raise IntegrationInvalidRequestException("add() requires a 'content' field in data") if self.vector_field_name: @@ -428,12 +519,18 @@ def _validate_filters(self) -> None: raise ValueError("At least one of the filters: agent_id, user_id, or application_id is required.") async def search_all(self, page_size: int = 200) -> list[dict[str, Any]]: - """Returns all documents in the index.""" - from redisvl.query import FilterQuery + """Returns all documents matching the provider's configured scopes.""" + await self._ensure_index() + self._validate_filters() + combined_filter = self._build_filter_from_dict({ + "application_id": self.application_id, + "agent_id": self.agent_id, + "user_id": self.user_id, + }) out: list[dict[str, Any]] = [] async for batch in self.redis_index.paginate( - FilterQuery(FilterExpression("*"), return_fields=[], num_results=page_size), + FilterQuery(combined_filter or FilterExpression("*"), return_fields=[], num_results=page_size), page_size=page_size, ): out.extend(batch) diff --git a/python/packages/redis/tests/test_context_provider_edges.py b/python/packages/redis/tests/test_context_provider_edges.py index ea6287ffe37..e8e343d7669 100644 --- a/python/packages/redis/tests/test_context_provider_edges.py +++ b/python/packages/redis/tests/test_context_provider_edges.py @@ -14,7 +14,13 @@ from redisvl.redis.utils import buffer_to_array from redisvl.utils.vectorize import BaseVectorizer -from agent_framework_redis._context_provider import RedisContextProvider +from agent_framework_redis._context_provider import ( + _SCOPE_FIELDS, + _SCOPE_TAG_SEPARATOR, + RedisContextProvider, + _scope_token, + _scope_token_field, +) @pytest.fixture @@ -38,15 +44,31 @@ def patch_index(mock_index: AsyncMock) -> Generator[MagicMock]: def test_build_filter_from_dict_combines_multiple_tags( patch_index: MagicMock, # noqa: ARG001 ) -> None: - provider = RedisContextProvider(source_id="ctx", user_id="user-1") + provider = RedisContextProvider(source_id="ctx", user_id="user|1") combined = provider._build_filter_from_dict({ - "application_id": "app-1", + "application_id": "App,1", "agent_id": None, - "user_id": "user-1", + "user_id": "user|1", }) assert combined is not None - assert str(combined) == "(@application_id:{app\\-1} @user_id:{user\\-1})" + assert str(combined) == ( + f"(@{_scope_token_field('application_id')}:{{{_scope_token('application_id', 'App,1')}}} " + f"@{_scope_token_field('user_id')}:{{{_scope_token('user_id', 'user|1')}}})" + ) + assert "App,1" not in str(combined) + assert "user|1" not in str(combined) + + +def test_scope_tokens_preserve_identity_and_field_domain() -> None: + assert _scope_token("user_id", "Alice") != _scope_token("user_id", "alice") + assert _scope_token("user_id", "same") != _scope_token("agent_id", "same") + assert set(_scope_token("user_id", r"a,b|c\ d{}*")) <= set("0123456789abcdef") + + +def test_scope_token_field_rejects_unknown_names() -> None: + with pytest.raises(ValueError, match="Unsupported Redis context scope field"): + _scope_token_field("content") def test_schema_dict_includes_vector_configuration( @@ -76,6 +98,21 @@ def test_schema_dict_includes_vector_configuration( } +def test_schema_uses_versioned_identity_preserving_scope_tags( + patch_index: MagicMock, # noqa: ARG001 +) -> None: + provider = RedisContextProvider(source_id="ctx", user_id="user-1") + fields = {field["name"]: field for field in provider.schema_dict["fields"]} + + for name in _SCOPE_FIELDS: + assert name not in fields + assert fields[_scope_token_field(name)] == { + "name": _scope_token_field(name), + "type": "tag", + "attrs": {"case_sensitive": True, "separator": _SCOPE_TAG_SEPARATOR}, + } + + async def test_ensure_index_short_circuits_after_first_initialization( mock_index: AsyncMock, patch_index: MagicMock, # noqa: ARG001 @@ -104,6 +141,21 @@ async def test_ensure_index_validates_existing_schema_before_create( assert provider._index_initialized is True +async def test_ensure_index_rebuilds_and_migrates_when_explicitly_enabled( + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 +) -> None: + mock_index.exists.return_value = True + provider = RedisContextProvider(source_id="ctx", user_id="user-1", overwrite_index=True) + + with patch.object(provider, "_migrate_scope_tokens", AsyncMock()) as migrate: + await provider._ensure_index() + + mock_index.create.assert_awaited_once_with(overwrite=True, drop=False) + migrate.assert_awaited_once() + assert provider._index_initialized is True + + async def test_validate_schema_compatibility_raises_for_significant_mismatch( patch_index: MagicMock, ) -> None: @@ -117,7 +169,7 @@ async def test_validate_schema_compatibility_raises_for_significant_mismatch( ) patch_index.from_existing = AsyncMock(return_value=existing_index) - with pytest.raises(ValueError, match="overwrite_index=True"): + with pytest.raises(ValueError, match="coordinated migration"): await provider._validate_schema_compatibility() @@ -154,8 +206,16 @@ async def test_add_vectorizes_documents_and_applies_defaults( with patch.object(provider, "_ensure_index", AsyncMock()): await provider._add( data=[ - {"content": "first"}, - {"content": "second", "conversation_id": "custom-conversation"}, + { + "content": "first", + "application_id": "caller-app", + "agent_id": "caller-agent", + "user_id": "caller-user", + "thread_id": "caller-thread", + "conversation_id": "caller-conversation", + _scope_token_field("user_id"): "caller-token", + }, + {"content": "second"}, ], session_id="session-1", ) @@ -167,11 +227,66 @@ async def test_add_vectorizes_documents_and_applies_defaults( assert loaded_docs[0]["user_id"] == "user-1" assert loaded_docs[0]["thread_id"] == "session-1" assert loaded_docs[0]["conversation_id"] == "session-1" + for name, value in { + "application_id": "app-1", + "agent_id": "agent-1", + "user_id": "user-1", + "thread_id": "session-1", + "conversation_id": "session-1", + }.items(): + assert loaded_docs[0][_scope_token_field(name)] == _scope_token(name, value) assert isinstance(loaded_docs[0]["embedding"], bytes) assert isinstance(loaded_docs[1]["embedding"], bytes) vectorizer.aembed_many.assert_awaited_once_with(["first", "second"], batch_size=2) +async def test_migrate_scope_tokens_backfills_preserved_hashes( + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 +) -> None: + client = MagicMock() + + async def scan_iter(**kwargs: Any) -> AsyncIterator[bytes]: + assert kwargs == {"_type": "hash"} + yield b"context:one" + yield b"other:two" + + client.scan_iter = scan_iter + client.hgetall = AsyncMock( + side_effect=[ + { + b"application_id": b"App", + b"agent_id": b"Agent", + b"user_id": b"Alice", + b"thread_id": b"Thread", + b"conversation_id": b"Conversation", + } + ] + ) + pipeline = MagicMock() + pipeline.__aenter__ = AsyncMock(return_value=pipeline) + pipeline.__aexit__ = AsyncMock(return_value=False) + pipeline.hset = MagicMock() + pipeline.execute = AsyncMock(return_value=[1]) + client.pipeline.return_value = pipeline + mock_index.client = client + mock_index.info = AsyncMock(return_value={"num_docs": 1}) + provider = RedisContextProvider(source_id="ctx", user_id="user-1") + + await provider._migrate_scope_tokens() + + mapping = pipeline.hset.call_args.kwargs["mapping"] + for name, value in { + "application_id": "App", + "agent_id": "Agent", + "user_id": "Alice", + "thread_id": "Thread", + "conversation_id": "Conversation", + }.items(): + assert mapping[_scope_token_field(name)] == _scope_token(name, value) + mock_index.info.assert_awaited_once() + + @pytest.mark.parametrize("dtype, expected_size", [("float16", 4), ("float32", 8), ("float64", 16), ("bfloat16", 4)]) async def test_after_run_stores_vectors_in_configured_dtype( dtype: str, @@ -257,12 +372,18 @@ async def test_search_all_collects_paginated_batches( patch_index: MagicMock, # noqa: ARG001 ) -> None: provider = RedisContextProvider(source_id="ctx", user_id="user-1") + captured_query: list[Any] = [] + + def paginate(*args: Any, **kwargs: Any) -> AsyncIterator[list[dict[str, str]]]: # noqa: ARG001 + captured_query.append(args[0]) + + async def batches() -> AsyncIterator[list[dict[str, str]]]: + yield [{"content": "first"}] + yield [{"content": "second"}, {"content": "third"}] - async def paginate(*args: Any, **kwargs: Any) -> AsyncIterator[list[dict[str, str]]]: # noqa: ARG001 - yield [{"content": "first"}] - yield [{"content": "second"}, {"content": "third"}] + return batches() - mock_index.paginate = MagicMock(return_value=paginate()) + mock_index.paginate = MagicMock(side_effect=paginate) results = await provider.search_all(page_size=2) @@ -271,3 +392,5 @@ async def paginate(*args: Any, **kwargs: Any) -> AsyncIterator[list[dict[str, st {"content": "second"}, {"content": "third"}, ] + assert _scope_token_field("user_id") in str(captured_query[0]._filter_expression) + assert "user-1" not in str(captured_query[0]._filter_expression) diff --git a/python/packages/redis/tests/test_context_provider_integration.py b/python/packages/redis/tests/test_context_provider_integration.py new file mode 100644 index 00000000000..a511c7447ab --- /dev/null +++ b/python/packages/redis/tests/test_context_provider_integration.py @@ -0,0 +1,214 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live Redis isolation and migration coverage for ``RedisContextProvider``.""" + +from __future__ import annotations + +import os +from collections.abc import AsyncIterator, Awaitable +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from redis.asyncio import Redis +from redisvl.index import AsyncSearchIndex +from redisvl.utils.vectorize import BaseVectorizer + +from agent_framework_redis import RedisContextProvider +from agent_framework_redis._context_provider import _scope_token_field + + +@pytest.fixture +async def live_redis() -> AsyncIterator[tuple[str, Redis, str, str]]: + url = os.getenv("REDIS_VECTOR_TEST_URL") + if not url: + pytest.skip("Set REDIS_VECTOR_TEST_URL to a disposable Redis Search server.") + suffix = uuid4().hex + index_name = f"af-context-test-{suffix}" + prefix = f"af-context-test-{suffix}:" + client = Redis.from_url(url) + try: + yield url, client, index_name, prefix + finally: + indexes = cast(list[bytes], await client.execute_command("FT._LIST")) + encoded_name = index_name.encode() + if encoded_name in indexes or index_name in indexes: + await client.execute_command("FT.DROPINDEX", index_name) + keys = [key async for key in client.scan_iter(match=f"{prefix}*")] + if keys: + await client.delete(*keys) + await client.aclose() + + +def _provider( + url: str, + index_name: str, + prefix: str, + *, + application_id: str = "app", + agent_id: str = "agent", + user_id: str = "user", + **kwargs: Any, +) -> RedisContextProvider: + return RedisContextProvider( + redis_url=url, + index_name=index_name, + prefix=prefix, + application_id=application_id, + agent_id=agent_id, + user_id=user_id, + **kwargs, + ) + + +@pytest.mark.integration +@pytest.mark.parametrize( + ("victim_id", "other_id"), + [ + ("Alice", "alice"), + ("a,b", "a"), + ("victim", "x|victim"), + (r"a,b|c\ d{}*", r"a,b|c\ d{}?"), + ], +) +async def test_live_user_scope_preserves_exact_identity( + live_redis: tuple[str, Redis, str, str], victim_id: str, other_id: str +) -> None: + url, _, index_name, prefix = live_redis + victim = _provider(url, index_name, prefix, user_id=victim_id) + other = _provider(url, index_name, prefix, user_id=other_id) + + await victim._add(data={"content": "shared marker victim memory"}, session_id="victim-session") + + assert [item["content"] for item in await victim._redis_search("shared marker")] == ["shared marker victim memory"] + assert await other._redis_search("shared marker") == [] + assert len(await victim.search_all()) == 1 + assert await other.search_all() == [] + + +@pytest.mark.integration +@pytest.mark.parametrize("scope_name", ["application_id", "agent_id", "user_id"]) +async def test_live_each_provider_scope_is_independent( + live_redis: tuple[str, Redis, str, str], scope_name: str +) -> None: + url, _, index_name, prefix = live_redis + victim_scopes = {"application_id": "app", "agent_id": "agent", "user_id": "user"} + other_scopes = dict(victim_scopes) + other_scopes[scope_name] = victim_scopes[scope_name] + ",other|scope" + victim = _provider(url, index_name, prefix, **victim_scopes) + other = _provider(url, index_name, prefix, **other_scopes) + + await victim._add(data={"content": f"scope marker {scope_name}"}, session_id="session") + + assert len(await victim._redis_search(f"scope marker {scope_name}")) == 1 + assert await other._redis_search(f"scope marker {scope_name}") == [] + + +@pytest.mark.integration +async def test_live_session_scopes_and_document_overrides_are_provider_owned( + live_redis: tuple[str, Redis, str, str], +) -> None: + url, client, index_name, prefix = live_redis + provider = _provider(url, index_name, prefix, user_id="authoritative-user") + await provider._add( + data={ + "content": "session marker", + "application_id": "caller-app", + "agent_id": "caller-agent", + "user_id": "caller-user", + "thread_id": "caller-thread", + "conversation_id": "caller-conversation", + _scope_token_field("user_id"): "caller-token", + }, + session_id="authoritative-session", + ) + + assert len(await provider._redis_search("session marker", session_id="authoritative-session")) == 1 + assert await provider._redis_search("session marker", session_id="other-session") == [] + caller = _provider( + url, index_name, prefix, application_id="caller-app", agent_id="caller-agent", user_id="caller-user" + ) + assert await caller._redis_search("session marker") == [] + + keys = [key async for key in client.scan_iter(match=f"{prefix}*")] + assert len(keys) == 1 + record = await cast(Awaitable[dict[bytes, bytes]], client.hgetall(keys[0])) + assert record[b"application_id"] == b"app" + assert record[b"agent_id"] == b"agent" + assert record[b"user_id"] == b"authoritative-user" + assert record[b"thread_id"] == b"authoritative-session" + assert record[b"conversation_id"] == b"authoritative-session" + assert record[_scope_token_field("user_id").encode()] != b"caller-token" + + +@pytest.mark.integration +async def test_live_hybrid_search_preserves_scope_isolation( + live_redis: tuple[str, Redis, str, str], +) -> None: + url, _, index_name, prefix = live_redis + vectorizer = MagicMock(spec=BaseVectorizer) + vectorizer.dims = 2 + vectorizer.dtype = "float32" + vectorizer.aembed_many = AsyncMock(return_value=[[1.0, 0.0]]) + vectorizer.aembed = AsyncMock(return_value=[1.0, 0.0]) + victim = _provider( + url, + index_name, + prefix, + user_id="victim", + redis_vectorizer=vectorizer, + vector_field_name="embedding", + ) + other = _provider( + url, + index_name, + prefix, + user_id="x|victim", + redis_vectorizer=vectorizer, + vector_field_name="embedding", + ) + + await victim._add(data={"content": "hybrid marker"}, session_id="session") + + assert len(await victim._redis_search("hybrid marker")) == 1 + assert await other._redis_search("hybrid marker") == [] + + +@pytest.mark.integration +async def test_live_legacy_index_requires_explicit_rebuild_and_preserves_exact_scopes( + live_redis: tuple[str, Redis, str, str], +) -> None: + url, client, index_name, prefix = live_redis + legacy_schema = { + "index": {"name": index_name, "prefix": prefix, "key_separator": ":", "storage_type": "hash"}, + "fields": [ + {"name": "content", "type": "text"}, + {"name": "application_id", "type": "tag"}, + {"name": "agent_id", "type": "tag"}, + {"name": "user_id", "type": "tag"}, + {"name": "thread_id", "type": "tag"}, + {"name": "conversation_id", "type": "tag"}, + ], + } + legacy_index = AsyncSearchIndex.from_dict(legacy_schema, redis_client=client) + await legacy_index.create() + await legacy_index.load([ + { + "content": "migration marker", + "application_id": "app", + "agent_id": "agent", + "user_id": "Alice", + "thread_id": "session", + "conversation_id": "session", + } + ]) + + incompatible = _provider(url, index_name, prefix, user_id="Alice") + with pytest.raises(ValueError, match="coordinated migration"): + await incompatible._redis_search("migration marker") + + migrated = _provider(url, index_name, prefix, user_id="Alice", overwrite_index=True) + assert len(await migrated._redis_search("migration marker")) == 1 + assert await _provider(url, index_name, prefix, user_id="alice")._redis_search("migration marker") == [] + assert await _provider(url, index_name, prefix, user_id="Alice,other")._redis_search("migration marker") == [] diff --git a/python/packages/redis/tests/test_context_provider_sdk.py b/python/packages/redis/tests/test_context_provider_sdk.py index c8a497e0578..a0f496bb380 100644 --- a/python/packages/redis/tests/test_context_provider_sdk.py +++ b/python/packages/redis/tests/test_context_provider_sdk.py @@ -18,6 +18,7 @@ from redisvl.utils.vectorize import BaseVectorizer from agent_framework_redis import RedisContextProvider +from agent_framework_redis._context_provider import _scope_token, _scope_token_field class _TransportConnection(Connection): @@ -122,10 +123,14 @@ async def test_after_run_encodes_messages_with_optional_fields_omitted( assert record["user_id"] == "user-1" assert record["conversation_id"] == "session-1" assert record["thread_id"] == "session-1" + assert record[_scope_token_field("user_id")] == _scope_token("user_id", "user-1") + assert record[_scope_token_field("conversation_id")] == _scope_token("conversation_id", "session-1") + assert record[_scope_token_field("thread_id")] == _scope_token("thread_id", "session-1") assert "message_id" not in record assert "author_name" not in record for key, value in scopes.items(): assert record[key] == value + assert record[_scope_token_field(key)] == _scope_token(key, value) if not all_scopes: assert "application_id" not in record assert "agent_id" not in record @@ -163,11 +168,13 @@ async def test_borrowed_index_omits_null_values_without_changing_nonnull_values_ await provider._add(data=document) record = write.call_args.args[2] - assert "application_id" not in record + assert record["application_id"] == "configured-app" + assert record["agent_id"] == "configured-agent" + assert record[_scope_token_field("application_id")] == _scope_token("application_id", "configured-app") + assert record[_scope_token_field("agent_id")] == _scope_token("agent_id", "configured-agent") assert "thread_id" not in record assert "optional_metadata" not in record - assert record["agent_id"] == "0" - assert record["conversation_id"] == "0" + assert "conversation_id" not in record assert record["message_id"] == "0" assert record["author_name"] == "" assert record["mime_type"] == 0 diff --git a/python/packages/redis/tests/test_providers.py b/python/packages/redis/tests/test_providers.py index 711c63e7686..5f554fb5beb 100644 --- a/python/packages/redis/tests/test_providers.py +++ b/python/packages/redis/tests/test_providers.py @@ -14,7 +14,7 @@ from agent_framework._sessions import AgentSession, SessionContext from redis.asyncio import Redis -from agent_framework_redis._context_provider import RedisContextProvider +from agent_framework_redis._context_provider import _SCOPE_FIELDS, RedisContextProvider, _scope_token_field from agent_framework_redis._feature_usage import FeatureIndex from agent_framework_redis._history_provider import RedisHistoryProvider, _redis_result @@ -136,8 +136,11 @@ def test_schema_has_expected_fields(self, patch_index_from_dict: MagicMock): # provider = RedisContextProvider(source_id="ctx", user_id="u1") schema = provider.schema_dict field_names = [f["name"] for f in schema["fields"]] - for expected in ("role", "content", "conversation_id", "message_id", "application_id", "agent_id", "user_id"): + for expected in ("role", "content", "message_id"): assert expected in field_names + for scope_field in _SCOPE_FIELDS: + assert scope_field not in field_names + assert _scope_token_field(scope_field) in field_names assert schema["index"]["name"] == "context" assert schema["index"]["prefix"] == "context" From a6fcf087e214758e33c050afd9fda1fe7a54653b Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 19 Oct 2026 10:14:30 +0200 Subject: [PATCH 2/6] Python: make Redis scope migration resumable --- python/packages/redis/README.md | 15 +- .../_context_provider.py | 322 +++++++++++++----- .../tests/test_context_provider_edges.py | 128 +++++-- .../test_context_provider_integration.py | 25 +- .../redis/tests/test_context_provider_sdk.py | 12 +- python/packages/redis/tests/test_providers.py | 4 + 6 files changed, 377 insertions(+), 129 deletions(-) diff --git a/python/packages/redis/README.md b/python/packages/redis/README.md index c7b2ce8ffbc..776eb66374a 100644 --- a/python/packages/redis/README.md +++ b/python/packages/redis/README.md @@ -172,12 +172,15 @@ context_provider = RedisContextProvider( ) ``` -The rebuild replaces the index definition first, then backfills the current -scope tokens from the exact identifiers stored in each matching Redis hash. -Existing memory content remains in place. After the migration completes, -remove `overwrite_index=True`; leaving it enabled rebuilds the index whenever a -new provider process initializes. Review or remove records whose historical -scope metadata is not authoritative before migration. +The migration records resumable progress outside the indexed key namespace, +backfills the current scope tokens from the exact identifiers stored in each +matching Redis hash, and only then replaces the index definition. Each pass is +bounded; if a large namespace needs another pass, initialization raises an +actionable error and the next provider with `overwrite_index=True` resumes from +the saved cursor. Normal providers reject an incomplete migration. Existing +memory content remains in place. After the migration completes, remove +`overwrite_index=True`. Review or remove records whose historical scope +metadata is not authoritative before migration. ## Store and search documents diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index 16fbf3f68fb..5e3bf9ac59b 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -11,6 +11,7 @@ import asyncio import hashlib import json +import secrets import sys from typing import TYPE_CHECKING, Any, ClassVar, Literal @@ -46,7 +47,15 @@ _SCOPE_FIELDS = ("application_id", "agent_id", "user_id", "thread_id", "conversation_id") _SCOPE_TAG_SEPARATOR = "\x1f" _SCOPE_MIGRATION_BATCH_SIZE = 100 +_SCOPE_MIGRATION_MAX_RECORDS = 1000 _SCOPE_MIGRATION_TIMEOUT_SECONDS = 10.0 +_SCOPE_MIGRATION_LOCK_SECONDS = 30 +_SCOPE_MIGRATION_BACKFILL = "backfill" +_SCOPE_MIGRATION_REINDEXING = "reindexing" +_SCOPE_MIGRATION_COMPLETE = "complete" + +_TAG_DEFAULTS = {"separator": ",", "case_sensitive": False, "withsuffixtrie": False} +_TEXT_DEFAULTS = {"weight": 1.0, "no_stem": False} def _scope_token_field(name: str) -> str: @@ -60,6 +69,53 @@ def _scope_token(name: str, value: str) -> str: return hashlib.sha256(payload).hexdigest() +def _scope_migration_key(index_name: str, prefix: str) -> str: + identity = hashlib.sha256(f"{index_name}\0{prefix}".encode()).hexdigest() + return f"agent-framework:redis-context:scope-schema:{identity}" + + +def _scope_key_namespace(prefix: str, key_separator: str) -> str: + if not prefix or not key_separator: + raise ValueError("Redis context index migration requires a non-empty key prefix and key separator.") + normalized_prefix = prefix.rstrip(key_separator) + if not normalized_prefix: + raise ValueError("Redis context index migration requires a key prefix containing non-separator characters.") + return f"{normalized_prefix}{key_separator}" + + +def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]: + def significant_index(index: dict[str, Any]) -> dict[str, Any]: + return {key: index.get(key) for key in ("name", "prefix", "key_separator", "storage_type")} + + def tag_attrs(attrs: dict[str, Any] | None) -> dict[str, Any]: + merged = {**_TAG_DEFAULTS, **(attrs or {})} + return {key: merged[key] for key in ("separator", "case_sensitive", "withsuffixtrie")} + + def text_attrs(attrs: dict[str, Any] | None) -> dict[str, Any]: + merged = {**_TEXT_DEFAULTS, **(attrs or {})} + return {key: merged[key] for key in ("weight", "no_stem")} + + signature: dict[str, Any] = {"index": significant_index(schema.get("index", {})), "fields": {}} + for raw_field in schema.get("fields", []): + field: dict[str, Any] = raw_field + name, field_type = field.get("name"), field.get("type") + if not name: + continue + if field_type == "tag": + signature["fields"][name] = {"type": "tag", "attrs": tag_attrs(field.get("attrs"))} + elif field_type == "text": + signature["fields"][name] = {"type": "text", "attrs": text_attrs(field.get("attrs"))} + elif field_type == "vector": + attrs: dict[str, Any] = field.get("attrs") or {} + signature["fields"][name] = { + "type": "vector", + "attrs": {key: attrs.get(key) for key in ("algorithm", "dims", "distance_metric", "datatype")}, + } + else: + signature["fields"][name] = {"type": field_type} + return signature + + class RedisContextProvider(ContextProvider): """Redis context provider using the new ContextProvider hooks pattern. @@ -281,55 +337,18 @@ async def _ensure_index(self) -> None: if self._index_initialized: return index_exists = await self.redis_index.exists() - if index_exists and not self.overwrite_index: - await self._validate_schema_compatibility() - await self.redis_index.create(overwrite=self.overwrite_index, drop=False) - if index_exists and self.overwrite_index: + if not index_exists: + await self.redis_index.create() + await self._set_scope_migration_state(_SCOPE_MIGRATION_COMPLETE, cursor=0) + elif self.overwrite_index: await self._migrate_scope_tokens() + else: + await self._validate_schema_compatibility() self._index_initialized = True async def _validate_schema_compatibility(self) -> None: """Validate that existing index schema matches current configuration.""" - TAG_DEFAULTS = {"separator": ",", "case_sensitive": False, "withsuffixtrie": False} - TEXT_DEFAULTS = {"weight": 1.0, "no_stem": False} - - def _significant_index(i: dict[str, Any]) -> dict[str, Any]: - return {k: i.get(k) for k in ("name", "prefix", "key_separator", "storage_type")} - - def _sig_tag(attrs: dict[str, Any] | None) -> dict[str, Any]: - a = {**TAG_DEFAULTS, **(attrs or {})} - return {k: a[k] for k in ("separator", "case_sensitive", "withsuffixtrie")} - - def _sig_text(attrs: dict[str, Any] | None) -> dict[str, Any]: - a = {**TEXT_DEFAULTS, **(attrs or {})} - return {k: a[k] for k in ("weight", "no_stem")} - - def _sig_vector(attrs: dict[str, Any] | None) -> dict[str, Any]: - a = {**(attrs or {})} - return {k: a.get(k) for k in ("algorithm", "dims", "distance_metric", "datatype")} - - def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]: - sig: dict[str, Any] = {"index": _significant_index(schema.get("index", {})), "fields": {}} - for f in schema.get("fields", []): - name, ftype = f.get("name"), f.get("type") - if not name: - continue - if ftype == "tag": - sig["fields"][name] = {"type": "tag", "attrs": _sig_tag(f.get("attrs"))} - elif ftype == "text": - sig["fields"][name] = {"type": "text", "attrs": _sig_text(f.get("attrs"))} - elif ftype == "vector": - sig["fields"][name] = {"type": "vector", "attrs": _sig_vector(f.get("attrs"))} - else: - sig["fields"][name] = {"type": ftype} - return sig - - existing_client = self.redis_index.client - existing_index: Any = await AsyncSearchIndex.from_existing( # pyright: ignore[reportUnknownMemberType] - self.index_name, - redis_client=existing_client, - redis_url=None if existing_client is not None else self.redis_url, - ) + existing_index = await self._existing_index() existing_schema = existing_index.schema.to_dict() current_schema = self.schema_dict existing_sig = _schema_signature(existing_schema) @@ -342,60 +361,189 @@ def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]: "Set overwrite_index=True during a coordinated migration to rebuild the index and populate " "versioned scope tokens while preserving stored hashes." ) + state = await self._scope_migration_state() + if state.get("version") != str(_SCOPE_SCHEMA_VERSION) or state.get("state") != _SCOPE_MIGRATION_COMPLETE: + raise ValueError( + "Redis context scope migration is incomplete. Set overwrite_index=True to resume the coordinated " + "migration before using this index." + ) + + async def _existing_index(self) -> Any: + existing_client = self.redis_index.client + return await AsyncSearchIndex.from_existing( # pyright: ignore[reportUnknownMemberType] + self.index_name, + redis_client=existing_client, + redis_url=None if existing_client is not None else self.redis_url, + ) + + async def _scope_migration_state(self) -> dict[str, str]: + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index did not initialize a client for scope migration state.") + values = await client.hmget( + _scope_migration_key(self.index_name, self.prefix), + "version", + "state", + "cursor", + "sample_field", + "sample_token", + ) + decoded = [ + value.decode() if isinstance(value, bytes) else str(value) if value is not None else "" for value in values + ] + return dict(zip(("version", "state", "cursor", "sample_field", "sample_token"), decoded, strict=True)) + + async def _set_scope_migration_state( + self, + state: str, + *, + cursor: int, + sample_field: str = "", + sample_token: str = "", + ) -> None: + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index did not initialize a client for scope migration state.") + await client.hset( + _scope_migration_key(self.index_name, self.prefix), + mapping={ + "version": str(_SCOPE_SCHEMA_VERSION), + "state": state, + "cursor": str(cursor), + "sample_field": sample_field, + "sample_token": sample_token, + }, + ) async def _migrate_scope_tokens(self) -> None: - """Populate versioned scope tokens in hashes preserved during an index rebuild.""" + """Resume scope-token backfill, rebuild the index, and mark the migration complete.""" client = self.redis_index.client if client is None: raise RuntimeError("Redis index rebuild did not initialize a client for scope migration.") + marker_key = _scope_migration_key(self.index_name, self.prefix) + lock_key = f"{marker_key}:lock" + lock_value = secrets.token_hex(16) + if not await client.set(lock_key, lock_value, nx=True, ex=_SCOPE_MIGRATION_LOCK_SECONDS): + raise RuntimeError("Another Redis context scope migration is already running.") - prefix = self.prefix.encode() - pending: list[tuple[Any, dict[str, str]]] = [] - migrated_documents = 0 - - async def flush() -> None: - nonlocal migrated_documents - if not pending: + try: + state = await self._scope_migration_state() + existing_index = await self._existing_index() + schema_is_current = _schema_signature(existing_index.schema.to_dict()) == _schema_signature( + self.schema_dict + ) + if state.get("state") == _SCOPE_MIGRATION_COMPLETE and schema_is_current: return - async with client.pipeline(transaction=False) as pipeline: - for key, mapping in pending: - pipeline.hset(key, mapping=mapping) - await pipeline.execute() - migrated_documents += len(pending) - pending.clear() - - async for key in client.scan_iter(_type="hash"): - encoded_key = key if isinstance(key, bytes) else str(key).encode() - if not encoded_key.startswith(prefix): - continue - record = await client.hgetall(key) - tokens: dict[str, str] = {} - for name in _SCOPE_FIELDS: - raw_value = record.get(name) - if raw_value is None: - raw_value = record.get(name.encode()) - if raw_value is None: - continue - value = raw_value.decode() if isinstance(raw_value, bytes) else str(raw_value) - if value: - tokens[_scope_token_field(name)] = _scope_token(name, value) - if tokens: - pending.append((key, tokens)) - if len(pending) >= _SCOPE_MIGRATION_BATCH_SIZE: - await flush() - await flush() - - if not migrated_documents: - return + if state.get("state") == _SCOPE_MIGRATION_REINDEXING and schema_is_current: + await self._wait_for_scope_index(state.get("sample_field", ""), state.get("sample_token", "")) + await self._set_scope_migration_state( + _SCOPE_MIGRATION_COMPLETE, + cursor=0, + sample_field=state.get("sample_field", ""), + sample_token=state.get("sample_token", ""), + ) + return + + namespace = _scope_key_namespace(self.prefix, self.redis_index.key_separator) + cursor = int(state.get("cursor") or 0) if state.get("state") == _SCOPE_MIGRATION_BACKFILL else 0 + sample_field = state.get("sample_field", "") + sample_token = state.get("sample_token", "") + await self._set_scope_migration_state( + _SCOPE_MIGRATION_BACKFILL, + cursor=cursor, + sample_field=sample_field, + sample_token=sample_token, + ) + + deadline = asyncio.get_running_loop().time() + _SCOPE_MIGRATION_TIMEOUT_SECONDS + processed = 0 + while True: + if processed >= _SCOPE_MIGRATION_MAX_RECORDS or asyncio.get_running_loop().time() >= deadline: + raise RuntimeError( + "Redis context scope migration paused after a bounded backfill pass. " + "Rerun with overwrite_index=True to resume." + ) + cursor, keys = await client.scan( + cursor=cursor, + match=f"{namespace}*", + count=min(_SCOPE_MIGRATION_BATCH_SIZE, _SCOPE_MIGRATION_MAX_RECORDS - processed), + _type="hash", + ) + if keys: + async with client.pipeline(transaction=False) as read_pipeline: + for key in keys: + read_pipeline.hmget(key, *_SCOPE_FIELDS) + records = await read_pipeline.execute() + writes: list[tuple[Any, dict[str, str]]] = [] + for key, values in zip(keys, records, strict=True): + tokens: dict[str, str] = {} + for name, raw_value in zip(_SCOPE_FIELDS, values, strict=True): + if raw_value is None: + continue + value = raw_value.decode() if isinstance(raw_value, bytes) else str(raw_value) + if value: + tokens[_scope_token_field(name)] = _scope_token(name, value) + if tokens: + writes.append((key, tokens)) + if not sample_token: + sample_field, sample_token = next(iter(tokens.items())) + if writes: + async with client.pipeline(transaction=False) as write_pipeline: + for key, mapping in writes: + write_pipeline.hset(key, mapping=mapping) + await write_pipeline.execute() + processed += len(keys) + await self._set_scope_migration_state( + _SCOPE_MIGRATION_BACKFILL, + cursor=int(cursor), + sample_field=sample_field, + sample_token=sample_token, + ) + if int(cursor) == 0: + break + + await self._set_scope_migration_state( + _SCOPE_MIGRATION_REINDEXING, + cursor=0, + sample_field=sample_field, + sample_token=sample_token, + ) + await self.redis_index.create(overwrite=True, drop=False) + await self._wait_for_scope_index(sample_field, sample_token) + await self._set_scope_migration_state( + _SCOPE_MIGRATION_COMPLETE, + cursor=0, + sample_field=sample_field, + sample_token=sample_token, + ) + finally: + await client.eval( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", + 1, + lock_key, + lock_value, + ) + + async def _wait_for_scope_index(self, sample_field: str, sample_token: str) -> None: deadline = asyncio.get_running_loop().time() + _SCOPE_MIGRATION_TIMEOUT_SECONDS while True: info = await self.redis_index.info() - if int(info.get("num_docs", 0)) >= migrated_documents: - return - if asyncio.get_running_loop().time() >= deadline: - raise RuntimeError( - f"Redis context index rebuild timed out after migrating {migrated_documents} stored hashes." + percent_indexed = float(info.get("percent_indexed", 1.0)) + indexing = int(info.get("indexing", 0)) + if percent_indexed >= 1.0 and indexing == 0: + if not sample_token: + return + results = await self.redis_index.query( + FilterQuery( + FilterExpression(f"@{sample_field}:{{{sample_token}}}"), + return_fields=[], + num_results=1, + ) ) + if results: + return + if asyncio.get_running_loop().time() >= deadline: + raise RuntimeError("Redis context index rebuild timed out before scope tokens became queryable.") await asyncio.sleep(0.05) async def _add( diff --git a/python/packages/redis/tests/test_context_provider_edges.py b/python/packages/redis/tests/test_context_provider_edges.py index e8e343d7669..8accdecfc57 100644 --- a/python/packages/redis/tests/test_context_provider_edges.py +++ b/python/packages/redis/tests/test_context_provider_edges.py @@ -16,8 +16,11 @@ from agent_framework_redis._context_provider import ( _SCOPE_FIELDS, + _SCOPE_MIGRATION_BACKFILL, + _SCOPE_MIGRATION_COMPLETE, _SCOPE_TAG_SEPARATOR, RedisContextProvider, + _scope_migration_key, _scope_token, _scope_token_field, ) @@ -26,6 +29,10 @@ @pytest.fixture def mock_index() -> AsyncMock: index = AsyncMock() + client = MagicMock() + client.hset = AsyncMock() + client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b""]) + index.client = client index.create = AsyncMock() index.exists = AsyncMock(return_value=False) index.load = AsyncMock() @@ -137,7 +144,7 @@ async def test_ensure_index_validates_existing_schema_before_create( await provider._ensure_index() validate_schema.assert_awaited_once() - mock_index.create.assert_awaited_once_with(overwrite=False, drop=False) + mock_index.create.assert_not_awaited() assert provider._index_initialized is True @@ -151,7 +158,7 @@ async def test_ensure_index_rebuilds_and_migrates_when_explicitly_enabled( with patch.object(provider, "_migrate_scope_tokens", AsyncMock()) as migrate: await provider._ensure_index() - mock_index.create.assert_awaited_once_with(overwrite=True, drop=False) + mock_index.create.assert_not_awaited() migrate.assert_awaited_once() assert provider._index_initialized is True @@ -173,6 +180,20 @@ async def test_validate_schema_compatibility_raises_for_significant_mismatch( await provider._validate_schema_compatibility() +async def test_validate_schema_compatibility_rejects_incomplete_migration( + mock_index: AsyncMock, + patch_index: MagicMock, +) -> None: + provider = RedisContextProvider(source_id="ctx", user_id="user-1") + existing_index = AsyncMock() + existing_index.schema.to_dict = MagicMock(return_value=provider.schema_dict) + patch_index.from_existing = AsyncMock(return_value=existing_index) + mock_index.client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"12", b"", b""]) + + with pytest.raises(ValueError, match="migration is incomplete"): + await provider._validate_schema_compatibility() + + async def test_add_requires_content_field( patch_index: MagicMock, # noqa: ARG001 ) -> None: @@ -245,37 +266,36 @@ async def test_migrate_scope_tokens_backfills_preserved_hashes( patch_index: MagicMock, # noqa: ARG001 ) -> None: client = MagicMock() - - async def scan_iter(**kwargs: Any) -> AsyncIterator[bytes]: - assert kwargs == {"_type": "hash"} - yield b"context:one" - yield b"other:two" - - client.scan_iter = scan_iter - client.hgetall = AsyncMock( - side_effect=[ - { - b"application_id": b"App", - b"agent_id": b"Agent", - b"user_id": b"Alice", - b"thread_id": b"Thread", - b"conversation_id": b"Conversation", - } - ] - ) - pipeline = MagicMock() - pipeline.__aenter__ = AsyncMock(return_value=pipeline) - pipeline.__aexit__ = AsyncMock(return_value=False) - pipeline.hset = MagicMock() - pipeline.execute = AsyncMock(return_value=[1]) - client.pipeline.return_value = pipeline + client.set = AsyncMock(return_value=True) + client.eval = AsyncMock(return_value=1) + client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"7", b"", b""]) + client.hset = AsyncMock() + client.scan = AsyncMock(return_value=(0, [b"context:one"])) + read_pipeline = MagicMock() + read_pipeline.__aenter__ = AsyncMock(return_value=read_pipeline) + read_pipeline.__aexit__ = AsyncMock(return_value=False) + read_pipeline.hmget = MagicMock() + read_pipeline.execute = AsyncMock(return_value=[[b"App", b"Agent", b"Alice", b"Thread", b"Conversation"]]) + write_pipeline = MagicMock() + write_pipeline.__aenter__ = AsyncMock(return_value=write_pipeline) + write_pipeline.__aexit__ = AsyncMock(return_value=False) + write_pipeline.hset = MagicMock() + write_pipeline.execute = AsyncMock(return_value=[1]) + client.pipeline.side_effect = [read_pipeline, write_pipeline] mock_index.client = client - mock_index.info = AsyncMock(return_value={"num_docs": 1}) + mock_index.key_separator = ":" + mock_index.info = AsyncMock(return_value={"percent_indexed": 1, "indexing": 0}) + mock_index.query = AsyncMock(return_value=[{"content": "migrated"}]) provider = RedisContextProvider(source_id="ctx", user_id="user-1") + existing_index = AsyncMock() + existing_index.schema.to_dict = MagicMock(return_value={"index": {}, "fields": []}) + patch_index.from_existing = AsyncMock(return_value=existing_index) await provider._migrate_scope_tokens() - mapping = pipeline.hset.call_args.kwargs["mapping"] + client.scan.assert_awaited_once_with(cursor=7, match="context:*", count=100, _type="hash") + read_pipeline.hmget.assert_called_once_with(b"context:one", *_SCOPE_FIELDS) + mapping = write_pipeline.hset.call_args.kwargs["mapping"] for name, value in { "application_id": "App", "agent_id": "Agent", @@ -284,7 +304,57 @@ async def scan_iter(**kwargs: Any) -> AsyncIterator[bytes]: "conversation_id": "Conversation", }.items(): assert mapping[_scope_token_field(name)] == _scope_token(name, value) - mock_index.info.assert_awaited_once() + mock_index.create.assert_awaited_once_with(overwrite=True, drop=False) + mock_index.query.assert_awaited_once() + complete_state = client.hset.await_args_list[-1].kwargs["mapping"] + assert complete_state["state"] == _SCOPE_MIGRATION_COMPLETE + client.eval.assert_awaited_once() + + +async def test_migrate_scope_tokens_pauses_with_resumable_cursor( + mock_index: AsyncMock, + patch_index: MagicMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = MagicMock() + client.set = AsyncMock(return_value=True) + client.eval = AsyncMock(return_value=1) + client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"9", b"", b""]) + client.hset = AsyncMock() + client.scan = AsyncMock(return_value=(11, [b"context:one"])) + read_pipeline = MagicMock() + read_pipeline.__aenter__ = AsyncMock(return_value=read_pipeline) + read_pipeline.__aexit__ = AsyncMock(return_value=False) + read_pipeline.hmget = MagicMock() + read_pipeline.execute = AsyncMock(return_value=[[None, None, b"Alice", None, None]]) + write_pipeline = MagicMock() + write_pipeline.__aenter__ = AsyncMock(return_value=write_pipeline) + write_pipeline.__aexit__ = AsyncMock(return_value=False) + write_pipeline.hset = MagicMock() + write_pipeline.execute = AsyncMock(return_value=[1]) + client.pipeline.side_effect = [read_pipeline, write_pipeline] + mock_index.client = client + mock_index.key_separator = ":" + provider = RedisContextProvider(source_id="ctx", user_id="user-1") + existing_index = AsyncMock() + existing_index.schema.to_dict = MagicMock(return_value={"index": {}, "fields": []}) + patch_index.from_existing = AsyncMock(return_value=existing_index) + monkeypatch.setattr("agent_framework_redis._context_provider._SCOPE_MIGRATION_MAX_RECORDS", 1) + + with pytest.raises(RuntimeError, match="resume"): + await provider._migrate_scope_tokens() + + persisted = client.hset.await_args_list[-1].kwargs["mapping"] + assert persisted["state"] == _SCOPE_MIGRATION_BACKFILL + assert persisted["cursor"] == "11" + mock_index.create.assert_not_awaited() + client.eval.assert_awaited_once() + + +def test_scope_migration_marker_is_outside_index_namespace() -> None: + marker = _scope_migration_key("context", "context") + + assert not marker.startswith("context:") @pytest.mark.parametrize("dtype, expected_size", [("float16", 4), ("float32", 8), ("float64", 16), ("bfloat16", 4)]) diff --git a/python/packages/redis/tests/test_context_provider_integration.py b/python/packages/redis/tests/test_context_provider_integration.py index a511c7447ab..15966f15a63 100644 --- a/python/packages/redis/tests/test_context_provider_integration.py +++ b/python/packages/redis/tests/test_context_provider_integration.py @@ -16,7 +16,12 @@ from redisvl.utils.vectorize import BaseVectorizer from agent_framework_redis import RedisContextProvider -from agent_framework_redis._context_provider import _scope_token_field +from agent_framework_redis._context_provider import ( + _SCOPE_MIGRATION_BACKFILL, + _SCOPE_MIGRATION_COMPLETE, + _scope_migration_key, + _scope_token_field, +) @pytest.fixture @@ -35,9 +40,11 @@ async def live_redis() -> AsyncIterator[tuple[str, Redis, str, str]]: encoded_name = index_name.encode() if encoded_name in indexes or index_name in indexes: await client.execute_command("FT.DROPINDEX", index_name) - keys = [key async for key in client.scan_iter(match=f"{prefix}*")] + keys = [key async for key in client.scan_iter(match=f"{prefix.rstrip(':')}*")] if keys: await client.delete(*keys) + marker_key = _scope_migration_key(index_name, prefix) + await client.delete(marker_key, f"{marker_key}:lock") await client.aclose() @@ -203,6 +210,11 @@ async def test_live_legacy_index_requires_explicit_rebuild_and_preserves_exact_s "conversation_id": "session", } ]) + adjacent_key = f"{prefix.rstrip(':')}-other:foreign" + await cast( + Awaitable[int], + client.hset(adjacent_key, mapping={"content": "foreign", "user_id": "Alice"}), + ) incompatible = _provider(url, index_name, prefix, user_id="Alice") with pytest.raises(ValueError, match="coordinated migration"): @@ -212,3 +224,12 @@ async def test_live_legacy_index_requires_explicit_rebuild_and_preserves_exact_s assert len(await migrated._redis_search("migration marker")) == 1 assert await _provider(url, index_name, prefix, user_id="alice")._redis_search("migration marker") == [] assert await _provider(url, index_name, prefix, user_id="Alice,other")._redis_search("migration marker") == [] + adjacent_record = await cast(Awaitable[dict[bytes, bytes]], client.hgetall(adjacent_key)) + assert _scope_token_field("user_id").encode() not in adjacent_record + + marker_key = _scope_migration_key(index_name, prefix) + marker = await cast(Awaitable[dict[bytes, bytes]], client.hgetall(marker_key)) + assert marker[b"state"] == _SCOPE_MIGRATION_COMPLETE.encode() + await cast(Awaitable[int], client.hset(marker_key, mapping={"state": _SCOPE_MIGRATION_BACKFILL})) + with pytest.raises(ValueError, match="migration is incomplete"): + await _provider(url, index_name, prefix, user_id="Alice")._redis_search("migration marker") diff --git a/python/packages/redis/tests/test_context_provider_sdk.py b/python/packages/redis/tests/test_context_provider_sdk.py index a0f496bb380..3c5d5fdb3a7 100644 --- a/python/packages/redis/tests/test_context_provider_sdk.py +++ b/python/packages/redis/tests/test_context_provider_sdk.py @@ -67,6 +67,8 @@ def _provider(client: redis.Redis, *, borrowed: bool, **kwargs: Any) -> RedisCon # Borrowed indexes use RedisVL's default validate_on_load=False. index = AsyncSearchIndex.from_dict(provider.schema_dict, redis_client=client) provider = RedisContextProvider(redis_index=index, **kwargs) + else: + provider.redis_index._redis_client = client # pyright: ignore[reportPrivateUsage] provider.redis_index.exists = AsyncMock(return_value=False) provider.redis_index.create = AsyncMock() return provider @@ -134,8 +136,8 @@ async def test_after_run_encodes_messages_with_optional_fields_omitted( if not all_scopes: assert "application_id" not in record assert "agent_id" not in record - assert len(_TransportConnection.packets) == 1 - assert b"HSET" in _TransportConnection.packets[0] + assert len(_TransportConnection.packets) == 2 + assert all(b"HSET" in packet for packet in _TransportConnection.packets) assert message.to_dict() == before @@ -179,7 +181,7 @@ async def test_borrowed_index_omits_null_values_without_changing_nonnull_values_ assert record["author_name"] == "" assert record["mime_type"] == 0 assert record["embedding"] == b"\x00\x00\x80?" - assert len(_TransportConnection.packets) == 1 + assert len(_TransportConnection.packets) == 2 assert document == before @@ -196,7 +198,7 @@ async def test_borrowed_index_omits_unfilled_vector_placeholder( assert "embedding" not in record assert "conversation_id" not in record assert "thread_id" not in record - assert len(_TransportConnection.packets) == 1 + assert len(_TransportConnection.packets) == 2 assert document == {"role": "user", "content": "Remember tea"} @@ -218,5 +220,5 @@ async def test_borrowed_index_retains_generated_vector_buffer( vectorizer.aembed_many.assert_awaited_once_with(["Remember tea"], batch_size=1) assert write.call_args.args[2]["embedding"] == array_to_buffer([1.0, 0.0], dtype="float32") - assert len(_TransportConnection.packets) == 1 + assert len(_TransportConnection.packets) == 2 assert document == before diff --git a/python/packages/redis/tests/test_providers.py b/python/packages/redis/tests/test_providers.py index 5f554fb5beb..e0fb07430ae 100644 --- a/python/packages/redis/tests/test_providers.py +++ b/python/packages/redis/tests/test_providers.py @@ -36,6 +36,10 @@ async def test_empty_history_save_marks_redis_used(mock_redis_client: MagicMock) @pytest.fixture def mock_index() -> AsyncMock: idx = AsyncMock() + client = MagicMock() + client.hset = AsyncMock() + client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b""]) + idx.client = client idx.create = AsyncMock() idx.load = AsyncMock() idx.query = AsyncMock(return_value=[]) From 010aaa2d49fb950c53af99651d41c2687b8fcd5e Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Wed, 21 Oct 2026 14:23:52 +0200 Subject: [PATCH 3/6] Python: trim Redis migration documentation --- python/packages/redis/AGENTS.md | 4 ---- python/packages/redis/README.md | 39 +++++---------------------------- 2 files changed, 5 insertions(+), 38 deletions(-) diff --git a/python/packages/redis/AGENTS.md b/python/packages/redis/AGENTS.md index d1921183dd4..00aff40a2f6 100644 --- a/python/packages/redis/AGENTS.md +++ b/python/packages/redis/AGENTS.md @@ -10,10 +10,6 @@ Redis-based storage for agent threads and context. an owned client from a URL or credential provider, or borrow a standalone async Redis client configured with `decode_responses=True`; borrowed clients remain caller-owned. - **`RedisContextProvider`** - Context provider with Redis-backed retrieval - Scope identifiers are stored as raw metadata but indexed through versioned, - domain-separated tokens in case-sensitive TAG fields. Existing legacy - indexes fail closed unless `overwrite_index=True` is explicitly used for a - coordinated in-place rebuild and token backfill. - **`RedisSettings`** - TypedDict connection settings for vector stores, resolved with core `load_settings` from explicit URL overrides, an optional .env file, or `REDIS_URL`. URLs use `SecretString` to mask credentials. - **`RedisCollection` / `RedisStore`** - Experimental generic vector storage and search over native HASH and JSON diff --git a/python/packages/redis/README.md b/python/packages/redis/README.md index 776eb66374a..9f8481a0d07 100644 --- a/python/packages/redis/README.md +++ b/python/packages/redis/README.md @@ -53,6 +53,11 @@ JSON collections additionally require RedisJSON. Redis Search indexes require logical database 0, so vector-store URLs and clients must select database 0 (the default). +When upgrading an existing `RedisContextProvider` index, stop older writers and +initialize once with `overwrite_index=True` to migrate its scope schema while +preserving stored hashes. Repeat if prompted to resume a bounded migration, +then remove the option; normal providers reject incomplete migrations. + Choose `storage_type="hash"` for non-null records with binary vector storage, or `"json"` for nullable fields/vectors and native JSON data. HASH rejects `None`, including `vector=None`; JSON preserves explicit null. Both support @@ -148,40 +153,6 @@ and session into the corresponding scoped key using an application-owned migration process. After verifying the copied history, remove legacy keys separately according to the application's retention policy. -## Isolate retrieved memories - -`RedisContextProvider` stores exact application, agent, user, and session -identifiers as metadata, but indexes only deterministic versioned tokens for -those values. Scope comparisons therefore preserve case and treat punctuation -as ordinary identifier content. Values supplied inside documents cannot -override the provider's configured scopes. - -Indexes created by earlier releases use a different scope schema. A provider -using the current schema fails closed when it encounters such an index. During -a coordinated deployment, stop older writers and rebuild the index once while -preserving its hashes: - -```python -context_provider = RedisContextProvider( - redis_url="redis://localhost:6379", - index_name="context", - prefix="context", - application_id="support-app", - user_id="user-123", - overwrite_index=True, -) -``` - -The migration records resumable progress outside the indexed key namespace, -backfills the current scope tokens from the exact identifiers stored in each -matching Redis hash, and only then replaces the index definition. Each pass is -bounded; if a large namespace needs another pass, initialization raises an -actionable error and the next provider with `overwrite_index=True` resumes from -the saved cursor. Normal providers reject an incomplete migration. Existing -memory content remains in place. After the migration completes, remove -`overwrite_index=True`. Review or remove records whose historical scope -metadata is not authoritative before migration. - ## Store and search documents This example uses precomputed vectors, so no embedding service is needed. From a211770106a96dd7f7683f9b7d925ad15e187b1e Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Thu, 22 Oct 2026 13:31:32 +0200 Subject: [PATCH 4/6] Python: harden Redis scope migration control --- .../_context_provider.py | 323 ++++++++++++++---- .../tests/test_context_provider_edges.py | 173 ++++++++-- .../test_context_provider_integration.py | 65 +++- python/packages/redis/tests/test_providers.py | 2 +- 4 files changed, 466 insertions(+), 97 deletions(-) diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index 5e3bf9ac59b..fc22d7bfc5f 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -9,6 +9,8 @@ from __future__ import annotations import asyncio +import base64 +import contextlib import hashlib import json import secrets @@ -53,6 +55,9 @@ _SCOPE_MIGRATION_BACKFILL = "backfill" _SCOPE_MIGRATION_REINDEXING = "reindexing" _SCOPE_MIGRATION_COMPLETE = "complete" +_SCOPE_MIGRATION_PAUSED_MESSAGE = ( + "Redis context scope migration paused after a bounded backfill pass. Rerun with overwrite_index=True to resume." +) _TAG_DEFAULTS = {"separator": ",", "case_sensitive": False, "withsuffixtrie": False} _TEXT_DEFAULTS = {"weight": 1.0, "no_stem": False} @@ -83,6 +88,18 @@ def _scope_key_namespace(prefix: str, key_separator: str) -> str: return f"{normalized_prefix}{key_separator}" +def _redis_scan_literal(value: str) -> str: + return "".join(f"\\{char}" if char in "\\*?[]" else char for char in value) + + +def _encode_redis_keys(keys: list[Any]) -> str: + return json.dumps([base64.b64encode(key if isinstance(key, bytes) else str(key).encode()).decode() for key in keys]) + + +def _decode_redis_keys(value: str) -> list[bytes]: + return [base64.b64decode(item) for item in json.loads(value or "[]")] + + def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]: def significant_index(index: dict[str, Any]) -> dict[str, Any]: return {key: index.get(key) for key in ("name", "prefix", "key_separator", "storage_type")} @@ -387,11 +404,27 @@ async def _scope_migration_state(self) -> dict[str, str]: "cursor", "sample_field", "sample_token", + "pending_cursor", + "pending_keys", ) decoded = [ value.decode() if isinstance(value, bytes) else str(value) if value is not None else "" for value in values ] - return dict(zip(("version", "state", "cursor", "sample_field", "sample_token"), decoded, strict=True)) + return dict( + zip( + ( + "version", + "state", + "cursor", + "sample_field", + "sample_token", + "pending_cursor", + "pending_keys", + ), + decoded, + strict=True, + ) + ) async def _set_scope_migration_state( self, @@ -400,20 +433,211 @@ async def _set_scope_migration_state( cursor: int, sample_field: str = "", sample_token: str = "", + pending_cursor: int = 0, + pending_keys: list[Any] | None = None, + lock_key: str | None = None, + lock_value: str | None = None, ) -> None: client = self.redis_index.client if client is None: raise RuntimeError("Redis index did not initialize a client for scope migration state.") - await client.hset( - _scope_migration_key(self.index_name, self.prefix), - mapping={ - "version": str(_SCOPE_SCHEMA_VERSION), - "state": state, - "cursor": str(cursor), - "sample_field": sample_field, - "sample_token": sample_token, - }, + mapping = { + "version": str(_SCOPE_SCHEMA_VERSION), + "state": state, + "cursor": str(cursor), + "sample_field": sample_field, + "sample_token": sample_token, + "pending_cursor": str(pending_cursor), + "pending_keys": _encode_redis_keys(pending_keys or []), + } + marker_key = _scope_migration_key(self.index_name, self.prefix) + if lock_key is None or lock_value is None: + await client.hset(marker_key, mapping=mapping) + return + arguments = [lock_value] + for name, value in mapping.items(): + arguments.extend((name, value)) + result = await client.eval( + "if redis.call('get', KEYS[1]) ~= ARGV[1] then return 0 end " + "for i = 2, #ARGV, 2 do redis.call('hset', KEYS[2], ARGV[i], ARGV[i + 1]) end return 1", + 2, + lock_key, + marker_key, + *arguments, + ) + if not result: + raise RuntimeError("Redis context scope migration lost its distributed lock.") + + async def _assert_scope_migration_lock(self, lock_key: str, lock_value: str) -> None: + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index did not initialize a client for scope migration.") + current = await client.get(lock_key) + current_value = current.decode() if isinstance(current, bytes) else current + if current_value != lock_value: + raise RuntimeError("Redis context scope migration lost its distributed lock.") + + async def _renew_scope_migration_lock(self, lock_key: str, lock_value: str) -> None: + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index did not initialize a client for scope migration.") + while True: + await asyncio.sleep(max(_SCOPE_MIGRATION_LOCK_SECONDS / 3, 0.1)) + renewed = await client.eval( + "if redis.call('get', KEYS[1]) == ARGV[1] " + "then return redis.call('expire', KEYS[1], ARGV[2]) else return 0 end", + 1, + lock_key, + lock_value, + _SCOPE_MIGRATION_LOCK_SECONDS, + ) + if not renewed: + raise RuntimeError("Redis context scope migration lost its distributed lock.") + + async def _process_scope_migration_keys( + self, + keys: list[bytes], + *, + sample_field: str, + sample_token: str, + lock_key: str, + lock_value: str, + ) -> tuple[str, str]: + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index did not initialize a client for scope migration.") + async with client.pipeline(transaction=False) as read_pipeline: + for key in keys: + read_pipeline.hmget(key, *_SCOPE_FIELDS) + records = await read_pipeline.execute() + writes: list[tuple[bytes, dict[str, str]]] = [] + for key, values in zip(keys, records, strict=True): + tokens: dict[str, str] = {} + for name, raw_value in zip(_SCOPE_FIELDS, values, strict=True): + if raw_value is None: + continue + value = raw_value.decode() if isinstance(raw_value, bytes) else str(raw_value) + if value: + tokens[_scope_token_field(name)] = _scope_token(name, value) + if tokens: + writes.append((key, tokens)) + if not sample_token: + sample_field, sample_token = next(iter(tokens.items())) + if writes: + await self._assert_scope_migration_lock(lock_key, lock_value) + async with client.pipeline(transaction=False) as write_pipeline: + for key, mapping in writes: + write_pipeline.hset(key, mapping=mapping) + await write_pipeline.execute() + return sample_field, sample_token + + async def _backfill_scope_tokens( + self, + state: dict[str, str], + *, + lock_key: str, + lock_value: str, + ) -> tuple[str, str]: + client = self.redis_index.client + if client is None: + raise RuntimeError("Redis index did not initialize a client for scope migration.") + state_is_current = state.get("version") == str(_SCOPE_SCHEMA_VERSION) + cursor = ( + int(state.get("cursor") or 0) if state_is_current and state.get("state") == _SCOPE_MIGRATION_BACKFILL else 0 ) + sample_field = state.get("sample_field", "") if state_is_current else "" + sample_token = state.get("sample_token", "") if state_is_current else "" + pending_cursor = int(state.get("pending_cursor") or 0) if state_is_current else 0 + pending_keys = _decode_redis_keys(state.get("pending_keys", "")) if state_is_current else [] + namespace = _scope_key_namespace(self.prefix, self.redis_index.key_separator) + encoded_namespace = namespace.encode() + match_pattern = f"{_redis_scan_literal(namespace)}*" + processed = 0 + + await self._set_scope_migration_state( + _SCOPE_MIGRATION_BACKFILL, + cursor=cursor, + sample_field=sample_field, + sample_token=sample_token, + pending_cursor=pending_cursor, + pending_keys=pending_keys, + lock_key=lock_key, + lock_value=lock_value, + ) + while True: + await self._assert_scope_migration_lock(lock_key, lock_value) + remaining = _SCOPE_MIGRATION_MAX_RECORDS - processed + if remaining <= 0: + raise RuntimeError(_SCOPE_MIGRATION_PAUSED_MESSAGE) + + if not pending_keys: + scan_cursor = cursor + next_cursor, scanned_keys = await client.scan( + cursor=scan_cursor, + match=match_pattern, + count=min(_SCOPE_MIGRATION_BATCH_SIZE, remaining), + _type="hash", + ) + pending_keys = [ + key if isinstance(key, bytes) else str(key).encode() + for key in scanned_keys + if (key if isinstance(key, bytes) else str(key).encode()).startswith(encoded_namespace) + ] + pending_cursor = int(next_cursor) + await self._set_scope_migration_state( + _SCOPE_MIGRATION_BACKFILL, + cursor=scan_cursor, + sample_field=sample_field, + sample_token=sample_token, + pending_cursor=pending_cursor, + pending_keys=pending_keys, + lock_key=lock_key, + lock_value=lock_value, + ) + if not pending_keys: + cursor = pending_cursor + pending_cursor = 0 + await self._set_scope_migration_state( + _SCOPE_MIGRATION_BACKFILL, + cursor=cursor, + sample_field=sample_field, + sample_token=sample_token, + lock_key=lock_key, + lock_value=lock_value, + ) + if cursor == 0: + return sample_field, sample_token + continue + + batch = pending_keys[:remaining] + sample_field, sample_token = await self._process_scope_migration_keys( + batch, + sample_field=sample_field, + sample_token=sample_token, + lock_key=lock_key, + lock_value=lock_value, + ) + processed += len(batch) + pending_keys = pending_keys[len(batch) :] + if not pending_keys: + cursor = pending_cursor + pending_cursor = 0 + await self._set_scope_migration_state( + _SCOPE_MIGRATION_BACKFILL, + cursor=cursor, + sample_field=sample_field, + sample_token=sample_token, + pending_cursor=pending_cursor, + pending_keys=pending_keys, + lock_key=lock_key, + lock_value=lock_value, + ) + if pending_keys: + raise RuntimeError(_SCOPE_MIGRATION_PAUSED_MESSAGE) + if cursor == 0: + return sample_field, sample_token + if processed >= _SCOPE_MIGRATION_MAX_RECORDS: + raise RuntimeError(_SCOPE_MIGRATION_PAUSED_MESSAGE) async def _migrate_scope_tokens(self) -> None: """Resume scope-token backfill, rebuild the index, and mark the migration complete.""" @@ -426,88 +650,44 @@ async def _migrate_scope_tokens(self) -> None: if not await client.set(lock_key, lock_value, nx=True, ex=_SCOPE_MIGRATION_LOCK_SECONDS): raise RuntimeError("Another Redis context scope migration is already running.") + renewal_task = asyncio.create_task(self._renew_scope_migration_lock(lock_key, lock_value)) try: state = await self._scope_migration_state() existing_index = await self._existing_index() schema_is_current = _schema_signature(existing_index.schema.to_dict()) == _schema_signature( self.schema_dict ) - if state.get("state") == _SCOPE_MIGRATION_COMPLETE and schema_is_current: + state_is_current = state.get("version") == str(_SCOPE_SCHEMA_VERSION) + if state_is_current and state.get("state") == _SCOPE_MIGRATION_COMPLETE and schema_is_current: return - if state.get("state") == _SCOPE_MIGRATION_REINDEXING and schema_is_current: + if state_is_current and state.get("state") == _SCOPE_MIGRATION_REINDEXING and schema_is_current: await self._wait_for_scope_index(state.get("sample_field", ""), state.get("sample_token", "")) await self._set_scope_migration_state( _SCOPE_MIGRATION_COMPLETE, cursor=0, sample_field=state.get("sample_field", ""), sample_token=state.get("sample_token", ""), + lock_key=lock_key, + lock_value=lock_value, ) return - namespace = _scope_key_namespace(self.prefix, self.redis_index.key_separator) - cursor = int(state.get("cursor") or 0) if state.get("state") == _SCOPE_MIGRATION_BACKFILL else 0 - sample_field = state.get("sample_field", "") - sample_token = state.get("sample_token", "") - await self._set_scope_migration_state( - _SCOPE_MIGRATION_BACKFILL, - cursor=cursor, - sample_field=sample_field, - sample_token=sample_token, - ) - - deadline = asyncio.get_running_loop().time() + _SCOPE_MIGRATION_TIMEOUT_SECONDS - processed = 0 - while True: - if processed >= _SCOPE_MIGRATION_MAX_RECORDS or asyncio.get_running_loop().time() >= deadline: - raise RuntimeError( - "Redis context scope migration paused after a bounded backfill pass. " - "Rerun with overwrite_index=True to resume." - ) - cursor, keys = await client.scan( - cursor=cursor, - match=f"{namespace}*", - count=min(_SCOPE_MIGRATION_BATCH_SIZE, _SCOPE_MIGRATION_MAX_RECORDS - processed), - _type="hash", - ) - if keys: - async with client.pipeline(transaction=False) as read_pipeline: - for key in keys: - read_pipeline.hmget(key, *_SCOPE_FIELDS) - records = await read_pipeline.execute() - writes: list[tuple[Any, dict[str, str]]] = [] - for key, values in zip(keys, records, strict=True): - tokens: dict[str, str] = {} - for name, raw_value in zip(_SCOPE_FIELDS, values, strict=True): - if raw_value is None: - continue - value = raw_value.decode() if isinstance(raw_value, bytes) else str(raw_value) - if value: - tokens[_scope_token_field(name)] = _scope_token(name, value) - if tokens: - writes.append((key, tokens)) - if not sample_token: - sample_field, sample_token = next(iter(tokens.items())) - if writes: - async with client.pipeline(transaction=False) as write_pipeline: - for key, mapping in writes: - write_pipeline.hset(key, mapping=mapping) - await write_pipeline.execute() - processed += len(keys) - await self._set_scope_migration_state( - _SCOPE_MIGRATION_BACKFILL, - cursor=int(cursor), - sample_field=sample_field, - sample_token=sample_token, + try: + sample_field, sample_token = await asyncio.wait_for( + self._backfill_scope_tokens(state, lock_key=lock_key, lock_value=lock_value), + timeout=_SCOPE_MIGRATION_TIMEOUT_SECONDS, ) - if int(cursor) == 0: - break - + except TimeoutError as exc: + raise RuntimeError(_SCOPE_MIGRATION_PAUSED_MESSAGE) from exc await self._set_scope_migration_state( _SCOPE_MIGRATION_REINDEXING, cursor=0, sample_field=sample_field, sample_token=sample_token, + lock_key=lock_key, + lock_value=lock_value, ) + await self._assert_scope_migration_lock(lock_key, lock_value) await self.redis_index.create(overwrite=True, drop=False) await self._wait_for_scope_index(sample_field, sample_token) await self._set_scope_migration_state( @@ -515,8 +695,13 @@ async def _migrate_scope_tokens(self) -> None: cursor=0, sample_field=sample_field, sample_token=sample_token, + lock_key=lock_key, + lock_value=lock_value, ) finally: + renewal_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await renewal_task await client.eval( "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", 1, diff --git a/python/packages/redis/tests/test_context_provider_edges.py b/python/packages/redis/tests/test_context_provider_edges.py index 8accdecfc57..6ad54e2c463 100644 --- a/python/packages/redis/tests/test_context_provider_edges.py +++ b/python/packages/redis/tests/test_context_provider_edges.py @@ -18,8 +18,11 @@ _SCOPE_FIELDS, _SCOPE_MIGRATION_BACKFILL, _SCOPE_MIGRATION_COMPLETE, + _SCOPE_MIGRATION_LOCK_SECONDS, _SCOPE_TAG_SEPARATOR, RedisContextProvider, + _encode_redis_keys, + _redis_scan_literal, _scope_migration_key, _scope_token, _scope_token_field, @@ -31,7 +34,7 @@ def mock_index() -> AsyncMock: index = AsyncMock() client = MagicMock() client.hset = AsyncMock() - client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b""]) + client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b"", b"0", b"[]"]) index.client = client index.create = AsyncMock() index.exists = AsyncMock(return_value=False) @@ -188,7 +191,7 @@ async def test_validate_schema_compatibility_rejects_incomplete_migration( existing_index = AsyncMock() existing_index.schema.to_dict = MagicMock(return_value=provider.schema_dict) patch_index.from_existing = AsyncMock(return_value=existing_index) - mock_index.client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"12", b"", b""]) + mock_index.client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"12", b"", b"", b"0", b"[]"]) with pytest.raises(ValueError, match="migration is incomplete"): await provider._validate_schema_compatibility() @@ -263,12 +266,14 @@ async def test_add_vectorizes_documents_and_applies_defaults( async def test_migrate_scope_tokens_backfills_preserved_hashes( mock_index: AsyncMock, - patch_index: MagicMock, # noqa: ARG001 + patch_index: MagicMock, + monkeypatch: pytest.MonkeyPatch, ) -> None: client = MagicMock() client.set = AsyncMock(return_value=True) client.eval = AsyncMock(return_value=1) - client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"7", b"", b""]) + client.get = AsyncMock(return_value=b"lock") + client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"7", b"", b"", b"0", b"[]"]) client.hset = AsyncMock() client.scan = AsyncMock(return_value=(0, [b"context:one"])) read_pipeline = MagicMock() @@ -290,8 +295,10 @@ async def test_migrate_scope_tokens_backfills_preserved_hashes( existing_index = AsyncMock() existing_index.schema.to_dict = MagicMock(return_value={"index": {}, "fields": []}) patch_index.from_existing = AsyncMock(return_value=existing_index) + monkeypatch.setattr("agent_framework_redis._context_provider.secrets.token_hex", lambda _: "lock") - await provider._migrate_scope_tokens() + with patch.object(provider, "_set_scope_migration_state", wraps=provider._set_scope_migration_state) as set_state: + await provider._migrate_scope_tokens() client.scan.assert_awaited_once_with(cursor=7, match="context:*", count=100, _type="hash") read_pipeline.hmget.assert_called_once_with(b"context:one", *_SCOPE_FIELDS) @@ -306,22 +313,18 @@ async def test_migrate_scope_tokens_backfills_preserved_hashes( assert mapping[_scope_token_field(name)] == _scope_token(name, value) mock_index.create.assert_awaited_once_with(overwrite=True, drop=False) mock_index.query.assert_awaited_once() - complete_state = client.hset.await_args_list[-1].kwargs["mapping"] - assert complete_state["state"] == _SCOPE_MIGRATION_COMPLETE - client.eval.assert_awaited_once() + assert any(call.args[0] == _SCOPE_MIGRATION_COMPLETE for call in set_state.await_args_list) -async def test_migrate_scope_tokens_pauses_with_resumable_cursor( +async def test_backfill_stages_oversized_scan_page_without_advancing_cursor( mock_index: AsyncMock, - patch_index: MagicMock, + patch_index: MagicMock, # noqa: ARG001 monkeypatch: pytest.MonkeyPatch, ) -> None: client = MagicMock() - client.set = AsyncMock(return_value=True) client.eval = AsyncMock(return_value=1) - client.hmget = AsyncMock(return_value=[b"1", b"backfill", b"9", b"", b""]) - client.hset = AsyncMock() - client.scan = AsyncMock(return_value=(11, [b"context:one"])) + client.get = AsyncMock(return_value=b"lock") + client.scan = AsyncMock(return_value=(0, [b"context:one", b"context:two"])) read_pipeline = MagicMock() read_pipeline.__aenter__ = AsyncMock(return_value=read_pipeline) read_pipeline.__aexit__ = AsyncMock(return_value=False) @@ -336,19 +339,139 @@ async def test_migrate_scope_tokens_pauses_with_resumable_cursor( mock_index.client = client mock_index.key_separator = ":" provider = RedisContextProvider(source_id="ctx", user_id="user-1") + monkeypatch.setattr("agent_framework_redis._context_provider._SCOPE_MIGRATION_MAX_RECORDS", 1) + + with ( + patch.object(provider, "_set_scope_migration_state", wraps=provider._set_scope_migration_state) as set_state, + pytest.raises(RuntimeError, match="resume"), + ): + await provider._backfill_scope_tokens( + { + "version": "1", + "state": "backfill", + "cursor": "0", + "sample_field": "", + "sample_token": "", + "pending_cursor": "0", + "pending_keys": "[]", + }, + lock_key="lock-key", + lock_value="lock", + ) + + final_state = set_state.await_args_list[-1] + assert final_state.kwargs["cursor"] == 0 + assert final_state.kwargs["pending_cursor"] == 0 + assert final_state.kwargs["pending_keys"] == [b"context:two"] + read_pipeline.hmget.assert_called_once_with(b"context:one", *_SCOPE_FIELDS) + + +@pytest.mark.parametrize("state_name", [_SCOPE_MIGRATION_COMPLETE, _SCOPE_MIGRATION_BACKFILL]) +async def test_backfill_restarts_wrong_version_state( + state_name: str, + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 +) -> None: + client = MagicMock() + client.eval = AsyncMock(return_value=1) + client.get = AsyncMock(return_value=b"lock") + client.scan = AsyncMock(return_value=(0, [])) + mock_index.client = client + mock_index.key_separator = ":" + provider = RedisContextProvider(source_id="ctx", user_id="user-1") + + await provider._backfill_scope_tokens( + { + "version": "0", + "state": state_name, + "cursor": "99", + "sample_field": "old-field", + "sample_token": "old-token", + "pending_cursor": "88", + "pending_keys": _encode_redis_keys([b"context:skipped"]), + }, + lock_key="lock-key", + lock_value="lock", + ) + + client.scan.assert_awaited_once_with(cursor=0, match="context:*", count=100, _type="hash") + + +async def test_migrate_wrong_version_complete_marker_does_not_short_circuit( + mock_index: AsyncMock, + patch_index: MagicMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = MagicMock() + client.set = AsyncMock(return_value=True) + client.eval = AsyncMock(return_value=1) + client.get = AsyncMock(return_value=b"lock") + client.hmget = AsyncMock(return_value=[b"0", b"complete", b"99", b"old", b"old", b"88", b"[]"]) + client.scan = AsyncMock(return_value=(0, [])) + mock_index.client = client + mock_index.key_separator = ":" + mock_index.info = AsyncMock(return_value={"percent_indexed": 1, "indexing": 0}) + provider = RedisContextProvider(source_id="ctx", user_id="user-1") existing_index = AsyncMock() - existing_index.schema.to_dict = MagicMock(return_value={"index": {}, "fields": []}) + existing_index.schema.to_dict = MagicMock(return_value=provider.schema_dict) patch_index.from_existing = AsyncMock(return_value=existing_index) - monkeypatch.setattr("agent_framework_redis._context_provider._SCOPE_MIGRATION_MAX_RECORDS", 1) + monkeypatch.setattr("agent_framework_redis._context_provider.secrets.token_hex", lambda _: "lock") - with pytest.raises(RuntimeError, match="resume"): - await provider._migrate_scope_tokens() + await provider._migrate_scope_tokens() - persisted = client.hset.await_args_list[-1].kwargs["mapping"] - assert persisted["state"] == _SCOPE_MIGRATION_BACKFILL - assert persisted["cursor"] == "11" - mock_index.create.assert_not_awaited() - client.eval.assert_awaited_once() + client.scan.assert_awaited_once_with(cursor=0, match="context:*", count=100, _type="hash") + mock_index.create.assert_awaited_once_with(overwrite=True, drop=False) + + +async def test_backfill_escapes_scan_pattern_and_rechecks_namespace( + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 +) -> None: + client = MagicMock() + client.eval = AsyncMock(return_value=1) + client.get = AsyncMock(return_value=b"lock") + client.scan = AsyncMock(return_value=(0, [b"ctx[*]:one", b"ctxX:foreign"])) + read_pipeline = MagicMock() + read_pipeline.__aenter__ = AsyncMock(return_value=read_pipeline) + read_pipeline.__aexit__ = AsyncMock(return_value=False) + read_pipeline.hmget = MagicMock() + read_pipeline.execute = AsyncMock(return_value=[[None, None, b"Alice", None, None]]) + write_pipeline = MagicMock() + write_pipeline.__aenter__ = AsyncMock(return_value=write_pipeline) + write_pipeline.__aexit__ = AsyncMock(return_value=False) + write_pipeline.hset = MagicMock() + write_pipeline.execute = AsyncMock(return_value=[1]) + client.pipeline.side_effect = [read_pipeline, write_pipeline] + mock_index.client = client + mock_index.key_separator = ":" + provider = RedisContextProvider(source_id="ctx", prefix="ctx[*]", user_id="user-1") + + await provider._backfill_scope_tokens({}, lock_key="lock-key", lock_value="lock") + + client.scan.assert_awaited_once_with(cursor=0, match=r"ctx\[\*\]:*", count=100, _type="hash") + read_pipeline.hmget.assert_called_once_with(b"ctx[*]:one", *_SCOPE_FIELDS) + + +async def test_migration_lock_renews_only_for_owner( + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = MagicMock() + client.eval = AsyncMock(side_effect=[1, 0]) + mock_index.client = client + provider = RedisContextProvider(source_id="ctx", user_id="user-1") + + async def immediate_sleep(_: float) -> None: + return + + monkeypatch.setattr("agent_framework_redis._context_provider.asyncio.sleep", immediate_sleep) + + with pytest.raises(RuntimeError, match="lost"): + await provider._renew_scope_migration_lock("lock-key", "owner") + + assert client.eval.await_count == 2 + assert client.eval.await_args_list[0].args[-2:] == ("owner", _SCOPE_MIGRATION_LOCK_SECONDS) def test_scope_migration_marker_is_outside_index_namespace() -> None: @@ -357,6 +480,10 @@ def test_scope_migration_marker_is_outside_index_namespace() -> None: assert not marker.startswith("context:") +def test_redis_scan_literal_escapes_glob_metacharacters() -> None: + assert _redis_scan_literal(r"ctx[*]?\:") == r"ctx\[\*\]\?\\:" + + @pytest.mark.parametrize("dtype, expected_size", [("float16", 4), ("float32", 8), ("float64", 16), ("bfloat16", 4)]) async def test_after_run_stores_vectors_in_configured_dtype( dtype: str, diff --git a/python/packages/redis/tests/test_context_provider_integration.py b/python/packages/redis/tests/test_context_provider_integration.py index 15966f15a63..4e60b36dd32 100644 --- a/python/packages/redis/tests/test_context_provider_integration.py +++ b/python/packages/redis/tests/test_context_provider_integration.py @@ -4,6 +4,7 @@ from __future__ import annotations +import asyncio import os from collections.abc import AsyncIterator, Awaitable from typing import Any, cast @@ -31,7 +32,7 @@ async def live_redis() -> AsyncIterator[tuple[str, Redis, str, str]]: pytest.skip("Set REDIS_VECTOR_TEST_URL to a disposable Redis Search server.") suffix = uuid4().hex index_name = f"af-context-test-{suffix}" - prefix = f"af-context-test-{suffix}:" + prefix = f"af-context-test-[*]-{suffix}:" client = Redis.from_url(url) try: yield url, client, index_name, prefix @@ -40,7 +41,12 @@ async def live_redis() -> AsyncIterator[tuple[str, Redis, str, str]]: encoded_name = index_name.encode() if encoded_name in indexes or index_name in indexes: await client.execute_command("FT.DROPINDEX", index_name) - keys = [key async for key in client.scan_iter(match=f"{prefix.rstrip(':')}*")] + encoded_prefix = prefix.encode() + keys = [ + key + async for key in client.scan_iter(_type="hash") + if (key if isinstance(key, bytes) else str(key).encode()).startswith(encoded_prefix) + ] if keys: await client.delete(*keys) marker_key = _scope_migration_key(index_name, prefix) @@ -138,9 +144,15 @@ async def test_live_session_scopes_and_document_overrides_are_provider_owned( ) assert await caller._redis_search("session marker") == [] - keys = [key async for key in client.scan_iter(match=f"{prefix}*")] + encoded_prefix = prefix.encode() + keys = [ + key + async for key in client.scan_iter(_type="hash") + if (key if isinstance(key, bytes) else str(key).encode()).startswith(encoded_prefix) + ] assert len(keys) == 1 - record = await cast(Awaitable[dict[bytes, bytes]], client.hgetall(keys[0])) + key = keys[0].decode() if isinstance(keys[0], bytes) else str(keys[0]) + record = await cast(Awaitable[dict[bytes, bytes]], client.hgetall(key)) assert record[b"application_id"] == b"app" assert record[b"agent_id"] == b"agent" assert record[b"user_id"] == b"authoritative-user" @@ -233,3 +245,48 @@ async def test_live_legacy_index_requires_explicit_rebuild_and_preserves_exact_s await cast(Awaitable[int], client.hset(marker_key, mapping={"state": _SCOPE_MIGRATION_BACKFILL})) with pytest.raises(ValueError, match="migration is incomplete"): await _provider(url, index_name, prefix, user_id="Alice")._redis_search("migration marker") + + +@pytest.mark.integration +async def test_live_migration_renews_lock_during_slow_scan( + live_redis: tuple[str, Redis, str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, client, index_name, prefix = live_redis + legacy_schema = { + "index": {"name": index_name, "prefix": prefix, "key_separator": ":", "storage_type": "hash"}, + "fields": [ + {"name": "content", "type": "text"}, + {"name": "application_id", "type": "tag"}, + {"name": "agent_id", "type": "tag"}, + {"name": "user_id", "type": "tag"}, + ], + } + legacy_index = AsyncSearchIndex.from_dict(legacy_schema, redis_client=client) + await legacy_index.create() + await legacy_index.load([ + {"content": "lease marker", "application_id": "app", "agent_id": "agent", "user_id": "Alice"} + ]) + monkeypatch.setattr("agent_framework_redis._context_provider._SCOPE_MIGRATION_LOCK_SECONDS", 1) + first = _provider(url, index_name, prefix, user_id="Alice", overwrite_index=True) + await first.redis_index.exists() + first_client = first.redis_index.client + assert first_client is not None + original_scan = first_client.scan + scan_started = asyncio.Event() + + async def slow_scan(*args: Any, **kwargs: Any) -> Any: + scan_started.set() + await asyncio.sleep(1.5) + return await original_scan(*args, **kwargs) + + monkeypatch.setattr(first_client, "scan", slow_scan) + migration = asyncio.create_task(first._redis_search("lease marker")) + await scan_started.wait() + await asyncio.sleep(1.1) + + second = _provider(url, index_name, prefix, user_id="Alice", overwrite_index=True) + with pytest.raises(RuntimeError, match="already running"): + await second._redis_search("lease marker") + + assert len(await migration) == 1 diff --git a/python/packages/redis/tests/test_providers.py b/python/packages/redis/tests/test_providers.py index e0fb07430ae..a598b606681 100644 --- a/python/packages/redis/tests/test_providers.py +++ b/python/packages/redis/tests/test_providers.py @@ -38,7 +38,7 @@ def mock_index() -> AsyncMock: idx = AsyncMock() client = MagicMock() client.hset = AsyncMock() - client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b""]) + client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b"", b"0", b"[]"]) idx.client = client idx.create = AsyncMock() idx.load = AsyncMock() From a1f9151662489c9ce7a79afffad31dbe85eaca33 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Thu, 22 Oct 2026 13:55:44 +0200 Subject: [PATCH 5/6] Python: clarify Redis migration task cleanup --- .../packages/redis/agent_framework_redis/_context_provider.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index fc22d7bfc5f..58f204c1bcd 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -701,7 +701,7 @@ async def _migrate_scope_tokens(self) -> None: finally: renewal_task.cancel() with contextlib.suppress(asyncio.CancelledError): - await renewal_task + _ = await renewal_task await client.eval( "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", 1, From 6d6b768af77b67a122faf5eb672a37bee8f31712 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Thu, 22 Oct 2026 15:03:22 +0200 Subject: [PATCH 6/6] Python: align Redis context lifecycle tests --- python/packages/redis/tests/test_context_provider_sdk.py | 3 --- python/packages/redis/tests/test_providers.py | 1 + 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/python/packages/redis/tests/test_context_provider_sdk.py b/python/packages/redis/tests/test_context_provider_sdk.py index 3c5d5fdb3a7..d0c92cf5170 100644 --- a/python/packages/redis/tests/test_context_provider_sdk.py +++ b/python/packages/redis/tests/test_context_provider_sdk.py @@ -92,9 +92,6 @@ async def test_context_manager_closes_owned_index_only( for borrowed, expected_waits in ((False, 1), (True, 0)): closes.reset_mock() provider = _provider(redis_client, borrowed=borrowed) - # The index connects lazily on first use, so model an index that has been - # used (before_run runs exists/create through it) before the provider exits. - await provider.redis_index.set_client(redis_client) async with provider: pass assert closes.await_count == expected_waits diff --git a/python/packages/redis/tests/test_providers.py b/python/packages/redis/tests/test_providers.py index a598b606681..7b84b4be9e7 100644 --- a/python/packages/redis/tests/test_providers.py +++ b/python/packages/redis/tests/test_providers.py @@ -39,6 +39,7 @@ def mock_index() -> AsyncMock: client = MagicMock() client.hset = AsyncMock() client.hmget = AsyncMock(return_value=[b"1", b"complete", b"0", b"", b"", b"0", b"[]"]) + client.aclose = AsyncMock() idx.client = client idx.create = AsyncMock() idx.load = AsyncMock()