From e0d616a235e0b262d25b89c175c277ea20e23162 Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Mon, 26 Aug 2019 18:14:33 -0700
Subject: [PATCH 1/6] EOD Push

Most work on feature is completed. Currently attempting to fix bug where UI does not properly render visualizations for selected component.
---
 .../viewers/VisualizationCreator.tsx          |  10 +-
 frontend/src/lib/Apis.ts                      |  23 +-
 frontend/src/pages/RunDetails.test.tsx        |   4 +-
 frontend/src/pages/RunDetails.tsx             |  92 ++++++-
 .../__snapshots__/RunDetails.test.tsx.snap    | 247 ++++++++++++++++++
 5 files changed, 371 insertions(+), 5 deletions(-)

diff --git a/frontend/src/components/viewers/VisualizationCreator.tsx b/frontend/src/components/viewers/VisualizationCreator.tsx
index 52aa2eada25..c21d160ce02 100644
--- a/frontend/src/components/viewers/VisualizationCreator.tsx
+++ b/frontend/src/components/viewers/VisualizationCreator.tsx
@@ -151,7 +151,15 @@ class VisualizationCreator extends Viewer<VisualizationCreatorProps, Visualizati
       <BusyButton title='Generate Visualization' busy={isBusy} disabled={!canGenerate}
         onClick={() => {
           if (onGenerate && selectedType) {
-            onGenerate(_arguments, source, selectedType);
+            const specifiedArguments: any = JSON.parse(_arguments || '{}');
+            if (selectedType === ApiVisualizationType.CUSTOM) {
+              specifiedArguments.code = code.split('\n');
+            }
+            onGenerate(
+              JSON.stringify(specifiedArguments),
+              source,
+              selectedType
+            );
           }
         }} />
     </div>;
diff --git a/frontend/src/lib/Apis.ts b/frontend/src/lib/Apis.ts
index 70f9b710c80..46b71a2a95e 100644
--- a/frontend/src/lib/Apis.ts
+++ b/frontend/src/lib/Apis.ts
@@ -18,7 +18,9 @@ import { JobServiceApi } from '../apis/job';
 import { RunServiceApi } from '../apis/run';
 import { PipelineServiceApi, ApiPipeline } from '../apis/pipeline';
 import { StoragePath } from './WorkflowParser';
-import { VisualizationServiceApi } from '../apis/visualization';
+import { VisualizationServiceApi, ApiVisualization } from '../apis/visualization';
+import { HTMLViewerConfig } from 'src/components/viewers/HTMLViewer';
+import { PlotType } from '../components/viewers/Viewer';
 
 const v1beta1Prefix = 'apis/v1beta1';
 
@@ -52,6 +54,25 @@ export class Apis {
     return customVisualizationsAllowed;
   }
 
+  public static async buildPythonVisualizationConfig(visualizationData: ApiVisualization): Promise<HTMLViewerConfig> {
+    const visualization = await Apis.visualizationServiceApi.createVisualization(visualizationData);
+    if (visualization.html) {
+      const htmlContent = visualization.html
+        // Fixes issue with TFX components (and other iframe based
+        // visualizations), where the method in which javascript interacts
+        // with embedded iframes is not allowed when embedded in an additional
+        // iframe. This is resolved by setting the srcdoc value rather that
+        // manipulating the document directly.
+        .replace('contentWindow.document.write', 'srcdoc=');
+      return {
+        htmlContent,
+        type: PlotType.WEB_APP,
+      } as HTMLViewerConfig;
+    } else {
+      throw new Error('Unable to generate visualization!');
+    }
+  }
+
   /**
    * Get pod logs
    */
diff --git a/frontend/src/pages/RunDetails.test.tsx b/frontend/src/pages/RunDetails.test.tsx
index 96a183168a9..4be6e41f199 100644
--- a/frontend/src/pages/RunDetails.test.tsx
+++ b/frontend/src/pages/RunDetails.test.tsx
@@ -40,12 +40,13 @@ describe('RunDetails', () => {
   const getClusterNameSpy = jest.spyOn(Apis, 'getClusterName');
   const getRunSpy = jest.spyOn(Apis.runServiceApi, 'getRun');
   const getExperimentSpy = jest.spyOn(Apis.experimentServiceApi, 'getExperiment');
+  const isCustomVisualizationsAllowedSpy = jest.spyOn(Apis, 'areCustomVisualizationsAllowed');
   const getPodLogsSpy = jest.spyOn(Apis, 'getPodLogs');
   const pathsParser = jest.spyOn(WorkflowParser, 'loadNodeOutputPaths');
   const pathsWithStepsParser = jest.spyOn(WorkflowParser, 'loadAllOutputPathsWithStepNames');
   const loaderSpy = jest.spyOn(OutputArtifactLoader, 'load');
   // We mock this because it uses toLocaleDateString, which causes mismatches between local and CI
-  // test enviroments
+  // test environments
   const formatDateStringSpy = jest.spyOn(Utils, 'formatDateString');
 
   let testRun: ApiRunDetail = {};
@@ -93,6 +94,7 @@ describe('RunDetails', () => {
     getClusterNameSpy.mockImplementation(() => Promise.resolve('some-cluster'));
     getRunSpy.mockImplementation(() => Promise.resolve(testRun));
     getExperimentSpy.mockImplementation(() => Promise.resolve({ id: 'some-experiment-id', name: 'some experiment' }));
+    isCustomVisualizationsAllowedSpy.mockImplementation(() => Promise.resolve(false));
     getPodLogsSpy.mockImplementation(() => 'test logs');
     pathsParser.mockImplementation(() => []);
     pathsWithStepsParser.mockImplementation(() => []);
diff --git a/frontend/src/pages/RunDetails.tsx b/frontend/src/pages/RunDetails.tsx
index de1a4d1687f..f24afdf7967 100644
--- a/frontend/src/pages/RunDetails.tsx
+++ b/frontend/src/pages/RunDetails.tsx
@@ -39,7 +39,7 @@ import { OutputArtifactLoader } from '../lib/OutputArtifactLoader';
 import { Page } from './Page';
 import { RoutePage, RouteParams } from '../components/Router';
 import { ToolbarProps } from '../components/Toolbar';
-import { ViewerConfig } from '../components/viewers/Viewer';
+import { ViewerConfig, PlotType } from '../components/viewers/Viewer';
 import { Workflow } from '../../third_party/argo-ui/argo_template';
 import { classes, stylesheet } from 'typestyle';
 import { commonCss, padding, color, fonts, fontsize } from '../Css';
@@ -47,6 +47,9 @@ import { componentMap } from '../components/viewers/ViewerContainer';
 import { flatten } from 'lodash';
 import { formatDateString, getRunDurationFromWorkflow, logger, errorToMessage } from '../lib/Utils';
 import { statusToIcon } from './Status';
+import VisualizationCreator, { VisualizationCreatorConfig } from '../components/viewers/VisualizationCreator';
+import { ApiVisualization, ApiVisualizationType } from '../apis/visualization';
+import { HTMLViewerConfig } from '../components/viewers/HTMLViewer';
 
 enum SidePaneTab {
   ARTIFACTS,
@@ -72,9 +75,17 @@ interface AnnotatedConfig {
   stepName: string;
 }
 
+interface GeneratedVisualization {
+  config: HTMLViewerConfig;
+  nodeId: string;
+}
+
 interface RunDetailsState {
   allArtifactConfigs: AnnotatedConfig[];
+  allowCustomVisualizations: boolean;
   experiment?: ApiExperiment;
+  generatedVisualizations: GeneratedVisualization[];
+  isGeneratingVisualization: boolean;
   legacyStackdriverUrl: string;
   logsBannerAdditionalInfo: string;
   logsBannerMessage: string;
@@ -135,6 +146,9 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
 
     this.state = {
       allArtifactConfigs: [],
+      allowCustomVisualizations: false,
+      generatedVisualizations: [],
+      isGeneratingVisualization: false,
       legacyStackdriverUrl: '',
       logsBannerAdditionalInfo: '',
       logsBannerMessage: '',
@@ -170,7 +184,9 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
   public render(): JSX.Element {
     const {
       allArtifactConfigs,
+      allowCustomVisualizations,
       graph,
+      isGeneratingVisualization,
       legacyStackdriverUrl,
       runFinished,
       runMetadata,
@@ -185,6 +201,14 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     const workflowParameters = WorkflowParser.getParameters(workflow);
     const nodeInputOutputParams = WorkflowParser.getNodeInputOutputParams(workflow, selectedNodeId);
     const hasMetrics = runMetadata && runMetadata.metrics && runMetadata.metrics.length > 0;
+    const visualizationCreatorConfig: VisualizationCreatorConfig = {
+      allowCustomVisualizations,
+      isBusy: isGeneratingVisualization,
+      onGenerate: (visualizationArguments: string, source: string, type: ApiVisualizationType) => {
+        this._onGenerate(visualizationArguments, source, type);
+      },
+      type: PlotType.VISUALIZATION_CREATOR,
+    };
 
     return (
       <div className={classes(commonCss.page, padding(20, 't'))}>
@@ -215,7 +239,14 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
                         <div className={commonCss.page}>
                           {sidepanelSelectedTab === SidePaneTab.ARTIFACTS && (
                             <div className={commonCss.page}>
-                              {(selectedNodeDetails.viewerConfigs || []).map((config, i) => {
+                              <div className={padding(20, 'lrt')}>
+                                <PlotCard
+                                  configs={[visualizationCreatorConfig]}
+                                  title={VisualizationCreator.prototype.getDisplayName()}
+                                  maxDimension={500} />
+                                <Hr />
+                              </div>
+                              {this._getSelectedNodeConfigs().map((config, i) => {
                                 const title = componentMap[config.type].prototype.getDisplayName();
                                 return (
                                   <div key={i} className={padding(20, 'lrt')}>
@@ -395,6 +426,13 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     this.clearBanner();
     const runId = this.props.match.params[RouteParams.runId];
 
+    try {
+      const allowCustomVisualizations = await Apis.areCustomVisualizationsAllowed();
+      this.setState({ allowCustomVisualizations });
+    } catch (err) {
+      // Ignore error
+    }
+
     try {
       const runDetail = await Apis.runServiceApi.getRun(runId);
 
@@ -641,6 +679,56 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
       this.setStateSafe({ sidepanelBusy: false });
     }
   }
+
+  private _getSelectedNodeConfigs(): ViewerConfig[] {
+    const { generatedVisualizations, selectedNodeDetails } = this.state;
+    let selectedNodeConfigs: ViewerConfig[] = [];
+    if (selectedNodeDetails !== null) {
+      const nodeConfigs = selectedNodeDetails.viewerConfigs || [];
+      const generatedConfigs = generatedVisualizations
+        .filter(visualization => visualization.nodeId === selectedNodeDetails.id)
+        .map(visualization => visualization.config);
+        selectedNodeConfigs = nodeConfigs.concat(generatedConfigs);
+    }
+    return selectedNodeConfigs;
+  }
+
+  private async _onGenerate(visualizationArguments: string, source: string, type: ApiVisualizationType): Promise<void> {
+    const { selectedNodeDetails } = this.state;
+    const nodeId = selectedNodeDetails ? selectedNodeDetails.id : '';
+    if (nodeId.length === 0) {
+      this.showPageError('Unable to generate visualization, no component selected.');
+      return;
+    }
+
+    if (visualizationArguments.length) {
+      try {
+        JSON.parse(visualizationArguments);
+      } catch (err) {
+        this.showPageError('Unable to generate visualization, invalid JSON provided.', err);
+        return;
+      }
+    }
+    this.setState({ isGeneratingVisualization: true });
+    const visualizationData: ApiVisualization = {
+      arguments: visualizationArguments,
+      source,
+      type,
+    };
+    try {
+      const config = await Apis.buildPythonVisualizationConfig(visualizationData);
+      const { generatedVisualizations } = this.state;
+      generatedVisualizations.push({
+        config,
+        nodeId,
+      });
+      this.setState({ generatedVisualizations });
+    } catch (err) {
+      this.showPageError('Unable to generate visualization, an unexpected error was encountered.', err);
+    } finally {
+      this.setState({ isGeneratingVisualization: false });
+    }
+  }
 }
 
 export default RunDetails;
diff --git a/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap b/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap
index c5c8bbe1f31..3b976d3a4d2 100644
--- a/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap
+++ b/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap
@@ -574,6 +574,25 @@ exports[`RunDetails logs tab does not load logs if clicked node status is skippe
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -685,6 +704,25 @@ exports[`RunDetails logs tab keeps side pane open and on same tab when logs chan
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -796,6 +834,25 @@ exports[`RunDetails logs tab loads and shows logs in side pane 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -907,6 +964,25 @@ exports[`RunDetails logs tab shows error banner atop logs area if fetching logs
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -1016,6 +1092,25 @@ exports[`RunDetails logs tab shows error banner atop logs area if fetching logs
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -1125,6 +1220,25 @@ exports[`RunDetails logs tab shows warning banner and link to Stackdriver in log
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -1254,6 +1368,25 @@ exports[`RunDetails logs tab switches to logs tab in side pane 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -1365,6 +1498,25 @@ exports[`RunDetails opens side panel when graph node is clicked 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 />
@@ -1589,6 +1741,25 @@ exports[`RunDetails shows clicked node message in side panel 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 />
@@ -1691,6 +1862,25 @@ exports[`RunDetails shows clicked node output in side pane 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className="page"
                 >
@@ -2140,6 +2330,25 @@ exports[`RunDetails switches to inputs/outputs tab in side pane 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className=""
                 >
@@ -2269,6 +2478,25 @@ exports[`RunDetails switches to manifest tab in side pane 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className=""
                 >
@@ -2416,6 +2644,25 @@ exports[`RunDetails switches to volumes tab in side pane 1`] = `
               <div
                 className="page"
               >
+                <div
+                  className=""
+                >
+                  <PlotCard
+                    configs={
+                      Array [
+                        Object {
+                          "allowCustomVisualizations": false,
+                          "isBusy": false,
+                          "onGenerate": [Function],
+                          "type": "visualization-creator",
+                        },
+                      ]
+                    }
+                    maxDimension={500}
+                    title="Visualization Creator"
+                  />
+                  <Component />
+                </div>
                 <div
                   className=""
                 >

From 6bc83dbb206527a0226b1ca40a9b4902c4c0553f Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Wed, 28 Aug 2019 11:11:27 -0700
Subject: [PATCH 2/6] Fixed bug where switching selected node would not update
 visualizations in artifacts tab

---
 frontend/src/pages/RunDetails.tsx | 45 ++++++++++++++++---------------
 1 file changed, 23 insertions(+), 22 deletions(-)

diff --git a/frontend/src/pages/RunDetails.tsx b/frontend/src/pages/RunDetails.tsx
index f24afdf7967..e44087302fe 100644
--- a/frontend/src/pages/RunDetails.tsx
+++ b/frontend/src/pages/RunDetails.tsx
@@ -246,7 +246,7 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
                                   maxDimension={500} />
                                 <Hr />
                               </div>
-                              {this._getSelectedNodeConfigs().map((config, i) => {
+                              {(selectedNodeDetails.viewerConfigs || []).map((config, i) => {
                                 const title = componentMap[config.type].prototype.getDisplayName();
                                 return (
                                   <div key={i} className={padding(20, 'lrt')}>
@@ -623,7 +623,7 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
   }
 
   private async _loadSelectedNodeOutputs(): Promise<void> {
-    const selectedNodeDetails = this.state.selectedNodeDetails;
+    const { generatedVisualizations, selectedNodeDetails } = this.state;
     if (!selectedNodeDetails) {
       return;
     }
@@ -632,12 +632,21 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     if (workflow && workflow.status && workflow.status.nodes) {
       // Load runtime outputs from the selected Node
       const outputPaths = WorkflowParser.loadNodeOutputPaths(workflow.status.nodes[selectedNodeDetails.id]);
-
       // Load the viewer configurations from the output paths
       let viewerConfigs: ViewerConfig[] = [];
       for (const path of outputPaths) {
         viewerConfigs = viewerConfigs.concat(await OutputArtifactLoader.load(path));
       }
+      // tslint:disable-next-line: no-console
+      console.log('selected node:', selectedNodeDetails);
+      // tslint:disable-next-line: no-console
+      console.log('loaded configs:', viewerConfigs);
+      const generatedConfigs = generatedVisualizations
+        .filter(visualization => visualization.nodeId === selectedNodeDetails.id)
+        .map(visualization => visualization.config);
+      // tslint:disable-next-line: no-console
+      console.log('generated configs:', generatedConfigs);
+      viewerConfigs = viewerConfigs.concat(generatedConfigs);
 
       selectedNodeDetails.viewerConfigs = viewerConfigs;
       this.setStateSafe({ selectedNodeDetails });
@@ -680,22 +689,8 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     }
   }
 
-  private _getSelectedNodeConfigs(): ViewerConfig[] {
-    const { generatedVisualizations, selectedNodeDetails } = this.state;
-    let selectedNodeConfigs: ViewerConfig[] = [];
-    if (selectedNodeDetails !== null) {
-      const nodeConfigs = selectedNodeDetails.viewerConfigs || [];
-      const generatedConfigs = generatedVisualizations
-        .filter(visualization => visualization.nodeId === selectedNodeDetails.id)
-        .map(visualization => visualization.config);
-        selectedNodeConfigs = nodeConfigs.concat(generatedConfigs);
-    }
-    return selectedNodeConfigs;
-  }
-
   private async _onGenerate(visualizationArguments: string, source: string, type: ApiVisualizationType): Promise<void> {
-    const { selectedNodeDetails } = this.state;
-    const nodeId = selectedNodeDetails ? selectedNodeDetails.id : '';
+    const nodeId = this.state.selectedNodeDetails ? this.state.selectedNodeDetails.id : '';
     if (nodeId.length === 0) {
       this.showPageError('Unable to generate visualization, no component selected.');
       return;
@@ -717,12 +712,18 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
     };
     try {
       const config = await Apis.buildPythonVisualizationConfig(visualizationData);
-      const { generatedVisualizations } = this.state;
-      generatedVisualizations.push({
+      const { generatedVisualizations, selectedNodeDetails } = this.state;
+      const generatedVisualization: GeneratedVisualization = {
         config,
         nodeId,
-      });
-      this.setState({ generatedVisualizations });
+      };
+      generatedVisualizations.push(generatedVisualization);
+      if (selectedNodeDetails) {
+        const viewerConfigs = selectedNodeDetails.viewerConfigs || []; 
+        viewerConfigs.push(generatedVisualization.config);
+        selectedNodeDetails.viewerConfigs = viewerConfigs;
+      }
+      this.setState({ generatedVisualizations, selectedNodeDetails });
     } catch (err) {
       this.showPageError('Unable to generate visualization, an unexpected error was encountered.', err);
     } finally {

From 1e6455a39dc6762b6df8b7a5958dbc8e033d56c0 Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Wed, 28 Aug 2019 11:12:19 -0700
Subject: [PATCH 3/6] Remove debugging code

---
 frontend/src/pages/RunDetails.tsx | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/frontend/src/pages/RunDetails.tsx b/frontend/src/pages/RunDetails.tsx
index e44087302fe..ff9f51478a9 100644
--- a/frontend/src/pages/RunDetails.tsx
+++ b/frontend/src/pages/RunDetails.tsx
@@ -637,15 +637,9 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
       for (const path of outputPaths) {
         viewerConfigs = viewerConfigs.concat(await OutputArtifactLoader.load(path));
       }
-      // tslint:disable-next-line: no-console
-      console.log('selected node:', selectedNodeDetails);
-      // tslint:disable-next-line: no-console
-      console.log('loaded configs:', viewerConfigs);
       const generatedConfigs = generatedVisualizations
         .filter(visualization => visualization.nodeId === selectedNodeDetails.id)
         .map(visualization => visualization.config);
-      // tslint:disable-next-line: no-console
-      console.log('generated configs:', generatedConfigs);
       viewerConfigs = viewerConfigs.concat(generatedConfigs);
 
       selectedNodeDetails.viewerConfigs = viewerConfigs;

From 4b432a224779c7ce638ab128580b50ea96d13634 Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Wed, 28 Aug 2019 11:17:05 -0700
Subject: [PATCH 4/6] Updated RunDetails.test.tsx.snap

---
 .../__snapshots__/RunDetails.test.tsx.snap    | 306 ++++--------------
 1 file changed, 59 insertions(+), 247 deletions(-)

diff --git a/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap b/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap
index 3b976d3a4d2..1fe280a7f12 100644
--- a/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap
+++ b/frontend/src/pages/__snapshots__/RunDetails.test.tsx.snap
@@ -574,25 +574,6 @@ exports[`RunDetails logs tab does not load logs if clicked node status is skippe
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -704,25 +685,6 @@ exports[`RunDetails logs tab keeps side pane open and on same tab when logs chan
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -834,25 +796,6 @@ exports[`RunDetails logs tab loads and shows logs in side pane 1`] = `
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -964,25 +907,6 @@ exports[`RunDetails logs tab shows error banner atop logs area if fetching logs
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -1092,25 +1016,6 @@ exports[`RunDetails logs tab shows error banner atop logs area if fetching logs
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -1220,25 +1125,6 @@ exports[`RunDetails logs tab shows warning banner and link to Stackdriver in log
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -1368,25 +1254,6 @@ exports[`RunDetails logs tab switches to logs tab in side pane 1`] = `
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
@@ -1499,27 +1366,28 @@ exports[`RunDetails opens side panel when graph node is clicked 1`] = `
                 className="page"
               >
                 <div
-                  className=""
+                  className="page"
                 >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
+                  <div
+                    className=""
+                  >
+                    <PlotCard
+                      configs={
+                        Array [
+                          Object {
+                            "allowCustomVisualizations": false,
+                            "isBusy": false,
+                            "onGenerate": [Function],
+                            "type": "visualization-creator",
+                          },
+                        ]
+                      }
+                      maxDimension={500}
+                      title="Visualization Creator"
+                    />
+                    <Component />
+                  </div>
                 </div>
-                <div
-                  className="page"
-                />
               </div>
             </div>
           </SidePanel>
@@ -1742,27 +1610,28 @@ exports[`RunDetails shows clicked node message in side panel 1`] = `
                 className="page"
               >
                 <div
-                  className=""
+                  className="page"
                 >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
+                  <div
+                    className=""
+                  >
+                    <PlotCard
+                      configs={
+                        Array [
+                          Object {
+                            "allowCustomVisualizations": false,
+                            "isBusy": false,
+                            "onGenerate": [Function],
+                            "type": "visualization-creator",
+                          },
+                        ]
+                      }
+                      maxDimension={500}
+                      title="Visualization Creator"
+                    />
+                    <Component />
+                  </div>
                 </div>
-                <div
-                  className="page"
-                />
               </div>
             </div>
           </SidePanel>
@@ -1862,28 +1731,28 @@ exports[`RunDetails shows clicked node output in side pane 1`] = `
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className="page"
                 >
+                  <div
+                    className=""
+                  >
+                    <PlotCard
+                      configs={
+                        Array [
+                          Object {
+                            "allowCustomVisualizations": false,
+                            "isBusy": false,
+                            "onGenerate": [Function],
+                            "type": "visualization-creator",
+                          },
+                        ]
+                      }
+                      maxDimension={500}
+                      title="Visualization Creator"
+                    />
+                    <Component />
+                  </div>
                   <div
                     className=""
                     key="0"
@@ -2330,25 +2199,6 @@ exports[`RunDetails switches to inputs/outputs tab in side pane 1`] = `
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className=""
                 >
@@ -2478,25 +2328,6 @@ exports[`RunDetails switches to manifest tab in side pane 1`] = `
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className=""
                 >
@@ -2644,25 +2475,6 @@ exports[`RunDetails switches to volumes tab in side pane 1`] = `
               <div
                 className="page"
               >
-                <div
-                  className=""
-                >
-                  <PlotCard
-                    configs={
-                      Array [
-                        Object {
-                          "allowCustomVisualizations": false,
-                          "isBusy": false,
-                          "onGenerate": [Function],
-                          "type": "visualization-creator",
-                        },
-                      ]
-                    }
-                    maxDimension={500}
-                    title="Visualization Creator"
-                  />
-                  <Component />
-                </div>
                 <div
                   className=""
                 >

From 4c79c83d74bc22ee85bb274e8becb35f8cb1f1c5 Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Thu, 29 Aug 2019 09:35:14 -0700
Subject: [PATCH 5/6] Addressed PR comments

---
 frontend/src/pages/RunDetails.test.tsx | 13 +++++++++++++
 frontend/src/pages/RunDetails.tsx      |  3 ++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/frontend/src/pages/RunDetails.test.tsx b/frontend/src/pages/RunDetails.test.tsx
index 4be6e41f199..17629b780cf 100644
--- a/frontend/src/pages/RunDetails.test.tsx
+++ b/frontend/src/pages/RunDetails.test.tsx
@@ -647,6 +647,19 @@ describe('RunDetails', () => {
     });
   });
 
+  it('shows an error banner if the custom visualizations state API fails', async () => {
+    TestUtils.makeErrorResponseOnce(isCustomVisualizationsAllowedSpy, 'woops');
+    tree = shallow(<RunDetails {...generateProps()} />);
+    await isCustomVisualizationsAllowedSpy;
+    await TestUtils.flushPromises();
+    expect(updateBannerSpy).toHaveBeenCalledTimes(2); // Once initially to clear
+    expect(updateBannerSpy).toHaveBeenLastCalledWith(expect.objectContaining({
+      additionalInfo: 'woops',
+      message:'Error: Unable to enable custom visualizations. Click Details for more information.',
+      mode: 'error',
+    }));
+  });
+
   describe('logs tab', () => {
     it('switches to logs tab in side pane', async () => {
       testRun.pipeline_runtime!.workflow_manifest = JSON.stringify({
diff --git a/frontend/src/pages/RunDetails.tsx b/frontend/src/pages/RunDetails.tsx
index ff9f51478a9..74e3a9279c5 100644
--- a/frontend/src/pages/RunDetails.tsx
+++ b/frontend/src/pages/RunDetails.tsx
@@ -430,7 +430,7 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
       const allowCustomVisualizations = await Apis.areCustomVisualizationsAllowed();
       this.setState({ allowCustomVisualizations });
     } catch (err) {
-      // Ignore error
+      this.showPageError('Error: Unable to enable custom visualizations.', err);
     }
 
     try {
@@ -692,6 +692,7 @@ class RunDetails extends Page<RunDetailsProps, RunDetailsState> {
 
     if (visualizationArguments.length) {
       try {
+        // Attempts to validate JSON, if attempt fails an error is displayed.
         JSON.parse(visualizationArguments);
       } catch (err) {
         this.showPageError('Unable to generate visualization, invalid JSON provided.', err);

From 76081681e31738977ea2db14e1ec5aa8def70ef1 Mon Sep 17 00:00:00 2001
From: Kirin Patel <kirinpatel@gmail.com>
Date: Thu, 29 Aug 2019 09:55:49 -0700
Subject: [PATCH 6/6] Improved error message readability

---
 frontend/src/lib/Apis.ts | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/frontend/src/lib/Apis.ts b/frontend/src/lib/Apis.ts
index 46b71a2a95e..64360871dd1 100644
--- a/frontend/src/lib/Apis.ts
+++ b/frontend/src/lib/Apis.ts
@@ -69,7 +69,9 @@ export class Apis {
         type: PlotType.WEB_APP,
       } as HTMLViewerConfig;
     } else {
-      throw new Error('Unable to generate visualization!');
+      // This should never be thrown as the html property of a generated
+      // visualization is always set for successful visualization generations.
+      throw new Error('Visualization was generated successfully but generated HTML was not found.');
     }
   }