Merge pull request #1630 from langchain-ai/nc/5sep/scheduler-kafka

Implement LangGraph Scheduler for Kafka

.github/workflows/_lint.yml

@@ -44,7 +44,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
           poetry-version: ${{ env.POETRY_VERSION }}
           working-directory: ${{ inputs.working-directory }}
-          cache-key: lint-with-extras
+          cache-key: lint-${{ inputs.working-directory }}
 
       - name: Check Poetry File
         if: steps.changed-files.outputs.all

.github/workflows/_test.yml

@@ -34,7 +34,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
           poetry-version: ${{ env.POETRY_VERSION }}
           working-directory: ${{ inputs.working-directory }}
-          cache-key: core
+          cache-key: test-${{ inputs.working-directory }}
 
       - name: Install dependencies
         shell: bash

.github/workflows/ci.yml

@@ -1,114 +1,104 @@
 ---
-    name: CI
+name: CI
 
-    on:
-      push:
-        branches: [main]
-      pull_request:
+on:
+  push:
+    branches: [main]
+  pull_request:
 
-    # If another push to the same PR or branch happens while this workflow is still running,
-    # cancel the earlier run in favor of the next run.
-    #
-    # There's no point in testing an outdated version of the code. GitHub only allows
-    # a limited number of job runners to be active at the same time, so it's better to cancel
-    # pointless jobs early so that more useful jobs can run sooner.
-    concurrency:
-      group: ${{ github.workflow }}-${{ github.ref }}
-      cancel-in-progress: true
+# If another push to the same PR or branch happens while this workflow is still running,
+# cancel the earlier run in favor of the next run.
+#
+# There's no point in testing an outdated version of the code. GitHub only allows
+# a limited number of job runners to be active at the same time, so it's better to cancel
+# pointless jobs early so that more useful jobs can run sooner.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
 
+env:
+  POETRY_VERSION: "1.7.1"
+
+jobs:
+  lint:
+    name: cd ${{ matrix.working-directory }}
+    strategy:
+      matrix:
+        working-directory:
+          [
+            "libs/langgraph",
+            "libs/sdk-py",
+            "libs/cli",
+            "libs/checkpoint",
+            "libs/checkpoint-sqlite",
+            "libs/checkpoint-postgres",
+            "libs/scheduler-kafka",
+          ]
+    uses: ./.github/workflows/_lint.yml
+    with:
+      working-directory: ${{ matrix.working-directory }}
+    secrets: inherit
+
+  test:
+    name: cd ${{ matrix.working-directory }}
+    strategy:
+      matrix:
+        working-directory: [
+            "libs/langgraph",
+            "libs/cli",
+            "libs/checkpoint",
+            "libs/checkpoint-sqlite",
+            "libs/checkpoint-postgres",
+            # "libs/scheduler-kafka"
+          ]
+    uses: ./.github/workflows/_test.yml
+    with:
+      working-directory: ${{ matrix.working-directory }}
+    secrets: inherit
+
+  integration-test:
+    name: CLI integration test
+    uses: ./.github/workflows/_integration_test.yml
+    secrets: inherit
+
+  lint-js:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        working-directory:
+          - "libs/sdk-js"
+    defaults:
+      run:
+        working-directory: ${{ matrix.working-directory }}
+    steps:
+      - uses: actions/checkout@v3
+      - name: Setup Node.js (LTS)
+        uses: actions/setup-node@v3
+        with:
+          node-version: "20"
+          cache: "yarn"
+          cache-dependency-path: ${{ matrix.working-directory }}/yarn.lock
+      - name: Install dependencies
+        run: yarn install
+      - name: Run lint
+        run: yarn lint
+      - name: Build
+        run: yarn build
+
+  ci_success:
+    name: "CI Success"
+    needs: [lint, lint-js, test, integration-test]
+    if: |
+      always()
+    runs-on: ubuntu-latest
     env:
-      POETRY_VERSION: "1.7.1"
-
-    jobs:
-      build:
-        runs-on: ubuntu-latest
-        steps:
-          - uses: actions/checkout@v4
-          - uses: actions/setup-python@v5
-            with:
-              python-version: '3.10'
-
-      lint:
-        name: cd ${{ matrix.working-directory }}
-        needs: [ build ]
-        strategy:
-          matrix:
-            working-directory: [
-                "libs/langgraph",
-                "libs/sdk-py",
-                "libs/cli",
-                "libs/checkpoint",
-                "libs/checkpoint-sqlite",
-                "libs/checkpoint-postgres"
-            ]
-        uses: ./.github/workflows/_lint.yml
-        with:
-          working-directory: ${{ matrix.working-directory }}
-        secrets: inherit
-
-      test:
-        name: cd ${{ matrix.working-directory }}
-        needs: [ build ]
-        strategy:
-          matrix:
-            working-directory: [
-                "libs/langgraph",
-                "libs/cli",
-                "libs/checkpoint",
-                "libs/checkpoint-sqlite",
-                "libs/checkpoint-postgres"
-            ]
-        uses: ./.github/workflows/_test.yml
-        with:
-          working-directory: ${{ matrix.working-directory }}
-        secrets: inherit
-
-      integration-test:
-        name: CLI integration test
-        needs: [ build ]
-        uses: ./.github/workflows/_integration_test.yml
-        secrets: inherit
-
-      lint-js:
-        runs-on: ubuntu-latest
-        needs: [ build ]
-        strategy:
-          matrix:
-            working-directory:
-              - "libs/sdk-js"
-        defaults:
-          run:
-            working-directory: ${{ matrix.working-directory }}
-        steps:
-          - uses: actions/checkout@v3
-          - name: Setup Node.js (LTS)
-            uses: actions/setup-node@v3
-            with:
-              node-version: "20"
-              cache: "yarn"
-              cache-dependency-path: ${{ matrix.working-directory }}/yarn.lock
-          - name: Install dependencies
-            run: yarn install
-          - name: Run lint
-            run: yarn lint
-          - name: Build
-            run: yarn build
-
-      ci_success:
-        name: "CI Success"
-        needs: [build, lint, lint-js, test, integration-test]
-        if: |
-          always()
-        runs-on: ubuntu-latest
-        env:
-          JOBS_JSON: ${{ toJSON(needs) }}
-          RESULTS_JSON: ${{ toJSON(needs.*.result) }}
-          EXIT_CODE: ${{!contains(needs.*.result, 'failure') && !contains(needs.*.result, 'cancelled') && '0' || '1'}}
-        steps:
-          - name: "CI Success"
-            run: |
-              echo $JOBS_JSON
-              echo $RESULTS_JSON
-              echo "Exiting with $EXIT_CODE"
-              exit $EXIT_CODE
-    
+      JOBS_JSON: ${{ toJSON(needs) }}
+      RESULTS_JSON: ${{ toJSON(needs.*.result) }}
+      EXIT_CODE: ${{!contains(needs.*.result, 'failure') && !contains(needs.*.result, 'cancelled') && '0' || '1'}}
+    steps:
+      - name: "CI Success"
+        run: |
+          echo $JOBS_JSON
+          echo $RESULTS_JSON
+          echo "Exiting with $EXIT_CODE"
+          exit $EXIT_CODE

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -10,6 +10,7 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -337,9 +338,14 @@ class PostgresSaver(BasePostgresSaver):
             writes (List[Tuple[str, Any]]): List of writes to store.
             task_id (str): Identifier for the task creating the writes.
         """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
         with self._cursor(pipeline=True) as cur:
             cur.executemany(
-                self.UPSERT_CHECKPOINT_WRITES_SQL,
+                query,
                 self._dump_writes(
                     config["configurable"]["thread_id"],
                     config["configurable"]["checkpoint_ns"],

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -10,6 +10,7 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -292,18 +293,21 @@ class AsyncPostgresSaver(BasePostgresSaver):
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
         """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        params = await asyncio.to_thread(
+            self._dump_writes,
+            config["configurable"]["thread_id"],
+            config["configurable"]["checkpoint_ns"],
+            config["configurable"]["checkpoint_id"],
+            task_id,
+            writes,
+        )
         async with self._cursor(pipeline=True) as cur:
-            await cur.executemany(
-                self.UPSERT_CHECKPOINT_WRITES_SQL,
-                await asyncio.to_thread(
-                    self._dump_writes,
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    writes,
-                ),
-            )
+            await cur.executemany(query, params)
 
     @asynccontextmanager
     async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -108,6 +108,15 @@ UPSERT_CHECKPOINTS_SQL = """
 """
 
 UPSERT_CHECKPOINT_WRITES_SQL = """
+    INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
+    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
+        channel = EXCLUDED.channel,
+        type = EXCLUDED.type,
+        blob = EXCLUDED.blob;
+"""
+
+INSERT_CHECKPOINT_WRITES_SQL = """
     INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
     VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
@@ -120,6 +129,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
+    INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
 

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py

@@ -435,9 +435,14 @@ class SqliteSaver(BaseCheckpointSaver):
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
         """
+        query = (
+            "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+        )
         with self.cursor() as cur:
             cur.executemany(
-                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                query,
                 [
                     (
                         str(config["configurable"]["thread_id"]),

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py

@@ -476,10 +476,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
         """
+        query = (
+            "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
+        )
         await self.setup()
         async with self.lock, self.conn.cursor() as cur:
             await cur.executemany(
-                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                query,
                 [
                     (
                         str(config["configurable"]["thread_id"]),

libs/langgraph/langgraph/constants.py

@@ -9,20 +9,27 @@ CONFIG_KEY_STREAM = "__pregel_stream"
 CONFIG_KEY_STORE = "__pregel_store"
 CONFIG_KEY_RESUMING = "__pregel_resuming"
 CONFIG_KEY_TASK_ID = "__pregel_task_id"
+CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks"
+CONFIG_KEY_ENSURE_LATEST = "__pregel_ensure_latest"
+CONFIG_KEY_DELEGATE = "__pregel_delegate"
 # this one part of public API so more readable
 CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map"
 INTERRUPT = "__interrupt__"
 ERROR = "__error__"
+NO_WRITES = "__no_writes__"
 SCHEDULED = "__scheduled__"
-TASKS = "__pregel_tasks"
-SUBSCRIPTIONS = "__pregel_subscriptions"
+TASKS = "__pregel_tasks"  # for backwards compat, this is the original name of PUSH
+PUSH = "__pregel_push"
+PULL = "__pregel_pull"
 RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__"
 RESERVED = {
     SCHEDULED,
     INTERRUPT,
     ERROR,
+    NO_WRITES,
     TASKS,
-    SUBSCRIPTIONS,
+    PUSH,
+    PULL,
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
     CONFIG_KEY_CHECKPOINTER,
@@ -30,6 +37,9 @@ RESERVED = {
     CONFIG_KEY_STORE,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
+    CONFIG_KEY_DEDUPE_TASKS,
+    CONFIG_KEY_ENSURE_LATEST,
+    CONFIG_KEY_DELEGATE,
     INPUT,
     RUNTIME_PLACEHOLDER,
 }

libs/langgraph/langgraph/errors.py

@@ -43,6 +43,13 @@ class NodeInterrupt(GraphInterrupt):
         super().__init__([Interrupt(value)])
 
 
+class GraphDelegate(Exception):
+    """Raised when a graph is delegated."""
+
+    def __init__(self, *args: dict[str, Any]) -> None:
+        super().__init__(*args)
+
+
 class EmptyInputError(Exception):
     """Raised when graph receives an empty input."""
 
@@ -55,6 +62,12 @@ class TaskNotFound(Exception):
     pass
 
 
+class CheckpointNotLatest(Exception):
+    """Raised when the checkpoint is not the latest version."""
+
+    pass
+
+
 __all__ = [
     "GraphRecursionError",
     "InvalidUpdateError",

libs/langgraph/langgraph/managed/is_last_step.py

@@ -5,7 +5,7 @@ from langgraph.managed.base import ManagedValue
 
 class IsLastStepManager(ManagedValue[bool]):
     def __call__(self, step: int) -> bool:
-        return step == self.config["recursion_limit"] - 1
+        return step == self.config.get("recursion_limit", 0) - 1
 
 
 IsLastStep = Annotated[bool, IsLastStepManager]

libs/langgraph/langgraph/pregel/__init__.py

@@ -21,7 +21,6 @@ from typing import (
 from uuid import UUID, uuid5
 
 from langchain_core.globals import get_debug
-from langchain_core.load.dump import dumpd
 from langchain_core.runnables import (
     Runnable,
     RunnableLambda,
@@ -782,6 +781,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
             last_seen_by_node = sorted(
                 (v, n)
                 for n, seen in checkpoint["versions_seen"].items()
+                if n in self.nodes
                 for v in seen.values()
             )
             # if two nodes updated the state at the same time, it's ambiguous
@@ -939,6 +939,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
             last_seen_by_node = sorted(
                 (v, n)
                 for n, seen in checkpoint["versions_seen"].items()
+                if n in self.nodes
                 for v in seen.values()
             )
             # if two nodes updated the state at the same time, it's ambiguous
@@ -1158,7 +1159,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
         config = ensure_config(merge_configs(self.config, config))
         callback_manager = get_callback_manager_for_config(config)
         run_manager = callback_manager.on_chain_start(
-            dumpd(self),
+            None,
             input,
             name=config.get("run_name", self.get_name()),
             run_id=config.get("run_id"),
@@ -1339,7 +1340,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
         config = ensure_config(merge_configs(self.config, config))
         callback_manager = get_async_callback_manager_for_config(config)
         run_manager = await callback_manager.on_chain_start(
-            dumpd(self),
+            None,
             input,
             name=config.get("run_name", self.get_name()),
             run_id=config.get("run_id"),

libs/langgraph/langgraph/pregel/algo.py

@@ -25,13 +25,14 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINTER,
     CONFIG_KEY_READ,
-    CONFIG_KEY_RESUMING,
     CONFIG_KEY_SEND,
     CONFIG_KEY_TASK_ID,
     INTERRUPT,
+    NO_WRITES,
     NS_SEP,
+    PULL,
+    PUSH,
     RESERVED,
-    SUBSCRIPTIONS,
     TAG_HIDDEN,
     TASKS,
     Send,
@@ -196,7 +197,9 @@ def apply_writes(
     pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
-            if chan == TASKS:
+            if chan == NO_WRITES:
+                pass
+            elif chan == TASKS:
                 checkpoint["pending_sends"].append(val)
             elif chan in channels:
                 pending_writes_by_channel[chan].append(val)
@@ -242,7 +245,6 @@ def prepare_next_tasks(
     step: int,
     *,
     for_execution: Literal[False],
-    is_resuming: bool = False,
     checkpointer: Literal[None] = None,
     manager: Literal[None] = None,
 ) -> dict[str, PregelTask]: ...
@@ -258,7 +260,6 @@ def prepare_next_tasks(
     step: int,
     *,
     for_execution: Literal[True],
-    is_resuming: bool,
     checkpointer: Optional[BaseCheckpointSaver],
     manager: Union[None, ParentRunManager, AsyncParentRunManager],
 ) -> dict[str, PregelExecutableTask]: ...
@@ -273,7 +274,6 @@ def prepare_next_tasks(
     step: int,
     *,
     for_execution: bool,
-    is_resuming: bool = False,
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
 ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
@@ -281,7 +281,7 @@ def prepare_next_tasks(
     # Consume pending packets
     for idx, _ in enumerate(checkpoint["pending_sends"]):
         if task := prepare_single_task(
-            (TASKS, idx),
+            (PUSH, idx),
             None,
             checkpoint=checkpoint,
             processes=processes,
@@ -290,7 +290,6 @@ def prepare_next_tasks(
             config=config,
             step=step,
             for_execution=for_execution,
-            is_resuming=is_resuming,
             checkpointer=checkpointer,
             manager=manager,
         ):
@@ -299,7 +298,7 @@ def prepare_next_tasks(
     # If so, prepare the values to be passed to them
     for name in processes:
         if task := prepare_single_task(
-            (SUBSCRIPTIONS, name),
+            (PULL, name),
             None,
             checkpoint=checkpoint,
             processes=processes,
@@ -308,7 +307,6 @@ def prepare_next_tasks(
             config=config,
             step=step,
             for_execution=for_execution,
-            is_resuming=is_resuming,
             checkpointer=checkpointer,
             manager=manager,
         ):
@@ -327,7 +325,6 @@ def prepare_single_task(
     config: RunnableConfig,
     step: int,
     for_execution: bool,
-    is_resuming: bool = False,
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
 ) -> Union[None, PregelTask, PregelExecutableTask]:
@@ -335,8 +332,10 @@ def prepare_single_task(
     configurable = config.get("configurable", {})
     parent_ns = configurable.get("checkpoint_ns", "")
 
-    if task_path[0] == TASKS:
+    if task_path[0] == PUSH:
         idx = int(task_path[1])
+        if idx >= len(checkpoint["pending_sends"]):
+            return
         packet = checkpoint["pending_sends"][idx]
         if not isinstance(packet, Send):
             logger.warning(
@@ -347,7 +346,7 @@ def prepare_single_task(
             logger.warning(f"Ignoring unknown node name {packet.node} in pending sends")
             return
         # create task id
-        triggers = [TASKS]
+        triggers = [PUSH]
         metadata = {
             "langgraph_step": step,
             "langgraph_node": packet.node,
@@ -362,7 +361,7 @@ def prepare_single_task(
             checkpoint_ns,
             str(step),
             packet.node,
-            TASKS,
+            PUSH,
             str(idx),
         )
         if task_id_checksum is not None:
@@ -416,7 +415,6 @@ def prepare_single_task(
                                 **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
                                 parent_ns: checkpoint["id"],
                             },
-                            CONFIG_KEY_RESUMING: is_resuming,
                             "checkpoint_id": None,
                             "checkpoint_ns": task_checkpoint_ns,
                         },
@@ -429,9 +427,11 @@ def prepare_single_task(
                 )
 
         else:
-            return PregelTask(task_id, packet.node)
-    elif task_path[0] == SUBSCRIPTIONS:
+            return PregelTask(task_id, packet.node, task_path)
+    elif task_path[0] == PULL:
         name = str(task_path[1])
+        if name not in processes:
+            return
         proc = processes[name]
         version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
         null_version = version_type()
@@ -470,7 +470,7 @@ def prepare_single_task(
                 checkpoint_ns,
                 str(step),
                 name,
-                SUBSCRIPTIONS,
+                PULL,
                 *triggers,
             )
             if task_id_checksum is not None:
@@ -525,7 +525,7 @@ def prepare_single_task(
                                     **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
                                     parent_ns: checkpoint["id"],
                                 },
-                                CONFIG_KEY_RESUMING: is_resuming,
+                                "checkpoint_id": None,
                                 "checkpoint_ns": task_checkpoint_ns,
                             },
                         ),
@@ -536,7 +536,7 @@ def prepare_single_task(
                         task_path,
                     )
             else:
-                return PregelTask(task_id, name)
+                return PregelTask(task_id, name, task_path)
 
 
 def _proc_input(

libs/langgraph/langgraph/pregel/debug.py

@@ -210,6 +210,7 @@ def tasks_w_writes(
         PregelTask(
             task.id,
             task.name,
+            task.path,
             next(
                 (
                     exc

libs/langgraph/langgraph/pregel/executor.py

@@ -31,6 +31,7 @@ class Submit(Protocol[P, T]):
         *args: P.args,
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]: ...
 
@@ -39,7 +40,7 @@ class BackgroundExecutor(ContextManager):
     def __init__(self, config: RunnableConfig) -> None:
         self.stack = ExitStack()
         self.executor = self.stack.enter_context(get_executor_for_config(config))
-        self.tasks: dict[concurrent.futures.Future, bool] = {}
+        self.tasks: dict[concurrent.futures.Future, tuple[bool, bool]] = {}
 
     def submit(
         self,
@@ -47,10 +48,11 @@ class BackgroundExecutor(ContextManager):
         *args: P.args,
         __name__: Optional[str] = None,  # currently not used in sync version
         __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]:
         task = self.executor.submit(fn, *args, **kwargs)
-        self.tasks[task] = __cancel_on_exit__
+        self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task
 
@@ -76,7 +78,7 @@ class BackgroundExecutor(ContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # cancel all tasks that should be cancelled
-        for task, cancel in self.tasks.items():
+        for task, (cancel, _) in self.tasks.items():
             if cancel:
                 task.cancel()
         # wait for all tasks to finish
@@ -87,7 +89,9 @@ class BackgroundExecutor(ContextManager):
         # re-raise the first exception that occurred in a task
         if exc_type is None:
             # if there's already an exception being raised, don't raise another one
-            for task in self.tasks:
+            for task, (_, reraise) in self.tasks.items():
+                if not reraise:
+                    continue
                 try:
                     task.result()
                 except concurrent.futures.CancelledError:
@@ -97,7 +101,7 @@ class BackgroundExecutor(ContextManager):
 class AsyncBackgroundExecutor(AsyncContextManager):
     def __init__(self) -> None:
         self.context_not_supported = sys.version_info < (3, 11)
-        self.tasks: dict[asyncio.Task, bool] = {}
+        self.tasks: dict[asyncio.Task, tuple[bool, bool]] = {}
         self.sentinel = object()
         self.loop = asyncio.get_running_loop()
 
@@ -107,6 +111,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         *args: P.args,
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> asyncio.Task[T]:
         coro = fn(*args, **kwargs)
@@ -114,7 +119,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
             task = self.loop.create_task(coro, name=__name__)
         else:
             task = self.loop.create_task(coro, name=__name__, context=copy_context())
-        self.tasks[task] = __cancel_on_exit__
+        self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task
 
@@ -133,14 +138,14 @@ class AsyncBackgroundExecutor(AsyncContextManager):
     async def __aenter__(self) -> Submit:
         return self.submit
 
-    async def exit(
+    async def __aexit__(
         self,
         exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> None:
         # cancel all tasks that should be cancelled
-        for task, cancel in self.tasks.items():
+        for task, (cancel, _) in self.tasks.items():
             if cancel:
                 task.cancel(self.sentinel)
         # wait for all tasks to finish
@@ -149,21 +154,11 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         # if there's already an exception being raised, don't raise another one
         if exc_type is None:
             # re-raise the first exception that occurred in a task
-            for task in self.tasks:
+            for task, (_, reraise) in self.tasks.items():
+                if not reraise:
+                    continue
                 try:
                     if exc := task.exception():
                         raise exc
                 except asyncio.CancelledError:
                     pass
-
-    async def __aexit__(
-        self,
-        exc_type: Optional[type[BaseException]],
-        exc_value: Optional[BaseException],
-        traceback: Optional[TracebackType],
-    ) -> Optional[bool]:
-        # we cannot use `await` outside of asyncio.shield, as this code can run
-        # after owning task is cancelled, so pulling async logic to separate method
-
-        # wait for all background tasks to finish, shielded from cancellation
-        await asyncio.shield(self.exit(exc_type, exc_value, traceback))

libs/langgraph/langgraph/pregel/loop.py

@@ -39,6 +39,9 @@ from langgraph.checkpoint.base import (
 )
 from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_MAP,
+    CONFIG_KEY_DEDUPE_TASKS,
+    CONFIG_KEY_DELEGATE,
+    CONFIG_KEY_ENSURE_LATEST,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_STREAM,
     CONFIG_KEY_TASK_ID,
@@ -48,7 +51,12 @@ from langgraph.constants import (
     SCHEDULED,
     TAG_HIDDEN,
 )
-from langgraph.errors import EmptyInputError, GraphInterrupt
+from langgraph.errors import (
+    CheckpointNotLatest,
+    EmptyInputError,
+    GraphDelegate,
+    GraphInterrupt,
+)
 from langgraph.managed.base import (
     ManagedValueMapping,
     ManagedValueSpec,
@@ -179,7 +187,10 @@ class PregelLoop:
         self.output_keys = output_keys
         self.stream_keys = stream_keys
         self.is_nested = CONFIG_KEY_TASK_ID in self.config.get("configurable", {})
-        self.skip_done_tasks = "checkpoint_id" not in config["configurable"]
+        self.skip_done_tasks = (
+            "checkpoint_id" not in config["configurable"]
+            or CONFIG_KEY_DEDUPE_TASKS in config["configurable"]
+        )
         self.debug = debug
         if CONFIG_KEY_STREAM in config["configurable"]:
             self.stream = DuplexStream(
@@ -187,7 +198,7 @@ class PregelLoop:
             )
         if not self.is_nested and config["configurable"].get("checkpoint_ns"):
             self.config = patch_configurable(
-                config, {"checkpoint_ns": "", "checkpoint_id": None}
+                self.config, {"checkpoint_ns": "", "checkpoint_id": None}
             )
         if (
             CONFIG_KEY_CHECKPOINT_MAP in self.config["configurable"]
@@ -238,7 +249,6 @@ class PregelLoop:
     ) -> bool:
         """Execute a single iteration of the Pregel loop.
         Returns True if more iterations are needed."""
-
         if self.status != "pending":
             raise RuntimeError("Cannot tick when status is no longer 'pending'")
 
@@ -310,7 +320,6 @@ class PregelLoop:
             for_execution=True,
             manager=manager,
             checkpointer=self.checkpointer,
-            is_resuming=self.input is INPUT_RESUMING,
         )
 
         # produce debug output
@@ -334,6 +343,18 @@ class PregelLoop:
             self.status = "done"
             return False
 
+        # check if we should delegate (used by subgraphs in distributed mode)
+        if self.config["configurable"].get(CONFIG_KEY_DELEGATE):
+            assert self.input is INPUT_RESUMING
+            raise GraphDelegate(
+                {
+                    "config": patch_configurable(
+                        self.config, {CONFIG_KEY_DELEGATE: False}
+                    ),
+                    "input": None,
+                }
+            )
+
         # if there are pending writes from a previous loop, apply them
         if self.skip_done_tasks and self.checkpoint_pending_writes:
             for tid, k, v in self.checkpoint_pending_writes:
@@ -341,7 +362,13 @@ class PregelLoop:
                     continue
                 if task := self.tasks.get(tid):
                     if k == SCHEDULED:
-                        self.tasks[tid] = task._replace(scheduled=True)
+                        if v == max(
+                            self.checkpoint["versions_seen"]
+                            .get(INTERRUPT, {})
+                            .values(),
+                            default=None,
+                        ):
+                            self.tasks[tid] = task._replace(scheduled=True)
                     else:
                         task.writes.append((k, v))
             # print output for any tasks we applied previous writes to
@@ -384,9 +411,9 @@ class PregelLoop:
         # resuming from previous checkpoint requires
         # - finding a previous checkpoint
         # - receiving None input (outer graph) or RESUMING flag (subgraph)
+        configurable = self.config.get("configurable", {})
         is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
-            self.config.get("configurable", {}).get(CONFIG_KEY_RESUMING)
-            or self.input is None
+            configurable.get(CONFIG_KEY_RESUMING, self.input is None)
         )
 
         # proceed past previous checkpoint
@@ -403,6 +430,16 @@ class PregelLoop:
             )
         # map inputs to channel updates
         elif input_writes := deque(map_input(input_keys, self.input)):
+            # check if we should delegate (used by subgraphs in distributed mode)
+            if self.config["configurable"].get(CONFIG_KEY_DELEGATE):
+                raise GraphDelegate(
+                    {
+                        "config": patch_configurable(
+                            self.config, {CONFIG_KEY_DELEGATE: False}
+                        ),
+                        "input": self.input,
+                    }
+                )
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
@@ -423,10 +460,15 @@ class PregelLoop:
             ), "Can't write to SharedValues in graph input"
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": dict(input_writes)})
-        else:
+        elif CONFIG_KEY_RESUMING not in configurable:
             raise EmptyInputError(f"Received no input for {input_keys}")
         # done with input
         self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
+        # update config
+        if not self.is_nested:
+            self.config = patch_configurable(
+                self.config, {CONFIG_KEY_RESUMING: is_resuming}
+            )
 
     def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
         # assign step
@@ -586,15 +628,31 @@ class SyncPregelLoop(PregelLoop, ContextManager):
     # context manager
 
     def __enter__(self) -> Self:
-        saved = (
-            self.checkpointer.get_tuple(self.checkpoint_config)
-            if self.checkpointer
-            else None
-        ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
+        if self.config.get("configurable", {}).get(
+            CONFIG_KEY_ENSURE_LATEST
+        ) and self.checkpoint_config["configurable"].get("checkpoint_id"):
+            saved = self.checkpointer.get_tuple(
+                patch_configurable(self.checkpoint_config, {"checkpoint_id": None})
+            )
+            if (
+                saved is None
+                or saved.checkpoint["id"]
+                != self.checkpoint_config["configurable"]["checkpoint_id"]
+            ):
+                raise CheckpointNotLatest
+        elif self.checkpointer:
+            saved = self.checkpointer.get_tuple(self.checkpoint_config)
+        else:
+            saved = None
+        if saved is None:
+            saved = CheckpointTuple(
+                self.config, empty_checkpoint(), {"step": -2}, None, []
+            )
         self.checkpoint_config = {
             **self.config,
             **saved.config,
             "configurable": {
+                "checkpoint_ns": "",
                 **self.config.get("configurable", {}),
                 **saved.config.get("configurable", {}),
             },
@@ -688,15 +746,31 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
     # context manager
 
     async def __aenter__(self) -> Self:
-        saved = (
-            await self.checkpointer.aget_tuple(self.checkpoint_config)
-            if self.checkpointer
-            else None
-        ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
+        if self.config.get("configurable", {}).get(
+            CONFIG_KEY_ENSURE_LATEST
+        ) and self.checkpoint_config["configurable"].get("checkpoint_id"):
+            saved = await self.checkpointer.aget_tuple(
+                patch_configurable(self.checkpoint_config, {"checkpoint_id": None})
+            )
+            if (
+                saved is None
+                or saved.checkpoint["id"]
+                != self.checkpoint_config["configurable"]["checkpoint_id"]
+            ):
+                raise CheckpointNotLatest
+        elif self.checkpointer:
+            saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
+        else:
+            saved = None
+        if saved is None:
+            saved = CheckpointTuple(
+                self.config, empty_checkpoint(), {"step": -2}, None, []
+            )
         self.checkpoint_config = {
             **self.config,
             **saved.config,
             "configurable": {
+                "checkpoint_ns": "",
                 **self.config.get("configurable", {}),
                 **saved.config.get("configurable", {}),
             },

libs/langgraph/langgraph/pregel/runner.py

@@ -12,8 +12,8 @@ from typing import (
     Union,
 )
 
-from langgraph.constants import ERROR, INTERRUPT
-from langgraph.errors import GraphInterrupt
+from langgraph.constants import ERROR, INTERRUPT, NO_WRITES
+from langgraph.errors import GraphDelegate, GraphInterrupt
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
 from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
@@ -35,6 +35,7 @@ class PregelRunner:
         self,
         tasks: list[PregelExecutableTask],
         *,
+        reraise: bool = True,
         timeout: Optional[float] = None,
         retry_policy: Optional[RetryPolicy] = None,
     ) -> Iterator[None]:
@@ -48,6 +49,7 @@ class PregelRunner:
                 run_with_retry,
                 task,
                 retry_policy,
+                __reraise_on_exit__=reraise,
             ): task
             for task in tasks
             if not task.writes
@@ -67,12 +69,17 @@ class PregelRunner:
                 if exc := _exception(fut):
                     if isinstance(exc, GraphInterrupt):
                         # save interrupt to checkpointer
-                        self.put_writes(task.id, [(INTERRUPT, i) for i in exc.args[0]])
+                        if interrupts := [(INTERRUPT, i) for i in exc.args[0]]:
+                            self.put_writes(task.id, interrupts)
+                    elif isinstance(exc, GraphDelegate):
+                        raise exc
                     else:
                         # save error to checkpointer
                         self.put_writes(task.id, [(ERROR, exc)])
-
                 else:
+                    if not task.writes:
+                        # add no writes marker
+                        task.writes.append((NO_WRITES, None))
                     # save task writes to checkpointer
                     self.put_writes(task.id, task.writes)
             else:
@@ -84,12 +91,13 @@ class PregelRunner:
             # give control back to the caller
             yield
         # panic on failure or timeout
-        _panic_or_proceed(all_futures)
+        _panic_or_proceed(all_futures, panic=reraise)
 
     async def atick(
         self,
         tasks: list[PregelExecutableTask],
         *,
+        reraise: bool = True,
         timeout: Optional[float] = None,
         retry_policy: Optional[RetryPolicy] = None,
     ) -> AsyncIterator[None]:
@@ -107,6 +115,7 @@ class PregelRunner:
                 stream=self.use_astream,
                 __name__=task.name,
                 __cancel_on_exit__=True,
+                __reraise_on_exit__=reraise,
             ): task
             for task in tasks
             if not task.writes
@@ -126,11 +135,17 @@ class PregelRunner:
                 if exc := _exception(fut):
                     if isinstance(exc, GraphInterrupt):
                         # save interrupt to checkpointer
-                        self.put_writes(task.id, [(INTERRUPT, i) for i in exc.args[0]])
+                        if interrupts := [(INTERRUPT, i) for i in exc.args[0]]:
+                            self.put_writes(task.id, interrupts)
+                    elif isinstance(exc, GraphDelegate):
+                        raise exc
                     else:
                         # save error to checkpointer
                         self.put_writes(task.id, [(ERROR, exc)])
                 else:
+                    if not task.writes:
+                        # add no writes marker
+                        task.writes.append((NO_WRITES, None))
                     # save task writes to checkpointer
                     self.put_writes(task.id, task.writes)
             else:
@@ -142,7 +157,9 @@ class PregelRunner:
             # give control back to the caller
             yield
         # panic on failure or timeout
-        _panic_or_proceed(all_futures, asyncio.TimeoutError)
+        _panic_or_proceed(
+            all_futures, timeout_exc_cls=asyncio.TimeoutError, panic=reraise
+        )
 
 
 def _should_stop_others(
@@ -171,7 +188,9 @@ def _exception(
 
 def _panic_or_proceed(
     futs: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
+    *,
     timeout_exc_cls: Type[Exception] = TimeoutError,
+    panic: bool = True,
 ) -> None:
     done: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set()
     inflight: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set()
@@ -187,7 +206,10 @@ def _panic_or_proceed(
             while inflight:
                 inflight.pop().cancel()
             # raise the exception
-            raise exc
+            if panic:
+                raise exc
+            else:
+                return
     if inflight:
         # if we got here means we timed out
         while inflight:

libs/langgraph/langgraph/pregel/types.py

@@ -66,6 +66,7 @@ class CachePolicy(NamedTuple):
 class PregelTask(NamedTuple):
     id: str
     name: str
+    path: tuple[str, ...]
     error: Optional[Exception] = None
     interrupts: tuple[Interrupt, ...] = ()
     state: Union[None, RunnableConfig, "StateSnapshot"] = None

libs/langgraph/tests/compose-postgres.yml

@@ -1,3 +1,4 @@
+name: langgraph-tests
 services:
   postgres-test:
     image: postgres:16

libs/langgraph/tests/test_pregel.py

@@ -52,7 +52,7 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.constants import ERROR, Interrupt, Send
+from langgraph.constants import ERROR, PULL, PUSH, Interrupt, Send
 from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
@@ -741,7 +741,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
         StateSnapshot(
             values={"inbox": 4, "output": 4, "input": 3},
-            tasks=(PregelTask(AnyStr(), "two"),),
+            tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
             next=("two",),
             config={
                 "configurable": {
@@ -761,7 +761,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
         StateSnapshot(
             values={"inbox": 21, "output": 4, "input": 3},
-            tasks=(PregelTask(AnyStr(), "one"),),
+            tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
             next=("one",),
             config={
                 "configurable": {
@@ -781,7 +781,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
         StateSnapshot(
             values={"inbox": 21, "output": 4, "input": 20},
-            tasks=(PregelTask(AnyStr(), "two"),),
+            tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
             next=("two",),
             config={
                 "configurable": {
@@ -801,7 +801,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
         StateSnapshot(
             values={"inbox": 3, "output": 4, "input": 20},
-            tasks=(PregelTask(AnyStr(), "one"),),
+            tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
             next=("one",),
             config={
                 "configurable": {
@@ -836,7 +836,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
         StateSnapshot(
             values={"inbox": 3, "input": 2},
-            tasks=(PregelTask(AnyStr(), "two"),),
+            tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
             next=("two",),
             config={
                 "configurable": {
@@ -856,7 +856,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
         StateSnapshot(
             values={"input": 2},
-            tasks=(PregelTask(AnyStr(), "one"),),
+            tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
             next=("one",),
             config={
                 "configurable": {
@@ -941,7 +941,7 @@ def test_fork_always_re_runs_nodes(
         ),
         StateSnapshot(
             values=5,
-            tasks=(PregelTask(AnyStr(), "add_one"),),
+            tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
             next=("add_one",),
             config={
                 "configurable": {
@@ -961,7 +961,7 @@ def test_fork_always_re_runs_nodes(
         ),
         StateSnapshot(
             values=4,
-            tasks=(PregelTask(AnyStr(), "add_one"),),
+            tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
             next=("add_one",),
             config={
                 "configurable": {
@@ -981,7 +981,7 @@ def test_fork_always_re_runs_nodes(
         ),
         StateSnapshot(
             values=3,
-            tasks=(PregelTask(AnyStr(), "add_one"),),
+            tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
             next=("add_one",),
             config={
                 "configurable": {
@@ -1001,7 +1001,7 @@ def test_fork_always_re_runs_nodes(
         ),
         StateSnapshot(
             values=2,
-            tasks=(PregelTask(AnyStr(), "add_one"),),
+            tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
             next=("add_one",),
             config={
                 "configurable": {
@@ -1021,7 +1021,7 @@ def test_fork_always_re_runs_nodes(
         ),
         StateSnapshot(
             values=1,
-            tasks=(PregelTask(AnyStr(), "add_one"),),
+            tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
             next=("add_one",),
             config={
                 "configurable": {
@@ -1036,7 +1036,7 @@ def test_fork_always_re_runs_nodes(
         ),
         StateSnapshot(
             values=0,
-            tasks=(PregelTask(AnyStr(), "__start__"),),
+            tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
             next=("__start__",),
             config={
                 "configurable": {
@@ -1474,8 +1474,8 @@ def test_pending_writes_resume(
     assert state.values == {"value": 1}
     assert state.next == ("one", "two")
     assert state.tasks == (
-        PregelTask(AnyStr(), "one"),
-        PregelTask(AnyStr(), "two", 'ConnectionError("I\'m not good")'),
+        PregelTask(AnyStr(), "one", (PULL, "one")),
+        PregelTask(AnyStr(), "two", (PULL, "two"), 'ConnectionError("I\'m not good")'),
     )
     assert state.metadata == {
         "parents": {},
@@ -2261,7 +2261,7 @@ def test_conditional_graph(
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
@@ -2314,7 +2314,7 @@ def test_conditional_graph(
                 "input": "what is weather in sf",
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2487,7 +2487,7 @@ def test_conditional_graph(
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2534,7 +2534,7 @@ def test_conditional_graph(
                 "input": "what is weather in sf",
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2707,7 +2707,7 @@ def test_conditional_graph(
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3182,7 +3182,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3224,7 +3224,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3349,7 +3349,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3390,7 +3390,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3501,7 +3501,7 @@ def test_conditional_state_graph(
         values={
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "agent"),),
+        tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3526,7 +3526,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3580,7 +3580,7 @@ def test_conditional_state_graph(
                 ]
             ],
         },
-        tasks=(PregelTask(AnyStr(), "agent"),),
+        tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3645,7 +3645,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3699,7 +3699,7 @@ def test_conditional_state_graph(
                 ]
             ],
         },
-        tasks=(PregelTask(AnyStr(), "agent"),),
+        tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4482,7 +4482,7 @@ def test_state_graph_packets(
                 ),
             ]
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
         next=("tools",),
         config=(app_w_interrupt.checkpointer.get_tuple(config)).config,
         created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"],
@@ -4534,7 +4534,7 @@ def test_state_graph_packets(
                 ),
             ]
         },
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"],
@@ -4632,7 +4632,10 @@ def test_state_graph_packets(
                 ),
             ]
         },
-        tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")),
+        tasks=(
+            PregelTask(AnyStr(), "tools", (PUSH, 0)),
+            PregelTask(AnyStr(), "tools", (PUSH, 1)),
+        ),
         next=("tools", "tools"),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"],
@@ -4977,7 +4980,7 @@ def test_message_graph(
                 id="ai1",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5023,7 +5026,7 @@ def test_message_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=next_config,
         created_at=AnyStr(),
@@ -5104,7 +5107,7 @@ def test_message_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5207,7 +5210,7 @@ def test_message_graph(
                 id="ai1",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5253,7 +5256,7 @@ def test_message_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5334,7 +5337,7 @@ def test_message_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5428,7 +5431,7 @@ def test_message_graph(
             AIMessage(content="answer", id="ai2"),
             _AnyIdAIMessage(content="an extra message"),
         ],
-        tasks=(PregelTask(AnyStr(), "agent"),),
+        tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5703,7 +5706,7 @@ def test_root_graph(
                 id="ai1",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5749,7 +5752,7 @@ def test_root_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=next_config,
         created_at=AnyStr(),
@@ -5831,7 +5834,7 @@ def test_root_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5935,7 +5938,7 @@ def test_root_graph(
                 id="ai1",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5981,7 +5984,7 @@ def test_root_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6063,7 +6066,7 @@ def test_root_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools"),),
+        tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6156,7 +6159,7 @@ def test_root_graph(
             AIMessage(content="answer", id="ai2"),
             _AnyIdAIMessage(content="an extra message"),
         ],
-        tasks=(PregelTask(AnyStr(), "agent"),),
+        tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6228,7 +6231,7 @@ def test_root_graph(
                 _AnyIdAIMessage(content="an extra message"),
             ]
         },
-        tasks=(PregelTask(AnyStr(), "agent"),),
+        tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6581,6 +6584,7 @@ def test_dynamic_interrupt(
             PregelTask(
                 AnyStr(),
                 "tool_two",
+                (PULL, "tool_two"),
                 interrupts=(Interrupt("Just because..."),),
             ),
         ),
@@ -6673,7 +6677,7 @@ def test_start_branch_then(
     ]
     assert tool_two.get_state(thread1) == StateSnapshot(
         values={"my_key": "value ⛰️", "market": "DE"},
-        tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
         next=("tool_two_slow",),
         config=tool_two.checkpointer.get_tuple(thread1).config,
         created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6708,7 +6712,7 @@ def test_start_branch_then(
     }
     assert tool_two.get_state(thread2) == StateSnapshot(
         values={"my_key": "value", "market": "US"},
-        tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
         next=("tool_two_fast",),
         config=tool_two.checkpointer.get_tuple(thread2).config,
         created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6743,7 +6747,7 @@ def test_start_branch_then(
     }
     assert tool_two.get_state(thread3) == StateSnapshot(
         values={"my_key": "value", "market": "US"},
-        tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
         next=("tool_two_fast",),
         config=tool_two.checkpointer.get_tuple(thread3).config,
         created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6754,7 +6758,7 @@ def test_start_branch_then(
     tool_two.update_state(thread3, {"my_key": "key"})  # appends to my_key
     assert tool_two.get_state(thread3) == StateSnapshot(
         values={"my_key": "valuekey", "market": "US"},
-        tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
         next=("tool_two_fast",),
         config=tool_two.checkpointer.get_tuple(thread3).config,
         created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -7064,7 +7068,7 @@ def test_branch_then(
     }
     assert tool_two.get_state(thread1) == StateSnapshot(
         values={"my_key": "value prepared", "market": "DE"},
-        tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
         next=("tool_two_slow",),
         config=tool_two.checkpointer.get_tuple(thread1).config,
         created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7104,7 +7108,7 @@ def test_branch_then(
     }
     assert tool_two.get_state(thread2) == StateSnapshot(
         values={"my_key": "value prepared", "market": "US"},
-        tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
         next=("tool_two_fast",),
         config=tool_two.checkpointer.get_tuple(thread2).config,
         created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -7152,7 +7156,7 @@ def test_branch_then(
             "my_key": "value prepared slow",
             "market": "DE",
         },
-        tasks=(PregelTask(AnyStr(), "finish"),),
+        tasks=(PregelTask(AnyStr(), "finish", (PULL, "finish")),),
         next=("finish",),
         config=tool_two.checkpointer.get_tuple(thread1).config,
         created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7172,7 +7176,7 @@ def test_branch_then(
             "my_key": "value prepared slower",
             "market": "DE",
         },
-        tasks=(PregelTask(AnyStr(), "finish"),),
+        tasks=(PregelTask(AnyStr(), "finish", (PULL, "finish")),),
         next=("finish",),
         config=tool_two.checkpointer.get_tuple(thread1).config,
         created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7201,7 +7205,7 @@ def test_branch_then(
     }
     assert tool_two.get_state(thread1) == StateSnapshot(
         values={"my_key": "value prepared", "market": "DE"},
-        tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
         next=("tool_two_slow",),
         config=tool_two.checkpointer.get_tuple(thread1).config,
         created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7241,7 +7245,7 @@ def test_branch_then(
     }
     assert tool_two.get_state(thread2) == StateSnapshot(
         values={"my_key": "value prepared", "market": "US"},
-        tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
         next=("tool_two_fast",),
         config=tool_two.checkpointer.get_tuple(thread2).config,
         created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -7279,7 +7283,7 @@ def test_branch_then(
     # check current state
     assert tool_two.get_state(thread3) == StateSnapshot(
         values={"my_key": "key", "market": "DE"},
-        tasks=(PregelTask(AnyStr(), "prepare"),),
+        tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),),
         next=("prepare",),
         config=uconfig,
         created_at=AnyStr(),
@@ -7299,7 +7303,7 @@ def test_branch_then(
     # get state after first node
     assert tool_two.get_state(thread3) == StateSnapshot(
         values={"my_key": "key prepared", "market": "DE"},
-        tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+        tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
         next=("tool_two_slow",),
         config=tool_two.checkpointer.get_tuple(thread3).config,
         created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -7443,7 +7447,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(
             "query": "analyzed: query: what is weather in sf",
             "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
         },
-        tasks=(PregelTask(AnyStr(), "qa"),),
+        tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),),
         next=("qa",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -8631,6 +8635,7 @@ def test_nested_graph_state(
             PregelTask(
                 AnyStr(),
                 "inner",
+                (PULL, "inner"),
                 state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}},
             ),
         ),
@@ -8664,6 +8669,7 @@ def test_nested_graph_state(
             PregelTask(
                 AnyStr(),
                 "inner",
+                (PULL, "inner"),
                 state=StateSnapshot(
                     values={
                         "my_key": "hi my value here",
@@ -8672,8 +8678,8 @@ def test_nested_graph_state(
                     tasks=(
                         PregelTask(
                             AnyStr(),
-                            name="inner_2",
-                            error=None,
+                            "inner_2",
+                            (PULL, "inner_2"),
                         ),
                     ),
                     next=("inner_2",),
@@ -8743,6 +8749,7 @@ def test_nested_graph_state(
                 PregelTask(
                     AnyStr(),
                     "inner",
+                    (PULL, "inner"),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -8776,7 +8783,7 @@ def test_nested_graph_state(
         ),
         StateSnapshot(
             values={"my_key": "my value"},
-            tasks=(PregelTask(AnyStr(), "outer_1"),),
+            tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
             next=("outer_1",),
             config={
                 "configurable": {
@@ -8797,7 +8804,7 @@ def test_nested_graph_state(
         ),
         StateSnapshot(
             values={},
-            tasks=(PregelTask(AnyStr(), "__start__"),),
+            tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
             next=("__start__",),
             config={
                 "configurable": {
@@ -8851,7 +8858,7 @@ def test_nested_graph_state(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            tasks=(PregelTask(id=AnyStr(), name="inner_2"),),
+            tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),),
         ),
         StateSnapshot(
             values={"my_key": "hi my value"},
@@ -8880,7 +8887,7 @@ def test_nested_graph_state(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            tasks=(PregelTask(id=AnyStr(), name="inner_1"),),
+            tasks=(PregelTask(AnyStr(), "inner_1", (PULL, "inner_1")),),
         ),
         StateSnapshot(
             values={},
@@ -8903,7 +8910,7 @@ def test_nested_graph_state(
             },
             created_at=AnyStr(),
             parent_config=None,
-            tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+            tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
         ),
     ]
 
@@ -8971,7 +8978,7 @@ def test_nested_graph_state(
         ),
         StateSnapshot(
             values={"my_key": "hi my value here and there"},
-            tasks=(PregelTask(AnyStr(), "outer_2"),),
+            tasks=(PregelTask(AnyStr(), "outer_2", (PULL, "outer_2")),),
             next=("outer_2",),
             config={
                 "configurable": {
@@ -9001,6 +9008,7 @@ def test_nested_graph_state(
                 PregelTask(
                     AnyStr(),
                     "inner",
+                    (PULL, "inner"),
                     state={
                         "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}
                     },
@@ -9031,7 +9039,7 @@ def test_nested_graph_state(
         ),
         StateSnapshot(
             values={"my_key": "my value"},
-            tasks=(PregelTask(AnyStr(), "outer_1"),),
+            tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
             next=("outer_1",),
             config={
                 "configurable": {
@@ -9052,7 +9060,7 @@ def test_nested_graph_state(
         ),
         StateSnapshot(
             values={},
-            tasks=(PregelTask(AnyStr(), "__start__"),),
+            tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
             next=("__start__",),
             config={
                 "configurable": {
@@ -9149,6 +9157,7 @@ def test_doubly_nested_graph_state(
             PregelTask(
                 AnyStr(),
                 "child",
+                (PULL, "child"),
                 state={
                     "configurable": {
                         "thread_id": "1",
@@ -9189,6 +9198,7 @@ def test_doubly_nested_graph_state(
                 PregelTask(
                     AnyStr(),
                     "child_1",
+                    (PULL, "child_1"),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -9228,6 +9238,7 @@ def test_doubly_nested_graph_state(
             PregelTask(
                 AnyStr(),
                 "grandchild_2",
+                (PULL, "grandchild_2"),
             ),
         ),
         next=("grandchild_2",),
@@ -9272,18 +9283,21 @@ def test_doubly_nested_graph_state(
             PregelTask(
                 AnyStr(),
                 "child",
+                (PULL, "child"),
                 state=StateSnapshot(
                     values={"my_key": "hi my value"},
                     tasks=(
                         PregelTask(
                             AnyStr(),
                             "child_1",
+                            (PULL, "child_1"),
                             state=StateSnapshot(
                                 values={"my_key": "hi my value here"},
                                 tasks=(
                                     PregelTask(
                                         AnyStr(),
                                         "grandchild_2",
+                                        (PULL, "grandchild_2"),
                                     ),
                                 ),
                                 next=("grandchild_2",),
@@ -9476,7 +9490,13 @@ def test_doubly_nested_graph_state(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            tasks=(PregelTask(id=AnyStr(), name="parent_2"),),
+            tasks=(
+                PregelTask(
+                    id=AnyStr(),
+                    name="parent_2",
+                    path=(PULL, "parent_2"),
+                ),
+            ),
         ),
         StateSnapshot(
             values={"my_key": "hi my value"},
@@ -9484,6 +9504,7 @@ def test_doubly_nested_graph_state(
                 PregelTask(
                     AnyStr(),
                     "child",
+                    (PULL, "child"),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -9534,7 +9555,7 @@ def test_doubly_nested_graph_state(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            tasks=(PregelTask(id=AnyStr(), name="parent_1"),),
+            tasks=(PregelTask(id=AnyStr(), name="parent_1", path=(PULL, "parent_1")),),
         ),
         StateSnapshot(
             values={},
@@ -9554,7 +9575,9 @@ def test_doubly_nested_graph_state(
             },
             created_at=AnyStr(),
             parent_config=None,
-            tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+            tasks=(
+                PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
+            ),
         ),
     ]
     # get child graph history
@@ -9620,6 +9643,7 @@ def test_doubly_nested_graph_state(
                 PregelTask(
                     id=AnyStr(),
                     name="child_1",
+                    path=(PULL, "child_1"),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -9650,7 +9674,9 @@ def test_doubly_nested_graph_state(
             },
             created_at=AnyStr(),
             parent_config=None,
-            tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+            tasks=(
+                PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
+            ),
         ),
     ]
     # get grandchild graph history
@@ -9730,7 +9756,11 @@ def test_doubly_nested_graph_state(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            tasks=(PregelTask(id=AnyStr(), name="grandchild_2"),),
+            tasks=(
+                PregelTask(
+                    id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2")
+                ),
+            ),
         ),
         StateSnapshot(
             values={"my_key": "hi my value"},
@@ -9768,7 +9798,11 @@ def test_doubly_nested_graph_state(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            tasks=(PregelTask(id=AnyStr(), name="grandchild_1"),),
+            tasks=(
+                PregelTask(
+                    id=AnyStr(), name="grandchild_1", path=(PULL, "grandchild_1")
+                ),
+            ),
         ),
         StateSnapshot(
             values={},
@@ -9800,7 +9834,9 @@ def test_doubly_nested_graph_state(
             },
             created_at=AnyStr(),
             parent_config=None,
-            tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+            tasks=(
+                PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
+            ),
         ),
     ]
 
@@ -9875,6 +9911,7 @@ def test_send_to_nested_graphs(
             PregelTask(
                 AnyStr(),
                 "generate_joke",
+                (PUSH, 0),
                 state={
                     "configurable": {
                         "thread_id": "1",
@@ -9885,6 +9922,7 @@ def test_send_to_nested_graphs(
             PregelTask(
                 AnyStr(),
                 "generate_joke",
+                (PUSH, 1),
                 state={
                     "configurable": {
                         "thread_id": "1",
@@ -9942,7 +9980,7 @@ def test_send_to_nested_graphs(
                 "checkpoint_id": AnyStr(),
             }
         },
-        tasks=(PregelTask(id=AnyStr(""), name="generate"),),
+        tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
     )
     assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot(
         values={"subject": "dogs - hohoho", "jokes": []},
@@ -9974,7 +10012,7 @@ def test_send_to_nested_graphs(
                 "checkpoint_id": AnyStr(),
             }
         },
-        tasks=(PregelTask(id=AnyStr(""), name="generate"),),
+        tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
     )
     # update state of dogs joke graph
     graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
@@ -10069,6 +10107,7 @@ def test_send_to_nested_graphs(
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
+                    (PUSH, 0),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -10079,6 +10118,7 @@ def test_send_to_nested_graphs(
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
+                    (PUSH, 1),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -10107,7 +10147,7 @@ def test_send_to_nested_graphs(
         ),
         StateSnapshot(
             values={"jokes": []},
-            tasks=(PregelTask(AnyStr(), "__start__"),),
+            tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
             next=("__start__",),
             config={
                 "configurable": {
@@ -10284,6 +10324,7 @@ def test_weather_subgraph(
             PregelTask(
                 id=AnyStr(),
                 name="weather_graph",
+                path=(PULL, "weather_graph"),
                 state={
                     "configurable": {
                         "thread_id": "1",
@@ -10369,6 +10410,7 @@ def test_weather_subgraph(
             PregelTask(
                 id=AnyStr(),
                 name="weather_graph",
+                path=(PULL, "weather_graph"),
                 state=StateSnapshot(
                     values={
                         "messages": [
@@ -10404,7 +10446,13 @@ def test_weather_subgraph(
                             "checkpoint_id": AnyStr(),
                         }
                     },
-                    tasks=(PregelTask(id=AnyStr(), name="weather_node"),),
+                    tasks=(
+                        PregelTask(
+                            id=AnyStr(),
+                            name="weather_node",
+                            path=(PULL, "weather_node"),
+                        ),
+                    ),
                 ),
             ),
         ),
@@ -10446,6 +10494,7 @@ def test_weather_subgraph(
             PregelTask(
                 id=AnyStr(),
                 name="weather_graph",
+                path=(PULL, "weather_graph"),
                 state=StateSnapshot(
                     values={
                         "messages": [

libs/langgraph/tests/test_pregel_async.py

@@ -50,7 +50,7 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.constants import ERROR, Interrupt, Send
+from langgraph.constants import ERROR, PULL, PUSH, Interrupt, Send
 from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
@@ -298,6 +298,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "tool_two",
+                    (PULL, "tool_two"),
                     interrupts=(Interrupt("Just because..."),),
                 ),
             ),
@@ -927,7 +928,7 @@ async def test_invoke_two_processes_in_out_interrupt(
             ),
             StateSnapshot(
                 values={"inbox": 4, "output": 4, "input": 3},
-                tasks=(PregelTask(AnyStr(), "two"),),
+                tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
                 next=("two",),
                 config={
                     "configurable": {
@@ -947,7 +948,7 @@ async def test_invoke_two_processes_in_out_interrupt(
             ),
             StateSnapshot(
                 values={"inbox": 21, "output": 4, "input": 3},
-                tasks=(PregelTask(AnyStr(), "one"),),
+                tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
                 next=("one",),
                 config={
                     "configurable": {
@@ -967,7 +968,7 @@ async def test_invoke_two_processes_in_out_interrupt(
             ),
             StateSnapshot(
                 values={"inbox": 21, "output": 4, "input": 20},
-                tasks=(PregelTask(AnyStr(), "two"),),
+                tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
                 next=("two",),
                 config={
                     "configurable": {
@@ -987,7 +988,7 @@ async def test_invoke_two_processes_in_out_interrupt(
             ),
             StateSnapshot(
                 values={"inbox": 3, "output": 4, "input": 20},
-                tasks=(PregelTask(AnyStr(), "one"),),
+                tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
                 next=("one",),
                 config={
                     "configurable": {
@@ -1027,7 +1028,7 @@ async def test_invoke_two_processes_in_out_interrupt(
             ),
             StateSnapshot(
                 values={"inbox": 3, "input": 2},
-                tasks=(PregelTask(AnyStr(), "two"),),
+                tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
                 next=("two",),
                 config={
                     "configurable": {
@@ -1047,7 +1048,7 @@ async def test_invoke_two_processes_in_out_interrupt(
             ),
             StateSnapshot(
                 values={"input": 2},
-                tasks=(PregelTask(AnyStr(), "one"),),
+                tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
                 next=("one",),
                 config={
                     "configurable": {
@@ -1141,7 +1142,7 @@ async def test_fork_always_re_runs_nodes(
             ),
             StateSnapshot(
                 values=5,
-                tasks=(PregelTask(AnyStr(), "add_one"),),
+                tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
                 next=("add_one",),
                 config={
                     "configurable": {
@@ -1161,7 +1162,7 @@ async def test_fork_always_re_runs_nodes(
             ),
             StateSnapshot(
                 values=4,
-                tasks=(PregelTask(AnyStr(), "add_one"),),
+                tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
                 next=("add_one",),
                 config={
                     "configurable": {
@@ -1181,7 +1182,7 @@ async def test_fork_always_re_runs_nodes(
             ),
             StateSnapshot(
                 values=3,
-                tasks=(PregelTask(AnyStr(), "add_one"),),
+                tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
                 next=("add_one",),
                 config={
                     "configurable": {
@@ -1201,7 +1202,7 @@ async def test_fork_always_re_runs_nodes(
             ),
             StateSnapshot(
                 values=2,
-                tasks=(PregelTask(AnyStr(), "add_one"),),
+                tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
                 next=("add_one",),
                 config={
                     "configurable": {
@@ -1221,7 +1222,7 @@ async def test_fork_always_re_runs_nodes(
             ),
             StateSnapshot(
                 values=1,
-                tasks=(PregelTask(AnyStr(), "add_one"),),
+                tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
                 next=("add_one",),
                 config={
                     "configurable": {
@@ -1236,7 +1237,7 @@ async def test_fork_always_re_runs_nodes(
             ),
             StateSnapshot(
                 values=0,
-                tasks=(PregelTask(AnyStr(), "__start__"),),
+                tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
                 next=("__start__",),
                 config={
                     "configurable": {
@@ -1659,8 +1660,13 @@ async def test_pending_writes_resume(
         assert state.values == {"value": 1}
         assert state.next == ("one", "two")
         assert state.tasks == (
-            PregelTask(AnyStr(), "one"),
-            PregelTask(AnyStr(), "two", 'ConnectionError("I\'m not good")'),
+            PregelTask(AnyStr(), "one", (PULL, "one")),
+            PregelTask(
+                AnyStr(),
+                "two",
+                (PULL, "two"),
+                'ConnectionError("I\'m not good")',
+            ),
         )
         assert state.metadata == {
             "parents": {},
@@ -2524,7 +2530,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None:
                     ),
                 },
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -2575,7 +2581,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None:
                     "input": "what is weather in sf",
                 },
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -2763,7 +2769,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None:
                     ),
                 },
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -2814,7 +2820,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None:
                     "input": "what is weather in sf",
                 },
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -3002,7 +3008,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None:
                     ),
                 },
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -3427,7 +3433,7 @@ async def test_conditional_graph_state(
                 ),
                 "intermediate_steps": [],
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -3473,7 +3479,7 @@ async def test_conditional_graph_state(
                 ),
                 "intermediate_steps": [],
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -3610,7 +3616,7 @@ async def test_conditional_graph_state(
                 ),
                 "intermediate_steps": [],
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -3655,7 +3661,7 @@ async def test_conditional_graph_state(
                 ),
                 "intermediate_steps": [],
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -4297,7 +4303,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                     ),
                 ]
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
             created_at=(
@@ -4352,7 +4358,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                     ),
                 ]
             },
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
             next=("tools",),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -4452,7 +4458,10 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                     ),
                 ]
             },
-            tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")),
+            tasks=(
+                PregelTask(AnyStr(), "tools", (PUSH, 0)),
+                PregelTask(AnyStr(), "tools", (PUSH, 1)),
+            ),
             next=("tools", "tools"),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -4767,7 +4776,7 @@ async def test_message_graph(checkpointer_name: str) -> None:
                     id="ai1",
                 ),
             ],
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -4816,7 +4825,7 @@ async def test_message_graph(checkpointer_name: str) -> None:
                     ],
                 ),
             ],
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -4900,7 +4909,7 @@ async def test_message_graph(checkpointer_name: str) -> None:
                     id="ai2",
                 ),
             ],
-            tasks=(PregelTask(AnyStr(), "tools"),),
+            tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
             next=("tools",),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -5276,7 +5285,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None:
         ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value", "market": "DE"},
-            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -5319,7 +5328,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[
@@ -5362,7 +5371,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None:
         }
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -5377,7 +5386,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None:
         await tool_two.aupdate_state(thread3, {"my_key": "key"})  # appends to my_key
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "valuekey", "market": "US"},
-            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -5809,7 +5818,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
         ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value prepared", "market": "DE"},
-            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -5857,7 +5866,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value prepared", "market": "US"},
-            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[
@@ -5913,7 +5922,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
         }
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value prepared", "market": "DE"},
-            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -5961,7 +5970,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value prepared", "market": "US"},
-            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[
@@ -6009,7 +6018,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
         # check current state
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "key", "market": "DE"},
-            tasks=(PregelTask(AnyStr(), "prepare"),),
+            tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),),
             next=("prepare",),
             config=uconfig,
             created_at=AnyStr(),
@@ -6029,7 +6038,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
         # get state after first node
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "key prepared", "market": "DE"},
-            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -7236,6 +7245,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "inner",
+                    (PULL, "inner"),
                     state={
                         "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}
                     },
@@ -7271,6 +7281,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "inner",
+                    (PULL, "inner"),
                     state=StateSnapshot(
                         values={
                             "my_key": "hi my value here",
@@ -7280,6 +7291,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                             PregelTask(
                                 AnyStr(),
                                 name="inner_2",
+                                path=(PULL, "inner_2"),
                                 error=None,
                             ),
                         ),
@@ -7350,6 +7362,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "inner",
+                        (PULL, "inner"),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -7383,7 +7396,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
             ),
             StateSnapshot(
                 values={"my_key": "my value"},
-                tasks=(PregelTask(AnyStr(), "outer_1"),),
+                tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
                 next=("outer_1",),
                 config={
                     "configurable": {
@@ -7409,7 +7422,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
             ),
             StateSnapshot(
                 values={},
-                tasks=(PregelTask(AnyStr(), "__start__"),),
+                tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
                 next=("__start__",),
                 config={
                     "configurable": {
@@ -7465,7 +7478,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                         "checkpoint_id": AnyStr(),
                     }
                 },
-                tasks=(PregelTask(id=AnyStr(), name="inner_2"),),
+                tasks=(
+                    PregelTask(id=AnyStr(), name="inner_2", path=(PULL, "inner_2")),
+                ),
             ),
             StateSnapshot(
                 values={"my_key": "hi my value"},
@@ -7494,7 +7509,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                         "checkpoint_id": AnyStr(),
                     }
                 },
-                tasks=(PregelTask(id=AnyStr(), name="inner_1"),),
+                tasks=(
+                    PregelTask(id=AnyStr(), name="inner_1", path=(PULL, "inner_1")),
+                ),
             ),
             StateSnapshot(
                 values={},
@@ -7517,7 +7534,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                 },
                 created_at=AnyStr(),
                 parent_config=None,
-                tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+                tasks=(
+                    PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
+                ),
             ),
         ]
 
@@ -7587,7 +7606,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
             ),
             StateSnapshot(
                 values={"my_key": "hi my value here and there"},
-                tasks=(PregelTask(AnyStr(), "outer_2"),),
+                tasks=(PregelTask(AnyStr(), "outer_2", (PULL, "outer_2")),),
                 next=("outer_2",),
                 config={
                     "configurable": {
@@ -7617,6 +7636,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "inner",
+                        (PULL, "inner"),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -7650,7 +7670,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
             ),
             StateSnapshot(
                 values={"my_key": "my value"},
-                tasks=(PregelTask(AnyStr(), "outer_1"),),
+                tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
                 next=("outer_1",),
                 config={
                     "configurable": {
@@ -7676,7 +7696,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
             ),
             StateSnapshot(
                 values={},
-                tasks=(PregelTask(AnyStr(), "__start__"),),
+                tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
                 next=("__start__",),
                 config={
                     "configurable": {
@@ -7772,6 +7792,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "child",
+                    (PULL, "child"),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -7812,6 +7833,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "child_1",
+                        (PULL, "child_1"),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -7854,6 +7876,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "grandchild_2",
+                    (PULL, "grandchild_2"),
                 ),
             ),
             next=("grandchild_2",),
@@ -7898,18 +7921,21 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "child",
+                    (PULL, "child"),
                     state=StateSnapshot(
                         values={"my_key": "hi my value"},
                         tasks=(
                             PregelTask(
                                 AnyStr(),
                                 "child_1",
+                                (PULL, "child_1"),
                                 state=StateSnapshot(
                                     values={"my_key": "hi my value here"},
                                     tasks=(
                                         PregelTask(
                                             AnyStr(),
                                             "grandchild_2",
+                                            (PULL, "grandchild_2"),
                                         ),
                                     ),
                                     next=("grandchild_2",),
@@ -8113,7 +8139,11 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                             "checkpoint_id": AnyStr(),
                         }
                     },
-                    tasks=(PregelTask(id=AnyStr(), name="parent_2"),),
+                    tasks=(
+                        PregelTask(
+                            id=AnyStr(), name="parent_2", path=(PULL, "parent_2")
+                        ),
+                    ),
                 ),
                 StateSnapshot(
                     values={"my_key": "hi my value"},
@@ -8121,6 +8151,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         PregelTask(
                             AnyStr(),
                             "child",
+                            (PULL, "child"),
                             state={
                                 "configurable": {
                                     "thread_id": "1",
@@ -8176,7 +8207,11 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                             "checkpoint_id": AnyStr(),
                         }
                     },
-                    tasks=(PregelTask(id=AnyStr(), name="parent_1"),),
+                    tasks=(
+                        PregelTask(
+                            id=AnyStr(), name="parent_1", path=(PULL, "parent_1")
+                        ),
+                    ),
                 ),
                 StateSnapshot(
                     values={},
@@ -8196,7 +8231,11 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                     },
                     created_at=AnyStr(),
                     parent_config=None,
-                    tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+                    tasks=(
+                        PregelTask(
+                            id=AnyStr(), name="__start__", path=(PULL, "__start__")
+                        ),
+                    ),
                 ),
             ][0]
         )
@@ -8265,6 +8304,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                     PregelTask(
                         id=AnyStr(),
                         name="child_1",
+                        path=(PULL, "child_1"),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -8295,7 +8335,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                 },
                 created_at=AnyStr(),
                 parent_config=None,
-                tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+                tasks=(
+                    PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
+                ),
             ),
         ]
         # get grandchild graph history
@@ -8379,7 +8421,11 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         "checkpoint_id": AnyStr(),
                     }
                 },
-                tasks=(PregelTask(id=AnyStr(), name="grandchild_2"),),
+                tasks=(
+                    PregelTask(
+                        id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2")
+                    ),
+                ),
             ),
             StateSnapshot(
                 values={"my_key": "hi my value"},
@@ -8417,7 +8463,11 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         "checkpoint_id": AnyStr(),
                     }
                 },
-                tasks=(PregelTask(id=AnyStr(), name="grandchild_1"),),
+                tasks=(
+                    PregelTask(
+                        id=AnyStr(), name="grandchild_1", path=(PULL, "grandchild_1")
+                    ),
+                ),
             ),
             StateSnapshot(
                 values={},
@@ -8449,7 +8499,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                 },
                 created_at=AnyStr(),
                 parent_config=None,
-                tasks=(PregelTask(id=AnyStr(), name="__start__"),),
+                tasks=(
+                    PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
+                ),
             ),
         ]
 
@@ -8524,6 +8576,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
+                    (PUSH, 0),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -8534,6 +8587,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
+                    (PUSH, 1),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -8656,6 +8710,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "generate_joke",
+                        (PUSH, 0),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -8666,6 +8721,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "generate_joke",
+                        (PUSH, 1),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -8693,7 +8749,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
             ),
             StateSnapshot(
                 values={"jokes": []},
-                tasks=(PregelTask(AnyStr(), "__start__"),),
+                tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
                 next=("__start__",),
                 config={
                     "configurable": {
@@ -8879,6 +8935,7 @@ async def test_weather_subgraph(
                 PregelTask(
                     id=AnyStr(),
                     name="weather_graph",
+                    path=(PULL, "weather_graph"),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -8968,6 +9025,7 @@ async def test_weather_subgraph(
                 PregelTask(
                     id=AnyStr(),
                     name="weather_graph",
+                    path=(PULL, "weather_graph"),
                     state=StateSnapshot(
                         values={
                             "messages": [
@@ -9003,7 +9061,13 @@ async def test_weather_subgraph(
                                 "checkpoint_id": AnyStr(),
                             }
                         },
-                        tasks=(PregelTask(id=AnyStr(), name="weather_node"),),
+                        tasks=(
+                            PregelTask(
+                                id=AnyStr(),
+                                name="weather_node",
+                                path=(PULL, "weather_node"),
+                            ),
+                        ),
                     ),
                 ),
             ),
@@ -9045,6 +9109,7 @@ async def test_weather_subgraph(
                 PregelTask(
                     id=AnyStr(),
                     name="weather_graph",
+                    path=(PULL, "weather_graph"),
                     state=StateSnapshot(
                         values={
                             "messages": [

libs/scheduler-kafka/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 LangChain, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

libs/scheduler-kafka/Makefile

@@ -0,0 +1,46 @@
+.PHONY: test test_watch lint format
+
+######################
+# TESTING AND COVERAGE
+######################
+
+start-services:
+	docker compose -f tests/compose.yml up -V --force-recreate --wait --remove-orphans
+
+stop-services:
+	docker compose -f tests/compose.yml down
+
+test:
+	make start-services && poetry run pytest; \
+	EXIT_CODE=$$?; \
+	make stop-services; \
+	exit $$EXIT_CODE
+
+test_watch:
+	make start-services && poetry run ptw .; \
+	EXIT_CODE=$$?; \
+	make stop-services; \
+	exit $$EXIT_CODE
+
+######################
+# LINTING AND FORMATTING
+######################
+
+# Define a variable for Python and notebook files.
+PYTHON_FILES=.
+MYPY_CACHE=.mypy_cache
+lint format: PYTHON_FILES=.
+lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --relative --diff-filter=d main . | grep -E '\.py$$|\.ipynb$$')
+lint_package: PYTHON_FILES=langgraph
+lint_tests: PYTHON_FILES=tests
+lint_tests: MYPY_CACHE=.mypy_cache_test
+
+lint lint_diff lint_package lint_tests:
+	poetry run ruff check .
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
+	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
+	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
+
+format format_diff:
+	poetry run ruff format $(PYTHON_FILES)
+	poetry run ruff check --select I --fix $(PYTHON_FILES)

libs/scheduler-kafka/README.md

@@ -0,0 +1,134 @@
+# LangGraph Scheduler for Kafka
+
+This library implements a distributed scheduler for LangGraph using Kafka as the message broker.
+
+## Architecture
+
+![](./langgraph-distributed.png)
+
+- Combination of Kafka (at least once) with a LangGraph Checkpointer provides exactly once semantics both for orchestrator and executor messages
+- Checkpointer ensures writes for a given task are saved only once, even if the task is re-executed
+- Checkpointer is used to record whether each task in each step has been successfully published to Kafka, to ensure tasks aren't lost, or published more than once
+- Orchestrator and Executor manage commit of offsets manually to ensure tasks are marked as done only after finished processing
+- Orchestrator and Executor pick up from the earliest message not yet consumed when restarted, to ensure no message is lost, and avoid processing messages more than once
+- Orchestrator messages are keyed by thread ID and checkpoint NS, to ensure that no two consumers can process updates for same step of same thread concurrently
+- Executor messages are not keyed, as they can be processed concurrently
+- Orchestrator and Executor execute messages in configurable batches (up to N messages within space of X seconds), and dedupe messages intra-batch where appropriate (this is purely a performance optimization, with no impact on correctness whether applied or not)
+
+## Basic Usage
+
+Launch orchestrator and executor processes:
+
+`orchestrator.py`
+
+```python
+import asyncio
+import logging
+import os
+
+from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator
+from langgraph.scheduler.kafka.types import Topics
+
+from your_lib import graph # graph expected to be a compiled LangGraph graph
+
+logger = logging.getLogger(__name__)
+
+topics = Topics(
+    orchestrator: os.environ['KAFKA_TOPIC_ORCHESTRATOR'],
+    executor: os.environ['KAFKA_TOPIC_EXECUTOR'],
+    error: os.environ['KAFKA_TOPIC_ERROR'],
+)
+
+async def main():
+    async with KafkaOrchestrator(graph, topics) as orch:
+        async for msgs in orch:
+            logger.info('Procesed %d messages', len(msgs))
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(main())
+```
+
+`executor.py`
+
+```python
+import asyncio
+import logging
+import os
+
+from langgraph.scheduler.kafka.executor import KafkaExecutor
+from langgraph.scheduler.kafka.types import Topics
+
+from your_lib import graph # graph expected to be a compiled LangGraph graph
+
+logger = logging.getLogger(__name__)
+
+topics = Topics(
+    orchestrator: os.environ['KAFKA_TOPIC_ORCHESTRATOR'],
+    executor: os.environ['KAFKA_TOPIC_EXECUTOR'],
+    error: os.environ['KAFKA_TOPIC_ERROR'],
+)
+
+async def main():
+    async with KafkaExecutor(graph, topics) as orch:
+        async for msgs in orch:
+            logger.info('Procesed %d messages', len(msgs))
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(main())
+```
+
+```bash
+export KAFKA_TOPIC_ORCHESTRATOR='orchestrator'
+export KAFKA_TOPIC_EXECUTOR='executor'
+export KAFKA_TOPIC_ERROR='error'
+python orchestrator.py &
+python executor.py &
+```
+
+## Configuration
+
+You can pass any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor` to configure the consumer:
+
+- group_id (str): a name for the consumer group. Defaults to 'orchestrator' or 'executor', respectively.
+- batch_max_n (int): Maximum number of messages to include in a single batch. Default: 10.
+- batch_max_ms (int): Maximum time in milliseconds to wait for messages to include in a batch. Default: 1000.
+- retry_policy (langgraph.pregel.types.RetryPolicy): Controls which graph-level errors will be retried when processing messages. A good use for this is to retry database errors thrown by the checkpointer. Defaults to None.
+
+### Connection settings
+
+By default the orchestrator and executor will attempt to connect to a Kafka broker running on `localhost:9092`. You can change connection settings by passing any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor`:
+
+- bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
+  strings) that the consumer should contact to bootstrap initial
+  cluster metadata. This does not have to be the full node list.
+  It just needs to have at least one broker that will respond to
+  Metadata API Request. Default port is 9092. If no servers are
+  specified, will default to localhost:9092.
+- client_id (str): a name for this client. This string is passed in
+  each request to servers and can be used to identify specific
+  server-side log entries that correspond to this client. Also
+  submitted to GroupCoordinator for logging with respect to
+  consumer group administration. Default: 'aiokafka-{ver}'
+- request_timeout_ms (int): Client request timeout in milliseconds.
+  Default: 40000.
+- metadata_max_age_ms (int): The period of time in milliseconds after
+  which we force a refresh of metadata even if we haven't seen
+  any partition leadership changes to proactively discover any
+  new brokers or partitions. Default: 300000
+- retry_backoff_ms (int): Milliseconds to backoff when retrying on
+  errors. Default: 100.
+- api_version (str): specify which kafka API version to use.
+  AIOKafka supports Kafka API versions >=0.9 only.
+  If set to 'auto', will attempt to infer the broker version by
+  probing various APIs. Default: auto
+- security_protocol (str): Protocol used to communicate with brokers.
+  Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
+  Default: PLAINTEXT.
+- ssl_context (ssl.SSLContext): pre-configured SSLContext for wrapping
+  socket connections. For more information see :ref:`ssl_auth`.
+  Default: None.
+- connections_max_idle_ms (int): Close idle connections after the number
+  of milliseconds specified by this config. Specifying `None` will
+  disable idle checks. Default: 540000 (9 minutes).

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -0,0 +1,206 @@
+import asyncio
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from functools import partial
+from typing import Any, Optional, Sequence
+
+import aiokafka
+import orjson
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import Self
+
+import langgraph.scheduler.kafka.serde as serde
+from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP
+from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound
+from langgraph.pregel import Pregel
+from langgraph.pregel.algo import prepare_single_task
+from langgraph.pregel.executor import AsyncBackgroundExecutor, Submit
+from langgraph.pregel.manager import AsyncChannelsManager
+from langgraph.pregel.runner import PregelRunner
+from langgraph.pregel.types import RetryPolicy
+from langgraph.scheduler.kafka.retry import aretry
+from langgraph.scheduler.kafka.types import (
+    ErrorMessage,
+    MessageToExecutor,
+    MessageToOrchestrator,
+    Topics,
+)
+from langgraph.utils.config import patch_configurable
+
+
+class KafkaExecutor(AbstractAsyncContextManager):
+    def __init__(
+        self,
+        graph: Pregel,
+        topics: Topics,
+        *,
+        group_id: str = "executor",
+        batch_max_n: int = 10,
+        batch_max_ms: int = 1000,
+        retry_policy: Optional[RetryPolicy] = None,
+        consumer_kwargs: Optional[dict[str, Any]] = None,
+        producer_kwargs: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.graph = graph
+        self.topics = topics
+        self.stack = AsyncExitStack()
+        self.kwargs = kwargs
+        self.consumer_kwargs = consumer_kwargs or {}
+        self.producer_kwargs = producer_kwargs or {}
+        self.group_id = group_id
+        self.batch_max_n = batch_max_n
+        self.batch_max_ms = batch_max_ms
+        self.retry_policy = retry_policy
+
+    async def __aenter__(self) -> Self:
+        self.consumer = await self.stack.enter_async_context(
+            aiokafka.AIOKafkaConsumer(
+                self.topics.executor,
+                value_deserializer=serde.loads,
+                auto_offset_reset="earliest",
+                group_id=self.group_id,
+                enable_auto_commit=False,
+                **self.kwargs,
+            )
+        )
+        self.producer = await self.stack.enter_async_context(
+            aiokafka.AIOKafkaProducer(
+                key_serializer=serde.dumps,
+                value_serializer=serde.dumps,
+                **self.kwargs,
+            )
+        )
+        self.subgraphs = {
+            k: v async for k, v in self.graph.aget_subgraphs(recurse=True)
+        }
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        return await self.stack.__aexit__(*args)
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> Sequence[MessageToExecutor]:
+        # wait for next batch
+        try:
+            recs = await self.consumer.getmany(
+                timeout_ms=self.batch_max_ms, max_records=self.batch_max_n
+            )
+            msgs: list[MessageToExecutor] = [
+                msg.value for msgs in recs.values() for msg in msgs
+            ]
+        except aiokafka.ConsumerStoppedError:
+            raise StopAsyncIteration from None
+        # process batch
+        await asyncio.gather(*(self.each(msg) for msg in msgs))
+        # commit offsets
+        await self.consumer.commit()
+        # return message
+        return msgs
+
+    async def each(self, msg: MessageToExecutor) -> None:
+        try:
+            await aretry(self.retry_policy, self.attempt, msg)
+        except CheckpointNotLatest:
+            pass
+        except GraphDelegate as exc:
+            for arg in exc.args:
+                await self.producer.send_and_wait(
+                    self.topics.orchestrator,
+                    value=MessageToOrchestrator(
+                        config=arg["config"],
+                        input=orjson.Fragment(
+                            self.graph.checkpointer.serde.dumps(arg["input"])
+                        ),
+                        finally_executor=[msg],
+                    ),
+                    # use thread_id, checkpoint_ns as partition key
+                    key=(
+                        arg["config"]["configurable"]["thread_id"],
+                        arg["config"]["configurable"].get("checkpoint_ns"),
+                    ),
+                )
+        except Exception as exc:
+            await self.producer.send_and_wait(
+                self.topics.error,
+                value=ErrorMessage(
+                    topic=self.topics.executor,
+                    msg=msg,
+                    error=repr(exc),
+                ),
+            )
+
+    async def attempt(self, msg: MessageToExecutor) -> None:
+        # find graph
+        if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
+            # remove task_ids from checkpoint_ns
+            recast_checkpoint_ns = NS_SEP.join(
+                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
+            )
+            # find the subgraph with the matching name
+            if recast_checkpoint_ns in self.subgraphs:
+                graph = self.subgraphs[recast_checkpoint_ns]
+            else:
+                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+        else:
+            graph = self.graph
+        # process message
+        saved = await self.graph.checkpointer.aget_tuple(
+            patch_configurable(msg["config"], {"checkpoint_id": None})
+        )
+        if saved is None:
+            raise RuntimeError("Checkpoint not found")
+        if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
+            raise CheckpointNotLatest()
+        async with AsyncChannelsManager(
+            graph.channels, saved.checkpoint, msg["config"], self.graph.store
+        ) as (channels, managed), AsyncBackgroundExecutor() as submit:
+            if task := await asyncio.to_thread(
+                prepare_single_task,
+                msg["task"]["path"],
+                msg["task"]["id"],
+                checkpoint=saved.checkpoint,
+                processes=graph.nodes,
+                channels=channels,
+                managed=managed,
+                config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}),
+                step=saved.metadata["step"] + 1,
+                for_execution=True,
+                checkpointer=self.graph.checkpointer,
+            ):
+                # execute task, saving writes
+                runner = PregelRunner(
+                    submit=submit,
+                    put_writes=partial(self._put_writes, submit, msg["config"]),
+                )
+                async for _ in runner.atick([task], reraise=False):
+                    pass
+            else:
+                # task was not found
+                await self.graph.checkpointer.put_writes(
+                    msg["config"], [(ERROR, TaskNotFound())]
+                )
+        # notify orchestrator
+        await self.producer.send_and_wait(
+            self.topics.orchestrator,
+            value=MessageToOrchestrator(
+                input=None,
+                config=msg["config"],
+                finally_executor=msg.get("finally_executor"),
+            ),
+            # use thread_id, checkpoint_ns as partition key
+            key=(
+                msg["config"]["configurable"]["thread_id"],
+                msg["config"]["configurable"].get("checkpoint_ns"),
+            ),
+        )
+
+    def _put_writes(
+        self,
+        submit: Submit,
+        config: RunnableConfig,
+        task_id: str,
+        writes: list[tuple[str, Any]],
+    ) -> None:
+        return submit(self.graph.checkpointer.aput_writes, config, writes, task_id)

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -0,0 +1,205 @@
+import asyncio
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from typing import Any, Optional
+
+import aiokafka
+from langchain_core.runnables import ensure_config
+from typing_extensions import Self
+
+import langgraph.scheduler.kafka.serde as serde
+from langgraph.constants import (
+    CONFIG_KEY_DEDUPE_TASKS,
+    CONFIG_KEY_ENSURE_LATEST,
+    INTERRUPT,
+    NS_END,
+    NS_SEP,
+    SCHEDULED,
+)
+from langgraph.errors import CheckpointNotLatest, GraphInterrupt
+from langgraph.pregel import Pregel
+from langgraph.pregel.loop import AsyncPregelLoop
+from langgraph.pregel.types import RetryPolicy
+from langgraph.scheduler.kafka.retry import aretry
+from langgraph.scheduler.kafka.types import (
+    ErrorMessage,
+    ExecutorTask,
+    MessageToExecutor,
+    MessageToOrchestrator,
+    Topics,
+)
+from langgraph.utils.config import patch_configurable
+
+
+class KafkaOrchestrator(AbstractAsyncContextManager):
+    def __init__(
+        self,
+        graph: Pregel,
+        topics: Topics,
+        group_id: str = "orchestrator",
+        batch_max_n: int = 10,
+        batch_max_ms: int = 1000,
+        retry_policy: Optional[RetryPolicy] = None,
+        consumer_kwargs: Optional[dict[str, Any]] = None,
+        producer_kwargs: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.graph = graph
+        self.topics = topics
+        self.stack = AsyncExitStack()
+        self.kwargs = kwargs
+        self.consumer_kwargs = consumer_kwargs or {}
+        self.producer_kwargs = producer_kwargs or {}
+        self.group_id = group_id
+        self.batch_max_n = batch_max_n
+        self.batch_max_ms = batch_max_ms
+        self.retry_policy = retry_policy
+
+    async def __aenter__(self) -> Self:
+        self.consumer = await self.stack.enter_async_context(
+            aiokafka.AIOKafkaConsumer(
+                self.topics.orchestrator,
+                auto_offset_reset="earliest",
+                group_id=self.group_id,
+                enable_auto_commit=False,
+                **self.kwargs,
+                **self.consumer_kwargs,
+            )
+        )
+        self.producer = await self.stack.enter_async_context(
+            aiokafka.AIOKafkaProducer(
+                value_serializer=serde.dumps,
+                **self.kwargs,
+                **self.producer_kwargs,
+            )
+        )
+        self.subgraphs = {
+            k: v async for k, v in self.graph.aget_subgraphs(recurse=True)
+        }
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        return await self.stack.__aexit__(*args)
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> list[MessageToOrchestrator]:
+        # wait for next batch
+        try:
+            recs = await self.consumer.getmany(
+                timeout_ms=self.batch_max_ms, max_records=self.batch_max_n
+            )
+            # dedupe messages, eg. if multiple nodes finish around same time
+            uniq = set(msg.value for msgs in recs.values() for msg in msgs)
+            msgs: list[MessageToOrchestrator] = [serde.loads(msg) for msg in uniq]
+        except aiokafka.ConsumerStoppedError:
+            raise StopAsyncIteration from None
+        # process batch
+        await asyncio.gather(*(self.each(msg) for msg in msgs))
+        # commit offsets
+        await self.consumer.commit()
+        # return message
+        return msgs
+
+    async def each(self, msg: MessageToOrchestrator) -> None:
+        try:
+            await aretry(self.retry_policy, self.attempt, msg)
+        except CheckpointNotLatest:
+            pass
+        except GraphInterrupt:
+            pass
+        except Exception as exc:
+            await self.producer.send_and_wait(
+                self.topics.error,
+                value=ErrorMessage(
+                    topic=self.topics.orchestrator,
+                    msg=msg,
+                    error=repr(exc),
+                ),
+            )
+
+    async def attempt(self, msg: MessageToOrchestrator) -> None:
+        # find graph
+        if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
+            # remove task_ids from checkpoint_ns
+            recast_checkpoint_ns = NS_SEP.join(
+                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
+            )
+            # find the subgraph with the matching name
+            if recast_checkpoint_ns in self.subgraphs:
+                graph = self.subgraphs[recast_checkpoint_ns]
+            else:
+                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+        else:
+            graph = self.graph
+        # process message
+        async with AsyncPregelLoop(
+            msg["input"],
+            config=ensure_config(msg["config"]),
+            stream=None,
+            store=self.graph.store,
+            checkpointer=self.graph.checkpointer,
+            nodes=graph.nodes,
+            specs=graph.channels,
+            output_keys=graph.output_channels,
+            stream_keys=graph.stream_channels,
+        ) as loop:
+            if loop.tick(
+                input_keys=graph.input_channels,
+                interrupt_after=graph.interrupt_after_nodes,
+                interrupt_before=graph.interrupt_before_nodes,
+            ):
+                # wait for checkpoint to be saved
+                if hasattr(loop, "_put_checkpoint_fut"):
+                    await loop._put_checkpoint_fut
+                # schedule any new tasks
+                if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
+                    # send messages to executor
+                    futures: list[asyncio.Future] = await asyncio.gather(
+                        *(
+                            self.producer.send(
+                                self.topics.executor,
+                                value=MessageToExecutor(
+                                    config=patch_configurable(
+                                        loop.config,
+                                        {
+                                            **loop.checkpoint_config["configurable"],
+                                            CONFIG_KEY_DEDUPE_TASKS: True,
+                                            CONFIG_KEY_ENSURE_LATEST: True,
+                                        },
+                                    ),
+                                    task=ExecutorTask(id=task.id, path=task.path),
+                                    finally_executor=msg.get("finally_executor"),
+                                ),
+                            )
+                            for task in new_tasks
+                        )
+                    )
+                    # wait for messages to be sent
+                    await asyncio.gather(*futures)
+                    # mark as scheduled
+                    for task in new_tasks:
+                        loop.put_writes(
+                            task.id,
+                            [
+                                (
+                                    SCHEDULED,
+                                    max(
+                                        loop.checkpoint["versions_seen"]
+                                        .get(INTERRUPT, {})
+                                        .values(),
+                                        default=None,
+                                    ),
+                                )
+                            ],
+                        )
+            elif loop.status == "done" and msg.get("finally_executor"):
+                # schedule any finally_executor tasks
+                futs = await asyncio.gather(
+                    *(
+                        self.producer.send(self.topics.executor, value=m)
+                        for m in msg["finally_executor"]
+                    )
+                )
+                # wait for messages to be sent
+                await asyncio.gather(*futs)

libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py

@@ -0,0 +1,54 @@
+import asyncio
+import logging
+import random
+from typing import Awaitable, Callable, Optional
+
+from typing_extensions import ParamSpec
+
+from langgraph.pregel.types import RetryPolicy
+
+logger = logging.getLogger(__name__)
+P = ParamSpec("P")
+
+
+async def aretry(
+    retry_policy: Optional[RetryPolicy],
+    func: Callable[P, Awaitable[None]],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> None:
+    """Run a task asynchronously with retries."""
+    interval = retry_policy.initial_interval if retry_policy else 0
+    attempts = 0
+    while True:
+        try:
+            await func(*args, **kwargs)
+            # if successful, end
+            break
+        except Exception as exc:
+            if retry_policy is None:
+                raise
+            # increment attempts
+            attempts += 1
+            # check if we should retry
+            if callable(retry_policy.retry_on):
+                if not retry_policy.retry_on(exc):
+                    raise
+            elif not isinstance(exc, retry_policy.retry_on):
+                raise
+            # check if we should give up
+            if attempts >= retry_policy.max_attempts:
+                raise
+            # sleep before retrying
+            interval = min(
+                retry_policy.max_interval,
+                interval * retry_policy.backoff_factor,
+            )
+            await asyncio.sleep(
+                interval + random.uniform(0, 1) if retry_policy.jitter else interval
+            )
+            # log the retry
+            logger.info(
+                f"Retrying function {func} with {args} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}",
+                exc_info=exc,
+            )

libs/scheduler-kafka/langgraph/scheduler/kafka/serde.py

@@ -0,0 +1,16 @@
+from typing import Any
+
+import orjson
+
+
+def loads(v: bytes) -> Any:
+    return orjson.loads(v)
+
+
+def dumps(v: Any) -> bytes:
+    return orjson.dumps(v, default=_default)
+
+
+def _default(v: Any) -> Any:
+    # things we don't know how to serialize (eg. functions) ignore
+    return None

libs/scheduler-kafka/langgraph/scheduler/kafka/types.py

@@ -0,0 +1,32 @@
+from typing import Any, NamedTuple, Optional, Sequence, TypedDict, Union
+
+from langchain_core.runnables import RunnableConfig
+
+
+class Topics(NamedTuple):
+    orchestrator: str
+    executor: str
+    error: str
+
+
+class MessageToOrchestrator(TypedDict):
+    input: Optional[dict[str, Any]]
+    config: RunnableConfig
+    finally_executor: Optional[Sequence["MessageToExecutor"]]
+
+
+class ExecutorTask(TypedDict):
+    id: str
+    path: tuple[str, ...]
+
+
+class MessageToExecutor(TypedDict):
+    config: RunnableConfig
+    task: ExecutorTask
+    finally_executor: Optional[Sequence["MessageToExecutor"]]
+
+
+class ErrorMessage(TypedDict):
+    topic: str
+    error: str
+    msg: Union[MessageToExecutor, MessageToOrchestrator]

libs/scheduler-kafka/pyproject.toml

@@ -0,0 +1,58 @@
+[tool.poetry]
+name = "langgraph-scheduler-kafka"
+version = "1.0.0"
+description = "Library with Kafka-based work scheduler."
+authors = []
+license = "MIT"
+readme = "README.md"
+repository = "https://www.github.com/langchain-ai/langgraph"
+packages = [{ include = "langgraph" }]
+
+[tool.poetry.dependencies]
+python = "^3.9.0,<4.0"
+orjson = "^3.10.7"
+crc32c = "^2.7.post1"
+aiokafka = "^0.11.0"
+langgraph = "^0.2.19"
+
+[tool.poetry.group.dev.dependencies]
+ruff = "^0.6.2"
+codespell = "^2.2.0"
+pytest = "^7.2.1"
+pytest-mock = "^3.11.1"
+pytest-watcher = "^0.4.1"
+mypy = "^1.10.0"
+langgraph = {path = "../langgraph", develop = true}
+langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true}
+langgraph-checkpoint = {path = "../checkpoint", develop = true}
+kafka-python-ng = "^2.2.2"
+
+[tool.pytest.ini_options]
+# --strict-markers will raise errors on unknown marks.
+# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
+#
+# https://docs.pytest.org/en/7.1.x/reference/reference.html
+# --strict-config       any warnings encountered while parsing the `pytest`
+#                       section of the configuration file raise errors.
+addopts = "--strict-markers --strict-config --durations=5 -vv"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.ruff]
+lint.select = [
+  "E",  # pycodestyle
+  "F",  # Pyflakes
+  "UP", # pyupgrade
+  "B",  # flake8-bugbear
+  "I",  # isort
+]
+lint.ignore = ["E501", "B008", "UP007", "UP006"]
+
+[tool.pytest-watcher]
+now = true
+delay = 0.1
+runner_args = ["--ff", "-v", "--tb", "short", "-s"]
+patterns = ["*.py"]

libs/scheduler-kafka/tests/any.py

@@ -0,0 +1,37 @@
+import re
+from typing import Union
+
+
+class AnyStr(str):
+    def __init__(self, prefix: Union[str, re.Pattern] = "") -> None:
+        super().__init__()
+        self.prefix = prefix
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, str) and (
+            other.startswith(self.prefix)
+            if isinstance(self.prefix, str)
+            else self.prefix.match(other)
+        )
+
+    def __hash__(self) -> int:
+        return hash((str(self), self.prefix))
+
+
+class AnyDict(dict):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def __eq__(self, other: object) -> bool:
+        if not self and isinstance(other, dict):
+            return True
+        if not isinstance(other, dict) or len(self) != len(other):
+            return False
+        for k, v in self.items():
+            if kk := next((kk for kk in other if kk == k), None):
+                if v == other[kk]:
+                    continue
+                else:
+                    return False
+        else:
+            return True

libs/scheduler-kafka/tests/compose.yml

@@ -0,0 +1,21 @@
+name: scheduler-kafka-tests
+services:
+  broker:
+    image: apache/kafka:latest
+    ports:
+      - "9092:9092"
+  postgres:
+    image: postgres:16
+    ports:
+      - "5443:5432"
+    environment:
+      POSTGRES_DB: postgres
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+    healthcheck:
+      test: pg_isready -U postgres
+      start_period: 10s
+      timeout: 1s
+      retries: 5
+      interval: 60s
+      start_interval: 1s

libs/scheduler-kafka/tests/conftest.py

@@ -0,0 +1,62 @@
+from typing import AsyncIterator, Iterator
+from uuid import uuid4
+
+import kafka.admin
+import pytest
+from psycopg import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
+
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.scheduler.kafka.types import Topics
+
+DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5443/"
+
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"
+
+
+@pytest.fixture
+def topics() -> Iterator[Topics]:
+    o = f"test_o_{uuid4().hex[:16]}"
+    e = f"test_e_{uuid4().hex[:16]}"
+    z = f"test_z_{uuid4().hex[:16]}"
+    admin = kafka.admin.KafkaAdminClient()
+    # create topics
+    admin.create_topics(
+        [
+            kafka.admin.NewTopic(name=o, num_partitions=1, replication_factor=1),
+            kafka.admin.NewTopic(name=e, num_partitions=1, replication_factor=1),
+            kafka.admin.NewTopic(name=z, num_partitions=1, replication_factor=1),
+        ]
+    )
+    # yield topics
+    yield Topics(orchestrator=o, executor=e, error=z)
+    # delete topics
+    admin.delete_topics([o, e, z])
+    admin.close()
+
+
+@pytest.fixture
+async def checkpointer() -> AsyncIterator[AsyncPostgresSaver]:
+    database = f"test_{uuid4().hex[:16]}"
+    # create unique db
+    async with await AsyncConnection.connect(
+        DEFAULT_POSTGRES_URI, autocommit=True
+    ) as conn:
+        await conn.execute(f"CREATE DATABASE {database}")
+    try:
+        # yield checkpointer
+        async with AsyncConnectionPool(
+            DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
+        ) as pool:
+            checkpointer = AsyncPostgresSaver(pool)
+            await checkpointer.setup()
+            yield checkpointer
+    finally:
+        # drop unique db
+        async with await AsyncConnection.connect(
+            DEFAULT_POSTGRES_URI, autocommit=True
+        ) as conn:
+            await conn.execute(f"DROP DATABASE {database}")

libs/scheduler-kafka/tests/drain.py

@@ -0,0 +1,79 @@
+import asyncio
+from typing import Optional, TypeVar
+
+import anyio
+from aiokafka import AIOKafkaConsumer
+from typing_extensions import ParamSpec
+
+from langgraph.pregel import Pregel
+from langgraph.scheduler.kafka.executor import KafkaExecutor
+from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+
+C = ParamSpec("C")
+R = TypeVar("R")
+
+
+async def drain_topics(
+    topics: Topics, graph: Pregel, *, debug: bool = False
+) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]:
+    scope: Optional[anyio.CancelScope] = None
+    orch_msgs = []
+    exec_msgs = []
+    errors = []
+
+    def done() -> bool:
+        return (
+            len(orch_msgs) > 0
+            and len(exec_msgs) > 0
+            and not orch_msgs[-1]
+            and not exec_msgs[-1]
+        )
+
+    async def orchestrator() -> None:
+        async with KafkaOrchestrator(graph, topics) as orch:
+            async for msgs in orch:
+                orch_msgs.append(msgs)
+                if debug:
+                    print("\n---\norch", len(msgs), msgs)
+                if done():
+                    scope.cancel()
+
+    async def executor() -> None:
+        async with KafkaExecutor(graph, topics) as exec:
+            async for msgs in exec:
+                exec_msgs.append(msgs)
+                if debug:
+                    print("\n---\nexec", len(msgs), msgs)
+                if done():
+                    scope.cancel()
+
+    async def error_consumer() -> None:
+        async with AIOKafkaConsumer(topics.error) as consumer:
+            async for msg in consumer:
+                errors.append(msg)
+                if scope:
+                    scope.cancel()
+
+    # start error consumer
+    error_task = asyncio.create_task(error_consumer(), name="error_consumer")
+
+    # run the orchestrator and executor until break_when
+    async with anyio.create_task_group() as tg:
+        tg.cancel_scope.deadline = anyio.current_time() + 20
+        scope = tg.cancel_scope
+        tg.start_soon(orchestrator, name="orchestrator")
+        tg.start_soon(executor, name="executor")
+
+    # cancel error consumer
+    error_task.cancel()
+
+    try:
+        await error_task
+    except asyncio.CancelledError:
+        pass
+
+    # check no errors
+    assert not errors, errors
+
+    return [m for mm in orch_msgs for m in mm], [m for mm in exec_msgs for m in mm]

libs/scheduler-kafka/tests/test_fanout.py

@@ -0,0 +1,274 @@
+import asyncio
+import operator
+from typing import (
+    Annotated,
+    Sequence,
+    TypedDict,
+    Union,
+)
+
+import pytest
+from aiokafka import AIOKafkaProducer
+
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph.state import StateGraph
+from langgraph.pregel import Pregel
+from langgraph.scheduler.kafka import serde
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from tests.any import AnyDict
+from tests.drain import drain_topics
+
+pytestmark = pytest.mark.anyio
+
+
+def mk_fanout_graph(
+    checkpointer: BaseCheckpointSaver, interrupt_before: Sequence[str] = ()
+) -> Pregel:
+    # copied from test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    async def rewrite_query(data: State) -> State:
+        return {"query": f'query: {data["query"]}'}
+
+    async def retriever_picker(data: State) -> list[str]:
+        return ["analyzer_one", "retriever_two"]
+
+    async def analyzer_one(data: State) -> State:
+        return {"query": f'analyzed: {data["query"]}'}
+
+    async def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    async def retriever_two(data: State) -> State:
+        await asyncio.sleep(0.1)
+        return {"docs": ["doc3", "doc4"]}
+
+    async def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    async def decider(data: State) -> None:
+        return None
+
+    def decider_cond(data: State) -> str:
+        if data["query"].count("analyzed") > 1:
+            return "qa"
+        else:
+            return "rewrite_query"
+
+    builder = StateGraph(State)
+
+    builder.add_node("rewrite_query", rewrite_query)
+    builder.add_node("analyzer_one", analyzer_one)
+    builder.add_node("retriever_one", retriever_one)
+    builder.add_node("retriever_two", retriever_two)
+    builder.add_node("decider", decider)
+    builder.add_node("qa", qa)
+
+    builder.set_entry_point("rewrite_query")
+    builder.add_conditional_edges("rewrite_query", retriever_picker)
+    builder.add_edge("analyzer_one", "retriever_one")
+    builder.add_edge(["retriever_one", "retriever_two"], "decider")
+    builder.add_conditional_edges("decider", decider_cond)
+    builder.set_finish_point("qa")
+
+    return builder.compile(checkpointer, interrupt_before=interrupt_before)
+
+
+async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None:
+    input = {"query": "what is weather in sf"}
+    config = {"configurable": {"thread_id": "1"}}
+    graph = mk_fanout_graph(checkpointer)
+
+    # start a new run
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=input, config=config),
+        )
+
+    # drain topics
+    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+
+    # check state
+    state = await graph.aget_state(config)
+    assert state.next == ()
+    assert (
+        state.values
+        == await graph.ainvoke(input, {"configurable": {"thread_id": "2"}})
+        == {
+            "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+            "query": "analyzed: query: analyzed: query: what is weather in sf",
+            "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+        }
+    )
+
+    # check history
+    history = [c async for c in graph.aget_state_history(config)]
+    assert len(history) == 11
+
+    # check messages
+    assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "input": None,
+            "finally_executor": None,
+        }
+        for c in reversed(history)
+        for _ in c.tasks
+    ]
+    assert exec_msgs == [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "task": {
+                "id": t.id,
+                "path": list(t.path),
+            },
+            "finally_executor": None,
+        }
+        for c in reversed(history)
+        for t in c.tasks
+    ]
+
+
+async def test_fanout_graph_w_interrupt(
+    topics: Topics, checkpointer: BaseCheckpointSaver
+) -> None:
+    input = {"query": "what is weather in sf"}
+    config = {"configurable": {"thread_id": "1"}}
+    graph = mk_fanout_graph(checkpointer, interrupt_before=["qa"])
+
+    # start a new run
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=input, config=config),
+        )
+
+    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+
+    # check interrupted state
+    state = await graph.aget_state(config)
+    assert state.next == ("qa",)
+    assert (
+        state.values
+        == await graph.ainvoke(input, {"configurable": {"thread_id": "2"}})
+        == {
+            "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+            "query": "analyzed: query: analyzed: query: what is weather in sf",
+        }
+    )
+
+    # check history
+    history = [c async for c in graph.aget_state_history(config)]
+    assert len(history) == 10
+
+    # check messages
+    assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "input": None,
+            "finally_executor": None,
+        }
+        for c in reversed(history[1:])  # the last one wasn't executed
+        # orchestrator messages appear only after tasks for that checkpoint
+        # finish executing, ie. after executor sends message to resume checkpoint
+        for _ in c.tasks
+    ]
+    assert exec_msgs == [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "task": {
+                "id": t.id,
+                "path": list(t.path),
+            },
+            "finally_executor": None,
+        }
+        for c in reversed(history[1:])  # the last one wasn't executed
+        for t in c.tasks
+    ]
+
+    # resume the thread
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=None, config=config),
+        )
+
+    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+
+    # check final state
+    state = await graph.aget_state(config)
+    assert state.next == ()
+    assert (
+        state.values
+        == await graph.ainvoke(None, {"configurable": {"thread_id": "2"}})
+        == {
+            "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+            "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+            "query": "analyzed: query: analyzed: query: what is weather in sf",
+        }
+    )

libs/scheduler-kafka/tests/test_subgraph.py

@@ -0,0 +1,700 @@
+from typing import Literal, cast
+
+import pytest
+from aiokafka import AIOKafkaProducer
+from langchain_core.language_models.fake_chat_models import (
+    FakeMessagesListChatModel,
+)
+from langchain_core.messages import AIMessage, HumanMessage, ToolCall
+from langchain_core.tools import tool
+
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import END, START
+from langgraph.graph import MessagesState
+from langgraph.graph.state import StateGraph
+from langgraph.pregel import Pregel
+from langgraph.scheduler.kafka import serde
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from tests.any import AnyDict, AnyStr
+from tests.drain import drain_topics
+
+pytestmark = pytest.mark.anyio
+
+
+def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel:
+    # copied from test_weather_subgraph
+
+    # setup subgraph
+
+    @tool
+    def get_weather(city: str):
+        """Get the weather for a specific city"""
+        return f"I'ts sunny in {city}!"
+
+    weather_model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="",
+                tool_calls=[
+                    ToolCall(
+                        id="tool_call123",
+                        name="get_weather",
+                        args={"city": "San Francisco"},
+                    )
+                ],
+            )
+        ]
+    )
+
+    class SubGraphState(MessagesState):
+        city: str
+
+    def model_node(state: SubGraphState):
+        result = weather_model.invoke(state["messages"])
+        return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]}
+
+    def weather_node(state: SubGraphState):
+        result = get_weather.invoke({"city": state["city"]})
+        return {"messages": [{"role": "assistant", "content": result}]}
+
+    subgraph = StateGraph(SubGraphState)
+    subgraph.add_node(model_node)
+    subgraph.add_node(weather_node)
+    subgraph.add_edge(START, "model_node")
+    subgraph.add_edge("model_node", "weather_node")
+    subgraph.add_edge("weather_node", END)
+    subgraph = subgraph.compile(interrupt_before=["weather_node"])
+
+    # setup main graph
+
+    class RouterState(MessagesState):
+        route: Literal["weather", "other"]
+
+    router_model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="",
+                tool_calls=[
+                    ToolCall(
+                        id="tool_call123",
+                        name="router",
+                        args={"dest": "weather"},
+                    )
+                ],
+            )
+        ]
+    )
+
+    def router_node(state: RouterState):
+        system_message = "Classify the incoming query as either about weather or not."
+        messages = [{"role": "system", "content": system_message}] + state["messages"]
+        route = router_model.invoke(messages)
+        return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]}
+
+    def normal_llm_node(state: RouterState):
+        return {"messages": [AIMessage("Hello!")]}
+
+    def route_after_prediction(state: RouterState):
+        if state["route"] == "weather":
+            return "weather_graph"
+        else:
+            return "normal_llm_node"
+
+    async def weather_graph(state: RouterState):
+        return await subgraph.ainvoke(state)
+
+    graph = StateGraph(RouterState)
+    graph.add_node(router_node)
+    graph.add_node(normal_llm_node)
+    graph.add_node("weather_graph", weather_graph)
+    graph.add_edge(START, "router_node")
+    graph.add_conditional_edges("router_node", route_after_prediction)
+    graph.add_edge("normal_llm_node", END)
+    graph.add_edge("weather_graph", END)
+
+    return graph.compile(checkpointer=checkpointer)
+
+
+async def test_subgraph_w_interrupt(
+    topics: Topics, checkpointer: BaseCheckpointSaver
+) -> None:
+    input = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
+    config = {"configurable": {"thread_id": "1"}}
+    graph = mk_weather_graph(checkpointer)
+
+    # start a new run
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=input, config=config),
+        )
+
+    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+
+    # check interrupted state
+    state = await graph.aget_state(config)
+    assert state.next == ("weather_graph",)
+    assert state.values == {
+        "messages": [HumanMessage(id=AnyStr(), content="what's the weather in sf")],
+        "route": "weather",
+    }
+
+    # check outer history
+    history = [c async for c in graph.aget_state_history(config)]
+    assert len(history) == 3
+
+    # check child history
+    child_history = [
+        c async for c in graph.aget_state_history(history[0].tasks[0].state)
+    ]
+    assert len(child_history) == 3
+
+    # check messages
+    assert (
+        orch_msgs
+        == (
+            # initial message to outer graph
+            [MessageToOrchestrator(input=input, config=config)]
+            # outer graph messages, until interrupted
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": False,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_ns": "",
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "input": None,
+                    "finally_executor": None,
+                }
+                for c in reversed(history[1:])  # the last one wasn't executed
+                # orchestrator messages appear only after tasks for that checkpoint
+                # finish executing, ie. after executor sends message to resume checkpoint
+                for _ in c.tasks
+            ]
+            # initial message to child graph
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_checkpointer": None,
+                            "__pregel_delegate": False,
+                            "__pregel_read": None,
+                            "__pregel_send": None,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": False,
+                            "__pregel_task_id": history[0].tasks[0].id,
+                            "checkpoint_id": None,
+                            "checkpoint_map": {
+                                "": history[0].config["configurable"]["checkpoint_id"]
+                            },
+                            "checkpoint_ns": history[0]
+                            .tasks[0]
+                            .state["configurable"]["checkpoint_ns"],
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "input": {
+                        "messages": [
+                            {
+                                "id": [
+                                    "langchain",
+                                    "schema",
+                                    "messages",
+                                    "HumanMessage",
+                                ],
+                                "kwargs": {
+                                    "content": "what's the weather in sf",
+                                    "id": AnyStr(),
+                                    "type": "human",
+                                },
+                                "lc": 1,
+                                "type": "constructor",
+                            }
+                        ],
+                        "route": "weather",
+                    },
+                    "finally_executor": [
+                        {
+                            "config": {
+                                "callbacks": None,
+                                "configurable": {
+                                    "__pregel_dedupe_tasks": True,
+                                    "__pregel_ensure_latest": True,
+                                    "__pregel_resuming": False,
+                                    "checkpoint_id": history[0].config["configurable"][
+                                        "checkpoint_id"
+                                    ],
+                                    "checkpoint_ns": "",
+                                    "thread_id": "1",
+                                },
+                                "metadata": AnyDict(),
+                                "recursion_limit": 25,
+                                "tags": [],
+                            },
+                            "finally_executor": None,
+                            "task": {
+                                "id": history[0].tasks[0].id,
+                                "path": list(history[0].tasks[0].path),
+                            },
+                        }
+                    ],
+                }
+            ]
+            # child graph messages, until interrupted
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_checkpointer": None,
+                            "__pregel_delegate": False,
+                            "__pregel_read": None,
+                            "__pregel_send": None,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": False,
+                            "__pregel_task_id": history[0].tasks[0].id,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_map": {
+                                "": history[0].config["configurable"]["checkpoint_id"]
+                            },
+                            "checkpoint_ns": history[0]
+                            .tasks[0]
+                            .state["configurable"]["checkpoint_ns"],
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "input": None,
+                    "finally_executor": [
+                        {
+                            "config": {
+                                "callbacks": None,
+                                "configurable": {
+                                    "__pregel_dedupe_tasks": True,
+                                    "__pregel_ensure_latest": True,
+                                    "__pregel_resuming": False,
+                                    "checkpoint_id": history[0].config["configurable"][
+                                        "checkpoint_id"
+                                    ],
+                                    "checkpoint_ns": "",
+                                    "thread_id": "1",
+                                },
+                                "metadata": AnyDict(),
+                                "recursion_limit": 25,
+                                "tags": [],
+                            },
+                            "finally_executor": None,
+                            "task": {
+                                "id": history[0].tasks[0].id,
+                                "path": list(history[0].tasks[0].path),
+                            },
+                        }
+                    ],
+                }
+                for c in reversed(child_history[1:])  # the last one wasn't executed
+                # orchestrator messages appear only after tasks for that checkpoint
+                # finish executing, ie. after executor sends message to resume checkpoint
+                for _ in c.tasks
+            ]
+        )
+    )
+    assert (
+        exec_msgs
+        == (
+            # outer graph tasks
+            [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": False,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_ns": "",
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "task": {
+                        "id": t.id,
+                        "path": list(t.path),
+                    },
+                    "finally_executor": None,
+                }
+                for c in reversed(history)
+                for t in c.tasks
+            ]
+            # child graph tasks
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_checkpointer": None,
+                            "__pregel_delegate": False,
+                            "__pregel_read": None,
+                            "__pregel_send": None,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": False,
+                            "__pregel_task_id": history[0].tasks[0].id,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_map": {
+                                "": history[0].config["configurable"]["checkpoint_id"]
+                            },
+                            "checkpoint_ns": history[0]
+                            .tasks[0]
+                            .state["configurable"]["checkpoint_ns"],
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "task": {
+                        "id": t.id,
+                        "path": list(t.path),
+                    },
+                    "finally_executor": [
+                        {
+                            "config": {
+                                "callbacks": None,
+                                "configurable": {
+                                    "__pregel_dedupe_tasks": True,
+                                    "__pregel_ensure_latest": True,
+                                    "__pregel_resuming": False,
+                                    "checkpoint_id": history[0].config["configurable"][
+                                        "checkpoint_id"
+                                    ],
+                                    "checkpoint_ns": "",
+                                    "thread_id": "1",
+                                },
+                                "metadata": AnyDict(),
+                                "recursion_limit": 25,
+                                "tags": [],
+                            },
+                            "finally_executor": None,
+                            "task": {
+                                "id": history[0].tasks[0].id,
+                                "path": list(history[0].tasks[0].path),
+                            },
+                        }
+                    ],
+                }
+                for c in reversed(child_history[1:])  # the last one wasn't executed
+                for t in c.tasks
+            ]
+        )
+    )
+
+    # resume the thread
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=None, config=config),
+        )
+
+    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+
+    # check final state
+    state = await graph.aget_state(config)
+    assert state.next == ()
+    assert state.values == {
+        "messages": [
+            HumanMessage(id=AnyStr(), content="what's the weather in sf"),
+            AIMessage(content="I'ts sunny in San Francisco!", id=AnyStr()),
+        ],
+        "route": "weather",
+    }
+
+    # check outer history
+    history = [c async for c in graph.aget_state_history(config)]
+    assert len(history) == 4
+
+    # check child history
+    # accessing second to last checkpoint, since that's the one w/ subgraph task
+    child_history = [
+        c async for c in graph.aget_state_history(history[1].tasks[0].state)
+    ]
+    assert len(child_history) == 4
+
+    # check messages
+    assert (
+        orch_msgs
+        == (
+            # initial message to outer graph
+            [MessageToOrchestrator(input=None, config=config)]
+            # initial message to child graph
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_checkpointer": None,
+                            "__pregel_delegate": False,
+                            "__pregel_read": None,
+                            "__pregel_send": None,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": True,
+                            "__pregel_task_id": history[1].tasks[0].id,
+                            "checkpoint_id": None,
+                            "checkpoint_map": {
+                                "": history[1].config["configurable"]["checkpoint_id"]
+                            },
+                            "checkpoint_ns": history[1]
+                            .tasks[0]
+                            .state["configurable"]["checkpoint_ns"],
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "input": None,
+                    "finally_executor": [
+                        {
+                            "config": {
+                                "callbacks": None,
+                                "configurable": {
+                                    "__pregel_dedupe_tasks": True,
+                                    "__pregel_ensure_latest": True,
+                                    "__pregel_resuming": True,
+                                    "checkpoint_id": history[1].config["configurable"][
+                                        "checkpoint_id"
+                                    ],
+                                    "checkpoint_ns": "",
+                                    "thread_id": "1",
+                                },
+                                "metadata": AnyDict(),
+                                "recursion_limit": 25,
+                                "tags": [],
+                            },
+                            "finally_executor": None,
+                            "task": {
+                                "id": history[1].tasks[0].id,
+                                "path": list(history[1].tasks[0].path),
+                            },
+                        }
+                    ],
+                }
+            ]
+            # child graph messages, from previous last checkpoint onwards
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_checkpointer": None,
+                            "__pregel_delegate": False,
+                            "__pregel_read": None,
+                            "__pregel_send": None,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": True,
+                            "__pregel_task_id": history[1].tasks[0].id,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_map": {
+                                "": history[1].config["configurable"]["checkpoint_id"]
+                            },
+                            "checkpoint_ns": history[1]
+                            .tasks[0]
+                            .state["configurable"]["checkpoint_ns"],
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "input": None,
+                    "finally_executor": [
+                        {
+                            "config": {
+                                "callbacks": None,
+                                "configurable": {
+                                    "__pregel_dedupe_tasks": True,
+                                    "__pregel_ensure_latest": True,
+                                    "__pregel_resuming": True,
+                                    "checkpoint_id": history[1].config["configurable"][
+                                        "checkpoint_id"
+                                    ],
+                                    "checkpoint_ns": "",
+                                    "thread_id": "1",
+                                },
+                                "metadata": AnyDict(),
+                                "recursion_limit": 25,
+                                "tags": [],
+                            },
+                            "finally_executor": None,
+                            "task": {
+                                "id": history[1].tasks[0].id,
+                                "path": list(history[1].tasks[0].path),
+                            },
+                        }
+                    ],
+                }
+                for c in reversed(child_history[:2])
+                for _ in c.tasks
+            ]
+            # outer graph messages, from previous last checkpoint onwards
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": True,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_ns": "",
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "input": None,
+                    "finally_executor": None,
+                }
+                for c in reversed(history[:2])
+                for _ in c.tasks
+            ]
+        )
+    )
+    assert (
+        exec_msgs
+        == (
+            # outer graph tasks
+            [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": True,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_ns": "",
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "task": {
+                        "id": t.id,
+                        "path": list(t.path),
+                    },
+                    "finally_executor": None,
+                }
+                for c in reversed(history[:2])
+                for t in c.tasks
+            ]
+            # child graph tasks
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_checkpointer": None,
+                            "__pregel_delegate": False,
+                            "__pregel_read": None,
+                            "__pregel_send": None,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_resuming": True,
+                            "__pregel_task_id": history[1].tasks[0].id,
+                            "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                            "checkpoint_map": {
+                                "": history[1].config["configurable"]["checkpoint_id"]
+                            },
+                            "checkpoint_ns": history[1]
+                            .tasks[0]
+                            .state["configurable"]["checkpoint_ns"],
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "task": {
+                        "id": t.id,
+                        "path": list(t.path),
+                    },
+                    "finally_executor": [
+                        {
+                            "config": {
+                                "callbacks": None,
+                                "configurable": {
+                                    "__pregel_dedupe_tasks": True,
+                                    "__pregel_ensure_latest": True,
+                                    "__pregel_resuming": True,
+                                    "checkpoint_id": history[1].config["configurable"][
+                                        "checkpoint_id"
+                                    ],
+                                    "checkpoint_ns": "",
+                                    "thread_id": "1",
+                                },
+                                "metadata": AnyDict(),
+                                "recursion_limit": 25,
+                                "tags": [],
+                            },
+                            "finally_executor": None,
+                            "task": {
+                                "id": history[1].tasks[0].id,
+                                "path": list(history[1].tasks[0].path),
+                            },
+                        }
+                    ],
+                }
+                for c in reversed(child_history[:2])
+                for t in c.tasks
+            ]
+            # "finally" tasks
+            + [
+                {
+                    "config": {
+                        "callbacks": None,
+                        "configurable": {
+                            "__pregel_dedupe_tasks": True,
+                            "__pregel_ensure_latest": True,
+                            "__pregel_resuming": True,
+                            "checkpoint_id": history[1].config["configurable"][
+                                "checkpoint_id"
+                            ],
+                            "checkpoint_ns": "",
+                            "thread_id": "1",
+                        },
+                        "metadata": AnyDict(),
+                        "recursion_limit": 25,
+                        "tags": [],
+                    },
+                    "finally_executor": None,
+                    "task": {
+                        "id": history[1].tasks[0].id,
+                        "path": list(history[1].tasks[0].path),
+                    },
+                }
+            ]
+        )
+    )