From 9d538f97bfbb8995aca36f5f9446befd8111c53d Mon Sep 17 00:00:00 2001
From: Nico <Nico>
Date: Wed, 3 Jun 2026 10:39:15 -0700
Subject: [PATCH] fix: preserve map return context

---
 frontend/src/components/store/StoreTabs.jsx |  9 +++++++--
 frontend/src/pages/LocationMapManager.jsx   | 14 ++++++++++++--
 frontend/tests/store-selector.spec.ts       |  9 +++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/frontend/src/components/store/StoreTabs.jsx b/frontend/src/components/store/StoreTabs.jsx
index 51812fd..4ea6904 100644
--- a/frontend/src/components/store/StoreTabs.jsx
+++ b/frontend/src/components/store/StoreTabs.jsx
@@ -1,5 +1,5 @@
 import { useContext, useMemo, useState } from 'react';
-import { useNavigate } from 'react-router-dom';
+import { useLocation, useNavigate } from 'react-router-dom';
 import { StoreContext } from '../../context/StoreContext';
 import '../../styles/components/StoreTabs.css';
 
@@ -55,6 +55,7 @@ function buildStoreOptions(stores) {
 export default function StoreTabs() {
   const { stores, activeStore, setActiveStore, loading } = useContext(StoreContext);
   const navigate = useNavigate();
+  const routeLocation = useLocation();
   const [activePicker, setActivePicker] = useState(null);
   const storeOptions = useMemo(() => buildStoreOptions(stores || []), [stores]);
   const activeStoreKey = getStoreKey(activeStore);
@@ -99,7 +100,11 @@ export default function StoreTabs() {
 
   const handleManageMap = () => {
     if (!selectedLocation?.id || !selectedStoreOption?.key) return;
-    navigate(`/stores/${selectedStoreOption.key}/locations/${selectedLocation.id}/map`);
+    navigate(`/stores/${selectedStoreOption.key}/locations/${selectedLocation.id}/map`, {
+      state: {
+        returnTo: `${routeLocation.pathname}${routeLocation.search}${routeLocation.hash}`,
+      },
+    });
   };
 
   return (
diff --git a/frontend/src/pages/LocationMapManager.jsx b/frontend/src/pages/LocationMapManager.jsx
index 6efe2f6..c6f641a 100644
--- a/frontend/src/pages/LocationMapManager.jsx
+++ b/frontend/src/pages/LocationMapManager.jsx
@@ -1,5 +1,5 @@
 import { useCallback, useContext, useEffect, useRef, useState } from "react";
-import { useNavigate, useParams } from "react-router-dom";
+import { useLocation, useNavigate, useParams } from "react-router-dom";
 import {
   createBlankLocationMap,
   createLocationMapFromZones,
@@ -38,6 +38,8 @@ import "../styles/pages/LocationMapManager.css";
 export default function LocationMapManager() {
   const { storeId, locationId } = useParams();
   const navigate = useNavigate();
+  const routeLocation = useLocation();
+  const returnPath = routeLocation.state?.returnTo;
   const toast = useActionToast();
   const { username } = useContext(AuthContext);
   const { activeHousehold, loading: householdLoading, hasLoaded } = useContext(HouseholdContext);
@@ -400,6 +402,14 @@ export default function LocationMapManager() {
     setZoom(Number(nextZoom.toFixed(2)));
   };
 
+  const handleBack = useCallback(() => {
+    if (typeof returnPath === "string" && returnPath.startsWith("/")) {
+      navigate(returnPath, { replace: true });
+      return;
+    }
+    navigate("/");
+  }, [navigate, returnPath]);
+
   if (!hasLoaded || householdLoading || loading) {
     return (
       <div className="location-map-page">
@@ -421,7 +431,7 @@ export default function LocationMapManager() {
 
   return (
     <div className="location-map-page">
-      <LocationMapTopbar location={location} status={status} onBack={() => navigate(-1)} />
+      <LocationMapTopbar location={location} status={status} onBack={handleBack} />
 
       <main className="location-map-workspace">
         <LocationMapToolbar
diff --git a/frontend/tests/store-selector.spec.ts b/frontend/tests/store-selector.spec.ts
index 0fc4e5c..b020d40 100644
--- a/frontend/tests/store-selector.spec.ts
+++ b/frontend/tests/store-selector.spec.ts
@@ -182,6 +182,15 @@ test("grocery store selector shows one selected store and picks from a modal", a
 
   await page.getByRole("button", { name: "Manage Map" }).click();
   await expect(page).toHaveURL(/\/stores\/100\/locations\/13\/map$/);
+  await expect(page.getByRole("heading", { name: "Costco" })).toBeVisible();
+  await expect(page.getByText("Fontana")).toBeVisible();
+  await expect(page.getByRole("heading", { name: "No Map" })).toBeVisible();
+
+  await page.getByRole("button", { name: "Back" }).click();
+  await expect(page).toHaveURL(/\/$/);
+  await expect(storeSelector).toContainText("Costco");
+  await expect(locationSelector).toContainText("Fontana");
+  await expect(page.getByRole("button", { name: "Manage Map" })).toBeVisible();
 });
 
 test("single grocery store selector click does not open picker modal", async ({ page }) => {