diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index 4751bc16961..117a0dd19b7 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -115,6 +115,9 @@ def __init__( 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; + # an injected index stays under the caller's control (matching RedisHistoryProvider). + self._owns_index = redis_index is None index = redis_index or AsyncSearchIndex.from_dict( # pyright: ignore[reportUnknownMemberType] self.schema_dict, redis_url=self.redis_url, validate_on_load=True ) @@ -441,7 +444,15 @@ async def __aenter__(self) -> Self: return self async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: - """Async context manager exit.""" + """Async context manager exit; releases a provider-owned index connection.""" + await self.aclose() + + async def aclose(self) -> None: + """Close the Redis connection when this provider owns the index.""" + if self._owns_index: + client = getattr(self.redis_index, "client", None) + if client is not None: + await client.aclose() __all__ = ["RedisContextProvider"] diff --git a/python/packages/redis/tests/test_context_provider_sdk.py b/python/packages/redis/tests/test_context_provider_sdk.py index 31964a2c1e2..c8a497e0578 100644 --- a/python/packages/redis/tests/test_context_provider_sdk.py +++ b/python/packages/redis/tests/test_context_provider_sdk.py @@ -77,6 +77,26 @@ def _capture_hash_writes(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return write +@pytest.mark.asyncio +async def test_context_manager_closes_owned_index_only( + redis_client: redis.Redis, monkeypatch: pytest.MonkeyPatch +) -> None: + """Exiting the context manager closes the connection an owned index created and + leaves a borrowed index under the caller's control.""" + closes = AsyncMock(wraps=redis_client.aclose) + monkeypatch.setattr(redis_client, "aclose", closes) + + 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 + + @pytest.mark.parametrize("borrowed", [False, True], ids=["owned-index", "borrowed-index"]) @pytest.mark.parametrize("all_scopes", [False, True], ids=["user-scope", "all-scopes"]) async def test_after_run_encodes_messages_with_optional_fields_omitted(