mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-26 17:42:24 +02:00
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
This commit is contained in:
@@ -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))
|
||||
|
||||
@@ -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"}
|
||||
|
||||
@@ -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"}
|
||||
|
||||
Reference in New Issue
Block a user