WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-20 15:01:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(langgraph): forward parent resume to RemoteGraph subgraph nodes

Browse Source 
When a RemoteGraph runs as a subgraph node and the parent is resumed with
Command(resume=...), forward the resume value to the remote thread as a
command instead of re-sending the parent state as new input. Also surface
every interrupt from the remote run so they can be resumed by id.

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Dariel Dato-on
2026-10-08 21:17:38 +00:00
co-authored by open-swe[bot] <open-swe@users.noreply.github.com>
parent 93f5eaff21
commit 6be0ddf9ad
2 changed files with 187 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/remote.py
+59 -22
View File
@@ -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,
libs/langgraph/tests/test_remote_graph.py
+128 -1
View File
@@ -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