Fix subgraph test

libs/langgraph/langgraph/pregel/algo.py

@@ -518,6 +518,7 @@ def prepare_single_task(
                                     **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
                                     parent_ns: checkpoint["id"],
                                 },
+                                "checkpoint_id": None,
                                 "checkpoint_ns": task_checkpoint_ns,
                             },
                         ),

libs/langgraph/langgraph/pregel/loop.py

@@ -336,13 +336,6 @@ class PregelLoop:
             self.status = "done"
             return False
 
-        print(
-            self.step,
-            self.skip_done_tasks,
-            [(t.id, t.name) for t in self.tasks.values()],
-            self.checkpoint_pending_writes,
-        )
-
         # 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:
@@ -350,7 +343,7 @@ class PregelLoop:
                     continue
                 if task := self.tasks.get(tid):
                     if k == SCHEDULED:
-                        if True or v == max(
+                        if v == max(
                             self.checkpoint["versions_seen"]
                             .get(INTERRUPT, {})
                             .values(),

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

@@ -120,6 +120,7 @@ class KafkaExecutor(AbstractAsyncContextManager):
                 config=msg["config"],
                 step=saved.metadata["step"] + 1,
                 for_execution=True,
+                checkpointer=self.graph.checkpointer,
             ):
                 # execute task, saving writes
                 runner = PregelRunner(
@@ -146,5 +147,4 @@ class KafkaExecutor(AbstractAsyncContextManager):
         task_id: str,
         writes: list[tuple[str, Any]],
     ) -> None:
-        print("put_writes", task_id, writes)
         return submit(self.graph.checkpointer.aput_writes, config, writes, task_id)

libs/scheduler-kafka/tests/run.py

@@ -92,6 +92,6 @@ async def drain_topics(
         pass
 
     # check no errors
-    assert not errors
+    assert not errors, errors
 
     return orch_msgs, exec_msgs

libs/scheduler-kafka/tests/test_subgraph.py

@@ -1,11 +1,7 @@
-import asyncio
-import functools
-import re
-from typing import Callable, Literal, Optional, ParamSpec, TypeVar, Union, cast
+from typing import Literal, ParamSpec, TypeVar, cast
 
-import anyio
 import pytest
-from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+from aiokafka import AIOKafkaProducer
 from langchain_core.language_models.fake_chat_models import (
     FakeMessagesListChatModel,
 )
@@ -18,43 +14,15 @@ 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.executor import KafkaExecutor
-from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator
 from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from tests.any import AnyStr
+from tests.run import drain_topics
 
 pytestmark = pytest.mark.anyio
 C = ParamSpec("C")
 R = TypeVar("R")
 
 
-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))
-
-
-def timeout(delay: int):
-    def decorator(func: Callable[C, R]) -> Callable[C, R]:
-        @functools.wraps(func)
-        async def new_func(*args: C.args, **kwargs: C.kwargs) -> R:
-            async with asyncio.timeout(delay):
-                return await func(*args, **kwargs)
-
-        return new_func
-
-    return decorator
-
-
 def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel:
     # copied from test_weather_subgraph
 
@@ -149,45 +117,12 @@ def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel:
     return graph.compile(checkpointer=checkpointer)
 
 
-@timeout(10)
 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)
-    n_orch_msgs = 0
-    n_exec_msgs = 0
-    errors = []
-    scope: Optional[anyio.CancelScope] = None
-
-    async def orchestrator(expected: int) -> None:
-        nonlocal n_orch_msgs
-        async with KafkaOrchestrator(graph, topics) as orch:
-            async for msgs in orch:
-                n_orch_msgs += len(msgs)
-                print("orch", n_orch_msgs, msgs)
-                if n_orch_msgs == expected:
-                    break
-
-    async def executor(expected: int) -> None:
-        nonlocal n_exec_msgs
-        async with KafkaExecutor(graph, topics) as exec:
-            async for msgs in exec:
-                n_exec_msgs += len(msgs)
-                print("exec", n_exec_msgs, msgs)
-                if n_exec_msgs == expected:
-                    break
-
-    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")
 
     # start a new run
     async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
@@ -196,19 +131,17 @@ async def test_subgraph_w_interrupt(
             MessageToOrchestrator(input=input, config=config),
         )
 
-    # run the orchestrator and executor
-    async with anyio.create_task_group() as tg:
-        scope = tg.cancel_scope
-        tg.start_soon(orchestrator, 4, name="orchestrator")
-        tg.start_soon(executor, 3, name="executor")
-
-    # check no errors
-    assert not errors
+    orch_msgs, exec_msgs = await drain_topics(
+        topics,
+        graph,
+        config,
+        until=lambda state: state.next == ("weather_graph",),
+    )
 
     # check interrupted state
     state = await graph.aget_state(config)
-    assert n_orch_msgs == 4
-    assert n_exec_msgs == 3
+    assert len(orch_msgs) == 4
+    assert len(exec_msgs) == 3
     assert state.next == ("weather_graph",)
     assert state.values == {
         "messages": [HumanMessage(id=AnyStr(), content="what's the weather in sf")],
@@ -222,24 +155,19 @@ async def test_subgraph_w_interrupt(
             MessageToOrchestrator(input=None, config=config),
         )
 
-    # run the orchestrator and executor
-    async with anyio.create_task_group() as tg:
-        scope = tg.cancel_scope
-        tg.start_soon(orchestrator, 6, name="orchestrator")
-        tg.start_soon(executor, 4, name="executor")
-
-    # check no errors
-    assert not errors
+    orch_msgs, exec_msgs = await drain_topics(
+        topics, graph, config, until=lambda state: state.next == (), debug=True
+    )
 
     # check final state
     state = await graph.aget_state(config)
-    assert n_orch_msgs == 6
-    assert n_exec_msgs == 4
+    assert len(orch_msgs) == 2
+    assert len(exec_msgs) == 1
     assert state.next == ()
     assert state.values == {
-        "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",
+        "messages": [
+            HumanMessage(id=AnyStr(), content="what's the weather in sf"),
+            AIMessage(content="I'ts sunny in San Francisco!", id=AnyStr()),
+        ],
+        "route": "weather",
     }
-
-    error_task.cancel()