Remove removed annotations from additional selection collections

src/sidebar/store/modules/selection.js

@@ -266,12 +266,6 @@ const update = {
   },
 
   REMOVE_ANNOTATIONS: function (state, action) {
-    const selection = { ...state.selected };
-    action.annotationsToRemove.forEach(annotation => {
-      if (annotation.id) {
-        delete selection[annotation.id];
-      }
-    });
     let newTab = state.selectedTab;
     // If the orphans tab is selected but no remaining annotations are orphans,
     // switch back to annotations tab
@@ -281,9 +275,23 @@ const update = {
     ) {
       newTab = uiConstants.TAB_ANNOTATIONS;
     }
+
+    const removeAnns = collection => {
+      action.annotationsToRemove.forEach(annotation => {
+        if (annotation.id) {
+          delete collection[annotation.id];
+        }
+        if (annotation.$tag) {
+          delete collection[annotation.$tag];
+        }
+      });
+      return collection;
+    };
     return {
       ...setTab(newTab, state.selectedTab),
-      selected: selection,
+      expanded: removeAnns({ ...state.expanded }),
+      forcedVisible: removeAnns({ ...state.forcedVisible }),
+      selected: removeAnns({ ...state.selected }),
     };
   },
 };

src/sidebar/store/modules/test/selection-test.js

@@ -271,11 +271,17 @@ describe('sidebar/store/modules/selection', () => {
   describe('#REMOVE_ANNOTATIONS', function () {
     it('removing an annotation should also remove it from the selection', function () {
       store.selectAnnotations([1, 2, 3]);
+      store.setForcedVisible(2, true);
+      store.setForcedVisible(1, true);
+      store.setExpanded(1, true);
+      store.setExpanded(2, true);
       store.removeAnnotations([{ id: 2 }]);
       assert.deepEqual(getSelectionState().selected, {
         1: true,
         3: true,
       });
+      assert.deepEqual(store.forcedVisibleAnnotations(), ['1']);
+      assert.deepEqual(store.expandedMap(), { 1: true });
     });
   });