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/README.md b/python/packages/redis/README.md index 018fa4ecedf..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 diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index 117a0dd19b7..58f204c1bcd 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -8,7 +8,12 @@ from __future__ import annotations +import asyncio +import base64 +import contextlib +import hashlib import json +import secrets import sys from typing import TYPE_CHECKING, Any, ClassVar, Literal @@ -20,10 +25,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 +45,93 @@ 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_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" +_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} + + +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() + + +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 _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")} + + 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. @@ -93,7 +184,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 +204,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 +294,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 +322,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,50 +354,18 @@ 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 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() - await self.redis_index.create(overwrite=self.overwrite_index, drop=False) 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_index: Any = await AsyncSearchIndex.from_existing( # pyright: ignore[reportUnknownMemberType] - self.index_name, redis_url=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) @@ -309,8 +375,361 @@ 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." + ) + 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", + "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", + "pending_cursor", + "pending_keys", + ), + decoded, + strict=True, ) + ) + + async def _set_scope_migration_state( + self, + state: str, + *, + 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.") + 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.""" + 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.") + + 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 + ) + 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_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 + + 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, + ) + 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( + _SCOPE_MIGRATION_COMPLETE, + 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, + 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() + 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( self, @@ -327,11 +746,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 +852,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..6ad54e2c463 100644 --- a/python/packages/redis/tests/test_context_provider_edges.py +++ b/python/packages/redis/tests/test_context_provider_edges.py @@ -14,12 +14,28 @@ 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_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, +) @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"", b"0", b"[]"]) + index.client = client index.create = AsyncMock() index.exists = AsyncMock(return_value=False) index.load = AsyncMock() @@ -38,15 +54,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 +108,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 @@ -100,7 +147,22 @@ 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 + + +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_not_awaited() + migrate.assert_awaited_once() assert provider._index_initialized is True @@ -117,7 +179,21 @@ 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() + + +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"", b"0", b"[]"]) + + with pytest.raises(ValueError, match="migration is incomplete"): await provider._validate_schema_compatibility() @@ -154,8 +230,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 +251,239 @@ 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, + 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"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() + 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.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) + monkeypatch.setattr("agent_framework_redis._context_provider.secrets.token_hex", lambda _: "lock") + + 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) + mapping = write_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.create.assert_awaited_once_with(overwrite=True, drop=False) + mock_index.query.assert_awaited_once() + assert any(call.args[0] == _SCOPE_MIGRATION_COMPLETE for call in set_state.await_args_list) + + +async def test_backfill_stages_oversized_scan_page_without_advancing_cursor( + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = MagicMock() + client.eval = AsyncMock(return_value=1) + 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) + 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") + 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=provider.schema_dict) + 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() + + 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: + marker = _scope_migration_key("context", "context") + + 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, @@ -257,12 +569,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 +589,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..4e60b36dd32 --- /dev/null +++ b/python/packages/redis/tests/test_context_provider_integration.py @@ -0,0 +1,292 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live Redis isolation and migration coverage for ``RedisContextProvider``.""" + +from __future__ import annotations + +import asyncio +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_MIGRATION_BACKFILL, + _SCOPE_MIGRATION_COMPLETE, + _scope_migration_key, + _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) + 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) + await client.delete(marker_key, f"{marker_key}:lock") + 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") == [] + + 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 + 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" + 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", + } + ]) + 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"): + 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") == [] + 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") + + +@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_context_provider_sdk.py b/python/packages/redis/tests/test_context_provider_sdk.py index c8a497e0578..d0c92cf5170 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): @@ -66,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 @@ -89,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 @@ -122,15 +122,19 @@ 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 - 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 @@ -163,16 +167,18 @@ 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 assert record["embedding"] == b"\x00\x00\x80?" - assert len(_TransportConnection.packets) == 1 + assert len(_TransportConnection.packets) == 2 assert document == before @@ -189,7 +195,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"} @@ -211,5 +217,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 711c63e7686..7b84b4be9e7 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 @@ -36,6 +36,11 @@ 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"", b"0", b"[]"]) + client.aclose = AsyncMock() + idx.client = client idx.create = AsyncMock() idx.load = AsyncMock() idx.query = AsyncMock(return_value=[]) @@ -136,8 +141,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"