Merge pull request #1691 from langchain-ai/nc/11sep/kafka-sync

kafka: Implement sync orchestrator and executor classes

libs/scheduler-kafka/Makefile

@@ -10,14 +10,16 @@ start-services:
 stop-services:
 	docker compose -f tests/compose.yml down
 
+TEST_PATH ?= .
+
 test:
-	make start-services && poetry run pytest; \
+	make start-services && poetry run pytest $(TEST_PATH); \
 	EXIT_CODE=$$?; \
 	make stop-services; \
 	exit $$EXIT_CODE
 
 test_watch:
-	make start-services && poetry run ptw .; \
+	make start-services && poetry run ptw . -- $(TEST_PATH); \
 	EXIT_CODE=$$?; \
 	make stop-services; \
 	exit $$EXIT_CODE

libs/scheduler-kafka/README.md

@@ -26,7 +26,7 @@ import asyncio
 import logging
 import os
 
-from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator
+from langgraph.scheduler.kafka.orchestrator import AsyncKafkaOrchestrator
 from langgraph.scheduler.kafka.types import Topics
 
 from your_lib import graph # graph expected to be a compiled LangGraph graph
@@ -40,7 +40,7 @@ topics = Topics(
 )
 
 async def main():
-    async with KafkaOrchestrator(graph, topics) as orch:
+    async with AsyncKafkaOrchestrator(graph, topics) as orch:
         async for msgs in orch:
             logger.info('Procesed %d messages', len(msgs))
 
@@ -56,7 +56,7 @@ import asyncio
 import logging
 import os
 
-from langgraph.scheduler.kafka.executor import KafkaExecutor
+from langgraph.scheduler.kafka.executor import AsyncKafkaExecutor
 from langgraph.scheduler.kafka.types import Topics
 
 from your_lib import graph # graph expected to be a compiled LangGraph graph
@@ -70,7 +70,7 @@ topics = Topics(
 )
 
 async def main():
-    async with KafkaExecutor(graph, topics) as orch:
+    async with AsyncKafkaExecutor(graph, topics) as orch:
         async for msgs in orch:
             logger.info('Procesed %d messages', len(msgs))
 
@@ -89,6 +89,8 @@ python executor.py &
 
 ## Configuration
 
+We offer sync and async versions of the orchestrator and executor, `KafkaOrchestrator` and `AsyncKafkaOrchestrator`, and `KafkaExecutor` and `AsyncKafkaExecutor` respectively. The async versions are recommended, especially if you want to process tasks in batches. With the async classes we recommend using `uvloop` for better performance.
+
 You can pass any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor` to configure the consumer:
 
 - batch_max_n (int): Maximum number of messages to include in a single batch. Default: 10.
@@ -131,3 +133,7 @@ By default the orchestrator and executor will attempt to connect to a Kafka brok
 - 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).
+
+### Custom consumer/producer
+
+Both the orchestrator and executor accept a `consumer` and `producer` argument, which should implement the `Consumer` or `Producer` protocols respectively. We expect the consumer to have auto-commit disabled, and the producer and consumer to have no serializers/deserializers set.

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

@@ -1,15 +1,8 @@
-import dataclasses
-from typing import Any, Sequence
-
 import aiokafka
 
 
 class DefaultAsyncConsumer(aiokafka.AIOKafkaConsumer):
-    async def getmany(
-        self, timeout_ms: int, max_records: int
-    ) -> dict[str, Sequence[dict[str, Any]]]:
-        batch = await super().getmany(timeout_ms=timeout_ms, max_records=max_records)
-        return {t: [dataclasses.asdict(m) for m in msgs] for t, msgs in batch.items()}
+    pass
 
 
 class DefaultAsyncProducer(aiokafka.AIOKafkaProducer):

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

@@ -0,0 +1,39 @@
+import concurrent.futures
+from typing import Optional, Sequence
+
+from kafka import KafkaConsumer, KafkaProducer
+from langgraph.scheduler.kafka.types import ConsumerRecord, TopicPartition
+
+
+class DefaultConsumer(KafkaConsumer):
+    def getmany(
+        self, timeout_ms: int, max_records: int
+    ) -> dict[TopicPartition, Sequence[ConsumerRecord]]:
+        return self.poll(timeout_ms=timeout_ms, max_records=max_records)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+class DefaultProducer(KafkaProducer):
+    def send(
+        self,
+        topic: str,
+        *,
+        key: Optional[bytes] = None,
+        value: Optional[bytes] = None,
+    ) -> concurrent.futures.Future:
+        fut = concurrent.futures.Future()
+        kfut = super().send(topic, key=key, value=value)
+        kfut.add_callback(fut.set_result)
+        kfut.add_errback(fut.set_exception)
+        return fut
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()

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

@@ -1,5 +1,11 @@
 import asyncio
-from contextlib import AbstractAsyncContextManager, AsyncExitStack
+import concurrent.futures
+from contextlib import (
+    AbstractAsyncContextManager,
+    AbstractContextManager,
+    AsyncExitStack,
+    ExitStack,
+)
 from functools import partial
 from typing import Any, Optional, Sequence
 
@@ -12,23 +18,29 @@ 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.executor import (
+    AsyncBackgroundExecutor,
+    BackgroundExecutor,
+    Submit,
+)
+from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.runner import PregelRunner
 from langgraph.pregel.types import RetryPolicy
-from langgraph.scheduler.kafka.retry import aretry
+from langgraph.scheduler.kafka.retry import aretry, retry
 from langgraph.scheduler.kafka.types import (
     AsyncConsumer,
     AsyncProducer,
+    Consumer,
     ErrorMessage,
     MessageToExecutor,
     MessageToOrchestrator,
+    Producer,
     Topics,
 )
 from langgraph.utils.config import patch_configurable
 
 
-class KafkaExecutor(AbstractAsyncContextManager):
+class AsyncKafkaExecutor(AbstractAsyncContextManager):
     consumer: AsyncConsumer
 
     producer: AsyncProducer
@@ -93,7 +105,7 @@ class KafkaExecutor(AbstractAsyncContextManager):
             timeout_ms=self.batch_max_ms, max_records=self.batch_max_n
         )
         msgs: list[MessageToExecutor] = [
-            serde.loads(msg["value"]) for msgs in recs.values() for msg in msgs
+            serde.loads(msg.value) for msgs in recs.values() for msg in msgs
         ]
         # process batch
         await asyncio.gather(*(self.each(msg) for msg in msgs))
@@ -189,7 +201,7 @@ class KafkaExecutor(AbstractAsyncContextManager):
                     pass
             else:
                 # task was not found
-                await self.graph.checkpointer.put_writes(
+                await self.graph.checkpointer.aput_writes(
                     msg["config"], [(ERROR, TaskNotFound())]
                 )
         # notify orchestrator
@@ -220,3 +232,195 @@ class KafkaExecutor(AbstractAsyncContextManager):
         writes: list[tuple[str, Any]],
     ) -> None:
         return submit(self.graph.checkpointer.aput_writes, config, writes, task_id)
+
+
+class KafkaExecutor(AbstractContextManager):
+    consumer: Consumer
+
+    producer: Producer
+
+    def __init__(
+        self,
+        graph: Pregel,
+        topics: Topics,
+        *,
+        batch_max_n: int = 10,
+        batch_max_ms: int = 1000,
+        retry_policy: Optional[RetryPolicy] = None,
+        consumer: Optional[Consumer] = None,
+        producer: Optional[Producer] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.graph = graph
+        self.topics = topics
+        self.stack = ExitStack()
+        self.kwargs = kwargs
+        self.consumer = consumer
+        self.producer = producer
+        self.batch_max_n = batch_max_n
+        self.batch_max_ms = batch_max_ms
+        self.retry_policy = retry_policy
+
+    def __enter__(self) -> Self:
+        self.subgraphs = dict(self.graph.get_subgraphs(recurse=True))
+        self.submit = self.stack.enter_context(BackgroundExecutor({}))
+        if self.consumer is None:
+            from langgraph.scheduler.kafka.default_sync import DefaultConsumer
+
+            self.consumer = self.stack.enter_context(
+                DefaultConsumer(
+                    self.topics.executor,
+                    auto_offset_reset="earliest",
+                    group_id="executor",
+                    enable_auto_commit=False,
+                    **self.kwargs,
+                )
+            )
+        if self.producer is None:
+            from langgraph.scheduler.kafka.default_sync import DefaultProducer
+
+            self.producer = self.stack.enter_context(
+                DefaultProducer(
+                    **self.kwargs,
+                )
+            )
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        return self.stack.__exit__(*args)
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> Sequence[MessageToExecutor]:
+        # wait for next batch
+        recs = self.consumer.getmany(
+            timeout_ms=self.batch_max_ms, max_records=self.batch_max_n
+        )
+        msgs: list[MessageToExecutor] = [
+            serde.loads(msg.value) for msgs in recs.values() for msg in msgs
+        ]
+        # process batch
+        concurrent.futures.wait(self.submit(self.each, msg) for msg in msgs)
+        # commit offsets
+        self.consumer.commit()
+        # return message
+        return msgs
+
+    def each(self, msg: MessageToExecutor) -> None:
+        try:
+            retry(self.retry_policy, self.attempt, msg)
+        except CheckpointNotLatest:
+            pass
+        except GraphDelegate as exc:
+            for arg in exc.args:
+                fut = self.producer.send(
+                    self.topics.orchestrator,
+                    value=serde.dumps(
+                        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=serde.dumps(
+                        (
+                            arg["config"]["configurable"]["thread_id"],
+                            arg["config"]["configurable"].get("checkpoint_ns"),
+                        )
+                    ),
+                )
+                fut.result()
+        except Exception as exc:
+            fut = self.producer.send(
+                self.topics.error,
+                value=serde.dumps(
+                    ErrorMessage(
+                        topic=self.topics.executor,
+                        msg=msg,
+                        error=repr(exc),
+                    )
+                ),
+            )
+            fut.result()
+
+    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 = self.graph.checkpointer.get_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()
+        with ChannelsManager(
+            graph.channels, saved.checkpoint, msg["config"], self.graph.store
+        ) as (channels, managed), BackgroundExecutor({}) as submit:
+            if task := 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"]),
+                )
+                for _ in runner.tick([task], reraise=False):
+                    pass
+            else:
+                # task was not found
+                self.graph.checkpointer.put_writes(
+                    msg["config"], [(ERROR, TaskNotFound())]
+                )
+        # notify orchestrator
+        fut = self.producer.send(
+            self.topics.orchestrator,
+            value=serde.dumps(
+                MessageToOrchestrator(
+                    input=None,
+                    config=msg["config"],
+                    finally_executor=msg.get("finally_executor"),
+                )
+            ),
+            # use thread_id, checkpoint_ns as partition key
+            key=serde.dumps(
+                (
+                    msg["config"]["configurable"]["thread_id"],
+                    msg["config"]["configurable"].get("checkpoint_ns"),
+                )
+            ),
+        )
+        fut.result()
+
+    def _put_writes(
+        self,
+        submit: Submit,
+        config: RunnableConfig,
+        task_id: str,
+        writes: list[tuple[str, Any]],
+    ) -> None:
+        return submit(self.graph.checkpointer.put_writes, config, writes, task_id)

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

@@ -1,5 +1,11 @@
 import asyncio
-from contextlib import AbstractAsyncContextManager, AsyncExitStack
+import concurrent.futures
+from contextlib import (
+    AbstractAsyncContextManager,
+    AbstractContextManager,
+    AsyncExitStack,
+    ExitStack,
+)
 from typing import Any, Optional
 
 from langchain_core.runnables import ensure_config
@@ -16,22 +22,25 @@ from langgraph.constants import (
 )
 from langgraph.errors import CheckpointNotLatest, GraphInterrupt
 from langgraph.pregel import Pregel
-from langgraph.pregel.loop import AsyncPregelLoop
+from langgraph.pregel.executor import BackgroundExecutor, Submit
+from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
 from langgraph.pregel.types import RetryPolicy
-from langgraph.scheduler.kafka.retry import aretry
+from langgraph.scheduler.kafka.retry import aretry, retry
 from langgraph.scheduler.kafka.types import (
     AsyncConsumer,
     AsyncProducer,
+    Consumer,
     ErrorMessage,
     ExecutorTask,
     MessageToExecutor,
     MessageToOrchestrator,
+    Producer,
     Topics,
 )
 from langgraph.utils.config import patch_configurable
 
 
-class KafkaOrchestrator(AbstractAsyncContextManager):
+class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
     consumer: AsyncConsumer
 
     producer: AsyncProducer
@@ -95,7 +104,7 @@ class KafkaOrchestrator(AbstractAsyncContextManager):
             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)
+        uniq = set(msg.value for msgs in recs.values() for msg in msgs)
         msgs: list[MessageToOrchestrator] = [serde.loads(msg) for msg in uniq]
         # process batch
         await asyncio.gather(*(self.each(msg) for msg in msgs))
@@ -213,3 +222,183 @@ class KafkaOrchestrator(AbstractAsyncContextManager):
                 )
                 # wait for messages to be sent
                 await asyncio.gather(*futs)
+
+
+class KafkaOrchestrator(AbstractContextManager):
+    consumer: Consumer
+
+    producer: Producer
+
+    submit: Submit
+
+    def __init__(
+        self,
+        graph: Pregel,
+        topics: Topics,
+        batch_max_n: int = 10,
+        batch_max_ms: int = 1000,
+        retry_policy: Optional[RetryPolicy] = None,
+        consumer: Optional[Consumer] = None,
+        producer: Optional[Producer] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.graph = graph
+        self.topics = topics
+        self.stack = ExitStack()
+        self.kwargs = kwargs
+        self.consumer = consumer
+        self.producer = producer
+        self.batch_max_n = batch_max_n
+        self.batch_max_ms = batch_max_ms
+        self.retry_policy = retry_policy
+
+    def __enter__(self) -> Self:
+        self.subgraphs = dict(self.graph.get_subgraphs(recurse=True))
+        self.submit = self.stack.enter_context(BackgroundExecutor({}))
+        if self.consumer is None:
+            from langgraph.scheduler.kafka.default_sync import DefaultConsumer
+
+            self.consumer = self.stack.enter_context(
+                DefaultConsumer(
+                    self.topics.orchestrator,
+                    auto_offset_reset="earliest",
+                    group_id="orchestrator",
+                    enable_auto_commit=False,
+                    **self.kwargs,
+                )
+            )
+        if self.producer is None:
+            from langgraph.scheduler.kafka.default_sync import DefaultProducer
+
+            self.producer = self.stack.enter_context(
+                DefaultProducer(
+                    **self.kwargs,
+                )
+            )
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        return self.stack.__exit__(*args)
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> list[MessageToOrchestrator]:
+        # wait for next batch
+        recs = self.consumer.getmany(
+            timeout_ms=self.batch_max_ms, max_records=self.batch_max_n
+        )
+        print("orch.__next__", recs)
+        # 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]
+        # process batch
+        concurrent.futures.wait(self.submit(self.each, msg) for msg in msgs)
+        # commit offsets
+        self.consumer.commit()
+        # return message
+        return msgs
+
+    def each(self, msg: MessageToOrchestrator) -> None:
+        try:
+            retry(self.retry_policy, self.attempt, msg)
+        except CheckpointNotLatest:
+            pass
+        except GraphInterrupt:
+            pass
+        except Exception as exc:
+            fut = self.producer.send(
+                self.topics.error,
+                value=serde.dumps(
+                    ErrorMessage(
+                        topic=self.topics.orchestrator,
+                        msg=msg,
+                        error=repr(exc),
+                    )
+                ),
+            )
+            fut.result()
+
+    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
+        with SyncPregelLoop(
+            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"):
+                    loop._put_checkpoint_fut.result()
+                # schedule any new tasks
+                if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
+                    # send messages to executor
+                    futures = [
+                        self.producer.send(
+                            self.topics.executor,
+                            value=serde.dumps(
+                                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
+                    concurrent.futures.wait(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 = [
+                    self.producer.send(self.topics.executor, value=serde.dumps(m))
+                    for m in msg["finally_executor"]
+                ]
+                # wait for messages to be sent
+                concurrent.futures.wait(futs)

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

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 import random
+import time
 from typing import Awaitable, Callable, Optional
 
 from typing_extensions import ParamSpec
@@ -11,6 +12,49 @@ logger = logging.getLogger(__name__)
 P = ParamSpec("P")
 
 
+def retry(
+    retry_policy: Optional[RetryPolicy],
+    func: Callable[P, 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:
+            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,
+            )
+            time.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,
+            )
+
+
 async def aretry(
     retry_policy: Optional[RetryPolicy],
     func: Callable[P, Awaitable[None]],

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

@@ -34,10 +34,32 @@ class ErrorMessage(TypedDict):
     msg: Union[MessageToExecutor, MessageToOrchestrator]
 
 
+class TopicPartition(Protocol):
+    topic: str
+    partition: int
+
+
+class ConsumerRecord(Protocol):
+    topic: str
+    "The topic this record is received from"
+    partition: int
+    "The partition from which this record is received"
+    offset: int
+    "The position of this record in the corresponding Kafka partition."
+    timestamp: int
+    "The timestamp of this record"
+    timestamp_type: int
+    "The timestamp type of this record"
+    key: Optional[bytes]
+    "The key (or `None` if no key is specified)"
+    value: Optional[bytes]
+    "The value"
+
+
 class Consumer(Protocol):
     def getmany(
         self, timeout_ms: int, max_records: int
-    ) -> dict[str, Sequence[dict[str, Any]]]: ...
+    ) -> dict[TopicPartition, Sequence[ConsumerRecord]]: ...
 
     def commit(self) -> None: ...
 
@@ -45,7 +67,7 @@ class Consumer(Protocol):
 class AsyncConsumer(Protocol):
     async def getmany(
         self, timeout_ms: int, max_records: int
-    ) -> dict[str, Sequence[dict[str, Any]]]: ...
+    ) -> dict[TopicPartition, Sequence[ConsumerRecord]]: ...
 
     async def commit(self) -> None: ...
 

libs/scheduler-kafka/tests/conftest.py

@@ -3,9 +3,10 @@ from uuid import uuid4
 
 import kafka.admin
 import pytest
-from psycopg import AsyncConnection
-from psycopg_pool import AsyncConnectionPool
+from psycopg import AsyncConnection, Connection
+from psycopg_pool import AsyncConnectionPool, ConnectionPool
 
+from langgraph.checkpoint.postgres import PostgresSaver
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.scheduler.kafka.types import Topics
 
@@ -39,7 +40,7 @@ def topics() -> Iterator[Topics]:
 
 
 @pytest.fixture
-async def checkpointer() -> AsyncIterator[AsyncPostgresSaver]:
+async def acheckpointer() -> AsyncIterator[AsyncPostgresSaver]:
     database = f"test_{uuid4().hex[:16]}"
     # create unique db
     async with await AsyncConnection.connect(
@@ -60,3 +61,23 @@ async def checkpointer() -> AsyncIterator[AsyncPostgresSaver]:
             DEFAULT_POSTGRES_URI, autocommit=True
         ) as conn:
             await conn.execute(f"DROP DATABASE {database}")
+
+
+@pytest.fixture
+def checkpointer() -> Iterator[PostgresSaver]:
+    database = f"test_{uuid4().hex[:16]}"
+    # create unique db
+    with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
+        conn.execute(f"CREATE DATABASE {database}")
+    try:
+        # yield checkpointer
+        with ConnectionPool(
+            DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
+        ) as pool:
+            checkpointer = PostgresSaver(pool)
+            checkpointer.setup()
+            yield checkpointer
+    finally:
+        # drop unique db
+        with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
+            conn.execute(f"DROP DATABASE {database}")

libs/scheduler-kafka/tests/drain.py

@@ -1,4 +1,7 @@
 import asyncio
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
 from typing import Optional, TypeVar
 
 import anyio
@@ -6,15 +9,19 @@ 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.default_sync import DefaultConsumer
+from langgraph.scheduler.kafka.executor import AsyncKafkaExecutor, KafkaExecutor
+from langgraph.scheduler.kafka.orchestrator import (
+    AsyncKafkaOrchestrator,
+    KafkaOrchestrator,
+)
 from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
 
 C = ParamSpec("C")
 R = TypeVar("R")
 
 
-async def drain_topics(
+async def drain_topics_async(
     topics: Topics, graph: Pregel, *, debug: bool = False
 ) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]:
     scope: Optional[anyio.CancelScope] = None
@@ -31,7 +38,7 @@ async def drain_topics(
         )
 
     async def orchestrator() -> None:
-        async with KafkaOrchestrator(graph, topics) as orch:
+        async with AsyncKafkaOrchestrator(graph, topics) as orch:
             async for msgs in orch:
                 orch_msgs.append(msgs)
                 if debug:
@@ -40,7 +47,7 @@ async def drain_topics(
                     scope.cancel()
 
     async def executor() -> None:
-        async with KafkaExecutor(graph, topics) as exec:
+        async with AsyncKafkaExecutor(graph, topics) as exec:
             async for msgs in exec:
                 exec_msgs.append(msgs)
                 if debug:
@@ -77,3 +84,83 @@ async def drain_topics(
     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]
+
+
+def drain_topics(
+    topics: Topics, graph: Pregel, *, debug: bool = False
+) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]:
+    orch_msgs = []
+    exec_msgs = []
+    errors = []
+    event = threading.Event()
+
+    def done() -> bool:
+        return (
+            len(orch_msgs) > 0
+            and len(exec_msgs) > 0
+            and not orch_msgs[-1]
+            and not exec_msgs[-1]
+        )
+
+    def orchestrator() -> None:
+        try:
+            with KafkaOrchestrator(graph, topics) as orch:
+                for msgs in orch:
+                    orch_msgs.append(msgs)
+                    if debug:
+                        print("\n---\norch", len(msgs), msgs)
+                    if done():
+                        print("am i done? orchestrator")
+                        event.set()
+                    if event.is_set():
+                        break
+        except Exception as e:
+            errors.append(e)
+            event.set()
+
+    def executor() -> None:
+        try:
+            with KafkaExecutor(graph, topics) as exec:
+                for msgs in exec:
+                    exec_msgs.append(msgs)
+                    if debug:
+                        print("\n---\nexec", len(msgs), msgs)
+                    if done():
+                        print("am i done? executor")
+                        event.set()
+                    if event.is_set():
+                        break
+        except Exception as e:
+            errors.append(e)
+            event.set()
+
+    def error_consumer() -> None:
+        try:
+            with DefaultConsumer(topics.error) as consumer:
+                while not event.is_set():
+                    if msg := consumer.poll(timeout_ms=100):
+                        errors.append(msg)
+                        event.set()
+        except Exception as e:
+            errors.append(e)
+            event.set()
+
+    with ThreadPoolExecutor() as pool:
+        # start error consumer
+        pool.submit(error_consumer)
+
+        # run the orchestrator and executor until break_when
+        pool.submit(orchestrator)
+        pool.submit(executor)
+
+        # timeout
+        start = time.time()
+        while not event.is_set():
+            time.sleep(0.1)
+            if time.time() - start > 20:
+                event.set()
+
+    # 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

@@ -16,7 +16,7 @@ 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
+from tests.drain import drain_topics_async
 
 pytestmark = pytest.mark.anyio
 
@@ -86,10 +86,10 @@ def mk_fanout_graph(
     return builder.compile(checkpointer, interrupt_before=interrupt_before)
 
 
-async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None:
+async def test_fanout_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
     input = {"query": "what is weather in sf"}
     config = {"configurable": {"thread_id": "1"}}
-    graph = mk_fanout_graph(checkpointer)
+    graph = mk_fanout_graph(acheckpointer)
 
     # start a new run
     async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
@@ -99,7 +99,7 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -
         )
 
     # drain topics
-    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
 
     # check state
     state = await graph.aget_state(config)
@@ -169,11 +169,11 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -
 
 
 async def test_fanout_graph_w_interrupt(
-    topics: Topics, checkpointer: BaseCheckpointSaver
+    topics: Topics, acheckpointer: BaseCheckpointSaver
 ) -> None:
     input = {"query": "what is weather in sf"}
     config = {"configurable": {"thread_id": "1"}}
-    graph = mk_fanout_graph(checkpointer, interrupt_before=["qa"])
+    graph = mk_fanout_graph(acheckpointer, interrupt_before=["qa"])
 
     # start a new run
     async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
@@ -182,7 +182,7 @@ async def test_fanout_graph_w_interrupt(
             MessageToOrchestrator(input=input, config=config),
         )
 
-    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
 
     # check interrupted state
     state = await graph.aget_state(config)
@@ -258,7 +258,7 @@ async def test_fanout_graph_w_interrupt(
             MessageToOrchestrator(input=None, config=config),
         )
 
-    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
 
     # check final state
     state = await graph.aget_state(config)

libs/scheduler-kafka/tests/test_fanout_sync.py

@@ -0,0 +1,273 @@
+import operator
+import time
+from typing import (
+    Annotated,
+    Sequence,
+    TypedDict,
+    Union,
+)
+
+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.default_sync import DefaultProducer
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from tests.any import AnyDict
+from tests.drain import drain_topics
+
+
+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]
+
+    def rewrite_query(data: State) -> State:
+        return {"query": f'query: {data["query"]}'}
+
+    def retriever_picker(data: State) -> list[str]:
+        return ["analyzer_one", "retriever_two"]
+
+    def analyzer_one(data: State) -> State:
+        return {"query": f'analyzed: {data["query"]}'}
+
+    def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    def retriever_two(data: State) -> State:
+        time.sleep(0.1)
+        return {"docs": ["doc3", "doc4"]}
+
+    def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    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)
+
+
+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
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=input, config=config)),
+        )
+        producer.flush()
+
+    # drain topics
+    orch_msgs, exec_msgs = drain_topics(topics, graph, debug=1)
+
+    # check state
+    state = graph.get_state(config)
+    assert state.next == ()
+    assert (
+        state.values
+        == graph.invoke(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 for c in graph.get_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
+    ]
+
+
+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
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=input, config=config)),
+        )
+        producer.flush()
+
+    orch_msgs, exec_msgs = drain_topics(topics, graph, debug=1)
+
+    # check interrupted state
+    state = graph.get_state(config)
+    assert state.next == ("qa",)
+    assert (
+        state.values
+        == graph.invoke(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 for c in graph.get_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
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=None, config=config)),
+        )
+        producer.flush()
+
+    orch_msgs, exec_msgs = drain_topics(topics, graph)
+
+    # check final state
+    state = graph.get_state(config)
+    assert state.next == ()
+    assert (
+        state.values
+        == graph.invoke(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

@@ -16,7 +16,7 @@ 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
+from tests.drain import drain_topics_async
 from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
 
 pytestmark = pytest.mark.anyio
@@ -117,11 +117,11 @@ def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel:
 
 
 async def test_subgraph_w_interrupt(
-    topics: Topics, checkpointer: BaseCheckpointSaver
+    topics: Topics, acheckpointer: BaseCheckpointSaver
 ) -> None:
     input = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
     config = {"configurable": {"thread_id": "1"}}
-    graph = mk_weather_graph(checkpointer)
+    graph = mk_weather_graph(acheckpointer)
 
     # start a new run
     async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
@@ -130,7 +130,7 @@ async def test_subgraph_w_interrupt(
             MessageToOrchestrator(input=input, config=config),
         )
 
-    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
 
     # check interrupted state
     state = await graph.aget_state(config)
@@ -415,7 +415,7 @@ async def test_subgraph_w_interrupt(
             MessageToOrchestrator(input=None, config=config),
         )
 
-    orch_msgs, exec_msgs = await drain_topics(topics, graph)
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
 
     # check final state
     state = await graph.aget_state(config)

libs/scheduler-kafka/tests/test_subgraph_sync.py

@@ -0,0 +1,699 @@
+from typing import Literal, cast
+
+import pytest
+from langchain_core.language_models.fake_chat_models import (
+    FakeMessagesListChatModel,
+)
+from langchain_core.messages import AIMessage, 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.default_sync import DefaultProducer
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from tests.any import AnyDict, AnyStr
+from tests.drain import drain_topics
+from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
+
+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"
+
+    def weather_graph(state: RouterState):
+        return subgraph.invoke(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)
+
+
+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
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=input, config=config)),
+        )
+        producer.flush()
+
+    orch_msgs, exec_msgs = drain_topics(topics, graph)
+
+    # check interrupted state
+    state = graph.get_state(config)
+    assert state.next == ("weather_graph",)
+    assert state.values == {
+        "messages": [_AnyIdHumanMessage(content="what's the weather in sf")],
+        "route": "weather",
+    }
+
+    # check outer history
+    history = [c for c in graph.get_state_history(config)]
+    assert len(history) == 3
+
+    # check child history
+    child_history = [c for c in graph.get_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
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=None, config=config)),
+        )
+        producer.flush()
+
+    orch_msgs, exec_msgs = drain_topics(topics, graph)
+
+    # check final state
+    state = graph.get_state(config)
+    assert state.next == ()
+    assert state.values == {
+        "messages": [
+            _AnyIdHumanMessage(content="what's the weather in sf"),
+            _AnyIdAIMessage(content="I'ts sunny in San Francisco!"),
+        ],
+        "route": "weather",
+    }
+
+    # check outer history
+    history = [c for c in graph.get_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 for c in graph.get_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),
+                    },
+                }
+            ]
+        )
+    )