fix(redis): propagate listener startup failures to senders

ag2/streams/redis/stream.py

@@ -69,7 +69,7 @@ def __init__(
         self._serializer = serializer
         self._channel = f"{prefix}:pubsub:{self.id}"
         self._instance_id = str(uuid4())
-        self._listener_task: asyncio.Task | None = None
+        self._listener_task: asyncio.Task[None] | None = None
         self._listener_ready = asyncio.Event()
         self._pubsub_redis = aioredis.from_url(redis_url)
         self._publish_redis = aioredis.from_url(redis_url)
@@ -78,18 +78,21 @@ def __init__(
         # subscribers processing events from remote instances.
         self._base_context: Context | None = None
 
-    def _ensure_listener(self) -> None:
+    def _ensure_listener(self) -> asyncio.Task[None]:
         """Start the Redis Pub/Sub listener if not already running."""
         if self._listener_task is None or self._listener_task.done():
             self._listener_ready.clear()
             self._listener_task = asyncio.create_task(self._listen())
+        return self._listener_task
 
     async def _listen(self) -> None:
         """Listen for events on the Redis Pub/Sub channel and dispatch to local subscribers."""
         pubsub = self._pubsub_redis.pubsub()
-        await pubsub.subscribe(self._channel)
-        self._listener_ready.set()
+        subscribed = False
         try:
+            await pubsub.subscribe(self._channel)
+            subscribed = True
+            self._listener_ready.set()
             async for message in pubsub.listen():
                 if message["type"] != "message":
                     continue
@@ -106,8 +109,11 @@ async def _listen(self) -> None:
             # CancelledError (and any other exception) propagates out of the
             # try; this cleanup still runs on cancellation so the pub/sub
             # subscription is always released.
-            await pubsub.unsubscribe(self._channel)
-            await pubsub.aclose()
+            try:
+                if subscribed:
+                    await pubsub.unsubscribe(self._channel)
+            finally:
+                await pubsub.aclose()
 
     @staticmethod
     def _split_origin(raw: bytes) -> tuple[str | None, bytes]:
@@ -123,8 +129,9 @@ def _split_origin(raw: bytes) -> tuple[str | None, bytes]:
 
     async def send(self, event: BaseEvent, context: Context) -> None:
         """Persist the event, dispatch locally, and publish to Redis for remote listeners."""
-        self._ensure_listener()
-        await self._listener_ready.wait()
+        listener = self._ensure_listener()
+        if not self._listener_ready.is_set():
+            await _wait_for_listener(listener, self._listener_ready)
         # Capture the first caller's context so the listener can reuse its
         # dependency_provider and dependencies for remotely-received events.
         if self._base_context is None:
@@ -147,3 +154,15 @@ async def close(self) -> None:
         await self._pubsub_redis.aclose()
         await self._publish_redis.aclose()
         await self._redis_storage.close()
+
+
+async def _wait_for_listener(listener: asyncio.Task[None], ready: asyncio.Event) -> None:
+    """Wait for startup or propagate its failure without cancelling the shared listener."""
+    waiter = asyncio.create_task(ready.wait())
+    try:
+        done, _ = await asyncio.wait((waiter, listener), return_when=asyncio.FIRST_COMPLETED)
+        if listener in done:
+            await listener
+    finally:
+        waiter.cancel()
+        await asyncio.gather(waiter, return_exceptions=True)

test/stream/redis/test_redis_stream.py

@@ -10,9 +10,11 @@
 
 import pytest
 import pytest_asyncio
+from redis.exceptions import ConnectionError as RedisConnectionError
 
 from ag2 import Context
 from ag2.events import ImageInput, ModelMessage, TextInput, ToolCallEvent
+from ag2.streams.redis import RedisStream
 from ag2.streams.redis.serializer import Serializer
 
 pytestmark = [
@@ -108,6 +110,45 @@ async def aclose(self) -> None:
         pass
 
 
+class StartupRedis(MockRedis):
+    """Redis boundary with controllable subscription startup and cleanup."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.starting = asyncio.Event()
+        self.release = asyncio.Event()
+        self.error: Exception | None = None
+        self.unsubscribe_error: Exception | None = None
+        self.pubsubs: list[StartupPubSub] = []
+
+    def pubsub(self) -> "StartupPubSub":
+        pubsub = StartupPubSub(self)
+        self.pubsubs.append(pubsub)
+        return pubsub
+
+
+class StartupPubSub(MockPubSub):
+    def __init__(self, redis: StartupRedis) -> None:
+        super().__init__(redis)
+        self.server = redis
+        self.closed = False
+
+    async def subscribe(self, channel: str) -> None:
+        self.server.starting.set()
+        await self.server.release.wait()
+        if self.server.error is not None:
+            raise self.server.error
+        await super().subscribe(channel)
+
+    async def unsubscribe(self, channel: str) -> None:
+        if self.server.unsubscribe_error is not None:
+            raise self.server.unsubscribe_error
+        await super().unsubscribe(channel)
+
+    async def aclose(self) -> None:
+        self.closed = True
+
+
 # Shared fake Redis instance per test — all RedisStream instances in the same
 # test connect to the same "server" so cross-instance pub/sub works.
 @pytest.fixture()
@@ -117,8 +158,6 @@ def mock_redis():
 
 @pytest_asyncio.fixture(params=[Serializer.JSON, Serializer.PICKLE], ids=["json", "pickle"])
 async def redis_stream(mock_redis, request):
-    from ag2.streams.redis import RedisStream
-
     serializer = request.param
     streams: list = []
 
@@ -411,6 +450,102 @@ async def test_local_dispatch_survives_listener_failure(self, redis_stream):
         assert stream._listener_task is dead_task and dead_task.done()
 
 
+class TestRedisStartup:
+    @pytest.fixture()
+    def mock_redis(self):
+        return StartupRedis()
+
+    async def test_subscription_failure_reaches_sender(self, redis_stream, mock_redis):
+        stream = redis_stream()
+        mock_redis.error = RedisConnectionError("subscription unavailable")
+        mock_redis.release.set()
+
+        with pytest.raises(RedisConnectionError, match="subscription unavailable"):
+            await asyncio.wait_for(stream.send(ModelMessage("hello"), Context(stream)), timeout=1)
+
+        assert mock_redis.pubsubs[0].closed
+        assert await stream.history.get_events() == []
+
+    async def test_subscription_failure_releases_all_senders(self, redis_stream, mock_redis):
+        stream = redis_stream()
+        error = mock_redis.error = RedisConnectionError("subscription unavailable")
+        sends = [asyncio.create_task(stream.send(ModelMessage(text), Context(stream))) for text in ("a", "b")]
+        try:
+            await mock_redis.starting.wait()
+            mock_redis.release.set()
+            outcomes = await asyncio.wait_for(asyncio.gather(*sends, return_exceptions=True), timeout=1)
+        finally:
+            for send in sends:
+                send.cancel()
+            await asyncio.gather(*sends, return_exceptions=True)
+
+        assert outcomes == [error, error]
+        assert mock_redis.pubsubs[0].closed
+        assert await stream.history.get_events() == []
+
+    async def test_send_retries_subscription_after_startup_failure(self, redis_stream, mock_redis):
+        stream = redis_stream()
+        mock_redis.error = RedisConnectionError("subscription unavailable")
+        mock_redis.release.set()
+        with pytest.raises(RedisConnectionError):
+            await asyncio.wait_for(stream.send(ModelMessage("failed"), Context(stream)), timeout=1)
+
+        mock_redis.error = None
+        event = ModelMessage("recovered")
+        received = []
+        stream.subscribe(lambda ev: received.append(ev))
+        await asyncio.wait_for(stream.send(event, Context(stream)), timeout=1)
+
+        assert received == [event]
+        assert await stream.history.get_events() == [event]
+        assert mock_redis.pubsubs[0].closed
+        await stream.close()
+        await stream.close()
+        assert all(pubsub.closed for pubsub in mock_redis.pubsubs)
+
+    async def test_cancelling_one_sender_keeps_subscription_for_others(self, redis_stream, mock_redis):
+        stream = redis_stream()
+        first = asyncio.create_task(stream.send(ModelMessage("cancelled"), Context(stream)))
+        await mock_redis.starting.wait()
+        first.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await first
+
+        mock_redis.release.set()
+        event = ModelMessage("remaining")
+        await asyncio.wait_for(stream.send(event, Context(stream)), timeout=1)
+
+        assert await stream.history.get_events() == [event]
+        await stream.close()
+        assert all(pubsub.closed for pubsub in mock_redis.pubsubs)
+
+    async def test_close_during_startup_releases_sender_and_pubsub(self, redis_stream, mock_redis):
+        stream = redis_stream()
+        send = asyncio.create_task(stream.send(ModelMessage("cancelled"), Context(stream)))
+        try:
+            await mock_redis.starting.wait()
+            await stream.close()
+            with pytest.raises(asyncio.CancelledError):
+                await asyncio.wait_for(asyncio.shield(send), timeout=1)
+        finally:
+            send.cancel()
+            await asyncio.gather(send, return_exceptions=True)
+
+        assert mock_redis.pubsubs[0].closed
+        assert await stream.history.get_events() == []
+
+    async def test_unsubscribe_failure_still_closes_pubsub(self, redis_stream, mock_redis):
+        stream = redis_stream()
+        mock_redis.release.set()
+        await stream.send(ModelMessage("hello"), Context(stream))
+        mock_redis.unsubscribe_error = RedisConnectionError("unsubscribe unavailable")
+
+        with pytest.raises(RedisConnectionError, match="unsubscribe unavailable"):
+            await stream.close()
+
+        assert mock_redis.pubsubs[0].closed
+
+
 class TestBinaryRoundTrip:
     """Events carrying binary or provider-SDK fields must round-trip intact."""