From e5e78e41928e6b5c0bb6ff1c21ebf5785e8ffc5f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Jun 2025 16:02:35 -0700 Subject: [PATCH] Fix Command(graph=PARENT) when used together w checkpointer=True --- libs/langgraph/langgraph/pregel/retry.py | 4 +- libs/langgraph/tests/test_pregel.py | 47 ++++++++++++++++++++++- libs/langgraph/tests/test_pregel_async.py | 47 ++++++++++++++++++++++- 3 files changed, 92 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 42b73de36..225f6eef3 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -41,7 +41,7 @@ def run_with_retry( except ParentCommand as exc: ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] cmd = exc.args[0] - if cmd.graph == ns: + if cmd.graph in (ns, task.name): # this command is for the current graph, handle it for w in task.writers: w.invoke(cmd, config) @@ -137,7 +137,7 @@ async def arun_with_retry( except ParentCommand as exc: ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] cmd = exc.args[0] - if cmd.graph == ns: + if cmd.graph in (ns, task.name): # this command is for the current graph, handle it for w in task.writers: w.invoke(cmd, config) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 9336e53e4..287b3d376 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4747,7 +4747,10 @@ def test_runnable_passthrough_node_graph() -> None: assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() -def test_parent_command(sync_checkpointer: BaseCheckpointSaver) -> None: +@pytest.mark.parametrize("subgraph_persist", [True, False]) +def test_parent_command( + sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool +) -> None: from langchain_core.messages import BaseMessage from langchain_core.tools import tool @@ -4759,7 +4762,7 @@ def test_parent_command(sync_checkpointer: BaseCheckpointSaver) -> None: subgraph_builder = StateGraph(MessagesState) subgraph_builder.add_node("tool", get_user_name) subgraph_builder.add_edge(START, "tool") - subgraph = subgraph_builder.compile() + subgraph = subgraph_builder.compile(checkpointer=subgraph_persist) class CustomParentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] @@ -7915,3 +7918,43 @@ def test_imp_exception( {"my_task": 2}, {"my_workflow": "done"}, ] + + +@pytest.mark.parametrize("subgraph_persist", [True, False]) +def test_parent_command_goto( + sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool +) -> None: + class State(TypedDict): + dialog_state: Annotated[list[str], operator.add] + + def node_a_child(state): + return {"dialog_state": ["a_child_state"]} + + def node_b_child(state): + return Command( + graph=Command.PARENT, + goto="node_b_parent", + update={"dialog_state": ["b_child_state"]}, + ) + + sub_builder = StateGraph(State) + sub_builder.add_node(node_a_child) + sub_builder.add_node(node_b_child) + sub_builder.add_edge(START, "node_a_child") + sub_builder.add_edge("node_a_child", "node_b_child") + sub_graph = sub_builder.compile(checkpointer=subgraph_persist) + + def node_b_parent(state): + return {"dialog_state": ["node_b_parent"]} + + main_builder = StateGraph(State) + main_builder.add_node(node_b_parent) + main_builder.add_edge(START, "subgraph_node") + main_builder.add_node("subgraph_node", sub_graph, destinations=("node_b_parent",)) + + main_graph = main_builder.compile(sync_checkpointer, name="parent") + config = {"configurable": {"thread_id": 1}} + + assert main_graph.invoke(input={"dialog_state": ["init_state"]}, config=config) == { + "dialog_state": ["init_state", "b_child_state", "node_b_parent"] + } diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f2b7668fa..3bf50611b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6030,7 +6030,10 @@ async def test_debug_nested_subgraphs( assert stream_task.get("state") == history_task.state -async def test_parent_command(async_checkpointer: BaseCheckpointSaver) -> None: +@pytest.mark.parametrize("subgraph_persist", [True, False]) +async def test_parent_command( + async_checkpointer: BaseCheckpointSaver, subgraph_persist: bool +) -> None: from langchain_core.messages import BaseMessage from langchain_core.tools import tool @@ -6042,7 +6045,7 @@ async def test_parent_command(async_checkpointer: BaseCheckpointSaver) -> None: subgraph_builder = StateGraph(MessagesState) subgraph_builder.add_node("tool", get_user_name) subgraph_builder.add_edge(START, "tool") - subgraph = subgraph_builder.compile() + subgraph = subgraph_builder.compile(checkpointer=subgraph_persist) class CustomParentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] @@ -8650,3 +8653,43 @@ async def test_imp_exception( "parent_ids": [], }, ] + + +@pytest.mark.parametrize("subgraph_persist", [True, False]) +async def test_parent_command_goto( + async_checkpointer: BaseCheckpointSaver, subgraph_persist: bool +) -> None: + class State(TypedDict): + dialog_state: Annotated[list[str], operator.add] + + async def node_a_child(state): + return {"dialog_state": ["a_child_state"]} + + async def node_b_child(state): + return Command( + graph=Command.PARENT, + goto="node_b_parent", + update={"dialog_state": ["b_child_state"]}, + ) + + sub_builder = StateGraph(State) + sub_builder.add_node(node_a_child) + sub_builder.add_node(node_b_child) + sub_builder.add_edge(START, "node_a_child") + sub_builder.add_edge("node_a_child", "node_b_child") + sub_graph = sub_builder.compile(checkpointer=subgraph_persist) + + async def node_b_parent(state): + return {"dialog_state": ["node_b_parent"]} + + main_builder = StateGraph(State) + main_builder.add_node(node_b_parent) + main_builder.add_edge(START, "subgraph_node") + main_builder.add_node("subgraph_node", sub_graph, destinations=("node_b_parent",)) + + main_graph = main_builder.compile(async_checkpointer, name="parent") + config = {"configurable": {"thread_id": 1}} + + assert await main_graph.ainvoke( + input={"dialog_state": ["init_state"]}, config=config + ) == {"dialog_state": ["init_state", "b_child_state", "node_b_parent"]}