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..a71506d8a49 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). +The beta `RedisContextProvider` scope schema is not backward compatible with +indexes created by earlier releases. Delete the old index and its context +hashes, let the provider create a new index, and re-ingest the context data. +Incompatible indexes fail closed; `overwrite_index=True` does not migrate them. + 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..d4ef19eb1fa 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -8,6 +8,7 @@ from __future__ import annotations +import hashlib import json import sys from typing import TYPE_CHECKING, Any, ClassVar, Literal @@ -20,10 +21,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 +41,57 @@ 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" + +_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 _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 +144,9 @@ 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 recreate an existing compatible index + definition while preserving its stored hashes. This does not + migrate indexes created with an older scope schema. """ super().__init__(source_id) self.redis_url = redis_url @@ -112,7 +165,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 +255,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 +283,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 +315,15 @@ 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: await self._validate_schema_compatibility() - await self.redis_index.create(overwrite=self.overwrite_index, drop=False) + await self.redis_index.create(overwrite=self.overwrite_index if index_exists else False, drop=False) + 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_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,9 +333,19 @@ 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." + "This beta schema change cannot migrate existing context data safely. " + "Delete the existing index and its context hashes, recreate the provider index, " + "and re-ingest the context data." ) + 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 _add( self, *, @@ -327,11 +361,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 +467,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..1ae85cfe70a 100644 --- a/python/packages/redis/tests/test_context_provider_edges.py +++ b/python/packages/redis/tests/test_context_provider_edges.py @@ -14,12 +14,20 @@ 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 def mock_index() -> AsyncMock: index = AsyncMock() + client = MagicMock() + index.client = client index.create = AsyncMock() index.exists = AsyncMock(return_value=False) index.load = AsyncMock() @@ -31,22 +39,46 @@ def mock_index() -> AsyncMock: def patch_index(mock_index: AsyncMock) -> Generator[MagicMock]: with patch("agent_framework_redis._context_provider.AsyncSearchIndex") as mock_cls: mock_cls.from_dict = MagicMock(return_value=mock_index) - mock_cls.from_existing = AsyncMock() + + async def mock_from_existing(*args: Any, **kwargs: Any) -> AsyncMock: # noqa: ARG001 + existing = AsyncMock() + existing.schema.to_dict = MagicMock( + side_effect=lambda: mock_cls.from_dict.call_args.args[0] if mock_cls.from_dict.call_args else {} + ) + return existing + + mock_cls.from_existing = AsyncMock(side_effect=mock_from_existing) yield mock_cls 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 @@ -99,11 +146,48 @@ async def test_ensure_index_validates_existing_schema_before_create( with patch.object(provider, "_validate_schema_compatibility", AsyncMock()) as validate_schema: await provider._ensure_index() - validate_schema.assert_awaited_once() + assert validate_schema.await_count == 2 mock_index.create.assert_awaited_once_with(overwrite=False, drop=False) assert provider._index_initialized is True +async def test_ensure_index_validates_before_overwriting_compatible_index( + 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, "_validate_schema_compatibility", AsyncMock()) as validate_schema: + await provider._ensure_index() + + assert validate_schema.await_count == 2 + mock_index.create.assert_awaited_once_with(overwrite=True, drop=False) + assert provider._index_initialized is True + + +async def test_ensure_index_validates_schema_created_by_concurrent_initializer( + mock_index: AsyncMock, + patch_index: MagicMock, # noqa: ARG001 +) -> None: + mock_index.exists.return_value = False + provider = RedisContextProvider(source_id="ctx", user_id="user-1", overwrite_index=True) + + with ( + patch.object( + provider, + "_validate_schema_compatibility", + AsyncMock(side_effect=ValueError("incompatible concurrent schema")), + ) as validate_schema, + pytest.raises(ValueError, match="incompatible concurrent schema"), + ): + await provider._ensure_index() + + mock_index.create.assert_awaited_once_with(overwrite=False, drop=False) + validate_schema.assert_awaited_once() + assert provider._index_initialized is False + + async def test_validate_schema_compatibility_raises_for_significant_mismatch( patch_index: MagicMock, ) -> None: @@ -117,7 +201,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="Delete the existing index and its context hashes"): await provider._validate_schema_compatibility() @@ -154,8 +238,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,6 +259,14 @@ 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) @@ -189,6 +289,7 @@ async def test_after_run_stores_vectors_in_configured_dtype( vector_field_name="embedding", redis_index=mock_index, ) + object.__setattr__(provider, "_validate_schema_compatibility", AsyncMock()) session = AgentSession(session_id="session-1") context = SessionContext( session_id=session.session_id, @@ -257,12 +358,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 +378,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..1fc69ea3b09 --- /dev/null +++ b/python/packages/redis/tests/test_context_provider_integration.py @@ -0,0 +1,241 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live Redis isolation 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) + 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) + 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_fails_closed_until_recreated( + 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"}), + ) + + for overwrite_index in (False, True): + incompatible = _provider( + url, + index_name, + prefix, + user_id="Alice", + overwrite_index=overwrite_index, + ) + with pytest.raises(ValueError, match="Delete the existing index and its context hashes"): + await incompatible._redis_search("migration marker") + + await legacy_index.delete(drop=True) + recreated = _provider(url, index_name, prefix, user_id="Alice") + await recreated._add(data={"content": "re-ingested marker"}, session_id="session") + + assert len(await recreated._redis_search("re-ingested marker")) == 1 + assert await _provider(url, index_name, prefix, user_id="alice")._redis_search("re-ingested marker") == [] + adjacent_record = await cast(Awaitable[dict[bytes, bytes]], client.hgetall(adjacent_key)) + assert adjacent_record[b"content"] == b"foreign" diff --git a/python/packages/redis/tests/test_context_provider_sdk.py b/python/packages/redis/tests/test_context_provider_sdk.py index c8a497e0578..47af6d523a5 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,8 +67,11 @@ 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() + object.__setattr__(provider, "_validate_schema_compatibility", AsyncMock()) return provider @@ -89,9 +93,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 +123,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 all(b"HSET" in packet for packet in _TransportConnection.packets) assert message.to_dict() == before @@ -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..c8aba3a2e55 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,9 @@ async def test_empty_history_save_marks_redis_used(mock_redis_client: MagicMock) @pytest.fixture def mock_index() -> AsyncMock: idx = AsyncMock() + client = MagicMock() + client.aclose = AsyncMock() + idx.client = client idx.create = AsyncMock() idx.load = AsyncMock() idx.query = AsyncMock(return_value=[]) @@ -48,7 +51,7 @@ def patch_index_from_dict(mock_index: AsyncMock): with patch("agent_framework_redis._context_provider.AsyncSearchIndex") as mock_cls: mock_cls.from_dict = MagicMock(return_value=mock_index) - async def mock_from_existing(index_name: str, redis_url: str): # noqa: ARG001 + async def mock_from_existing(*args: Any, **kwargs: Any): # noqa: ARG001 mock_existing = AsyncMock() mock_existing.schema.to_dict = MagicMock( side_effect=lambda: mock_cls.from_dict.call_args[0][0] if mock_cls.from_dict.call_args else {} @@ -136,8 +139,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"