diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index b5c4245f0..3b8745a65 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -49,6 +49,9 @@ from langgraph._internal._constants import ( CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_RESUME_MAP, + CONFIG_KEY_RESUMING, + CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, INTERRUPT, @@ -81,6 +84,8 @@ _CONF_DROPLIST = frozenset( CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_TASK_ID, + CONFIG_KEY_RESUME_MAP, + CONFIG_KEY_RESUMING, ), ) @@ -433,6 +438,22 @@ class RemoteGraph(PregelProtocol): return sanitized + def _get_command( + self, input: dict[str, Any] | Any, config: RunnableConfig + ) -> CommandSDK | None: + """Build the remote command, forwarding a parent graph's resume value.""" + if isinstance(input, Command): + return cast(CommandSDK, asdict(input)) + conf = config.get(CONF, {}) + scratchpad = conf.get(CONFIG_KEY_SCRATCHPAD) + if scratchpad is None or not conf.get(CONFIG_KEY_RESUMING): + return None + if resume_map := conf.get(CONFIG_KEY_RESUME_MAP): + return {"resume": resume_map} + if (resume := scratchpad.get_null_resume(True)) is not None: + return {"resume": resume} + return None + def get_state( self, config: RunnableConfig, @@ -794,13 +815,12 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) - if isinstance(input, Command): - command: CommandSDK | None = cast(CommandSDK, asdict(input)) + command = self._get_command(input, merged_config) + if command is not None: input = None - else: - command = None thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None) + interrupts: dict[str, Interrupt] = {} for chunk in sync_client.runs.stream( thread_id=thread_id, assistant_id=self.assistant_id, @@ -832,16 +852,23 @@ class RemoteGraph(PregelProtocol): if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): caller_ns = tuple(caller_ns.split(NS_SEP)) ns = caller_ns + ns + is_interrupt = ( + chunk.event.startswith("updates") + and isinstance(chunk.data, dict) + and INTERRUPT in chunk.data + ) + if interrupts and not is_interrupt: + continue # stream to parent stream if stream is not None and mode in stream.modes: stream((ns, mode, chunk.data)) - # raise interrupt or errors - if chunk.event.startswith("updates"): - if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: - if caller_ns: - raise GraphInterrupt( - [Interrupt(**i) for i in chunk.data[INTERRUPT]] - ) + # collect interrupts or raise errors + if is_interrupt: + if caller_ns: + for i in chunk.data[INTERRUPT]: + item = Interrupt(**i) + interrupts.setdefault(item.id, item) + continue elif chunk.event.startswith("error"): raise RemoteException(chunk.data) # filter for what was actually requested @@ -874,6 +901,8 @@ class RemoteGraph(PregelProtocol): yield chunk.data else: yield chunk + if interrupts: + raise GraphInterrupt(list(interrupts.values())) @overload def astream( @@ -949,13 +978,12 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) - if isinstance(input, Command): - command: CommandSDK | None = cast(CommandSDK, asdict(input)) + command = self._get_command(input, merged_config) + if command is not None: input = None - else: - command = None thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None) + interrupts: dict[str, Interrupt] = {} async for chunk in client.runs.stream( thread_id=thread_id, assistant_id=self.assistant_id, @@ -987,16 +1015,23 @@ class RemoteGraph(PregelProtocol): if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): caller_ns = tuple(caller_ns.split(NS_SEP)) ns = caller_ns + ns + is_interrupt = ( + chunk.event.startswith("updates") + and isinstance(chunk.data, dict) + and INTERRUPT in chunk.data + ) + if interrupts and not is_interrupt: + continue # stream to parent stream if stream is not None and mode in stream.modes: stream((ns, mode, chunk.data)) - # raise interrupt or errors - if chunk.event.startswith("updates"): - if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: - if caller_ns: - raise GraphInterrupt( - [Interrupt(**i) for i in chunk.data[INTERRUPT]] - ) + # collect interrupts or raise errors + if is_interrupt: + if caller_ns: + for i in chunk.data[INTERRUPT]: + item = Interrupt(**i) + interrupts.setdefault(item.id, item) + continue elif chunk.event.startswith("error"): raise RemoteException(chunk.data) # filter for what was actually requested @@ -1029,6 +1064,8 @@ class RemoteGraph(PregelProtocol): yield chunk.data else: yield chunk + if interrupts: + raise GraphInterrupt(list(interrupts.values())) def stream_events( self, diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 8e10e209e..05bb2b2ed 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -20,7 +20,7 @@ from langgraph.errors import GraphInterrupt from langgraph.graph import END, START, MessagesState, StateGraph, add_messages from langgraph.pregel import Pregel from langgraph.pregel.remote import RemoteGraph -from langgraph.types import Interrupt, StateSnapshot +from langgraph.types import Command, Interrupt, StateSnapshot from tests.any_str import AnyStr from tests.conftest import NO_DOCKER from tests.example_app.example_graph import app @@ -884,6 +884,133 @@ def test_stream_sanitizes_thread_id(): assert not passed_config["configurable"] +class _ApprovalState(TypedDict, total=False): + request: str + approved: bool + result: str + + +_REMOTE_INTERRUPT_ID = "0123456789abcdef0123456789abcdef" + +_REMOTE_INTERRUPT_RUN = [ + StreamPart( + event="updates", + data={ + "__interrupt__": [ + {"value": {"question": "Approve X?"}, "id": _REMOTE_INTERRUPT_ID} + ] + }, + ) +] + +_REMOTE_DONE_RUN = [ + StreamPart( + event="values", data={"request": "X", "approved": True, "result": "done"} + ) +] + + +def _remote_interrupt_parent(remote_client: MagicMock, next_client: MagicMock): + return ( + StateGraph(_ApprovalState) + .add_node( + "agent_b", + RemoteGraph("agent_b", sync_client=remote_client, client=remote_client), + ) + .add_node( + "agent_c", + RemoteGraph("agent_c", sync_client=next_client, client=next_client), + ) + .add_edge(START, "agent_b") + .add_edge("agent_b", "agent_c") + .compile(checkpointer=InMemorySaver()) + ) + + +@pytest.mark.parametrize( + "resume", + [{"approved": True}, {_REMOTE_INTERRUPT_ID: {"approved": True}}], + ids=["value", "interrupt_id"], +) +def test_subgraph_resume_is_forwarded_to_remote_thread(resume): + remote_client = MagicMock() + remote_client.runs.stream.side_effect = [_REMOTE_INTERRUPT_RUN, _REMOTE_DONE_RUN] + next_client = MagicMock() + next_client.runs.stream.return_value = _REMOTE_DONE_RUN + parent = _remote_interrupt_parent(remote_client, next_client) + config = {"configurable": {"thread_id": "parent_thread"}} + + out = parent.invoke({"request": "X"}, config) + assert out["__interrupt__"] == [ + Interrupt(value={"question": "Approve X?"}, id=_REMOTE_INTERRUPT_ID) + ] + + out = parent.invoke(Command(resume=resume), config) + assert "__interrupt__" not in out + assert out["result"] == "done" + + first, second = remote_client.runs.stream.call_args_list + assert first.kwargs["input"] == {"request": "X"} + assert first.kwargs["command"] is None + assert second.kwargs["input"] is None + assert second.kwargs["command"] == {"resume": resume} + assert second.kwargs["thread_id"] == first.kwargs["thread_id"] + + next_kwargs = next_client.runs.stream.call_args.kwargs + assert next_kwargs["command"] is None + assert next_kwargs["input"] == {"request": "X", "approved": True, "result": "done"} + + +def test_subgraph_surfaces_all_remote_interrupts(): + second_id = "fedcba9876543210fedcba9876543210" + remote_client = MagicMock() + remote_client.runs.stream.return_value = [ + *_REMOTE_INTERRUPT_RUN, + StreamPart(event="values", data={"request": "X"}), + StreamPart( + event="updates", + data={"__interrupt__": [{"value": "second", "id": second_id}]}, + ), + ] + parent = _remote_interrupt_parent(remote_client, MagicMock()) + config = {"configurable": {"thread_id": "parent_thread"}} + + out = parent.invoke({"request": "X"}, config) + + assert out["__interrupt__"] == [ + Interrupt(value={"question": "Approve X?"}, id=_REMOTE_INTERRUPT_ID), + Interrupt(value="second", id=second_id), + ] + + +async def test_subgraph_resume_is_forwarded_to_remote_thread_async(): + remote_client = MagicMock() + remote_iters = [] + for parts in (_REMOTE_INTERRUPT_RUN, _REMOTE_DONE_RUN): + async_iter = MagicMock() + async_iter.__aiter__.return_value = parts + remote_iters.append(async_iter) + remote_client.runs.stream.side_effect = remote_iters + next_client = MagicMock() + next_iter = MagicMock() + next_iter.__aiter__.return_value = _REMOTE_DONE_RUN + next_client.runs.stream.return_value = next_iter + parent = _remote_interrupt_parent(remote_client, next_client) + config = {"configurable": {"thread_id": "parent_thread"}} + + out = await parent.ainvoke({"request": "X"}, config) + assert "__interrupt__" in out + + out = await parent.ainvoke(Command(resume={"approved": True}), config) + assert "__interrupt__" not in out + assert out["result"] == "done" + + second = remote_client.runs.stream.call_args_list[1] + assert second.kwargs["input"] is None + assert second.kwargs["command"] == {"resume": {"approved": True}} + assert next_client.runs.stream.call_args.kwargs["command"] is None + + @pytest.mark.anyio async def test_ainvoke(): # set up test