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)
2021-05-11 20:24:05 +02:00 · 2021-05-11 20:24:05 +02:00 · 861042c27e
parent 4ed8fe9e7d
commit 861042c27e
24 changed files with 683 additions and 222 deletions
--- a/less/Common/Constants.less
+++ b/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;
--- a/package-lock.json
+++ b/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=="
    }
  }
 }
--- a/package.json
+++ b/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",
--- a/src/CellOutputViewer/CellOutputViewer.tsx
+++ b/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
--- a/src/Explorer/Explorer.tsx
+++ b/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);
    }
  }
--- a/src/Explorer/Notebook/NotebookComponent/NotebookComponentAdapter.tsx
+++ b/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 => {
--- a/src/Explorer/Notebook/NotebookComponent/NotebookComponentBootstrapper.tsx
+++ b/src/Explorer/Notebook/NotebookComponent/NotebookComponentBootstrapper.tsx
@ -1,35 +1,29 @@
-import * as React from "react";
-
-import { NotebookComponent } from "./NotebookComponent";
-import { NotebookClientV2 } from "../NotebookClientV2";
-import { NotebookUtil } from "../NotebookUtil";
-
+import { CellId, CellType, ImmutableNotebook } from "@nteract/commutable";
 // 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 * as CdbActions from "./actions";
+import { AnyAction, Store } from "redux";
+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;
--- a/src/Explorer/Notebook/NotebookComponent/actions.ts
+++ b/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 interface UpdateNotebookParentDomEltAction {
-  type: "UPDATE_NOTEBOOK_PARENT_DOM_ELTS";
+export const STORE_CELL_OUTPUT_SNAPSHOT = "STORE_CELL_OUTPUT_SNAPSHOT";
+export interface StoreCellOutputSnapshotAction {
+  type: "STORE_CELL_OUTPUT_SNAPSHOT";
  payload: {
-    contentRef: ContentRef;
-    parentElt: HTMLElement;
+    cellId: string;
+    snapshot: SnapshotFragment;
  };
 }

-export const UpdateNotebookParentDomElt = (payload: {
-  contentRef: ContentRef;
-  parentElt: HTMLElement;
-}): UpdateNotebookParentDomEltAction => {
+export const storeCellOutputSnapshot = (payload: {
+  cellId: string;
+  snapshot: SnapshotFragment;
+}): 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,
  };
 };
--- a/src/Explorer/Notebook/NotebookComponent/reducers.ts
+++ b/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: {
-      const typedAction = action as cdbActions.UpdateNotebookParentDomEltAction;
-      var parentEltsMap = state.get("currentNotebookParentElements");
-      const contentRef = typedAction.payload.contentRef;
-      const parentElt = typedAction.payload.parentElt;
-      if (parentElt) {
-        parentEltsMap.set(contentRef, parentElt);
-      } else {
-        parentEltsMap.delete(contentRef);
+    case cdbActions.STORE_CELL_OUTPUT_SNAPSHOT: {
+      const typedAction = action as cdbActions.StoreCellOutputSnapshotAction;
+      state.cellOutputSnapshots.set(typedAction.payload.cellId, typedAction.payload.snapshot);
+      // TODO Simpler datastructure to instantiate new Map?
+      return state.set("cellOutputSnapshots", new Map(state.cellOutputSnapshots));
    }
-      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;
--- a/src/Explorer/Notebook/NotebookComponent/types.ts
+++ b/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,
 });
--- a/src/Explorer/Notebook/NotebookManager.tsx
+++ b/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
    );
  }

--- a/src/Explorer/Notebook/NotebookRenderer/NotebookRenderer.tsx
+++ b/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() {
-    this.props.updateNotebookParentDomElt(this.props.contentRef, this.notebookRendererRef.current);
+  async componentDidUpdate(): Promise<void> {
+    // 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 }) => {
-        return dispatch(
+      addTransform: (transform: React.ComponentType & { MIMETYPE: string }) =>
+        dispatch(
          actions.addTransform({
            mediaType: transform.MIMETYPE,
            component: transform,
          })
-        );
-      },
-      updateNotebookParentDomElt: (contentRef: ContentRef, parentElt: HTMLElement) => {
-        return dispatch(
-          cdbActions.UpdateNotebookParentDomElt({
-            contentRef,
-            parentElt,
-          })
-        );
-      },
+        ),
+      storeNotebookSnapshot: (imageSrc: string, requestId: string) =>
+        dispatch(cdbActions.storeNotebookSnapshot({ imageSrc, requestId })),
+      notebookSnapshotError: (error: string) => dispatch(cdbActions.notebookSnapshotError({ error })),
    };
  };
  return mapDispatchToProps;
 };

-export default connect(null, makeMapDispatchToProps)(BaseNotebookRenderer);
+export default connect(makeMapStateToProps, makeMapDispatchToProps)(BaseNotebookRenderer);
--- a/src/Explorer/Notebook/NotebookRenderer/Toolbar.tsx
+++ b/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 })
-      .cell_type;
+    const cell = selectors.cell.cellFromState(state, { id: ownProps.id, contentRef: ownProps.contentRef });
+    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;
--- a/src/Explorer/Notebook/NotebookRenderer/outputs/SandboxOutputs.tsx
+++ b/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;
--- a/src/Explorer/Notebook/NotebookUtil.test.ts
+++ b/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");
    });
  });
 });
--- a/src/Explorer/Notebook/NotebookUtil.ts
+++ b/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 {
-    let codeCellIndex = 0;
-    for (let i = 0; i < notebookObject.cellOrder.size; i++) {
-      const cellId = notebookObject.cellOrder.get(i);
-      if (cellId) {
+  public static hasCodeCellOutput(cell: ImmutableCodeCell): boolean {
+    return !!cell?.outputs?.find(
+      (output) =>
+        output.output_type === "display_data" ||
+        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(
-            (output) => output.output_type === "display_data" || output.output_type === "execute_result"
+        if (NotebookUtil.hasCodeCellOutput(cell as ImmutableCodeCell)) {
+          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);
  }
 }
--- a/src/Explorer/Panes/PanelContainerComponent.test.tsx
+++ b/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", () => {
--- a/src/Explorer/Panes/PublishNotebookPane/PublishNotebookPane.test.tsx
+++ b/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} />);
--- a/src/Explorer/Panes/PublishNotebookPane/PublishNotebookPane.tsx
+++ b/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}>
--- a/src/Explorer/Panes/PublishNotebookPane/PublishNotebookPaneComponent.tsx
+++ b/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 {
-          await takeScreenshot(notebookParentDomElement, screenshotErrorHandler);
-        } catch (error) {
-          screenshotErrorHandler(error);
-        }
+        onTakeSnapshot({
+          aspectRatio: cardHeightToWidthRatio,
+          requestId: new Date().getTime().toString(),
+          type: "notebook",
+          notebookContentRef,
+        });
      } else if (options.text === ImageTypes.UseFirstDisplayOutput) {
-        try {
-          await takeScreenshot(findFirstOutput(), firstOutputErrorHandler);
-        } catch (error) {
-          firstOutputErrorHandler(error);
+        const cellIds = NotebookUtil.findCodeCellWithDisplay(notebookObject);
+        if (cellIds.length > 0) {
+          onTakeSnapshot({
+            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 }}>
--- a/src/Explorer/Panes/PublishNotebookPane/snapshots/PublishNotebookPane.test.tsx.snap
+++ b/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>
--- a/src/Explorer/Tabs/NotebookV2Tab.ts
+++ b/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
    );
  };

--- a/src/hooks/useNotebookSnapshotStore.ts
+++ b/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 })),
+}));
--- a/src/hooks/useSidePanel.ts
+++ b/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 };