langgraph-runtime-inmem 0.9.0__py3-none-any.whl → 0.20.1__py3-none-any.whl

langgraph_runtime_inmem/__init__.py

@@ -9,7 +9,7 @@ from langgraph_runtime_inmem import (
     store,
 )
 
-__version__ = "0.9.0"
+__version__ = "0.20.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()
 
 

langgraph_runtime_inmem/inmem_stream.py

@@ -58,6 +58,7 @@ class StreamManager:
         )  # 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])
@@ -95,7 +96,7 @@ class StreamManager:
 
     async def put(
         self,
-        run_id: UUID | str,
+        run_id: UUID | str | None,
         thread_id: UUID | str | None,
         message: Message,
         resumable: bool = False,
@@ -107,9 +108,10 @@ class StreamManager:
             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)
-        topic = message.topic.decode()
         if "control" in topic:
             self.control_keys[thread_id][run_id] = message
             queues = self.control_queues[thread_id][run_id]
@@ -121,6 +123,20 @@ class StreamManager:
             if isinstance(result, Exception):
                 logger.exception(f"Failed to put message in queue: {result}")
 
+    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:
@@ -145,6 +161,12 @@ class StreamManager:
         self.control_queues[thread_id][run_id].append(queue)
         return 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
     ):

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