improvement: Some improvements done. some bugs fixed.

frontend/src/components/ui/treatment/TreatmentWorkspace.tsx

@@ -329,12 +329,70 @@ function mergeServerIdsIntoDetails(
   });
 }
 
+const PREVIEW_TREATMENT_ID = 'current-draft';
+
+function isPreviewTreatment(treatment: PastTreatment): boolean {
+  return treatment.id === PREVIEW_TREATMENT_ID;
+}
+
+function serializeLabCases(drafts: LabCaseDraft[]): string {
+  return JSON.stringify(
+    [...drafts]
+      .sort((a, b) => a.clientId.localeCompare(b.clientId))
+      .map((lc) => ({
+        clientId: lc.clientId,
+        id: lc.id ?? null,
+        destinationOrganizationId: lc.destinationOrganizationId ?? null,
+        detailClientId: lc.detailClientId,
+        dueDate: lc.dueDate ?? null,
+        attachmentIds: [...lc.attachmentIds].sort(),
+        toothProsthesis: [...lc.toothProsthesis]
+          .map((tp) => ({
+            detailClientId: tp.detailClientId,
+            tooth: tp.tooth,
+            prosthesisTypeCode: tp.prosthesisTypeCode,
+            selectionGroupId: tp.selectionGroupId ?? '',
+          }))
+          .sort(
+            (a, b) =>
+              a.tooth.localeCompare(b.tooth) ||
+              a.prosthesisTypeCode.localeCompare(b.prosthesisTypeCode),
+          ),
+      })),
+  );
+}
+
+function isLabCasesDirty(drafts: LabCaseDraft[], savedSnapshot: string | null): boolean {
+  if (savedSnapshot === null) {
+    return drafts.some(
+      (lc) =>
+        lc.toothProsthesis.length > 0 ||
+        Boolean(lc.destinationOrganizationId) ||
+        Boolean(lc.dueDate) ||
+        lc.attachmentIds.length > 0,
+    );
+  }
+  return serializeLabCases(drafts) !== savedSnapshot;
+}
+
+function mergeServerIdsIntoLabCases(
+  local: LabCaseDraft[],
+  fromServer: LabCaseDraft[],
+): LabCaseDraft[] {
+  const byClientId = new Map(fromServer.map((lc) => [lc.clientId, lc]));
+  return local.map((lc) => {
+    const s = byClientId.get(lc.clientId);
+    if (!s) return lc;
+    return { ...lc, id: s.id ?? lc.id };
+  });
+}
+
 function detailsToPreviewTreatment(
   details: TreatmentDetailDraft[],
   meta: { title: string; patientId: string; treatmentAt: string; id?: string },
 ): PastTreatment {
   return {
-    id: meta.id ?? 'current-draft',
+    id: meta.id ?? PREVIEW_TREATMENT_ID,
     patientId: meta.patientId,
     title: meta.title,
     treatmentAt: meta.treatmentAt,
@@ -434,6 +492,7 @@ export function TreatmentWorkspace({
   const [activeDetailId, setActiveDetailId] = useState<string>(() => details[0].clientId);
   const [activeLabCaseId, setActiveLabCaseId] = useState<string | null>(null);
   const [savedSnapshot, setSavedSnapshot] = useState<string | null>(null);
+  const [savedLabCasesSnapshot, setSavedLabCasesSnapshot] = useState<string | null>(null);
   const [saveStatus, setSaveStatus] = useState<'idle' | 'dirty' | 'saving' | 'saved' | 'error'>('idle');
   const [selectedPreviewId, setSelectedPreviewId] = useState<string | null>(null);
   const [workspaceMode, setWorkspaceMode] = useState<WorkspaceMode>('live');
@@ -445,6 +504,8 @@ export function TreatmentWorkspace({
   detailsRef.current = details;
   const savedSnapshotRef = useRef(savedSnapshot);
   savedSnapshotRef.current = savedSnapshot;
+  const savedLabCasesSnapshotRef = useRef(savedLabCasesSnapshot);
+  savedLabCasesSnapshotRef.current = savedLabCasesSnapshot;
   const autosaveTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const saveInFlightRef = useRef(false);
   const saveQueuedRef = useRef(false);
@@ -498,8 +559,10 @@ export function TreatmentWorkspace({
   }, [labCaseDrafts, activeDetailId]);
 
   const isDirty = useMemo(
-    () => isDetailsDirty(details, savedSnapshot),
-    [details, savedSnapshot],
+    () =>
+      isDetailsDirty(details, savedSnapshot) ||
+      isLabCasesDirty(labCaseDrafts, savedLabCasesSnapshot),
+    [details, savedSnapshot, labCaseDrafts, savedLabCasesSnapshot],
   );
 
   const AUTOSAVE_DEBOUNCE_MS = 600;
@@ -799,7 +862,11 @@ export function TreatmentWorkspace({
     const mappedLabCases = withoutEmptyLabCaseDrafts(
       (treatment.labCases ?? []).map(mapLabCaseDraftFromApi),
     );
+    const labSnap = serializeLabCases(mappedLabCases);
+    labCaseDraftsRef.current = mappedLabCases;
+    savedLabCasesSnapshotRef.current = labSnap;
     setLabCaseDrafts(mappedLabCases);
+    setSavedLabCasesSnapshot(labSnap);
     setActiveLabCaseId(mappedLabCases[0]?.clientId ?? null);
     setOrganizationSearch('');
     setSaveStatus('idle');
@@ -1158,7 +1225,11 @@ export function TreatmentWorkspace({
         const mappedLabCases = withoutEmptyLabCaseDrafts(
           (response.data?.labCases ?? []).map(mapLabCaseDraftFromApi),
         );
+        const labSnap = serializeLabCases(mappedLabCases);
+        labCaseDraftsRef.current = mappedLabCases;
+        savedLabCasesSnapshotRef.current = labSnap;
         setLabCaseDrafts(mappedLabCases);
+        setSavedLabCasesSnapshot(labSnap);
         setActiveLabCaseId(mappedLabCases[0]?.clientId ?? null);
         setOrganizationSearch('');
         setSaveStatus('idle');
@@ -1241,12 +1312,15 @@ export function TreatmentWorkspace({
       // If the user changed details while this save was in flight (e.g. removed a
       // detail), do not clobber local state with the stale response.
       if (serializeDetails(localNow) === sentSnapshot) {
+        detailsRef.current = mapped;
         setDetails(mapped);
         setActiveDetailId((prev) => {
           const stillExists = mapped.some((d) => d.clientId === prev);
           return stillExists ? prev : (mapped[0]?.clientId ?? '');
         });
-        setSavedSnapshot(serializeDetails(mapped));
+        const snap = serializeDetails(mapped);
+        savedSnapshotRef.current = snap;
+        setSavedSnapshot(snap);
       } else {
         const merged = mergeServerIdsIntoDetails(localNow, mapped);
         detailsRef.current = merged;
@@ -1262,6 +1336,111 @@ export function TreatmentWorkspace({
     [selectedAppointment, selectedStandalone, t],
   );
 
+  const persistLabCases = useCallback(
+    async (savedTreatment: PastTreatment, draftsOverride?: LabCaseDraft[]) => {
+      if (!selectedAppointment && !selectedStandalone) throw new Error('No visit selected');
+      if (isPreviewTreatment(savedTreatment)) {
+        return savedTreatment;
+      }
+
+      const drafts = draftsOverride ?? labCaseDraftsRef.current;
+      const detailIdByClientId = new Map(
+        savedTreatment.details.map((d) => [d.clientId, d.id]),
+      );
+
+      const payload = drafts
+        .map((lc) => {
+          if (!lc.detailClientId) return null;
+          const treatmentDetailId = detailIdByClientId.get(lc.detailClientId);
+          if (!treatmentDetailId) return null;
+
+          return {
+            clientId: lc.clientId,
+            id: lc.id,
+            destinationOrganizationId: lc.destinationOrganizationId ?? undefined,
+            treatmentDetailId,
+            toothProsthesis: lc.toothProsthesis
+              .map((tp) => {
+                const detailId = detailIdByClientId.get(tp.detailClientId);
+                if (!detailId) return null;
+                return {
+                  treatmentDetailId: detailId,
+                  tooth: tp.tooth,
+                  prosthesisTypeCode: tp.prosthesisTypeCode,
+                  selectionGroupId: tp.selectionGroupId ?? '',
+                };
+              })
+              .filter(
+                (
+                  row,
+                ): row is {
+                  treatmentDetailId: string;
+                  tooth: string;
+                  prosthesisTypeCode: string;
+                  selectionGroupId: string;
+                } => row !== null,
+              ),
+            attachmentIds: lc.attachmentIds,
+            dueDate: lc.dueDate ?? null,
+          };
+        })
+        .filter((row): row is NonNullable<typeof row> => row !== null);
+
+      const localJobCount = drafts.reduce((n, lc) => n + lc.toothProsthesis.length, 0);
+      const payloadJobCount = payload.reduce((n, lc) => n + lc.toothProsthesis.length, 0);
+      if (payloadJobCount < localJobCount) {
+        throw new Error('Lab case jobs could not be mapped to saved details');
+      }
+
+      const applySavedDrafts = (mapped: LabCaseDraft[]) => {
+        const sent = serializeLabCases(drafts);
+        const localNow = labCaseDraftsRef.current;
+        if (serializeLabCases(localNow) === sent) {
+          const snap = serializeLabCases(mapped);
+          labCaseDraftsRef.current = mapped;
+          savedLabCasesSnapshotRef.current = snap;
+          setLabCaseDrafts(mapped);
+          setSavedLabCasesSnapshot(snap);
+          setActiveLabCaseId((prev) => {
+            if (prev && mapped.some((lc) => lc.clientId === prev)) return prev;
+            return mapped[0]?.clientId ?? null;
+          });
+          return;
+        }
+        const merged = mergeServerIdsIntoLabCases(localNow, mapped);
+        labCaseDraftsRef.current = merged;
+        setLabCaseDrafts(merged);
+      };
+
+      const emptySnap = serializeLabCases([]);
+      if (payload.length === 0) {
+        if (drafts.length > 0) {
+          throw new Error('Lab case drafts could not be mapped to saved details');
+        }
+        if (
+          !savedLabCasesSnapshotRef.current ||
+          savedLabCasesSnapshotRef.current === emptySnap
+        ) {
+          savedLabCasesSnapshotRef.current = emptySnap;
+          setSavedLabCasesSnapshot(emptySnap);
+          return savedTreatment;
+        }
+      }
+
+      const response = selectedAppointment
+        ? await treatmentsApi.saveLabCases(selectedAppointment.id, {
+            labCases: payload,
+          })
+        : await treatmentsApi.saveLabCasesByTreatment(selectedStandalone!.id, {
+            labCases: payload,
+          });
+      const mapped = withoutEmptyLabCaseDrafts(response.data.labCases.map(mapLabCaseDraftFromApi));
+      applySavedDrafts(mapped);
+      return response.data;
+    },
+    [selectedAppointment, selectedStandalone],
+  );
+
   const refreshHistory = useCallback(async (patientId: string, options?: { silentLabCases?: boolean }) => {
     const requestId = ++historyRequestRef.current;
     try {
@@ -1281,17 +1460,28 @@ export function TreatmentWorkspace({
       return;
     }
 
-    if (
-      !isDetailsDirty(detailsRef.current, savedSnapshotRef.current) ||
-      !areDetailsPersistable(detailsRef.current)
-    ) {
+    const detailsDirty = isDetailsDirty(detailsRef.current, savedSnapshotRef.current);
+    const labDirty = isLabCasesDirty(
+      labCaseDraftsRef.current,
+      savedLabCasesSnapshotRef.current,
+    );
+    if (!detailsDirty && !labDirty) {
+      return;
+    }
+    if (!areDetailsPersistable(detailsRef.current)) {
       return;
     }
 
     saveInFlightRef.current = true;
     setSaveStatus('saving');
+    const draftsAtSave = labCaseDraftsRef.current;
     try {
-      await persistDraft();
+      const saved = await persistDraft({ force: labDirty && !detailsDirty });
+      if (isPreviewTreatment(saved)) {
+        setSaveStatus('dirty');
+        return;
+      }
+      await persistLabCases(saved, draftsAtSave);
       setSaveStatus('saved');
       if (historyPatientId) {
         await refreshHistory(historyPatientId);
@@ -1304,12 +1494,15 @@ export function TreatmentWorkspace({
       saveInFlightRef.current = false;
       if (saveQueuedRef.current) {
         saveQueuedRef.current = false;
-        if (isDetailsDirty(detailsRef.current, savedSnapshotRef.current)) {
+        if (
+          isDetailsDirty(detailsRef.current, savedSnapshotRef.current) ||
+          isLabCasesDirty(labCaseDraftsRef.current, savedLabCasesSnapshotRef.current)
+        ) {
           void runDraftSave();
         }
       }
     }
-  }, [hasLiveContext, persistDraft, showError, t, historyPatientId, refreshHistory]);
+  }, [hasLiveContext, persistDraft, persistLabCases, showError, t, historyPatientId, refreshHistory]);
 
   const flushDraftSave = useCallback(async (): Promise<boolean> => {
     if (autosaveTimerRef.current) {
@@ -1325,7 +1518,10 @@ export function TreatmentWorkspace({
       await new Promise((resolve) => setTimeout(resolve, 50));
     }
 
-    if (!isDetailsDirty(detailsRef.current, savedSnapshotRef.current)) {
+    if (
+      !isDetailsDirty(detailsRef.current, savedSnapshotRef.current) &&
+      !isLabCasesDirty(labCaseDraftsRef.current, savedLabCasesSnapshotRef.current)
+    ) {
       return true;
     }
 
@@ -1850,79 +2046,11 @@ export function TreatmentWorkspace({
     [canEditTreatmentForDay, isDetailLocked, hasLiveContext, selectedAppointment, selectedStandalone, showSuccess, showError, t, tErrors],
   );
 
-  const persistLabCases = useCallback(
-    async (savedTreatment: PastTreatment, draftsOverride?: LabCaseDraft[]) => {
-      if (!selectedAppointment && !selectedStandalone) throw new Error('No visit selected');
-
-      const drafts = draftsOverride ?? labCaseDrafts;
-      const detailIdByClientId = new Map(
-        savedTreatment.details.map((d) => [d.clientId, d.id]),
-      );
-
-      const payload = drafts
-        .map((lc) => {
-          if (!lc.detailClientId) return null;
-          const treatmentDetailId = detailIdByClientId.get(lc.detailClientId);
-          if (!treatmentDetailId) return null;
-
-          return {
-            clientId: lc.clientId,
-            id: lc.id,
-            destinationOrganizationId: lc.destinationOrganizationId ?? undefined,
-            treatmentDetailId,
-            toothProsthesis: lc.toothProsthesis
-              .map((tp) => {
-                const detailId = detailIdByClientId.get(tp.detailClientId);
-                if (!detailId) return null;
-                return {
-                  treatmentDetailId: detailId,
-                  tooth: tp.tooth,
-                  prosthesisTypeCode: tp.prosthesisTypeCode,
-                  selectionGroupId: tp.selectionGroupId ?? '',
-                };
-              })
-              .filter(
-                (
-                  row,
-                ): row is {
-                  treatmentDetailId: string;
-                  tooth: string;
-                  prosthesisTypeCode: string;
-                  selectionGroupId: string;
-                } => row !== null,
-              ),
-            attachmentIds: lc.attachmentIds,
-            dueDate: lc.dueDate ?? null,
-          };
-        })
-        .filter((row): row is NonNullable<typeof row> => row !== null);
-
-      if (payload.length === 0) {
-        return savedTreatment;
-      }
-
-      const response = selectedAppointment
-        ? await treatmentsApi.saveLabCases(selectedAppointment.id, {
-            labCases: payload,
-          })
-        : await treatmentsApi.saveLabCasesByTreatment(selectedStandalone!.id, {
-            labCases: payload,
-          });
-      const mapped = withoutEmptyLabCaseDrafts(response.data.labCases.map(mapLabCaseDraftFromApi));
-      setLabCaseDrafts(mapped);
-      setActiveLabCaseId((prev) => {
-        if (prev && mapped.some((lc) => lc.clientId === prev)) return prev;
-        return mapped[0]?.clientId ?? null;
-      });
-      return response.data;
-    },
-    [labCaseDrafts, selectedAppointment, selectedStandalone],
-  );
-
   const handleLabCasesChange = useCallback(
     (next: LabCaseDraft[]) => {
       const prevCleaned = withoutEmptyLabCaseDrafts(labCaseDrafts);
       const cleaned = withoutEmptyLabCaseDrafts(next);
+      labCaseDraftsRef.current = cleaned;
       setLabCaseDrafts(cleaned);
 
       if (!cleaned.some((lc) => lc.detailClientId === activeDetailId)) {
@@ -2036,9 +2164,7 @@ export function TreatmentWorkspace({
         const updatedLabCases = [...labCaseDrafts, draft];
         setLabCaseDrafts(updatedLabCases);
 
-        // Autosave only watches `details`, so a lab draft left in state alone loses the
-        // lab, the due date and the prosthesis map on reload — silently, because the
-        // detail itself survives.
+        // Persist the new detail first so lab-case rows can use real treatmentDetailIds.
         void (async () => {
           try {
             const saved = await persistDraft({ force: true });