feat(langchain): deterministic durability for parallel StateGraph nodes (auto-enabled)

python/restate/ext/langchain/__init__.py

@@ -17,12 +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 durable_scope, enable
 
 
 def restate_context() -> Context:
@@ -51,4 +53,16 @@ def restate_object_context() -> ObjectContext:
     "RestateMiddleware",
     "restate_context",
     "restate_object_context",
+    "durable_scope",
+    "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

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
 
@@ -101,7 +102,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(
@@ -120,9 +126,32 @@ 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.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)

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: