langgraph-runtime-inmem 0.8.1__tar.gz → 0.9.0__tar.gz

{langgraph_runtime_inmem-0.8.1 → langgraph_runtime_inmem-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-runtime-inmem
-Version: 0.8.1
+Version: 0.9.0
 Summary: Inmem implementation for the LangGraph API server.
 Author-email: Will Fu-Hinthorn <will@langchain.dev>
 License: Elastic-2.0

{langgraph_runtime_inmem-0.8.1 → langgraph_runtime_inmem-0.9.0}/langgraph_runtime_inmem/__init__.py

@@ -9,7 +9,7 @@ from langgraph_runtime_inmem import (
     store,
 )
 
-__version__ = "0.8.1"
+__version__ = "0.9.0"
 __all__ = [
     "ops",
     "database",

{langgraph_runtime_inmem-0.8.1 → langgraph_runtime_inmem-0.9.0}/langgraph_runtime_inmem/database.py

@@ -208,6 +208,6 @@ async def healthcheck() -> None:
     pass
 
 
-def pool_stats() -> dict[str, dict[str, int]]:
+def pool_stats(*args, **kwargs) -> dict[str, dict[str, int]]:
     # TODO??
     return {}

langgraph_runtime_inmem-0.9.0/langgraph_runtime_inmem/inmem_stream.py

@@ -0,0 +1,247 @@
+import asyncio
+import logging
+import time
+from collections import defaultdict
+from collections.abc import Iterator
+from dataclasses import dataclass
+from uuid import UUID
+
+logger = logging.getLogger(__name__)
+
+
+def _ensure_uuid(id: str | UUID) -> UUID:
+    return UUID(id) if isinstance(id, str) else id
+
+
+def _generate_ms_seq_id() -> str:
+    """Generate a Redis-like millisecond-sequence ID (e.g., '1234567890123-0')"""
+    # Get current time in milliseconds
+    ms = int(time.time() * 1000)
+    # For simplicity, always use sequence 0 since we're not handling high throughput
+    return f"{ms}-0"
+
+
+@dataclass
+class Message:
+    topic: bytes
+    data: bytes
+    id: bytes | None = None
+
+
+class ContextQueue(asyncio.Queue):
+    """Queue that supports async context manager protocol"""
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object | None,
+    ) -> None:
+        # Clear the queue
+        while not self.empty():
+            try:
+                self.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+
+
+THREADLESS_KEY = "no-thread"
+
+
+class StreamManager:
+    def __init__(self):
+        self.queues = defaultdict(
+            lambda: defaultdict(list)
+        )  # Dict[str, List[asyncio.Queue]]
+        self.control_keys = defaultdict(lambda: defaultdict())
+        self.control_queues = defaultdict(lambda: defaultdict(list))
+
+        self.message_stores = defaultdict(
+            lambda: defaultdict(list[Message])
+        )  # Dict[str, List[Message]]
+
+    def get_queues(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> list[asyncio.Queue]:
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        return self.queues[thread_id][run_id]
+
+    def get_control_queues(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> list[asyncio.Queue]:
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        return self.control_queues[thread_id][run_id]
+
+    def get_control_key(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> Message | None:
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        return self.control_keys.get(thread_id, {}).get(run_id)
+
+    async def put(
+        self,
+        run_id: UUID | str,
+        thread_id: UUID | str | None,
+        message: Message,
+        resumable: bool = False,
+    ) -> None:
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+
+        message.id = _generate_ms_seq_id().encode()
+        if resumable:
+            self.message_stores[thread_id][run_id].append(message)
+        topic = message.topic.decode()
+        if "control" in topic:
+            self.control_keys[thread_id][run_id] = message
+            queues = self.control_queues[thread_id][run_id]
+        else:
+            queues = self.queues[thread_id][run_id]
+        coros = [queue.put(message) for queue in queues]
+        results = await asyncio.gather(*coros, return_exceptions=True)
+        for result in results:
+            if isinstance(result, Exception):
+                logger.exception(f"Failed to put message in queue: {result}")
+
+    async def add_queue(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> asyncio.Queue:
+        run_id = _ensure_uuid(run_id)
+        queue = ContextQueue()
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        self.queues[thread_id][run_id].append(queue)
+        return queue
+
+    async def add_control_queue(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> asyncio.Queue:
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        queue = ContextQueue()
+        self.control_queues[thread_id][run_id].append(queue)
+        return queue
+
+    async def remove_queue(
+        self, run_id: UUID | str, thread_id: UUID | str | None, queue: asyncio.Queue
+    ):
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        if thread_id in self.queues and run_id in self.queues[thread_id]:
+            self.queues[thread_id][run_id].remove(queue)
+            if not self.queues[thread_id][run_id]:
+                del self.queues[thread_id][run_id]
+
+    async def remove_control_queue(
+        self, run_id: UUID | str, thread_id: UUID | str | None, queue: asyncio.Queue
+    ):
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        if (
+            thread_id in self.control_queues
+            and run_id in self.control_queues[thread_id]
+        ):
+            self.control_queues[thread_id][run_id].remove(queue)
+            if not self.control_queues[thread_id][run_id]:
+                del self.control_queues[thread_id][run_id]
+
+    def restore_messages(
+        self, run_id: UUID | str, thread_id: UUID | str | None, message_id: str | None
+    ) -> Iterator[Message]:
+        """Get a stored message by ID for resumable streams."""
+        run_id = _ensure_uuid(run_id)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        if message_id is None:
+            return
+        try:
+            # Handle ms-seq format (e.g., "1234567890123-0")
+            if thread_id in self.message_stores:
+                for message in self.message_stores[thread_id][run_id]:
+                    if message.id.decode() > message_id:
+                        yield message
+        except TypeError:
+            # Try integer format if ms-seq fails
+            message_idx = int(message_id) + 1
+            if run_id in self.message_stores:
+                yield from self.message_stores[thread_id][run_id][message_idx:]
+
+    def get_queues_by_thread_id(self, thread_id: UUID | str) -> list[asyncio.Queue]:
+        """Get all queues for a specific thread_id across all runs."""
+        all_queues = []
+        # Search through all stored queue keys for ones ending with the thread_id
+        thread_id = _ensure_uuid(thread_id)
+        if thread_id in self.queues:
+            for run_id in self.queues[thread_id]:
+                all_queues.extend(self.queues[thread_id][run_id])
+
+        return all_queues
+
+
+# Global instance
+stream_manager = StreamManager()
+
+
+async def start_stream() -> None:
+    """Initialize the queue system.
+    In this in-memory implementation, we just need to ensure we have a clean StreamManager instance.
+    """
+    global stream_manager
+    stream_manager = StreamManager()
+
+
+async def stop_stream() -> None:
+    """Clean up the queue system.
+    Clear all queues and stored control messages."""
+    global stream_manager
+
+    # Send 'done' message to all active queues before clearing
+    for run_id in list(stream_manager.queues.keys()):
+        control_message = Message(topic=f"run:{run_id}:control".encode(), data=b"done")
+
+        for queue in stream_manager.queues[run_id]:
+            try:
+                await queue.put(control_message)
+            except (Exception, RuntimeError):
+                pass  # Ignore errors during shutdown
+
+    # Clear all stored data
+    stream_manager.queues.clear()
+    stream_manager.control_queues.clear()
+    stream_manager.message_stores.clear()
+
+
+def get_stream_manager() -> StreamManager:
+    """Get the global stream manager instance."""
+    return stream_manager

{langgraph_runtime_inmem-0.8.1 → langgraph_runtime_inmem-0.9.0}/langgraph_runtime_inmem/ops.py

@@ -27,7 +27,11 @@ from starlette.exceptions import HTTPException
 
 from langgraph_runtime_inmem.checkpoint import Checkpointer
 from langgraph_runtime_inmem.database import InMemConnectionProto, connect
-from langgraph_runtime_inmem.inmem_stream import Message, get_stream_manager
+from langgraph_runtime_inmem.inmem_stream import (
+    THREADLESS_KEY,
+    Message,
+    get_stream_manager,
+)
 
 if typing.TYPE_CHECKING:
     from langgraph_api.asyncio import ValueEvent
@@ -406,19 +410,17 @@ class Assistants(Authenticated):
             else 1
         )
 
-        # Update assistant_versions table
-        if metadata:
-            metadata = {
-                **assistant["metadata"],
-                **metadata,
-            }
         new_version_entry = {
             "assistant_id": assistant_id,
             "version": new_version,
             "graph_id": graph_id if graph_id is not None else assistant["graph_id"],
             "config": config if config else assistant["config"],
             "context": context if context is not None else assistant.get("context", {}),
-            "metadata": metadata if metadata is not None else assistant["metadata"],
+            "metadata": (
+                {**assistant["metadata"], **metadata}
+                if metadata is not None
+                else assistant["metadata"]
+            ),
             "created_at": now,
             "name": name if name is not None else assistant["name"],
             "description": (
@@ -1611,6 +1613,151 @@ class Threads(Authenticated):
 
             return []
 
+    class Stream:
+        @staticmethod
+        async def subscribe(
+            conn: InMemConnectionProto | AsyncConnectionProto,
+            thread_id: UUID,
+            seen_runs: set[UUID],
+        ) -> list[tuple[UUID, asyncio.Queue]]:
+            """Subscribe to the thread stream, creating queues for unseen runs."""
+            stream_manager = get_stream_manager()
+            queues = []
+
+            # Create new queues only for runs not yet seen
+            thread_id = _ensure_uuid(thread_id)
+            for run in conn.store["runs"]:
+                if run["thread_id"] == thread_id:
+                    run_id = run["run_id"]
+                    if run_id not in seen_runs:
+                        queue = await stream_manager.add_queue(run_id, thread_id)
+                        queues.append((run_id, queue))
+                        seen_runs.add(run_id)
+
+            return queues
+
+        @staticmethod
+        async def join(
+            thread_id: UUID,
+            *,
+            last_event_id: str | None = None,
+        ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
+            """Stream the thread output."""
+            from langgraph_api.serde import json_loads
+
+            stream_manager = get_stream_manager()
+            seen_runs: set[UUID] = set()
+            created_queues: list[tuple[UUID, asyncio.Queue]] = []
+
+            try:
+                async with connect() as conn:
+                    await logger.ainfo(
+                        "Joined thread stream",
+                        thread_id=str(thread_id),
+                    )
+
+                    # Restore messages if resuming from a specific event
+                    if last_event_id is not None:
+                        # Collect all events from all message stores for this thread
+                        all_events = []
+                        for run_id in stream_manager.message_stores.get(
+                            str(thread_id), []
+                        ):
+                            for message in stream_manager.restore_messages(
+                                run_id, thread_id, last_event_id
+                            ):
+                                all_events.append((message, run_id))
+
+                        # Sort by message ID (which is ms-seq format)
+                        all_events.sort(key=lambda x: x[0].id.decode())
+
+                        # Yield sorted events
+                        for message, run_id in all_events:
+                            data = json_loads(message.data)
+                            event_name = data["event"]
+                            message_content = data["message"]
+
+                            if event_name == "control":
+                                if message_content == b"done":
+                                    yield (
+                                        b"metadata",
+                                        orjson.dumps(
+                                            {"status": "run_done", "run_id": run_id}
+                                        ),
+                                        message.id,
+                                    )
+                            else:
+                                yield (
+                                    event_name.encode(),
+                                    base64.b64decode(message_content),
+                                    message.id,
+                                )
+
+                    # Listen for live messages from all queues
+                    while True:
+                        # Refresh queues to pick up any new runs that joined this thread
+                        new_queue_tuples = await Threads.Stream.subscribe(
+                            conn, thread_id, seen_runs
+                        )
+                        # Track new queues for cleanup
+                        for run_id, queue in new_queue_tuples:
+                            created_queues.append((run_id, queue))
+
+                        for run_id, queue in created_queues:
+                            try:
+                                message = await asyncio.wait_for(
+                                    queue.get(), timeout=0.2
+                                )
+                                data = json_loads(message.data)
+                                event_name = data["event"]
+                                message_content = data["message"]
+
+                                if event_name == "control":
+                                    if message_content == b"done":
+                                        # Extract run_id from topic
+                                        topic = message.topic.decode()
+                                        run_id = topic.split("run:")[1].split(":")[0]
+                                        yield (
+                                            b"metadata",
+                                            orjson.dumps(
+                                                {"status": "run_done", "run_id": run_id}
+                                            ),
+                                            message.id,
+                                        )
+                                else:
+                                    yield (
+                                        event_name.encode(),
+                                        base64.b64decode(message_content),
+                                        message.id,
+                                    )
+
+                            except TimeoutError:
+                                continue
+                            except (ValueError, KeyError):
+                                continue
+
+                        # Yield execution to other tasks to prevent event loop starvation
+                        await asyncio.sleep(0)
+
+            except WrappedHTTPException as e:
+                raise e.http_exception from None
+            except asyncio.CancelledError:
+                await logger.awarning(
+                    "Thread stream client disconnected",
+                    thread_id=str(thread_id),
+                )
+                raise
+            except:
+                raise
+            finally:
+                # Clean up all created queues
+                for run_id, queue in created_queues:
+                    try:
+                        await stream_manager.remove_queue(run_id, thread_id, queue)
+                    except Exception:
+                        # Ignore cleanup errors
+                        pass
+
     @staticmethod
     async def count(
         conn: InMemConnectionProto,
@@ -1769,7 +1916,7 @@ class Runs(Authenticated):
     @asynccontextmanager
     @staticmethod
     async def enter(
-        run_id: UUID, loop: asyncio.AbstractEventLoop
+        run_id: UUID, thread_id: UUID | None, loop: asyncio.AbstractEventLoop
     ) -> AsyncIterator[ValueEvent]:
         """Enter a run, listen for cancellation while running, signal when done."
         This method should be called as a context manager by a worker executing a run.
@@ -1777,12 +1924,14 @@ class Runs(Authenticated):
         from langgraph_api.asyncio import SimpleTaskGroup, ValueEvent
 
         stream_manager = get_stream_manager()
-        # Get queue for this run
-        queue = await stream_manager.add_control_queue(run_id)
+        # Get control queue for this run (normal queue is created during run creation)
+        control_queue = await stream_manager.add_control_queue(run_id, thread_id)
 
         async with SimpleTaskGroup(cancel=True, taskgroup_name="Runs.enter") as tg:
             done = ValueEvent()
-            tg.create_task(listen_for_cancellation(queue, run_id, done))
+            tg.create_task(
+                listen_for_cancellation(control_queue, run_id, thread_id, done)
+            )
 
             # Give done event to caller
             yield done
@@ -1790,17 +1939,17 @@ class Runs(Authenticated):
             control_message = Message(
                 topic=f"run:{run_id}:control".encode(), data=b"done"
             )
-            await stream_manager.put(run_id, control_message)
+            await stream_manager.put(run_id, thread_id, control_message)
 
             # Signal done to all subscribers
             stream_message = Message(
                 topic=f"run:{run_id}:stream".encode(),
                 data={"event": "control", "message": b"done"},
             )
-            await stream_manager.put(run_id, stream_message)
+            await stream_manager.put(run_id, thread_id, stream_message)
 
-            # Remove the queue
-            await stream_manager.remove_control_queue(run_id, queue)
+            # Remove the control_queue (normal queue is cleaned up during run deletion)
+            await stream_manager.remove_control_queue(run_id, thread_id, control_queue)
 
     @staticmethod
     async def sweep() -> None:
@@ -2088,6 +2237,7 @@ class Runs(Authenticated):
             if not thread:
                 return _empty_generator()
         _delete_checkpoints_for_thread(thread_id, conn, run_id=run_id)
+
         found = False
         for i, run in enumerate(conn.store["runs"]):
             if run["run_id"] == run_id and run["thread_id"] == thread_id:
@@ -2270,9 +2420,9 @@ class Runs(Authenticated):
                 topic=f"run:{run_id}:control".encode(),
                 data=action.encode(),
             )
-            coros.append(stream_manager.put(run_id, control_message))
+            coros.append(stream_manager.put(run_id, thread_id, control_message))
 
-            queues = stream_manager.get_queues(run_id)
+            queues = stream_manager.get_queues(run_id, thread_id)
 
             if run["status"] in ("pending", "running"):
                 cancelable_runs.append(run)
@@ -2387,15 +2537,25 @@ class Runs(Authenticated):
         @staticmethod
         async def subscribe(
             run_id: UUID,
+            thread_id: UUID | None = None,
         ) -> asyncio.Queue:
             """Subscribe to the run stream, returning a queue."""
             stream_manager = get_stream_manager()
-            queue = await stream_manager.add_queue(_ensure_uuid(run_id))
+            queue = await stream_manager.add_queue(_ensure_uuid(run_id), thread_id)
 
             # If there's a control message already stored, send it to the new subscriber
-            if control_messages := stream_manager.control_queues.get(run_id):
-                for control_msg in control_messages:
-                    await queue.put(control_msg)
+            if thread_id is None:
+                thread_id = THREADLESS_KEY
+            if control_queues := stream_manager.control_queues.get(thread_id, {}).get(
+                run_id
+            ):
+                for control_queue in control_queues:
+                    try:
+                        while True:
+                            control_msg = control_queue.get()
+                            await queue.put(control_msg)
+                    except asyncio.QueueEmpty:
+                        pass
             return queue
 
         @staticmethod
@@ -2417,7 +2577,7 @@ class Runs(Authenticated):
             queue = (
                 stream_channel
                 if stream_channel
-                else await Runs.Stream.subscribe(run_id)
+                else await Runs.Stream.subscribe(run_id, thread_id)
             )
 
             try:
@@ -2440,7 +2600,7 @@ class Runs(Authenticated):
                     run = await Runs.get(conn, run_id, thread_id=thread_id, ctx=ctx)
 
                     for message in get_stream_manager().restore_messages(
-                        run_id, last_event_id
+                        run_id, thread_id, last_event_id
                     ):
                         data, id = message.data, message.id
 
@@ -2531,7 +2691,7 @@ class Runs(Authenticated):
                 raise
             finally:
                 stream_manager = get_stream_manager()
-                await stream_manager.remove_queue(run_id, queue)
+                await stream_manager.remove_queue(run_id, thread_id, queue)
 
         @staticmethod
         async def publish(
@@ -2539,6 +2699,7 @@ class Runs(Authenticated):
             event: str,
             message: bytes,
             *,
+            thread_id: UUID | str | None = None,
             resumable: bool = False,
         ) -> None:
             """Publish a message to all subscribers of the run stream."""
@@ -2555,17 +2716,19 @@ class Runs(Authenticated):
                 }
             )
             await stream_manager.put(
-                run_id, Message(topic=topic, data=payload), resumable
+                run_id, thread_id, Message(topic=topic, data=payload), resumable
             )
 
 
-async def listen_for_cancellation(queue: asyncio.Queue, run_id: UUID, done: ValueEvent):
+async def listen_for_cancellation(
+    queue: asyncio.Queue, run_id: UUID, thread_id: UUID | None, done: ValueEvent
+):
     """Listen for cancellation messages and set the done event accordingly."""
     from langgraph_api.errors import UserInterrupt, UserRollback
 
     stream_manager = get_stream_manager()
 
-    if control_key := stream_manager.get_control_key(run_id):
+    if control_key := stream_manager.get_control_key(run_id, thread_id):
         payload = control_key.data
         if payload == b"rollback":
             done.set(UserRollback())

langgraph_runtime_inmem-0.8.1/langgraph_runtime_inmem/inmem_stream.py

@@ -1,159 +0,0 @@
-import asyncio
-import logging
-from collections import defaultdict
-from collections.abc import Iterator
-from dataclasses import dataclass
-from uuid import UUID
-
-logger = logging.getLogger(__name__)
-
-
-def _ensure_uuid(id: str | UUID) -> UUID:
-    return UUID(id) if isinstance(id, str) else id
-
-
-@dataclass
-class Message:
-    topic: bytes
-    data: bytes
-    id: bytes | None = None
-
-
-class ContextQueue(asyncio.Queue):
-    """Queue that supports async context manager protocol"""
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: object | None,
-    ) -> None:
-        # Clear the queue
-        while not self.empty():
-            try:
-                self.get_nowait()
-            except asyncio.QueueEmpty:
-                break
-
-
-class StreamManager:
-    def __init__(self):
-        self.queues = defaultdict(list)  # Dict[UUID, List[asyncio.Queue]]
-        self.control_keys = defaultdict()
-        self.control_queues = defaultdict(list)
-
-        self.message_stores = defaultdict(list)  # Dict[UUID, List[Message]]
-        self.message_next_idx = defaultdict(int)  # Dict[UUID, int]
-
-    def get_queues(self, run_id: UUID | str) -> list[asyncio.Queue]:
-        run_id = _ensure_uuid(run_id)
-        return self.queues[run_id]
-
-    def get_control_queues(self, run_id: UUID | str) -> list[asyncio.Queue]:
-        run_id = _ensure_uuid(run_id)
-        return self.control_queues[run_id]
-
-    def get_control_key(self, run_id: UUID | str) -> Message | None:
-        run_id = _ensure_uuid(run_id)
-        return self.control_keys.get(run_id)
-
-    async def put(
-        self, run_id: UUID | str, message: Message, resumable: bool = False
-    ) -> None:
-        run_id = _ensure_uuid(run_id)
-        message.id = str(self.message_next_idx[run_id]).encode()
-        self.message_next_idx[run_id] += 1
-        if resumable:
-            self.message_stores[run_id].append(message)
-        topic = message.topic.decode()
-        if "control" in topic:
-            self.control_keys[run_id] = message
-            queues = self.control_queues[run_id]
-        else:
-            queues = self.queues[run_id]
-        coros = [queue.put(message) for queue in queues]
-        results = await asyncio.gather(*coros, return_exceptions=True)
-        for result in results:
-            if isinstance(result, Exception):
-                logger.exception(f"Failed to put message in queue: {result}")
-
-    async def add_queue(self, run_id: UUID | str) -> asyncio.Queue:
-        run_id = _ensure_uuid(run_id)
-        queue = ContextQueue()
-        self.queues[run_id].append(queue)
-        return queue
-
-    async def add_control_queue(self, run_id: UUID | str) -> asyncio.Queue:
-        run_id = _ensure_uuid(run_id)
-        queue = ContextQueue()
-        self.control_queues[run_id].append(queue)
-        return queue
-
-    async def remove_queue(self, run_id: UUID | str, queue: asyncio.Queue):
-        run_id = _ensure_uuid(run_id)
-        if run_id in self.queues:
-            self.queues[run_id].remove(queue)
-            if not self.queues[run_id]:
-                del self.queues[run_id]
-
-    async def remove_control_queue(self, run_id: UUID | str, queue: asyncio.Queue):
-        run_id = _ensure_uuid(run_id)
-        if run_id in self.control_queues:
-            self.control_queues[run_id].remove(queue)
-            if not self.control_queues[run_id]:
-                del self.control_queues[run_id]
-
-    def restore_messages(
-        self, run_id: UUID | str, message_id: str | None
-    ) -> Iterator[Message]:
-        """Get a stored message by ID for resumable streams."""
-        run_id = _ensure_uuid(run_id)
-        message_idx = int(message_id) + 1 if message_id else None
-
-        if message_idx is None:
-            yield from []
-            return
-
-        if run_id in self.message_stores:
-            yield from self.message_stores[run_id][message_idx:]
-
-
-# Global instance
-stream_manager = StreamManager()
-
-
-async def start_stream() -> None:
-    """Initialize the queue system.
-    In this in-memory implementation, we just need to ensure we have a clean StreamManager instance.
-    """
-    global stream_manager
-    stream_manager = StreamManager()
-
-
-async def stop_stream() -> None:
-    """Clean up the queue system.
-    Clear all queues and stored control messages."""
-    global stream_manager
-
-    # Send 'done' message to all active queues before clearing
-    for run_id in list(stream_manager.queues.keys()):
-        control_message = Message(topic=f"run:{run_id}:control".encode(), data=b"done")
-
-        for queue in stream_manager.queues[run_id]:
-            try:
-                await queue.put(control_message)
-            except (Exception, RuntimeError):
-                pass  # Ignore errors during shutdown
-
-    # Clear all stored data
-    stream_manager.queues.clear()
-    stream_manager.control_queues.clear()
-    stream_manager.message_stores.clear()
-
-
-def get_stream_manager() -> StreamManager:
-    """Get the global stream manager instance."""
-    return stream_manager