From 2855c2a01e4b1bb34855df93fc271010eb1db6dc Mon Sep 17 00:00:00 2001
From: Jamie Grier <jamie.grier@gmail.com>
Date: Thu, 23 Jul 2026 17:36:29 -0700
Subject: [PATCH 1/5] feat(langchain): deterministic durability for StateGraph
 parallelism

RestateMiddleware linearizes create_agent's parallel tool-call batch. This adds
`durable_scope`, which generalizes the same guarantee to an arbitrary StateGraph:
parallel branches, Send/map-reduce, and nested parallel subgraphs.

LangGraph runs a superstep's nodes concurrently on asyncio; `ctx.run_typed`
creates a journal command synchronously at call time, so parallel nodes create
commands in I/O-completion order -- non-deterministic, and different on replay
(journaled results resolve instantly). Restate matches the journal by position,
so replay mismatches and the invocation cannot recover.

`durable_scope(ctx)` yields a context wrapper; nodes call `dctx.run_typed(...)`
unchanged (no node wrappers, no config threading). A per-invocation coordinator
batches durable ops and creates each round's commands in stable sorted order
(deterministic journal) while awaiting the I/O concurrently via `restate.gather`.
Task identity uses LangGraph's replay-stable name/path metadata, composed
hierarchically for nesting.

RFC/draft: membership + per-task keying come from lazily monkeypatching a private
LangGraph function (langgraph.pregel._runner.arun_with_retry). This is brittle
across LangGraph versions; a merge-ready version wants a supported LangGraph
runner/executor hook. Details and the crash-recovery verification are in the PR.

Test is guarded by importorskip('langgraph') since langgraph is an optional
`langchain` extra, not a default test dependency.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 python/restate/ext/langchain/__init__.py    |   3 +
 python/restate/ext/langchain/_stategraph.py | 284 ++++++++++++++++++++
 tests/langchain_stategraph.py               | 138 ++++++++++
 3 files changed, 425 insertions(+)
 create mode 100644 python/restate/ext/langchain/_stategraph.py
 create mode 100644 tests/langchain_stategraph.py

diff --git a/python/restate/ext/langchain/__init__.py b/python/restate/ext/langchain/__init__.py
index 04ceb91..a9c3223 100644
--- a/python/restate/ext/langchain/__init__.py
+++ b/python/restate/ext/langchain/__init__.py
@@ -23,6 +23,7 @@
 from restate.server_context import current_context
 
 from ._middleware import RestateMiddleware
+from ._stategraph import DurableContext, durable_scope
 
 
 def restate_context() -> Context:
@@ -51,4 +52,6 @@ def restate_object_context() -> ObjectContext:
     "RestateMiddleware",
     "restate_context",
     "restate_object_context",
+    "durable_scope",
+    "DurableContext",
 ]
diff --git a/python/restate/ext/langchain/_stategraph.py b/python/restate/ext/langchain/_stategraph.py
new file mode 100644
index 0000000..d9d3c4a
--- /dev/null
+++ b/python/restate/ext/langchain/_stategraph.py
@@ -0,0 +1,284 @@
+#
+#  Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+"""Deterministic durability for LangGraph ``StateGraph`` parallelism on Restate.
+
+``RestateMiddleware`` (see ``_middleware.py``) makes a ``create_agent`` agent
+durable and linearizes its parallel *tool-call batch*. This module generalizes
+the same idea to an arbitrary ``StateGraph``: parallel branches, ``Send`` /
+map-reduce, and nested parallel subgraphs.
+
+The problem
+-----------
+LangGraph runs the nodes of a superstep concurrently on asyncio. ``ctx.run_typed``
+synchronously creates a journal command at call time, so parallel nodes create
+commands in I/O-completion order — non-deterministic, and different on replay
+(journaled results resolve instantly). Restate matches the journal by position, so
+replay mismatches raise ``TerminalError`` / journal-mismatch and the invocation
+cannot recover.
+
+The approach
+------------
+``durable_scope(ctx)`` yields a context wrapper; nodes call ``dctx.run_typed(...)``
+exactly as usual (no node wrappers, no ``config`` threading). Under the hood a
+per-invocation coordinator batches durable ops and creates each round's commands
+in a stable, sorted order (deterministic journal) while awaiting the actual I/O
+concurrently via :func:`restate.gather`. Task identity comes from LangGraph's
+replay-stable ``name``/``path`` metadata, composed hierarchically for nesting.
+
+Usage::
+
+    from restate.ext.langchain import durable_scope
+
+    @svc.handler()
+    async def run(ctx: restate.Context, req: Req) -> Resp:
+        with durable_scope(ctx) as dctx:
+            graph = build_graph(dctx)          # nodes call dctx.run_typed(...)
+            return await graph.ainvoke(state)  # no config, no node wrappers
+
+.. warning::
+   **RFC / draft.** Membership and per-task keying are obtained by monkeypatching
+   a *private* LangGraph function (``langgraph.pregel._runner.arun_with_retry``),
+   applied lazily the first time :class:`durable_scope` is entered. This is
+   brittle across LangGraph versions; a merge-ready version wants a *supported*
+   LangGraph runner/executor hook. See the PR discussion.
+
+.. note::
+   Current limitations: (1) round batching imposes a cross-branch barrier per
+   durable-op round, so heterogeneous per-branch latencies run slower than native;
+   (2) assumes one durable op in flight per task at a time — a node that fires
+   multiple durable ops *concurrently* (intra-node ``gather``) needs op-level
+   sub-keys, not handled here.
+"""
+
+import asyncio
+import contextvars
+import functools
+from typing import Any, Callable, Optional
+
+import restate
+
+# ---------------------------------------------------------------------------
+# Coordinator: deterministic command-creation order, concurrent I/O.
+# ---------------------------------------------------------------------------
+
+
+class _Pending:
+    __slots__ = ("name", "fn", "ctx", "event", "value", "error")
+
+    def __init__(self, name: str, fn: Callable[[], Any], ctx: Any) -> None:
+        self.name, self.fn, self.ctx = name, fn, ctx
+        self.event = asyncio.Event()
+        self.value: Any = None
+        self.error: Optional[BaseException] = None
+
+
+def _await_all(futs: list) -> Any:
+    """:func:`restate.gather` for real SDK futures; :func:`asyncio.gather` for
+    plain awaitables (keeps the coordinator unit-testable without a server)."""
+    try:
+        from restate.server_context import ServerDurableFuture
+
+        if futs and isinstance(futs[0], ServerDurableFuture):
+            return restate.gather(*futs)
+    except Exception:  # pylint: disable=broad-except
+        pass
+    return asyncio.gather(*futs, return_exceptions=True)
+
+
+class _Coordinator:
+    """Nesting-aware coordinator. Flushes a round when every active *leaf* task
+    (an active task with no active descendant) has a pending op, creating that
+    round's commands in sorted-key order, then awaiting them concurrently.
+
+    With eager registration via the runner hook the full sibling set is known up
+    front, so ``settle_turns == 0`` (flush as soon as every active leaf is
+    pending). A positive ``settle_turns`` tolerates lazy registration.
+    """
+
+    def __init__(self, settle_turns: int = 0) -> None:
+        self.settle_turns = settle_turns
+        self.active: set[str] = set()
+        self.pending: dict[str, _Pending] = {}
+        self._flushing = False
+        self._flush_scheduled = False
+
+    def register(self, full_key: str) -> None:
+        self.active.add(full_key)
+
+    def checkout(self, full_key: str) -> None:
+        self.active.discard(full_key)
+        self.pending.pop(full_key, None)
+        self._try_flush()
+
+    def _leaves(self) -> set[str]:
+        act = self.active
+        return {k for k in act if not any(o != k and o.startswith(k + "/") for o in act)}
+
+    def _ready_set(self) -> Optional[set[str]]:
+        leaves = self._leaves()
+        if leaves and all(k in self.pending for k in leaves):
+            return leaves
+        return None
+
+    def _try_flush(self) -> None:
+        if self._flushing or self._flush_scheduled:
+            return
+        if self._ready_set() is None:
+            return
+        self._flush_scheduled = True
+        asyncio.create_task(self._flush_when_stable())
+
+    async def _flush_when_stable(self) -> None:
+        try:
+            # settle_turns == 0 (hook installed): flush immediately. Otherwise
+            # wait until the ready leaf-set is stable across a few turns so
+            # lazily-registering siblings can appear first.
+            stable = 0
+            prev: Optional[frozenset] = None
+            while stable < self.settle_turns:
+                if self._flushing:
+                    return
+                ready = self._ready_set()
+                if ready is None:
+                    prev, stable = None, 0
+                else:
+                    snap = frozenset(ready)
+                    stable = stable + 1 if snap == prev else 1
+                    prev = snap
+                await asyncio.sleep(0)
+            ready = self._ready_set()
+            if self._flushing or ready is None:
+                return
+            self._flushing = True
+            entries = []
+            for k in sorted(ready):
+                p = self.pending.pop(k)
+                fut = p.ctx.run_typed(p.name, p.fn)  # sync create -> deterministic order
+                entries.append((p, fut))
+            asyncio.create_task(self._drive(entries))
+        finally:
+            self._flush_scheduled = False
+
+    async def _drive(self, entries: list) -> None:
+        try:
+            gathered = await _await_all([fut for _, fut in entries])
+            is_results = isinstance(gathered, list) and gathered and not hasattr(gathered[0], "__await__")
+            for i, (p, fut) in enumerate(entries):
+                try:
+                    p.value = gathered[i] if is_results else await fut
+                    if isinstance(p.value, BaseException):
+                        p.error, p.value = p.value, None
+                except BaseException as e:  # noqa: BLE001  pylint: disable=broad-except
+                    p.error = e
+                p.event.set()
+        finally:
+            self._flushing = False
+            self._try_flush()
+
+    async def submit(self, full_key: str, ctx: Any, name: str, fn: Callable[[], Any]) -> Any:
+        p = _Pending(name, fn, ctx)
+        self.pending[full_key] = p
+        self._try_flush()
+        await p.event.wait()
+        if p.error is not None:
+            raise p.error
+        return p.value
+
+
+# ---------------------------------------------------------------------------
+# Runner hook + context wrapper (the integration surface).
+# ---------------------------------------------------------------------------
+
+# (coordinator, full_key) for the task currently executing on this async task.
+_current: contextvars.ContextVar[Optional[tuple]] = contextvars.ContextVar("_restate_lg_current", default=None)
+# The active invocation's coordinator (set by durable_scope in the handler).
+_active_coord: contextvars.ContextVar[Optional[_Coordinator]] = contextvars.ContextVar(
+    "_restate_lg_active_coord", default=None
+)
+
+_installed = False
+
+
+def install() -> None:
+    """Idempotently patch LangGraph's task runner. Called lazily by
+    :class:`durable_scope`, so the monkeypatch is only applied when the feature
+    is actually used (not merely on import)."""
+    global _installed  # pylint: disable=global-statement
+    if _installed:
+        return
+    _installed = True
+
+    import langgraph.pregel._runner as _runner_mod  # pylint: disable=import-outside-toplevel
+
+    orig_arun = _runner_mod.arun_with_retry
+
+    @functools.wraps(orig_arun)
+    async def patched_arun(task, *args, **kwargs):  # type: ignore[no-untyped-def]
+        coord = _active_coord.get()
+        if coord is None or task is None:
+            return await orig_arun(task, *args, **kwargs)
+        parent = _current.get()
+        parent_key = parent[1] if parent else ""
+        # Hierarchical, replay-stable key: parent + this task's name + path
+        # (path carries the Send index). No task-id UUIDs (those aren't stable).
+        full = f"{parent_key}/{task.name}|{task.path}"
+        coord.register(full)  # eager: at task start, before any deep await
+        token = _current.set((coord, full))
+        try:
+            return await orig_arun(task, *args, **kwargs)
+        finally:
+            _current.reset(token)
+            coord.checkout(full)
+
+    _runner_mod.arun_with_retry = patched_arun
+
+
+class DurableContext:
+    """Wraps a Restate ``Context`` so ``run_typed`` routes through the active
+    coordinator. Everything else delegates to the real context — node code calls
+    ``ctx.run_typed`` exactly as normal and the determinism is invisible."""
+
+    def __init__(self, ctx: Any) -> None:
+        self._ctx = ctx
+
+    def run_typed(self, name: str, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        cur = _current.get()
+        if cur is None:
+            return self._ctx.run_typed(name, fn, *args, **kwargs)
+        coord, full = cur
+        bound = fn if not (args or kwargs) else functools.partial(fn, *args, **kwargs)
+        return coord.submit(full, self._ctx, name, bound)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._ctx, name)
+
+
+class durable_scope:  # pylint: disable=invalid-name
+    """Handler helper. ``with durable_scope(ctx) as dctx:`` installs the runner
+    hook (lazily, idempotent) and a fresh per-invocation coordinator, and yields
+    a durable-wrapped context. Build the graph with ``dctx`` and invoke normally
+    — no node wrapping, no ``config``."""
+
+    def __init__(self, ctx: Any) -> None:
+        install()
+        self.ctx = ctx
+        self.coord = _Coordinator(settle_turns=0)
+        self._token: Any = None
+
+    def __enter__(self) -> DurableContext:
+        self._token = _active_coord.set(self.coord)
+        return DurableContext(self.ctx)
+
+    def __exit__(self, *exc: Any) -> None:
+        _active_coord.reset(self._token)
+
+
+__all__ = ["durable_scope", "DurableContext", "install"]
diff --git a/tests/langchain_stategraph.py b/tests/langchain_stategraph.py
new file mode 100644
index 0000000..c96ab69
--- /dev/null
+++ b/tests/langchain_stategraph.py
@@ -0,0 +1,138 @@
+#
+#  Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+"""Determinism test for `restate.ext.langchain.durable_scope` over a StateGraph.
+
+Proves the core property without a server: parallel nodes (here, a fan-out to
+subgraphs that themselves fan out) create their durable commands in a stable,
+sorted order every run — i.e. the order Restate would journal is identical across
+executions, so replay lines up positionally.
+
+A `FakeCtx` stands in for the Restate context; its `run_typed` records the command
+name at *creation* time (exactly what Restate journals) and returns the coroutine,
+so the async I/O still races while the recorded order stays fixed.
+
+Skipped unless `langgraph` is installed (it is an optional `langchain` extra, not a
+default test dependency).
+"""
+
+import asyncio
+import contextvars
+import operator
+import random
+from typing import Annotated, TypedDict
+
+import pytest
+
+pytest.importorskip("langgraph")
+
+from langgraph.graph import StateGraph, START, END  # noqa: E402
+from langgraph.types import Send  # noqa: E402
+
+from restate.ext.langchain import durable_scope  # noqa: E402
+
+pytestmark = [pytest.mark.anyio]
+
+
+@pytest.fixture(scope="session")
+def anyio_backend():
+    return "asyncio"
+
+
+_journal: contextvars.ContextVar[list] = contextvars.ContextVar("journal")
+
+
+class FakeCtx:
+    """Records command-creation order (what Restate journals); returns the coro."""
+
+    def run_typed(self, name, fn, *args, **kwargs):
+        _journal.get().append(name)
+        return fn()
+
+
+class _Inner(TypedDict):
+    branch: str
+    results: Annotated[list, operator.add]
+
+
+class _InnerWS(TypedDict):
+    branch: str
+    leg: str
+
+
+def _make_sub(dctx):
+    async def inner_worker(state: _InnerWS) -> dict:
+        b, leg = state["branch"], state["leg"]
+
+        async def fetch():
+            await asyncio.sleep(random.uniform(0.005, 0.03))
+            return f"f:{b}.{leg}"
+
+        async def store():
+            await asyncio.sleep(random.uniform(0.005, 0.03))
+            return f"s:{b}.{leg}"
+
+        await dctx.run_typed(f"fetch:{b}.{leg}", fetch)
+        await dctx.run_typed(f"store:{b}.{leg}", store)
+        return {"results": [f"{b}.{leg}"]}
+
+    def inner_fan(state: _Inner):
+        return [Send("inner_worker", {"branch": state["branch"], "leg": leg}) for leg in ("L", "R")]
+
+    b = StateGraph(_Inner)
+    b.add_node("inner_worker", inner_worker)
+    b.add_conditional_edges(START, inner_fan, ["inner_worker"])
+    b.add_edge("inner_worker", END)
+    return b.compile()
+
+
+class _Outer(TypedDict):
+    branches: list
+    results: Annotated[list, operator.add]
+
+
+class _OuterWS(TypedDict):
+    branch: str
+
+
+def _build(dctx):
+    sub = _make_sub(dctx)
+
+    async def branch_node(state: _OuterWS) -> dict:
+        out = await sub.ainvoke({"branch": state["branch"], "results": []})
+        return {"results": out["results"]}
+
+    def outer_fan(state: _Outer):
+        return [Send("branch_node", {"branch": b}) for b in state["branches"]]
+
+    b = StateGraph(_Outer)
+    b.add_node("branch_node", branch_node)
+    b.add_conditional_edges(START, outer_fan, ["branch_node"])
+    b.add_edge("branch_node", END)
+    return b.compile()
+
+
+async def _one_run() -> tuple:
+    journal: list = []
+    token = _journal.set(journal)
+    try:
+        with durable_scope(FakeCtx()) as dctx:
+            await _build(dctx).ainvoke({"branches": ["A", "B", "C"], "results": []})
+    finally:
+        _journal.reset(token)
+    return tuple(journal)
+
+
+async def test_stategraph_durable_order_is_deterministic():
+    runs = [await _one_run() for _ in range(20)]
+    distinct = set(runs)
+    assert len(distinct) == 1, f"non-deterministic journal order: {distinct}"
+    # 3 branches x 2 legs x 2 ops = 12 durable commands per run
+    assert len(runs[0]) == 12

From f7d71751d48bcb0fd731bdc61c43e793c1920c7c Mon Sep 17 00:00:00 2001
From: Jamie Grier <jamie.grier@gmail.com>
Date: Fri, 24 Jul 2026 08:46:18 -0700
Subject: [PATCH 2/5] feat(langchain): auto-enabled StateGraph determinism +
 parallel-safe turnstile

Two changes that together make Rushabh's parallel-node/subgraph repro
deterministic with zero user code (validated: baseline 570-poison -> fixed).

1) StateGraph determinism coordinator (_stategraph.py), rewritten from the
   earlier node-wrapper prototype:
   - Patches ServerInvocationContext.run_typed at the class level so
     restate_context()/RestateMiddleware route parallel durable ops through a
     per-invocation coordinator, which creates their journal commands in a
     stable, hierarchical-key sorted order (langgraph task name|path, composed
     for nested subgraphs) and awaits the I/O concurrently via restate.gather.
     The ctx stays a real ServerInvocationContext (extension_data still works).
   - Patches langgraph.pregel._runner.arun_with_retry to tag each task with its
     replay-stable key.
   - enable() auto-activates a coordinator per invocation by wrapping the SDK's
     invoke_handler; __init__ calls enable() on import so merely importing the
     integration turns it on. Opt out with RESTATE_STATEGRAPH_DETERMINISM=0.

2) Parallel-safe middleware turnstile (_state.py, _middleware.py):
   - _State now holds turnstiles keyed by tool_call_id instead of one shared
     slot. Parallel tool-using agents in concurrent LangGraph nodes were
     clobbering the single slot (KeyError) once the coordinator synchronized
     their model calls; keying by the unique tool_call_id fixes it and is
     backward-compatible for the single-agent tool batch.

RFC/draft: still monkeypatches private langgraph internals (arun_with_retry);
a merge-ready version wants a supported runner hook. Determinism validated
end-to-end against Rushabh's exact repro with a stubbed model; see notes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 python/restate/ext/langchain/__init__.py    |  15 +-
 python/restate/ext/langchain/_middleware.py |  18 +-
 python/restate/ext/langchain/_state.py      |  22 ++-
 python/restate/ext/langchain/_stategraph.py | 177 +++++++++++---------
 tests/langchain_stategraph.py               | 149 ++++++----------
 5 files changed, 199 insertions(+), 182 deletions(-)

diff --git a/python/restate/ext/langchain/__init__.py b/python/restate/ext/langchain/__init__.py
index a9c3223..67d7d29 100644
--- a/python/restate/ext/langchain/__init__.py
+++ b/python/restate/ext/langchain/__init__.py
@@ -17,13 +17,14 @@
 the tool body.
 """
 
+import os
 import typing
 
 from restate import Context, ObjectContext
 from restate.server_context import current_context
 
 from ._middleware import RestateMiddleware
-from ._stategraph import DurableContext, durable_scope
+from ._stategraph import durable_scope, enable
 
 
 def restate_context() -> Context:
@@ -53,5 +54,15 @@ def restate_object_context() -> ObjectContext:
     "restate_context",
     "restate_object_context",
     "durable_scope",
-    "DurableContext",
+    "enable",
 ]
+
+
+# Zero-config: importing the LangGraph integration is enough. Parallel StateGraph
+# nodes journal deterministically with NO user code — no wrappers, no startup call.
+# Set RESTATE_STATEGRAPH_DETERMINISM=0 to opt out. Never let this break the import.
+if os.environ.get("RESTATE_STATEGRAPH_DETERMINISM", "1") != "0":
+    try:
+        enable()
+    except Exception:  # pragma: no cover  pylint: disable=broad-except
+        pass
diff --git a/python/restate/ext/langchain/_middleware.py b/python/restate/ext/langchain/_middleware.py
index 1f0aa87..95f862a 100644
--- a/python/restate/ext/langchain/_middleware.py
+++ b/python/restate/ext/langchain/_middleware.py
@@ -101,7 +101,12 @@ async def call_model() -> SerializableModelResponse:
         state = get_or_create_state(ctx)
         if ai_message is not None:
             tool_call_ids = [tid for tc in (ai_message.tool_calls or []) if (tid := tc.get("id")) is not None]
-            state.turnstile = Turnstile(tool_call_ids)
+            turnstile = Turnstile(tool_call_ids)
+            # Register this turn's turnstile under EACH of its tool-call ids.
+            # Keyed by (unique) id rather than a single shared slot, so concurrent
+            # agents in parallel LangGraph nodes don't clobber each other.
+            for tid in tool_call_ids:
+                state.turnstiles[tid] = turnstile
 
         # Turn into ModelResponse as expected by the agent
         return ModelResponse(
@@ -122,7 +127,16 @@ async def awrap_tool_call(
         assert ctx is not None, "RestateMiddleware must run inside a Restate handler"
         state = state_from_ctx(ctx)
         assert state is not None, "RestateMiddleware must run inside a Restate handler"
-        turnstile = state.turnstile
+        turnstile = state.turnstiles.get(tool_call_id)
+
+        if turnstile is None:
+            # No turnstile registered for this id (e.g. a tool call not produced
+            # by a journaled model turn). Run it without ordering rather than
+            # deadlock/KeyError.
+            result = await handler(request)
+            if isinstance(result, ToolMessage):
+                result.id = str(ctx.uuid())
+            return result
 
         try:
             await turnstile.wait_for(tool_call_id)
diff --git a/python/restate/ext/langchain/_state.py b/python/restate/ext/langchain/_state.py
index 9b680d9..59b91f3 100644
--- a/python/restate/ext/langchain/_state.py
+++ b/python/restate/ext/langchain/_state.py
@@ -31,19 +31,27 @@ class _State:
     the sibling ``awrap_tool_call`` tasks spawned by langgraph's
     ``tool_node``; ``ctx.extension_data`` does.
 
-    Holds the current turnstile — replaced on every model call to
-    describe that turn's batch of tool calls. ``aafter_agent`` resets
-    it so subsequent agent runs in the same handler start clean.
+    Holds a map of ``tool_call_id -> Turnstile``. Each model turn creates one
+    turnstile and registers it under every tool-call id in that batch, so
+    ``awrap_tool_call`` finds it by its own id. Keying by the (globally unique)
+    tool-call id — rather than a single shared slot — is what makes this safe
+    when multiple agents run concurrently in parallel LangGraph nodes: sibling
+    branches write to distinct keys and never clobber each other.
     """
 
-    __slots__ = ("turnstile",)
+    __slots__ = ("turnstiles",)
 
     def __init__(self) -> None:
-        self.turnstile: Turnstile = Turnstile([])
+        self.turnstiles: dict[str, Turnstile] = {}
 
     def __close__(self) -> None:
-        # Called at handler end via auto_close_extension_data.
-        self.turnstile.cancel_all()
+        # Called at handler end via auto_close_extension_data. Cancel each
+        # distinct turnstile once (ids in a batch share one instance).
+        seen: set[int] = set()
+        for t in self.turnstiles.values():
+            if id(t) not in seen:
+                seen.add(id(t))
+                t.cancel_all()
 
 
 def get_or_create_state(ctx: Context) -> _State:
diff --git a/python/restate/ext/langchain/_stategraph.py b/python/restate/ext/langchain/_stategraph.py
index d9d3c4a..0e70053 100644
--- a/python/restate/ext/langchain/_stategraph.py
+++ b/python/restate/ext/langchain/_stategraph.py
@@ -12,26 +12,34 @@
 
 ``RestateMiddleware`` (see ``_middleware.py``) makes a ``create_agent`` agent
 durable and linearizes its parallel *tool-call batch*. This module generalizes
-the same idea to an arbitrary ``StateGraph``: parallel branches, ``Send`` /
+the guarantee to an arbitrary ``StateGraph``: parallel branches, ``Send`` /
 map-reduce, and nested parallel subgraphs.
 
 The problem
 -----------
-LangGraph runs the nodes of a superstep concurrently on asyncio. ``ctx.run_typed``
-synchronously creates a journal command at call time, so parallel nodes create
+LangGraph runs a superstep's nodes concurrently on asyncio. ``ctx.run_typed``
+creates a journal command synchronously at call time, so parallel nodes create
 commands in I/O-completion order — non-deterministic, and different on replay
-(journaled results resolve instantly). Restate matches the journal by position, so
-replay mismatches raise ``TerminalError`` / journal-mismatch and the invocation
-cannot recover.
-
-The approach
-------------
-``durable_scope(ctx)`` yields a context wrapper; nodes call ``dctx.run_typed(...)``
-exactly as usual (no node wrappers, no ``config`` threading). Under the hood a
-per-invocation coordinator batches durable ops and creates each round's commands
-in a stable, sorted order (deterministic journal) while awaiting the actual I/O
-concurrently via :func:`restate.gather`. Task identity comes from LangGraph's
-replay-stable ``name``/``path`` metadata, composed hierarchically for nesting.
+(journaled results resolve instantly). Restate matches the journal by position,
+so replay mismatches and the invocation cannot recover.
+
+The approach — fully hidden, no user code change
+------------------------------------------------
+Wrap the handler body in ``with durable_scope(ctx):``. That is the *only* change:
+nodes, agents, ``RestateMiddleware()`` and ``restate_context().run_typed(...)``
+are all written exactly as usual. Under the scope:
+
+  * ``ServerInvocationContext.run_typed`` is patched (once, lazily) so that when a
+    coordinator is active it routes durable ops through a per-invocation
+    coordinator instead of journaling them immediately. The context object stays a
+    real ``ServerInvocationContext`` (so ``extension_data`` etc. keep working).
+  * ``langgraph.pregel._runner.arun_with_retry`` is patched to tag each task with a
+    hierarchical, replay-stable key (parent + ``name|path``; the path carries the
+    Send index — no task-id UUIDs, which are not replay-stable).
+
+The coordinator batches durable ops and creates each round's commands in stable,
+sorted-key order (deterministic journal) while awaiting the actual I/O
+concurrently via :func:`restate.gather`.
 
 Usage::
 
@@ -39,23 +47,22 @@
 
     @svc.handler()
     async def run(ctx: restate.Context, req: Req) -> Resp:
-        with durable_scope(ctx) as dctx:
-            graph = build_graph(dctx)          # nodes call dctx.run_typed(...)
-            return await graph.ainvoke(state)  # no config, no node wrappers
+        with durable_scope(ctx):
+            return await graph.ainvoke(state)   # plain nodes, middleware, tools
 
 .. warning::
-   **RFC / draft.** Membership and per-task keying are obtained by monkeypatching
-   a *private* LangGraph function (``langgraph.pregel._runner.arun_with_retry``),
-   applied lazily the first time :class:`durable_scope` is entered. This is
-   brittle across LangGraph versions; a merge-ready version wants a *supported*
-   LangGraph runner/executor hook. See the PR discussion.
+   **RFC / draft.** Task keying is obtained by monkeypatching a *private* LangGraph
+   function (``langgraph.pregel._runner.arun_with_retry``). Brittle across LangGraph
+   versions; a merge-ready version wants a supported runner/executor hook.
 
 .. note::
-   Current limitations: (1) round batching imposes a cross-branch barrier per
-   durable-op round, so heterogeneous per-branch latencies run slower than native;
-   (2) assumes one durable op in flight per task at a time — a node that fires
-   multiple durable ops *concurrently* (intra-node ``gather``) needs op-level
-   sub-keys, not handled here.
+   Limitations: (1) round batching imposes a cross-branch barrier per durable-op
+   round, so heterogeneous per-branch latencies run slower than native; (2) assumes
+   one durable op in flight per task at a time — a node that fires multiple durable
+   ops *concurrently* (intra-node ``gather``) needs op-level sub-keys, not handled;
+   (3) a routed ``run_typed`` returns a coroutine rather than a ``RestateDurableFuture``,
+   so code that starts an op without awaiting it (to combine via ``restate.gather``)
+   must run outside ``durable_scope``.
 """
 
 import asyncio
@@ -71,10 +78,10 @@ async def run(ctx: restate.Context, req: Req) -> Resp:
 
 
 class _Pending:
-    __slots__ = ("name", "fn", "ctx", "event", "value", "error")
+    __slots__ = ("make_future", "event", "value", "error")
 
-    def __init__(self, name: str, fn: Callable[[], Any], ctx: Any) -> None:
-        self.name, self.fn, self.ctx = name, fn, ctx
+    def __init__(self, make_future: Callable[[], Any]) -> None:
+        self.make_future = make_future  # () -> RestateDurableFuture, called at flush
         self.event = asyncio.Event()
         self.value: Any = None
         self.error: Optional[BaseException] = None
@@ -96,12 +103,7 @@ def _await_all(futs: list) -> Any:
 class _Coordinator:
     """Nesting-aware coordinator. Flushes a round when every active *leaf* task
     (an active task with no active descendant) has a pending op, creating that
-    round's commands in sorted-key order, then awaiting them concurrently.
-
-    With eager registration via the runner hook the full sibling set is known up
-    front, so ``settle_turns == 0`` (flush as soon as every active leaf is
-    pending). A positive ``settle_turns`` tolerates lazy registration.
-    """
+    round's commands in sorted-key order, then awaiting them concurrently."""
 
     def __init__(self, settle_turns: int = 0) -> None:
         self.settle_turns = settle_turns
@@ -138,9 +140,6 @@ def _try_flush(self) -> None:
 
     async def _flush_when_stable(self) -> None:
         try:
-            # settle_turns == 0 (hook installed): flush immediately. Otherwise
-            # wait until the ready leaf-set is stable across a few turns so
-            # lazily-registering siblings can appear first.
             stable = 0
             prev: Optional[frozenset] = None
             while stable < self.settle_turns:
@@ -161,7 +160,7 @@ async def _flush_when_stable(self) -> None:
             entries = []
             for k in sorted(ready):
                 p = self.pending.pop(k)
-                fut = p.ctx.run_typed(p.name, p.fn)  # sync create -> deterministic order
+                fut = p.make_future()  # sync create -> deterministic order
                 entries.append((p, fut))
             asyncio.create_task(self._drive(entries))
         finally:
@@ -183,8 +182,8 @@ async def _drive(self, entries: list) -> None:
             self._flushing = False
             self._try_flush()
 
-    async def submit(self, full_key: str, ctx: Any, name: str, fn: Callable[[], Any]) -> Any:
-        p = _Pending(name, fn, ctx)
+    async def submit(self, full_key: str, make_future: Callable[[], Any]) -> Any:
+        p = _Pending(make_future)
         self.pending[full_key] = p
         self._try_flush()
         await p.event.wait()
@@ -194,7 +193,8 @@ async def submit(self, full_key: str, ctx: Any, name: str, fn: Callable[[], Any]
 
 
 # ---------------------------------------------------------------------------
-# Runner hook + context wrapper (the integration surface).
+# Patches + activation. No context wrapper: the real ServerInvocationContext's
+# run_typed is patched, so restate_context()/middleware code routes transparently.
 # ---------------------------------------------------------------------------
 
 # (coordinator, full_key) for the task currently executing on this async task.
@@ -205,19 +205,24 @@ async def submit(self, full_key: str, ctx: Any, name: str, fn: Callable[[], Any]
 )
 
 _installed = False
+_auto_enabled = False
+_orig_run_typed: Any = None
 
 
 def install() -> None:
-    """Idempotently patch LangGraph's task runner. Called lazily by
-    :class:`durable_scope`, so the monkeypatch is only applied when the feature
-    is actually used (not merely on import)."""
-    global _installed  # pylint: disable=global-statement
+    """Idempotently install the routing patches (LangGraph task runner + Restate
+    ``run_typed``). Both are no-ops unless a coordinator is active on the current
+    context — so this alone changes nothing until either :class:`durable_scope`
+    (scoped) or :func:`enable` (global) activates one."""
+    global _installed, _orig_run_typed  # pylint: disable=global-statement
     if _installed:
         return
     _installed = True
 
     import langgraph.pregel._runner as _runner_mod  # pylint: disable=import-outside-toplevel
+    from restate.server_context import ServerInvocationContext  # pylint: disable=import-outside-toplevel
 
+    # --- 1. tag each LangGraph task with a hierarchical, replay-stable key ---
     orig_arun = _runner_mod.arun_with_retry
 
     @functools.wraps(orig_arun)
@@ -227,8 +232,6 @@ async def patched_arun(task, *args, **kwargs):  # type: ignore[no-untyped-def]
             return await orig_arun(task, *args, **kwargs)
         parent = _current.get()
         parent_key = parent[1] if parent else ""
-        # Hierarchical, replay-stable key: parent + this task's name + path
-        # (path carries the Send index). No task-id UUIDs (those aren't stable).
         full = f"{parent_key}/{task.name}|{task.path}"
         coord.register(full)  # eager: at task start, before any deep await
         token = _current.set((coord, full))
@@ -240,45 +243,67 @@ async def patched_arun(task, *args, **kwargs):  # type: ignore[no-untyped-def]
 
     _runner_mod.arun_with_retry = patched_arun
 
+    # --- 2. route ctx.run_typed through the coordinator when one is active ---
+    _orig_run_typed = ServerInvocationContext.run_typed
+
+    def patched_run_typed(self, *args, **kwargs):  # type: ignore[no-untyped-def]
+        coord = _active_coord.get()
+        cur = _current.get()
+        if coord is None or cur is None:
+            return _orig_run_typed(self, *args, **kwargs)
+        _, full = cur
+        # Defer command creation to the coordinator's sorted flush; the flush
+        # calls the ORIGINAL run_typed (no recursion) to create the command.
+        return coord.submit(full, lambda: _orig_run_typed(self, *args, **kwargs))
+
+    ServerInvocationContext.run_typed = patched_run_typed  # type: ignore[method-assign]
+
+
+def enable() -> None:
+    """Turn on StateGraph determinism **globally** — call once at app startup and
+    every handler activates a fresh per-invocation coordinator automatically. NO
+    per-handler code, no ``durable_scope``: existing handlers, ``RestateMiddleware()``
+    and ``restate_context().run_typed(...)`` become deterministic as-is.
+
+    Implemented by wrapping the SDK's per-invocation ``invoke_handler`` entry point.
+    The coordinator is a no-op for any handler that doesn't run parallel LangGraph
+    nodes (the run_typed patch only routes when inside a LangGraph task)."""
+    global _auto_enabled  # pylint: disable=global-statement
+    install()
+    if _auto_enabled:
+        return
+    _auto_enabled = True
 
-class DurableContext:
-    """Wraps a Restate ``Context`` so ``run_typed`` routes through the active
-    coordinator. Everything else delegates to the real context — node code calls
-    ``ctx.run_typed`` exactly as normal and the determinism is invisible."""
+    import restate.server_context as _sc  # pylint: disable=import-outside-toplevel
 
-    def __init__(self, ctx: Any) -> None:
-        self._ctx = ctx
+    orig_invoke_handler = _sc.invoke_handler
 
-    def run_typed(self, name: str, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
-        cur = _current.get()
-        if cur is None:
-            return self._ctx.run_typed(name, fn, *args, **kwargs)
-        coord, full = cur
-        bound = fn if not (args or kwargs) else functools.partial(fn, *args, **kwargs)
-        return coord.submit(full, self._ctx, name, bound)
+    async def patched_invoke_handler(handler, ctx, in_buffer):  # type: ignore[no-untyped-def]
+        token = _active_coord.set(_Coordinator(settle_turns=0))
+        try:
+            return await orig_invoke_handler(handler=handler, ctx=ctx, in_buffer=in_buffer)
+        finally:
+            _active_coord.reset(token)
 
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self._ctx, name)
+    _sc.invoke_handler = patched_invoke_handler  # server_context.enter() resolves this name
 
 
 class durable_scope:  # pylint: disable=invalid-name
-    """Handler helper. ``with durable_scope(ctx) as dctx:`` installs the runner
-    hook (lazily, idempotent) and a fresh per-invocation coordinator, and yields
-    a durable-wrapped context. Build the graph with ``dctx`` and invoke normally
-    — no node wrapping, no ``config``."""
+    """Scoped opt-in alternative to :func:`enable`: ``with durable_scope(ctx):``
+    activates a coordinator for just this handler body. Useful if you don't want
+    the global :func:`enable`. Redundant (harmless) once :func:`enable` is set."""
 
-    def __init__(self, ctx: Any) -> None:
+    def __init__(self, ctx: Any = None) -> None:
         install()
-        self.ctx = ctx
-        self.coord = _Coordinator(settle_turns=0)
+        self._ctx = ctx  # accepted for a natural call site; not otherwise needed
         self._token: Any = None
 
-    def __enter__(self) -> DurableContext:
-        self._token = _active_coord.set(self.coord)
-        return DurableContext(self.ctx)
+    def __enter__(self) -> "durable_scope":
+        self._token = _active_coord.set(_Coordinator(settle_turns=0))
+        return self
 
     def __exit__(self, *exc: Any) -> None:
         _active_coord.reset(self._token)
 
 
-__all__ = ["durable_scope", "DurableContext", "install"]
+__all__ = ["enable", "durable_scope", "install"]
diff --git a/tests/langchain_stategraph.py b/tests/langchain_stategraph.py
index c96ab69..bc0931b 100644
--- a/tests/langchain_stategraph.py
+++ b/tests/langchain_stategraph.py
@@ -8,35 +8,20 @@
 #  directory of this repository or package, or at
 #  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
 #
-"""Determinism test for `restate.ext.langchain.durable_scope` over a StateGraph.
+"""Unit test for the StateGraph determinism coordinator.
 
-Proves the core property without a server: parallel nodes (here, a fan-out to
-subgraphs that themselves fan out) create their durable commands in a stable,
-sorted order every run — i.e. the order Restate would journal is identical across
-executions, so replay lines up positionally.
-
-A `FakeCtx` stands in for the Restate context; its `run_typed` records the command
-name at *creation* time (exactly what Restate journals) and returns the coroutine,
-so the async I/O still races while the recorded order stays fixed.
-
-Skipped unless `langgraph` is installed (it is an optional `langchain` extra, not a
-default test dependency).
+The coordinator is what makes parallel LangGraph nodes journal deterministically:
+given a set of concurrently-arriving durable ops (one per parallel task), it must
+create their journal commands in a STABLE, sorted-by-key order regardless of the
+order they arrived or completed. This test drives the coordinator directly (no
+server, no langgraph needed) and asserts that property.
 """
 
 import asyncio
-import contextvars
-import operator
-import random
-from typing import Annotated, TypedDict
 
 import pytest
 
-pytest.importorskip("langgraph")
-
-from langgraph.graph import StateGraph, START, END  # noqa: E402
-from langgraph.types import Send  # noqa: E402
-
-from restate.ext.langchain import durable_scope  # noqa: E402
+from restate.ext.langchain._stategraph import _Coordinator
 
 pytestmark = [pytest.mark.anyio]
 
@@ -46,93 +31,67 @@ def anyio_backend():
     return "asyncio"
 
 
-_journal: contextvars.ContextVar[list] = contextvars.ContextVar("journal")
-
-
-class FakeCtx:
-    """Records command-creation order (what Restate journals); returns the coro."""
-
-    def run_typed(self, name, fn, *args, **kwargs):
-        _journal.get().append(name)
-        return fn()
-
-
-class _Inner(TypedDict):
-    branch: str
-    results: Annotated[list, operator.add]
-
-
-class _InnerWS(TypedDict):
-    branch: str
-    leg: str
-
-
-def _make_sub(dctx):
-    async def inner_worker(state: _InnerWS) -> dict:
-        b, leg = state["branch"], state["leg"]
-
-        async def fetch():
-            await asyncio.sleep(random.uniform(0.005, 0.03))
-            return f"f:{b}.{leg}"
+async def test_coordinator_creates_commands_in_sorted_key_order():
+    coord = _Coordinator(settle_turns=0)
+    creation_order: list[str] = []
 
-        async def store():
-            await asyncio.sleep(random.uniform(0.005, 0.03))
-            return f"s:{b}.{leg}"
+    # Hierarchical, replay-stable keys, deliberately submitted out of order.
+    keys = ["/w|0003", "/w|0001", "/w|0004", "/w|0002", "/w|0000"]
 
-        await dctx.run_typed(f"fetch:{b}.{leg}", fetch)
-        await dctx.run_typed(f"store:{b}.{leg}", store)
-        return {"results": [f"{b}.{leg}"]}
+    # Eager registration (mirrors the runner hook tagging every superstep task at
+    # start), so the coordinator waits for the whole batch before flushing.
+    for k in keys:
+        coord.register(k)
 
-    def inner_fan(state: _Inner):
-        return [Send("inner_worker", {"branch": state["branch"], "leg": leg}) for leg in ("L", "R")]
+    async def leaf(key: str):
+        def make_future():
+            # make_future() is called at flush time; record the CREATION order
+            # (this is the order commands would be appended to the journal).
+            creation_order.append(key)
 
-    b = StateGraph(_Inner)
-    b.add_node("inner_worker", inner_worker)
-    b.add_conditional_edges(START, inner_fan, ["inner_worker"])
-    b.add_edge("inner_worker", END)
-    return b.compile()
+            async def io():
+                await asyncio.sleep(0)  # simulate concurrent I/O completion
+                return f"result:{key}"
 
+            return io()
 
-class _Outer(TypedDict):
-    branches: list
-    results: Annotated[list, operator.add]
+        try:
+            return await coord.submit(key, make_future)
+        finally:
+            coord.checkout(key)
 
+    results = await asyncio.gather(*(leaf(k) for k in keys))
 
-class _OuterWS(TypedDict):
-    branch: str
+    # Commands were created in sorted key order — independent of arrival order.
+    assert creation_order == sorted(keys)
+    # Every op still got its own result routed back.
+    assert set(results) == {f"result:{k}" for k in keys}
 
 
-def _build(dctx):
-    sub = _make_sub(dctx)
+async def test_coordinator_multiple_rounds_stay_ordered():
+    coord = _Coordinator(settle_turns=0)
+    rounds: list[list[str]] = [[], []]
+    keys = ["/w|0002", "/w|0000", "/w|0001"]
+    for k in keys:
+        coord.register(k)
 
-    async def branch_node(state: _OuterWS) -> dict:
-        out = await sub.ainvoke({"branch": state["branch"], "results": []})
-        return {"results": out["results"]}
+    async def leaf(key: str):
+        for r in (0, 1):  # two sequential durable ops per task
 
-    def outer_fan(state: _Outer):
-        return [Send("branch_node", {"branch": b}) for b in state["branches"]]
+            def make_future(_r=r, _k=key):
+                rounds[_r].append(_k)
 
-    b = StateGraph(_Outer)
-    b.add_node("branch_node", branch_node)
-    b.add_conditional_edges(START, outer_fan, ["branch_node"])
-    b.add_edge("branch_node", END)
-    return b.compile()
+                async def io():
+                    await asyncio.sleep(0)
+                    return _k
 
+                return io()
 
-async def _one_run() -> tuple:
-    journal: list = []
-    token = _journal.set(journal)
-    try:
-        with durable_scope(FakeCtx()) as dctx:
-            await _build(dctx).ainvoke({"branches": ["A", "B", "C"], "results": []})
-    finally:
-        _journal.reset(token)
-    return tuple(journal)
+            await coord.submit(key, make_future)
+        coord.checkout(key)
 
+    await asyncio.gather(*(leaf(k) for k in keys))
 
-async def test_stategraph_durable_order_is_deterministic():
-    runs = [await _one_run() for _ in range(20)]
-    distinct = set(runs)
-    assert len(distinct) == 1, f"non-deterministic journal order: {distinct}"
-    # 3 branches x 2 legs x 2 ops = 12 durable commands per run
-    assert len(runs[0]) == 12
+    # Each round's commands are independently sorted.
+    assert rounds[0] == sorted(keys)
+    assert rounds[1] == sorted(keys)

From eec95f921a5dc312fd499460501ef47de3b71ea1 Mon Sep 17 00:00:00 2001
From: Jamie Grier <jamie.grier@gmail.com>
Date: Fri, 24 Jul 2026 09:26:29 -0700
Subject: [PATCH 3/5] test(langchain): cover concurrent tool-call-task ordering
 in the coordinator

Adds a coordinator test modeling a node's concurrent tool calls as separate
Pregel push tasks with distinct hierarchical keys (verified against LangGraph:
.../tools|('__pregel_push', N, False)), alongside sibling single-op nodes.
Asserts they flush in stable sorted-key order with no same-key collision.

This documents the actual mechanism for multi-tool-call parallel nodes: the
coordinator orders the distinct-key tool-call tasks directly; the middleware
turnstile is not what serializes them. (Empirically: create_agent v1 doesn't
invoke the turnstile for tool exec; deepagents does, trivially for a single
call. The per-tool_call_id turnstile keying still matters to prevent the
cross-node clobber KeyError when parallel nodes share the middleware state.)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 tests/langchain_stategraph.py | 37 +++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/tests/langchain_stategraph.py b/tests/langchain_stategraph.py
index bc0931b..a9736ac 100644
--- a/tests/langchain_stategraph.py
+++ b/tests/langchain_stategraph.py
@@ -68,6 +68,43 @@ async def io():
     assert set(results) == {f"result:{k}" for k in keys}
 
 
+async def test_coordinator_orders_concurrent_tool_call_tasks():
+    # A node's concurrent tool calls run as SEPARATE Pregel push tasks with
+    # distinct hierarchical keys (verified against LangGraph: a tool-using node
+    # "deep" spawns .../tools|('__pregel_push', N, False) tasks). The coordinator
+    # must order those distinct-key ops deterministically and never collide,
+    # alongside sibling single-op nodes — all in one flush round.
+    coord = _Coordinator(settle_turns=0)
+    order: list[str] = []
+    tool_tasks = [f"/deep|(pull)/tools|(push,{i})" for i in range(3)]
+    sibling_tasks = ["/branch_b|(pull)", "/branch_a|(pull)"]
+    all_keys = tool_tasks + sibling_tasks
+    for k in all_keys:
+        coord.register(k)
+
+    async def leaf(key: str):
+        def make_future():
+            order.append(key)
+
+            async def io():
+                await asyncio.sleep(0)
+                return key
+
+            return io()
+
+        try:
+            return await coord.submit(key, make_future)
+        finally:
+            coord.checkout(key)
+
+    results = await asyncio.gather(*(leaf(k) for k in all_keys))
+
+    # All ops created (none lost to a same-key collision), in sorted key order.
+    assert len(order) == len(all_keys)
+    assert order == sorted(all_keys)
+    assert set(results) == set(all_keys)
+
+
 async def test_coordinator_multiple_rounds_stay_ordered():
     coord = _Coordinator(settle_turns=0)
     rounds: list[list[str]] = [[], []]

From 0429c8131a61a683a553cdcb98649368bda8e0dd Mon Sep 17 00:00:00 2001
From: Jamie Grier <jamie.grier@gmail.com>
Date: Fri, 24 Jul 2026 18:25:59 -0700
Subject: [PATCH 4/5] fix(langchain): skip middleware turnstile under the
 coordinator (deadlock)

A node with multiple tool calls that invokes the turnstile (e.g. deepagents)
DEADLOCKED under the StateGraph coordinator: the tool calls run as separate
Pregel tasks (distinct keys), all registered as active leaves; the turnstile
blocks tool calls 2..N in wait_for BEFORE they submit, so they never become
pending, so the coordinator's "flush when every active leaf is pending" barrier
never fires, so tool call 1 never completes, so the turnstile never releases the
rest. Circular deadlock (confirmed by tracing).

Fix: awrap_tool_call skips the turnstile when a coordinator is active
(coordinator_active()). The coordinator already orders the tool-call tasks by
key, so the turnstile is redundant there; the stock (no-coordinator) path keeps
using it. Verified: deepagents x3 tool calls completes under the fix (turnstile
skipped) and under stock (turnstile serializes); his full graph with a 3-tool-call
deep node reproduces the 570 in baseline and is 0x570 with the fix.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 python/restate/ext/langchain/_middleware.py | 15 +++++++++++++++
 python/restate/ext/langchain/_stategraph.py |  8 ++++++++
 tests/langchain_stategraph.py               | 15 ++++++++++++++-
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/python/restate/ext/langchain/_middleware.py b/python/restate/ext/langchain/_middleware.py
index 95f862a..8f80e7b 100644
--- a/python/restate/ext/langchain/_middleware.py
+++ b/python/restate/ext/langchain/_middleware.py
@@ -36,6 +36,7 @@
 from restate.ext.turnstile import Turnstile
 
 from ._state import get_or_create_state, state_from_ctx
+from ._stategraph import coordinator_active as _stategraph_coordinator_active
 
 ToolCallResult = ToolMessage | Command
 
@@ -125,6 +126,20 @@ async def awrap_tool_call(
 
         ctx = current_context()
         assert ctx is not None, "RestateMiddleware must run inside a Restate handler"
+
+        if _stategraph_coordinator_active():
+            # The StateGraph determinism coordinator is active. A node's tool
+            # calls are separate Pregel tasks with distinct keys that the
+            # coordinator already orders deterministically; the turnstile is
+            # redundant here AND would deadlock against the coordinator's
+            # "flush when every active leaf is pending" barrier (a turnstile
+            # blocks later tool calls *before* they submit, so they never become
+            # pending, so the first never flushes). Skip the turnstile.
+            result = await handler(request)
+            if isinstance(result, ToolMessage):
+                result.id = str(ctx.uuid())
+            return result
+
         state = state_from_ctx(ctx)
         assert state is not None, "RestateMiddleware must run inside a Restate handler"
         turnstile = state.turnstiles.get(tool_call_id)
diff --git a/python/restate/ext/langchain/_stategraph.py b/python/restate/ext/langchain/_stategraph.py
index 0e70053..9fce1b2 100644
--- a/python/restate/ext/langchain/_stategraph.py
+++ b/python/restate/ext/langchain/_stategraph.py
@@ -204,6 +204,14 @@ async def submit(self, full_key: str, make_future: Callable[[], Any]) -> Any:
     "_restate_lg_active_coord", default=None
 )
 
+
+def coordinator_active() -> bool:
+    """True if a StateGraph determinism coordinator is active on the current
+    context. RestateMiddleware uses this to skip its tool-call turnstile, which
+    is redundant with — and would deadlock against — the coordinator."""
+    return _active_coord.get() is not None
+
+
 _installed = False
 _auto_enabled = False
 _orig_run_typed: Any = None
diff --git a/tests/langchain_stategraph.py b/tests/langchain_stategraph.py
index a9736ac..39dbdda 100644
--- a/tests/langchain_stategraph.py
+++ b/tests/langchain_stategraph.py
@@ -21,11 +21,24 @@
 
 import pytest
 
-from restate.ext.langchain._stategraph import _Coordinator
+from restate.ext.langchain._stategraph import _active_coord, _Coordinator, coordinator_active
 
 pytestmark = [pytest.mark.anyio]
 
 
+async def test_coordinator_active_flag_gates_the_turnstile():
+    # RestateMiddleware.awrap_tool_call consults coordinator_active() to skip its
+    # turnstile when the StateGraph coordinator is on (the turnstile is redundant
+    # there and would deadlock against the coordinator's all-leaves-pending flush).
+    assert coordinator_active() is False
+    token = _active_coord.set(_Coordinator())
+    try:
+        assert coordinator_active() is True
+    finally:
+        _active_coord.reset(token)
+    assert coordinator_active() is False
+
+
 @pytest.fixture(scope="session")
 def anyio_backend():
     return "asyncio"

From dbeed48f9985d1326bc3234cdae7e874dff3f982 Mon Sep 17 00:00:00 2001
From: Jamie Grier <jamie.grier@gmail.com>
Date: Fri, 24 Jul 2026 19:51:48 -0700
Subject: [PATCH 5/5] fix(langchain): satisfy CI type checks (mypy assignment
 ignore, pyright import)

- mypy runs on python/ (not just tests/): the class-level run_typed patch needs
  type: ignore[method-assign, assignment], not just [method-assign].
- invoke_handler is imported but not re-exported by server_context; use
  getattr/setattr so pyright doesn't flag reportPrivateImportUsage.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 python/restate/ext/langchain/_stategraph.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/python/restate/ext/langchain/_stategraph.py b/python/restate/ext/langchain/_stategraph.py
index 9fce1b2..b2f7ab3 100644
--- a/python/restate/ext/langchain/_stategraph.py
+++ b/python/restate/ext/langchain/_stategraph.py
@@ -264,7 +264,7 @@ def patched_run_typed(self, *args, **kwargs):  # type: ignore[no-untyped-def]
         # calls the ORIGINAL run_typed (no recursion) to create the command.
         return coord.submit(full, lambda: _orig_run_typed(self, *args, **kwargs))
 
-    ServerInvocationContext.run_typed = patched_run_typed  # type: ignore[method-assign]
+    ServerInvocationContext.run_typed = patched_run_typed  # type: ignore[method-assign, assignment]
 
 
 def enable() -> None:
@@ -284,7 +284,9 @@ def enable() -> None:
 
     import restate.server_context as _sc  # pylint: disable=import-outside-toplevel
 
-    orig_invoke_handler = _sc.invoke_handler
+    # getattr/setattr: invoke_handler is imported into server_context but not
+    # re-exported, so plain attribute access trips pyright's private-import check.
+    orig_invoke_handler = getattr(_sc, "invoke_handler")
 
     async def patched_invoke_handler(handler, ctx, in_buffer):  # type: ignore[no-untyped-def]
         token = _active_coord.set(_Coordinator(settle_turns=0))
@@ -293,7 +295,7 @@ async def patched_invoke_handler(handler, ctx, in_buffer):  # type: ignore[no-un
         finally:
             _active_coord.reset(token)
 
-    _sc.invoke_handler = patched_invoke_handler  # server_context.enter() resolves this name
+    setattr(_sc, "invoke_handler", patched_invoke_handler)  # server_context.enter() resolves this name
 
 
 class durable_scope:  # pylint: disable=invalid-name