From f637d52b93c8ac2eb0d0c8c07b7d9dc8a9c13e1c Mon Sep 17 00:00:00 2001 From: Dariel Dato-on <293921406+darieldatoon@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:37:48 -0500 Subject: [PATCH] fix(langgraph): keep RemoteGraph interrupts and errors intact on retry - Raise when a remote pauses with an empty interrupt list (interrupt_before/after). - Raise remote errors that arrive after an interrupt instead of skipping them. - Forward each resume map answer to a remote at most once per run. A retry of the node or any ancestor re-raises the pending interrupts instead of resending it. Co-Authored-By: Claude Opus 5.5 --- .../langgraph/_internal/_constants.py | 3 + libs/langgraph/langgraph/pregel/_loop.py | 2 + libs/langgraph/langgraph/pregel/remote.py | 22 ++- libs/langgraph/tests/test_remote_graph.py | 183 +++++++++++++++++- 4 files changed, 203 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index 7a369b891..7386d8cc5 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -74,6 +74,8 @@ CONFIG_KEY_RUNTIME = sys.intern("__pregel_runtime") # holds a `Runtime` instance with context, store, stream writer, etc. CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map") # holds a mapping of task ns -> resume value for resuming tasks +CONFIG_KEY_RESUME_MAP_SENT = sys.intern("__pregel_resume_map_sent") +# holds the set of resume map ids already forwarded to a remote graph in this run CONFIG_KEY_STREAM_MESSAGES_V2 = sys.intern("__pregel_stream_messages_v2") # when True, attach StreamMessagesHandlerV2 so content-block (v2) events # flow through stream_mode="messages"; set by StreamingHandler only. @@ -134,6 +136,7 @@ RESERVED = { CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_TIMED_ATTEMPT_OBSERVER, CONFIG_KEY_RESUME_MAP, + CONFIG_KEY_RESUME_MAP_SENT, CONFIG_KEY_STREAM_MESSAGES_V2, # other constants PUSH, diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 008aefdb1..a0459d176 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -43,6 +43,7 @@ from langgraph._internal._constants import ( CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_REPLAY_STATE, CONFIG_KEY_RESUME_MAP, + CONFIG_KEY_RESUME_MAP_SENT, CONFIG_KEY_RESUMING, CONFIG_KEY_RUNTIME, CONFIG_KEY_SCRATCHPAD, @@ -988,6 +989,7 @@ class PregelLoop: and all(is_xxh3_128_hexdigest(k) for k in resume) ): self.config[CONF][CONFIG_KEY_RESUME_MAP] = resume + self.config[CONF][CONFIG_KEY_RESUME_MAP_SENT] = set() else: if len(self._pending_interrupts()) > 1: raise RuntimeError( diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index d596ebf42..26940aaab 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -52,6 +52,7 @@ from langgraph._internal._constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUME_MAP, + CONFIG_KEY_RESUME_MAP_SENT, CONFIG_KEY_RESUMING, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_STREAM, @@ -87,6 +88,7 @@ _CONF_DROPLIST = frozenset( CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_TASK_ID, CONFIG_KEY_RESUME_MAP, + CONFIG_KEY_RESUME_MAP_SENT, CONFIG_KEY_RESUMING, ), ) @@ -462,10 +464,16 @@ class RemoteGraph(PregelProtocol): return None conf = config[CONF] if resume_map := conf.get(CONFIG_KEY_RESUME_MAP): + sent: set[str] = conf.get(CONFIG_KEY_RESUME_MAP_SENT, set()) # Do not forward the whole map: it can hold answers for other remotes. + # Send each answer once per run: a later interrupt in the same remote + # node reuses the id, so a retry would apply the old answer to it. if scoped := { - i.id: resume_map[i.id] for i in pending if i.id in resume_map + i.id: resume_map[i.id] + for i in pending + if i.id in resume_map and i.id not in sent }: + sent.update(scoped) return {"resume": scoped} elif (resume := conf[CONFIG_KEY_SCRATCHPAD].get_null_resume(True)) is not None: return {"resume": resume} @@ -845,6 +853,7 @@ class RemoteGraph(PregelProtocol): if command is not None: input = None + interrupted = False interrupts: dict[str, Interrupt] = {} for chunk in sync_client.runs.stream( thread_id=thread_id, @@ -882,7 +891,7 @@ class RemoteGraph(PregelProtocol): and isinstance(chunk.data, dict) and INTERRUPT in chunk.data ) - if interrupts and not is_interrupt: + if interrupted and not is_interrupt and not chunk.event.startswith("error"): continue # stream to parent stream if stream is not None and mode in stream.modes: @@ -890,6 +899,7 @@ class RemoteGraph(PregelProtocol): # collect interrupts or raise errors if is_interrupt: if caller_ns: + interrupted = True for i in chunk.data[INTERRUPT]: item = Interrupt(**i) interrupts.setdefault(item.id, item) @@ -926,7 +936,7 @@ class RemoteGraph(PregelProtocol): yield chunk.data else: yield chunk - if interrupts: + if interrupted: raise GraphInterrupt(list(interrupts.values())) @overload @@ -1016,6 +1026,7 @@ class RemoteGraph(PregelProtocol): if command is not None: input = None + interrupted = False interrupts: dict[str, Interrupt] = {} async for chunk in client.runs.stream( thread_id=thread_id, @@ -1053,7 +1064,7 @@ class RemoteGraph(PregelProtocol): and isinstance(chunk.data, dict) and INTERRUPT in chunk.data ) - if interrupts and not is_interrupt: + if interrupted and not is_interrupt and not chunk.event.startswith("error"): continue # stream to parent stream if stream is not None and mode in stream.modes: @@ -1061,6 +1072,7 @@ class RemoteGraph(PregelProtocol): # collect interrupts or raise errors if is_interrupt: if caller_ns: + interrupted = True for i in chunk.data[INTERRUPT]: item = Interrupt(**i) interrupts.setdefault(item.id, item) @@ -1097,7 +1109,7 @@ class RemoteGraph(PregelProtocol): yield chunk.data else: yield chunk - if interrupts: + if interrupted: raise GraphInterrupt(list(interrupts.values())) def stream_events( diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 42b8b93a6..1435bf07b 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -16,10 +16,12 @@ from langgraph_sdk.schema import StreamPart from pydantic import BaseModel from typing_extensions import TypedDict +from langgraph.channels.last_value import LastValue from langgraph.errors import GraphInterrupt +from langgraph.func import task from langgraph.graph import END, START, MessagesState, StateGraph, add_messages -from langgraph.pregel import Pregel -from langgraph.pregel.remote import RemoteGraph +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.pregel.remote import RemoteException, RemoteGraph from langgraph.types import Command, Interrupt, RetryPolicy, StateSnapshot from tests.any_str import AnyStr from tests.conftest import NO_DOCKER @@ -1081,6 +1083,124 @@ def test_subgraph_resume_map_is_scoped_to_each_remote(): } +def _resumed_then_dropped(): + # A remote node's next interrupt() reuses the id of the one just resumed. + yield from _REMOTE_INTERRUPT_RUN + raise ValueError("transient") + + +def test_subgraph_retry_does_not_resend_resume_map(): + remote_client, next_client = _sync_clients( + [_REMOTE_INTERRUPT_RUN, _resumed_then_dropped()], [] + ) + parent = _remote_interrupt_parent(remote_client, next_client) + config = {"configurable": {"thread_id": "parent_thread"}} + parent.invoke({"request": "X"}, config) + + out = parent.invoke( + Command(resume={_REMOTE_INTERRUPT_ID: {"approved": True}}), config + ) + + assert out["__interrupt__"] == [ + Interrupt(value={"question": "Approve X?"}, id=_REMOTE_INTERRUPT_ID) + ] + assert remote_client.runs.stream.call_count == 2 + + +def _ancestor_retry_parent(remote_client: MagicMock, via: str = "node"): + # A low-level Pregel subgraph has no START step, so an ancestor retry's + # RESUMING flag reaches the RemoteGraph task with a fresh node_attempt. + remote = RemoteGraph("agent_b", sync_client=remote_client, client=remote_client) + inner = Pregel( + nodes={ + "agent_b": NodeBuilder() + .subscribe_only("request") + .do(remote) + .write_to("result") + }, + channels={"request": LastValue(str), "result": LastValue(object)}, + input_channels=["request"], + output_channels=["result"], + ) + if via == "node": + p = inner + elif via == "task": + + @task + def child(request: str): + return inner.invoke({"request": request})["result"] + + def p(state: _ApprovalState): + return {"result": child(state["request"]).result()} + + else: + + @task + async def child(request: str): + return (await inner.ainvoke({"request": request}))["result"] + + async def p(state: _ApprovalState): + return {"result": await child(state["request"])} + + return ( + StateGraph(_ApprovalState) + .add_node("p", p, retry_policy=_RETRY) + .add_edge(START, "p") + .compile(checkpointer=InMemorySaver()) + ) + + +@pytest.mark.parametrize("via", ["node", "task"]) +def test_subgraph_ancestor_retry_does_not_resend_resume_map(via): + remote_client, _ = _sync_clients( + [_REMOTE_INTERRUPT_RUN, _resumed_then_dropped(), _REMOTE_DONE_RUN], [] + ) + parent = _ancestor_retry_parent(remote_client, via) + config = {"configurable": {"thread_id": "parent_thread"}} + parent.invoke({"request": "X"}, config) + + out = parent.invoke(Command(resume={_REMOTE_INTERRUPT_ID: "first"}), config) + assert out["__interrupt__"] == [ + Interrupt(value={"question": "Approve X?"}, id=_REMOTE_INTERRUPT_ID) + ] + + parent.invoke(Command(resume={_REMOTE_INTERRUPT_ID: "second"}), config) + assert [c.kwargs["command"] for c in remote_client.runs.stream.call_args_list] == [ + None, + {"resume": {_REMOTE_INTERRUPT_ID: "first"}}, + {"resume": {_REMOTE_INTERRUPT_ID: "second"}}, + ] + + +def test_subgraph_empty_remote_interrupt_pauses_parent(): + remote_client, next_client = _sync_clients( + [[StreamPart(event="updates", data={"__interrupt__": []})]], [] + ) + parent = _remote_interrupt_parent(remote_client, next_client) + config = {"configurable": {"thread_id": "parent_thread"}} + + parent.invoke({"request": "X"}, config) + + assert parent.get_state(config).next == ("agent_b",) + next_client.runs.stream.assert_not_called() + + +def test_subgraph_remote_error_after_interrupt_is_raised(): + remote_client, next_client = _sync_clients( + [ + [ + *_REMOTE_INTERRUPT_RUN, + StreamPart(event="error", data={"error": "boom", "message": "x"}), + ] + ], + [], + ) + parent = _remote_interrupt_parent(remote_client, next_client) + + with pytest.raises(RemoteException): + parent.invoke({"request": "X"}, {"configurable": {"thread_id": "t"}}) + + def test_subgraph_surfaces_all_remote_interrupts(): second_id = "fedcba9876543210fedcba9876543210" remote_client, next_client = _sync_clients( @@ -1140,6 +1260,65 @@ async def test_subgraph_resume_is_forwarded_to_remote_thread_async(): assert next_client.runs.stream.call_args.kwargs["command"] is None +async def test_subgraph_retry_does_not_resend_resume_map_async(): + async def run(*parts: StreamPart, error: Exception | None = None): + for part in parts: + yield part + if error: + raise error + + remote_client = MagicMock() + remote_client.runs.stream.side_effect = [ + run(*_REMOTE_INTERRUPT_RUN), + run(*_REMOTE_INTERRUPT_RUN, error=ValueError("transient")), + ] + remote_client.threads.get = AsyncMock(return_value=_INTERRUPTED_THREAD) + parent = _remote_interrupt_parent(remote_client, MagicMock()) + config = {"configurable": {"thread_id": "parent_thread"}} + await parent.ainvoke({"request": "X"}, config) + + resume = {_REMOTE_INTERRUPT_ID: {"approved": True}} + out = await parent.ainvoke(Command(resume=resume), config) + + assert out["__interrupt__"] == [ + Interrupt(value={"question": "Approve X?"}, id=_REMOTE_INTERRUPT_ID) + ] + assert remote_client.runs.stream.call_count == 2 + assert remote_client.runs.stream.call_args.kwargs["command"] == {"resume": resume} + + +@pytest.mark.parametrize("via", ["node", "atask"]) +async def test_subgraph_ancestor_retry_does_not_resend_resume_map_async(via): + async def run(*parts: StreamPart, error: Exception | None = None): + for part in parts: + yield part + if error: + raise error + + remote_client = MagicMock() + remote_client.runs.stream.side_effect = [ + run(*_REMOTE_INTERRUPT_RUN), + run(*_REMOTE_INTERRUPT_RUN, error=ValueError("transient")), + run(*_REMOTE_DONE_RUN), + ] + remote_client.threads.get = AsyncMock(return_value=_INTERRUPTED_THREAD) + parent = _ancestor_retry_parent(remote_client, via) + config = {"configurable": {"thread_id": "parent_thread"}} + await parent.ainvoke({"request": "X"}, config) + + out = await parent.ainvoke(Command(resume={_REMOTE_INTERRUPT_ID: "first"}), config) + assert out["__interrupt__"] == [ + Interrupt(value={"question": "Approve X?"}, id=_REMOTE_INTERRUPT_ID) + ] + + await parent.ainvoke(Command(resume={_REMOTE_INTERRUPT_ID: "second"}), config) + assert [c.kwargs["command"] for c in remote_client.runs.stream.call_args_list] == [ + None, + {"resume": {_REMOTE_INTERRUPT_ID: "first"}}, + {"resume": {_REMOTE_INTERRUPT_ID: "second"}}, + ] + + @pytest.mark.anyio async def test_ainvoke(): # set up test