diff --git a/src/lib/AnnotationControls.ts b/src/lib/AnnotationControls.ts
index d22f6a123c..4e09cc7855 100644
--- a/src/lib/AnnotationControls.ts
+++ b/src/lib/AnnotationControls.ts
@@ -16,12 +16,11 @@ export enum AnnotationMode {
     NONE = 'none',
     REGION = 'region',
 }
-export type ClickHandler = ({ event }: { event: MouseEvent }) => void;
+export type ClickHandler = ({ event, mode }: { event: MouseEvent; mode: AnnotationMode }) => void;
 export type Options = {
     fileId: string;
+    onClick?: ClickHandler;
     onEscape?: () => void;
-    onHighlightClick?: ClickHandler;
-    onRegionClick?: ClickHandler;
     showHighlightText: boolean;
 };
 
@@ -65,6 +64,8 @@ export default class AnnotationControls {
 
     private hasInit = false;
 
+    private isModeTemp = false;
+
     private onEscape: () => void = noop;
 
     /**
@@ -144,7 +145,7 @@ export default class AnnotationControls {
      * Set the mode. If the mode is different from what is currently saved in state,
      * then reset the current controls and apply the active state based on the provided mode.
      */
-    public setMode(mode: AnnotationMode): void {
+    public setTempMode(mode: AnnotationMode): void {
         // Only update buttons if mode has changed
         if (this.currentMode === mode) {
             return;
@@ -153,6 +154,16 @@ export default class AnnotationControls {
         this.resetControls();
         this.currentMode = mode;
         this.updateButton(mode);
+
+        this.isModeTemp = mode !== AnnotationMode.NONE;
+    }
+
+    public resetTempMode(): void {
+        if (!this.isModeTemp) {
+            return;
+        }
+
+        this.setTempMode(AnnotationMode.NONE);
     }
 
     /**
@@ -161,13 +172,14 @@ export default class AnnotationControls {
     private handleClick = (onClick: ClickHandler, mode: AnnotationMode) => (event: MouseEvent): void => {
         const prevMode = this.currentMode;
         this.resetControls();
+        this.isModeTemp = false;
 
         if (prevMode !== mode) {
             this.currentMode = mode as AnnotationMode;
             this.updateButton(mode);
         }
 
-        onClick({ event });
+        onClick({ event, mode: this.currentMode });
     };
 
     /**
@@ -210,22 +222,16 @@ export default class AnnotationControls {
     /**
      * Initialize the annotation controls with options.
      */
-    public init({
-        fileId,
-        onEscape = noop,
-        onRegionClick = noop,
-        onHighlightClick = noop,
-        showHighlightText = false,
-    }: Options): void {
+    public init({ fileId, onEscape = noop, onClick = noop, showHighlightText = false }: Options): void {
         if (this.hasInit) {
             return;
         }
         const groupElement = this.controls.addGroup(CLASS_ANNOTATIONS_GROUP);
         groupElement.setAttribute('data-resin-feature', 'annotations');
 
-        this.addButton(AnnotationMode.REGION, onRegionClick, groupElement, fileId);
+        this.addButton(AnnotationMode.REGION, onClick, groupElement, fileId);
         if (showHighlightText) {
-            this.addButton(AnnotationMode.HIGHLIGHT, onHighlightClick, groupElement, fileId);
+            this.addButton(AnnotationMode.HIGHLIGHT, onClick, groupElement, fileId);
         }
 
         this.onEscape = onEscape;
diff --git a/src/lib/viewers/BaseViewer.js b/src/lib/viewers/BaseViewer.js
index 91329f5b19..f7ebe093fa 100644
--- a/src/lib/viewers/BaseViewer.js
+++ b/src/lib/viewers/BaseViewer.js
@@ -151,12 +151,11 @@ class BaseViewer extends EventEmitter {
         this.mobileZoomEndHandler = this.mobileZoomEndHandler.bind(this);
         this.handleAnnotatorEvents = this.handleAnnotatorEvents.bind(this);
         this.handleAnnotationCreateEvent = this.handleAnnotationCreateEvent.bind(this);
-        this.handleAnnotationModeChangeEvent = this.handleAnnotationModeChangeEvent.bind(this);
+        this.handleAnnotationControlsClick = this.handleAnnotationControlsClick.bind(this);
         this.handleAnnotationControlsEscape = this.handleAnnotationControlsEscape.bind(this);
+        this.handleAnnotationStagedChangeEvent = this.handleAnnotationStagedChangeEvent.bind(this);
         this.handleFullscreenEnter = this.handleFullscreenEnter.bind(this);
         this.handleFullscreenExit = this.handleFullscreenExit.bind(this);
-        this.handleHighlightClick = this.handleHighlightClick.bind(this);
-        this.handleRegionClick = this.handleRegionClick.bind(this);
         this.createAnnotator = this.createAnnotator.bind(this);
         this.viewerLoadHandler = this.viewerLoadHandler.bind(this);
         this.initAnnotations = this.initAnnotations.bind(this);
@@ -1016,7 +1015,7 @@ class BaseViewer extends EventEmitter {
 
         if (this.areNewAnnotationsEnabled() && this.annotationControls) {
             this.annotator.addListener('annotations_create', this.handleAnnotationCreateEvent);
-            this.annotator.addListener('annotations_mode_change', this.handleAnnotationModeChangeEvent);
+            this.annotator.addListener('annotations_staged_change', this.handleAnnotationStagedChangeEvent);
         }
     }
 
@@ -1071,23 +1070,14 @@ class BaseViewer extends EventEmitter {
     }
 
     /**
-     * Handler for annotation toolbar highlight text button click event.
+     * Handler for annotation controls button click event.
      *
      * @private
+     * @param {AnnotationMode} mode one of annotation controls modes
      * @return {void}
      */
-    handleHighlightClick() {
-        this.annotator.toggleAnnotationMode(AnnotationMode.HIGHLIGHT);
-    }
-
-    /**
-     * Handler for annotation toolbar region comment button click event.
-     *
-     * @private
-     * @return {void}
-     */
-    handleRegionClick() {
-        this.annotator.toggleAnnotationMode(AnnotationMode.REGION);
+    handleAnnotationControlsClick({ mode }) {
+        this.annotator.toggleAnnotationMode(mode);
     }
 
     /**
@@ -1259,12 +1249,21 @@ class BaseViewer extends EventEmitter {
         // we remain in create mode
         if (status === 'success') {
             this.annotator.emit('annotations_active_set', id);
+
+            if (this.annotationControls) {
+                this.annotationControls.resetTempMode();
+            }
         }
     }
 
-    handleAnnotationModeChangeEvent({ mode }) {
-        if (this.annotationControls) {
-            this.annotationControls.setMode(mode);
+    handleAnnotationStagedChangeEvent({ status, type }) {
+        if (!this.annotationControls) {
+            return;
+        }
+        if (!status || status === 'cancel') {
+            this.annotationControls.resetTempMode();
+        } else {
+            this.annotationControls.setTempMode(type);
         }
     }
 
diff --git a/src/lib/viewers/doc/DocBaseViewer.js b/src/lib/viewers/doc/DocBaseViewer.js
index 2436cfa6c8..034ebb1ba0 100644
--- a/src/lib/viewers/doc/DocBaseViewer.js
+++ b/src/lib/viewers/doc/DocBaseViewer.js
@@ -1101,8 +1101,7 @@ class DocBaseViewer extends BaseViewer {
             this.annotationControls.init({
                 fileId: this.options.file.id,
                 onEscape: this.handleAnnotationControlsEscape,
-                onHighlightClick: this.handleHighlightClick,
-                onRegionClick: this.handleRegionClick,
+                onClick: this.handleAnnotationControlsClick,
                 showHighlightText: this.options.showAnnotationsHighlightText,
             });
         }