Bug fixes: Class deletion across sample & undo/redo for polygon+mag lasso

frontend/src/app/pages/AnnotatePage.tsx

@@ -9,6 +9,7 @@ import { useAnnotationStore } from '@/stores/annotationStore';
 import { useToolStore } from '@/stores/toolStore';
 import { useClassStore } from '@/stores/classStore';
 import { useDraftSync } from '@/hooks/useDraftSync';
+import { clearHistory } from '@/hooks/editHistory';
 import { useGuideLoad } from '@/hooks/useGuideSync';
 import { useSave, type VersionPayload } from '@/hooks/useSave';
 import { buildSourceKey } from '@/lib/sourceKey';
@@ -87,6 +88,8 @@ export default function AnnotatePage() {
 
   // Crash-recovery autosave (local draft only, no Tiled sync)
   useDraftSync(sourceKey);
+  // Undo/redo is per-sample: reset the region history + class-delete journal on switch.
+  useEffect(() => { clearHistory(); }, [sourceKey]);
   // Load the dataset's annotation guide (read-only) for class suggestions/examples.
   useGuideLoad(sourceKey);
 

frontend/src/components/annotate/ClassManager/index.tsx

@@ -11,19 +11,22 @@ import { useToolStore } from '@/stores/toolStore';
 import { useDatasetStore } from '@/stores/datasetStore';
 import { useReferenceGuideStore, type GuideClass } from '@/stores/referenceGuideStore';
 import { useSettingsStore } from '@/stores/settingsStore';
+import { markClassDelete } from '@/hooks/editHistory';
 import { getClassPalette } from '@/lib/classColors';
 import { buildSourceKey } from '@/lib/sourceKey';
 import { cn } from '@/lib/utils';
 import { Copy } from '@phosphor-icons/react';
 
-/** Counts all shapes assigned to a class across every image/slice in the annotation store. */
-function countShapesForClass(classId: number): number {
-  const { byImage } = useAnnotationStore.getState();
+/** Counts shapes assigned to a class within a single sample (*sourceKey*), across its
+ *  slices. Scoped to the current sample so the delete prompt never counts (or deletes)
+ *  annotations belonging to other samples that happen to reuse the same classId. */
+function countShapesForClass(classId: number, sourceKey: string | null): number {
+  if (!sourceKey) return 0;
+  const slices = useAnnotationStore.getState().byImage[sourceKey];
+  if (!slices) return 0;
   let total = 0;
-  for (const slices of Object.values(byImage)) {
-    for (const shapes of Object.values(slices)) {
-      total += shapes.filter((sh) => sh.classId === classId).length;
-    }
+  for (const shapes of Object.values(slices)) {
+    total += shapes.filter((sh) => sh.classId === classId).length;
   }
   return total;
 }
@@ -32,7 +35,9 @@ interface ClassRowProps {
   cls: AnnotationClass;
   isActive: boolean;
   onActivate: () => void;
-  onClassDeleted: (deletedClassId: number) => void;
+  /** Delete this class (with confirmation + undo) — handled by the parent so the
+   *  undo affordance can outlive this row. */
+  onDelete: () => void;
   /** Duplicate this class + all its shapes into a new class. */
   onDuplicate: () => void;
   /** Matching guide entry (description + example crops), if the guide defines this class. */
@@ -42,10 +47,8 @@ interface ClassRowProps {
 }
 
 /** Renders a single class row with inline rename, visibility toggle, and delete. */
-function ClassRow({ cls, isActive, onActivate, onClassDeleted, onDuplicate, guide, hotkey }: ClassRowProps) {
-  const { updateClass, deleteClass, toggleVisibility } = useClassStore();
-  const removeShapesByClassId = useAnnotationStore((s) => s.removeShapesByClassId);
-  const setSelectedShapeId = useToolStore((s) => s.setSelectedShapeId);
+function ClassRow({ cls, isActive, onActivate, onDelete, onDuplicate, guide, hotkey }: ClassRowProps) {
+  const { updateClass, toggleVisibility } = useClassStore();
   const [editing, setEditing] = useState(false);
   const [showGuide, setShowGuide] = useState(false);
   const [labelInput, setLabelInput] = useState(cls.label);
@@ -69,28 +72,6 @@ function ClassRow({ cls, isActive, onActivate, onClassDeleted, onDuplicate, guid
     setEditing(false);
   };
 
-  /** Confirms with the user, then removes the class and all of its shapes from the stores. */
-  const handleDelete = (e: React.MouseEvent) => {
-    e.stopPropagation();
-    const shapeCount = countShapesForClass(cls.classId);
-    const annotationNote =
-      shapeCount === 0
-        ? 'This class has no annotations.'
-        : shapeCount === 1
-          ? 'This will permanently delete 1 annotation.'
-          : `This will permanently delete ${shapeCount} annotations.`;
-
-    const confirmed = window.confirm(
-      `Do you really want to delete "${cls.label}" and its annotations?\n\n${annotationNote}`
-    );
-    if (!confirmed) return;
-
-    removeShapesByClassId(cls.classId);
-    deleteClass(cls.classId);
-    setSelectedShapeId(null);
-    onClassDeleted(cls.classId);
-  };
-
   return (
     <div>
       <div
@@ -184,7 +165,7 @@ function ClassRow({ cls, isActive, onActivate, onClassDeleted, onDuplicate, guid
         <button
           aria-label="Delete class and its annotations"
           className="shrink-0 p-0.5 hover:text-red-500"
-          onClick={handleDelete}
+          onClick={(e) => { e.stopPropagation(); onDelete(); }}
         >
           <Trash size={14} />
         </button>
@@ -222,8 +203,11 @@ const DEFAULT_SUGGESTED_CLASSES = ['air', 'sample', 'void', 'pore', 'background'
 
 /** Renders the class list, add form, and quick-add suggestion chips. */
 export default function ClassManager({ activeClassId, onActivate, onClassDeleted }: ClassManagerProps) {
-  const { classes, addClass } = useClassStore();
+  const { classes, addClass, deleteClass } = useClassStore();
   const duplicateClassShapes = useAnnotationStore((s) => s.duplicateClassShapes);
+  const removeShapesByClassIdInSource = useAnnotationStore((s) => s.removeShapesByClassIdInSource);
+  const touchHistory = useAnnotationStore((s) => s.touchHistory);
+  const setSelectedShapeId = useToolStore((s) => s.setSelectedShapeId);
   const { source, kind, serverUri } = useDatasetStore();
   const sourceKey = source && kind ? buildSourceKey(kind as 'tiled' | 'local', source, serverUri) : null;
   const guideEntries = useReferenceGuideStore((s) => s.entries);
@@ -244,14 +228,34 @@ export default function ClassManager({ activeClassId, onActivate, onClassDeleted
     guideEntries.filter((g) => g.label.trim()).map((g) => [g.label.trim().toLowerCase(), g]),
   );
 
-  /** Notifies the parent of a deletion and re-activates the first remaining class if the active one was removed. */
-  const handleClassDeleted = (deletedClassId: number) => {
-    onClassDeleted?.(deletedClassId);
-    if (activeClassId === deletedClassId) {
+  /** Confirms, then removes the class and its regions (scoped to the current sample) and
+   *  re-activates the first remaining class. Registered for undo so Ctrl/Cmd+Z restores
+   *  the class entry along with its regions (see editHistory). */
+  const handleDeleteClass = (cls: AnnotationClass) => {
+    const shapeCount = countShapesForClass(cls.classId, sourceKey);
+    const note =
+      shapeCount === 0
+        ? 'This class has no annotations in this sample.'
+        : shapeCount === 1
+          ? 'This will delete 1 annotation in this sample.'
+          : `This will delete ${shapeCount} annotations in this sample.`;
+    if (!window.confirm(`Delete "${cls.label}"?\n\n${note}\n\nYou can undo with Ctrl/Cmd+Z.`)) return;
+
+    const index = classes.findIndex((c) => c.classId === cls.classId);
+    if (sourceKey) {
+      // Put the deletion on the undo timeline (see editHistory). Removing regions is a
+      // tracked edit; a region-less class still needs one entry, so we touch history.
+      markClassDelete(sourceKey, cls, index < 0 ? classes.length : index);
+      if (shapeCount > 0) removeShapesByClassIdInSource(sourceKey, cls.classId);
+      else touchHistory();
+    }
+    deleteClass(cls.classId);
+    setSelectedShapeId(null);
+
+    onClassDeleted?.(cls.classId);
+    if (activeClassId === cls.classId) {
       const remaining = useClassStore.getState().classes;
-      if (remaining.length > 0) {
-        onActivate(remaining[0].classId);
-      }
+      if (remaining.length > 0) onActivate(remaining[0].classId);
     }
   };
 
@@ -392,7 +396,7 @@ export default function ClassManager({ activeClassId, onActivate, onClassDeleted
             cls={cls}
             isActive={cls.classId === activeClassId}
             onActivate={() => onActivate(cls.classId)}
-            onClassDeleted={handleClassDeleted}
+            onDelete={() => handleDeleteClass(cls)}
             onDuplicate={() => handleDuplicate(cls)}
             guide={guideByLabel.get(cls.label.trim().toLowerCase())}
             hotkey={idx < 9 ? idx + 1 : undefined}

frontend/src/components/annotate/Toolbar/index.tsx

@@ -7,6 +7,7 @@ import { Hand, Cursor, Polygon, MagnetStraight, MagicWand, Rectangle, Circle, Pa
 import { useStore } from 'zustand';
 import { useToolStore, type Tool } from '@/stores/toolStore';
 import { useAnnotationStore } from '@/stores/annotationStore';
+import * as editHistory from '@/hooks/editHistory';
 import { cn } from '@/lib/utils';
 import DebouncedSlider from '@/components/common/DebouncedSlider';
 import { useSam } from '@/hooks/useSam';
@@ -95,7 +96,8 @@ export default function Toolbar({ disabled = false }: ToolbarProps) {
     selectScope, setSelectScope,
   } = useToolStore();
   const sam = useSam(tool === 'magic' && magicEngine === 'sam');
-  const { undo, redo } = useStore(useAnnotationStore.temporal);
+  // Undo/redo route through editHistory so a class deletion replays alongside its region
+  // change; canUndo/canRedo still reflect the (1:1) zundo stack.
   const canUndo = useStore(useAnnotationStore.temporal, (s) => s.pastStates.length > 0);
   const canRedo = useStore(useAnnotationStore.temporal, (s) => s.futureStates.length > 0);
 
@@ -118,7 +120,7 @@ export default function Toolbar({ disabled = false }: ToolbarProps) {
       <div className="flex gap-1">
         <button
           type="button"
-          onClick={() => undo()}
+          onClick={() => editHistory.undo()}
           disabled={!canUndo}
           title="Undo (Ctrl/⌘+Z)"
           aria-label="Undo"
@@ -134,7 +136,7 @@ export default function Toolbar({ disabled = false }: ToolbarProps) {
         </button>
         <button
           type="button"
-          onClick={() => redo()}
+          onClick={() => editHistory.redo()}
           disabled={!canRedo}
           title="Redo (Ctrl+Shift+Z)"
           aria-label="Redo"

frontend/src/hooks/editHistory.ts

@@ -0,0 +1,98 @@
+/**
+ * editHistory — makes class deletion a real, ordered, redoable Ctrl/Cmd+Z action.
+ *
+ * Region edits already live on zundo's temporal (undo) stack. Class *entries* live in a
+ * separate, non-temporal store, so this module keeps a small journal in lockstep with
+ * zundo: one entry per tracked edit (via annotationStore's `setEditListener`), tagged as
+ * either a plain region edit (`null`) or a class deletion. All undo/redo call sites route
+ * through `undo()`/`redo()` here, which drive zundo AND replay the paired class op — so a
+ * class delete undoes/redoes as one step, correctly ordered with every other edit.
+ *
+ * Because each journal entry maps 1:1 to a zundo entry, `undoOps.length` stays equal to
+ * `pastStates.length`; we re-trim after each edit to mirror zundo's `limit`.
+ */
+import { useAnnotationStore, setEditListener } from '@/stores/annotationStore';
+import { useClassStore, type AnnotationClass } from '@/stores/classStore';
+import { useDatasetStore } from '@/stores/datasetStore';
+import { buildSourceKey } from '@/lib/sourceKey';
+
+interface ClassDeleteEntry {
+  kind: 'classDelete';
+  sourceKey: string;
+  cls: AnnotationClass;
+  index: number;
+}
+/** A journal entry: a class deletion to replay, or `null` for a plain edit (region edit or
+ *  a polygon/lasso draft node/close — those are fully handled by zundo's tracked `draft`). */
+type Entry = ClassDeleteEntry | null;
+
+let undoOps: Entry[] = [];
+let redoOps: Entry[] = [];
+
+// Set just before the annotation-store edit it describes, so the next tracked edit is
+// tagged with it. Consumed by onTrackedEdit().
+let pending: Entry = null;
+
+/** Current sample key, or null when nothing is open. */
+function currentSourceKey(): string | null {
+  const { source, kind, serverUri } = useDatasetStore.getState();
+  return source && kind ? buildSourceKey(kind as 'tiled' | 'local', source, serverUri) : null;
+}
+
+/** Tag the next tracked edit as a class deletion (call immediately before the edit). */
+export function markClassDelete(sourceKey: string, cls: AnnotationClass, index: number) {
+  pending = { kind: 'classDelete', sourceKey, cls, index };
+}
+
+// Matches the annotationStore temporal `limit` so the journal drops its oldest entry in
+// lockstep with zundo dropping its oldest pastState.
+const LIMIT = 200;
+
+/** Fired by annotationStore after each tracked edit — records one journal entry.
+ *  NOTE: zundo invokes onSave BEFORE it appends the new pastState, so we cannot compare
+ *  against pastStates.length here (it's still the pre-edit count) — just cap at LIMIT. */
+function onTrackedEdit() {
+  const entry = pending;
+  pending = null;
+  undoOps.push(entry);
+  redoOps = [];
+  if (undoOps.length > LIMIT) undoOps.shift();
+}
+
+// Register once when this module first loads (imported by the undo/redo call sites).
+setEditListener(onTrackedEdit);
+
+/** Undo one step: zundo region undo + (if the entry is a class delete) re-insert the class. */
+export function undo() {
+  const temporal = useAnnotationStore.temporal.getState();
+  if (temporal.pastStates.length === 0) return;
+  const entry = undoOps.length ? undoOps.pop()! : null;
+  temporal.undo();
+  if (entry?.kind === 'classDelete' && entry.sourceKey === currentSourceKey()) {
+    useClassStore.getState().insertClass(entry.cls, entry.index);
+  }
+  redoOps.push(entry);
+}
+
+/** Redo one step: zundo region redo + (if the entry is a class delete) re-remove the class. */
+export function redo() {
+  const temporal = useAnnotationStore.temporal.getState();
+  if (temporal.futureStates.length === 0) return;
+  const entry = redoOps.length ? redoOps.pop()! : null;
+  temporal.redo();
+  if (entry?.kind === 'classDelete' && entry.sourceKey === currentSourceKey()) {
+    useClassStore.getState().deleteClass(entry.cls.classId);
+  }
+  undoOps.push(entry);
+}
+
+/** Reset both the region (zundo) history and the class-delete journal. Called on sample
+ *  switch so undo/redo is scoped per sample. */
+export function clearHistory() {
+  // Drop any in-progress draft too (its nodes lived on the temporal stack we're clearing),
+  // so no draft survives a sample switch. The set is wiped by clear() immediately below.
+  useAnnotationStore.getState().clearDraft();
+  undoOps = [];
+  redoOps = [];
+  useAnnotationStore.temporal.getState().clear();
+}