From be483bb58e70abeef46965a017972916fab37062 Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Wed, 23 Oct 2024 21:55:30 -0700
Subject: [PATCH] Correctly restore ui state when selectedNodeId's parent is
 not in deepestExpandedGroupNodeIds.

PiperOrigin-RevId: 689231805
---
 src/ui/src/components/visualizer/webgl_renderer.ts | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/ui/src/components/visualizer/webgl_renderer.ts b/src/ui/src/components/visualizer/webgl_renderer.ts
index acf381c5..37393390 100644
--- a/src/ui/src/components/visualizer/webgl_renderer.ts
+++ b/src/ui/src/components/visualizer/webgl_renderer.ts
@@ -886,6 +886,20 @@ export class WebglRenderer implements OnInit, OnDestroy {
           );
           deepestExpandedGroupNodeIds = groupNodeIds;
         }
+        // Add the parent node of the selected node if it is not set in
+        // deepestExpandedGroupNodeIds.
+        else {
+          const selectedNode =
+            this.curModelGraph.nodesById[paneState.selectedNodeId];
+          const nsParentId = selectedNode?.nsParentId || '';
+          if (
+            selectedNode &&
+            nsParentId &&
+            !deepestExpandedGroupNodeIds.includes(nsParentId)
+          ) {
+            deepestExpandedGroupNodeIds.push(nsParentId);
+          }
+        }
         if (
           paneState.selectedNodeId !== '' ||
           deepestExpandedGroupNodeIds.length > 0