From 409fe4e6668512577048f44456cbb636913b4764 Mon Sep 17 00:00:00 2001
From: Yin Zhang <yinzz@google.com>
Date: Thu, 14 Nov 2024 14:44:29 -0800
Subject: [PATCH] Fix seenGroupNodeIds mismatch with modelGraph nodes data.

looks like the `seenGroupNodeIds` data set may contain group node id that does not belong to the current `modelGraph` nodes data, causing the lookup for node from modelGraph returns undefined (NodeA/NodeB), thus the group layer comparison fails.

PiperOrigin-RevId: 696656783
---
 src/ui/src/components/visualizer/worker/graph_expander.ts | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/ui/src/components/visualizer/worker/graph_expander.ts b/src/ui/src/components/visualizer/worker/graph_expander.ts
index 2fdf0c73..4ecd3fd2 100644
--- a/src/ui/src/components/visualizer/worker/graph_expander.ts
+++ b/src/ui/src/components/visualizer/worker/graph_expander.ts
@@ -129,7 +129,10 @@ export class GraphExpander {
     const queue: string[] = [...groupNodeIds];
     while (queue.length > 0) {
       const curGroupNodeId = queue.shift()!;
-      if (seenGroupNodeIds.has(curGroupNodeId)) {
+      if (
+        seenGroupNodeIds.has(curGroupNodeId) ||
+        !this.modelGraph.nodesById[curGroupNodeId]
+      ) {
         continue;
       }
       seenGroupNodeIds.add(curGroupNodeId);