Fix

libs/langgraph/langgraph/pregel/__init__.py

@@ -115,6 +115,7 @@ from langgraph.utils.config import (
     patch_checkpoint_map,
     patch_config,
     patch_configurable,
+    recast_checkpoint_ns,
 )
 from langgraph.utils.fields import get_enhanced_type_hints
 from langgraph.utils.pydantic import create_model
@@ -694,19 +695,15 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            for _, pregel in self.get_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
                 return pregel.get_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     subgraphs=subgraphs,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(self.config, config) if self.config else config
         saved = checkpointer.get_tuple(config)
@@ -731,19 +728,15 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            async for _, pregel in self.aget_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
                 return await pregel.aget_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     subgraphs=subgraphs,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(self.config, config) if self.config else config
         saved = await checkpointer.aget_tuple(config)
@@ -774,13 +767,9 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            for _, pregel in self.get_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
                 yield from pregel.get_state_history(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     filter=filter,
@@ -789,7 +778,7 @@ class Pregel(PregelProtocol):
                 )
                 return
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(
             self.config,
@@ -824,13 +813,9 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            async for _, pregel in self.aget_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
                 async for state in pregel.aget_state_history(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     filter=filter,
@@ -840,7 +825,7 @@ class Pregel(PregelProtocol):
                     yield state
                 return
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(
             self.config,
@@ -879,20 +864,16 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            for _, pregel in self.get_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
                 return pregel.update_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     values,
                     as_node,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         # get last checkpoint
         config = ensure_config(self.config, config)
@@ -1163,20 +1144,16 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            async for _, pregel in self.aget_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
                 return await pregel.aupdate_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     values,
                     as_node,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         # get last checkpoint
         config = ensure_config(self.config, config)

libs/langgraph/langgraph/pregel/retry.py

@@ -48,7 +48,10 @@ def run_with_retry(
                 break
             elif cmd.graph == Command.PARENT:
                 # this command is for the parent graph, assign it to the parent
-                parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
+                parts = ns.split(NS_SEP)
+                if parts[-1].isdigit():
+                    parts.pop()
+                parent_ns = NS_SEP.join(parts[:-1])
                 exc.args = (replace(cmd, graph=parent_ns),)
             # bubble up
             raise
@@ -133,7 +136,10 @@ async def arun_with_retry(
                 break
             elif cmd.graph == Command.PARENT:
                 # this command is for the parent graph, assign it to the parent
-                parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
+                parts = ns.split(NS_SEP)
+                if parts[-1].isdigit():
+                    parts.pop()
+                parent_ns = NS_SEP.join(parts[:-1])
                 exc.args = (replace(cmd, graph=parent_ns),)
             # bubble up
             raise

libs/langgraph/langgraph/utils/config.py

@@ -23,9 +23,25 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_ID,
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINT_NS,
+    NS_END,
+    NS_SEP,
 )
 
 
+def recast_checkpoint_ns(ns: str) -> str:
+    """Remove task IDs from checkpoint namespace.
+
+    Args:
+        ns (str): The checkpoint namespace with task IDs.
+
+    Returns:
+        str: The checkpoint namespace without task IDs.
+    """
+    return NS_SEP.join(
+        part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit()
+    )
+
+
 def patch_configurable(
     config: Optional[RunnableConfig], patch: dict[str, Any]
 ) -> RunnableConfig:

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

@@ -1,5 +1,6 @@
 import asyncio
 import concurrent.futures
+from collections.abc import Sequence
 from contextlib import (
     AbstractAsyncContextManager,
     AbstractContextManager,
@@ -7,7 +8,7 @@ from contextlib import (
     ExitStack,
 )
 from functools import partial
-from typing import Any, Optional, Sequence
+from typing import Any, Optional
 from uuid import UUID
 
 import orjson
@@ -15,7 +16,7 @@ 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.constants import CONFIG_KEY_DELEGATE, ERROR
 from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound
 from langgraph.pregel import Pregel
 from langgraph.pregel.algo import prepare_single_task
@@ -39,7 +40,7 @@ from langgraph.scheduler.kafka.types import (
     Topics,
 )
 from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
-from langgraph.utils.config import patch_configurable
+from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
 
 
 class AsyncKafkaExecutor(AbstractAsyncContextManager):
@@ -165,14 +166,12 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
         # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -183,16 +182,19 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
             raise RuntimeError("Checkpoint not found")
         if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
             raise CheckpointNotLatest()
-        async with AsyncChannelsManager(
-            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), AsyncBackgroundExecutor(msg["config"]) as submit:
+        async with (
+            AsyncChannelsManager(
+                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),
+            AsyncBackgroundExecutor(msg["config"]) as submit,
+        ):
             if task := await asyncio.to_thread(
                 prepare_single_task,
                 msg["task"]["path"],
@@ -378,14 +380,12 @@ class KafkaExecutor(AbstractContextManager):
         # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -396,16 +396,19 @@ class KafkaExecutor(AbstractContextManager):
             raise RuntimeError("Checkpoint not found")
         if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
             raise CheckpointNotLatest()
-        with ChannelsManager(
-            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:
+        with (
+            ChannelsManager(
+                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"],
                 msg["task"]["id"],

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

@@ -16,8 +16,6 @@ from langgraph.constants import (
     CONFIG_KEY_DEDUPE_TASKS,
     CONFIG_KEY_ENSURE_LATEST,
     INTERRUPT,
-    NS_END,
-    NS_SEP,
     SCHEDULED,
 )
 from langgraph.errors import CheckpointNotLatest, GraphInterrupt
@@ -37,7 +35,7 @@ from langgraph.scheduler.kafka.types import (
     Topics,
 )
 from langgraph.types import RetryPolicy
-from langgraph.utils.config import patch_configurable
+from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
 
 
 class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
@@ -140,14 +138,12 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
         # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -329,14 +325,12 @@ class KafkaOrchestrator(AbstractContextManager):
         # 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)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message