fix(langgraph): correct ParentCommand bubbling when checkpoint_ns includes numeric task segments (#6864)

Fixes incorrect `Command.PARENT` bubbling when checkpoint namespaces
include numeric task-disambiguation segments like `|1`. In some
nested-invoke/fanout scenarios, the runtime inserts a purely-numeric
namespace segment between `name:task_id` segments (e.g.
`parent_first:<id>|1|node:<id>`). The previous ParentCommand rewrite
logic only handled numeric segments at the end of the namespace, which
could produce a malformed parent graph identifier (e.g.
`parent_first:<id>|1`) and prevent the command from routing to the
intended parent node.

This change normalizes checkpoint namespaces by dropping numeric
segments before computing the parent namespace in both sync and async
retry paths. Added a minimal regression test that exercises the
nested-invoke case and asserts that `Command(graph=Command.PARENT,
goto=...)` reliably routes to the parent graph, regardless of whether
the jump comes from the first or second nested invocation.

libs/langgraph/langgraph/pregel/_retry.py

@@ -23,6 +23,35 @@ logger = logging.getLogger(__name__)
 SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
 
 
+def _checkpoint_ns_for_parent_command(ns: str) -> str:
+    """Return the checkpoint namespace for the parent graph.
+
+    The checkpoint namespace is a `|`-separated path. Each segment is usually
+    of the form `name:task_id` (e.g. `parent_first:<uuid>|node:<uuid>`), but the
+    runtime may also insert a purely-numeric segment (e.g. `|1`) to disambiguate
+    concurrent tasks (e.g. `parent_first:<uuid>|1|node:<uuid>`).
+
+    Numeric segments are not real path levels, so we drop them before computing
+    the parent namespace.
+    """
+
+    parts = ns.split(NS_SEP)
+
+    # Drop any trailing numeric selectors for the current frame (e.g. `...|node:<id>|1`).
+    while parts and parts[-1].isdigit():
+        parts.pop()
+
+    # Drop the current frame segment itself (e.g. the `node:<id>`).
+    if parts:
+        parts.pop()
+
+    # Drop any trailing numeric selectors for the parent frame (e.g. `...|1|node:<id>`).
+    while parts and parts[-1].isdigit():
+        parts.pop()
+
+    return NS_SEP.join(parts)
+
+
 def run_with_retry(
     task: PregelExecutableTask,
     retry_policy: Sequence[RetryPolicy] | None,
@@ -50,12 +79,8 @@ def run_with_retry(
                     w.invoke(cmd, config)
                 break
             elif cmd.graph == Command.PARENT:
-                # this command is for the parent graph, assign it to the parent
-                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),)
+                # this command is for the parent graph, assign it to the parent.
+                exc.args = (replace(cmd, graph=_checkpoint_ns_for_parent_command(ns)),)
             # bubble up
             raise
         except GraphBubbleUp:
@@ -146,12 +171,8 @@ async def arun_with_retry(
                     w.invoke(cmd, config)
                 break
             elif cmd.graph == Command.PARENT:
-                # this command is for the parent graph, assign it to the parent
-                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),)
+                # this command is for the parent graph, assign it to the parent.
+                exc.args = (replace(cmd, graph=_checkpoint_ns_for_parent_command(ns)),)
             # bubble up
             raise
         except GraphBubbleUp:

libs/langgraph/tests/test_parent_command.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing_extensions import TypedDict
+
+from langgraph.graph import END, START, StateGraph
+from langgraph.types import Command
+
+
+def test_parent_command_from_nested_subgraph() -> None:
+    class ParentState(TypedDict):
+        jump_from_idx: int
+
+    class ChildState(TypedDict):
+        jump: bool
+
+    child_builder: StateGraph[ChildState] = StateGraph(ChildState)
+
+    def child_node(state: ChildState) -> Command | ChildState:
+        if state["jump"]:
+            return Command(graph=Command.PARENT, goto="parent_second")
+        return state
+
+    child_builder.add_node("node", child_node)
+    child_builder.add_edge(START, "node")
+
+    child_0 = child_builder.compile()
+    child_1 = child_builder.compile()
+
+    parent_builder: StateGraph[ParentState] = StateGraph(ParentState)
+
+    def parent_first(state: ParentState) -> ParentState:
+        child_0.invoke({"jump": state["jump_from_idx"] == 1})
+        if state["jump_from_idx"] == 1:
+            raise AssertionError("Shouldn't be here")
+
+        child_1.invoke({"jump": state["jump_from_idx"] == 2})
+        if state["jump_from_idx"] == 2:
+            raise AssertionError("Shouldn't be here")
+
+        return state
+
+    def parent_second(state: ParentState) -> ParentState:
+        return state
+
+    parent_builder.add_node("parent_first", parent_first)
+    parent_builder.add_node("parent_second", parent_second)
+    parent_builder.add_edge(START, "parent_first")
+    parent_builder.add_edge("parent_second", END)
+
+    graph = parent_builder.compile()
+
+    assert graph.invoke({"jump_from_idx": 1}) == {"jump_from_idx": 1}
+    assert graph.invoke({"jump_from_idx": 2}) == {"jump_from_idx": 2}

libs/langgraph/tests/test_parent_command_async.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import pytest
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import TypedDict
+
+from langgraph.graph import END, START, StateGraph
+from langgraph.types import Command
+
+pytestmark = pytest.mark.anyio
+
+
+async def test_parent_command_from_nested_subgraph() -> None:
+    class ParentState(TypedDict):
+        jump_from_idx: int
+
+    class ChildState(TypedDict):
+        jump: bool
+
+    child_builder: StateGraph[ChildState] = StateGraph(ChildState)
+
+    async def child_node(state: ChildState) -> Command | ChildState:
+        if state["jump"]:
+            return Command(graph=Command.PARENT, goto="parent_second")
+        return state
+
+    child_builder.add_node("node", child_node)
+    child_builder.add_edge(START, "node")
+
+    child_0 = child_builder.compile()
+    child_1 = child_builder.compile()
+
+    parent_builder: StateGraph[ParentState] = StateGraph(ParentState)
+
+    async def parent_first(state: ParentState, config: RunnableConfig) -> ParentState:
+        await child_0.ainvoke({"jump": state["jump_from_idx"] == 1}, config)
+        if state["jump_from_idx"] == 1:
+            raise AssertionError("Shouldn't be here")
+
+        await child_1.ainvoke({"jump": state["jump_from_idx"] == 2}, config)
+        if state["jump_from_idx"] == 2:
+            raise AssertionError("Shouldn't be here")
+
+        return state
+
+    async def parent_second(state: ParentState) -> ParentState:
+        return state
+
+    parent_builder.add_node("parent_first", parent_first)
+    parent_builder.add_node("parent_second", parent_second)
+    parent_builder.add_edge(START, "parent_first")
+    parent_builder.add_edge("parent_second", END)
+
+    graph = parent_builder.compile().with_config(recursion_limit=10)
+
+    assert await graph.ainvoke({"jump_from_idx": 1}) == {"jump_from_idx": 1}
+    assert await graph.ainvoke({"jump_from_idx": 2}) == {"jump_from_idx": 2}

libs/langgraph/tests/test_retry.py

@@ -4,7 +4,7 @@ import pytest
 from typing_extensions import TypedDict
 
 from langgraph.graph import START, StateGraph
-from langgraph.pregel._retry import _should_retry_on
+from langgraph.pregel._retry import _checkpoint_ns_for_parent_command, _should_retry_on
 from langgraph.types import RetryPolicy
 
 
@@ -78,6 +78,22 @@ def test_should_retry_on_empty_sequence():
     assert _should_retry_on(policy, ValueError("test error")) is False
 
 
+def test_checkpoint_ns_for_parent_command() -> None:
+    assert _checkpoint_ns_for_parent_command("") == ""
+    assert _checkpoint_ns_for_parent_command("node:1") == ""
+    assert _checkpoint_ns_for_parent_command("node:1|child:2") == "node:1"
+    assert _checkpoint_ns_for_parent_command("node:1|1|child:2") == "node:1"
+    assert _checkpoint_ns_for_parent_command("node:1|1|child:2|1") == "node:1"
+    assert (
+        _checkpoint_ns_for_parent_command("parent:1|1|child:1|1|node:1|1")
+        == "parent:1|1|child:1"
+    )
+    assert (
+        _checkpoint_ns_for_parent_command("parent:1|1|child:1|1|node:1")
+        == "parent:1|1|child:1"
+    )
+
+
 def test_should_retry_default_retry_on():
     """Test the default retry_on function."""
     import httpx