fix(work-orders): persist extra media placement

src/app/(protected)/workorders/_components/detail/use-slide-over-media-state.ts

@@ -1,4 +1,4 @@
-import { useCallback, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 import { toast } from "react-toastify";
 import type { MediaUploadRequest } from "@/app/(protected)/workorders/_components/media/media-uploader";
 import { useMediaUploadSession } from "@/app/(protected)/workorders/_hooks/use-media-upload-session";
@@ -14,6 +14,7 @@ import type {
 import {
   isCompletionSurfaceMedia,
   isExtrasSurfaceMedia,
+  type MediaSurfaceOwnership,
 } from "@/domain/work-orders/utils/partition-slide-over-media";
 
 type UseSlideOverMediaStateArgs = {
@@ -37,34 +38,83 @@ function dropOwnedId(previous: Set<string>, id: string): Set<string> {
   return next;
 }
 
+const OWNERSHIP_STORAGE_PREFIX = "work-order-media-surface:";
+
+function loadOwnership(workOrderId: string | number | undefined): MediaSurfaceOwnership {
+  if (workOrderId == null) return { completionOwnedIds: new Set(), extraOwnedIds: new Set() };
+  try {
+    const stored = sessionStorage.getItem(`${OWNERSHIP_STORAGE_PREFIX}${workOrderId}`);
+    if (!stored) return { completionOwnedIds: new Set(), extraOwnedIds: new Set() };
+    const parsed = JSON.parse(stored) as { completion?: unknown; extra?: unknown };
+    return {
+      completionOwnedIds: new Set(Array.isArray(parsed.completion) ? parsed.completion : []),
+      extraOwnedIds: new Set(Array.isArray(parsed.extra) ? parsed.extra : []),
+    };
+  } catch {
+    return { completionOwnedIds: new Set(), extraOwnedIds: new Set() };
+  }
+}
+
+function saveOwnership(
+  workOrderId: string | number | undefined,
+  ownership: MediaSurfaceOwnership,
+): void {
+  if (workOrderId == null) return;
+  try {
+    sessionStorage.setItem(
+      `${OWNERSHIP_STORAGE_PREFIX}${workOrderId}`,
+      JSON.stringify({
+        completion: [...ownership.completionOwnedIds],
+        extra: [...ownership.extraOwnedIds],
+      }),
+    );
+  } catch {
+    return;
+  }
+}
+
 export function useSlideOverMediaState({
   detailForRow,
   uploadMediaMutation,
   deleteMediaMutation,
   categorizeMediaMutation,
 }: UseSlideOverMediaStateArgs) {
+  const workOrderId = detailForRow?.info.id;
   const [mediaBatchPending, setMediaBatchPending] = useState(false);
-  const [completionOwnedIds, setCompletionOwnedIds] = useState<Set<string>>(() => new Set());
-  const [extraOwnedIds, setExtraOwnedIds] = useState<Set<string>>(() => new Set());
+  const [ownership, setOwnershipState] = useState<MediaSurfaceOwnership>(() =>
+    loadOwnership(workOrderId),
+  );
+
+  useEffect(() => {
+    setOwnershipState(loadOwnership(workOrderId));
+  }, [workOrderId]);
+
+  const updateOwnership = useCallback(
+    (update: (current: MediaSurfaceOwnership) => MediaSurfaceOwnership) => {
+      setOwnershipState((current) => {
+        const next = update(current);
+        saveOwnership(workOrderId, next);
+        return next;
+      });
+    },
+    [workOrderId],
+  );
 
   const persistCompletionUpload = useCallback(
     async (input: UploadWorkOrderMediaInput) => {
       const item = await uploadMediaMutation.mutateAsync(input);
-      setCompletionOwnedIds((previous) => addOwnedId(previous, item.id));
-      setExtraOwnedIds((previous) => dropOwnedId(previous, item.id));
+      updateOwnership((current) => ({
+        completionOwnedIds: addOwnedId(new Set(current.completionOwnedIds), item.id),
+        extraOwnedIds: dropOwnedId(new Set(current.extraOwnedIds), item.id),
+      }));
       return item;
     },
-    [uploadMediaMutation],
+    [updateOwnership, uploadMediaMutation],
   );
 
   const { uploads, startUploads, retryUpload, cancelUpload, removeLocal, isLocal } =
     useMediaUploadSession(detailForRow, persistCompletionUpload);
 
-  const ownership = useMemo(
-    () => ({ completionOwnedIds, extraOwnedIds }),
-    [completionOwnedIds, extraOwnedIds],
-  );
-
   const completionMedia = useMemo(
     () => uploads.filter((item) => isCompletionSurfaceMedia(item, ownership)),
     [uploads, ownership],
@@ -96,8 +146,10 @@ export function useSlideOverMediaState({
       for (const file of files) {
         // Upload stays Extra; Aveta is applied only by an explicit Extra Docs PATCH.
         const item = await uploadMediaMutation.mutateAsync({ file, category: "Extra" });
-        setExtraOwnedIds((previous) => addOwnedId(previous, item.id));
-        setCompletionOwnedIds((previous) => dropOwnedId(previous, item.id));
+        updateOwnership((current) => ({
+          completionOwnedIds: dropOwnedId(new Set(current.completionOwnedIds), item.id),
+          extraOwnedIds: addOwnedId(new Set(current.extraOwnedIds), item.id),
+        }));
       }
     } finally {
       setMediaBatchPending(false);
@@ -110,14 +162,20 @@ export function useSlideOverMediaState({
       toast.error("Missing work order version. Refresh and try again.");
       return;
     }
-    if (category === "Extra") {
-      setExtraOwnedIds((previous) => addOwnedId(previous, mediaId));
-      setCompletionOwnedIds((previous) => dropOwnedId(previous, mediaId));
-    } else {
-      setExtraOwnedIds((previous) => dropOwnedId(previous, mediaId));
-      setCompletionOwnedIds((previous) => dropOwnedId(previous, mediaId));
-    }
-    categorizeMediaMutation.mutate({ mediaId, category, workOrderVersion });
+    categorizeMediaMutation.mutate(
+      { mediaId, category, workOrderVersion },
+      {
+        onSuccess: () => {
+          updateOwnership((current) => ({
+            completionOwnedIds: dropOwnedId(new Set(current.completionOwnedIds), mediaId),
+            extraOwnedIds:
+              category === "Extra"
+                ? addOwnedId(new Set(current.extraOwnedIds), mediaId)
+                : dropOwnedId(new Set(current.extraOwnedIds), mediaId),
+          }));
+        },
+      },
+    );
   };
 
   const removeMedia = (id: string) => {

src/app/(protected)/workorders/_components/media/media-uploader-grid-item-meta.tsx

@@ -38,7 +38,7 @@ export function MediaUploaderGridItemMeta({
       {needsCategory && (
         <select
           aria-label={`Category for ${file.name}`}
-          value={file.purpose ?? ""}
+          value={file.purpose === "Extra" ? "" : (file.purpose ?? "")}
           disabled={pending}
           onChange={(event) => {
             const value = event.target.value as UploadableMediaCategory;

src/test/app/(protected)/workorders/media-uploader-grid-item-meta.test.tsx

@@ -0,0 +1,35 @@
+import { fireEvent, render, screen } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { MediaUploaderGridItemMeta } from "@/app/(protected)/workorders/_components/media/media-uploader-grid-item-meta";
+import type { MediaFile } from "@/domain/work-orders/types/work-order-media";
+
+const extraPhoto: MediaFile = {
+  id: "10",
+  name: "site.jpg",
+  fileType: "image",
+  url: "https://cdn.example.com/site.jpg",
+  category: "Photo",
+  purpose: "Extra",
+  isLegacy: false,
+};
+
+describe("MediaUploaderGridItemMeta", () => {
+  it("allows an Extra photo to be explicitly assigned to Extra Docs", () => {
+    const onCategorize = vi.fn();
+    render(
+      <MediaUploaderGridItemMeta
+        file={extraPhoto}
+        readOnly={false}
+        pending={false}
+        onCategorize={onCategorize}
+      />,
+    );
+
+    const category = screen.getByLabelText("Category for site.jpg");
+    expect(category).toHaveValue("");
+
+    fireEvent.change(category, { target: { value: "Extra" } });
+
+    expect(onCategorize).toHaveBeenCalledWith("10", "Extra");
+  });
+});

src/test/app/(protected)/workorders/use-slide-over-media-state.test.tsx

@@ -1,5 +1,5 @@
 import { act, renderHook } from "@testing-library/react";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { useSlideOverMediaState } from "@/app/(protected)/workorders/_components/detail/use-slide-over-media-state";
 import type { WorkOrderBoardDetail } from "@/domain/work-orders/types/work-order-board-detail";
 import type { WorkOrderTableRow } from "@/domain/work-orders/types/work-order-table-row";
@@ -85,6 +85,10 @@ function idleMutation() {
 }
 
 describe("useSlideOverMediaState completion filter", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+  });
+
   it("keeps Before/After/Completion after refetch mapped purpose (not file category)", () => {
     const detail = detailWithMedia([
       {
@@ -214,7 +218,7 @@ describe("useSlideOverMediaState completion filter", () => {
     expect(result.current.extrasUploads[1]?.category).toBe("Aveta");
   });
 
-  it("moves an Extra photo to Extra Docs after assigning Extra", () => {
+  it("keeps an Extra photo in Extra Docs after assigning Extra and remounting", () => {
     const detail = detailWithMedia([
       {
         id: "10",
@@ -225,14 +229,16 @@ describe("useSlideOverMediaState completion filter", () => {
         isLegacy: false,
       },
     ]);
-    const categorizeMutate = vi.fn();
+    const categorizeMutate = vi.fn((_input: unknown, options?: { onSuccess?: () => void }) =>
+      options?.onSuccess?.(),
+    );
     const categorize = {
       mutate: categorizeMutate,
       mutateAsync: vi.fn().mockResolvedValue({}),
       isPending: false,
     } as never;
 
-    const { result } = renderHook(() =>
+    const { result, unmount } = renderHook(() =>
       useSlideOverMediaState({
         detailForRow: detail,
         uploadMediaMutation: idleMutation(),
@@ -250,10 +256,26 @@ describe("useSlideOverMediaState completion filter", () => {
 
     expect(result.current.completionMedia).toHaveLength(0);
     expect(result.current.extrasUploads.map((item) => item.id)).toEqual(["10"]);
-    expect(categorizeMutate).toHaveBeenCalledWith({
-      mediaId: "10",
-      category: "Extra",
-      workOrderVersion: "1",
-    });
+    expect(categorizeMutate).toHaveBeenCalledWith(
+      {
+        mediaId: "10",
+        category: "Extra",
+        workOrderVersion: "1",
+      },
+      expect.objectContaining({ onSuccess: expect.any(Function) }),
+    );
+
+    unmount();
+    const remounted = renderHook(() =>
+      useSlideOverMediaState({
+        detailForRow: detail,
+        uploadMediaMutation: idleMutation(),
+        deleteMediaMutation: idleMutation(),
+        categorizeMediaMutation: idleMutation(),
+      }),
+    );
+
+    expect(remounted.result.current.completionMedia).toHaveLength(0);
+    expect(remounted.result.current.extrasUploads.map((item) => item.id)).toEqual(["10"]);
   });
 });