fix(work-orders): harden CompDoc state, persist rollback, and close gate

src/app/(protected)/workorders/_components/comp-doc/comp-doc-dialog-footer.tsx

@@ -35,8 +35,13 @@ export function CompDocDialogFooter({
   onDownloadPdf,
   onUploadPdf,
 }: CompDocDialogFooterProps) {
+  // Exclusive footer branches (prototype: one docState at a time). A signed PDF
+  // on file implies Generated even if local docState has not caught up yet — never
+  // render Draft actions alongside Generated.
+  const showNotRequired = docState === "not-required";
   const showGenerated =
-    (docState === "generated" || Boolean(signOffAttachment)) && docState !== "not-required";
+    !showNotRequired && (docState === "generated" || Boolean(signOffAttachment));
+  const showDraft = !showNotRequired && !showGenerated && docState === "draft";
 
   return (
     <div
@@ -57,7 +62,7 @@ export function CompDocDialogFooter({
           event.target.value = "";
         }}
       />
-      {docState === "draft" && !isDocumentsReadOnly && (
+      {Boolean(showDraft && !isDocumentsReadOnly) && (
         <CompDocDialogFooterDraftActions
           rowVersion={rowVersion}
           isUploading={isUploading}
@@ -78,7 +83,7 @@ export function CompDocDialogFooter({
           onDownloadPdf={onDownloadPdf}
         />
       )}
-      {docState === "not-required" && (
+      {showNotRequired && (
         <CompDocDialogFooterNotRequiredActions
           isDocumentsReadOnly={isDocumentsReadOnly}
           onUndo={onUndo}

src/app/(protected)/workorders/_components/comp-doc/use-comp-doc-dialog.tsx

@@ -62,8 +62,8 @@ export function useCompDocDialog({
   const { data: locationDetail } = useLocationDetail(locId || undefined);
   const { data: boardDetail } = useWorkOrderBoardDetail(row.id, open);
   const uploadCompletionMutation = useUploadCompletionDoc(row.id);
-  const isUploading = uploadCompletionMutation.isPending;
   const pdfInputRef = useRef<HTMLInputElement>(null);
+  const [isFinalizing, setIsFinalizing] = useState(false);
 
   const info = boardDetail?.info ?? row;
   const completion = boardDetail?.completion;
@@ -78,9 +78,18 @@ export function useCompDocDialog({
   const [docState, setDocState] = useState<DocState>("draft");
   const [generatedAt, setGeneratedAt] = useState("");
   const [editDateCompleted, setEditDateCompleted] = useState("");
+  const seedKeyRef = useRef<string | null>(null);
 
+  // Prototype seeds once per open/WO session. Detail refetches must not wipe
+  // in-progress PM comments or date edits (Bugbot: dialog effect clears edits).
   useEffect(() => {
-    if (!open) return;
+    if (!open) {
+      seedKeyRef.current = null;
+      return;
+    }
+    const key = String(info.id);
+    if (seedKeyRef.current === key) return;
+    seedKeyRef.current = key;
     setNotes(data?.notes ?? "");
     setDocState(docStateOf(completion?.docStatus || info.docStatus));
     setGeneratedAt(data?.generatedAt ?? "");
@@ -100,6 +109,7 @@ export function useCompDocDialog({
   const cfg = DOC_STATE_CFG[docState];
   const isNotRequired = docState === "not-required";
   const isGenerated = docState === "generated" || Boolean(signOffAttachment);
+  const isBusy = uploadCompletionMutation.isPending || isFinalizing;
 
   const docFields: DocFields = {
     street: site.street,
@@ -127,14 +137,36 @@ export function useCompDocDialog({
       onDraftSaved: notifyDraftSaved,
     });
 
-  function handleMarkNotRequired() {
+  async function handleMarkNotRequired() {
+    const previous = docState;
     setDocState("not-required");
-    void persist({}, "NN");
+    setIsFinalizing(true);
+    try {
+      await persist({}, "NN");
+    } catch (error) {
+      setDocState(previous);
+      refreshAuthoritativeState();
+      toast.error(
+        error instanceof Error ? error.message : "Failed to mark document as not required",
+      );
+    } finally {
+      setIsFinalizing(false);
+    }
   }
 
-  function handleUndo() {
+  async function handleUndo() {
+    const previous = docState;
     setDocState("draft");
-    void persist({}, "No");
+    setIsFinalizing(true);
+    try {
+      await persist({}, "No");
+    } catch (error) {
+      setDocState(previous);
+      refreshAuthoritativeState();
+      toast.error(error instanceof Error ? error.message : "Failed to undo not-required");
+    } finally {
+      setIsFinalizing(false);
+    }
   }
 
   const { handleDownloadPdf, handleUploadPdf } = useCompDocPdfActions({
@@ -150,6 +182,7 @@ export function useCompDocDialog({
     setGeneratedAt,
     markGenerated: () => setDocState("generated"),
     onPersistFailure: refreshAuthoritativeState,
+    onFinalizingChange: setIsFinalizing,
   });
 
   return {
@@ -170,7 +203,7 @@ export function useCompDocDialog({
     cfg,
     isNotRequired,
     isGenerated,
-    isUploading,
+    isUploading: isBusy,
     pdfInputRef,
     handleMarkNotRequired,
     handleUndo,

src/app/(protected)/workorders/_components/comp-doc/use-comp-doc-pdf-actions.ts

@@ -26,6 +26,7 @@ type UseCompDocPdfActionsArgs = {
   setGeneratedAt: (value: string) => void;
   markGenerated: () => void;
   onPersistFailure?: () => void;
+  onFinalizingChange?: (finalizing: boolean) => void;
 };
 
 export function useCompDocPdfActions({
@@ -41,6 +42,7 @@ export function useCompDocPdfActions({
   setGeneratedAt,
   markGenerated,
   onPersistFailure,
+  onFinalizingChange,
 }: UseCompDocPdfActionsArgs) {
   function handleDownloadPdf() {
     if (signOffAttachment) {
@@ -69,6 +71,8 @@ export function useCompDocPdfActions({
       { file, workOrderVersion: rowVersion },
       {
         onSuccess: () => {
+          // Keep close disabled across the gap between upload settle and docStatus persist.
+          onFinalizingChange?.(true);
           void (async () => {
             const ts = fmtDocNow();
             try {
@@ -81,6 +85,8 @@ export function useCompDocPdfActions({
               toast.error(
                 error instanceof Error ? error.message : "Failed to finalize completion document",
               );
+            } finally {
+              onFinalizingChange?.(false);
             }
           })();
         },

src/app/(protected)/workorders/_hooks/use-work-orders-list-page.ts

@@ -87,7 +87,13 @@ export function useWorkOrdersListPage() {
     data: CompDocData,
     status: CompDocSaveStatus,
   ) => {
-    setCompletionDocs((prev) => ({ ...prev, [String(row.id)]: data }));
+    // Persist PM comments + generatedAt in session state (no dedicated BE field yet)
+    // before patching board docStatus — Save draft / Generate / NN all share this path.
+    const nextData: CompDocData = {
+      notes: data.notes,
+      generatedAt: data.generatedAt,
+    };
+    setCompletionDocs((prev) => ({ ...prev, [String(row.id)]: nextData }));
     await patchFieldAsync(row.id, { docStatus: status });
   };
 

src/test/app/(protected)/workorders/comp-doc-dialog-completion-contract.test.tsx

@@ -194,4 +194,36 @@ describe("CompDocDialog Save Draft / Generate completion contract", () => {
     expect(callOrder).toEqual(["date"]);
     expect(onSave).not.toHaveBeenCalled();
   });
+
+  it("Save draft persists PM comments notes through onSave", async () => {
+    const row = baseRow();
+    renderDialog(row);
+
+    const notes = screen.getByPlaceholderText(/add findings/i);
+    fireEvent.change(notes, { target: { value: "No issues noted" } });
+    fireEvent.click(screen.getByRole("button", { name: /save draft/i }));
+
+    await waitFor(() => {
+      expect(onSave).toHaveBeenCalledWith(
+        expect.objectContaining({ notes: "No issues noted" }),
+        "No",
+      );
+    });
+  });
+
+  it("Mark as not required rolls back UI when persist fails", async () => {
+    onSave.mockRejectedValueOnce(new Error("persist failed"));
+    const row = baseRow();
+    renderDialog(row);
+
+    fireEvent.click(screen.getByRole("button", { name: /mark as not required/i }));
+
+    await waitFor(() => {
+      expect(onSave).toHaveBeenCalledWith(expect.any(Object), "NN");
+    });
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /mark as not required/i })).toBeInTheDocument();
+    });
+    expect(screen.queryByRole("button", { name: /^undo$/i })).not.toBeInTheDocument();
+  });
 });

src/test/app/(protected)/workorders/comp-doc-dialog-footer-exclusive.test.tsx

@@ -0,0 +1,53 @@
+import { createRef } from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { CompDocDialogFooter } from "@/app/(protected)/workorders/_components/comp-doc/comp-doc-dialog-footer";
+
+describe("CompDocDialogFooter exclusive states", () => {
+  const pdfInputRef = createRef<HTMLInputElement | null>();
+  const handlers = {
+    onMarkNotRequired: vi.fn(),
+    onSaveDraft: vi.fn(),
+    onGenerate: vi.fn(),
+    onUndo: vi.fn(),
+    onDownloadPdf: vi.fn(),
+    onUploadPdf: vi.fn(),
+  };
+
+  it("does not render draft actions when a signed PDF is on file while docState is still draft", () => {
+    render(
+      <CompDocDialogFooter
+        docState="draft"
+        isDocumentsReadOnly={false}
+        rowVersion="1"
+        isUploading={false}
+        signOffAttachment="https://example.com/doc.pdf"
+        generatedAt=""
+        pdfInputRef={pdfInputRef}
+        {...handlers}
+      />,
+    );
+
+    expect(screen.queryByRole("button", { name: /save draft/i })).not.toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: /mark as not required/i })).not.toBeInTheDocument();
+    expect(screen.getByText(/document on file/i)).toBeInTheDocument();
+  });
+
+  it("renders only draft actions for a clean draft without attachment", () => {
+    render(
+      <CompDocDialogFooter
+        docState="draft"
+        isDocumentsReadOnly={false}
+        rowVersion="1"
+        isUploading={false}
+        signOffAttachment={null}
+        generatedAt=""
+        pdfInputRef={pdfInputRef}
+        {...handlers}
+      />,
+    );
+
+    expect(screen.getByRole("button", { name: /save draft/i })).toBeInTheDocument();
+    expect(screen.queryByText(/document on file/i)).not.toBeInTheDocument();
+  });
+});