langgraph-runtime-inmem 0.6.4__py3-none-any.whl → 0.18.1__py3-none-any.whl

langgraph_runtime_inmem/__init__.py

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

langgraph_runtime_inmem/database.py

@@ -142,7 +142,9 @@ class InMemConnectionProto:
 
 
 @asynccontextmanager
-async def connect(*, __test__: bool = False) -> AsyncIterator["AsyncConnectionProto"]:
+async def connect(
+    *, supports_core_api: bool = False, __test__: bool = False
+) -> AsyncIterator["AsyncConnectionProto"]:
     yield InMemConnectionProto()
 
 
@@ -182,6 +184,8 @@ async def start_pool() -> None:
         for a in GLOBAL_STORE["assistants"]:
             if a["metadata"].get("created_by") == "system":
                 GLOBAL_STORE["assistants"].remove(a)
+            if "context" not in a:
+                a["context"] = {}
     for k in ["crons"]:
         if not GLOBAL_STORE.get(k):
             GLOBAL_STORE[k] = {}
@@ -206,6 +210,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/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,72 +48,187 @@ class ContextQueue(asyncio.Queue):
                 break
 
 
+THREADLESS_KEY = "no-thread"
+
+
 class StreamManager:
     def __init__(self):
-        self.queues = defaultdict(list)  # Dict[UUID, List[asyncio.Queue]]
-        self.control_queues = defaultdict(list)
+        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(list)  # Dict[UUID, List[Message]]
-        self.message_next_idx = defaultdict(int)  # Dict[UUID, int]
+        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_queues(self, run_id: UUID | str) -> list[asyncio.Queue]:
+    def get_control_key(
+        self, run_id: UUID | str, thread_id: UUID | str | None
+    ) -> Message | None:
         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.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 resumable:
-            self.message_stores[run_id].append(message)
+        if thread_id is None:
+            thread_id = THREADLESS_KEY
+        else:
+            thread_id = _ensure_uuid(thread_id)
+
+        message.id = _generate_ms_seq_id().encode()
+        # For resumable run streams, embed the generated message ID into the frame
         topic = message.topic.decode()
+        if resumable:
+            self.message_stores[thread_id][run_id].append(message)
         if "control" in topic:
-            self.control_queues[run_id].append(message)
-        queues = self.queues.get(run_id, [])
+            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) -> 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)
-        for control_msg in self.control_queues[run_id]:
-            try:
-                await queue.put(control_msg)
-            except Exception:
-                logger.exception(
-                    f"Failed to put control message in queue: {control_msg}"
-                )
+        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, queue: asyncio.Queue):
+    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_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]
-                if run_id in self.message_stores:
-                    del self.message_stores[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, 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
+        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])
 
-        if run_id in self.message_stores:
-            yield from self.message_stores[run_id][message_idx:]
+        return all_queues
 
 
 # Global instance

langgraph_runtime_inmem/lifespan.py

@@ -14,9 +14,17 @@ from langgraph_runtime_inmem.database import start_pool, stop_pool
 logger = structlog.stdlib.get_logger(__name__)
 
 
+_LAST_LIFESPAN_ERROR: BaseException | None = None
+
+
+def get_last_error() -> BaseException | None:
+    return _LAST_LIFESPAN_ERROR
+
+
 @asynccontextmanager
 async def lifespan(
     app: Starlette | None = None,
+    cancel_event: asyncio.Event | None = None,
     taskset: set[asyncio.Task] | None = None,
     **kwargs: Any,
 ):
@@ -41,13 +49,31 @@ async def lifespan(
     except RuntimeError:
         await logger.aerror("Failed to set loop")
 
+    global _LAST_LIFESPAN_ERROR
+    _LAST_LIFESPAN_ERROR = None
+
     await start_http_client()
     await start_pool()
     await start_ui_bundler()
+
+    async def _log_graph_load_failure(err: graph.GraphLoadError) -> None:
+        cause = err.__cause__ or err.cause
+        log_fields = err.log_fields()
+        log_fields["action"] = "fix_user_graph"
+        await logger.aerror(
+            f"Graph '{err.spec.id}' failed to load: {err.cause_message}",
+            **log_fields,
+        )
+        await logger.adebug(
+            "Full graph load failure traceback (internal)",
+            **{k: v for k, v in log_fields.items() if k != "user_traceback"},
+            exc_info=cause,
+        )
+
     try:
         async with SimpleTaskGroup(
             cancel=True,
-            taskset=taskset,
+            cancel_event=cancel_event,
             taskgroup_name="Lifespan",
         ) as tg:
             tg.create_task(metadata_loop())
@@ -76,11 +102,21 @@ async def lifespan(
             var_child_runnable_config.set(langgraph_config)
 
             # Keep after the setter above so users can access the store from within the factory function
-            await graph.collect_graphs_from_env(True)
+            try:
+                await graph.collect_graphs_from_env(True)
+            except graph.GraphLoadError as exc:
+                _LAST_LIFESPAN_ERROR = exc
+                await _log_graph_load_failure(exc)
+                raise
             if config.N_JOBS_PER_WORKER > 0:
                 tg.create_task(queue_with_signal())
 
             yield
+    except graph.GraphLoadError as exc:
+        _LAST_LIFESPAN_ERROR = exc
+        raise
+    except asyncio.CancelledError:
+        pass
     finally:
         await api_store.exit_store()
         await stop_ui_bundler()
@@ -97,3 +133,6 @@ async def queue_with_signal():
     except Exception as exc:
         logger.exception("Queue failed. Signaling shutdown", exc_info=exc)
         signal.raise_signal(signal.SIGINT)
+
+
+lifespan.get_last_error = get_last_error  # type: ignore[attr-defined]

langgraph_runtime_inmem/metrics.py

@@ -1,7 +1,7 @@
 from langgraph_runtime_inmem.queue import get_num_workers
 
 
-def get_metrics() -> dict[str, int]:
+def get_metrics() -> dict[str, dict[str, int]]:
     from langgraph_api import config
 
     workers_max = config.N_JOBS_PER_WORKER