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