langgraph-runtime-inmem 0.8.2__py3-none-any.whl → 0.10.0__py3-none-any.whl

langgraph_runtime_inmem/__init__.py

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

langgraph_runtime_inmem/inmem_stream.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import time
 from collections import defaultdict
 from collections.abc import Iterator
 from dataclasses import dataclass
@@ -12,6 +13,14 @@ 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
@@ -39,86 +48,186 @@ class ContextQueue(asyncio.Queue):
                 break
 
 
-class StreamManager:
-    def __init__(self):
-        self.queues = defaultdict(list)  # Dict[UUID, List[asyncio.Queue]]
-        self.control_keys = defaultdict()
-        self.control_queues = defaultdict(list)
+THREADLESS_KEY = "no-thread"
 
-        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]:
+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.thread_streams = 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)
-        return self.queues[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) -> list[asyncio.Queue]:
+    def get_control_queues(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> list[asyncio.Queue]:
         run_id = _ensure_uuid(run_id)
-        return self.control_queues[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) -> Message | None:
+    def get_control_key(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> Message | None:
         run_id = _ensure_uuid(run_id)
-        return self.control_keys.get(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, message: Message, resumable: bool = False
+        self,
+        run_id: UUID | str | None,
+        thread_id: UUID | str | None,
+        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 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[run_id].append(message)
+            self.message_stores[thread_id][run_id].append(message)
         topic = message.topic.decode()
         if "control" in topic:
-            self.control_keys[run_id] = message
-            queues = self.control_queues[run_id]
+            self.control_keys[thread_id][run_id] = message
+            queues = self.control_queues[thread_id][run_id]
         else:
-            queues = self.queues[run_id]
+            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) -> asyncio.Queue:
+    async def put_thread(
+        self,
+        thread_id: UUID | str,
+        message: Message,
+    ) -> None:
+        thread_id = _ensure_uuid(thread_id)
+        message.id = _generate_ms_seq_id().encode()
+        queues = self.thread_streams[thread_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()
-        self.queues[run_id].append(queue)
+        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) -> asyncio.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[run_id].append(queue)
+        self.control_queues[thread_id][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 add_thread_stream(self, thread_id: UUID | str) -> asyncio.Queue:
+        thread_id = _ensure_uuid(thread_id)
+        queue = ContextQueue()
+        self.thread_streams[thread_id].append(queue)
+        return queue
 
-    async def remove_control_queue(self, run_id: UUID | str, queue: asyncio.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 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]
+        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, message_id: str | None
+        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)
-        message_idx = int(message_id) + 1 if message_id else None
-
-        if message_idx is None:
-            yield from []
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+        if message_id is None:
             return
-
-        if run_id in self.message_stores:
-            yield from self.message_stores[run_id][message_idx:]
+        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

langgraph_runtime_inmem/ops.py

@@ -27,7 +27,12 @@ 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,
+    ContextQueue,
+    Message,
+    get_stream_manager,
+)
 
 if typing.TYPE_CHECKING:
     from langgraph_api.asyncio import ValueEvent
@@ -54,6 +59,7 @@ if typing.TYPE_CHECKING:
         Thread,
         ThreadSelectField,
         ThreadStatus,
+        ThreadStreamMode,
         ThreadUpdateResponse,
     )
     from langgraph_api.schema import Interrupt as InterruptSchema
@@ -734,6 +740,7 @@ class Threads(Authenticated):
     async def search(
         conn: InMemConnectionProto,
         *,
+        ids: list[str] | list[UUID] | None = None,
         metadata: MetadataInput,
         values: MetadataInput,
         status: ThreadStatus | None,
@@ -761,7 +768,19 @@ class Threads(Authenticated):
         )
 
         # Apply filters
+        id_set: set[UUID] | None = None
+        if ids:
+            id_set = set()
+            for i in ids:
+                try:
+                    id_set.add(_ensure_uuid(i))
+                except Exception:
+                    raise HTTPException(
+                        status_code=400, detail="Invalid thread ID " + str(i)
+                    ) from None
         for thread in threads:
+            if id_set is not None and thread.get("thread_id") not in id_set:
+                continue
             if filters and not _check_filter_match(thread["metadata"], filters):
                 continue
 
@@ -1323,7 +1342,14 @@ class Threads(Authenticated):
                 )
 
             metadata = thread.get("metadata", {})
-            thread_config = thread.get("config", {})
+            thread_config = cast(dict[str, Any], thread.get("config", {}))
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
 
             # Fallback to graph_id from run if not in thread metadata
             graph_id = metadata.get("graph_id")
@@ -1410,6 +1436,13 @@ class Threads(Authenticated):
                     status_code=409,
                     detail=f"Thread {thread_id} has in-flight runs: {pending_runs}",
                 )
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
 
             # Fallback to graph_id from run if not in thread metadata
             graph_id = metadata.get("graph_id")
@@ -1450,6 +1483,19 @@ class Threads(Authenticated):
                             thread["values"] = state.values
                             break
 
+                    # Publish state update event
+                    from langgraph_api.serde import json_dumpb
+
+                    event_data = {
+                        "state": state,
+                        "thread_id": str(thread_id),
+                    }
+                    await Threads.Stream.publish(
+                        thread_id,
+                        "state_update",
+                        json_dumpb(event_data),
+                    )
+
                     return ThreadUpdateResponse(
                         checkpoint=next_config["configurable"],
                         # Including deprecated fields
@@ -1492,7 +1538,14 @@ class Threads(Authenticated):
                 thread_iter, not_found_detail=f"Thread {thread_id} not found."
             )
 
-            thread_config = thread["config"]
+            thread_config = cast(dict[str, Any], thread["config"])
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
             metadata = thread["metadata"]
 
             if not thread:
@@ -1539,6 +1592,19 @@ class Threads(Authenticated):
                             thread["values"] = state.values
                             break
 
+                    # Publish state update event
+                    from langgraph_api.serde import json_dumpb
+
+                    event_data = {
+                        "state": state,
+                        "thread_id": str(thread_id),
+                    }
+                    await Threads.Stream.publish(
+                        thread_id,
+                        "state_update",
+                        json_dumpb(event_data),
+                    )
+
                     return ThreadUpdateResponse(
                         checkpoint=next_config["configurable"],
                     )
@@ -1580,7 +1646,14 @@ class Threads(Authenticated):
             if not _check_filter_match(thread_metadata, filters):
                 return []
 
-            thread_config = thread["config"]
+            thread_config = cast(dict[str, Any], thread["config"])
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
             # If graph_id exists, get state history
             if graph_id := thread_metadata.get("graph_id"):
                 async with get_graph(
@@ -1609,6 +1682,222 @@ 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)
+
+            # Add thread stream queue
+            if thread_id not in seen_runs:
+                queue = await stream_manager.add_thread_stream(thread_id)
+                queues.append((thread_id, queue))
+                seen_runs.add(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,
+            stream_modes: list[ThreadStreamMode],
+        ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
+            """Stream the thread output."""
+
+            def should_filter_event(event_name: str, message_bytes: bytes) -> bool:
+                """Check if an event should be filtered out based on stream_modes."""
+                if "run_modes" in stream_modes and event_name != "state_update":
+                    return False
+                if "state_update" in stream_modes and event_name == "state_update":
+                    return False
+                if "lifecycle" in stream_modes and event_name == "metadata":
+                    try:
+                        message_data = orjson.loads(message_bytes)
+                        if message_data.get("status") == "run_done":
+                            return False
+                        if "attempt" in message_data and "run_id" in message_data:
+                            return False
+                    except (orjson.JSONDecodeError, TypeError):
+                        pass
+                return True
+
+            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":
+                                    event_bytes = b"metadata"
+                                    message_bytes = orjson.dumps(
+                                        {"status": "run_done", "run_id": run_id}
+                                    )
+                                    # Filter events based on stream_modes
+                                    if not should_filter_event(
+                                        "metadata", message_bytes
+                                    ):
+                                        yield (
+                                            event_bytes,
+                                            message_bytes,
+                                            message.id,
+                                        )
+                            else:
+                                event_bytes = event_name.encode()
+                                message_bytes = base64.b64decode(message_content)
+                                # Filter events based on stream_modes
+                                if not should_filter_event(event_name, message_bytes):
+                                    yield (
+                                        event_bytes,
+                                        message_bytes,
+                                        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]
+                                        event_bytes = b"metadata"
+                                        message_bytes = orjson.dumps(
+                                            {"status": "run_done", "run_id": run_id}
+                                        )
+                                        # Filter events based on stream_modes
+                                        if not should_filter_event(
+                                            "metadata", message_bytes
+                                        ):
+                                            yield (
+                                                event_bytes,
+                                                message_bytes,
+                                                message.id,
+                                            )
+                                else:
+                                    event_bytes = event_name.encode()
+                                    message_bytes = base64.b64decode(message_content)
+                                    # Filter events based on stream_modes
+                                    if not should_filter_event(
+                                        event_name, message_bytes
+                                    ):
+                                        yield (
+                                            event_bytes,
+                                            message_bytes,
+                                            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 publish(
+            thread_id: UUID | str,
+            event: str,
+            message: bytes,
+        ) -> None:
+            """Publish a thread-level event to the thread stream."""
+            from langgraph_api.serde import json_dumpb
+
+            topic = f"thread:{thread_id}:stream".encode()
+
+            stream_manager = get_stream_manager()
+            # Send to thread stream topic
+            payload = json_dumpb(
+                {
+                    "event": event,
+                    "message": message,
+                }
+            )
+            await stream_manager.put_thread(
+                str(thread_id), Message(topic=topic, data=payload)
+            )
+
     @staticmethod
     async def count(
         conn: InMemConnectionProto,
@@ -1767,7 +2056,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.
@@ -1775,12 +2064,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
@@ -1788,17 +2079,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:
@@ -1853,6 +2144,7 @@ class Runs(Authenticated):
         run_id = _ensure_uuid(run_id) if run_id else None
         metadata = metadata if metadata is not None else {}
         config = kwargs.get("config", {})
+        temporary = kwargs.get("temporary", False)
 
         # Handle thread creation/update
         existing_thread = next(
@@ -1862,7 +2154,7 @@ class Runs(Authenticated):
             ctx,
             "create_run",
             Auth.types.RunsCreate(
-                thread_id=thread_id,
+                thread_id=None if temporary else thread_id,
                 assistant_id=assistant_id,
                 run_id=run_id,
                 status=status,
@@ -2086,6 +2378,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:
@@ -2268,9 +2561,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)
@@ -2385,15 +2678,25 @@ class Runs(Authenticated):
         @staticmethod
         async def subscribe(
             run_id: UUID,
-        ) -> asyncio.Queue:
+            thread_id: UUID | None = None,
+        ) -> ContextQueue:
             """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
@@ -2415,7 +2718,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:
@@ -2438,7 +2741,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
 
@@ -2529,7 +2832,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(
@@ -2537,6 +2840,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."""
@@ -2553,17 +2857,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.2.dist-info → langgraph_runtime_inmem-0.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-runtime-inmem
-Version: 0.8.2
+Version: 0.10.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.2.dist-info → langgraph_runtime_inmem-0.10.0.dist-info}/RECORD

@@ -1,13 +1,13 @@
-langgraph_runtime_inmem/__init__.py,sha256=AgGhozyDAnBy1osTaFV8oxSvO7Is7Rx0ASBE6XpUMDE,310
+langgraph_runtime_inmem/__init__.py,sha256=4xhdO3o6RduCHDXSNh42I51Wwq7Kcnt3JK1U1IhP-BU,311
 langgraph_runtime_inmem/checkpoint.py,sha256=nc1G8DqVdIu-ibjKTqXfbPfMbAsKjPObKqegrSzo6Po,4432
 langgraph_runtime_inmem/database.py,sha256=QgaA_WQo1IY6QioYd8r-e6-0B0rnC5anS0muIEJWby0,6364
-langgraph_runtime_inmem/inmem_stream.py,sha256=UWk1srLF44HZPPbRdArGGhsy0MY0UOJKSIxBSO7Hosc,5138
+langgraph_runtime_inmem/inmem_stream.py,sha256=utL1OlOJsy6VDkSGAA6eX9nETreZlM6K6nhfNoubmRQ,9011
 langgraph_runtime_inmem/lifespan.py,sha256=t0w2MX2dGxe8yNtSX97Z-d2pFpllSLS4s1rh2GJDw5M,3557
 langgraph_runtime_inmem/metrics.py,sha256=HhO0RC2bMDTDyGBNvnd2ooLebLA8P1u5oq978Kp_nAA,392
-langgraph_runtime_inmem/ops.py,sha256=Qtu4rSNd6uFjpYrrLxXMFbDDK4Z7PMYtqcRK_ikPmQA,97862
+langgraph_runtime_inmem/ops.py,sha256=54jiyWhfbSu9z9pca6AQdNuaIBmD0WMrQ7xGQcLPDF4,111183
 langgraph_runtime_inmem/queue.py,sha256=33qfFKPhQicZ1qiibllYb-bTFzUNSN2c4bffPACP5es,9952
 langgraph_runtime_inmem/retry.py,sha256=XmldOP4e_H5s264CagJRVnQMDFcEJR_dldVR1Hm5XvM,763
 langgraph_runtime_inmem/store.py,sha256=rTfL1JJvd-j4xjTrL8qDcynaWF6gUJ9-GDVwH0NBD_I,3506
-langgraph_runtime_inmem-0.8.2.dist-info/METADATA,sha256=gcjizfhnsByvf34DwFigAdThjuyZxVD09uMbA1yQB9U,565
-langgraph_runtime_inmem-0.8.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langgraph_runtime_inmem-0.8.2.dist-info/RECORD,,
+langgraph_runtime_inmem-0.10.0.dist-info/METADATA,sha256=gdjdQjZF2KjDtwA9rDiW53pG4FYNfv8TkT1U8t2lftQ,566
+langgraph_runtime_inmem-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langgraph_runtime_inmem-0.10.0.dist-info/RECORD,,