Fix graph traversing for multi-branch DAGs (GH-8)

ArtyomVancyan · web-flow · commit 6ad87710490a · 2026-02-28T20:03:02.000+04:00
resolves #6
diff --git a/examples/fanout_agent.py b/examples/fanout_agent.py
@@ -0,0 +1,83 @@
+import asyncio
+import operator
+import random
+from typing import Annotated, TypedDict
+
+from langgraph.graph import END, START, StateGraph
+from langgraph.types import Send
+
+from langgraphics import watch
+
+RANDOM_FANOUT = False
+FANOUT_NODES = ["fanout_a", "fanout_b", "fanout_c"]
+FANOUT_CHANCE = 3 / 4
+
+
+class GraphState(TypedDict):
+    initial_result: str
+    fanout_results: Annotated[list[str], operator.add]
+    final_message: str
+
+
+async def initial_node(state: GraphState) -> dict:
+    await asyncio.sleep(2)
+    return {"initial_result": "initial done"}
+
+
+async def fanout_a(state: GraphState) -> dict:
+    await asyncio.sleep(2)
+    return {"fanout_results": ["fanout_a done"]}
+
+
+async def fanout_b(state: GraphState) -> dict:
+    await asyncio.sleep(2)
+    return {"fanout_results": ["fanout_b done"]}
+
+
+async def fanout_c(state: GraphState) -> dict:
+    await asyncio.sleep(2)
+    return {"fanout_results": ["fanout_c done"]}
+
+
+async def final_node(state: GraphState) -> dict:
+    await asyncio.sleep(2)
+    return {"final_message": "\n".join(state["fanout_results"])}
+
+
+def route_fanout(state: GraphState) -> list[Send]:
+    chosen = [n for n in FANOUT_NODES if random.random() < FANOUT_CHANCE] or [random.choice(FANOUT_NODES)]
+    return [Send(name, state) for name in chosen]
+
+
+builder = StateGraph(GraphState)
+
+builder.add_node("initial", initial_node)
+builder.add_node("fanout_a", fanout_a)
+builder.add_node("fanout_b", fanout_b)
+builder.add_node("fanout_c", fanout_c)
+builder.add_node("final", final_node)
+
+builder.add_edge(START, "initial")
+
+if RANDOM_FANOUT:
+    builder.add_conditional_edges("initial", route_fanout, {n: n for n in FANOUT_NODES})
+else:
+    builder.add_edge("initial", "fanout_a")
+    builder.add_edge("initial", "fanout_b")
+    builder.add_edge("initial", "fanout_c")
+
+builder.add_edge("fanout_a", "final")
+builder.add_edge("fanout_b", "final")
+builder.add_edge("fanout_c", "final")
+builder.add_edge("final", END)
+
+graph = builder.compile()
+graph = watch(graph)
+
+
+async def main() -> None:
+    await graph.ainvoke({"initial_result": "", "fanout_results": [], "final_message": ""})
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/langgraphics-web/src/components/GraphCanvas.tsx b/langgraphics-web/src/components/GraphCanvas.tsx
@@ -12,7 +12,7 @@ interface GraphCanvasProps {
     nodes: Node<NodeData>[];
     edges: Edge<EdgeData>[];
     events: ExecutionEvent[];
-    activeNodeId: string | null;
+    activeNodeIds: string[];
     inspect: ReactNode;
     initialMode: ViewMode;
     initialRankDir?: RankDir;
@@ -21,11 +21,11 @@ interface GraphCanvasProps {
     onRankDirChange?: (v: RankDir) => void;
 }
 
-export function GraphCanvas({nodes, edges, events, activeNodeId, inspect, initialMode = "auto", initialInspect = "off", initialColorMode = "system", initialRankDir = "TB", onRankDirChange}: GraphCanvasProps) {
+export function GraphCanvas({nodes, edges, events, activeNodeIds, inspect, initialMode = "auto", initialInspect = "off", initialColorMode = "system", initialRankDir = "TB", onRankDirChange}: GraphCanvasProps) {
     const [rankDir, setRankDir] = useState<RankDir>(initialRankDir);
     const [colorMode, setColorMode] = useState<ColorMode>(initialColorMode);
     const [inspectorMode, setInspectorMode] = useState<InspectorMode>(initialInspect);
-    const {isManual, goAuto, goManual, fitContent} = useFocus({nodes, edges, activeNodeId, rankDir, initialMode});
+    const {isManual, goAuto, goManual, fitContent} = useFocus({nodes, edges, activeNodeIds, rankDir, initialMode});
 
     const handleRankDirChange = useCallback(async (v: RankDir) => {
         setRankDir(v);
diff --git a/langgraphics-web/src/hooks/useFocus.ts b/langgraphics-web/src/hooks/useFocus.ts
@@ -7,7 +7,7 @@ import type {RankDir} from "../layout";
 interface UseFocusOptions {
     nodes: Node<NodeData>[];
     edges: Edge<EdgeData>[];
-    activeNodeId: string | null;
+    activeNodeIds: string[];
     rankDir?: RankDir;
     initialMode?: ViewMode;
 }
@@ -37,12 +37,12 @@ function getNeighbourIds(nodeId: string, nodes: Node<NodeData>[], edges: Edge<Ed
     return [before?.id, after?.id].filter((id): id is string => id !== undefined).map((id) => ({id}));
 }
 
-export function useFocus({nodes, edges, activeNodeId, rankDir = "TB", initialMode = "auto"}: UseFocusOptions) {
+export function useFocus({nodes, edges, activeNodeIds, rankDir = "TB", initialMode = "auto"}: UseFocusOptions) {
     const {fitView} = useReactFlow();
     const [mode, setMode] = useState<"auto" | "manual">(initialMode);
     const prevMode = useRef<"auto" | "manual">(mode);
     const initialDone = useRef(false);
-    const prevFocusId = useRef<string | null>(null);
+    const prevFocusKey = useRef<string>("");
 
     const isManual = useMemo(() => mode === "manual", [mode]);
     const isHorizontal = useMemo(() => ["LR", "RL"].includes(rankDir), [rankDir]);
@@ -72,28 +72,34 @@ export function useFocus({nodes, edges, activeNodeId, rankDir = "TB", initialMod
                     duration: 0,
                 }).then();
             }
-            prevFocusId.current = null;
+            prevFocusKey.current = "";
             prevMode.current = mode;
             return;
         }
 
-        if (mode === "auto" && prevMode.current !== "auto") prevFocusId.current = null;
+        if (mode === "auto" && prevMode.current !== "auto") prevFocusKey.current = "";
         prevMode.current = mode;
 
         if (mode !== "auto") return;
 
-        if (activeNodeId && activeNodeId !== prevFocusId.current) {
-            prevFocusId.current = activeNodeId;
-
-            const activeNode = nodes.find((n) => n.id === activeNodeId);
-            if (activeNode?.data.nodeType === "node") {
+        const focusKey = [...activeNodeIds].sort().join(",");
+        if (activeNodeIds.length > 0 && focusKey !== prevFocusKey.current) {
+            prevFocusKey.current = focusKey;
+
+            const activeNodes = nodes.filter(
+                (n) => activeNodeIds.includes(n.id) && n.data.nodeType === "node",
+            );
+            if (activeNodes.length > 0) {
+                const neighbours = activeNodes.flatMap((n) =>
+                    getNeighbourIds(n.id, nodes, edges, isHorizontal),
+                );
                 fitView({
-                    nodes: [{id: activeNodeId}, ...getNeighbourIds(activeNodeId, nodes, edges, isHorizontal)],
+                    nodes: [...activeNodes.map((n) => ({id: n.id})), ...neighbours],
                     duration: FIT_VIEW_DURATION,
                 }).then();
             }
         }
-    }, [nodes, edges, activeNodeId, fitView, mode, isHorizontal]);
+    }, [nodes, edges, activeNodeIds, fitView, mode, isHorizontal]);
 
     return {isManual, goAuto, goManual, fitContent};
 }
diff --git a/langgraphics-web/src/hooks/useGraphState.ts b/langgraphics-web/src/hooks/useGraphState.ts
@@ -9,20 +9,27 @@ export function computeStatuses(events: ExecutionEvent[]): {
 } {
     const nodeStatuses = new Map<string, NodeStatus>();
     const edgeStatuses = new Map<string, EdgeStatus>();
+    const edgeInfo = new Map<string, {source: string; target: string}>();
 
     for (const event of events) {
         if (event.type === "run_start") {
             nodeStatuses.clear();
             edgeStatuses.clear();
+            edgeInfo.clear();
         } else if (event.type === "edge_active") {
-            for (const [id, status] of edgeStatuses) {
-                if (status === "active") edgeStatuses.set(id, "traversed");
+            edgeInfo.set(event.edge_id, {source: event.source, target: event.target});
+            if (nodeStatuses.get(event.source) === "active") {
+                nodeStatuses.set(event.source, "completed");
             }
-            for (const [id, status] of nodeStatuses) {
-                if (status === "active") nodeStatuses.set(id, "completed");
+            for (const [id, info] of edgeInfo) {
+                if (info.target === event.source && edgeStatuses.get(id) === "active") {
+                    edgeStatuses.set(id, "traversed");
+                }
             }
             edgeStatuses.set(event.edge_id, "active");
-            nodeStatuses.set(event.target, "active");
+            if (nodeStatuses.get(event.target) !== "error") {
+                nodeStatuses.set(event.target, "active");
+            }
         } else if (event.type === "error") {
             for (const [id, status] of edgeStatuses) {
                 if (status === "active") edgeStatuses.set(id, "traversed");
@@ -52,14 +59,14 @@ export function useGraphState(topology: GraphMessage | null, events: ExecutionEv
     }, [topology, rankDir]);
 
     return useMemo(() => {
-        if (events.length === 0) return {nodes: base.nodes, edges: base.edges, activeNodeId: null as string | null};
+        if (events.length === 0) return {nodes: base.nodes, edges: base.edges, activeNodeIds: [] as string[]};
 
         const {nodeStatuses, edgeStatuses} = computeStatuses(events);
 
-        let activeNodeId: string | null = null;
+        const activeNodeIds: string[] = [];
         const nodes = base.nodes.map((node) => {
             const status = nodeStatuses.get(node.id);
-            if (status === "active") activeNodeId = node.id;
+            if (status === "active") activeNodeIds.push(node.id);
             return {...node, className: status};
         });
 
@@ -81,6 +88,6 @@ export function useGraphState(topology: GraphMessage | null, events: ExecutionEv
             return {...edge, className, animated: status === "active", markerEnd};
         });
 
-        return {nodes, edges, activeNodeId};
+        return {nodes, edges, activeNodeIds};
     }, [base, events]);
 }
diff --git a/langgraphics-web/src/main.tsx b/langgraphics-web/src/main.tsx
@@ -32,7 +32,7 @@ const {theme, mode, inspect, direction} = parseParams();
 function Index() {
     const [rankDir, setRankDir] = useState<RankDir>(direction);
     const {topology, events, nodeEntries} = useWebSocket(WS_URL);
-    const {nodes, edges, activeNodeId} = useGraphState(topology, events, rankDir);
+    const {nodes, edges, activeNodeIds} = useGraphState(topology, events, rankDir);
 
     return (
         <ReactFlowProvider>
@@ -44,7 +44,7 @@ function Index() {
                 initialInspect={inspect}
                 initialColorMode={theme}
                 initialRankDir={direction}
-                activeNodeId={activeNodeId}
+                activeNodeIds={activeNodeIds}
                 onRankDirChange={setRankDir}
                 inspect={<InspectPanel nodeEntries={nodeEntries}/>}
             />
diff --git a/langgraphics/streamer.py b/langgraphics/streamer.py
@@ -96,10 +96,14 @@ def __init__(
         self.edge_lookup = edge_lookup
         self.http_server = http_server
         self.node_names: set[str] = set()
+        self.predecessors: dict[str, set[str]] = {}
         for src, tgt in edge_lookup:
             self.node_names.add(src)
             self.node_names.add(tgt)
+            self.predecessors.setdefault(tgt, set()).add(src)
         self.node_names -= {"__start__", "__end__"}
+        self.generation: dict[str, int] = {"__start__": 0}
+        self.linked: set[tuple[str, int, str]] = set()
 
     def __getattr__(self, name: str) -> Any:
         return getattr(self.graph, name)
@@ -118,17 +122,23 @@ async def broadcast(self, message: dict[str, Any]) -> None:
         except Exception:
             pass
 
-    async def _emit_edge(self, source: str, target: str) -> None:
-        edge_id = self.edge_lookup.get((source, target))
-        if edge_id:
-            await self.broadcast(
-                {
-                    "type": "edge_active",
-                    "source": source,
-                    "target": target,
-                    "edge_id": edge_id,
-                }
-            )
+    async def _emit_edge(self, target: str) -> None:
+        for source in self.predecessors.get(target, set()):
+            if src_gen := self.generation.get(source) is None:
+                continue
+            if (key := (source, src_gen, target)) in self.linked:
+                continue
+            self.linked.add(key)
+            if edge_id := self.edge_lookup.get((source, target)):
+                await self.broadcast(
+                    {
+                        "type": "edge_active",
+                        "source": source,
+                        "target": target,
+                        "edge_id": edge_id,
+                    }
+                )
+        self.generation[target] = self.generation.get(target, -1) + 1
 
     async def _emit_error(self, last_node: str) -> None:
         for (src, tgt), eid in self.edge_lookup.items():
@@ -166,14 +176,14 @@ async def ainvoke(self, input: Any, config: Any = None, **kwargs: Any) -> Any:
                 input, config=merged_config, stream_mode="updates", **kwargs
             ):
                 if isinstance(chunk, dict):
-                    for node_name in chunk:
+                    for node_name, node_result in chunk.items():
                         if node_name == "__metadata__":
                             continue
-                        await self._emit_edge(last_node, node_name)
+                        await self._emit_edge(node_name)
                         last_node = node_name
-                        result = chunk[node_name]
+                        result = node_result
 
-            await self._emit_edge(last_node, "__end__")
+            await self._emit_edge("__end__")
             await self.broadcast({"type": "run_end", "run_id": run_id})
         except Exception:
             await self._emit_error(last_node)
@@ -204,12 +214,12 @@ async def astream(
                     for node_name in chunk:
                         if node_name == "__metadata__":
                             continue
-                        await self._emit_edge(last_node, node_name)
+                        await self._emit_edge(node_name)
                         last_node = node_name
                 yield chunk
 
-            if last_node != "__start__":
-                await self._emit_edge(last_node, "__end__")
+            if len(self.generation) > 1:
+                await self._emit_edge("__end__")
 
             await self.broadcast({"type": "run_end", "run_id": run_id})
         except Exception:
