Fix “Back” button sometimes redirecting out of Mastodon

app/javascript/mastodon/components/column_back_button.jsx

@@ -23,9 +23,7 @@ export default class ColumnBackButton extends PureComponent {
 
     if (onClick) {
       onClick();
-    // Check if there is a previous page in the app to go back to per https://stackoverflow.com/a/70532858/9703201
-    // When upgrading to V6, check `location.key !== 'default'` instead per https://github.com/remix-run/history/blob/main/docs/api-reference.md#location
-    } else if (router.route.location.key) {
+    } else if (router.history.location?.state?.fromMastodon) {
       router.history.goBack();
     } else {
       router.history.push('/');

app/javascript/mastodon/components/column_header.jsx

@@ -63,10 +63,12 @@ class ColumnHeader extends PureComponent {
   };
 
   handleBackClick = () => {
-    if (window.history && window.history.state) {
-      this.context.router.history.goBack();
+    const { router } = this.context;
+
+    if (router.history.location?.state?.fromMastodon) {
+      router.history.goBack();
     } else {
-      this.context.router.history.push('/');
+      router.history.push('/');
     }
   };
 
@@ -83,6 +85,7 @@ class ColumnHeader extends PureComponent {
   };
 
   render () {
+    const { router } = this.context;
     const { title, icon, active, children, pinned, multiColumn, extraButton, showBackButton, intl: { formatMessage }, placeholder, appendContent, collapseIssues } = this.props;
     const { collapsed, animating } = this.state;
 
@@ -126,7 +129,7 @@ class ColumnHeader extends PureComponent {
       pinButton = <button key='pin-button' className='text-btn column-header__setting-btn' onClick={this.handlePin}><Icon id='plus' /> <FormattedMessage id='column_header.pin' defaultMessage='Pin' /></button>;
     }
 
-    if (!pinned && (multiColumn || showBackButton)) {
+    if (!pinned && ((multiColumn && router.history.location?.state?.fromMastodon) || showBackButton)) {
       backButton = (
         <button onClick={this.handleBackClick} className='column-header__back-button'>
           <Icon id='chevron-left' className='column-back-button__icon' fixedWidth />

app/javascript/mastodon/components/router.tsx

@@ -1,23 +1,46 @@
 import type { PropsWithChildren } from 'react';
 import React from 'react';
 
-import type { History } from 'history';
 import { createBrowserHistory } from 'history';
 import { Router as OriginalRouter } from 'react-router';
 
 import { layoutFromWindow } from 'mastodon/is_mobile';
 
-const browserHistory = createBrowserHistory();
+interface MastodonLocationState {
+  fromMastodon?: boolean;
+  mastodonModalKey?: string;
+}
+
+const browserHistory = createBrowserHistory<
+  MastodonLocationState | undefined
+>();
 const originalPush = browserHistory.push.bind(browserHistory);
+const originalReplace = browserHistory.replace.bind(browserHistory);
+
+browserHistory.push = (path: string, state?: MastodonLocationState) => {
+  state = state ?? {};
+  state.fromMastodon = true;
 
-browserHistory.push = (path: string, state: History.LocationState) => {
   if (layoutFromWindow() === 'multi-column' && !path.startsWith('/deck')) {
     originalPush(`/deck${path}`, state);
   } else {
     originalPush(path, state);
   }
 };
 
+browserHistory.replace = (path: string, state?: MastodonLocationState) => {
+  if (browserHistory.location.state?.fromMastodon) {
+    state = state ?? {};
+    state.fromMastodon = true;
+  }
+
+  if (layoutFromWindow() === 'multi-column' && !path.startsWith('/deck')) {
+    originalReplace(`/deck${path}`, state);
+  } else {
+    originalReplace(path, state);
+  }
+};
+
 export const Router: React.FC<PropsWithChildren> = ({ children }) => {
   return <OriginalRouter history={browserHistory}>{children}</OriginalRouter>;
 };

app/javascript/mastodon/features/ui/index.jsx

@@ -479,10 +479,12 @@ class UI extends PureComponent {
   };
 
   handleHotkeyBack = () => {
-    if (window.history && window.history.state) {
-      this.context.router.history.goBack();
+    const { router } = this.context;
+
+    if (router.history.location?.state?.fromMastodon) {
+      router.history.goBack();
     } else {
-      this.context.router.history.push('/');
+      router.history.push('/');
     }
   };