Python: close the RedisContextProvider index connection when the provider owns it

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"]

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(