From f178eb821e52906e1705c9cc02533bb88854b409 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Thu, 26 Feb 2026 13:26:48 -0500 Subject: [PATCH] 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:|1|node:`). 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:|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 | 45 +++++++++++---- libs/langgraph/tests/test_parent_command.py | 53 +++++++++++++++++ .../tests/test_parent_command_async.py | 57 +++++++++++++++++++ libs/langgraph/tests/test_retry.py | 18 +++++- 4 files changed, 160 insertions(+), 13 deletions(-) create mode 100644 libs/langgraph/tests/test_parent_command.py create mode 100644 libs/langgraph/tests/test_parent_command_async.py diff --git a/libs/langgraph/langgraph/pregel/_retry.py b/libs/langgraph/langgraph/pregel/_retry.py index b42b63644..ba62bda6e 100644 --- a/libs/langgraph/langgraph/pregel/_retry.py +++ b/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:|node:`), but the + runtime may also insert a purely-numeric segment (e.g. `|1`) to disambiguate + concurrent tasks (e.g. `parent_first:|1|node:`). + + 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:|1`). + while parts and parts[-1].isdigit(): + parts.pop() + + # Drop the current frame segment itself (e.g. the `node:`). + if parts: + parts.pop() + + # Drop any trailing numeric selectors for the parent frame (e.g. `...|1|node:`). + 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: diff --git a/libs/langgraph/tests/test_parent_command.py b/libs/langgraph/tests/test_parent_command.py new file mode 100644 index 000000000..6b368232b --- /dev/null +++ b/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} diff --git a/libs/langgraph/tests/test_parent_command_async.py b/libs/langgraph/tests/test_parent_command_async.py new file mode 100644 index 000000000..39a077631 --- /dev/null +++ b/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} diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index ac37bea91..864affe2f 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/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