fix: more fixes

...

private_gpt/components/streaming/providers/in_memory_stream_service.py

@@ -175,6 +175,7 @@ class InMemoryStreamService(StreamService):
                 self._waiters[correlation_id].add(waiter)
 
         if waiter is not None:
+            assert block_ms is not None
             try:
                 await asyncio.wait_for(waiter.wait(), timeout=block_ms / 1000)
             except TimeoutError:

private_gpt/components/streaming/stream/stream_reader.py

@@ -3,7 +3,7 @@ import contextlib
 import logging
 from collections import defaultdict
 from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Final
 
 from injector import inject, singleton
 from pydantic import BaseModel
@@ -20,7 +20,9 @@ logger = logging.getLogger(__name__)
 
 DEFAULT_BLOCK_MS: Final[int] = 1000
 BATCH_SIZE: Final[int] = 5
-STATUS_CHECK_INTERVAL: Final[float] = 5.0
+# How often (seconds) to poll for terminal status when no events arrive.
+# Halves Redis metadata round-trips vs. checking every block cycle (1 s).
+STATUS_CHECK_INTERVAL: Final[float] = 2.0
 TERMINAL_STATUSES: Final[set[StreamStatus]] = {
     StreamStatus.COMPLETED,
     StreamStatus.CANCELLED,
@@ -28,39 +30,118 @@ TERMINAL_STATUSES: Final[set[StreamStatus]] = {
 }
 
 
-class StreamConsumer:
-    """A single consumer receiving events for one correlation id.
+class StreamBroadcast:
+    """Append-only batch log that broadcasts to N concurrent readers.
 
-    Events are buffered in an unbounded ``asyncio.Queue``.  A ``None`` entry
-    is the sentinel that signals the consumer the stream is finished.
+    A single producer calls ``push`` / ``close`` (both synchronous); each
+    subscriber holds an integer *cursor* into ``_batches`` and parks on a
+    ``asyncio.Future`` until new data arrives.  One ``push`` resolves *all*
+    pending futures in a single pass — no per-subscriber queues.
+
+    Safety: all operations are synchronous except ``read_from``, so no lock is
+    needed (asyncio cooperative threading guarantees no interleaving between
+    non-await statements).
     """
 
-    def __init__(self, correlation_id: str, event_handler: EventHandler):
-        self.correlation_id = correlation_id
-        self.event_handler = event_handler
-        self.queue: asyncio.Queue[BaseModel | None] = asyncio.Queue()
-        self.ref_count = 1
+    def __init__(self) -> None:
+        self._batches: list[list[BaseModel] | None] = []
+        self._waiters: list[asyncio.Future[None]] = []
+
+    def _notify(self) -> None:
+        waiters, self._waiters = self._waiters, []
+        for fut in waiters:
+            if not fut.done():
+                fut.set_result(None)
+
+    def push(self, batch: list[BaseModel]) -> None:
+        self._batches.append(batch)
+        self._notify()
 
     def close(self) -> None:
-        """Signal the consumer that no more events will arrive."""
-        self.queue.put_nowait(None)
+        """Append the terminal sentinel and wake all parked readers."""
+        self._batches.append(None)
+        self._notify()
+
+    def next_cursor(self) -> int:
+        """Cursor for a subscriber joining right now (no backfill)."""
+        return len(self._batches)
+
+    @property
+    def is_closed(self) -> bool:
+        return bool(self._batches) and self._batches[-1] is None
+
+    async def read_from(
+        self,
+        cursor: int,
+        stop_event: asyncio.Event | None = None,
+    ) -> AsyncGenerator[BaseModel, None]:
+        """Yield events starting at *cursor*, blocking until new ones arrive."""
+        while True:
+            # Drain already-available batches.  Safe without a lock: _batches is
+            # append-only and list operations are atomic between await points.
+            while cursor < len(self._batches):
+                batch = self._batches[cursor]
+                cursor += 1
+                if batch is None:
+                    return
+                for event in batch:
+                    yield event
+
+            if stop_event and stop_event.is_set():
+                return
+
+            # Park until the producer pushes more data.
+            fut: asyncio.Future[None] = asyncio.get_event_loop().create_future()
+            self._waiters.append(fut)
+            try:
+                await asyncio.wait_for(fut, timeout=1.0)
+            except TimeoutError:
+                pass  # Re-check stop_event and drain any new batches
+            finally:
+                # Clean up if push() hasn't already swapped us out.
+                if fut in self._waiters:
+                    self._waiters.remove(fut)
+
+
+class StreamConsumer:
+    """A subscriber on a ``StreamBroadcast`` with an individual read cursor."""
+
+    def __init__(
+        self,
+        correlation_id: str,
+        event_handler: EventHandler,
+        broadcast: StreamBroadcast,
+        cursor: int,
+    ) -> None:
+        self.correlation_id = correlation_id
+        self.event_handler = event_handler
+        self.broadcast = broadcast
+        self.cursor = cursor  # Index into broadcast._batches
 
 
 class StreamState:
-    """Read position and cache for a single multiplexed stream."""
+    """Read position and broadcast log for a single multiplexed stream."""
 
-    def __init__(self, last_id: str = "0"):
+    def __init__(self, last_id: str = "0") -> None:
         self.last_id = last_id
         self.cached_events: list[BaseModel] | None = None
+        # 0.0 → check on the first idle cycle (ancient enough to pass the gate)
         self.last_status_check: float = 0.0
+        self.broadcast = StreamBroadcast()
 
 
 @singleton
 class StreamReader:
-    """Reads events from the backend (Redis or in-memory) and deserializes them."""
+    """Reads events from the backend (Redis or in-memory) and deserializes them.
+
+    On Redis, ``read_events`` with ``block_ms > 0`` issues ``XREAD BLOCK`` which
+    parks the connection — the event loop stays free.  On the in-memory backend
+    the provider uses an ``asyncio.Event`` to avoid busy-waiting.
+    ``block_ms=None`` is always non-blocking (used for catch-up / drain reads).
+    """
 
     @inject
-    def __init__(self, stream_component: StreamComponent):
+    def __init__(self, stream_component: StreamComponent) -> None:
         self.stream_service = stream_component.stream
 
     async def read_events(
@@ -107,9 +188,9 @@ class StreamReader:
                 logger.error(f"Error checking cached status for {correlation_id}: {e}")
 
         try:
-            metadata: (
-                StreamMetadata | None
-            ) = await self.stream_service.get_stream_metadata(correlation_id)
+            metadata: StreamMetadata | None = (
+                await self.stream_service.get_stream_metadata(correlation_id)
+            )
             if metadata and metadata.status in TERMINAL_STATUSES:
                 return True
         except Exception as e:
@@ -140,7 +221,7 @@ class StreamReader:
         async def produce_events(current_last_id: str) -> None:
             try:
                 cached_events: list[BaseModel] | None = None
-                last_status_check: float = 0.0
+                last_status_check: float = 0.0  # 0 → check on first idle cycle
 
                 while not stop_event.is_set():
                     events, current_last_id = await self.read_events(
@@ -205,24 +286,22 @@ class StreamReader:
 class StreamMultiplexer:
     """Multiplexes many stream consumers onto a single worker task.
 
-    Each correlation id has one ``StreamState`` (read position) shared by all
-    consumers on that stream.  The worker loops over every active stream,
-    issues a blocking ``read_events`` and dispatches the batch to every
-    current consumer.
+    Each correlation id has one ``StreamState`` containing a ``StreamBroadcast``
+    shared by all subscribers.  The worker calls ``_process_stream`` once per
+    active stream per cycle: one Redis read, one ``broadcast.push`` — no
+    per-consumer queues.
 
     Key invariants
     ---------------
-    * Dispatch and ``last_id`` advance happen atomically under the lock, and
-      the consumer list is re-fetched immediately before dispatch.  This means
-      a consumer that joins while the worker is blocked inside ``read_events``
-      still receives the batch about to be dispatched.
-    * ``add_consumer`` accepts a ``last_id`` and initialises the stream state
-      from it, so a reconnecting consumer never receives a replay from ``"0"``.
-    * ``stop`` cancels the worker, clears all state and closes every consumer;
-      a subsequent ``start`` begins from a clean slate.
+    * ``broadcast.push`` is synchronous and resolves all parked futures in one
+      pass, so all N subscribers wake atomically after a single call.
+    * State is updated under the lock; ``broadcast.push`` happens outside it so
+      the lock is held for as short a time as possible.
+    * ``close`` on the broadcast terminates all ``read_from`` generators, making
+      ``stop`` and terminal-status cleanup self-contained.
     """
 
-    def __init__(self, stream_reader: StreamReader):
+    def __init__(self, stream_reader: StreamReader) -> None:
         self.stream_reader = stream_reader
         self.consumers: dict[str, list[StreamConsumer]] = defaultdict(list)
         self.states: dict[str, StreamState] = {}
@@ -237,7 +316,7 @@ class StreamMultiplexer:
             logger.info("Stream multiplexer started")
 
     async def stop(self) -> None:
-        """Cancel the worker, clear all state, and close every consumer."""
+        """Cancel the worker, clear all state, and close every broadcast."""
         self.shutdown_event.set()
         task = self.worker_task
         self.worker_task = None
@@ -247,11 +326,12 @@ class StreamMultiplexer:
                 await task
 
         async with self._lock:
-            all_consumers = [c for cl in self.consumers.values() for c in cl]
+            states_to_close = list(self.states.values())
             self.consumers.clear()
             self.states.clear()
-        for consumer in all_consumers:
-            consumer.close()
+
+        for state in states_to_close:
+            state.broadcast.close()
 
     async def add_consumer(
         self,
@@ -259,47 +339,40 @@ class StreamMultiplexer:
         event_handler: EventHandler,
         last_id: str = "0",
     ) -> StreamConsumer:
-        """Register a consumer for *correlation_id*.
+        """Register a new subscriber for *correlation_id*.
 
-        If a consumer with the same handler type already exists its ref-count
-        is incremented and the existing consumer is returned (reconnect case).
-        Otherwise a new consumer is created.  When this is the first consumer
-        for the stream, the read position is initialised to *last_id*.
+        Each call returns a distinct ``StreamConsumer`` with its own cursor so
+        all concurrent subscribers receive every event independently (true
+        fan-out).  When this is the first consumer for the stream the read
+        position is initialised to *last_id*.
         """
         async with self._lock:
-            for consumer in self.consumers.get(correlation_id, []):
-                if type(consumer.event_handler) is type(event_handler):
-                    consumer.ref_count += 1
-                    return consumer
-
-            consumer = StreamConsumer(correlation_id, event_handler)
-            self.consumers[correlation_id].append(consumer)
-
             if correlation_id not in self.states:
                 self.states[correlation_id] = StreamState(last_id=last_id)
 
+            state = self.states[correlation_id]
+            cursor = state.broadcast.next_cursor()
+            consumer = StreamConsumer(
+                correlation_id, event_handler, state.broadcast, cursor
+            )
+            self.consumers[correlation_id].append(consumer)
             return consumer
 
     async def remove_consumer(self, consumer: StreamConsumer) -> None:
-        should_close = False
+        broadcast_to_close: StreamBroadcast | None = None
         async with self._lock:
             correlation_id = consumer.correlation_id
-            consumer.ref_count -= 1
+            cl = self.consumers.get(correlation_id)
+            if cl and consumer in cl:
+                cl.remove(consumer)
+            if correlation_id in self.consumers and not self.consumers[correlation_id]:
+                del self.consumers[correlation_id]
+                state = self.states.pop(correlation_id, None)
+                if state:
+                    broadcast_to_close = state.broadcast
 
-            if consumer.ref_count <= 0:
-                cl = self.consumers.get(correlation_id)
-                if cl and consumer in cl:
-                    cl.remove(consumer)
-                if (
-                    correlation_id in self.consumers
-                    and not self.consumers[correlation_id]
-                ):
-                    del self.consumers[correlation_id]
-                    self.states.pop(correlation_id, None)
-                should_close = True
-
-        if should_close:
-            consumer.close()
+        if broadcast_to_close is not None:
+            broadcast_to_close.close()
 
     async def _run_worker(self) -> None:
         try:
@@ -317,9 +390,12 @@ class StreamMultiplexer:
         except Exception as e:
             logger.error(f"Worker error: {e}", exc_info=True)
         finally:
-            for cl in list(self.consumers.values()):
-                for consumer in cl:
-                    consumer.close()
+            async with self._lock:
+                states_to_close = list(self.states.values())
+                self.consumers.clear()
+                self.states.clear()
+            for state in states_to_close:
+                state.broadcast.close()
 
     async def _process_stream(self, correlation_id: str) -> None:
         try:
@@ -330,9 +406,10 @@ class StreamMultiplexer:
                 consumer_list = list(self.consumers.get(correlation_id, []))
                 if not consumer_list:
                     return
+                event_handler = consumer_list[0].event_handler
 
             events, new_last_id = await self.stream_reader.read_events(
-                event_handler=consumer_list[0].event_handler,
+                event_handler=event_handler,
                 correlation_id=correlation_id,
                 last_id=state.last_id,
                 block_ms=DEFAULT_BLOCK_MS,
@@ -349,11 +426,11 @@ class StreamMultiplexer:
                 state.last_status_check = current_time
                 if await self.stream_reader.check_terminal_status(
                     correlation_id,
-                    consumer_list[0].event_handler,
+                    event_handler,
                     state.cached_events,
                 ):
                     drain_events, drain_last_id = await self.stream_reader.read_events(
-                        event_handler=consumer_list[0].event_handler,
+                        event_handler=event_handler,
                         correlation_id=correlation_id,
                         last_id=state.last_id,
                         block_ms=None,
@@ -380,43 +457,45 @@ class StreamMultiplexer:
         new_last_id: str,
         current_time: float,
     ) -> None:
-        """Dispatch *events* to every current consumer and advance ``last_id``.
+        """Advance stream state then broadcast *events* to all subscribers.
 
-        The consumer list is re-fetched under the lock so that consumers that
-        joined or left during the preceding blocking ``read_events`` are
-        reflected.  Dispatch and ``last_id`` advance are atomic: no consumer
-        can join between the two and miss the batch.
+        State is updated under the lock; ``broadcast.push`` is synchronous and
+        happens outside so the lock is held for the minimum time.
         """
         async with self._lock:
             if correlation_id not in self.states:
                 return
-            consumer_list = list(self.consumers.get(correlation_id, []))
-            if not consumer_list:
+            if not self.consumers.get(correlation_id):
                 return
-            for consumer in consumer_list:
-                for event in events:
-                    consumer.queue.put_nowait(event)
             state.last_id = new_last_id
             state.cached_events = events
             state.last_status_check = current_time
 
+        state.broadcast.push(events)
+
     async def _close_all_consumers(self, correlation_id: str) -> None:
         async with self._lock:
-            consumer_list = list(self.consumers.get(correlation_id, []))
+            state = self.states.pop(correlation_id, None)
             self.consumers.pop(correlation_id, None)
-            self.states.pop(correlation_id, None)
-        for consumer in consumer_list:
-            consumer.close()
+        if state:
+            state.broadcast.close()
 
 
 @singleton
 class AdaptiveStreamReader:
-    """Switches between the direct reader and the multiplexer based on load.
+    """Routes stream consumers between direct and multiplexed paths.
 
-    When the number of concurrently active streams reaches
-    ``multiplexing_threshold`` the multiplexed path is used; otherwise the
-    direct ``StreamReader.stream_events`` path is used.  Both paths must
-    produce identical results for the same stream and ``last_id``.
+    Direct path: one XREAD BLOCK per consumer → one Redis connection held.
+    Multiplexed path: one shared worker for all active streams on this process.
+
+    When concurrent direct-path readers reach ``multiplexing_threshold`` new
+    consumers are routed to the multiplexer instead of opening another Redis
+    connection.  When the multiplexer drains to zero active streams it is
+    stopped, so load-shedding is fully reversible — consumers that arrive when
+    load is back below the threshold go direct again.
+
+    ``multiplexing_threshold=None`` disables multiplexing entirely.
+    ``enable_multiplexing=False`` is a hard kill-switch (e.g. for tests).
     """
 
     @inject
@@ -424,13 +503,22 @@ class AdaptiveStreamReader:
         self,
         settings: Settings,
         stream_reader: StreamReader,
-    ):
+    ) -> None:
         self.stream_reader = stream_reader
         self.multiplexer: StreamMultiplexer | None = None
-        self._active_count = 0
         self._lock = asyncio.Lock()
         self.enable_multiplexing = True
-        self.multiplexing_threshold = settings.chat.multiplexing_threshold or None
+        self.multiplexing_threshold: int | None = settings.chat.multiplexing_threshold
+        # Number of consumers currently on the direct path.
+        # Each one holds an XREAD BLOCK connection; this is what we throttle.
+        self._direct_count = 0
+
+    def _should_multiplex(self) -> bool:
+        return (
+            self.enable_multiplexing
+            and self.multiplexing_threshold is not None
+            and self._direct_count >= self.multiplexing_threshold
+        )
 
     async def read_events(
         self,
@@ -457,16 +545,14 @@ class AdaptiveStreamReader:
         stop_event: asyncio.Event | None = None,
     ) -> AsyncGenerator[BaseModel, None]:
         async def gen() -> AsyncGenerator[BaseModel, None]:
+            # Decide path under the lock so _direct_count is consistent.
             async with self._lock:
-                self._active_count += 1
-                should_multiplex = (
-                    self.enable_multiplexing
-                    and self.multiplexing_threshold is not None
-                    and self._active_count >= self.multiplexing_threshold
-                )
+                use_mux = self._should_multiplex()
+                if not use_mux:
+                    self._direct_count += 1
 
             try:
-                if should_multiplex:
+                if use_mux:
                     async for event in await self._stream_multiplexed(
                         event_handler, correlation_id, last_id, stop_event
                     ):
@@ -477,8 +563,9 @@ class AdaptiveStreamReader:
                     ):
                         yield event
             finally:
-                async with self._lock:
-                    self._active_count -= 1
+                if not use_mux:
+                    async with self._lock:
+                        self._direct_count -= 1
 
         return gen()
 
@@ -490,41 +577,38 @@ class AdaptiveStreamReader:
         stop_event: asyncio.Event | None,
     ) -> AsyncGenerator[BaseModel, None]:
         async def gen() -> AsyncGenerator[BaseModel, None]:
-            if self.multiplexer is None:
-                self.multiplexer = StreamMultiplexer(self.stream_reader)
-                await self.multiplexer.start()
-                logger.info("Multiplexed streaming activated")
+            # Start the multiplexer on first use; capture the reference so the
+            # finally block targets the right instance even if it is replaced.
+            async with self._lock:
+                if self.multiplexer is None:
+                    self.multiplexer = StreamMultiplexer(self.stream_reader)
+                    await self.multiplexer.start()
+                    logger.info("Multiplexer started (threshold reached)")
+                mux = self.multiplexer
 
-            consumer = await self.multiplexer.add_consumer(
+            consumer = await mux.add_consumer(
                 correlation_id, event_handler, last_id=last_id
             )
-
             try:
-                events_processed = 0
-                while True:
-                    if stop_event and stop_event.is_set():
-                        break
-
-                    try:
-                        event: Any | None = await asyncio.wait_for(
-                            consumer.queue.get(), timeout=1.0
-                        )
-                        if event is None:
-                            break
-                        yield event
-
-                        events_processed += 1
-                        if events_processed % BATCH_SIZE == 0:
-                            await asyncio.sleep(0)
-
-                    except TimeoutError:
-                        continue
+                async for event in consumer.broadcast.read_from(
+                    cursor=consumer.cursor,
+                    stop_event=stop_event,
+                ):
+                    yield event
             finally:
-                await self.multiplexer.remove_consumer(consumer)
+                await mux.remove_consumer(consumer)
+                # Stop the multiplexer when it has no more active streams so the
+                # next consumer below the threshold can go direct again.
+                async with self._lock:
+                    if self.multiplexer is mux and not mux.states:
+                        await mux.stop()
+                        self.multiplexer = None
+                        logger.info("Multiplexer stopped (no active streams)")
 
         return gen()
 
     async def shutdown(self) -> None:
-        if self.multiplexer:
-            await self.multiplexer.stop()
-            self.multiplexer = None
+        async with self._lock:
+            mux, self.multiplexer = self.multiplexer, None
+        if mux:
+            await mux.stop()

tests/components/streaming/test_stream_multiplexer.py

@@ -38,6 +38,7 @@ from private_gpt.components.streaming.providers.redis_stream_service import (
 from private_gpt.components.streaming.providers.stream_service import StreamService
 from private_gpt.components.streaming.stream.stream_reader import (
     DEFAULT_BLOCK_MS,
+    StreamConsumer,
     StreamMultiplexer,
     StreamReader,
 )
@@ -76,16 +77,18 @@ class MockStreamComponent:
         self.stream = service
 
 
-async def collect_until_closed(
-    queue: asyncio.Queue[BaseModel | None],
+async def collect_from_consumer(
+    consumer: StreamConsumer,
     timeout: float = 10.0,
 ) -> list[BaseModel]:
+    """Collect all events from a consumer's broadcast until it is closed."""
     events: list[BaseModel] = []
-    while True:
-        event = await asyncio.wait_for(queue.get(), timeout=timeout)
-        if event is None:
-            break
-        events.append(event)
+
+    async def drain() -> None:
+        async for event in consumer.broadcast.read_from(cursor=consumer.cursor):
+            events.append(event)
+
+    await asyncio.wait_for(drain(), timeout=timeout)
     return events
 
 
@@ -113,7 +116,7 @@ async def run_multiplexed(
     last_id: str = "0",
 ) -> list[BaseModel]:
     consumer = await multiplexer.add_consumer(cid, handler, last_id=last_id)
-    return await collect_until_closed(consumer.queue, timeout=30.0)
+    return await collect_from_consumer(consumer, timeout=30.0)
 
 
 async def push_n_and_complete(
@@ -396,9 +399,9 @@ async def test_late_joiner_receives_inflight_batch(
     await service.update_stream_status(cid, StreamStatus.COMPLETED)
 
     try:
-        recv1 = await collect_until_closed(c1.queue, timeout=10.0)
+        recv1 = await collect_from_consumer(c1, timeout=10.0)
         c2 = late_box["c2"]
-        recv2 = await collect_until_closed(c2.queue, timeout=10.0)
+        recv2 = await collect_from_consumer(c2, timeout=10.0)
     finally:
         await multiplexer.stop()
         stream_reader.read_events = original_read  # type: ignore[method-assign]
@@ -450,7 +453,7 @@ async def test_concurrent_remove_preserves_delivery(
     await service.update_stream_status(cid, StreamStatus.COMPLETED)
 
     try:
-        recv1 = await collect_until_closed(c1.queue, timeout=5.0)
+        recv1 = await collect_from_consumer(c1, timeout=5.0)
     finally:
         await multiplexer.stop()
         stream_reader.read_events = original_read  # type: ignore[method-assign]
@@ -475,8 +478,8 @@ async def test_multiple_consumers_all_receive(
     c1 = await multiplexer.add_consumer(cid, handler)
     c2 = await multiplexer.add_consumer(cid, handler2)
     try:
-        recv1 = await collect_until_closed(c1.queue, timeout=10.0)
-        recv2 = await collect_until_closed(c2.queue, timeout=10.0)
+        recv1 = await collect_from_consumer(c1, timeout=10.0)
+        recv2 = await collect_from_consumer(c2, timeout=10.0)
     finally:
         await multiplexer.stop()
 
@@ -505,8 +508,6 @@ async def test_events_in_terminal_gap_drained(
         correlation_id: str,
         event_handler: Any,
         cached_events: Any,
-        last_flush_time: Any,
-        cache_flush_interval: Any,
     ) -> bool:
         await service.push_event(correlation_id, json.dumps({"n": 1}))
         await service.push_event(correlation_id, json.dumps({"n": 2}))
@@ -518,7 +519,7 @@ async def test_events_in_terminal_gap_drained(
     await multiplexer.start()
     consumer = await multiplexer.add_consumer(cid, handler)
     try:
-        received = await collect_until_closed(consumer.queue, timeout=10.0)
+        received = await collect_from_consumer(consumer, timeout=10.0)
     finally:
         await multiplexer.stop()
         stream_reader.check_terminal_status = original_check  # type: ignore[method-assign]
@@ -542,17 +543,19 @@ async def test_stop_clears_state_and_closes_consumers(
 
     await multiplexer.start()
     consumer = await multiplexer.add_consumer(cid, handler)
-    await asyncio.wait_for(consumer.queue.get(), timeout=5.0)
+
+    # Wait until the worker has pushed at least one batch to the broadcast.
+    async def get_one() -> None:
+        async for _ in consumer.broadcast.read_from(cursor=consumer.cursor):
+            break
+
+    await asyncio.wait_for(get_one(), timeout=5.0)
 
     await multiplexer.stop()
 
     assert not multiplexer.consumers, "consumers dict must be empty after stop()"
     assert not multiplexer.states, "states dict must be empty after stop()"
-
-    remaining: list[Any] = []
-    while not consumer.queue.empty():
-        remaining.append(consumer.queue.get_nowait())
-    assert None in remaining, "consumer queue must have a None sentinel"
+    assert consumer.broadcast.is_closed, "broadcast must be closed after stop()"
 
 
 async def test_restart_after_stop(
@@ -577,7 +580,7 @@ async def test_restart_after_stop(
     await multiplexer.start()
     try:
         consumer = await multiplexer.add_consumer(cid2, handler)
-        received = await collect_until_closed(consumer.queue, timeout=10.0)
+        received = await collect_from_consumer(consumer, timeout=10.0)
     finally:
         await multiplexer.stop()