Python: Add Azure Managed Redis Support with Credential Provider

python/packages/redis/agent_framework_redis/_chat_message_store.py

@@ -55,6 +55,11 @@ class RedisChatMessageStore:
     def __init__(
         self,
         redis_url: str | None = None,
+        credential_provider: Any | None = None,
+        host: str | None = None,
+        port: int = 6380,
+        ssl: bool = True,
+        username: str | None = None,
         thread_id: str | None = None,
         key_prefix: str = "chat_messages",
         max_messages: int | None = None,
@@ -63,12 +68,19 @@ def __init__(
         """Initialize the Redis chat message store.
 
         Creates a Redis-backed chat message store for a specific conversation thread.
-        The store will automatically create a Redis connection and manage message
-        persistence using Redis List operations.
+        Supports both traditional URL-based authentication and Azure Managed Redis
+        with credential provider.
 
         Args:
             redis_url: Redis connection URL (e.g., "redis://localhost:6379").
-                      Required for establishing Redis connection.
+                      Used for traditional authentication. Mutually exclusive with credential_provider.
+            credential_provider: Credential provider for Azure AD authentication.
+                                Requires host parameter. Mutually exclusive with redis_url.
+            host: Redis host name (e.g., "myredis.redis.cache.windows.net").
+                 Required when using credential_provider.
+            port: Redis port number. Defaults to 6380 (Azure Redis SSL port).
+            ssl: Enable SSL/TLS connection. Defaults to True.
+            username: Redis username. Defaults to None.
             thread_id: Unique identifier for this conversation thread.
                       If not provided, a UUID will be auto-generated.
                       This becomes part of the Redis key: {key_prefix}:{thread_id}
@@ -82,23 +94,58 @@ def __init__(
                      Useful for resuming conversations or seeding with context.
 
         Raises:
-            ValueError: If redis_url is None (Redis connection is required).
-            redis.ConnectionError: If unable to connect to Redis server.
+            ValueError: If neither redis_url nor credential_provider is provided.
+            ValueError: If both redis_url and credential_provider are provided.
+            ValueError: If credential_provider is used without host parameter.
+
+        Examples:
+            Traditional connection:
+                store = RedisChatMessageStore(
+                    redis_url="redis://localhost:6379",
+                    thread_id="conversation_123"
+                )
+
+            Azure Managed Redis with credential provider:
+                from redis.credentials import CredentialProvider
+                from azure.identity.aio import DefaultAzureCredential
+
+                store = RedisChatMessageStore(
+                    credential_provider=CredentialProvider(DefaultAzureCredential()),
+                    host="myredis.redis.cache.windows.net",
+                    thread_id="conversation_123"
+                )
+        """
+        # Validate connection parameters
+        if redis_url is None and credential_provider is None:
+            raise ValueError("Either redis_url or credential_provider must be provided")
 
+        if redis_url is not None and credential_provider is not None:
+            raise ValueError("redis_url and credential_provider are mutually exclusive")
 
-        """
-        # Validate required parameters
-        if redis_url is None:
-            raise ValueError("redis_url is required for Redis connection")
+        if credential_provider is not None and host is None:
+            raise ValueError("host is required when using credential_provider")
 
         # Store configuration
-        self.redis_url = redis_url
         self.thread_id = thread_id or f"thread_{uuid4()}"
         self.key_prefix = key_prefix
         self.max_messages = max_messages
 
-        # Initialize Redis client with connection pooling and async support
-        self._redis_client = redis.from_url(redis_url, decode_responses=True)  # type: ignore[no-untyped-call]
+        # Initialize Redis client based on authentication method
+        if credential_provider is not None and host is not None:
+            # Azure AD authentication with credential provider
+            self.redis_url = None  # Not using URL-based auth
+            self._redis_client = redis.Redis(
+                host=host,
+                port=port,
+                ssl=ssl,
+                username=username,
+                credential_provider=credential_provider,
+                decode_responses=True,
+            )
+        else:
+            # Traditional URL-based authentication
+            self.redis_url = redis_url
+            self._redis_client = redis.from_url(redis_url, decode_responses=True)  # type: ignore[no-untyped-call]
 
         # Handle initial messages (will be moved to Redis on first access)
         self._initial_messages = list(messages) if messages else []

python/packages/redis/tests/test_redis_chat_message_store.py

@@ -93,11 +93,118 @@ def test_init_with_max_messages(self):
         assert store.max_messages == 100
 
     def test_init_with_redis_url_required(self):
-        """Test that redis_url is required for initialization."""
-        with pytest.raises(ValueError, match="redis_url is required for Redis connection"):
-            # Should raise an exception since redis_url is required
+        """Test that either redis_url or credential_provider is required."""
+        with pytest.raises(ValueError, match="Either redis_url or credential_provider must be provided"):
             RedisChatMessageStore(thread_id="test123")
 
+    def test_init_with_credential_provider(self):
+        """Test initialization with credential_provider."""
+        mock_credential_provider = MagicMock()
+
+        with patch("agent_framework_redis._chat_message_store.redis.Redis") as mock_redis_class:
+            mock_redis_instance = MagicMock()
+            mock_redis_class.return_value = mock_redis_instance
+
+            store = RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                thread_id="test123",
+            )
+
+            # Verify Redis.Redis was called with correct parameters
+            mock_redis_class.assert_called_once_with(
+                host="myredis.redis.cache.windows.net",
+                port=6380,
+                ssl=True,
+                username=None,
+                credential_provider=mock_credential_provider,
+                decode_responses=True,
+            )
+            # Verify store instance is properly initialized
+            assert store.thread_id == "test123"
+            assert store.redis_url is None  # Should be None for credential provider auth
+            assert store.key_prefix == "chat_messages"
+            assert store.max_messages is None
+
+    def test_init_with_credential_provider_custom_port(self):
+        """Test initialization with credential_provider and custom port."""
+        mock_credential_provider = MagicMock()
+
+        with patch("agent_framework_redis._chat_message_store.redis.Redis") as mock_redis_class:
+            mock_redis_instance = MagicMock()
+            mock_redis_class.return_value = mock_redis_instance
+
+            store = RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                port=6379,
+                ssl=False,
+                username="admin",
+                thread_id="test123",
+            )
+
+            # Verify custom parameters were passed
+            mock_redis_class.assert_called_once_with(
+                host="myredis.redis.cache.windows.net",
+                port=6379,
+                ssl=False,
+                username="admin",
+                credential_provider=mock_credential_provider,
+                decode_responses=True,
+            )
+            # Verify store instance is properly initialized
+            assert store.thread_id == "test123"
+            assert store.redis_url is None  # Should be None for credential provider auth
+            assert store.key_prefix == "chat_messages"
+
+    def test_init_credential_provider_requires_host(self):
+        """Test that credential_provider requires host parameter."""
+        mock_credential_provider = MagicMock()
+
+        with pytest.raises(ValueError, match="host is required when using credential_provider"):
+            RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                thread_id="test123",
+            )
+
+    def test_init_mutually_exclusive_params(self):
+        """Test that redis_url and credential_provider are mutually exclusive."""
+        mock_credential_provider = MagicMock()
+
+        with pytest.raises(ValueError, match="redis_url and credential_provider are mutually exclusive"):
+            RedisChatMessageStore(
+                redis_url="redis://localhost:6379",
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                thread_id="test123",
+            )
+
+    async def test_serialize_with_credential_provider(self):
+        """Test that serialization works correctly with credential provider authentication."""
+        mock_credential_provider = MagicMock()
+
+        with patch("agent_framework_redis._chat_message_store.redis.Redis") as mock_redis_class:
+            mock_redis_instance = MagicMock()
+            mock_redis_class.return_value = mock_redis_instance
+
+            store = RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                thread_id="test123",
+                key_prefix="custom_prefix",
+                max_messages=100,
+            )
+
+            # Serialize the store state
+            state = await store.serialize()
+
+            # Verify serialization includes correct values
+            assert state["thread_id"] == "test123"
+            assert state["redis_url"] is None  # Should be None for credential provider auth
+            assert state["key_prefix"] == "custom_prefix"
+            assert state["max_messages"] == 100
+            assert state["type"] == "redis_store_state"
+
     def test_init_with_initial_messages(self, sample_messages):
         """Test initialization with initial messages."""
         with patch("agent_framework_redis._chat_message_store.redis.from_url"):