From affaa90d2a7f1174251a7a4a16b04da26ee6834a Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Tue, 23 Sep 2025 12:47:50 -0700 Subject: [PATCH] fix(langgraph): fix graph rendering for defer=True (#6130) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description Some graphs with `defer=True` nodes rendered incorrectly. E.g.: * edge C2 -> E1 is missing and edge C2 -> END should not appear in #5772 * edge E3 -> END is missing and edge E -> END should not appear in #5182 * extra edge #5369 Fix: * Record the destinations declared by get_static_writes for each node. Build step_sources as a union of the runtime writes and the static writes (instead of just runtime writes). * Label deferred nodes with 'deferred' ### https://github.com/langchain-ai/langgraph/issues/5772 'Before' is how they were rendered before this PR | No defer | Before (defer `E1`) | After (defer `E1`) | -------- | ------- | ------- | | defer_after | defer_before | defer_after | Before: * For deferred joins (NamedBarrierValueAfterFinish), a writer from an upstream node may not produce a runtime task.writes entry until the barrier opens. draw_graph() builds edges from task.writes, so one side of the join (here C2) never gets recorded as a source, and C2 is seen as a sink, so there is an implicit edge: C2 -> END edge added. After: * C2's write to the join channel is recorded even if the barrier hasn’t opened. When E1 finally schedules, we correctly find both sources B2 and C2 for the same trigger and emit edges: B2 -> E1 and C2 -> E1. With C2 -> E1 present, C2 is no longer a terminus, so the unexpected edge: C2 -> END is not added. ### Other graphs Graphs for the most part remain unchanged. See: ### #5182 | No defer | Before (defer `d`) | After (defer `d`) | -------- | ------- | ------- | | defer_after | defer_before | defer_after | ### https://github.com/langchain-ai/langgraph/issues/5369 | No defer | Before (defer `595577`, `52642`) | After (defer `595577`, `52642`) | -------- | ------- | ------- | | defer_after | defer_before | defer_after | --- libs/langgraph/langgraph/pregel/_draw.py | 66 +++++-- .../tests/__snapshots__/test_pregel.ambr | 186 +++++++++--------- 2 files changed, 142 insertions(+), 110 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_draw.py b/libs/langgraph/langgraph/pregel/_draw.py index a6d33d4a9..922450df9 100644 --- a/libs/langgraph/langgraph/pregel/_draw.py +++ b/libs/langgraph/langgraph/pregel/_draw.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections import defaultdict from collections.abc import Mapping, Sequence -from typing import Any, cast +from typing import Any, NamedTuple, cast from langchain_core.runnables.config import RunnableConfig from langchain_core.runnables.graph import Graph, Node @@ -10,6 +10,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, INPUT from langgraph.channels.base import BaseChannel +from langgraph.channels.last_value import LastValueAfterFinish from langgraph.constants import END, START from langgraph.managed.base import ManagedValueSpec from langgraph.pregel._algo import ( @@ -25,6 +26,19 @@ from langgraph.pregel._write import ChannelWrite from langgraph.types import All, Checkpointer +class Edge(NamedTuple): + source: str + target: str + conditional: bool + data: str | None + + +class TriggerEdge(NamedTuple): + source: str + conditional: bool + data: str | None + + def draw_graph( config: RunnableConfig, *, @@ -49,7 +63,7 @@ def draw_graph( The graph for this Pregel instance. """ # (src, dest, is_conditional, label) - edges: set[tuple[str, str, bool, str | None]] = set() + edges: set[Edge] = set() step = -1 checkpoint = empty_checkpoint() @@ -63,8 +77,9 @@ def draw_graph( checkpoint, ) static_seen: set[Any] = set() - sources: dict[str, set[tuple[str, bool, str | None]]] = {} - step_sources: dict[str, set[tuple[str, bool, str | None]]] = {} + sources: dict[str, set[TriggerEdge]] = {} + step_sources: dict[str, set[TriggerEdge]] = {} + static_declared_writes: dict[str, set[TriggerEdge]] = defaultdict(set) # remove node mappers nodes = { k: v.copy(update={"mapper": None}) if v.mapper is not None else v @@ -123,32 +138,36 @@ def draw_graph( # END writes are not written, but become edges directly for t in writes: if t[0] == END: - edges.add((task.name, t[0], True, t[2])) + edges.add(Edge(task.name, t[0], True, t[2])) writes = [t for t in writes if t[0] != END] conditionals.update( {(task.name, t[0], t[1] or None): t[2] for t in writes} ) + # record static writes for edge creation + for t in writes: + static_declared_writes[task.name].add( + TriggerEdge(t[0], True, t[2]) + ) task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes]) # collect sources - step_sources = { - task.name: { - ( + step_sources = {} + for task in tasks.values(): + task_edges = { + TriggerEdge( w[0], (task.name, w[0], w[1] or None) in conditionals, conditionals.get((task.name, w[0], w[1] or None)), ) for w in task.writes } - for task in tasks.values() - } + task_edges |= static_declared_writes.get(task.name, set()) + step_sources[task.name] = task_edges sources.update(step_sources) # invert triggers - trigger_to_sources: dict[str, set[tuple[str, bool, str | None]]] = defaultdict( - set - ) + trigger_to_sources: dict[str, set[TriggerEdge]] = defaultdict(set) for src, triggers in sources.items(): for trigger, cond, label in triggers: - trigger_to_sources[trigger].add((src, cond, label)) + trigger_to_sources[trigger].add(TriggerEdge(src, cond, label)) # apply writes updated_channels = apply_writes( checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes @@ -170,26 +189,39 @@ def draw_graph( trigger_to_nodes=trigger_to_nodes, updated_channels=updated_channels, ) + # collect deferred nodes + deferred_nodes: set[str] = set() + edges_to_deferred_nodes: set[Edge] = set() + for channel, item in channels.items(): + if isinstance(item, LastValueAfterFinish): + deferred_node = channel.split(":", 2)[-1] + deferred_nodes.add(deferred_node) # collect edges for task in tasks.values(): added = False for trigger in task.triggers: for src, cond, label in sorted(trigger_to_sources[trigger]): - edges.add((src, task.name, cond, label)) + # record edge to be reviewed later + if task.name in deferred_nodes: + edges_to_deferred_nodes.add(Edge(src, task.name, cond, label)) + edges.add(Edge(src, task.name, cond, label)) # if the edge is from this step, skip adding the implicit edges if (trigger, cond, label) in step_sources.get(src, set()): added = True else: - sources[src].discard((trigger, cond, label)) + sources[src].discard(TriggerEdge(trigger, cond, label)) # if no edges from this step, add implicit edges from all previous tasks if not added: for src in step_sources: - edges.add((src, task.name, True, None)) + edges.add(Edge(src, task.name, True, None)) + # assemble the graph graph = Graph() # add nodes for name, node in nodes.items(): metadata = dict(node.metadata or {}) + if name in deferred_nodes: + metadata["defer"] = True if name in interrupt_before_nodes and name in interrupt_after_nodes: metadata["__interrupt"] = "before,after" elif name in interrupt_before_nodes: diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index b7dfbcd0e..d3808a63a 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -307,6 +307,99 @@ ''' # --- +# name: test_get_graph_nonterminal_last_step_source + ''' + { + "edges": [ + { + "source": "__start__", + "target": "human" + }, + { + "conditional": true, + "source": "chatbot", + "target": "human" + }, + { + "conditional": true, + "source": "chatbot", + "target": "tools" + }, + { + "conditional": true, + "source": "human", + "target": "__end__" + }, + { + "conditional": true, + "source": "human", + "target": "chatbot" + }, + { + "source": "tools", + "target": "chatbot" + } + ], + "nodes": [ + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "__start__" + }, + "id": "__start__", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "chatbot" + }, + "id": "chatbot", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "id": "tools", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "human" + }, + "id": "human", + "type": "runnable" + }, + { + "id": "__end__" + } + ] + } + ''' +# --- # name: test_get_graph_root_channel ''' { @@ -795,99 +888,6 @@ ''' # --- -# name: test_get_graph_nonterminal_last_step_source - ''' - { - "edges": [ - { - "source": "__start__", - "target": "human" - }, - { - "conditional": true, - "source": "chatbot", - "target": "human" - }, - { - "conditional": true, - "source": "chatbot", - "target": "tools" - }, - { - "conditional": true, - "source": "human", - "target": "__end__" - }, - { - "conditional": true, - "source": "human", - "target": "chatbot" - }, - { - "source": "tools", - "target": "chatbot" - } - ], - "nodes": [ - { - "data": { - "id": [ - "langgraph", - "_internal", - "_runnable", - "RunnableCallable" - ], - "name": "__start__" - }, - "id": "__start__", - "type": "runnable" - }, - { - "data": { - "id": [ - "langgraph", - "_internal", - "_runnable", - "RunnableCallable" - ], - "name": "chatbot" - }, - "id": "chatbot", - "type": "runnable" - }, - { - "data": { - "id": [ - "langgraph", - "_internal", - "_runnable", - "RunnableCallable" - ], - "name": "tools" - }, - "id": "tools", - "type": "runnable" - }, - { - "data": { - "id": [ - "langgraph", - "_internal", - "_runnable", - "RunnableCallable" - ], - "name": "human" - }, - "id": "human", - "type": "runnable" - }, - { - "id": "__end__" - } - ] - } - ''' -# --- # name: test_repeat_condition ''' graph TD;