Fix bug publish screenshot (#762)

[Preview this branch](https://cosmos-explorer-preview.azurewebsites.net/pull/762?feature.someFeatureFlagYouMightNeed=true)

The main change in this PR fixes the snapshot functionality in the Publish pane-related components. Because the code cell outputs are now rendered in their own iframes for security reasons, a single snapshot of the notebook is no longer possible: each cell output takes its own snapshot and the snapshots are collated on the main notebook snapshot.
- Move the snapshot functionality to notebook components: this removes the reference of the notebook DOM node that we must pass to the Publish pane via explorer.
- Add slice in the state and actions in notebook redux for notebook snapshot requests and result
- Add post robot message to take snapshots and receive results
- Add logic in `NotebookRenderer` to wait for all output snapshots done before taking the main one collating.
- Use `zustand` to share snapshot between Redux world and React world. This solves the issue of keeping the `PanelContainer` component generic, while being able to update its children (`PublishPanel` component) with the new snapshot.

Additional changes:
- Add `local()` in `@font-face` to check if font is already installed before downloading the font (must be done for Safari, but not Edge/Chrome)
- Add "Export output to image" menu item in notebook cell, since each cell output can take its own snapshot (which can be downloaded)
![image](https://user-images.githubusercontent.com/21954022/117454706-b5f16600-af46-11eb-8535-6bf99f3d9170.png)

less/Common/Constants.less

@@ -4,7 +4,7 @@
 
 @font-face {
   font-family: wf_segoe-ui_normal;
-  src: url("../../fonts/segoe-ui/west-european/normal/latest.woff");
+  src: local("Segoe UI"), url("../../fonts/segoe-ui/west-european/normal/latest.woff");
 }
 
 @DataExplorerFont: wf_segoe-ui_normal, "Segoe UI", "Segoe WP", Tahoma, Arial, sans-serif;

package-lock.json

@@ -5255,6 +5255,12 @@
         "@types/d3-selection": "*"
       }
     },
+    "@types/dom-to-image": {
+      "version": "2.6.2",
+      "resolved": "https://registry.npmjs.org/@types/dom-to-image/-/dom-to-image-2.6.2.tgz",
+      "integrity": "sha512-Yxbwmz/glNwRIXfBI8efG2bgIxrFAKV1MdfpqbUDq25ULMot7U7FYXPiso5G8DlBExSP+AakuG0mNus9yw4RZQ==",
+      "dev": true
+    },
     "@types/dom4": {
       "version": "2.0.1",
       "resolved": "https://registry.npmjs.org/@types/dom4/-/dom4-2.0.1.tgz",
@@ -9667,6 +9673,11 @@
         }
       }
     },
+    "dom-to-image": {
+      "version": "2.6.0",
+      "resolved": "https://registry.npmjs.org/dom-to-image/-/dom-to-image-2.6.0.tgz",
+      "integrity": "sha1-ilA2CAiMh7HCL5A0rgMuGJiVWGc="
+    },
     "dom-walk": {
       "version": "0.1.2",
       "resolved": "https://registry.npmjs.org/dom-walk/-/dom-walk-0.1.2.tgz",
@@ -25918,6 +25929,11 @@
       "version": "1.0.46",
       "resolved": "https://registry.npmjs.org/zalgo-promise/-/zalgo-promise-1.0.46.tgz",
       "integrity": "sha512-tzPpQRqaQQavxl17TY98nznvmr+judUg3My7ugsUcRDbdqisYOE2z79HNNDgXnyX3eA0mf2bMOJrqHptt00npg=="
+    },
+    "zustand": {
+      "version": "3.5.0",
+      "resolved": "https://registry.npmjs.org/zustand/-/zustand-3.5.0.tgz",
+      "integrity": "sha512-fwZfax2c954pWB+RYXHXG0HWyuoUj8YiNykRjZC/w6L7ay9fPQ7M90mgDVP1KJsRqxLsDILBWSNxuzw5BkCpxA=="
     }
   }
 }

package.json

@@ -56,6 +56,7 @@
     "datatables.net-dt": "1.10.19",
     "date-fns": "1.29.0",
     "dayjs": "1.8.19",
+    "dom-to-image": "2.6.0",
     "dotenv": "8.2.0",
     "eslint-plugin-jest": "23.13.2",
     "eslint-plugin-react": "7.20.0",
@@ -97,7 +98,8 @@
     "swr": "0.4.0",
     "terser-webpack-plugin": "3.1.0",
     "underscore": "1.9.1",
-    "utility-types": "3.10.0"
+    "utility-types": "3.10.0",
+    "zustand": "3.5.0"
   },
   "devDependencies": {
     "@babel/core": "7.9.0",
@@ -109,6 +111,7 @@
     "@types/codemirror": "0.0.56",
     "@types/crossroads": "0.0.30",
     "@types/d3": "5.9.2",
+    "@types/dom-to-image": "2.6.2",
     "@types/enzyme": "3.10.7",
     "@types/enzyme-adapter-react-16": "1.0.6",
     "@types/hasher": "0.0.31",

src/CellOutputViewer/CellOutputViewer.tsx

@@ -9,11 +9,17 @@ import postRobot from "post-robot";
 import * as React from "react";
 import * as ReactDOM from "react-dom";
 import "../../externals/iframeResizer.contentWindow.min.js"; // Required for iFrameResizer to work
+import { SnapshotRequest } from "../Explorer/Notebook/NotebookComponent/types";
 import "../Explorer/Notebook/NotebookRenderer/base.css";
 import "../Explorer/Notebook/NotebookRenderer/default.css";
+import { NotebookUtil } from "../Explorer/Notebook/NotebookUtil";
 import "./CellOutputViewer.less";
 import { TransformMedia } from "./TransformMedia";
 
+export interface SnapshotResponse {
+  imageSrc: string;
+  requestId: string;
+}
 export interface CellOutputViewerProps {
   id: string;
   contentRef: ContentRef;
@@ -62,6 +68,36 @@ const onInit = async () => {
       ReactDOM.render(outputs, document.getElementById("cellOutput"));
     }
   );
+
+  postRobot.on(
+    "snapshotRequest",
+    {
+      window: window.parent,
+      domain: window.location.origin,
+    },
+    async (event): Promise<SnapshotResponse> => {
+      const topNode = document.getElementById("cellOutput");
+      if (!topNode) {
+        const errorMsg = "No top node to snapshot";
+        return Promise.reject(new Error(errorMsg));
+      }
+
+      // Typescript definition for event is wrong. So read props by casting to <any>
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const snapshotRequest = (event as any).data as SnapshotRequest;
+      const result = await NotebookUtil.takeScreenshotDomToImage(
+        topNode,
+        snapshotRequest.aspectRatio,
+        undefined,
+        snapshotRequest.downloadFilename
+      );
+
+      return {
+        imageSrc: result.imageSrc,
+        requestId: snapshotRequest.requestId,
+      };
+    }
+  );
 };
 
 // Entry point

src/Explorer/Explorer.tsx

@@ -45,6 +45,7 @@ import { GalleryTab as GalleryTabKind } from "./Controls/NotebookGallery/Gallery
 import { CommandBarComponentAdapter } from "./Menus/CommandBar/CommandBarComponentAdapter";
 import { ConsoleData } from "./Menus/NotificationConsole/NotificationConsoleComponent";
 import * as FileSystemUtil from "./Notebook/FileSystemUtil";
+import { SnapshotRequest } from "./Notebook/NotebookComponent/types";
 import { NotebookContentItem, NotebookContentItemType } from "./Notebook/NotebookContentItem";
 import type NotebookManager from "./Notebook/NotebookManager";
 import type { NotebookPaneContent } from "./Notebook/NotebookManager";
@@ -91,7 +92,7 @@ export interface ExplorerParams {
   setIsNotificationConsoleExpanded: (isExpanded: boolean) => void;
   setNotificationConsoleData: (consoleData: ConsoleData) => void;
   setInProgressConsoleDataIdToBeDeleted: (id: string) => void;
-  openSidePanel: (headerText: string, panelContent: JSX.Element) => void;
+  openSidePanel: (headerText: string, panelContent: JSX.Element, onClose?: () => void) => void;
   closeSidePanel: () => void;
   closeDialog: () => void;
   openDialog: (props: DialogProps) => void;
@@ -125,7 +126,7 @@ export default class Explorer {
 
   // Panes
   public contextPanes: ContextualPaneBase[];
-  public openSidePanel: (headerText: string, panelContent: JSX.Element) => void;
+  public openSidePanel: (headerText: string, panelContent: JSX.Element, onClose?: () => void) => void;
   public closeSidePanel: () => void;
 
   // Resource Tree
@@ -1270,10 +1271,18 @@ export default class Explorer {
   public async publishNotebook(
     name: string,
     content: NotebookPaneContent,
-    parentDomElement?: HTMLElement
+    notebookContentRef?: string,
+    onTakeSnapshot?: (request: SnapshotRequest) => void,
+    onClosePanel?: () => void
   ): Promise<void> {
     if (this.notebookManager) {
-      await this.notebookManager.openPublishNotebookPane(name, content, parentDomElement);
+      await this.notebookManager.openPublishNotebookPane(
+        name,
+        content,
+        notebookContentRef,
+        onTakeSnapshot,
+        onClosePanel
+      );
       this.isPublishNotebookPaneEnabled(true);
     }
   }

src/Explorer/Notebook/NotebookComponent/NotebookComponentAdapter.tsx

@@ -1,14 +1,11 @@
-import * as React from "react";
-
-import { ReactAdapter } from "../../../Bindings/ReactBindingHandler";
-import { NotebookClientV2 } from "../NotebookClientV2";
-
 // Vendor modules
 import { actions, createContentRef, createKernelRef, selectors } from "@nteract/core";
-import VirtualCommandBarComponent from "./VirtualCommandBarComponent";
+import * as React from "react";
+import { ReactAdapter } from "../../../Bindings/ReactBindingHandler";
+import { NotebookClientV2 } from "../NotebookClientV2";
 import { NotebookContentItem } from "../NotebookContentItem";
 import { NotebookComponentBootstrapper } from "./NotebookComponentBootstrapper";
-import { CdbAppState } from "./types";
+import VirtualCommandBarComponent from "./VirtualCommandBarComponent";
 
 export interface NotebookComponentAdapterOptions {
   contentItem: NotebookContentItem;
@@ -19,7 +16,6 @@ export interface NotebookComponentAdapterOptions {
 
 export class NotebookComponentAdapter extends NotebookComponentBootstrapper implements ReactAdapter {
   private onUpdateKernelInfo: () => void;
-  public getNotebookParentElement: () => HTMLElement;
   public parameters: any;
 
   constructor(options: NotebookComponentAdapterOptions) {
@@ -46,11 +42,6 @@ export class NotebookComponentAdapter extends NotebookComponentBootstrapper impl
         })
       );
     }
-
-    this.getNotebookParentElement = () => {
-      const cdbAppState = this.getStore().getState() as CdbAppState;
-      return cdbAppState.cdb.currentNotebookParentElements.get(this.contentRef);
-    };
   }
 
   protected renderExtraComponent = (): JSX.Element => {

src/Explorer/Notebook/NotebookComponent/NotebookComponentBootstrapper.tsx

@@ -1,35 +1,29 @@
-import * as React from "react";
+import { CellId, CellType, ImmutableNotebook } from "@nteract/commutable";
-
-import { NotebookComponent } from "./NotebookComponent";
-import { NotebookClientV2 } from "../NotebookClientV2";
-import { NotebookUtil } from "../NotebookUtil";
-
 // Vendor modules
 import {
   actions,
   AppState,
-  createKernelRef,
-  DocumentRecordProps,
   ContentRef,
+  DocumentRecordProps,
   KernelRef,
   NotebookContentRecord,
   selectors,
 } from "@nteract/core";
-import * as Immutable from "immutable";
-import { Provider } from "react-redux";
-import { CellType, CellId, ImmutableNotebook } from "@nteract/commutable";
-import { Store, AnyAction } from "redux";
-
-import "./NotebookComponent.less";
-
-import "codemirror/addon/hint/show-hint.css";
-import "codemirror/lib/codemirror.css";
 import "@nteract/styles/editor-overrides.css";
 import "@nteract/styles/global-variables.css";
+import "codemirror/addon/hint/show-hint.css";
+import "codemirror/lib/codemirror.css";
+import * as Immutable from "immutable";
+import * as React from "react";
+import { Provider } from "react-redux";
 import "react-table/react-table.css";
-
+import { AnyAction, Store } from "redux";
-import * as CdbActions from "./actions";
+import { NotebookClientV2 } from "../NotebookClientV2";
+import { NotebookUtil } from "../NotebookUtil";
 import * as NteractUtil from "../NTeractUtil";
+import * as CdbActions from "./actions";
+import { NotebookComponent } from "./NotebookComponent";
+import "./NotebookComponent.less";
 
 export interface NotebookComponentBootstrapperOptions {
   notebookClient: NotebookClientV2;
@@ -37,7 +31,7 @@ export interface NotebookComponentBootstrapperOptions {
 }
 
 export class NotebookComponentBootstrapper {
-  protected contentRef: ContentRef;
+  public contentRef: ContentRef;
   protected renderExtraComponent: () => JSX.Element;
 
   private notebookClient: NotebookClientV2;

src/Explorer/Notebook/NotebookComponent/actions.ts

@@ -1,6 +1,7 @@
 import { CellId } from "@nteract/commutable";
 import { ContentRef } from "@nteract/core";
 import { Action } from "../../../Shared/Telemetry/TelemetryConstants";
+import { SnapshotFragment, SnapshotRequest } from "./types";
 
 export const CLOSE_NOTEBOOK = "CLOSE_NOTEBOOK";
 export interface CloseNotebookAction {
@@ -85,21 +86,68 @@ export const traceNotebookTelemetry = (payload: {
   };
 };
 
-export const UPDATE_NOTEBOOK_PARENT_DOM_ELTS = "UPDATE_NOTEBOOK_PARENT_DOM_ELTS";
+export const STORE_CELL_OUTPUT_SNAPSHOT = "STORE_CELL_OUTPUT_SNAPSHOT";
-export interface UpdateNotebookParentDomEltAction {
+export interface StoreCellOutputSnapshotAction {
-  type: "UPDATE_NOTEBOOK_PARENT_DOM_ELTS";
+  type: "STORE_CELL_OUTPUT_SNAPSHOT";
   payload: {
-    contentRef: ContentRef;
+    cellId: string;
-    parentElt: HTMLElement;
+    snapshot: SnapshotFragment;
   };
 }
 
-export const UpdateNotebookParentDomElt = (payload: {
+export const storeCellOutputSnapshot = (payload: {
-  contentRef: ContentRef;
+  cellId: string;
-  parentElt: HTMLElement;
+  snapshot: SnapshotFragment;
-}): UpdateNotebookParentDomEltAction => {
+}): StoreCellOutputSnapshotAction => {
   return {
-    type: UPDATE_NOTEBOOK_PARENT_DOM_ELTS,
+    type: STORE_CELL_OUTPUT_SNAPSHOT,
+    payload,
+  };
+};
+
+export const STORE_NOTEBOOK_SNAPSHOT = "STORE_NOTEBOOK_SNAPSHOT";
+export interface StoreNotebookSnapshotAction {
+  type: "STORE_NOTEBOOK_SNAPSHOT";
+  payload: {
+    imageSrc: string;
+    requestId: string;
+  };
+}
+
+export const storeNotebookSnapshot = (payload: {
+  imageSrc: string;
+  requestId: string;
+}): StoreNotebookSnapshotAction => {
+  return {
+    type: STORE_NOTEBOOK_SNAPSHOT,
+    payload,
+  };
+};
+
+export const TAKE_NOTEBOOK_SNAPSHOT = "TAKE_NOTEBOOK_SNAPSHOT";
+export interface TakeNotebookSnapshotAction {
+  type: "TAKE_NOTEBOOK_SNAPSHOT";
+  payload: SnapshotRequest;
+}
+
+export const takeNotebookSnapshot = (payload: SnapshotRequest): TakeNotebookSnapshotAction => {
+  return {
+    type: TAKE_NOTEBOOK_SNAPSHOT,
+    payload,
+  };
+};
+
+export const NOTEBOOK_SNAPSHOT_ERROR = "NOTEBOOK_SNAPSHOT_ERROR";
+export interface NotebookSnapshotErrorAction {
+  type: "NOTEBOOK_SNAPSHOT_ERROR";
+  payload: {
+    error: string;
+  };
+}
+
+export const notebookSnapshotError = (payload: { error: string }): NotebookSnapshotErrorAction => {
+  return {
+    type: NOTEBOOK_SNAPSHOT_ERROR,
     payload,
   };
 };

src/Explorer/Notebook/NotebookComponent/reducers.ts

@@ -70,17 +70,32 @@ export const cdbReducer = (state: CdbRecord, action: Action) => {
       return state.set("hoveredCellId", typedAction.payload.cellId);
     }
 
-    case cdbActions.UPDATE_NOTEBOOK_PARENT_DOM_ELTS: {
+    case cdbActions.STORE_CELL_OUTPUT_SNAPSHOT: {
-      const typedAction = action as cdbActions.UpdateNotebookParentDomEltAction;
+      const typedAction = action as cdbActions.StoreCellOutputSnapshotAction;
-      var parentEltsMap = state.get("currentNotebookParentElements");
+      state.cellOutputSnapshots.set(typedAction.payload.cellId, typedAction.payload.snapshot);
-      const contentRef = typedAction.payload.contentRef;
+      // TODO Simpler datastructure to instantiate new Map?
-      const parentElt = typedAction.payload.parentElt;
+      return state.set("cellOutputSnapshots", new Map(state.cellOutputSnapshots));
-      if (parentElt) {
-        parentEltsMap.set(contentRef, parentElt);
-      } else {
-        parentEltsMap.delete(contentRef);
     }
-      return state.set("currentNotebookParentElements", parentEltsMap);
+
+    case cdbActions.STORE_NOTEBOOK_SNAPSHOT: {
+      const typedAction = action as cdbActions.StoreNotebookSnapshotAction;
+      // Clear pending request
+      return state.set("notebookSnapshot", typedAction.payload).set("pendingSnapshotRequest", undefined);
+    }
+
+    case cdbActions.TAKE_NOTEBOOK_SNAPSHOT: {
+      const typedAction = action as cdbActions.TakeNotebookSnapshotAction;
+      // Clear previous snapshots
+      return state
+        .set("cellOutputSnapshots", new Map())
+        .set("notebookSnapshot", undefined)
+        .set("notebookSnapshotError", undefined)
+        .set("pendingSnapshotRequest", typedAction.payload);
+    }
+
+    case cdbActions.NOTEBOOK_SNAPSHOT_ERROR: {
+      const typedAction = action as cdbActions.NotebookSnapshotErrorAction;
+      return state.set("notebookSnapshotError", typedAction.payload.error);
     }
   }
   return state;

src/Explorer/Notebook/NotebookComponent/types.ts

@@ -1,15 +1,40 @@
-import * as Immutable from "immutable";
-import { AppState, ContentRef } from "@nteract/core";
-
-import { Notebook } from "../../../Common/Constants";
 import { CellId } from "@nteract/commutable";
+import { AppState } from "@nteract/core";
+import * as Immutable from "immutable";
+import { Notebook } from "../../../Common/Constants";
+
+export interface SnapshotFragment {
+  image: HTMLImageElement;
+  boundingClientRect: DOMRect;
+  requestId: string;
+}
+
+export type SnapshotRequest = NotebookSnapshotRequest | CellSnapshotRequest;
+interface NotebookSnapshotRequestBase {
+  requestId: string;
+  aspectRatio: number;
+  notebookContentRef: string; // notebook redux contentRef
+  downloadFilename?: string; // Optional: will download as a file
+}
+
+interface NotebookSnapshotRequest extends NotebookSnapshotRequestBase {
+  type: "notebook";
+}
+
+interface CellSnapshotRequest extends NotebookSnapshotRequestBase {
+  type: "celloutput";
+  cellId: string;
+}
 
 export interface CdbRecordProps {
   databaseAccountName: string | undefined;
   defaultExperience: string | undefined;
   kernelRestartDelayMs: number;
   hoveredCellId: CellId | undefined;
-  currentNotebookParentElements: Map<ContentRef, HTMLElement>;
+  cellOutputSnapshots: Map<string, SnapshotFragment>;
+  notebookSnapshot?: { imageSrc: string; requestId: string };
+  pendingSnapshotRequest?: SnapshotRequest;
+  notebookSnapshotError?: string;
 }
 
 export type CdbRecord = Immutable.RecordOf<CdbRecordProps>;
@@ -23,5 +48,8 @@ export const makeCdbRecord = Immutable.Record<CdbRecordProps>({
   defaultExperience: undefined,
   kernelRestartDelayMs: Notebook.kernelRestartInitialDelayMs,
   hoveredCellId: undefined,
-  currentNotebookParentElements: new Map<ContentRef, HTMLElement>(),
+  cellOutputSnapshots: new Map(),
+  notebookSnapshot: undefined,
+  pendingSnapshotRequest: undefined,
+  notebookSnapshotError: undefined,
 });

src/Explorer/Notebook/NotebookManager.tsx

@@ -26,6 +26,7 @@ import { CopyNotebookPane } from "../Panes/CopyNotebookPane/CopyNotebookPane";
 import { PublishNotebookPane } from "../Panes/PublishNotebookPane/PublishNotebookPane";
 import { ResourceTreeAdapter } from "../Tree/ResourceTreeAdapter";
 import { NotebookContentProvider } from "./NotebookComponent/NotebookContentProvider";
+import { SnapshotRequest } from "./NotebookComponent/types";
 import { NotebookContainerClient } from "./NotebookContainerClient";
 import { NotebookContentClient } from "./NotebookContentClient";
 
@@ -112,11 +113,13 @@ export default class NotebookManager {
   public async openPublishNotebookPane(
     name: string,
     content: NotebookPaneContent,
-    parentDomElement: HTMLElement
+    notebookContentRef: string,
+    onTakeSnapshot: (request: SnapshotRequest) => void,
+    onClosePanel: () => void
   ): Promise<void> {
     const explorer = this.params.container;
     explorer.openSidePanel(
-      "New Collection",
+      "Publish Notebook",
       <PublishNotebookPane
         explorer={this.params.container}
         junoClient={this.junoClient}
@@ -125,8 +128,10 @@ export default class NotebookManager {
         name={name}
         author={getFullName()}
         notebookContent={content}
-        parentDomElement={parentDomElement}
+        notebookContentRef={notebookContentRef}
-      />
+        onTakeSnapshot={onTakeSnapshot}
+      />,
+      onClosePanel
     );
   }
 

src/Explorer/Notebook/NotebookRenderer/NotebookRenderer.tsx

@@ -1,6 +1,6 @@
 import { CellId } from "@nteract/commutable";
 import { CellType } from "@nteract/commutable/src";
-import { actions, ContentRef } from "@nteract/core";
+import { actions, ContentRef, selectors } from "@nteract/core";
 import { Cells, CodeCell, RawCell } from "@nteract/stateful-components";
 import MonacoEditor from "@nteract/stateful-components/lib/inputs/connected-editors/monacoEditor";
 import { PassedEditorProps } from "@nteract/stateful-components/lib/inputs/editor";
@@ -12,6 +12,8 @@ import { Dispatch } from "redux";
 import { userContext } from "../../../UserContext";
 import * as cdbActions from "../NotebookComponent/actions";
 import loadTransform from "../NotebookComponent/loadTransform";
+import { CdbAppState, SnapshotFragment, SnapshotRequest } from "../NotebookComponent/types";
+import { NotebookUtil } from "../NotebookUtil";
 import { AzureTheme } from "./AzureTheme";
 import "./base.css";
 import CellCreator from "./decorators/CellCreator";
@@ -32,10 +34,18 @@ export interface NotebookRendererBaseProps {
 }
 
 interface NotebookRendererDispatchProps {
-  updateNotebookParentDomElt: (contentRef: ContentRef, parentElt: HTMLElement) => void;
+  storeNotebookSnapshot: (imageSrc: string, requestId: string) => void;
+  notebookSnapshotError: (error: string) => void;
 }
 
-type NotebookRendererProps = NotebookRendererBaseProps & NotebookRendererDispatchProps;
+interface StateProps {
+  pendingSnapshotRequest: SnapshotRequest;
+  cellOutputSnapshots: Map<string, SnapshotFragment>;
+  notebookSnapshot: { imageSrc: string; requestId: string };
+  nbCodeCells: number;
+}
+
+type NotebookRendererProps = NotebookRendererBaseProps & NotebookRendererDispatchProps & StateProps;
 
 const decorate = (id: string, contentRef: ContentRef, cell_type: CellType, children: React.ReactNode) => {
   const Cell = () => (
@@ -60,27 +70,37 @@ const decorate = (id: string, contentRef: ContentRef, cell_type: CellType, child
 class BaseNotebookRenderer extends React.Component<NotebookRendererProps> {
   private notebookRendererRef = React.createRef<HTMLDivElement>();
 
-  constructor(props: NotebookRendererProps) {
-    super(props);
-
-    this.state = {
-      hoveredCellId: undefined,
-    };
-  }
-
   componentDidMount() {
     if (!userContext.features.sandboxNotebookOutputs) {
       loadTransform(this.props as any);
     }
-    this.props.updateNotebookParentDomElt(this.props.contentRef, this.notebookRendererRef.current);
   }
 
-  componentDidUpdate() {
+  async componentDidUpdate(): Promise<void> {
-    this.props.updateNotebookParentDomElt(this.props.contentRef, this.notebookRendererRef.current);
+    // Take a snapshot if there's a pending request and all the outputs are also saved
+    if (
+      this.props.pendingSnapshotRequest &&
+      this.props.pendingSnapshotRequest.type === "notebook" &&
+      this.props.pendingSnapshotRequest.notebookContentRef === this.props.contentRef &&
+      (!this.props.notebookSnapshot ||
+        this.props.pendingSnapshotRequest.requestId !== this.props.notebookSnapshot.requestId) &&
+      this.props.cellOutputSnapshots.size === this.props.nbCodeCells
+    ) {
+      try {
+        // Use Html2Canvas because it is much more reliable and fast than dom-to-file
+        const result = await NotebookUtil.takeScreenshotHtml2Canvas(
+          this.notebookRendererRef.current,
+          this.props.pendingSnapshotRequest.aspectRatio,
+          [...this.props.cellOutputSnapshots.values()],
+          this.props.pendingSnapshotRequest.downloadFilename
+        );
+        this.props.storeNotebookSnapshot(result.imageSrc, this.props.pendingSnapshotRequest.requestId);
+      } catch (error) {
+        this.props.notebookSnapshotError(error.message);
+      } finally {
+        this.setState({ processedSnapshotRequest: undefined });
+      }
     }
-
-  componentWillUnmount() {
-    this.props.updateNotebookParentDomElt(this.props.contentRef, undefined);
   }
 
   render(): JSX.Element {
@@ -156,28 +176,40 @@ class BaseNotebookRenderer extends React.Component<NotebookRendererProps> {
   }
 }
 
+export const makeMapStateToProps = (
+  initialState: CdbAppState,
+  ownProps: NotebookRendererProps
+): ((state: CdbAppState) => StateProps) => {
+  const mapStateToProps = (state: CdbAppState): StateProps => {
+    const { contentRef } = ownProps;
+    const model = selectors.model(state, { contentRef });
+
+    let nbCodeCells;
+    if (model && model.type === "notebook") {
+      nbCodeCells = NotebookUtil.findCodeCellWithDisplay(model.notebook).length;
+    }
+    const { pendingSnapshotRequest, cellOutputSnapshots, notebookSnapshot } = state.cdb;
+    return { pendingSnapshotRequest, cellOutputSnapshots, notebookSnapshot, nbCodeCells };
+  };
+  return mapStateToProps;
+};
+
 const makeMapDispatchToProps = (initialDispatch: Dispatch, initialProps: NotebookRendererBaseProps) => {
   const mapDispatchToProps = (dispatch: Dispatch) => {
     return {
-      addTransform: (transform: React.ComponentType & { MIMETYPE: string }) => {
+      addTransform: (transform: React.ComponentType & { MIMETYPE: string }) =>
-        return dispatch(
+        dispatch(
           actions.addTransform({
             mediaType: transform.MIMETYPE,
             component: transform,
           })
-        );
+        ),
-      },
+      storeNotebookSnapshot: (imageSrc: string, requestId: string) =>
-      updateNotebookParentDomElt: (contentRef: ContentRef, parentElt: HTMLElement) => {
+        dispatch(cdbActions.storeNotebookSnapshot({ imageSrc, requestId })),
-        return dispatch(
+      notebookSnapshotError: (error: string) => dispatch(cdbActions.notebookSnapshotError({ error })),
-          cdbActions.UpdateNotebookParentDomElt({
-            contentRef,
-            parentElt,
-          })
-        );
-      },
     };
   };
   return mapDispatchToProps;
 };
 
-export default connect(null, makeMapDispatchToProps)(BaseNotebookRenderer);
+export default connect(makeMapStateToProps, makeMapDispatchToProps)(BaseNotebookRenderer);

src/Explorer/Notebook/NotebookRenderer/Toolbar.tsx

@@ -1,5 +1,5 @@
 import { ContextualMenuItemType, DirectionalHint, IconButton, IContextualMenuItem } from "@fluentui/react";
-import { CellId, CellType } from "@nteract/commutable";
+import { CellId, CellType, ImmutableCodeCell } from "@nteract/commutable";
 import { actions, AppState, DocumentRecordProps } from "@nteract/core";
 import * as selectors from "@nteract/selectors";
 import { CellToolbarContext } from "@nteract/stateful-components";
@@ -10,6 +10,8 @@ import { connect } from "react-redux";
 import { Dispatch } from "redux";
 import { Action, ActionModifiers } from "../../../Shared/Telemetry/TelemetryConstants";
 import * as cdbActions from "../NotebookComponent/actions";
+import { SnapshotRequest } from "../NotebookComponent/types";
+import { NotebookUtil } from "../NotebookUtil";
 
 export interface ComponentProps {
   contentRef: ContentRef;
@@ -26,12 +28,14 @@ interface DispatchProps {
   clearOutputs: () => void;
   deleteCell: () => void;
   traceNotebookTelemetry: (action: Action, actionModifier?: string, data?: any) => void;
+  takeNotebookSnapshot: (payload: SnapshotRequest) => void;
 }
 
 interface StateProps {
   cellType: CellType;
   cellIdAbove: CellId;
   cellIdBelow: CellId;
+  hasCodeOutput: boolean;
 }
 
 class BaseToolbar extends React.PureComponent<ComponentProps & DispatchProps & StateProps> {
@@ -58,11 +62,29 @@ class BaseToolbar extends React.PureComponent<ComponentProps & DispatchProps & S
             this.props.traceNotebookTelemetry(Action.NotebooksClearOutputsFromMenu, ActionModifiers.Mark);
           },
         },
-        {
+      ]);
+
+      if (this.props.hasCodeOutput) {
+        items.push({
+          key: "Export output to image",
+          text: "Export output to image",
+          onClick: () => {
+            this.props.takeNotebookSnapshot({
+              requestId: new Date().getTime().toString(),
+              aspectRatio: undefined,
+              type: "celloutput",
+              cellId: this.props.id,
+              notebookContentRef: this.props.contentRef,
+              downloadFilename: `celloutput-${this.props.contentRef}_${this.props.id}.png`,
+            });
+          },
+        });
+      }
+
+      items.push({
         key: "Divider",
         itemType: ContextualMenuItemType.Divider,
-        },
+      });
-      ]);
     }
 
     items = items.concat([
@@ -183,12 +205,13 @@ const mapDispatchToProps = (
   deleteCell: () => dispatch(actions.deleteCell({ id, contentRef })),
   traceNotebookTelemetry: (action: Action, actionModifier?: string, data?: any) =>
     dispatch(cdbActions.traceNotebookTelemetry({ action, actionModifier, data })),
+  takeNotebookSnapshot: (request: SnapshotRequest) => dispatch(cdbActions.takeNotebookSnapshot(request)),
 });
 
 const makeMapStateToProps = (state: AppState, ownProps: ComponentProps): ((state: AppState) => StateProps) => {
   const mapStateToProps = (state: AppState) => {
-    const cellType = selectors.cell.cellFromState(state, { id: ownProps.id, contentRef: ownProps.contentRef })
+    const cell = selectors.cell.cellFromState(state, { id: ownProps.id, contentRef: ownProps.contentRef });
-      .cell_type;
+    const cellType = cell.cell_type;
     const model = selectors.model(state, { contentRef: ownProps.contentRef });
     const cellOrder = selectors.notebook.cellOrder(model as RecordOf<DocumentRecordProps>);
     const cellIndex = cellOrder.indexOf(ownProps.id);
@@ -199,6 +222,7 @@ const makeMapStateToProps = (state: AppState, ownProps: ComponentProps): ((state
       cellType,
       cellIdAbove,
       cellIdBelow,
+      hasCodeOutput: cellType === "code" && NotebookUtil.hasCodeCellOutput(cell as ImmutableCodeCell),
     };
   };
   return mapStateToProps;

src/Explorer/Notebook/NotebookRenderer/outputs/SandboxOutputs.tsx

@@ -7,7 +7,9 @@ import postRobot from "post-robot";
 import React from "react";
 import { connect } from "react-redux";
 import { Dispatch } from "redux";
-import { CellOutputViewerProps } from "../../../../CellOutputViewer/CellOutputViewer";
+import { CellOutputViewerProps, SnapshotResponse } from "../../../../CellOutputViewer/CellOutputViewer";
+import * as cdbActions from "../../NotebookComponent/actions";
+import { CdbAppState, SnapshotFragment, SnapshotRequest } from "../../NotebookComponent/types";
 
 // Adapted from https://github.com/nteract/nteract/blob/main/packages/stateful-components/src/outputs/index.tsx
 // to add support for sandboxing using <iframe>
@@ -24,18 +26,35 @@ interface StateProps {
   expanded: boolean;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   outputs: Immutable.List<any>;
+
+  pendingSnapshotRequest: SnapshotRequest;
 }
 
 interface DispatchProps {
   onMetadataChange?: (metadata: JSONObject, mediaType: string, index?: number) => void;
+  storeNotebookSnapshot: (imageSrc: string, requestId: string) => void;
+  storeSnapshotFragment: (cellId: string, snapshotFragment: SnapshotFragment) => void;
+  notebookSnapshotError: (error: string) => void;
 }
 
-export class SandboxOutputs extends React.PureComponent<ComponentProps & StateProps & DispatchProps> {
+type SandboxOutputsProps = ComponentProps & StateProps & DispatchProps;
+
+export class SandboxOutputs extends React.Component<SandboxOutputsProps> {
   private childWindow: Window;
+  private nodeRef = React.createRef<HTMLDivElement>();
+
+  constructor(props: SandboxOutputsProps) {
+    super(props);
+
+    this.state = {
+      processedSnapshotRequest: undefined,
+    };
+  }
 
   render(): JSX.Element {
     // Using min-width to set the width of the iFrame, works around an issue in iOS that can prevent the iFrame from sizing correctly.
-    return (
+    return this.props.outputs && this.props.outputs.size > 0 ? (
+      <div ref={this.nodeRef}>
         <IframeResizer
           checkOrigin={false}
           loading="lazy"
@@ -45,6 +64,9 @@ export class SandboxOutputs extends React.PureComponent<ComponentProps & StatePr
           style={{ height: "1px", width: "1px", minWidth: "100%", border: "none" }}
           sandbox="allow-downloads allow-popups allow-forms allow-pointer-lock allow-scripts allow-popups-to-escape-sandbox"
         />
+      </div>
+    ) : (
+      <></>
     );
   }
 
@@ -76,8 +98,48 @@ export class SandboxOutputs extends React.PureComponent<ComponentProps & StatePr
     this.sendPropsToFrame();
   }
 
-  componentDidUpdate(): void {
+  async componentDidUpdate(prevProps: SandboxOutputsProps): Promise<void> {
     this.sendPropsToFrame();
+
+    if (
+      this.props.pendingSnapshotRequest &&
+      prevProps.pendingSnapshotRequest !== this.props.pendingSnapshotRequest &&
+      this.props.pendingSnapshotRequest.notebookContentRef === this.props.contentRef &&
+      this.nodeRef?.current
+    ) {
+      const boundingClientRect = this.nodeRef.current.getBoundingClientRect();
+
+      try {
+        const { data } = (await postRobot.send(
+          this.childWindow,
+          "snapshotRequest",
+          this.props.pendingSnapshotRequest
+        )) as { data: SnapshotResponse };
+        if (this.props.pendingSnapshotRequest.type === "notebook") {
+          if (data.imageSrc === undefined) {
+            this.props.storeSnapshotFragment(this.props.id, {
+              image: undefined,
+              boundingClientRect: boundingClientRect,
+              requestId: data.requestId,
+            });
+            return;
+          }
+          const image = new Image();
+          image.src = data.imageSrc;
+          image.onload = () => {
+            this.props.storeSnapshotFragment(this.props.id, {
+              image,
+              boundingClientRect: boundingClientRect,
+              requestId: data.requestId,
+            });
+          };
+        } else if (this.props.pendingSnapshotRequest.type === "celloutput") {
+          this.props.storeNotebookSnapshot(data.imageSrc, this.props.pendingSnapshotRequest.requestId);
+        }
+      } catch (error) {
+        this.props.notebookSnapshotError(error.message);
+      }
+    }
   }
 }
 
@@ -85,7 +147,7 @@ export const makeMapStateToProps = (
   initialState: AppState,
   ownProps: ComponentProps
 ): ((state: AppState) => StateProps) => {
-  const mapStateToProps = (state: AppState): StateProps => {
+  const mapStateToProps = (state: CdbAppState): StateProps => {
     let outputs = Immutable.List();
     let hidden = false;
     let expanded = false;
@@ -102,7 +164,17 @@ export const makeMapStateToProps = (
       }
     }
 
-    return { outputs, hidden, expanded };
+    // Determine whether to take a snapshot or not
+    let pendingSnapshotRequest = state.cdb.pendingSnapshotRequest;
+    if (
+      pendingSnapshotRequest &&
+      pendingSnapshotRequest.type === "celloutput" &&
+      pendingSnapshotRequest.cellId !== id
+    ) {
+      pendingSnapshotRequest = undefined;
+    }
+
+    return { outputs, hidden, expanded, pendingSnapshotRequest };
   };
   return mapStateToProps;
 };
@@ -125,6 +197,11 @@ export const makeMapDispatchToProps = (
           })
         );
       },
+      storeSnapshotFragment: (cellId: string, snapshot: SnapshotFragment) =>
+        dispatch(cdbActions.storeCellOutputSnapshot({ cellId, snapshot })),
+      storeNotebookSnapshot: (imageSrc: string, requestId: string) =>
+        dispatch(cdbActions.storeNotebookSnapshot({ imageSrc, requestId })),
+      notebookSnapshotError: (error: string) => dispatch(cdbActions.notebookSnapshotError({ error })),
     };
   };
   return mapDispatchToProps;

src/Explorer/Notebook/NotebookUtil.test.ts

@@ -1,15 +1,15 @@
-import { NotebookUtil } from "./NotebookUtil";
-import * as GitHubUtils from "../../Utils/GitHubUtils";
 import {
-  ImmutableNotebook,
-  MediaBundle,
   CodeCellParams,
-  MarkdownCellParams,
+  ImmutableNotebook,
   makeCodeCell,
   makeMarkdownCell,
   makeNotebookRecord,
+  MarkdownCellParams,
+  MediaBundle,
 } from "@nteract/commutable";
 import { List, Map } from "immutable";
+import * as GitHubUtils from "../../Utils/GitHubUtils";
+import { NotebookUtil } from "./NotebookUtil";
 
 const fileName = "file";
 const notebookName = "file.ipynb";
@@ -131,7 +131,7 @@ describe("NotebookUtil", () => {
   describe("findFirstCodeCellWithDisplay", () => {
     it("works for Notebook file", () => {
       const notebookObject = notebookRecord as ImmutableNotebook;
-      expect(NotebookUtil.findFirstCodeCellWithDisplay(notebookObject)).toEqual(1);
+      expect(NotebookUtil.findCodeCellWithDisplay(notebookObject)[0]).toEqual("1");
     });
   });
 });

src/Explorer/Notebook/NotebookUtil.ts

@@ -1,8 +1,11 @@
+import { ImmutableCodeCell, ImmutableNotebook } from "@nteract/commutable";
+import domtoimage from "dom-to-image";
+import Html2Canvas from "html2canvas";
 import path from "path";
-import { ImmutableNotebook, ImmutableCodeCell } from "@nteract/commutable";
-import { NotebookContentItem, NotebookContentItemType } from "./NotebookContentItem";
-import * as StringUtils from "../../Utils/StringUtils";
 import * as GitHubUtils from "../../Utils/GitHubUtils";
+import * as StringUtils from "../../Utils/StringUtils";
+import { SnapshotFragment } from "./NotebookComponent/types";
+import { NotebookContentItem, NotebookContentItemType } from "./NotebookContentItem";
 
 // Must match rx-jupyter' FileType
 export type FileType = "directory" | "file" | "notebook";
@@ -141,23 +144,175 @@ export class NotebookUtil {
     return `${basePath}${newName}`;
   }
 
-  public static findFirstCodeCellWithDisplay(notebookObject: ImmutableNotebook): number {
+  public static hasCodeCellOutput(cell: ImmutableCodeCell): boolean {
-    let codeCellIndex = 0;
+    return !!cell?.outputs?.find(
-    for (let i = 0; i < notebookObject.cellOrder.size; i++) {
+      (output) =>
-      const cellId = notebookObject.cellOrder.get(i);
+        output.output_type === "display_data" ||
-      if (cellId) {
+        output.output_type === "execute_result" ||
+        output.output_type === "stream"
+    );
+  }
+
+  /**
+   * Find code cells with display
+   * @param notebookObject
+   * @returns array of cell ids
+   */
+  public static findCodeCellWithDisplay(notebookObject: ImmutableNotebook): string[] {
+    return notebookObject.cellOrder.reduce((accumulator: string[], cellId) => {
       const cell = notebookObject.cellMap.get(cellId);
       if (cell?.cell_type === "code") {
-          const displayOutput = (cell as ImmutableCodeCell)?.outputs?.find(
+        if (NotebookUtil.hasCodeCellOutput(cell as ImmutableCodeCell)) {
-            (output) => output.output_type === "display_data" || output.output_type === "execute_result"
+          accumulator.push(cellId);
+        }
+      }
+      return accumulator;
+    }, []);
+  }
+
+  public static takeScreenshotHtml2Canvas = (
+    target: HTMLElement,
+    aspectRatio: number,
+    subSnapshots: SnapshotFragment[],
+    downloadFilename?: string
+  ): Promise<{ imageSrc: string | undefined }> => {
+    return new Promise(async (resolve, reject) => {
+      try {
+        // target.scrollIntoView();
+        const canvas = await Html2Canvas(target, {
+          useCORS: true,
+          allowTaint: true,
+          scale: 1,
+          logging: false,
+        });
+
+        //redraw canvas to fit aspect ratio
+        const originalImageData = canvas.toDataURL();
+        const width = parseInt(canvas.style.width.split("px")[0]);
+        if (aspectRatio) {
+          canvas.height = width * aspectRatio;
+        }
+
+        if (originalImageData === "data:,") {
+          // Empty output
+          resolve({ imageSrc: undefined });
+          return;
+        }
+
+        const context = canvas.getContext("2d");
+        const image = new Image();
+        image.src = originalImageData;
+        image.onload = () => {
+          if (!context) {
+            reject(new Error("No context to draw on"));
+            return;
+          }
+          context.drawImage(image, 0, 0);
+
+          // draw sub images
+          if (subSnapshots) {
+            const parentRect = target.getBoundingClientRect();
+            subSnapshots.forEach((snapshot) => {
+              if (snapshot.image) {
+                context.drawImage(
+                  snapshot.image,
+                  snapshot.boundingClientRect.x - parentRect.x,
+                  snapshot.boundingClientRect.y - parentRect.y
                 );
-          if (displayOutput) {
-            return codeCellIndex;
               }
-          codeCellIndex++;
+            });
           }
+
+          resolve({ imageSrc: canvas.toDataURL() });
+
+          if (downloadFilename) {
+            NotebookUtil.downloadFile(
+              downloadFilename,
+              canvas.toDataURL("image/png").replace("image/png", "image/octet-stream")
+            );
           }
+        };
+      } catch (error) {
+        return reject(error);
       }
-    throw new Error("Output does not exist for any of the cells.");
+    });
+  };
+
+  public static takeScreenshotDomToImage = (
+    target: HTMLElement,
+    aspectRatio: number,
+    subSnapshots: SnapshotFragment[],
+    downloadFilename?: string
+  ): Promise<{ imageSrc?: string }> => {
+    return new Promise(async (resolve, reject) => {
+      // target.scrollIntoView();
+      try {
+        const filter = (node: Node): boolean => {
+          const excludedList = ["IMG", "CANVAS"];
+          return !excludedList.includes((node as HTMLElement).tagName);
+        };
+
+        const originalImageData = await domtoimage.toPng(target, { filter });
+        if (originalImageData === "data:,") {
+          // Empty output
+          resolve({});
+          return;
+        }
+
+        const baseImage = new Image();
+        baseImage.src = originalImageData;
+        baseImage.onload = () => {
+          const canvas = document.createElement("canvas");
+          canvas.width = baseImage.width;
+          canvas.height = aspectRatio !== undefined ? baseImage.width * aspectRatio : baseImage.width;
+
+          const context = canvas.getContext("2d");
+          if (!context) {
+            reject(new Error("No Canvas to draw on"));
+            return;
+          }
+
+          // White background otherwise image is transparent
+          context.fillStyle = "white";
+          context.fillRect(0, 0, baseImage.width, baseImage.height);
+
+          context.drawImage(baseImage, 0, 0);
+
+          // draw sub images
+          if (subSnapshots) {
+            const parentRect = target.getBoundingClientRect();
+            subSnapshots.forEach((snapshot) => {
+              if (snapshot.image) {
+                context.drawImage(
+                  snapshot.image,
+                  snapshot.boundingClientRect.x - parentRect.x,
+                  snapshot.boundingClientRect.y - parentRect.y
+                );
+              }
+            });
+          }
+
+          resolve({ imageSrc: canvas.toDataURL() });
+
+          if (downloadFilename) {
+            NotebookUtil.downloadFile(
+              downloadFilename,
+              canvas.toDataURL("image/png").replace("image/png", "image/octet-stream")
+            );
+          }
+        };
+      } catch (error) {
+        reject(error);
+      }
+    });
+  };
+
+  private static downloadFile(filename: string, content: string): void {
+    const link = document.createElement("a");
+    link.href = content;
+    link.download = filename;
+    document.body.appendChild(link);
+    link.click();
+    document.body.removeChild(link);
   }
 }

src/Explorer/Panes/PanelContainerComponent.test.tsx

@@ -1,5 +1,5 @@
-import React from "react";
 import { shallow } from "enzyme";
+import React from "react";
 import { PanelContainerComponent, PanelContainerProps } from "./PanelContainerComponent";
 
 describe("PaneContainerComponent test", () => {

src/Explorer/Panes/PublishNotebookPane/PublishNotebookPane.test.tsx

@@ -12,13 +12,14 @@ describe("PublishNotebookPaneComponent", () => {
       notebookAuthor: "CosmosDB",
       notebookCreatedDate: "2020-07-17T00:00:00Z",
       notebookObject: undefined,
-      notebookParentDomElement: undefined,
+      notebookContentRef: undefined,
       setNotebookName: undefined,
       setNotebookDescription: undefined,
       setNotebookTags: undefined,
       setImageSrc: undefined,
       onError: undefined,
       clearFormError: undefined,
+      onTakeSnapshot: undefined,
     };
 
     const wrapper = shallow(<PublishNotebookPaneComponent {...props} />);

src/Explorer/Panes/PublishNotebookPane/PublishNotebookPane.tsx

@@ -2,6 +2,7 @@ import { ImmutableNotebook, toJS } from "@nteract/commutable";
 import React, { FunctionComponent, useEffect, useState } from "react";
 import { HttpStatusCodes } from "../../../Common/Constants";
 import { getErrorMessage, getErrorStack, handleError } from "../../../Common/ErrorHandlingUtils";
+import { useNotebookSnapshotStore } from "../../../hooks/useNotebookSnapshotStore";
 import { JunoClient } from "../../../Juno/JunoClient";
 import { Action } from "../../../Shared/Telemetry/TelemetryConstants";
 import { traceFailure, traceStart, traceSuccess } from "../../../Shared/Telemetry/TelemetryProcessor";
@@ -10,6 +11,7 @@ import { CodeOfConductComponent } from "../../Controls/NotebookGallery/CodeOfCon
 import { GalleryTab } from "../../Controls/NotebookGallery/GalleryViewerComponent";
 import Explorer from "../../Explorer";
 import * as FileSystemUtil from "../../Notebook/FileSystemUtil";
+import { SnapshotRequest } from "../../Notebook/NotebookComponent/types";
 import {
   GenericRightPaneComponent,
   GenericRightPaneProps,
@@ -24,7 +26,8 @@ export interface PublishNotebookPaneAProps {
   name: string;
   author: string;
   notebookContent: string | ImmutableNotebook;
-  parentDomElement: HTMLElement;
+  notebookContentRef: string;
+  onTakeSnapshot: (request: SnapshotRequest) => void;
 }
 export const PublishNotebookPane: FunctionComponent<PublishNotebookPaneAProps> = ({
   explorer: container,
@@ -33,7 +36,8 @@ export const PublishNotebookPane: FunctionComponent<PublishNotebookPaneAProps> =
   name,
   author,
   notebookContent,
-  parentDomElement,
+  notebookContentRef,
+  onTakeSnapshot,
 }: PublishNotebookPaneAProps): JSX.Element => {
   const [isCodeOfConductAccepted, setIsCodeOfConductAccepted] = useState<boolean>(false);
   const [content, setContent] = useState<string>("");
@@ -45,6 +49,7 @@ export const PublishNotebookPane: FunctionComponent<PublishNotebookPaneAProps> =
   const [notebookDescription, setNotebookDescription] = useState<string>("");
   const [notebookTags, setNotebookTags] = useState<string>("");
   const [imageSrc, setImageSrc] = useState<string>();
+  const { snapshot: notebookSnapshot, error: notebookSnapshotError } = useNotebookSnapshotStore();
 
   const CodeOfConductAccepted = async () => {
     try {
@@ -74,6 +79,14 @@ export const PublishNotebookPane: FunctionComponent<PublishNotebookPaneAProps> =
     setContent(newContent);
   }, []);
 
+  useEffect(() => {
+    setImageSrc(notebookSnapshot);
+  }, [notebookSnapshot]);
+
+  useEffect(() => {
+    setFormError(notebookSnapshotError);
+  }, [notebookSnapshotError]);
+
   const submit = async (): Promise<void> => {
     const clearPublishingMessage = NotificationConsoleUtils.logConsoleProgress(`Publishing ${name} to gallery`);
     setIsExecuting(true);
@@ -178,13 +191,14 @@ export const PublishNotebookPane: FunctionComponent<PublishNotebookPaneAProps> =
     notebookAuthor: author,
     notebookCreatedDate: new Date().toISOString(),
     notebookObject: notebookObject,
-    notebookParentDomElement: parentDomElement,
+    notebookContentRef,
     onError: createFormError,
     clearFormError: clearFormError,
     setNotebookName,
     setNotebookDescription,
     setNotebookTags,
     setImageSrc,
+    onTakeSnapshot,
   };
   return (
     <GenericRightPaneComponent {...props}>

src/Explorer/Panes/PublishNotebookPane/PublishNotebookPaneComponent.tsx

@@ -1,9 +1,9 @@
 import { Dropdown, IDropdownProps, ITextFieldProps, Stack, Text, TextField } from "@fluentui/react";
 import { ImmutableNotebook } from "@nteract/commutable";
-import Html2Canvas from "html2canvas";
 import React, { FunctionComponent, useState } from "react";
 import { GalleryCardComponent } from "../../Controls/NotebookGallery/Cards/GalleryCardComponent";
 import * as FileSystemUtil from "../../Notebook/FileSystemUtil";
+import { SnapshotRequest } from "../../Notebook/NotebookComponent/types";
 import { NotebookUtil } from "../../Notebook/NotebookUtil";
 import "./styled.less";
 
@@ -11,17 +11,19 @@ export interface PublishNotebookPaneProps {
   notebookName: string;
   notebookAuthor: string;
   notebookTags: string;
-  imageSrc: string;
   notebookDescription: string;
   notebookCreatedDate: string;
   notebookObject: ImmutableNotebook;
-  notebookParentDomElement?: HTMLElement;
+  notebookContentRef: string;
+  imageSrc: string;
+
   onError: (formError: string, formErrorDetail: string, area: string) => void;
   clearFormError: () => void;
   setNotebookName: (newValue: string) => void;
   setNotebookDescription: (newValue: string) => void;
   setNotebookTags: (newValue: string) => void;
   setImageSrc: (newValue: string) => void;
+  onTakeSnapshot: (request: SnapshotRequest) => void;
 }
 
 enum ImageTypes {
@@ -34,18 +36,19 @@ enum ImageTypes {
 export const PublishNotebookPaneComponent: FunctionComponent<PublishNotebookPaneProps> = ({
   notebookName,
   notebookTags,
-  imageSrc,
   notebookDescription,
   notebookAuthor,
   notebookCreatedDate,
   notebookObject,
-  notebookParentDomElement,
+  notebookContentRef,
+  imageSrc,
   onError,
   clearFormError,
   setNotebookName,
   setNotebookDescription,
   setNotebookTags,
   setImageSrc,
+  onTakeSnapshot,
 }: PublishNotebookPaneProps) => {
   const [type, setType] = useState<string>(ImageTypes.CustomImage);
   const CARD_WIDTH = 256;
@@ -63,25 +66,40 @@ export const PublishNotebookPaneComponent: FunctionComponent<PublishNotebookPane
   )}" to the gallery?`;
 
   const options: ImageTypes[] = [ImageTypes.CustomImage, ImageTypes.Url];
+  if (onTakeSnapshot) {
+    options.push(ImageTypes.TakeScreenshot);
+    if (notebookObject) {
+      options.push(ImageTypes.UseFirstDisplayOutput);
+    }
+  }
+
   const thumbnailSelectorProps: IDropdownProps = {
     label: "Cover image",
-    defaultSelectedKey: ImageTypes.CustomImage,
+    selectedKey: type,
     ariaLabel: "Cover image",
     options: options.map((value: string) => ({ text: value, key: value })),
     onChange: async (event, options) => {
       setImageSrc("");
       clearFormError();
       if (options.text === ImageTypes.TakeScreenshot) {
-        try {
+        onTakeSnapshot({
-          await takeScreenshot(notebookParentDomElement, screenshotErrorHandler);
+          aspectRatio: cardHeightToWidthRatio,
-        } catch (error) {
+          requestId: new Date().getTime().toString(),
-          screenshotErrorHandler(error);
+          type: "notebook",
-        }
+          notebookContentRef,
+        });
       } else if (options.text === ImageTypes.UseFirstDisplayOutput) {
-        try {
+        const cellIds = NotebookUtil.findCodeCellWithDisplay(notebookObject);
-          await takeScreenshot(findFirstOutput(), firstOutputErrorHandler);
+        if (cellIds.length > 0) {
-        } catch (error) {
+          onTakeSnapshot({
-          firstOutputErrorHandler(error);
+            aspectRatio: cardHeightToWidthRatio,
+            requestId: new Date().getTime().toString(),
+            type: "celloutput",
+            cellId: cellIds[0],
+            notebookContentRef,
+          });
+        } else {
+          firstOutputErrorHandler(new Error("Output does not exist for any of the cells."));
         }
       }
       setType(options.text);
@@ -97,13 +115,6 @@ export const PublishNotebookPaneComponent: FunctionComponent<PublishNotebookPane
     },
   };
 
-  const screenshotErrorHandler = (error: Error) => {
-    const formError = "Failed to take screen shot";
-    const formErrorDetail = `${error}`;
-    const area = "PublishNotebookPaneComponent/takeScreenshot";
-    onError(formError, formErrorDetail, area);
-  };
-
   const firstOutputErrorHandler = (error: Error) => {
     const formError = "Failed to capture first output";
     const formErrorDetail = `${error}`;
@@ -111,13 +122,6 @@ export const PublishNotebookPaneComponent: FunctionComponent<PublishNotebookPane
     onError(formError, formErrorDetail, area);
   };
 
-  if (notebookParentDomElement) {
-    options.push(ImageTypes.TakeScreenshot);
-    if (notebookObject) {
-      options.push(ImageTypes.UseFirstDisplayOutput);
-    }
-  }
-
   const imageToBase64 = (file: File, updateImageSrc: (result: string) => void) => {
     const reader = new FileReader();
     reader.readAsDataURL(file);
@@ -133,36 +137,6 @@ export const PublishNotebookPaneComponent: FunctionComponent<PublishNotebookPane
     };
   };
 
-  const takeScreenshot = (target: HTMLElement, onError: (error: Error) => void): void => {
-    const updateImageSrcWithScreenshot = (canvasUrl: string): void => {
-      setImageSrc(canvasUrl);
-    };
-
-    target.scrollIntoView();
-    Html2Canvas(target, {
-      useCORS: true,
-      allowTaint: true,
-      scale: 1,
-      logging: true,
-    })
-      .then((canvas) => {
-        //redraw canvas to fit Card Cover Image dimensions
-        const originalImageData = canvas.toDataURL();
-        const requiredHeight = parseInt(canvas.style.width.split("px")[0]) * cardHeightToWidthRatio;
-        canvas.height = requiredHeight;
-        const context = canvas.getContext("2d");
-        const image = new Image();
-        image.src = originalImageData;
-        image.onload = () => {
-          context.drawImage(image, 0, 0);
-          updateImageSrcWithScreenshot(canvas.toDataURL());
-        };
-      })
-      .catch((error) => {
-        onError(error);
-      });
-  };
-
   const renderThumbnailSelectors = (type: string) => {
     switch (type) {
       case ImageTypes.Url:
@@ -198,12 +172,6 @@ export const PublishNotebookPaneComponent: FunctionComponent<PublishNotebookPane
     }
   };
 
-  const findFirstOutput = (): HTMLElement => {
-    const indexOfFirstCodeCellWithDisplay = NotebookUtil.findFirstCodeCellWithDisplay(notebookObject);
-    const cellOutputDomElements = notebookParentDomElement.querySelectorAll<HTMLElement>(".nteract-cell-outputs");
-    return cellOutputDomElements[indexOfFirstCodeCellWithDisplay];
-  };
-
   return (
     <div className="publishNotebookPanelContent">
       <Stack className="panelMainContent" tokens={{ childrenGap: 20 }}>

src/Explorer/Panes/PublishNotebookPane/__snapshots__/PublishNotebookPane.test.tsx.snap

@@ -52,7 +52,6 @@ exports[`PublishNotebookPaneComponent renders 1`] = `
     <StackItem>
       <Dropdown
         ariaLabel="Cover image"
-        defaultSelectedKey="Custom Image"
         label="Cover image"
         onChange={[Function]}
         options={
@@ -67,6 +66,7 @@ exports[`PublishNotebookPaneComponent renders 1`] = `
             },
           ]
         }
+        selectedKey="Custom Image"
       />
     </StackItem>
     <StackItem>

src/Explorer/Tabs/NotebookV2Tab.ts

@@ -15,6 +15,7 @@ import SaveIcon from "../../../images/save-cosmos.svg";
 import { ArmApiVersions } from "../../Common/Constants";
 import { configContext } from "../../ConfigContext";
 import * as DataModels from "../../Contracts/DataModels";
+import { useNotebookSnapshotStore } from "../../hooks/useNotebookSnapshotStore";
 import { trackEvent } from "../../Shared/appInsights";
 import { Action, ActionModifiers, Source } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
@@ -24,7 +25,9 @@ import { logConsoleInfo } from "../../Utils/NotificationConsoleUtils";
 import { CommandButtonComponentProps } from "../Controls/CommandButton/CommandButtonComponent";
 import * as CommandBarComponentButtonFactory from "../Menus/CommandBar/CommandBarComponentButtonFactory";
 import { KernelSpecsDisplay } from "../Notebook/NotebookClientV2";
+import * as CdbActions from "../Notebook/NotebookComponent/actions";
 import { NotebookComponentAdapter } from "../Notebook/NotebookComponent/NotebookComponentAdapter";
+import { CdbAppState, SnapshotRequest } from "../Notebook/NotebookComponent/types";
 import { NotebookContentItem } from "../Notebook/NotebookContentItem";
 import NotebookTabBase, { NotebookTabBaseOptions } from "./NotebookTabBase";
 
@@ -458,11 +461,32 @@ export default class NotebookTabV2 extends NotebookTabBase {
       source: Source.CommandBarMenu,
     });
 
+    const notebookReduxStore = NotebookTabV2.clientManager.getStore();
+    const unsubscribe = notebookReduxStore.subscribe(() => {
+      const cdbState = (notebookReduxStore.getState() as CdbAppState).cdb;
+      useNotebookSnapshotStore.setState({
+        snapshot: cdbState.notebookSnapshot?.imageSrc,
+        error: cdbState.notebookSnapshotError,
+      });
+    });
+
     const notebookContent = this.notebookComponentAdapter.getContent();
+    const notebookContentRef = this.notebookComponentAdapter.contentRef;
+    const onPanelClose = (): void => {
+      unsubscribe();
+      useNotebookSnapshotStore.setState({
+        snapshot: undefined,
+        error: undefined,
+      });
+      notebookReduxStore.dispatch(CdbActions.takeNotebookSnapshot(undefined));
+    };
+
     await this.container.publishNotebook(
       notebookContent.name,
       notebookContent.content,
-      this.notebookComponentAdapter.getNotebookParentElement()
+      notebookContentRef,
+      (request: SnapshotRequest) => notebookReduxStore.dispatch(CdbActions.takeNotebookSnapshot(request)),
+      onPanelClose
     );
   };
 

src/hooks/useNotebookSnapshotStore.ts

@@ -0,0 +1,15 @@
+import create, { UseStore } from "zustand";
+
+export interface NotebookSnapshotHooks {
+  snapshot: string;
+  error: string;
+  setSnapshot: (imageSrc: string) => void;
+  setError: (error: string) => void;
+}
+
+export const useNotebookSnapshotStore: UseStore<NotebookSnapshotHooks> = create((set) => ({
+  snapshot: undefined,
+  error: undefined,
+  setSnapshot: (imageSrc: string) => set((state) => ({ ...state, snapshot: imageSrc })),
+  setError: (error: string) => set((state) => ({ ...state, error })),
+}));

src/hooks/useSidePanel.ts

@@ -4,7 +4,7 @@ export interface SidePanelHooks {
   isPanelOpen: boolean;
   panelContent: JSX.Element;
   headerText: string;
-  openSidePanel: (headerText: string, panelContent: JSX.Element) => void;
+  openSidePanel: (headerText: string, panelContent: JSX.Element, onClose?: () => void) => void;
   closeSidePanel: () => void;
 }
 
@@ -12,17 +12,23 @@ export const useSidePanel = (): SidePanelHooks => {
   const [isPanelOpen, setIsPanelOpen] = useState<boolean>(false);
   const [panelContent, setPanelContent] = useState<JSX.Element>();
   const [headerText, setHeaderText] = useState<string>();
+  const [onCloseCallback, setOnCloseCallback] = useState<{ callback: () => void }>();
 
-  const openSidePanel = (headerText: string, panelContent: JSX.Element): void => {
+  const openSidePanel = (headerText: string, panelContent: JSX.Element, onClose?: () => void): void => {
     setHeaderText(headerText);
     setPanelContent(panelContent);
     setIsPanelOpen(true);
+    setOnCloseCallback({ callback: onClose });
   };
 
   const closeSidePanel = (): void => {
     setHeaderText("");
     setPanelContent(undefined);
     setIsPanelOpen(false);
+    if (onCloseCallback) {
+      onCloseCallback.callback();
+      setOnCloseCallback(undefined);
+    }
   };
 
   return { isPanelOpen, panelContent, headerText, openSidePanel, closeSidePanel };