diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a50de8466..83d0f792b 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4910,6 +4910,38 @@ def test_interrupt_functional( f"checkpointer_{checkpointer_name}" ) + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + @entrypoint(checkpointer=checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**fut_foo.result(), "b": value} + fut_bar = bar(bar_input) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + graph.invoke({"a": ""}, config) + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_interrupt_task_functional( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + @task def foo(state: dict) -> dict: return {"a": state["a"] + "foo"} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1662ac478..541781c74 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6287,6 +6287,34 @@ async def test_interrupt_functional(checkpointer_name: str) -> None: async def foo(state: dict) -> dict: return {"a": state["a"] + "foo"} + @task + async def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @entrypoint(checkpointer=checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**foo_result, "b": value} + bar_result = await bar(bar_input) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_interrupt_task_functional(checkpointer_name: str) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + @task async def bar(state: dict) -> dict: value = interrupt("Provide value for bar:") @@ -6305,7 +6333,7 @@ async def test_interrupt_functional(checkpointer_name: str) -> None: await graph.ainvoke({"a": ""}, config) # Resume with an answer res = await graph.ainvoke(Command(resume="bar"), config) - assert res == {"a": "foobar"} + assert res == {"a": "foobar"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)