From 21874a7d2909f774071175b1904af60668b46057 Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Thu, 29 Aug 2024 15:58:00 -0700
Subject: [PATCH] Add support for reusing running server.

The frontend subscribes to a server-sent-event stream and will refresh the page
when instructed.

PiperOrigin-RevId: 669092712
---
 .../visualizer/model_graph_visualizer.ts      |  6 ++
 .../visualizer/server_director_service.ts     | 67 +++++++++++++++++++
 2 files changed, 73 insertions(+)
 create mode 100644 src/ui/src/components/visualizer/server_director_service.ts

diff --git a/src/ui/src/components/visualizer/model_graph_visualizer.ts b/src/ui/src/components/visualizer/model_graph_visualizer.ts
index 93fe6231..a62a186d 100644
--- a/src/ui/src/components/visualizer/model_graph_visualizer.ts
+++ b/src/ui/src/components/visualizer/model_graph_visualizer.ts
@@ -51,6 +51,7 @@ import {type VisualizerUiState} from './common/visualizer_ui_state';
 import {ExtensionService} from './extension_service';
 import {NodeDataProviderExtensionService} from './node_data_provider_extension_service';
 import {NodeStylerService} from './node_styler_service';
+import {ServerDirectorService} from './server_director_service';
 import {SplitPanesContainer} from './split_panes_container';
 import {ThreejsService} from './threejs_service';
 import {TitleBar} from './title_bar';
@@ -69,6 +70,7 @@ import {WorkerService} from './worker_service';
     ExtensionService,
     NodeDataProviderExtensionService,
     NodeStylerService,
+    ServerDirectorService,
     UiStateService,
     WorkerService,
   ],
@@ -130,7 +132,11 @@ export class ModelGraphVisualizer implements OnInit, OnDestroy, OnChanges {
     private readonly uiStateService: UiStateService,
     private readonly nodeDataProviderExtensionService: NodeDataProviderExtensionService,
     private readonly nodeStylerService: NodeStylerService,
+    private readonly serverDirectorService: ServerDirectorService,
   ) {
+
+    this.serverDirectorService.init();
+
     effect(() => {
       const curUiState = this.uiStateService.curUiState();
       if (!curUiState) {
diff --git a/src/ui/src/components/visualizer/server_director_service.ts b/src/ui/src/components/visualizer/server_director_service.ts
new file mode 100644
index 00000000..906daa26
--- /dev/null
+++ b/src/ui/src/components/visualizer/server_director_service.ts
@@ -0,0 +1,67 @@
+/**
+ * @license
+ * Copyright 2024 The Model Explorer Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ==============================================================================
+ */
+
+import {safeLocation} from 'safevalues/dom';
+import {IS_EXTERNAL} from '../../common/flags';
+
+import {Injectable} from '@angular/core';
+
+enum DirectiveName {
+  RefreshPage = 'refreshPage',
+}
+
+declare interface DirectiveBase {
+  name: DirectiveName;
+}
+
+declare interface RefreshPageDirective extends DirectiveBase {
+  name: DirectiveName.RefreshPage;
+  url: string;
+}
+
+type Directive = RefreshPageDirective;
+
+/** A service for handling directives streaming from ME server. */
+@Injectable()
+export class ServerDirectorService {
+  init() {
+    if (IS_EXTERNAL) {
+      // Listen to the streaming events (directives) from the following source
+      // that the server has established.
+      const eventSource = new EventSource('/apistream/server_director');
+      eventSource.addEventListener('message', (e) => {
+        if (!e.data) {
+          return;
+        }
+        const directive = JSON.parse(e.data) as Directive;
+        switch (directive.name) {
+          // Refresh page with the given url.
+          case DirectiveName.RefreshPage:
+            safeLocation.setHref(
+              window.location,
+              directive.url,
+            );
+
+            break;
+          default:
+            break;
+        }
+      });
+    }
+  }
+}