Update kafka

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

@@ -37,7 +37,7 @@ from langgraph.scheduler.kafka.types import (
     Sendable,
     Topics,
 )
-from langgraph.types import RetryPolicy
+from langgraph.types import LoopProtocol, RetryPolicy
 from langgraph.utils.config import patch_configurable
 
 
@@ -183,7 +183,15 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
         if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
             raise CheckpointNotLatest()
         async with AsyncChannelsManager(
-            graph.channels, saved.checkpoint, msg["config"], self.graph.store
+            graph.channels,
+            saved.checkpoint,
+            LoopProtocol(
+                config=msg["config"],
+                store=self.graph.store,
+                step=saved.metadata["step"] + 1,
+                stop=saved.metadata["step"] + 2,
+            ),
+            self.graph.store,
         ) as (channels, managed), AsyncBackgroundExecutor() as submit:
             if task := await asyncio.to_thread(
                 prepare_single_task,
@@ -379,7 +387,14 @@ class KafkaExecutor(AbstractContextManager):
         if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
             raise CheckpointNotLatest()
         with ChannelsManager(
-            graph.channels, saved.checkpoint, msg["config"], self.graph.store
+            graph.channels,
+            saved.checkpoint,
+            LoopProtocol(
+                config=msg["config"],
+                store=self.graph.store,
+                step=saved.metadata["step"] + 1,
+                stop=saved.metadata["step"] + 2,
+            ),
         ) as (channels, managed), BackgroundExecutor({}) as submit:
             if task := prepare_single_task(
                 msg["task"]["path"],