langgraph-runtime-inmem 0.38.0.dev13__tar.gz → 0.38.0.dev14__tar.gz

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/PKG-INFO

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

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/langgraph_runtime_inmem/__init__.py

@@ -11,7 +11,7 @@ from langgraph_runtime_inmem import (
     store,
 )
 
-__version__ = "0.38.0.dev13"
+__version__ = "0.38.0.dev14"
 __all__ = [
     "ops",
     "database",

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/langgraph_runtime_inmem/database.py

@@ -185,12 +185,27 @@ def _load_ops_file() -> None:
     orphaned = [r for r in GLOBAL_STORE.get("runs") or [] if r["status"] == "running"]
     now = datetime.now(UTC)
     for run in orphaned:
-        run["status"] = "pending"
+        # An interrupt recorded before the crash stands, so that run is not retried. A
+        # rollback needs the checkpointer, which starts later: the worker that picks the
+        # run up applies it before the graph runs.
+        run["status"] = (
+            "interrupted" if run.get("cancel_action") == "interrupt" else "pending"
+        )
         run["updated_at"] = now
-    if orphaned:
+    # No worker will finish those runs, so nothing else moves their threads off busy.
+    stopped = {r["thread_id"] for r in orphaned if r["status"] == "interrupted"}
+    if stopped:
+        from langgraph_runtime_inmem.ops import _thread_status  # noqa: PLC0415
+
+        conn = InMemConnectionProto()
+        for thread in GLOBAL_STORE.get("threads") or []:
+            if thread["thread_id"] in stopped:
+                thread["status"] = _thread_status(conn, thread["thread_id"])
+    retried = [str(r["run_id"]) for r in orphaned if r["status"] == "pending"]
+    if retried:
         logger.warning(
             "Retrying runs left running by a previous process that did not shut down cleanly",
-            run_ids=[str(r["run_id"]) for r in orphaned],
+            run_ids=retried,
         )
 
 
@@ -206,6 +221,19 @@ def upgrade_thread_status_facts(store: dict[str, Any]) -> None:
     store.setdefault("thread_failed", {})
 
 
+def upgrade_thread_priority_settings(store: dict[str, Any]) -> None:
+    """Give threads saved by older versions, before priority settings, none yet."""
+    for thread in store.get("threads") or ():
+        for key in ("run_priority", "fairness_group"):
+            thread.setdefault(key, None)
+
+
+def upgrade_run_priorities(store: dict[str, Any]) -> None:
+    """Give runs saved by older versions, before runs had a priority, the default level."""
+    for run in store.get("runs") or ():
+        run.setdefault("priority", 0)
+
+
 def _set_aside_ops_file(reason: Exception) -> None:
     """Rename an unloadable ops file to ``<file>.bak-<UTC timestamp>`` and log why."""
     from langgraph_api.graph import graph_file_module_names  # noqa: PLC0415
@@ -293,6 +321,8 @@ async def start_pool() -> None:
         if not GLOBAL_STORE.get(k):
             GLOBAL_STORE[k] = []
     upgrade_thread_status_facts(GLOBAL_STORE)
+    upgrade_run_priorities(GLOBAL_STORE)
+    upgrade_thread_priority_settings(GLOBAL_STORE)
     # Remove agents that were created by the system to avoid removed graphs
     if GLOBAL_STORE.get("assistants"):
         for a in GLOBAL_STORE["assistants"]:

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/langgraph_runtime_inmem/inmem_stream.py

@@ -13,12 +13,19 @@ 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"
+def stream_id_key(stream_id: str | bytes) -> tuple[int, int]:
+    """Numeric order of a `<ms>-<seq>` stream id, for comparing resume cursors.
+
+    Clients send `-` or a negative id such as `-1` to replay everything, so
+    those sort first, and `<ms>` or `<ms>-*` to replay that whole millisecond.
+    Raises ValueError for an id without a numeric `<ms>`.
+    """
+    if isinstance(stream_id, bytes):
+        stream_id = stream_id.decode()
+    if stream_id.startswith("-"):
+        return (-1, -1)
+    ms, _, seq = stream_id.partition("-")
+    return int(ms), int(seq) if seq.isdigit() else -1
 
 
 @dataclass
@@ -65,6 +72,14 @@ class StreamManager:
         self.message_stores = defaultdict(
             lambda: defaultdict(list[Message])
         )  # Dict[str, List[Message]]
+        self._last_id = (0, -1)
+
+    def _next_id(self) -> bytes:
+        """Strictly increasing, as a resume replays only the ids above its cursor."""
+        ms = int(time.time() * 1000)
+        last_ms, last_seq = self._last_id
+        self._last_id = (ms, 0) if ms > last_ms else (last_ms, last_seq + 1)
+        return f"{self._last_id[0]}-{self._last_id[1]}".encode()
 
     def get_queues(
         self, run_id: UUID | str, thread_id: UUID | str | None
@@ -109,8 +124,7 @@ class StreamManager:
         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
+        message.id = self._next_id()
         topic = message.topic.decode()
         if resumable:
             self.message_stores[thread_id][run_id].append(message)
@@ -131,8 +145,8 @@ class StreamManager:
         message: Message,
     ) -> None:
         thread_id = _ensure_uuid(thread_id)
-        message.id = _generate_ms_seq_id().encode()
-        queues = self.thread_streams[thread_id]
+        message.id = self._next_id()
+        queues = self.thread_streams.get(thread_id, ())
         coros = [queue.put(message) for queue in queues]
         results = await asyncio.gather(*coros, return_exceptions=True)
         for result in results:
@@ -182,6 +196,13 @@ class StreamManager:
             if not self.queues[thread_id][run_id]:
                 del self.queues[thread_id][run_id]
 
+    async def remove_thread_stream(self, thread_id: UUID | str, queue: asyncio.Queue):
+        thread_id = _ensure_uuid(thread_id)
+        if thread_id in self.thread_streams:
+            self.thread_streams[thread_id].remove(queue)
+            if not self.thread_streams[thread_id]:
+                del self.thread_streams[thread_id]
+
     async def remove_control_queue(
         self, run_id: UUID | str, thread_id: UUID | str | None, queue: asyncio.Queue
     ):
@@ -207,19 +228,15 @@ class StreamManager:
             thread_id = THREADLESS_KEY
         else:
             thread_id = _ensure_uuid(thread_id)
-        if message_id is None:
+        if message_id is None or thread_id not in self.message_stores:
             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:]
+            after = stream_id_key(message_id)
+        except ValueError:
+            return
+        for message in self.message_stores[thread_id][run_id]:
+            if stream_id_key(message.id) > after:
+                yield message
 
     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."""

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/langgraph_runtime_inmem/lifespan.py

@@ -136,6 +136,20 @@ async def lifespan(
                 _LAST_LIFESPAN_ERROR = exc
                 await _log_graph_load_failure(exc)
                 raise
+            try:
+                from langgraph_api.priority_settings import (  # noqa: PLC0415
+                    get_priority_callback,
+                )
+            except ImportError:
+                # An older langgraph-api has no priority callback.
+                pass
+            else:
+                # Loaded now so a bad path fails startup, not every run create.
+                if get_priority_callback() is not None:
+                    await logger.awarning(
+                        "The in-memory runtime stores each thread's priority settings"
+                        " but starts runs in the order they were created"
+                    )
             if config.N_JOBS_PER_WORKER > 0:
                 tg.create_task(queue_with_signal())
 

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/langgraph_runtime_inmem/ops.py

@@ -32,6 +32,7 @@ from langgraph_runtime_inmem.inmem_stream import (
     ContextQueue,
     Message,
     get_stream_manager,
+    stream_id_key,
 )
 
 if typing.TYPE_CHECKING:
@@ -1404,6 +1405,8 @@ class Threads(Authenticated):
             "status": "idle",
             "config": {},
             "values": None,
+            "run_priority": None,
+            "fairness_group": None,
         }
 
         # Add to store
@@ -1560,6 +1563,9 @@ class Threads(Authenticated):
 
         now = datetime.now(UTC)
 
+        if run.get("cancel_action") == "rollback":
+            # A rollback recorded while the run ran wins over the signal its worker saw first, as in Go.
+            run_status, checkpoint, exception = "rollback", None, None
         if run_status == "rollback":
             await Runs.delete(conn, run_id, thread_id=run["thread_id"])
         else:
@@ -1771,6 +1777,9 @@ class Threads(Authenticated):
                 "values": copy.deepcopy(original_thread.get("values")),
                 "interrupts": copy.deepcopy(original_thread.get("interrupts", {})),
                 "error": copy.deepcopy(original_thread.get("error")),
+                # A copy gets its own priority settings at its first run.
+                "run_priority": None,
+                "fairness_group": None,
             }
             # Add new thread to store
             conn.store["threads"].append(
@@ -2412,8 +2421,7 @@ class Threads(Authenticated):
                             ):
                                 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())
+                        all_events.sort(key=lambda x: stream_id_key(x[0].id))
 
                         # Yield sorted events
                         for message, run_id in all_events:
@@ -2535,9 +2543,14 @@ class Threads(Authenticated):
                 raise
             finally:
                 # Clean up all created queues
+                thread_key = _ensure_uuid(thread_id)
                 for run_id, queue in created_queues:
                     try:
-                        await stream_manager.remove_queue(run_id, thread_id, queue)
+                        # subscribe() files the thread-level queue under the thread id.
+                        if run_id == thread_key:
+                            await stream_manager.remove_thread_stream(thread_id, queue)
+                        else:
+                            await stream_manager.remove_queue(run_id, thread_id, queue)
                     except Exception:
                         # Ignore cleanup errors
                         pass
@@ -2877,6 +2890,17 @@ class Runs(Authenticated):
         from langgraph_api.schema import Run, Thread  # noqa: PLC0415
         from langgraph_api.utils import uuid7  # noqa: PLC0415
 
+        try:
+            from langgraph_api.priority_settings import (  # noqa: PLC0415
+                assign_priority_settings,
+                get_priority_callback,
+            )
+        except ImportError:
+            # An older langgraph-api has no priority callback.
+            priority_callback = None
+        else:
+            priority_callback = get_priority_callback()
+
         assistant_id = _ensure_uuid(assistant_id)
         assistant = next(
             (a for a in conn.store["assistants"] if a["assistant_id"] == assistant_id),
@@ -2949,6 +2973,53 @@ class Runs(Authenticated):
 
         _reject_existing_run(conn.store["runs"], run_id)
 
+        # A thread gets its priority settings at its first run, from the deployment's
+        # priority callback, before the thread is created or updated. They're kept as
+        # the callback returned them; inmem never orders runs by them.
+        priority_settings = None
+        creates_thread = not existing_thread and (
+            thread_id is None or if_not_exists == "create"
+        )
+        if priority_callback is not None and (
+            creates_thread
+            or (
+                existing_thread is not None
+                and existing_thread.get("run_priority") is None
+            )
+        ):
+            ids = {"graph_id": assistant["graph_id"], "assistant_id": str(assistant_id)}
+            priority_settings = await assign_priority_settings(
+                priority_callback,
+                thread_id=str(thread_id) if thread_id is not None else None,
+                thread_metadata=(
+                    {**ids, **(config.get("metadata") or {}), **metadata}
+                    if creates_thread
+                    else Runs._merge_jsonb(existing_thread["metadata"], ids)
+                ),
+                assistant_id=str(assistant_id),
+                graph_id=assistant["graph_id"],
+                assistant_metadata=assistant.get("metadata") or {},
+                run_metadata=metadata,
+                run_kwargs=kwargs,
+            )
+            # Another first run on the thread may have gone ahead while the callback
+            # ran: read the thread again, and keep its settings if it has them now.
+            if thread_id is not None:
+                existing_thread = next(
+                    (t for t in conn.store["threads"] if t["thread_id"] == thread_id),
+                    None,
+                )
+                if encryption_enabled and existing_thread:
+                    existing_thread = await _decrypt_from_storage(
+                        existing_thread, "thread"
+                    )
+                if (
+                    existing_thread
+                    and filters
+                    and not _check_filter_match(existing_thread["metadata"], filters)
+                ):
+                    return _empty_generator()
+
         if not existing_thread and (thread_id is None or if_not_exists == "create"):
             # Create new thread
             if thread_id is None:
@@ -2975,10 +3046,18 @@ class Runs(Authenticated):
                 created_at=datetime.now(UTC),
                 updated_at=datetime.now(UTC),
                 values=None,
+                run_priority=priority_settings["priority"]
+                if priority_settings
+                else None,
+                fairness_group=(
+                    priority_settings["fairness_group"] if priority_settings else None
+                ),
             )
 
-            await logger.ainfo("Creating thread", thread_id=thread_id)
+            # Append before the first await, so a concurrent first run that read the
+            # thread after its priority callback sees this one instead of adding another.
             conn.store["threads"].append(thread)
+            await logger.ainfo("Creating thread", thread_id=thread_id)
         elif existing_thread:
             # Update existing thread
             if existing_thread["status"] != "busy":
@@ -3024,6 +3103,30 @@ class Runs(Authenticated):
 
                 return _return_inflight()
 
+        # As in Go core, an existing thread gets the callback's result only once a run
+        # is created on it, and only if it still has no settings; while the deployment
+        # has a callback, every run takes the thread's priority. The check reads the
+        # stored thread, not a decrypted copy, which a concurrent first run's settings
+        # may have reached first.
+        if existing_thread:
+            stored = next(
+                (t for t in conn.store["threads"] if t["thread_id"] == thread_id),
+                existing_thread,
+            )
+            if priority_settings and stored.get("run_priority") is None:
+                stored.update(
+                    run_priority=priority_settings["priority"],
+                    fairness_group=priority_settings["fairness_group"],
+                )
+            existing_thread.update(
+                {key: stored.get(key) for key in ("run_priority", "fairness_group")}
+            )
+        thread_priority = (
+            (existing_thread or thread).get("run_priority")
+            if priority_callback is not None
+            else None
+        )
+
         # Create new run
         configurable = Runs._merge_jsonb(
             Runs._get_configurable(assistant["config"]),
@@ -3085,6 +3188,7 @@ class Runs(Authenticated):
             updated_at=datetime.now(UTC),
             langsmith_session_name=langsmith_session_name,
             cancel_reason=None,
+            priority=thread_priority if thread_priority is not None else 0,
         )
         stored_run = new_run
         encrypted_thread = None
@@ -3341,8 +3445,8 @@ class Runs(Authenticated):
         """Interrupt every in-flight run accepted at or before `accepted_before`.
 
         Applies no auth filters, so only the operator route may call it. As with
-        `cancel`, running runs are marked interrupted when they're signalled, so
-        `finishing` only counts runs whose worker is still applying a rollback.
+        `cancel`, a signalled running run stays running until its worker stops it,
+        so `finishing` counts the running runs still being stopped.
         """
         boundary_ms = (accepted_before - _EPOCH) // timedelta(milliseconds=1)
 
@@ -3403,9 +3507,13 @@ class Runs(Authenticated):
                 cancelable_runs.append(run)
                 if run["status"] == "pending":
                     pending_thread_ids.add(run["thread_id"])
-                if action == "rollback" and run["status"] == "running":
-                    # Its worker finalizes the rollback; staying "running" holds the thread's next run until then.
-                    run["cancel_action"] = "rollback"
+                if run["status"] == "running":
+                    # Its worker finalizes the cancel; staying "running" holds the thread's next run until then.
+                    # A later interrupt keeps the first one's reason, which names who stopped the run.
+                    if action == "rollback" or not run.get("cancel_action"):
+                        run["cancel_reason"] = reason
+                    run["cancel_action"] = action
+                    run["updated_at"] = datetime.now(tz=UTC)
                 elif queues or action != "rollback":
                     run["status"] = "interrupted"
                     run["cancel_reason"] = reason
@@ -3717,15 +3825,22 @@ async def listen_for_cancellation(
 
     if control_key := stream_manager.get_control_key(run_id, thread_id):
         _handle_control_message(control_key)
+    else:
+        # Control messages do not survive a restart; a cancel recorded on the run does.
+        runs = InMemConnectionProto().store["runs"]
+        run = next((r for r in runs if r["run_id"] == run_id), None)
+        if run and run.get("cancel_action"):
+            _handle_control_message(
+                Message(
+                    topic=f"run:{run_id}:control".encode(),
+                    data=run["cancel_action"].encode(),
+                    reason=run.get("cancel_reason"),
+                )
+            )
 
+    # Runs.enter cancels this task when the run ends.
     while not done.is_set():
-        try:
-            # This task gets cancelled when Runs.enter exits anyway,
-            # so we can have a pretty lengthy timeout here
-            message = await asyncio.wait_for(queue.get(), timeout=240)
-            if _handle_control_message(message):
-                break
-        except TimeoutError:
+        if _handle_control_message(await queue.get()):
             break
 
 

{langgraph_runtime_inmem-0.38.0.dev13 → langgraph_runtime_inmem-0.38.0.dev14}/pyproject.toml

@@ -65,6 +65,10 @@ exclude = [
 
 
 [tool.ruff]
-lint.select = ["E", "F", "UP", "B", "I", "PLC0415"]
+lint.select = ["E", "F", "UP", "B", "I", "PLC0415", "TID251"]
 lint.ignore = ["E501", "B008"]
 target-version = "py311"
+
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"asyncio.set_event_loop_policy".msg = "Pass a loop_factory to asyncio.Runner instead."
+"uvloop.install".msg = "Use langgraph_api.asyncio.run_main_loop instead."