fix(langgraph): fix graph rendering for defer=True (#6130)

### 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`)
| -------- | ------- | ------- |
| <img height="400" alt="defer_after"
src="https://github.com/user-attachments/assets/0a9fc992-1b6a-4c6d-8752-de54c703c329"
/> | <img height="400" alt="defer_before"
src="https://github.com/user-attachments/assets/825b09fc-3fb8-461a-9928-20c8d9cfc533"
/> | <img height="400" alt="defer_after"
src="https://github.com/user-attachments/assets/ce5334f7-b469-47b0-8f1e-35bda2544a4e"
/> |

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`)
| -------- | ------- | ------- |
| <img height="400" alt="defer_after"
src="https://github.com/user-attachments/assets/3509d25c-f3ad-473c-b877-c155b8008cd5"
/> | <img height="400" alt="defer_before"
src="https://github.com/user-attachments/assets/7af38e77-eb70-414d-b8fe-667da943f9e0"
/> | <img height="400" alt="defer_after"
src="https://github.com/user-attachments/assets/bc87a19f-b4fb-42d3-a6ee-5b0982d9af71"
/> |

### https://github.com/langchain-ai/langgraph/issues/5369

| No defer | Before (defer `595577`, `52642`) | After (defer `595577`,
`52642`)
| -------- | ------- | ------- |
| <img height="400" alt="defer_after"
src="https://github.com/user-attachments/assets/7c0824ce-3921-4dce-bc16-278f64289d28"
/> | <img height="400" alt="defer_before"
src="https://github.com/user-attachments/assets/28661079-7502-4912-874b-c086c0204a87"
/> | <img height="400" alt="defer_after"
src="https://github.com/user-attachments/assets/2a04956a-ed79-40e5-98b8-f6ecb2597a2e"
/> |
This commit is contained in:
Caspar Broekhuizen
2025-09-23 12:47:50 -07:00
committed by GitHub
parent 9f969f5fe1
commit affaa90d2a
2 changed files with 142 additions and 110 deletions
+49 -17
View File
@@ -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:
@@ -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;