From 20a69b711a4c2a3bd315b6c2006cee953f3efee4 Mon Sep 17 00:00:00 2001
From: Tom Moor <tom.moor@gmail.com>
Date: Mon, 4 Apr 2022 19:04:28 -0700
Subject: [PATCH] fix: Some spots where navigation state was not preserved fix:
 Collection in main nav pops open when moving from starred collection

---
 .../Sidebar/components/CollectionLink.tsx       |  2 +-
 .../components/DraggableCollectionLink.tsx      | 17 +++--------------
 app/scenes/Collection.tsx                       |  6 ++++--
 app/scenes/Document/components/Document.tsx     |  2 +-
 app/stores/UiStore.ts                           |  5 ++---
 5 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/app/components/Sidebar/components/CollectionLink.tsx b/app/components/Sidebar/components/CollectionLink.tsx
index 23ef51dfd..aecabab02 100644
--- a/app/components/Sidebar/components/CollectionLink.tsx
+++ b/app/components/Sidebar/components/CollectionLink.tsx
@@ -54,7 +54,7 @@ const CollectionLink: React.FC<Props> = ({
       await collection.save({
         name,
       });
-      history.push(collection.url);
+      history.replace(collection.url, history.location.state);
     },
     [collection, history]
   );
diff --git a/app/components/Sidebar/components/DraggableCollectionLink.tsx b/app/components/Sidebar/components/DraggableCollectionLink.tsx
index 00c50cb62..3815129e8 100644
--- a/app/components/Sidebar/components/DraggableCollectionLink.tsx
+++ b/app/components/Sidebar/components/DraggableCollectionLink.tsx
@@ -13,7 +13,6 @@ import CollectionLinkChildren from "./CollectionLinkChildren";
 import DropCursor from "./DropCursor";
 import Relative from "./Relative";
 import { DragObject } from "./SidebarLink";
-import { useStarredContext } from "./StarredContext";
 
 type Props = {
   collection: Collection;
@@ -37,10 +36,8 @@ function DraggableCollectionLink({
 }: Props) {
   const locationStateStarred = useLocationStateStarred();
   const { ui, collections } = useStores();
-  const inStarredSection = useStarredContext();
   const [expanded, setExpanded] = React.useState(
-    collection.id === ui.activeCollectionId &&
-      locationStateStarred === inStarredSection
+    collection.id === ui.activeCollectionId && !locationStateStarred
   );
   const can = usePolicy(collection.id);
   const belowCollectionIndex = belowCollection ? belowCollection.index : null;
@@ -88,18 +85,10 @@ function DraggableCollectionLink({
   // If the current collection is active and relevant to the sidebar section we
   // are in then expand it automatically
   React.useEffect(() => {
-    if (
-      collection.id === ui.activeCollectionId &&
-      locationStateStarred === inStarredSection
-    ) {
+    if (collection.id === ui.activeCollectionId && !locationStateStarred) {
       setExpanded(true);
     }
-  }, [
-    collection.id,
-    ui.activeCollectionId,
-    locationStateStarred,
-    inStarredSection,
-  ]);
+  }, [collection.id, ui.activeCollectionId, locationStateStarred]);
 
   const handleDisclosureClick = React.useCallback((ev) => {
     ev.preventDefault();
diff --git a/app/scenes/Collection.tsx b/app/scenes/Collection.tsx
index dad856073..d4242e06f 100644
--- a/app/scenes/Collection.tsx
+++ b/app/scenes/Collection.tsx
@@ -55,15 +55,17 @@ function CollectionScene() {
       const canonicalUrl = updateCollectionUrl(match.url, collection);
 
       if (match.url !== canonicalUrl) {
-        history.replace(canonicalUrl);
+        history.replace(canonicalUrl, history.location.state);
       }
     }
   }, [collection, collection?.name, history, id, match.url]);
 
   React.useEffect(() => {
     if (collection) {
-      ui.setActiveCollection(collection);
+      ui.setActiveCollection(collection.id);
     }
+
+    return () => ui.setActiveCollection(undefined);
   }, [ui, collection]);
 
   React.useEffect(() => {
diff --git a/app/scenes/Document/components/Document.tsx b/app/scenes/Document/components/Document.tsx
index a14231ad2..c9d233d97 100644
--- a/app/scenes/Document/components/Document.tsx
+++ b/app/scenes/Document/components/Document.tsx
@@ -200,7 +200,7 @@ class DocumentScene extends React.Component<Props> {
     if (response) {
       this.replaceDocument(response.data);
       toasts.showToast(t("Document restored"));
-      history.replace(this.props.document.url);
+      history.replace(this.props.document.url, history.location.state);
     }
   };
 
diff --git a/app/stores/UiStore.ts b/app/stores/UiStore.ts
index b93a7d86c..79c1eedd1 100644
--- a/app/stores/UiStore.ts
+++ b/app/stores/UiStore.ts
@@ -1,6 +1,5 @@
 import { action, autorun, computed, observable } from "mobx";
 import { light as defaultTheme } from "@shared/theme";
-import Collection from "~/models/Collection";
 import Document from "~/models/Document";
 import { ConnectionStatus } from "~/scenes/Document/components/MultiplayerEditor";
 
@@ -150,8 +149,8 @@ class UiStore {
   };
 
   @action
-  setActiveCollection = (collection: Collection): void => {
-    this.activeCollectionId = collection.id;
+  setActiveCollection = (collectionId: string | undefined): void => {
+    this.activeCollectionId = collectionId;
   };
 
   @action