langgraph-runtime-inmem 0.8.2__py3-none-any.whl → 0.9.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.9.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,165 @@ 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.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,
+        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 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):
+    async def remove_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.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):
+        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,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
@@ -1609,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,
@@ -1767,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.
@@ -1775,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
@@ -1788,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:
@@ -2086,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:
@@ -2268,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)
@@ -2385,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
@@ -2415,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:
@@ -2438,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
 
@@ -2529,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(
@@ -2537,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."""
@@ -2553,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.2.dist-info → langgraph_runtime_inmem-0.9.0.dist-info}/METADATA

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

@@ -1,13 +1,13 @@
-langgraph_runtime_inmem/__init__.py,sha256=AgGhozyDAnBy1osTaFV8oxSvO7Is7Rx0ASBE6XpUMDE,310
+langgraph_runtime_inmem/__init__.py,sha256=f-VPPHH1-hKFwEreffg7dNATe9IdcYwQedcSx2MiZog,310
 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=pUEiHW-1uXQrVTcwEYPwO8YXaYm5qZbpRWawt67y6Lw,8187
 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=0Jx65S3PCvvHlIpA0XYpl-UnDEo_AiGWXRE2QiFSocY,105165
 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.9.0.dist-info/METADATA,sha256=ptwW1Ei-Xln53P81eJK1aPcFozU8D192OCZBuC_y5EQ,565
+langgraph_runtime_inmem-0.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langgraph_runtime_inmem-0.9.0.dist-info/RECORD,,