From 314115513c2575ef5db3fbd9dc57374e06135ca6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Jun 2025 17:20:33 -0700 Subject: [PATCH] Fix step_timeout causing ParentCommand/GraphInterrupt exception to bubble up - We should not re-raise exceptions in commit() as that is now called in a future done callback - panic_or_proceed takes care of re-raising exceptions as needed anyway --- libs/langgraph/langgraph/pregel/runner.py | 3 +- libs/langgraph/tests/test_pregel.py | 39 +++++++++++++++++++-- libs/langgraph/tests/test_pregel_async.py | 41 +++++++++++++++++++++-- 3 files changed, 76 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 59ede35f5..a10ac9cdc 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -431,7 +431,8 @@ class PregelRunner: writes.extend(resumes) self.put_writes()(task.id, writes) # type: ignore[misc] elif isinstance(exception, GraphBubbleUp): - raise exception + # exception will be raised in _panic_or_proceed + pass else: # save error to checkpointer task.writes.append((ERROR, exception)) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 287b3d376..360ee7fa6 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -43,7 +43,7 @@ from langgraph.checkpoint.base import ( from langgraph.checkpoint.memory import InMemorySaver from langgraph.config import get_stream_writer from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START -from langgraph.errors import InvalidUpdateError +from langgraph.errors import InvalidUpdateError, ParentCommand from langgraph.func import entrypoint, task from langgraph.graph import END, StateGraph from langgraph.graph.message import MessagesState, add_messages @@ -7920,9 +7920,10 @@ def test_imp_exception( ] +@pytest.mark.parametrize("with_timeout", [False, "inner", "outer", "both"]) @pytest.mark.parametrize("subgraph_persist", [True, False]) def test_parent_command_goto( - sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool + sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool, with_timeout: bool ) -> None: class State(TypedDict): dialog_state: Annotated[list[str], operator.add] @@ -7943,6 +7944,8 @@ def test_parent_command_goto( 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) + if with_timeout in ("inner", "both"): + sub_graph.step_timeout = 1 def node_b_parent(state): return {"dialog_state": ["node_b_parent"]} @@ -7951,10 +7954,40 @@ def test_parent_command_goto( 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") + if with_timeout in ("outer", "both"): + main_graph.step_timeout = 1 + 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"] } + + +@pytest.mark.parametrize("with_timeout", [True, False]) +def test_timeout_with_parent_command( + sync_checkpointer: BaseCheckpointSaver, with_timeout: bool +) -> None: + """Test that parent commands are properly propagated during timeouts.""" + + class State(TypedDict): + value: str + + def parent_command_node(state: State) -> State: + time.sleep(0.1) # Add some delay before raising + return Command(graph=Command.PARENT, goto="test_cmd", update={"key": "value"}) + + builder = StateGraph(State) + builder.add_node("parent_cmd", parent_command_node) + builder.set_entry_point("parent_cmd") + graph = builder.compile(checkpointer=sync_checkpointer) + if with_timeout: + graph.step_timeout = 1 + + # Should propagate parent command, not timeout + thread1 = {"configurable": {"thread_id": "1"}} + with pytest.raises(ParentCommand) as exc_info: + graph.invoke({"value": "start"}, thread1) + assert exc_info.value.args[0].goto == "test_cmd" + assert exc_info.value.args[0].update == {"key": "value"} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3bf50611b..8556542ad 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -41,7 +41,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import InMemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START -from langgraph.errors import InvalidUpdateError, NodeInterrupt +from langgraph.errors import InvalidUpdateError, NodeInterrupt, ParentCommand from langgraph.func import entrypoint, task from langgraph.graph import END, StateGraph from langgraph.graph.message import MessagesState, add_messages @@ -8655,9 +8655,12 @@ async def test_imp_exception( ] +@pytest.mark.parametrize("with_timeout", [False, "inner", "outer", "both"]) @pytest.mark.parametrize("subgraph_persist", [True, False]) async def test_parent_command_goto( - async_checkpointer: BaseCheckpointSaver, subgraph_persist: bool + async_checkpointer: BaseCheckpointSaver, + subgraph_persist: bool, + with_timeout: Literal[False, "inner", "outer", "both"], ) -> None: class State(TypedDict): dialog_state: Annotated[list[str], operator.add] @@ -8678,6 +8681,8 @@ async def test_parent_command_goto( 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) + if with_timeout in ("inner", "both"): + sub_graph.step_timeout = 1 async def node_b_parent(state): return {"dialog_state": ["node_b_parent"]} @@ -8686,10 +8691,40 @@ async def test_parent_command_goto( 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") + if with_timeout in ("outer", "both"): + main_graph.step_timeout = 1 + 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"]} + + +@pytest.mark.parametrize("with_timeout", [True, False]) +async def test_timeout_with_parent_command( + async_checkpointer: BaseCheckpointSaver, with_timeout: bool +) -> None: + """Test that parent commands are properly propagated during timeouts.""" + + class State(TypedDict): + value: str + + async def parent_command_node(state: State) -> State: + await asyncio.sleep(0.1) # Add some delay before raising + return Command(graph=Command.PARENT, goto="test_cmd", update={"key": "value"}) + + builder = StateGraph(State) + builder.add_node("parent_cmd", parent_command_node) + builder.set_entry_point("parent_cmd") + graph = builder.compile(checkpointer=async_checkpointer) + if with_timeout: + graph.step_timeout = 1 + + # Should propagate parent command, not timeout + thread1 = {"configurable": {"thread_id": "1"}} + with pytest.raises(ParentCommand) as exc_info: + await graph.ainvoke({"value": "start"}, thread1) + assert exc_info.value.args[0].goto == "test_cmd" + assert exc_info.value.args[0].update == {"key": "value"}