Consolidate document draft management into updateDocument

Replace the three separate draft mutations (createDraftDocumentVersion,
updateDocumentVersion, deleteDraftDocumentVersion) with automatic draft
lifecycle management inside updateDocument. The backend now auto-creates
a draft when a published document is edited, updates the existing draft
on subsequent edits, and auto-deletes the draft when content reverts to
match the published version.

A new deleteDocumentDraft mutation provides explicit draft deletion.

Backend:
- Merge version-level fields (content, title, classification,
  documentType) into UpdateDocumentRequest
- Convert CreateDraft, UpdateVersion, DeleteDraft into private
  transaction helpers called from Update
- Update returns (*Document, *DocumentVersion, error) with the version
  present only when a draft exists

Frontend:
- Remove all create/update/delete draft mutations from components
- Auto-save via updateDocument with layout refetch on draft status
  transitions while preserving editor cursor (data-generation key)
- Title, type, and classification editable on published versions
  (backend auto-creates draft)
- Forms use react-hook-form values option to stay synced with Relay
  fragment data across draft/publish transitions

API surface (GraphQL, MCP, CLI, n8n) updated consistently:
- Removed: createDraftDocumentVersion, updateDocumentVersion,
  deleteDraftDocumentVersion
- Added: deleteDocumentDraft (document-level)
- Updated: updateDocument accepts content, classification, documentType

Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>

apps/console/src/hooks/graph/DocumentGraph.ts

@@ -16,7 +16,6 @@ import { useTranslate } from "@probo/i18n";
 import { graphql } from "relay-runtime";
 
 import type { DocumentGraphBulkExportDocumentsMutation } from "#/__generated__/core/DocumentGraphBulkExportDocumentsMutation.graphql";
-import type { DocumentGraphDeleteDraftMutation } from "#/__generated__/core/DocumentGraphDeleteDraftMutation.graphql";
 import type { DocumentGraphDeleteMutation } from "#/__generated__/core/DocumentGraphDeleteMutation.graphql";
 import type { DocumentGraphSendSigningNotificationsMutation } from "#/__generated__/core/DocumentGraphSendSigningNotificationsMutation.graphql";
 
@@ -47,29 +46,6 @@ export function useDeleteDocumentMutation() {
   );
 }
 
-const deleteDraftDocumentVersionMutation = graphql`
-  mutation DocumentGraphDeleteDraftMutation(
-    $input: DeleteDraftDocumentVersionInput!
-    $connections: [ID!]!
-  ) {
-    deleteDraftDocumentVersion(input: $input) {
-      deletedDocumentVersionId @deleteEdge(connections: $connections)
-    }
-  }
-`;
-
-export function useDeleteDraftDocumentVersionMutation() {
-  const { __ } = useTranslate();
-
-  return useMutationWithToasts<DocumentGraphDeleteDraftMutation>(
-    deleteDraftDocumentVersionMutation,
-    {
-      successMessage: __("Draft deleted successfully."),
-      errorMessage: __("Failed to delete draft"),
-    },
-  );
-}
-
 const bulkDeleteDocumentsMutation = graphql`
   mutation DocumentGraphBulkDeleteDocumentsMutation(
     $input: BulkDeleteDocumentsInput!

apps/console/src/pages/organizations/documents/DocumentLayout.tsx

@@ -122,12 +122,18 @@ export function DocumentLayout(props: { queryRef: PreloadedQuery<DocumentLayoutQ
 
   const publishDialogRef = useRef<PublishDialogRef>(null);
   const [approvalRequestedAt, setApprovalRequestedAt] = useState(0);
+  const [versionChangedAt, setVersionChangedAt] = useState(0);
 
   const handlePublishOrApproval = useCallback(() => {
     onRefetch();
     setApprovalRequestedAt(Date.now());
   }, [onRefetch]);
 
+  const handleVersionChanged = useCallback(() => {
+    onRefetch();
+    setVersionChangedAt(Date.now());
+  }, [onRefetch]);
+
   const { document, version } = usePreloadedQuery<DocumentLayoutQuery>(documentLayoutQuery, queryRef);
   if (document.__typename !== "Document" || (version && version.__typename !== "DocumentVersion")) {
     throw new Error("invalid node type");
@@ -178,13 +184,20 @@ export function DocumentLayout(props: { queryRef: PreloadedQuery<DocumentLayoutQ
             <DocumentActionsDropdown
               documentFragmentRef={document}
               versionFragmentRef={currentVersion}
-              onRefetch={onRefetch}
+              onVersionChanged={handleVersionChanged}
             />
           </div>
         </div>
 
         <PageHeader
-          title={<DocumentTitleForm fKey={currentVersion} />}
+          title={(
+            <DocumentTitleForm
+              fKey={currentVersion}
+              documentId={document.id}
+              documentStatus={document.status}
+              onVersionChanged={handleVersionChanged}
+            />
+          )}
         />
 
         <Tabs>
@@ -207,18 +220,22 @@ export function DocumentLayout(props: { queryRef: PreloadedQuery<DocumentLayoutQ
             <TabLink to={`${urlPrefix}/signatures`}>
               {__("Signatures")}
               <TabBadge>
-                {currentVersion.signedSignatures.totalCount}
+                {currentVersion.signedSignatures?.totalCount ?? 0}
                 /
-                {currentVersion.signatures.totalCount}
+                {currentVersion.signatures?.totalCount ?? 0}
               </TabBadge>
             </TabLink>
           )}
         </Tabs>
 
-        <Outlet context={{ onRefetch, approvalRequestedAt }} />
+        <Outlet context={{ onRefetch, approvalRequestedAt, versionChangedAt }} />
       </div>
 
-      <DocumentLayoutDrawer documentFragmentRef={document} versionFragmentRef={currentVersion} />
+      <DocumentLayoutDrawer
+        documentFragmentRef={document}
+        versionFragmentRef={currentVersion}
+        onVersionChanged={handleVersionChanged}
+      />
 
       <PublishDialog
         ref={publishDialogRef}

apps/console/src/pages/organizations/documents/DocumentLayoutLoader.tsx

@@ -42,7 +42,7 @@ function DocumentLayoutQueryLoader() {
   const onRefetch = useCallback(() => {
     loadQuery(
       { documentId, versionId: versionId ?? "", versionSpecified: !!versionId },
-      { fetchPolicy: "network-only" },
+      { fetchPolicy: "store-and-network" },
     );
   }, [documentId, versionId, loadQuery]);
 

apps/console/src/pages/organizations/documents/_components/DocumentActionsDropdown.tsx

@@ -14,20 +14,20 @@
 
 import { formatError, sprintf } from "@probo/helpers";
 import { useTranslate } from "@probo/i18n";
-import { ActionDropdown, DropdownItem, IconArchive, IconArrowDown, IconPencil, IconTrashCan, useConfirm, useToast } from "@probo/ui";
+import { ActionDropdown, DropdownItem, IconArchive, IconArrowDown, IconTrashCan, useConfirm, useToast } from "@probo/ui";
 import { use, useRef } from "react";
 import { useFragment, useMutation } from "react-relay";
-import { useNavigate, useParams } from "react-router";
+import { useNavigate } from "react-router";
 import { ConnectionHandler, graphql } from "relay-runtime";
 
 import type { DocumentActionsDropdown_archiveMutation } from "#/__generated__/core/DocumentActionsDropdown_archiveMutation.graphql";
-import type { DocumentActionsDropdown_createDraftMutation } from "#/__generated__/core/DocumentActionsDropdown_createDraftMutation.graphql";
+import type { DocumentActionsDropdown_deleteDocumentDraftMutation } from "#/__generated__/core/DocumentActionsDropdown_deleteDocumentDraftMutation.graphql";
 import type { DocumentActionsDropdown_documentFragment$key } from "#/__generated__/core/DocumentActionsDropdown_documentFragment.graphql";
 import type { DocumentActionsDropdown_exportVersionMutation } from "#/__generated__/core/DocumentActionsDropdown_exportVersionMutation.graphql";
 import type { DocumentActionsDropdown_unarchiveMutation } from "#/__generated__/core/DocumentActionsDropdown_unarchiveMutation.graphql";
 import type { DocumentActionsDropdown_versionFragment$key } from "#/__generated__/core/DocumentActionsDropdown_versionFragment.graphql";
 import { PdfDownloadDialog, type PdfDownloadDialogRef } from "#/components/documents/PdfDownloadDialog";
-import { DocumentsConnectionKey, useDeleteDocumentMutation, useDeleteDraftDocumentVersionMutation } from "#/hooks/graph/DocumentGraph";
+import { DocumentsConnectionKey, useDeleteDocumentMutation } from "#/hooks/graph/DocumentGraph";
 import { useOrganizationId } from "#/hooks/useOrganizationId";
 import { CurrentUser } from "#/providers/CurrentUser";
 
@@ -35,49 +35,10 @@ const documentFragment = graphql`
   fragment DocumentActionsDropdown_documentFragment on Document {
     id
     status
-    canUpdate: permission(action: "core:document:update")
     canArchive: permission(action: "core:document:archive")
     canUnarchive: permission(action: "core:document:unarchive")
     canDelete: permission(action: "core:document:delete")
-    versions(first: 1 orderBy: { field: CREATED_AT, direction: DESC }) {
+    canDeleteDraft: permission(action: "core:document:delete-draft")
-      totalCount
-      edges {
-        node {
-          id
-          title
-          status
-        }
-      }
-    }
-  }
-`;
-
-const createDraftDocumentVersionMutation = graphql`
-  mutation DocumentActionsDropdown_createDraftMutation(
-    $input: CreateDraftDocumentVersionInput!
-    $connections: [ID!]!
-  ) {
-    createDraftDocumentVersion(input: $input) {
-      documentVersionEdge @prependEdge(connections: $connections) {
-        node {
-          id
-          content
-          status
-          publishedAt
-          major
-          minor
-          updatedAt
-          signatures(first: 100) {
-            edges {
-              node {
-                id
-                state
-              }
-            }
-          }
-        }
-      }
-    }
   }
 `;
 
@@ -90,7 +51,6 @@ const archiveDocumentMutation = graphql`
         id
         status
         archivedAt
-        canUpdate: permission(action: "core:document:update")
         canArchive: permission(action: "core:document:archive")
         canUnarchive: permission(action: "core:document:unarchive")
         canDelete: permission(action: "core:document:delete")
@@ -108,7 +68,6 @@ const unarchiveDocumentMutation = graphql`
         id
         status
         archivedAt
-        canUpdate: permission(action: "core:document:update")
         canArchive: permission(action: "core:document:archive")
         canUnarchive: permission(action: "core:document:unarchive")
         canDelete: permission(action: "core:document:delete")
@@ -117,6 +76,19 @@ const unarchiveDocumentMutation = graphql`
   }
 `;
 
+const deleteDocumentDraftMutation = graphql`
+  mutation DocumentActionsDropdown_deleteDocumentDraftMutation(
+    $input: DeleteDocumentDraftInput!
+  ) {
+    deleteDocumentDraft(input: $input) {
+      document {
+        id
+        status
+      }
+    }
+  }
+`;
+
 const versionFragment = graphql`
   fragment DocumentActionsDropdown_versionFragment on DocumentVersion {
     id
@@ -124,7 +96,6 @@ const versionFragment = graphql`
     major
     minor
     status
-    canDeleteDraft: permission(action: "core:document-version:delete-draft")
   }
 `;
 
@@ -141,13 +112,12 @@ const exportDocumentVersionMutation = graphql`
 export function DocumentActionsDropdown(props: {
   documentFragmentRef: DocumentActionsDropdown_documentFragment$key;
   versionFragmentRef: DocumentActionsDropdown_versionFragment$key;
-  onRefetch: () => void;
+  onVersionChanged: () => void;
 }) {
-  const { documentFragmentRef, versionFragmentRef, onRefetch } = props;
+  const { documentFragmentRef, versionFragmentRef, onVersionChanged } = props;
 
   const organizationId = useOrganizationId();
   const navigate = useNavigate();
-  const { versionId } = useParams();
   const { __ } = useTranslate();
   const { email: defaultEmail } = use(CurrentUser);
   const pdfDownloadDialogRef = useRef<PdfDownloadDialogRef>(null);
@@ -157,53 +127,16 @@ export function DocumentActionsDropdown(props: {
   const document = useFragment<DocumentActionsDropdown_documentFragment$key>(documentFragment, documentFragmentRef);
   const version = useFragment<DocumentActionsDropdown_versionFragment$key>(versionFragment, versionFragmentRef);
 
-  const lastVersion = document.versions.edges[0].node;
-  const isLastVersionPublished = lastVersion.status === "PUBLISHED";
-  const isDraft = version.status === "DRAFT";
-
-  const [createDraftDocumentVersion, isCreatingDraft]
-    = useMutation<DocumentActionsDropdown_createDraftMutation>(createDraftDocumentVersionMutation);
   const [deleteDocument, isDeleting] = useDeleteDocumentMutation();
   const [archiveDocument, isArchiving]
     = useMutation<DocumentActionsDropdown_archiveMutation>(archiveDocumentMutation);
   const [unarchiveDocument, isUnarchiving]
     = useMutation<DocumentActionsDropdown_unarchiveMutation>(unarchiveDocumentMutation);
-  const [deleteDraftDocumentVersion, isDeletingDraft]
+  const [deleteDocumentDraft, isDeletingDraft]
-    = useDeleteDraftDocumentVersionMutation();
+    = useMutation<DocumentActionsDropdown_deleteDocumentDraftMutation>(deleteDocumentDraftMutation);
   const [exportDocumentVersion, isExporting]
     = useMutation<DocumentActionsDropdown_exportVersionMutation>(exportDocumentVersionMutation);
 
-  const handleCreateDraft = () => {
-    const connectionId = ConnectionHandler.getConnectionID(document.id, "DocumentversionsDropdownMenu_versions");
-    createDraftDocumentVersion({
-      variables: {
-        input: {
-          documentID: document.id,
-        },
-        connections: [connectionId],
-      },
-      onCompleted: (response, errors) => {
-        if (errors) {
-          toast({
-            variant: "error",
-            title: __("Error creating draft"),
-            description:
-                  errors[0]?.message || __("An unknown error occurred"),
-          });
-          return;
-        }
-
-        const newVersionId
-          = response.createDraftDocumentVersion.documentVersionEdge.node.id;
-
-        void navigate(`/organizations/${organizationId}/documents/${document.id}/versions/${newVersionId}`);
-      },
-      onError(error) {
-        toast({ title: __("Error"), description: error.message, variant: "error" });
-      },
-    });
-  };
-
   const handleArchive = () => {
     confirm(
       () =>
@@ -227,7 +160,7 @@ export function DocumentActionsDropdown(props: {
       {
         message: sprintf(
           __("This will archive the document \"%s\". It will no longer be editable."),
-          lastVersion.title,
+          version.title,
         ),
         variant: "danger",
         label: __("Archive"),
@@ -251,6 +184,36 @@ export function DocumentActionsDropdown(props: {
     });
   };
 
+  const handleDeleteDraft = () => {
+    confirm(
+      () =>
+        new Promise<void>((resolve) => {
+          deleteDocumentDraft({
+            variables: { input: { documentId: document.id } },
+            onCompleted(_, errors) {
+              if (errors?.length) {
+                toast({ title: __("Error"), description: formatError(__("Failed to delete draft"), errors), variant: "error" });
+              } else {
+                toast({ title: __("Success"), description: __("Draft deleted successfully."), variant: "success" });
+                onVersionChanged();
+                void navigate(`/organizations/${organizationId}/documents/${document.id}/description`);
+              }
+              resolve();
+            },
+            onError(error) {
+              toast({ title: __("Error"), description: error.message, variant: "error" });
+              resolve();
+            },
+          });
+        }),
+      {
+        message: __("This will delete the current draft and revert to the last published version."),
+        variant: "danger",
+        label: __("Delete draft"),
+      },
+    );
+  };
+
   const handleDelete = () => {
     const connectionId = ConnectionHandler.getConnectionID(
       organizationId,
@@ -273,41 +236,7 @@ export function DocumentActionsDropdown(props: {
           __(
             "This will permanently delete the document \"%s\". This action cannot be undone.",
           ),
-          lastVersion.title,
+          version.title,
-        ),
-      },
-    );
-  };
-
-  const handleDeleteDraft = () => {
-    const versionsConnectionId = ConnectionHandler.getConnectionID(document.id, "DocumentversionsDropdownMenu_versions");
-    const lastVersionConnectionId = ConnectionHandler.getConnectionID(
-      document.id,
-      "DocumentversionsDropdownMenu_lastVersion",
-      { orderBy: { field: "CREATED_AT", direction: "DESC" } },
-    );
-    confirm(
-      () =>
-        deleteDraftDocumentVersion({
-          variables: {
-            input: { documentVersionId: version.id },
-            connections: [versionsConnectionId, lastVersionConnectionId],
-          },
-          onSuccess() {
-            if (versionId) {
-              void navigate(`/organizations/${organizationId}/documents/${document.id}`);
-            } else {
-              onRefetch();
-            }
-          },
-        }),
-      {
-        message: sprintf(
-          __(
-            "This will permanently delete the draft version %s of \"%s\". This action cannot be undone.",
-          ),
-          `${version.major}.${version.minor}`,
-          lastVersion.title,
         ),
       },
     );
@@ -362,26 +291,6 @@ export function DocumentActionsDropdown(props: {
         defaultEmail={defaultEmail}
       />
       <ActionDropdown variant="secondary">
-        {document.canUpdate && isLastVersionPublished && (
-          <DropdownItem
-            onClick={handleCreateDraft}
-            icon={IconPencil}
-            disabled={isCreatingDraft}
-          >
-            {__("Create new draft")}
-          </DropdownItem>
-        )}
-        {isDraft
-          && document.versions.totalCount > 1
-          && version.canDeleteDraft && (
-          <DropdownItem
-            onClick={handleDeleteDraft}
-            icon={IconTrashCan}
-            disabled={isDeletingDraft}
-          >
-            {__("Delete draft document")}
-          </DropdownItem>
-        )}
         <DropdownItem
           onClick={() => pdfDownloadDialogRef.current?.open()}
           icon={IconArrowDown}
@@ -389,6 +298,15 @@ export function DocumentActionsDropdown(props: {
         >
           {__("Download PDF")}
         </DropdownItem>
+        {document.canDeleteDraft && version.status === "DRAFT" && !(version.major === 0 && version.minor === 1) && (
+          <DropdownItem
+            icon={IconTrashCan}
+            disabled={isDeletingDraft}
+            onClick={handleDeleteDraft}
+          >
+            {__("Delete draft")}
+          </DropdownItem>
+        )}
         {document.canArchive && document.status === "ACTIVE" && (
           <DropdownItem
             icon={IconArchive}

apps/console/src/pages/organizations/documents/_components/DocumentLayoutDrawer.tsx

@@ -22,7 +22,6 @@ import { z } from "zod";
 
 import type { DocumentLayoutDrawer_documentFragment$key } from "#/__generated__/core/DocumentLayoutDrawer_documentFragment.graphql";
 import type { DocumentLayoutDrawer_updateApproversMutation } from "#/__generated__/core/DocumentLayoutDrawer_updateApproversMutation.graphql";
-import type { DocumentLayoutDrawer_updateClassificationMutation } from "#/__generated__/core/DocumentLayoutDrawer_updateClassificationMutation.graphql";
 import type { DocumentLayoutDrawer_versionFragment$key } from "#/__generated__/core/DocumentLayoutDrawer_versionFragment.graphql";
 import type { DocumentLayoutDrawerMutation } from "#/__generated__/core/DocumentLayoutDrawerMutation.graphql";
 import { ControlledField } from "#/components/form/ControlledField";
@@ -31,6 +30,7 @@ import { DocumentTypeOptions } from "#/components/form/DocumentTypeOptions";
 import { PeopleMultiSelectField } from "#/components/form/PeopleMultiSelectField";
 import { useFormWithSchema } from "#/hooks/useFormWithSchema";
 import { useOrganizationId } from "#/hooks/useOrganizationId";
+
 const documentFragment = graphql`
   fragment DocumentLayoutDrawer_documentFragment on Document {
     id
@@ -58,23 +58,21 @@ const versionFragment = graphql`
   }
 `;
 
-const updateDocumentTypeMutation = graphql`
+const updateDocumentMutation = graphql`
-  mutation DocumentLayoutDrawerMutation($input: UpdateDocumentVersionInput!) {
+  mutation DocumentLayoutDrawerMutation($input: UpdateDocumentInput!) {
-    updateDocumentVersion(input: $input) {
+    updateDocument(input: $input) {
+      document {
+        id
+      }
       documentVersion {
         id
         documentType
-      }
-    }
-  }
-`;
-
-const updateClassificationMutation = graphql`
-  mutation DocumentLayoutDrawer_updateClassificationMutation($input: UpdateDocumentVersionInput!) {
-    updateDocumentVersion(input: $input) {
-      documentVersion {
-        id
         classification
+        major
+        minor
+        status
+        updatedAt
+        publishedAt
       }
     }
   }
@@ -110,8 +108,9 @@ const approversSchema = z.object({
 export function DocumentLayoutDrawer(props: {
   documentFragmentRef: DocumentLayoutDrawer_documentFragment$key;
   versionFragmentRef: DocumentLayoutDrawer_versionFragment$key;
+  onVersionChanged: () => void;
 }) {
-  const { documentFragmentRef, versionFragmentRef } = props;
+  const { documentFragmentRef, versionFragmentRef, onVersionChanged } = props;
 
   const { __ } = useTranslate();
   const organizationId = useOrganizationId();
@@ -125,12 +124,12 @@ export function DocumentLayoutDrawer(props: {
   const version = useFragment<DocumentLayoutDrawer_versionFragment$key>(versionFragment, versionFragmentRef);
 
   const isDraft = version.status === "DRAFT";
-  const canEdit = document.canUpdate;
+  const canEdit = document.canUpdate && document.status !== "ARCHIVED";
 
   const { control, handleSubmit, reset } = useFormWithSchema(
     schema,
     {
-      defaultValues: {
+      values: {
         documentType: version.documentType,
       },
     },
@@ -143,7 +142,7 @@ export function DocumentLayoutDrawer(props: {
   } = useFormWithSchema(
     classificationSchema,
     {
-      defaultValues: {
+      values: {
         classification: version.classification,
       },
     },
@@ -156,17 +155,14 @@ export function DocumentLayoutDrawer(props: {
   } = useFormWithSchema(
     approversSchema,
     {
-      defaultValues: {
+      values: {
         approverIds: document.defaultApprovers.map(a => a.id),
       },
     },
   );
 
-  const [updateDocumentType, isUpdatingDocumentType]
+  const [updateDocument, isUpdatingDocument]
-    = useMutation<DocumentLayoutDrawerMutation>(updateDocumentTypeMutation);
+    = useMutation<DocumentLayoutDrawerMutation>(updateDocumentMutation);
-
-  const [updateClassification, isUpdatingClassification]
-    = useMutation<DocumentLayoutDrawer_updateClassificationMutation>(updateClassificationMutation);
 
   const [updateApprovers, isUpdatingApprovers]
     = useMutation<DocumentLayoutDrawer_updateApproversMutation>(updateApproversMutation);
@@ -174,15 +170,19 @@ export function DocumentLayoutDrawer(props: {
   const handleUpdateDocumentType = (data: {
     documentType: (typeof documentTypes)[number];
   }) => {
-    updateDocumentType({
+    updateDocument({
       variables: {
         input: {
-          documentVersionId: version.id,
+          id: document.id,
           documentType: data.documentType,
         },
       },
-      onCompleted: () => {
+      onCompleted: (data) => {
         setIsEditingType(false);
+        const draftReturned = !!data.updateDocument.documentVersion;
+        if (isDraft !== draftReturned) {
+          onVersionChanged();
+        }
         toast({
           title: __("Success"),
           description: __("Document type updated successfully"),
@@ -202,15 +202,19 @@ export function DocumentLayoutDrawer(props: {
   const handleUpdateClassification = (data: {
     classification: (typeof documentClassifications)[number];
   }) => {
-    updateClassification({
+    updateDocument({
       variables: {
         input: {
-          documentVersionId: version.id,
+          id: document.id,
           classification: data.classification,
         },
       },
-      onCompleted: () => {
+      onCompleted: (data) => {
         setIsEditingClassification(false);
+        const draftReturned = !!data.updateDocument.documentVersion;
+        if (isDraft !== draftReturned) {
+          onVersionChanged();
+        }
         toast({
           title: __("Success"),
           description: __("Document classification updated successfully"),
@@ -302,9 +306,9 @@ export function DocumentLayoutDrawer(props: {
                 onSave={() => void handleSubmit(handleUpdateDocumentType)()}
                 onCancel={() => {
                   setIsEditingType(false);
-                  reset();
+                  reset({ documentType: version.documentType });
                 }}
-                disabled={isUpdatingDocumentType}
+                disabled={isUpdatingDocument}
               >
                 <ControlledField
                   name="documentType"
@@ -318,7 +322,7 @@ export function DocumentLayoutDrawer(props: {
           : (
               <ReadOnlyPropertyContent
                 onEdit={() => setIsEditingType(true)}
-                canEdit={canEdit && isDraft}
+                canEdit={canEdit}
               >
                 <div className="text-sm text-txt-secondary">
                   {getDocumentTypeLabel(__, version.documentType)}
@@ -333,9 +337,9 @@ export function DocumentLayoutDrawer(props: {
                 onSave={() => void handleClassificationSubmit(handleUpdateClassification)()}
                 onCancel={() => {
                   setIsEditingClassification(false);
-                  resetClassification();
+                  resetClassification({ classification: version.classification });
                 }}
-                disabled={isUpdatingClassification}
+                disabled={isUpdatingDocument}
               >
                 <ControlledField
                   name="classification"
@@ -349,7 +353,7 @@ export function DocumentLayoutDrawer(props: {
           : (
               <ReadOnlyPropertyContent
                 onEdit={() => setIsEditingClassification(true)}
-                canEdit={canEdit && isDraft}
+                canEdit={canEdit}
               >
                 <div className="text-sm text-txt-secondary">
                   {getDocumentClassificationLabel(__, version.classification)}

apps/console/src/pages/organizations/documents/_components/DocumentTitleForm.tsx

@@ -24,9 +24,9 @@ import type { DocumentTitleFormFragment$key } from "#/__generated__/core/Documen
 import type { DocumentTitleFormMutation } from "#/__generated__/core/DocumentTitleFormMutation.graphql";
 import { useFormWithSchema } from "#/hooks/useFormWithSchema";
 
-const updateDocumentVersionTitleMutation = graphql`
+const updateDocumentTitleMutation = graphql`
-  mutation DocumentTitleFormMutation($input: UpdateDocumentVersionInput!) {
+  mutation DocumentTitleFormMutation($input: UpdateDocumentInput!) {
-    updateDocumentVersion(input: $input) {
+    updateDocument(input: $input) {
       documentVersion {
         ...DocumentTitleFormFragment
       }
@@ -36,10 +36,9 @@ const updateDocumentVersionTitleMutation = graphql`
 
 const fragment = graphql`
   fragment DocumentTitleFormFragment on DocumentVersion {
-    id
     title
     status
-    canUpdate: permission(action: "core:document-version:update")
+    canUpdate: permission(action: "core:document:update")
   }
 `;
 
@@ -47,40 +46,52 @@ const schema = z.object({
   title: z.string().min(1, "Title is required").max(255),
 });
 
-export function DocumentTitleForm(props: { fKey: DocumentTitleFormFragment$key }) {
+export function DocumentTitleForm(props: {
-  const { fKey } = props;
+  fKey: DocumentTitleFormFragment$key;
+  documentId: string;
+  documentStatus: string;
+  onVersionChanged: () => void;
+}) {
+  const { fKey, documentId, documentStatus, onVersionChanged } = props;
 
   const { __ } = useTranslate();
   const { toast } = useToast();
 
   const version = useFragment<DocumentTitleFormFragment$key>(fragment, fKey);
-  const [updateDocumentVersion, isUpdating]
+  const [updateDocument, isUpdating]
-    = useMutation<DocumentTitleFormMutation>(updateDocumentVersionTitleMutation);
+    = useMutation<DocumentTitleFormMutation>(updateDocumentTitleMutation);
 
   const [isEditingTitle, setIsEditingTitle] = useState(false);
   const { register, handleSubmit, reset } = useFormWithSchema(
     schema,
     {
-      defaultValues: {
+      values: {
         title: version.title,
       },
     },
   );
 
+  const isDraft = version.status === "DRAFT";
+  const canEdit = version.canUpdate && documentStatus !== "ARCHIVED";
+
   const handleUpdateTitle = (data: { title: string }) => {
-    updateDocumentVersion({
+    updateDocument({
       variables: {
         input: {
-          documentVersionId: version.id,
+          id: documentId,
           title: data.title,
         },
       },
-      onCompleted(_, errors) {
+      onCompleted(data, errors) {
         if (errors?.length) {
           toast({ title: __("Error"), description: formatError(__("Failed to update document"), errors), variant: "error" });
           return;
         }
         setIsEditingTitle(false);
+        const draftReturned = !!data.updateDocument.documentVersion;
+        if (isDraft !== draftReturned) {
+          onVersionChanged();
+        }
       },
       onError(error) {
         toast({ title: __("Error"), description: error.message, variant: "error" });
@@ -99,7 +110,7 @@ export function DocumentTitleForm(props: { fKey: DocumentTitleFormFragment$key }
             onKeyDown={(e) => {
               if (e.key === "Escape") {
                 setIsEditingTitle(false);
-                reset();
+                reset({ title: version.title });
               }
               if (e.key === "Enter") {
                 void handleSubmit(handleUpdateTitle)();
@@ -117,7 +128,7 @@ export function DocumentTitleForm(props: { fKey: DocumentTitleFormFragment$key }
             icon={IconCrossLargeX}
             onClick={() => {
               setIsEditingTitle(false);
-              reset();
+              reset({ title: version.title });
             }}
           />
         </div>
@@ -125,7 +136,7 @@ export function DocumentTitleForm(props: { fKey: DocumentTitleFormFragment$key }
     : (
         <div className="flex items-center gap-2">
           <span>{version.title}</span>
-          {version.canUpdate && version.status === "DRAFT" && (
+          {canEdit && (
             <Button
               variant="quaternary"
               icon={IconPencil}

apps/console/src/pages/organizations/documents/description/DocumentDescriptionPage.tsx

@@ -15,8 +15,9 @@
 import { formatError } from "@probo/helpers";
 import { useTranslate } from "@probo/i18n";
 import { RichEditor, useToast } from "@probo/ui";
-import { useCallback } from "react";
+import { useCallback, useState } from "react";
 import { type PreloadedQuery, useMutation, usePreloadedQuery } from "react-relay";
+import { useOutletContext, useParams } from "react-router";
 import { graphql } from "relay-runtime";
 import { useDebounceCallback } from "usehooks-ts";
 
@@ -39,6 +40,9 @@ export const documentDescriptionPageQuery = graphql`
     document: node(id: $documentId) {
       __typename
       ... on Document {
+        id
+        status
+        canUpdate: permission(action: "core:document:update")
         # We use this on /documents/:documentId/description
         lastVersion: versions(first: 1 orderBy: { field: CREATED_AT, direction: DESC }) @skip(if: $versionSpecified) {
           edges {
@@ -55,20 +59,30 @@ export const documentDescriptionPageQuery = graphql`
 `;
 
 const updateContentMutation = graphql`
-  mutation DocumentDescriptionPage_updateContentMutation($input: UpdateDocumentVersionInput!) {
+  mutation DocumentDescriptionPage_updateContentMutation($input: UpdateDocumentInput!) {
-    updateDocumentVersion(input: $input) {
+    updateDocument(input: $input) {
+      document {
+        id
+      }
       documentVersion {
+        id
         content
+        status
       }
     }
   }
 `;
 
-export function DocumentDescriptionPage(props: { queryRef: PreloadedQuery<DocumentDescriptionPageQuery> }) {
+export function DocumentDescriptionPage(props: {
-  const { queryRef } = props;
+  queryRef: PreloadedQuery<DocumentDescriptionPageQuery>;
+  versionChangedAt: number;
+}) {
+  const { queryRef, versionChangedAt } = props;
 
   const { __ } = useTranslate();
   const { toast } = useToast();
+  const { versionId } = useParams();
+  const { onRefetch } = useOutletContext<{ onRefetch: () => void }>();
 
   const { document, version } = usePreloadedQuery<DocumentDescriptionPageQuery>(
     documentDescriptionPageQuery,
@@ -81,18 +95,21 @@ export function DocumentDescriptionPage(props: { queryRef: PreloadedQuery<Docume
   const lastVersion = document.lastVersion?.edges[0].node;
   const currentVersion = lastVersion ?? version as NonNullable<typeof lastVersion | typeof version>;
 
-  const [updateContent, _] = useMutation<DocumentDescriptionPage_updateContentMutation>(updateContentMutation);
+  const [updateContent] = useMutation<DocumentDescriptionPage_updateContentMutation>(updateContentMutation);
+
+  const documentId = document.id;
+  const wasDraft = currentVersion.status === "DRAFT";
 
   const handleUpdate = useDebounceCallback(
     useCallback((content: string) => {
       updateContent({
         variables: {
           input: {
-            documentVersionId: currentVersion.id,
+            id: documentId,
             content,
           },
         },
-        onCompleted: (_, errors) => {
+        onCompleted: (data, errors) => {
           if (errors?.length) {
             toast({
               title: __("Error"),
@@ -102,6 +119,15 @@ export function DocumentDescriptionPage(props: { queryRef: PreloadedQuery<Docume
             return;
           }
 
+          // Refetch the layout when draft status changes (draft created
+          // or auto-deleted) so the drawer and header reflect the current
+          // version. This does NOT remount the editor because the editor
+          // key is based on versionChangedAt (explicit actions only).
+          const draftReturned = !!data.updateDocument.documentVersion;
+          if (wasDraft !== draftReturned) {
+            onRefetch();
+          }
+
           toast({
             title: __("Success"),
             description: __("Content saved"),
@@ -116,16 +142,60 @@ export function DocumentDescriptionPage(props: { queryRef: PreloadedQuery<Docume
           });
         },
       });
-    }, [currentVersion.id, updateContent, toast, __]),
+    }, [documentId, wasDraft, updateContent, toast, __, onRefetch]),
     autoSaveIntervalMs,
   );
 
+  // When viewing a specific historical version, the editor is read-only.
+  // When viewing the latest version, editing is allowed if the user has
+  // update permission and the document is not archived — the backend
+  // will auto-create a draft if needed.
+  const isViewingSpecificVersion = !!version;
+  const canEdit = !isViewingSpecificVersion
+    && document.canUpdate
+    && document.status !== "ARCHIVED";
+
+  // The editor key must change on explicit actions (delete draft, edit
+  // title/type) but NOT on auto-save side effects (cursor preservation).
+  // We track a "data generation" that only increments when an explicit
+  // action (versionChangedAt change) is followed by fresh data arriving
+  // (currentVersion.id change). This uses React's "adjust state during
+  // render" pattern so we avoid refs-during-render and setState-in-effects.
+  const [prevVCA, setPrevVCA] = useState(versionChangedAt);
+  const [prevVersionId, setPrevVersionId] = useState(currentVersion.id);
+  const [dataGeneration, setDataGeneration] = useState(0);
+  const [pendingExplicit, setPendingExplicit] = useState(false);
+
+  if (versionChangedAt !== prevVCA) {
+    setPrevVCA(versionChangedAt);
+    if (currentVersion.id !== prevVersionId) {
+      // Both changed at once — data was already available.
+      setPrevVersionId(currentVersion.id);
+      setDataGeneration(g => g + 1);
+      setPendingExplicit(false);
+    } else {
+      // Explicit action fired but data hasn't arrived yet.
+      setPendingExplicit(true);
+    }
+  } else if (currentVersion.id !== prevVersionId) {
+    setPrevVersionId(currentVersion.id);
+    if (pendingExplicit) {
+      // Fresh data arrived for a pending explicit action — remount.
+      setDataGeneration(g => g + 1);
+      setPendingExplicit(false);
+    }
+    // Otherwise auto-save changed the version — don't bump generation.
+  }
+
+  const editorKey = `${versionId ?? "latest"}-${dataGeneration}`;
+
   return (
     <RichEditor
+      key={editorKey}
       className="flex-1"
       content={currentVersion.content}
       data-theme="document"
-      disabled={currentVersion.status !== "DRAFT"}
+      disabled={!canEdit}
       onChangeContent={handleUpdate}
     />
   );

apps/console/src/pages/organizations/documents/description/DocumentDescriptionPageLoader.tsx

@@ -14,7 +14,7 @@
 
 import { useEffect } from "react";
 import { useQueryLoader } from "react-relay";
-import { useParams } from "react-router";
+import { useOutletContext, useParams } from "react-router";
 
 import type { DocumentDescriptionPageQuery } from "#/__generated__/core/DocumentDescriptionPageQuery.graphql";
 import { LinkCardSkeleton } from "#/components/skeletons/LinkCardSkeleton";
@@ -28,23 +28,26 @@ function DocumentDescriptionPageQueryLoader() {
     throw new Error(":documentId missing in route params");
   }
 
+  const { versionChangedAt } = useOutletContext<{ versionChangedAt: number }>();
   const [queryRef, loadQuery] = useQueryLoader<DocumentDescriptionPageQuery>(documentDescriptionPageQuery);
 
   useEffect(() => {
-    if (!queryRef) {
+    loadQuery(
-      loadQuery({
+      { documentId, versionId: versionId ?? "", versionSpecified: !!versionId },
-        documentId: documentId,
+      { fetchPolicy: versionChangedAt > 0 ? "network-only" : "store-or-network" },
-        versionId: versionId ?? "",
+    );
-        versionSpecified: !!versionId,
+  }, [documentId, versionId, versionChangedAt, loadQuery]);
-      });
-    }
-  });
 
   if (!queryRef) {
     return <LinkCardSkeleton />;
   }
 
-  return <DocumentDescriptionPage queryRef={queryRef} />;
+  return (
+    <DocumentDescriptionPage
+      queryRef={queryRef}
+      versionChangedAt={versionChangedAt}
+    />
+  );
 }
 
 export default function DocumentDescriptionPageLoader() {

e2e/console/document_test.go

@@ -290,18 +290,17 @@ func TestDocument_Update(t *testing.T) {
 	owner := testutil.NewClient(t, testutil.RoleOwner)
 
 	t.Run(
-		"update title via document version",
+		"update title via document",
 		func(t *testing.T) {
 			t.Parallel()
 
 			doc := factory.NewDocument(owner).
 				WithTitle("Document to Update")
-			doc.Create()
+			documentID := doc.Create()
-			versionID := doc.VersionID()
 
 			query := `
-				mutation UpdateDocumentVersion($input: UpdateDocumentVersionInput!) {
+				mutation UpdateDocument($input: UpdateDocumentInput!) {
-					updateDocumentVersion(input: $input) {
+					updateDocument(input: $input) {
 						documentVersion {
 							id
 							title
@@ -311,33 +310,32 @@ func TestDocument_Update(t *testing.T) {
 			`
 
 			var result struct {
-				UpdateDocumentVersion struct {
+				UpdateDocument struct {
 					DocumentVersion struct {
 						ID    string `json:"id"`
 						Title string `json:"title"`
 					} `json:"documentVersion"`
-				} `json:"updateDocumentVersion"`
+				} `json:"updateDocument"`
 			}
 
 			err := owner.Execute(query, map[string]any{
 				"input": map[string]any{
-					"documentVersionId": versionID,
+					"id":    documentID,
-					"title":             "Updated Document Title",
+					"title": "Updated Document Title",
 				},
 			}, &result)
 			require.NoError(t, err)
-			assert.Equal(t, "Updated Document Title", result.UpdateDocumentVersion.DocumentVersion.Title)
+			assert.Equal(t, "Updated Document Title", result.UpdateDocument.DocumentVersion.Title)
 		},
 	)
 }
 
-func TestDocumentVersion_Update_TitleValidation(t *testing.T) {
+func TestDocument_Update_TitleValidation(t *testing.T) {
 	t.Parallel()
 	owner := testutil.NewClient(t, testutil.RoleOwner)
 
 	doc := factory.NewDocument(owner).WithTitle("Validation Test Document")
-	doc.Create()
+	baseDocumentID := doc.Create()
-	baseVersionID := doc.VersionID()
 
 	tests := []struct {
 		name              string
@@ -347,41 +345,41 @@ func TestDocumentVersion_Update_TitleValidation(t *testing.T) {
 	}{
 		{
 			name:  "title with HTML tags",
-			setup: func() string { return baseVersionID },
+			setup: func() string { return baseDocumentID },
 			input: func(id string) map[string]any {
-				return map[string]any{"documentVersionId": id, "title": "<script>alert('xss')</script>"}
+				return map[string]any{"id": id, "title": "<script>alert('xss')</script>"}
 			},
 			wantErrorContains: "HTML",
 		},
 		{
 			name:  "title with newline",
-			setup: func() string { return baseVersionID },
+			setup: func() string { return baseDocumentID },
 			input: func(id string) map[string]any {
-				return map[string]any{"documentVersionId": id, "title": "Test\nDocument"}
+				return map[string]any{"id": id, "title": "Test\nDocument"}
 			},
 			wantErrorContains: "newline",
 		},
 		{
 			name:  "title with carriage return",
-			setup: func() string { return baseVersionID },
+			setup: func() string { return baseDocumentID },
 			input: func(id string) map[string]any {
-				return map[string]any{"documentVersionId": id, "title": "Test\rDocument"}
+				return map[string]any{"id": id, "title": "Test\rDocument"}
 			},
 			wantErrorContains: "carriage return",
 		},
 		{
 			name:  "title with null byte",
-			setup: func() string { return baseVersionID },
+			setup: func() string { return baseDocumentID },
 			input: func(id string) map[string]any {
-				return map[string]any{"documentVersionId": id, "title": "Test\x00Document"}
+				return map[string]any{"id": id, "title": "Test\x00Document"}
 			},
 			wantErrorContains: "control character",
 		},
 		{
 			name:  "title with zero-width space",
-			setup: func() string { return baseVersionID },
+			setup: func() string { return baseDocumentID },
 			input: func(id string) map[string]any {
-				return map[string]any{"documentVersionId": id, "title": "Test\u200BDocument"}
+				return map[string]any{"id": id, "title": "Test\u200BDocument"}
 			},
 			wantErrorContains: "zero-width",
 		},
@@ -389,11 +387,11 @@ func TestDocumentVersion_Update_TitleValidation(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			versionID := tt.setup()
+			documentID := tt.setup()
 
 			query := `
-				mutation UpdateDocumentVersion($input: UpdateDocumentVersionInput!) {
+				mutation UpdateDocument($input: UpdateDocumentInput!) {
-					updateDocumentVersion(input: $input) {
+					updateDocument(input: $input) {
 						documentVersion {
 							id
 						}
@@ -401,7 +399,7 @@ func TestDocumentVersion_Update_TitleValidation(t *testing.T) {
 				}
 			`
 
-			_, err := owner.Do(query, map[string]any{"input": tt.input(versionID)})
+			_, err := owner.Do(query, map[string]any{"input": tt.input(documentID)})
 			require.Error(t, err)
 			assert.Contains(t, err.Error(), tt.wantErrorContains)
 		})
@@ -979,14 +977,13 @@ func TestDocument_MaxLength_Validation(t *testing.T) {
 		assert.Contains(t, err.Error(), "title")
 	})
 
-	t.Run("update version", func(t *testing.T) {
+	t.Run("update document with long title", func(t *testing.T) {
 		doc := factory.NewDocument(owner).WithTitle("Max Length Test")
-		doc.Create()
+		documentID := doc.Create()
-		versionID := doc.VersionID()
 
 		query := `
-			mutation UpdateDocumentVersion($input: UpdateDocumentVersionInput!) {
+			mutation UpdateDocument($input: UpdateDocumentInput!) {
-				updateDocumentVersion(input: $input) {
+				updateDocument(input: $input) {
 					documentVersion { id }
 				}
 			}
@@ -994,8 +991,8 @@ func TestDocument_MaxLength_Validation(t *testing.T) {
 
 		_, err := owner.Do(query, map[string]any{
 			"input": map[string]any{
-				"documentVersionId": versionID,
+				"id":    documentID,
-				"title":             longTitle,
+				"title": longTitle,
 			},
 		})
 		require.Error(t, err)
@@ -1028,15 +1025,14 @@ func TestDocument_MaxLength_Validation(t *testing.T) {
 		assert.Contains(t, err.Error(), "content")
 	})
 
-	t.Run("update version with long content", func(t *testing.T) {
+	t.Run("update document with long content", func(t *testing.T) {
-		docID, versionID := createTestDocument(t, owner)
+		docID, _ := createTestDocument(t, owner)
 		require.NotEmpty(t, docID)
-		require.NotEmpty(t, versionID)
 
 		query := `
-			mutation UpdateDocumentVersion($input: UpdateDocumentVersionInput!) {
+			mutation UpdateDocument($input: UpdateDocumentInput!) {
-				updateDocumentVersion(input: $input) {
+				updateDocument(input: $input) {
-					documentVersion { id }
+					document { id }
 				}
 			}
 		`
@@ -1045,8 +1041,8 @@ func TestDocument_MaxLength_Validation(t *testing.T) {
 
 		_, err := owner.Do(query, map[string]any{
 			"input": map[string]any{
-				"documentVersionId": versionID,
+				"id":      docID,
-				"content":           longContent,
+				"content": longContent,
 			},
 		})
 		require.Error(t, err)

e2e/console/document_version_test.go

@@ -220,7 +220,7 @@ func TestDocumentVersion_PublishVersion(t *testing.T) {
 	assert.Equal(t, 0, result.Node.Versions.Edges[0].Node.Minor)
 }
 
-func TestDocumentVersion_CreateDraft(t *testing.T) {
+func TestDocumentVersion_AutoCreateDraft(t *testing.T) {
 	t.Parallel()
 	owner := testutil.NewClient(t, testutil.RoleOwner)
 
@@ -228,38 +228,112 @@ func TestDocumentVersion_CreateDraft(t *testing.T) {
 	docID, _ := createTestDocument(t, owner)
 	approveTestDocument(t, owner, docID)
 
+	// Updating content should auto-create a draft
 	query := `
-		mutation CreateDraftDocumentVersion($input: CreateDraftDocumentVersionInput!) {
+		mutation UpdateDocument($input: UpdateDocumentInput!) {
-			createDraftDocumentVersion(input: $input) {
+			updateDocument(input: $input) {
-				documentVersionEdge {
+				document {
-					node {
+					id
-						id
+				}
-						status
+				documentVersion {
-					}
+					id
+					status
+					content
 				}
 			}
 		}
 	`
 
 	var result struct {
-		CreateDraftDocumentVersion struct {
+		UpdateDocument struct {
-			DocumentVersionEdge struct {
+			Document struct {
-				Node struct {
+				ID string `json:"id"`
-					ID     string `json:"id"`
+			} `json:"document"`
-					Status string `json:"status"`
+			DocumentVersion *struct {
-				} `json:"node"`
+				ID      string `json:"id"`
-			} `json:"documentVersionEdge"`
+				Status  string `json:"status"`
-		} `json:"createDraftDocumentVersion"`
+				Content string `json:"content"`
+			} `json:"documentVersion"`
+		} `json:"updateDocument"`
 	}
 
 	err := owner.Execute(query, map[string]any{
 		"input": map[string]any{
-			"documentID": docID,
+			"id":      docID,
+			"content": testutil.ProseMirrorTextDoc("Updated content"),
 		},
 	}, &result)
 	require.NoError(t, err)
 
-	assert.Equal(t, "DRAFT", result.CreateDraftDocumentVersion.DocumentVersionEdge.Node.Status)
+	require.NotNil(t, result.UpdateDocument.DocumentVersion)
+	assert.Equal(t, "DRAFT", result.UpdateDocument.DocumentVersion.Status)
+}
+
+func TestDocumentVersion_AutoDeleteDraft(t *testing.T) {
+	t.Parallel()
+	owner := testutil.NewClient(t, testutil.RoleOwner)
+
+	// Create and approve a document (auto-publishes on approval)
+	docID, _ := createTestDocument(t, owner)
+	approveTestDocument(t, owner, docID)
+
+	// First update to create a draft
+	query := `
+		mutation UpdateDocument($input: UpdateDocumentInput!) {
+			updateDocument(input: $input) {
+				document {
+					id
+				}
+				documentVersion {
+					id
+					status
+				}
+			}
+		}
+	`
+
+	var createResult struct {
+		UpdateDocument struct {
+			Document struct {
+				ID string `json:"id"`
+			} `json:"document"`
+			DocumentVersion *struct {
+				ID     string `json:"id"`
+				Status string `json:"status"`
+			} `json:"documentVersion"`
+		} `json:"updateDocument"`
+	}
+
+	err := owner.Execute(query, map[string]any{
+		"input": map[string]any{
+			"id":      docID,
+			"content": testutil.ProseMirrorTextDoc("Updated content"),
+		},
+	}, &createResult)
+	require.NoError(t, err)
+	require.NotNil(t, createResult.UpdateDocument.DocumentVersion)
+
+	// Now revert content to match the published version — draft should be auto-deleted
+	var revertResult struct {
+		UpdateDocument struct {
+			Document struct {
+				ID string `json:"id"`
+			} `json:"document"`
+			DocumentVersion *struct {
+				ID string `json:"id"`
+			} `json:"documentVersion"`
+		} `json:"updateDocument"`
+	}
+
+	err = owner.Execute(query, map[string]any{
+		"input": map[string]any{
+			"id":      docID,
+			"content": testutil.ProseMirrorTextDoc("Initial content"),
+		},
+	}, &revertResult)
+	require.NoError(t, err)
+
+	assert.Nil(t, revertResult.UpdateDocument.DocumentVersion)
 }
 
 func TestDocumentVersion_RequestSignature(t *testing.T) {
@@ -536,18 +610,17 @@ func TestDocumentVersion_BulkPublishMinorSkipsPendingApproval(t *testing.T) {
 	docID, _ := createTestDocument(t, owner)
 	approveTestDocument(t, owner, docID)
 
-	// Create a draft so we can publish minor
+	// Create a draft by updating content (auto-creates draft)
 	_, err := owner.Do(`
-		mutation($input: CreateDraftDocumentVersionInput!) {
+		mutation($input: UpdateDocumentInput!) {
-			createDraftDocumentVersion(input: $input) {
+			updateDocument(input: $input) {
-				documentVersionEdge {
+				documentVersion { id }
-					node { id }
-				}
 			}
 		}
 	`, map[string]any{
 		"input": map[string]any{
-			"documentID": docID,
+			"id":      docID,
+			"content": testutil.ProseMirrorTextDoc("Updated content to create a draft"),
 		},
 	})
 	require.NoError(t, err)
@@ -645,6 +718,128 @@ func TestDocumentVersion_BulkRequestSignatures(t *testing.T) {
 	}
 }
 
+func TestDocumentVersion_AutoCreateDraftOnClassificationOrTypeUpdate(t *testing.T) {
+	t.Parallel()
+	owner := testutil.NewClient(t, testutil.RoleOwner)
+
+	t.Run("documentType update creates draft", func(t *testing.T) {
+		t.Parallel()
+
+		// Create and approve a document (auto-publishes on approval)
+		docID, _ := createTestDocument(t, owner)
+		approveTestDocument(t, owner, docID)
+
+		// Updating documentType should auto-create a draft
+		query := `
+			mutation UpdateDocument($input: UpdateDocumentInput!) {
+				updateDocument(input: $input) {
+					document {
+						id
+					}
+					documentVersion {
+						id
+						status
+					}
+				}
+			}
+		`
+
+		var result struct {
+			UpdateDocument struct {
+				Document struct {
+					ID string `json:"id"`
+				} `json:"document"`
+				DocumentVersion *struct {
+					ID     string `json:"id"`
+					Status string `json:"status"`
+				} `json:"documentVersion"`
+			} `json:"updateDocument"`
+		}
+
+		err := owner.Execute(query, map[string]any{
+			"input": map[string]any{
+				"id":           docID,
+				"documentType": "PROCEDURE",
+			},
+		}, &result)
+		require.NoError(t, err)
+
+		require.NotNil(t, result.UpdateDocument.DocumentVersion)
+		assert.Equal(t, "DRAFT", result.UpdateDocument.DocumentVersion.Status)
+	})
+
+	t.Run("classification update creates draft", func(t *testing.T) {
+		t.Parallel()
+
+		// Create and approve a document (auto-publishes on approval)
+		docID, _ := createTestDocument(t, owner)
+		approveTestDocument(t, owner, docID)
+
+		// Updating classification should auto-create a draft
+		query := `
+			mutation UpdateDocument($input: UpdateDocumentInput!) {
+				updateDocument(input: $input) {
+					document {
+						id
+					}
+					documentVersion {
+						id
+						status
+					}
+				}
+			}
+		`
+
+		var result struct {
+			UpdateDocument struct {
+				Document struct {
+					ID string `json:"id"`
+				} `json:"document"`
+				DocumentVersion *struct {
+					ID     string `json:"id"`
+					Status string `json:"status"`
+				} `json:"documentVersion"`
+			} `json:"updateDocument"`
+		}
+
+		err := owner.Execute(query, map[string]any{
+			"input": map[string]any{
+				"id":             docID,
+				"classification": "CONFIDENTIAL",
+			},
+		}, &result)
+		require.NoError(t, err)
+
+		require.NotNil(t, result.UpdateDocument.DocumentVersion)
+		assert.Equal(t, "DRAFT", result.UpdateDocument.DocumentVersion.Status)
+	})
+}
+
+func TestDocumentVersion_ViewerCannotUpdateDocument(t *testing.T) {
+	t.Parallel()
+	owner := testutil.NewClient(t, testutil.RoleOwner)
+	viewer := testutil.NewClientInOrg(t, testutil.RoleViewer, owner)
+
+	// Create and approve a document (auto-publishes on approval)
+	docID, _ := createTestDocument(t, owner)
+	approveTestDocument(t, owner, docID)
+
+	// Viewer attempts to update content on the published document
+	_, err := viewer.Do(`
+		mutation UpdateDocument($input: UpdateDocumentInput!) {
+			updateDocument(input: $input) {
+				document { id }
+			}
+		}
+	`, map[string]any{
+		"input": map[string]any{
+			"id":      docID,
+			"content": testutil.ProseMirrorTextDoc("Viewer updated content"),
+		},
+	})
+	testutil.RequireForbiddenError(t, err, "viewer should not be able to update document")
+}
+
 func TestDocumentVersion_BulkDelete(t *testing.T) {
 	t.Parallel()
 	owner := testutil.NewClient(t, testutil.RoleOwner)
@@ -1143,3 +1338,148 @@ func TestDocument_DefaultApprovers(t *testing.T) {
 		assert.Empty(t, result.UpdateDocument.Document.DefaultApprovers)
 	})
 }
+
+func TestDocumentVersion_DeleteDraft(t *testing.T) {
+	t.Parallel()
+	owner := testutil.NewClient(t, testutil.RoleOwner)
+
+	const query = `
+		mutation DeleteDocumentDraft($input: DeleteDocumentDraftInput!) {
+			deleteDocumentDraft(input: $input) {
+				document {
+					id
+				}
+			}
+		}
+	`
+
+	t.Run(
+		"delete draft after publishing",
+		func(t *testing.T) {
+			t.Parallel()
+
+			docID, _ := createTestDocument(t, owner)
+			approveTestDocument(t, owner, docID)
+
+			// Create a draft by updating content
+			updateQuery := `
+				mutation UpdateDocument($input: UpdateDocumentInput!) {
+					updateDocument(input: $input) {
+						document { id }
+						documentVersion { id status }
+					}
+				}
+			`
+			var updateResult struct {
+				UpdateDocument struct {
+					Document struct {
+						ID string `json:"id"`
+					} `json:"document"`
+					DocumentVersion *struct {
+						ID     string `json:"id"`
+						Status string `json:"status"`
+					} `json:"documentVersion"`
+				} `json:"updateDocument"`
+			}
+			err := owner.Execute(updateQuery, map[string]any{
+				"input": map[string]any{
+					"id":      docID,
+					"content": testutil.ProseMirrorTextDoc("Draft content"),
+				},
+			}, &updateResult)
+			require.NoError(t, err)
+			require.NotNil(t, updateResult.UpdateDocument.DocumentVersion)
+			assert.Equal(t, "DRAFT", updateResult.UpdateDocument.DocumentVersion.Status)
+
+			// Now delete the draft
+			var result struct {
+				DeleteDocumentDraft struct {
+					Document struct {
+						ID string `json:"id"`
+					} `json:"document"`
+				} `json:"deleteDocumentDraft"`
+			}
+			err = owner.Execute(query, map[string]any{
+				"input": map[string]any{"documentId": docID},
+			}, &result)
+			require.NoError(t, err)
+			assert.Equal(t, docID, result.DeleteDocumentDraft.Document.ID)
+		},
+	)
+
+	t.Run(
+		"cannot delete initial v0.1 draft",
+		func(t *testing.T) {
+			t.Parallel()
+
+			docID, _ := createTestDocument(t, owner)
+
+			var result struct{}
+			err := owner.Execute(query, map[string]any{
+				"input": map[string]any{"documentId": docID},
+			}, &result)
+			require.Error(t, err)
+		},
+	)
+
+	t.Run(
+		"cannot delete when latest is published",
+		func(t *testing.T) {
+			t.Parallel()
+
+			docID, _ := createTestDocument(t, owner)
+			approveTestDocument(t, owner, docID)
+
+			var result struct{}
+			err := owner.Execute(query, map[string]any{
+				"input": map[string]any{"documentId": docID},
+			}, &result)
+			require.Error(t, err)
+		},
+	)
+
+	t.Run(
+		"viewer cannot delete draft",
+		func(t *testing.T) {
+			t.Parallel()
+
+			viewer := testutil.NewClientInOrg(t, testutil.RoleViewer, owner)
+
+			docID, _ := createTestDocument(t, owner)
+			approveTestDocument(t, owner, docID)
+
+			// Create a draft
+			updateQuery := `
+				mutation UpdateDocument($input: UpdateDocumentInput!) {
+					updateDocument(input: $input) {
+						document { id }
+						documentVersion { id }
+					}
+				}
+			`
+			var updateResult struct {
+				UpdateDocument struct {
+					Document struct {
+						ID string `json:"id"`
+					} `json:"document"`
+					DocumentVersion *struct {
+						ID string `json:"id"`
+					} `json:"documentVersion"`
+				} `json:"updateDocument"`
+			}
+			err := owner.Execute(updateQuery, map[string]any{
+				"input": map[string]any{
+					"id":      docID,
+					"content": testutil.ProseMirrorTextDoc("Draft content"),
+				},
+			}, &updateResult)
+			require.NoError(t, err)
+
+			var result struct{}
+			err = viewer.Execute(query, map[string]any{
+				"input": map[string]any{"documentId": docID},
+			}, &result)
+			testutil.RequireForbiddenError(t, err, "viewer should not be able to delete document draft")
+		},
+	)
+}

pkg/cmd/document/create-draft/create_draft.go

@@ -1,111 +0,0 @@
-// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
-//
-// Permission to use, copy, modify, and/or distribute this software for any
-// purpose with or without fee is hereby granted, provided that the above
-// copyright notice and this permission notice appear in all copies.
-//
-// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
-// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
-// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
-// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
-// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
-// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
-// PERFORMANCE OF THIS SOFTWARE.
-
-package createdraft
-
-import (
-	"encoding/json"
-	"fmt"
-
-	"github.com/spf13/cobra"
-	"go.probo.inc/probo/pkg/cli/api"
-	"go.probo.inc/probo/pkg/cmd/cmdutil"
-)
-
-const createDraftMutation = `
-mutation($input: CreateDraftDocumentVersionInput!) {
-  createDraftDocumentVersion(input: $input) {
-    documentVersionEdge {
-      node {
-        id
-        title
-        major
-        minor
-        status
-      }
-    }
-  }
-}
-`
-
-type createDraftResponse struct {
-	CreateDraftDocumentVersion struct {
-		DocumentVersionEdge struct {
-			Node struct {
-				ID     string `json:"id"`
-				Title  string `json:"title"`
-				Major  int    `json:"major"`
-				Minor  int    `json:"minor"`
-				Status string `json:"status"`
-			} `json:"node"`
-		} `json:"documentVersionEdge"`
-	} `json:"createDraftDocumentVersion"`
-}
-
-func NewCmdCreateDraft(f *cmdutil.Factory) *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "create-draft <document-id>",
-		Short: "Create a new draft version of a document",
-		Args:  cobra.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			cfg, err := f.Config()
-			if err != nil {
-				return err
-			}
-
-			host, hc, err := cfg.DefaultHost()
-			if err != nil {
-				return err
-			}
-
-			client := api.NewClient(
-				host,
-				hc.Token,
-				"/api/console/v1/graphql",
-				cfg.HTTPTimeoutDuration(),
-			)
-
-			data, err := client.Do(
-				createDraftMutation,
-				map[string]any{
-					"input": map[string]any{
-						"documentID": args[0],
-					},
-				},
-			)
-			if err != nil {
-				return err
-			}
-
-			var resp createDraftResponse
-			if err := json.Unmarshal(data, &resp); err != nil {
-				return fmt.Errorf("cannot parse response: %w", err)
-			}
-
-			v := resp.CreateDraftDocumentVersion.DocumentVersionEdge.Node
-			_, _ = fmt.Fprintf(
-				f.IOStreams.Out,
-				"Created draft version %s (%s v%d.%d)\n",
-				v.ID,
-				v.Title,
-				v.Major,
-				v.Minor,
-			)
-
-			return nil
-		},
-	}
-
-	return cmd
-}

pkg/cmd/document/delete-draft/delete_draft.go

@@ -24,9 +24,11 @@ import (
 )
 
 const deleteDraftMutation = `
-mutation($input: DeleteDraftDocumentVersionInput!) {
+mutation($input: DeleteDocumentDraftInput!) {
-  deleteDraftDocumentVersion(input: $input) {
+  deleteDocumentDraft(input: $input) {
-    deletedDocumentVersionId
+    document {
+      id
+    }
   }
 }
 `
@@ -35,8 +37,8 @@ func NewCmdDeleteDraft(f *cmdutil.Factory) *cobra.Command {
 	var flagYes bool
 
 	cmd := &cobra.Command{
-		Use:   "delete-draft <document-version-id>",
+		Use:   "delete-draft <document-id>",
-		Short: "Delete a draft document version",
+		Short: "Delete the draft version of a document",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if !flagYes {
@@ -46,7 +48,7 @@ func NewCmdDeleteDraft(f *cmdutil.Factory) *cobra.Command {
 
 				var confirmed bool
 				err := huh.NewConfirm().
-					Title(fmt.Sprintf("Delete draft version %s?", args[0])).
+					Title(fmt.Sprintf("Delete draft for document %s?", args[0])).
 					Value(&confirmed).
 					Run()
 				if err != nil {
@@ -78,7 +80,7 @@ func NewCmdDeleteDraft(f *cmdutil.Factory) *cobra.Command {
 				deleteDraftMutation,
 				map[string]any{
 					"input": map[string]any{
-						"documentVersionId": args[0],
+						"documentId": args[0],
 					},
 				},
 			)
@@ -88,7 +90,7 @@ func NewCmdDeleteDraft(f *cmdutil.Factory) *cobra.Command {
 
 			_, _ = fmt.Fprintf(
 				f.IOStreams.Out,
-				"Deleted draft version %s\n",
+				"Deleted draft for document %s\n",
 				args[0],
 			)
 

pkg/cmd/document/document.go

@@ -19,7 +19,6 @@ import (
 	"go.probo.inc/probo/pkg/cmd/cmdutil"
 	"go.probo.inc/probo/pkg/cmd/document/archive"
 	"go.probo.inc/probo/pkg/cmd/document/create"
-	createdraft "go.probo.inc/probo/pkg/cmd/document/create-draft"
 	"go.probo.inc/probo/pkg/cmd/document/delete"
 	deletedraft "go.probo.inc/probo/pkg/cmd/document/delete-draft"
 	"go.probo.inc/probo/pkg/cmd/document/list"
@@ -28,7 +27,6 @@ import (
 	publishminor "go.probo.inc/probo/pkg/cmd/document/publish-minor"
 	"go.probo.inc/probo/pkg/cmd/document/unarchive"
 	"go.probo.inc/probo/pkg/cmd/document/update"
-	updateversion "go.probo.inc/probo/pkg/cmd/document/update-version"
 	"go.probo.inc/probo/pkg/cmd/document/view"
 	viewversion "go.probo.inc/probo/pkg/cmd/document/view-version"
 )
@@ -48,9 +46,7 @@ func NewCmdDocument(f *cmdutil.Factory) *cobra.Command {
 	cmd.AddCommand(unarchive.NewCmdUnarchive(f))
 	cmd.AddCommand(listversions.NewCmdListVersions(f))
 	cmd.AddCommand(viewversion.NewCmdViewVersion(f))
-	cmd.AddCommand(createdraft.NewCmdCreateDraft(f))
 	cmd.AddCommand(deletedraft.NewCmdDeleteDraft(f))
-	cmd.AddCommand(updateversion.NewCmdUpdateVersion(f))
 	cmd.AddCommand(publishmajor.NewCmdPublishMajor(f))
 	cmd.AddCommand(publishminor.NewCmdPublishMinor(f))
 

pkg/cmd/document/update-version/update_version.go

@@ -1,154 +0,0 @@
-// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
-//
-// Permission to use, copy, modify, and/or distribute this software for any
-// purpose with or without fee is hereby granted, provided that the above
-// copyright notice and this permission notice appear in all copies.
-//
-// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
-// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
-// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
-// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
-// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
-// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
-// PERFORMANCE OF THIS SOFTWARE.
-
-package updateversion
-
-import (
-	"encoding/json"
-	"fmt"
-
-	"github.com/spf13/cobra"
-	"go.probo.inc/probo/pkg/cli/api"
-	"go.probo.inc/probo/pkg/cmd/cmdutil"
-)
-
-const updateMutation = `
-mutation($input: UpdateDocumentVersionInput!) {
-  updateDocumentVersion(input: $input) {
-    documentVersion {
-      id
-      title
-      major
-      minor
-      status
-      documentType
-      classification
-    }
-  }
-}
-`
-
-type updateResponse struct {
-	UpdateDocumentVersion struct {
-		DocumentVersion struct {
-			ID             string `json:"id"`
-			Title          string `json:"title"`
-			Major          int    `json:"major"`
-			Minor          int    `json:"minor"`
-			Status         string `json:"status"`
-			DocumentType   string `json:"documentType"`
-			Classification string `json:"classification"`
-		} `json:"documentVersion"`
-	} `json:"updateDocumentVersion"`
-}
-
-func NewCmdUpdateVersion(f *cmdutil.Factory) *cobra.Command {
-	var (
-		flagTitle          string
-		flagContent        string
-		flagDocumentType   string
-		flagClassification string
-	)
-
-	cmd := &cobra.Command{
-		Use:   "update-version <document-version-id>",
-		Short: "Update a document version",
-		Args:  cobra.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			cfg, err := f.Config()
-			if err != nil {
-				return err
-			}
-
-			host, hc, err := cfg.DefaultHost()
-			if err != nil {
-				return err
-			}
-
-			client := api.NewClient(
-				host,
-				hc.Token,
-				"/api/console/v1/graphql",
-				cfg.HTTPTimeoutDuration(),
-			)
-
-			input := map[string]any{
-				"documentVersionId": args[0],
-			}
-
-			if cmd.Flags().Changed("title") {
-				input["title"] = flagTitle
-			}
-			if cmd.Flags().Changed("content") {
-				input["content"] = flagContent
-			}
-			if cmd.Flags().Changed("document-type") {
-				if err := cmdutil.ValidateEnum(
-					"document-type",
-					flagDocumentType,
-					[]string{"OTHER", "GOVERNANCE", "POLICY", "PROCEDURE", "PLAN", "REGISTER", "RECORD", "REPORT", "TEMPLATE"},
-				); err != nil {
-					return err
-				}
-				input["documentType"] = flagDocumentType
-			}
-			if cmd.Flags().Changed("classification") {
-				if err := cmdutil.ValidateEnum(
-					"classification",
-					flagClassification,
-					[]string{"PUBLIC", "INTERNAL", "CONFIDENTIAL", "SECRET"},
-				); err != nil {
-					return err
-				}
-				input["classification"] = flagClassification
-			}
-
-			if len(input) == 1 {
-				return fmt.Errorf("at least one field must be specified for update")
-			}
-
-			data, err := client.Do(
-				updateMutation,
-				map[string]any{"input": input},
-			)
-			if err != nil {
-				return err
-			}
-
-			var resp updateResponse
-			if err := json.Unmarshal(data, &resp); err != nil {
-				return fmt.Errorf("cannot parse response: %w", err)
-			}
-
-			v := resp.UpdateDocumentVersion.DocumentVersion
-			_, _ = fmt.Fprintf(
-				f.IOStreams.Out,
-				"Updated version %s (%s v%d.%d)\n",
-				v.ID,
-				v.Title,
-				v.Major,
-				v.Minor,
-			)
-
-			return nil
-		},
-	}
-
-	cmd.Flags().StringVar(&flagTitle, "title", "", "Version title")
-	cmd.Flags().StringVar(&flagContent, "content", "", "Version content")
-	cmd.Flags().StringVar(&flagDocumentType, "document-type", "", "Document type: OTHER, GOVERNANCE, POLICY, PROCEDURE, PLAN, REGISTER, RECORD, REPORT, TEMPLATE")
-	cmd.Flags().StringVar(&flagClassification, "classification", "", "Classification: PUBLIC, INTERNAL, CONFIDENTIAL, SECRET")
-
-	return cmd
-}

pkg/cmd/document/update/update.go

@@ -29,13 +29,15 @@ mutation($input: UpdateDocumentInput!) {
     document {
       id
       trustCenterVisibility
-      versions(first: 1) {
+    }
-        edges {
+    documentVersion {
-          node {
+      id
-            title
+      title
-          }
+      major
-        }
+      minor
-      }
+      status
+      documentType
+      classification
     }
   }
 }
@@ -46,19 +48,27 @@ type updateResponse struct {
 		Document struct {
 			ID                    string `json:"id"`
 			TrustCenterVisibility string `json:"trustCenterVisibility"`
-			Versions              struct {
-				Edges []struct {
-					Node struct {
-						Title string `json:"title"`
-					} `json:"node"`
-				} `json:"edges"`
-			} `json:"versions"`
 		} `json:"document"`
+		DocumentVersion *struct {
+			ID             string `json:"id"`
+			Title          string `json:"title"`
+			Major          int    `json:"major"`
+			Minor          int    `json:"minor"`
+			Status         string `json:"status"`
+			DocumentType   string `json:"documentType"`
+			Classification string `json:"classification"`
+		} `json:"documentVersion"`
 	} `json:"updateDocument"`
 }
 
 func NewCmdUpdate(f *cmdutil.Factory) *cobra.Command {
-	var flagTrustCenterVisibility string
+	var (
+		flagTitle                 string
+		flagContent               string
+		flagDocumentType          string
+		flagClassification        string
+		flagTrustCenterVisibility string
+	)
 
 	cmd := &cobra.Command{
 		Use:   "update <id>",
@@ -86,6 +96,32 @@ func NewCmdUpdate(f *cmdutil.Factory) *cobra.Command {
 				"id": args[0],
 			}
 
+			if cmd.Flags().Changed("title") {
+				input["title"] = flagTitle
+			}
+			if cmd.Flags().Changed("content") {
+				input["content"] = flagContent
+			}
+			if cmd.Flags().Changed("document-type") {
+				if err := cmdutil.ValidateEnum(
+					"document-type",
+					flagDocumentType,
+					[]string{"OTHER", "GOVERNANCE", "POLICY", "PROCEDURE", "PLAN", "REGISTER", "RECORD", "REPORT", "TEMPLATE"},
+				); err != nil {
+					return err
+				}
+				input["documentType"] = flagDocumentType
+			}
+			if cmd.Flags().Changed("classification") {
+				if err := cmdutil.ValidateEnum(
+					"classification",
+					flagClassification,
+					[]string{"PUBLIC", "INTERNAL", "CONFIDENTIAL", "SECRET"},
+				); err != nil {
+					return err
+				}
+				input["classification"] = flagClassification
+			}
 			if cmd.Flags().Changed("trust-center-visibility") {
 				if err := cmdutil.ValidateEnum(
 					"trust-center-visibility",
@@ -115,21 +151,31 @@ func NewCmdUpdate(f *cmdutil.Factory) *cobra.Command {
 			}
 
 			doc := resp.UpdateDocument.Document
-			title := doc.ID
+			if v := resp.UpdateDocument.DocumentVersion; v != nil {
-			if len(doc.Versions.Edges) > 0 {
+				_, _ = fmt.Fprintf(
-				title = doc.Versions.Edges[0].Node.Title
+					f.IOStreams.Out,
+					"Updated document %s (%s v%d.%d)\n",
+					doc.ID,
+					v.Title,
+					v.Major,
+					v.Minor,
+				)
+			} else {
+				_, _ = fmt.Fprintf(
+					f.IOStreams.Out,
+					"Updated document %s\n",
+					doc.ID,
+				)
 			}
-			_, _ = fmt.Fprintf(
-				f.IOStreams.Out,
-				"Updated document %s (%s)\n",
-				doc.ID,
-				title,
-			)
 
 			return nil
 		},
 	}
 
+	cmd.Flags().StringVar(&flagTitle, "title", "", "Document title")
+	cmd.Flags().StringVar(&flagContent, "content", "", "Document content")
+	cmd.Flags().StringVar(&flagDocumentType, "document-type", "", "Document type: OTHER, GOVERNANCE, POLICY, PROCEDURE, PLAN, REGISTER, RECORD, REPORT, TEMPLATE")
+	cmd.Flags().StringVar(&flagClassification, "classification", "", "Classification: PUBLIC, INTERNAL, CONFIDENTIAL, SECRET")
 	cmd.Flags().StringVar(&flagTrustCenterVisibility, "trust-center-visibility", "", "Trust center visibility: NONE, PRIVATE, PUBLIC")
 
 	return cmd

pkg/probo/actions.go

@@ -186,7 +186,7 @@ const (
 	ActionDocumentChangelogGenerate        = "core:document:generate-changelog"
 	ActionDocumentArchive                  = "core:document:archive"
 	ActionDocumentUnarchive                = "core:document:unarchive"
-	ActionDocumentDraftVersionCreate       = "core:document:create-draft-version"
+	ActionDocumentDeleteDraft              = "core:document:delete-draft"
 	ActionDocumentSendSigningNotifications = "core:document:send-signing-notifications"
 
 	// DocumentVersion actions
@@ -194,8 +194,6 @@ const (
 	ActionDocumentVersionList            = "core:document-version:list"
 	ActionDocumentVersionExportPDF       = "core:document-version:export-pdf"
 	ActionDocumentVersionSign            = "core:document-version:sign"
-	ActionDocumentVersionUpdate          = "core:document-version:update"
-	ActionDocumentVersionDeleteDraft     = "core:document-version:delete-draft"
 	ActionDocumentVersionRequestApproval = "core:document-version:request-approval"
 	ActionDocumentVersionVoidApproval    = "core:document-version:void-approval"
 	ActionDocumentVersionApprove         = "core:document-version:approve"

pkg/probo/document_service.go

@@ -71,6 +71,9 @@ type (
 	ErrDocumentArchived struct {
 	}
 
+	ErrDocumentDraftNotDeletable struct {
+	}
+
 	ErrDocumentNotArchived struct {
 	}
 
@@ -89,18 +92,14 @@ type (
 
 	UpdateDocumentRequest struct {
 		DocumentID            gid.GID
+		Title                 *string
+		Content               *string
+		Classification        *coredata.DocumentClassification
+		DocumentType          *coredata.DocumentType
 		TrustCenterVisibility *coredata.TrustCenterVisibility
 		DefaultApproverIDs    *[]gid.GID
 	}
 
-	UpdateDocumentVersionRequest struct {
-		ID             gid.GID
-		Title          *string
-		Content        *string
-		Classification *coredata.DocumentClassification
-		DocumentType   *coredata.DocumentType
-	}
-
 	RequestSignatureRequest struct {
 		DocumentVersionID gid.GID
 		Signatory         gid.GID
@@ -158,24 +157,16 @@ func (udr *UpdateDocumentRequest) Validate() error {
 			v.Check(item, "default_approver_ids", validator.GID(coredata.MembershipProfileEntityType))
 		})
 	}
-
+	v.Check(udr.Title, "title", validator.SafeTextNoNewLine(TitleMaxLength))
-	return v.Error()
+	v.Check(udr.Classification, "classification", validator.OneOfSlice(coredata.DocumentClassifications()))
-}
-
-func (udvr *UpdateDocumentVersionRequest) Validate() error {
-	v := validator.New()
-
-	v.Check(udvr.ID, "id", validator.Required(), validator.GID(coredata.DocumentVersionEntityType))
-	v.Check(udvr.Title, "title", validator.SafeTextNoNewLine(TitleMaxLength))
-	v.Check(udvr.Classification, "classification", validator.OneOfSlice(coredata.DocumentClassifications()))
 	v.Check(
-		udvr.Content,
+		udr.Content,
 		"content",
 		validator.MaxLen(documentContentMaxJSONBytes),
 		validator.ProseMirrorDocumentContent(),
 		validator.ProseMirrorDocumentMaxTextLength(documentContentMaxTextLength),
 	)
-	v.Check(udvr.DocumentType, "document_type", validator.OneOfSlice(coredata.DocumentTypes()))
+	v.Check(udr.DocumentType, "document_type", validator.OneOfSlice(coredata.DocumentTypes()))
 
 	return v.Error()
 }
@@ -214,6 +205,10 @@ func (e ErrDocumentArchived) Error() string {
 	return "cannot modify an archived document"
 }
 
+func (e ErrDocumentDraftNotDeletable) Error() string {
+	return "latest version is not a deletable draft"
+}
+
 func (e ErrDocumentNotArchived) Error() string {
 	return "cannot unarchive a document that is not archived"
 }
@@ -824,68 +819,39 @@ func (s *DocumentService) signDocumentVersionInTx(
 	return documentVersionSignature, nil
 }
 
-func (s *DocumentService) UpdateVersion(
+func (s *DocumentService) updateVersionInTx(
 	ctx context.Context,
-	req UpdateDocumentVersionRequest,
+	tx pg.Tx,
-) (*coredata.DocumentVersion, error) {
+	draftVersion *coredata.DocumentVersion,
-	documentVersion := &coredata.DocumentVersion{}
+	content *string,
-	document := &coredata.Document{}
+	classification *coredata.DocumentClassification,
-
+	documentType *coredata.DocumentType,
-	if err := req.Validate(); err != nil {
+	title *string,
-		return nil, err
+) error {
+	if content != nil {
+		sanitized, err := prosemirror.SanitizeDocumentJSON(*content)
+		if err != nil {
+			return fmt.Errorf("cannot sanitize document content: %w", err)
+		}
+		draftVersion.Content = sanitized
 	}
 
-	err := s.svc.pg.WithTx(
+	if title != nil {
-		ctx,
+		draftVersion.Title = *title
-		func(ctx context.Context, conn pg.Tx) error {
+	}
-			if err := documentVersion.LoadByID(ctx, conn, s.svc.scope, req.ID); err != nil {
+	if classification != nil {
-				return fmt.Errorf("cannot load document version %q: %w", req.ID, err)
+		draftVersion.Classification = *classification
-			}
+	}
+	if documentType != nil {
+		draftVersion.DocumentType = *documentType
+	}
+	draftVersion.UpdatedAt = time.Now()
 
-			if err := document.LoadByID(ctx, conn, s.svc.scope, documentVersion.DocumentID); err != nil {
+	if err := draftVersion.Update(ctx, tx, s.svc.scope); err != nil {
-				return fmt.Errorf("cannot load document %q: %w", documentVersion.DocumentID, err)
+		return fmt.Errorf("cannot update document version: %w", err)
-			}
-
-			if document.ArchivedAt != nil {
-				return &ErrDocumentArchived{}
-			}
-
-			if documentVersion.Status != coredata.DocumentVersionStatusDraft {
-				return &ErrDocumentVersionNotDraft{}
-			}
-
-			if req.Content != nil {
-				content, err := prosemirror.SanitizeDocumentJSON(*req.Content)
-				if err != nil {
-					return fmt.Errorf("cannot sanitize document content: %w", err)
-				}
-				documentVersion.Content = content
-			}
-
-			if req.Title != nil {
-				documentVersion.Title = *req.Title
-			}
-			if req.Classification != nil {
-				documentVersion.Classification = *req.Classification
-			}
-			if req.DocumentType != nil {
-				documentVersion.DocumentType = *req.DocumentType
-			}
-			documentVersion.UpdatedAt = time.Now()
-
-			if err := documentVersion.Update(ctx, conn, s.svc.scope); err != nil {
-				return fmt.Errorf("cannot update document version: %w", err)
-			}
-
-			return nil
-		},
-	)
-
-	if err != nil {
-		return nil, err
 	}
 
-	return documentVersion, nil
+	return nil
 }
 
 func (s *DocumentService) GetVersionSignature(
@@ -1087,101 +1053,46 @@ func (s *DocumentService) IsVersionSignedByUserEmail(
 	return signed, nil
 }
 
-func (s *DocumentService) CreateDraft(
+func (s *DocumentService) createDraftInTx(
 	ctx context.Context,
-	documentID gid.GID,
+	tx pg.Tx,
+	document *coredata.Document,
+	latestVersion *coredata.DocumentVersion,
 ) (*coredata.DocumentVersion, error) {
-	draftVersionID := gid.New(s.svc.scope.GetTenantID(), coredata.DocumentVersionEntityType)
-
-	latestVersion := &coredata.DocumentVersion{}
-	document := &coredata.Document{}
-	draftVersion := &coredata.DocumentVersion{}
 	now := time.Now()
 
-	err := s.svc.pg.WithTx(
+	draftVersion := &coredata.DocumentVersion{
-		ctx,
+		ID:             gid.New(s.svc.scope.GetTenantID(), coredata.DocumentVersionEntityType),
-		func(ctx context.Context, conn pg.Tx) error {
+		OrganizationID: document.OrganizationID,
-			if err := document.LoadByID(ctx, conn, s.svc.scope, documentID); err != nil {
+		DocumentID:     document.ID,
-				return fmt.Errorf("cannot load document: %w", err)
+		Title:          latestVersion.Title,
-			}
+		Major:          latestVersion.Major,
+		Minor:          latestVersion.Minor + 1,
+		Classification: latestVersion.Classification,
+		DocumentType:   latestVersion.DocumentType,
+		Content:        latestVersion.Content,
+		Status:         coredata.DocumentVersionStatusDraft,
+		CreatedAt:      now,
+		UpdatedAt:      now,
+	}
 
-			if document.ArchivedAt != nil {
+	if err := draftVersion.Insert(ctx, tx, s.svc.scope); err != nil {
-				return &ErrDocumentArchived{}
+		return nil, fmt.Errorf("cannot create draft: %w", err)
-			}
-
-			if err := latestVersion.LoadLatestVersion(ctx, conn, s.svc.scope, documentID); err != nil {
-				return fmt.Errorf("cannot load latest version: %w", err)
-			}
-
-			if latestVersion.Status != coredata.DocumentVersionStatusPublished {
-				return &ErrDocumentVersionNotPublished{}
-			}
-
-			draftVersion.ID = draftVersionID
-			draftVersion.OrganizationID = document.OrganizationID
-			draftVersion.DocumentID = documentID
-			draftVersion.Title = latestVersion.Title
-			draftVersion.Major = latestVersion.Major
-			draftVersion.Minor = latestVersion.Minor + 1
-			draftVersion.Classification = latestVersion.Classification
-			draftVersion.DocumentType = latestVersion.DocumentType
-			draftVersion.Content = latestVersion.Content
-			draftVersion.Status = coredata.DocumentVersionStatusDraft
-			draftVersion.CreatedAt = now
-			draftVersion.UpdatedAt = now
-
-			if err := draftVersion.Insert(ctx, conn, s.svc.scope); err != nil {
-				return fmt.Errorf("cannot create draft: %w", err)
-			}
-
-			return nil
-		},
-	)
-
-	if err != nil {
-		return nil, err
 	}
 
 	return draftVersion, nil
 }
 
-func (s *DocumentService) DeleteDraft(
+func (s *DocumentService) deleteDraftInTx(
 	ctx context.Context,
-	documentVersionID gid.GID,
+	tx pg.Tx,
+	draftVersion *coredata.DocumentVersion,
 ) error {
-	documentVersion := &coredata.DocumentVersion{}
+	if err := draftVersion.Delete(ctx, tx, s.svc.scope); err != nil {
+		return fmt.Errorf("cannot delete document version: %w", err)
+	}
 
-	return s.svc.pg.WithTx(
+	return nil
-		ctx,
-		func(ctx context.Context, conn pg.Tx) error {
-			if err := documentVersion.LoadByID(ctx, conn, s.svc.scope, documentVersionID); err != nil {
-				return fmt.Errorf("cannot load document version: %w", err)
-			}
-
-			document := &coredata.Document{}
-			if err := document.LoadByID(ctx, conn, s.svc.scope, documentVersion.DocumentID); err != nil {
-				return fmt.Errorf("cannot load document: %w", err)
-			}
-
-			if document.ArchivedAt != nil {
-				return &ErrDocumentArchived{}
-			}
-
-			if documentVersion.Status != coredata.DocumentVersionStatusDraft {
-				return fmt.Errorf("cannot delete published document version")
-			}
-
-			if documentVersion.Major == 0 && documentVersion.Minor == 1 {
-				return fmt.Errorf("cannot delete the first version of a document")
-			}
-
-			if err := documentVersion.Delete(ctx, conn, s.svc.scope); err != nil {
-				return fmt.Errorf("cannot delete document version: %w", err)
-			}
-
-			return nil
-		},
-	)
 }
 
 func (s *DocumentService) SoftDelete(
@@ -1698,12 +1609,14 @@ func (s *DocumentService) ListForMeasureID(
 func (s *DocumentService) Update(
 	ctx context.Context,
 	req UpdateDocumentRequest,
-) (*coredata.Document, error) {
+) (*coredata.Document, *coredata.DocumentVersion, bool, error) {
 	if err := req.Validate(); err != nil {
-		return nil, err
+		return nil, nil, false, err
 	}
 
 	document := &coredata.Document{}
+	var resultVersion *coredata.DocumentVersion
+	var draftCreated bool
 	now := time.Now()
 
 	err := s.svc.pg.WithTx(
@@ -1727,6 +1640,73 @@ func (s *DocumentService) Update(
 				return fmt.Errorf("cannot update document: %w", err)
 			}
 
+			// Handle draft version logic for title/content/classification/type changes.
+			latestVersion := &coredata.DocumentVersion{}
+			if err := latestVersion.LoadLatestVersion(ctx, tx, s.svc.scope, req.DocumentID); err != nil {
+				return fmt.Errorf("cannot load latest version: %w", err)
+			}
+
+			hasVersionChanges := req.Title != nil || req.Content != nil || req.Classification != nil || req.DocumentType != nil
+
+			if !hasVersionChanges {
+				if req.DefaultApproverIDs != nil {
+					defaultApprovers := &coredata.DocumentDefaultApprovers{}
+					if err := defaultApprovers.MergeByDocumentID(ctx, tx, s.svc.scope, req.DocumentID, document.OrganizationID, *req.DefaultApproverIDs); err != nil {
+						return fmt.Errorf("cannot update default approvers: %w", err)
+					}
+				}
+				return nil
+			}
+
+			if latestVersion.Status == coredata.DocumentVersionStatusDraft {
+				// Draft exists: update it with any new values.
+				if err := s.updateVersionInTx(ctx, tx, latestVersion, req.Content, req.Classification, req.DocumentType, req.Title); err != nil {
+					return err
+				}
+
+				// If there is a published version and the draft matches it, delete the draft.
+				// Never delete the initial draft (v0.1) since there's nothing to fall back to.
+				if document.CurrentPublishedMajor != nil && (latestVersion.Major != 0 || latestVersion.Minor != 1) {
+					publishedVersion := &coredata.DocumentVersion{}
+					if err := publishedVersion.LoadByDocumentIDAndVersion(
+						ctx,
+						tx,
+						s.svc.scope,
+						req.DocumentID,
+						*document.CurrentPublishedMajor,
+						*document.CurrentPublishedMinor,
+					); err != nil {
+						return fmt.Errorf("cannot load published version: %w", err)
+					}
+
+					if latestVersion.Title == publishedVersion.Title &&
+						latestVersion.Content == publishedVersion.Content &&
+						latestVersion.Classification == publishedVersion.Classification &&
+						latestVersion.DocumentType == publishedVersion.DocumentType {
+						if err := s.deleteDraftInTx(ctx, tx, latestVersion); err != nil {
+							return err
+						}
+						resultVersion = nil
+						return nil
+					}
+				}
+
+				resultVersion = latestVersion
+			} else {
+				// No draft exists: create one.
+				draftVersion, err := s.createDraftInTx(ctx, tx, document, latestVersion)
+				if err != nil {
+					return err
+				}
+
+				if err := s.updateVersionInTx(ctx, tx, draftVersion, req.Content, req.Classification, req.DocumentType, req.Title); err != nil {
+					return err
+				}
+
+				resultVersion = draftVersion
+				draftCreated = true
+			}
+
 			if req.DefaultApproverIDs != nil {
 				defaultApprovers := &coredata.DocumentDefaultApprovers{}
 				if err := defaultApprovers.MergeByDocumentID(ctx, tx, s.svc.scope, req.DocumentID, document.OrganizationID, *req.DefaultApproverIDs); err != nil {
@@ -1738,6 +1718,47 @@ func (s *DocumentService) Update(
 		},
 	)
 
+	if err != nil {
+		return nil, nil, false, err
+	}
+
+	return document, resultVersion, draftCreated, nil
+}
+
+func (s *DocumentService) DeleteDraft(
+	ctx context.Context,
+	documentID gid.GID,
+) (*coredata.Document, error) {
+	document := &coredata.Document{}
+
+	err := s.svc.pg.WithTx(
+		ctx,
+		func(ctx context.Context, tx pg.Tx) error {
+			if err := document.LoadByID(ctx, tx, s.svc.scope, documentID); err != nil {
+				return fmt.Errorf("cannot load document %q: %w", documentID, err)
+			}
+
+			if document.ArchivedAt != nil {
+				return &ErrDocumentArchived{}
+			}
+
+			latestVersion := &coredata.DocumentVersion{}
+			if err := latestVersion.LoadLatestVersion(ctx, tx, s.svc.scope, documentID); err != nil {
+				return fmt.Errorf("cannot load latest version: %w", err)
+			}
+
+			if latestVersion.Status != coredata.DocumentVersionStatusDraft {
+				return &ErrDocumentDraftNotDeletable{}
+			}
+
+			if latestVersion.Major == 0 && latestVersion.Minor == 1 {
+				return &ErrDocumentDraftNotDeletable{}
+			}
+
+			return s.deleteDraftInTx(ctx, tx, latestVersion)
+		},
+	)
+
 	if err != nil {
 		return nil, err
 	}

pkg/server/api/console/v1/schema.graphql

@@ -3892,6 +3892,7 @@ type Mutation {
     # Document mutations
     createDocument(input: CreateDocumentInput!): CreateDocumentPayload!
     updateDocument(input: UpdateDocumentInput!): UpdateDocumentPayload!
+    deleteDocumentDraft(input: DeleteDocumentDraftInput!): DeleteDocumentDraftPayload!
     archiveDocument(input: ArchiveDocumentInput!): ArchiveDocumentPayload!
     unarchiveDocument(input: UnarchiveDocumentInput!): UnarchiveDocumentPayload!
     deleteDocument(input: DeleteDocumentInput!): DeleteDocumentPayload!
@@ -3955,15 +3956,6 @@ type Mutation {
     generateDocumentChangelog(
         input: GenerateDocumentChangelogInput!
     ): GenerateDocumentChangelogPayload!
-    createDraftDocumentVersion(
-        input: CreateDraftDocumentVersionInput!
-    ): CreateDraftDocumentVersionPayload!
-    deleteDraftDocumentVersion(
-        input: DeleteDraftDocumentVersionInput!
-    ): DeleteDraftDocumentVersionPayload!
-    updateDocumentVersion(
-        input: UpdateDocumentVersionInput!
-    ): UpdateDocumentVersionPayload!
     requestSignature(input: RequestSignatureInput!): RequestSignaturePayload!
     bulkRequestSignatures(
         input: BulkRequestSignaturesInput!
@@ -4673,6 +4665,10 @@ input CreateDocumentInput {
 
 input UpdateDocumentInput {
     id: ID!
+    title: String
+    content: String
+    classification: DocumentClassification
+    documentType: DocumentType
     trustCenterVisibility: TrustCenterVisibility
     defaultApproverIds: [ID!]
 }
@@ -4703,6 +4699,10 @@ input ExportTransferImpactAssessmentsPDFInput {
     filter: TransferImpactAssessmentFilter
 }
 
+input DeleteDocumentDraftInput {
+    documentId: ID!
+}
+
 input ArchiveDocumentInput {
     documentId: ID!
 }
@@ -5442,6 +5442,12 @@ type ExportTransferImpactAssessmentsPDFPayload {
 
 type UpdateDocumentPayload {
     document: Document!
+    documentVersion: DocumentVersion
+    documentVersionEdge: DocumentVersionEdge
+}
+
+type DeleteDocumentDraftPayload {
+    document: Document!
 }
 
 type ArchiveDocumentPayload {
@@ -5922,38 +5928,10 @@ type BulkPublishDocumentVersionsPayload {
     documents: [Document!]!
 }
 
-type CreateDraftDocumentVersionPayload {
-    documentVersionEdge: DocumentVersionEdge!
-}
-
-type DeleteDraftDocumentVersionPayload {
-    deletedDocumentVersionId: ID!
-}
-
-input CreateDraftDocumentVersionInput {
-    documentID: ID!
-}
-
-input DeleteDraftDocumentVersionInput {
-    documentVersionId: ID!
-}
-
-input UpdateDocumentVersionInput {
-    documentVersionId: ID!
-    title: String
-    content: String
-    classification: DocumentClassification
-    documentType: DocumentType
-}
-
 input CancelSignatureRequestInput {
     documentVersionSignatureId: ID!
 }
 
-type UpdateDocumentVersionPayload {
-    documentVersion: DocumentVersion!
-}
-
 input SendSigningNotificationsInput {
     organizationId: ID!
 }

pkg/server/api/console/v1/v1_resolver.go

@@ -5303,16 +5303,23 @@ func (r *mutationResolver) UpdateDocument(ctx context.Context, input types.Updat
 		defaultApproverIDs = &input.DefaultApproverIds
 	}
 
-	document, err := prb.Documents.Update(
+	document, documentVersion, draftCreated, err := prb.Documents.Update(
 		ctx,
 		probo.UpdateDocumentRequest{
 			DocumentID:            input.ID,
+			Title:                 input.Title,
+			Content:               input.Content,
+			Classification:        input.Classification,
+			DocumentType:          input.DocumentType,
 			TrustCenterVisibility: input.TrustCenterVisibility,
 			DefaultApproverIDs:    defaultApproverIDs,
 		},
 	)
 
 	if err != nil {
+		if errors.Is(err, coredata.ErrResourceNotFound) {
+			return nil, gqlutils.NotFound(ctx, err)
+		}
 		if errArchived, ok := errors.AsType[*probo.ErrDocumentArchived](err); ok {
 			return nil, gqlutils.Conflict(ctx, errArchived)
 		}
@@ -5323,7 +5330,48 @@ func (r *mutationResolver) UpdateDocument(ctx context.Context, input types.Updat
 		return nil, gqlutils.Internal(ctx)
 	}
 
-	return &types.UpdateDocumentPayload{
+	payload := &types.UpdateDocumentPayload{
+		Document: types.NewDocument(document),
+	}
+
+	if documentVersion != nil {
+		payload.DocumentVersion = types.NewDocumentVersion(documentVersion)
+	}
+
+	if draftCreated {
+		payload.DocumentVersionEdge = types.NewDocumentVersionEdge(
+			documentVersion,
+			coredata.DocumentVersionOrderFieldCreatedAt,
+		)
+	}
+
+	return payload, nil
+}
+
+// DeleteDocumentDraft is the resolver for the deleteDocumentDraft field.
+func (r *mutationResolver) DeleteDocumentDraft(ctx context.Context, input types.DeleteDocumentDraftInput) (*types.DeleteDocumentDraftPayload, error) {
+	if err := r.authorize(ctx, input.DocumentID, probo.ActionDocumentDeleteDraft); err != nil {
+		return nil, err
+	}
+
+	prb := r.ProboService(ctx, input.DocumentID.TenantID())
+
+	document, err := prb.Documents.DeleteDraft(ctx, input.DocumentID)
+	if err != nil {
+		if errors.Is(err, coredata.ErrResourceNotFound) {
+			return nil, gqlutils.NotFound(ctx, err)
+		}
+		if errNotDeletable, ok := errors.AsType[*probo.ErrDocumentDraftNotDeletable](err); ok {
+			return nil, gqlutils.Conflict(ctx, errNotDeletable)
+		}
+		if errArchived, ok := errors.AsType[*probo.ErrDocumentArchived](err); ok {
+			return nil, gqlutils.Conflict(ctx, errArchived)
+		}
+		r.logger.ErrorCtx(ctx, "cannot delete document draft", log.Error(err))
+		return nil, gqlutils.Internal(ctx)
+	}
+
+	return &types.DeleteDocumentDraftPayload{
 		Document: types.NewDocument(document),
 	}, nil
 }
@@ -6038,98 +6086,6 @@ func (r *mutationResolver) GenerateDocumentChangelog(ctx context.Context, input
 	}, nil
 }
 
-// CreateDraftDocumentVersion is the resolver for the createDraftDocumentVersion field.
-func (r *mutationResolver) CreateDraftDocumentVersion(ctx context.Context, input types.CreateDraftDocumentVersionInput) (*types.CreateDraftDocumentVersionPayload, error) {
-	if err := r.authorize(ctx, input.DocumentID, probo.ActionDocumentDraftVersionCreate); err != nil {
-		return nil, err
-	}
-
-	prb := r.ProboService(ctx, input.DocumentID.TenantID())
-
-	documentVersion, err := prb.Documents.CreateDraft(ctx, input.DocumentID)
-	if err != nil {
-		if errArchived, ok := errors.AsType[*probo.ErrDocumentArchived](err); ok {
-			return nil, gqlutils.Conflict(ctx, errArchived)
-		}
-
-		if errNotPublished, ok := errors.AsType[*probo.ErrDocumentVersionNotPublished](err); ok {
-			return nil, gqlutils.Conflict(ctx, errNotPublished)
-		}
-
-		r.logger.ErrorCtx(ctx, "cannot create draft document version", log.Error(err))
-		return nil, gqlutils.Internal(ctx)
-	}
-
-	return &types.CreateDraftDocumentVersionPayload{
-		DocumentVersionEdge: types.NewDocumentVersionEdge(documentVersion, coredata.DocumentVersionOrderFieldCreatedAt),
-	}, nil
-}
-
-// DeleteDraftDocumentVersion is the resolver for the deleteDraftDocumentVersion field.
-func (r *mutationResolver) DeleteDraftDocumentVersion(ctx context.Context, input types.DeleteDraftDocumentVersionInput) (*types.DeleteDraftDocumentVersionPayload, error) {
-	if err := r.authorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionDeleteDraft); err != nil {
-		return nil, err
-	}
-
-	prb := r.ProboService(ctx, input.DocumentVersionID.TenantID())
-
-	err := prb.Documents.DeleteDraft(ctx, input.DocumentVersionID)
-	if err != nil {
-		if errArchived, ok := errors.AsType[*probo.ErrDocumentArchived](err); ok {
-			return nil, gqlutils.Conflict(ctx, errArchived)
-		}
-
-		r.logger.ErrorCtx(ctx, "cannot delete draft document version", log.Error(err))
-		return nil, gqlutils.Internal(ctx)
-	}
-
-	return &types.DeleteDraftDocumentVersionPayload{
-		DeletedDocumentVersionID: input.DocumentVersionID,
-	}, nil
-}
-
-// UpdateDocumentVersion is the resolver for the updateDocumentVersion field.
-func (r *mutationResolver) UpdateDocumentVersion(ctx context.Context, input types.UpdateDocumentVersionInput) (*types.UpdateDocumentVersionPayload, error) {
-	if err := r.authorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionUpdate); err != nil {
-		return nil, err
-	}
-
-	prb := r.ProboService(ctx, input.DocumentVersionID.TenantID())
-
-	documentVersion, err := prb.Documents.UpdateVersion(
-		ctx,
-		probo.UpdateDocumentVersionRequest{
-			ID:             input.DocumentVersionID,
-			Title:          input.Title,
-			Content:        input.Content,
-			Classification: input.Classification,
-			DocumentType:   input.DocumentType,
-		},
-	)
-	if err != nil {
-		if errors.Is(err, coredata.ErrResourceNotFound) {
-			return nil, gqlutils.NotFound(ctx, err)
-		}
-
-		if errNotDraft, ok := errors.AsType[*probo.ErrDocumentVersionNotDraft](err); ok {
-			return nil, gqlutils.Conflict(ctx, errNotDraft)
-		}
-
-		if errArchived, ok := errors.AsType[*probo.ErrDocumentArchived](err); ok {
-			return nil, gqlutils.Conflict(ctx, errArchived)
-		}
-		if validationErrors, ok := errors.AsType[validator.ValidationErrors](err); ok {
-			return nil, gqlutils.InvalidValidationErrors(ctx, validationErrors)
-		}
-		r.logger.ErrorCtx(ctx, "cannot update document version", log.Error(err))
-		return nil, gqlutils.Internal(ctx)
-	}
-
-	return &types.UpdateDocumentVersionPayload{
-		DocumentVersion: types.NewDocumentVersion(documentVersion),
-	}, nil
-}
-
 // RequestSignature is the resolver for the requestSignature field.
 func (r *mutationResolver) RequestSignature(ctx context.Context, input types.RequestSignatureInput) (*types.RequestSignaturePayload, error) {
 	if err := r.authorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionSignatureRequest); err != nil {

pkg/server/api/mcp/v1/schema.resolvers.go

@@ -2115,10 +2115,23 @@ func (r *Resolver) UpdateDocumentTool(ctx context.Context, req *mcp.CallToolRequ
 		defaultApproverIDs = &input.DefaultApproverIds
 	}
 
-	document, err := svc.Documents.Update(
+	var content *string
+	if input.Content != nil {
+		c, err := markdownToProseMirrorJSON(*input.Content)
+		if err != nil {
+			panic(fmt.Errorf("cannot convert markdown to prosemirror: %w", err))
+		}
+		content = &c
+	}
+
+	document, documentVersion, _, err := svc.Documents.Update(
 		ctx,
 		probo.UpdateDocumentRequest{
 			DocumentID:            input.ID,
+			Title:                 input.Title,
+			Content:               content,
+			Classification:        input.Classification,
+			DocumentType:          input.DocumentType,
 			TrustCenterVisibility: input.TrustCenterVisibility,
 			DefaultApproverIDs:    defaultApproverIDs,
 		},
@@ -2127,9 +2140,15 @@ func (r *Resolver) UpdateDocumentTool(ctx context.Context, req *mcp.CallToolRequ
 		panic(fmt.Errorf("cannot update document: %w", err))
 	}
 
-	return nil, types.UpdateDocumentOutput{
+	output := types.UpdateDocumentOutput{
 		Document: types.NewDocument(document),
-	}, nil
+	}
+
+	if documentVersion != nil {
+		output.DocumentVersion = types.NewDocumentVersion(documentVersion)
+	}
+
+	return nil, output, nil
 }
 
 func (r *Resolver) ListDocumentVersionsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDocumentVersionsInput) (*mcp.CallToolResult, types.ListDocumentVersionsOutput, error) {
@@ -2172,72 +2191,6 @@ func (r *Resolver) GetDocumentVersionTool(ctx context.Context, req *mcp.CallTool
 	}, nil
 }
 
-func (r *Resolver) CreateDraftDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.CreateDraftDocumentVersionInput) (*mcp.CallToolResult, types.CreateDraftDocumentVersionOutput, error) {
-	r.MustAuthorize(ctx, input.DocumentID, probo.ActionDocumentDraftVersionCreate)
-
-	svc := r.ProboService(ctx, input.DocumentID)
-
-	draftVersion, err := svc.Documents.CreateDraft(ctx, input.DocumentID)
-	if err != nil {
-		panic(fmt.Errorf("cannot create draft document version: %w", err))
-	}
-
-	if input.Content != nil {
-		content, err := markdownToProseMirrorJSON(*input.Content)
-		if err != nil {
-			panic(fmt.Errorf("cannot convert markdown to prosemirror: %w", err))
-		}
-
-		draftVersion, err = svc.Documents.UpdateVersion(
-			ctx,
-			probo.UpdateDocumentVersionRequest{
-				ID:      draftVersion.ID,
-				Content: &content,
-			},
-		)
-		if err != nil {
-			panic(fmt.Errorf("cannot update draft document version content: %w", err))
-		}
-	}
-
-	return nil, types.CreateDraftDocumentVersionOutput{
-		DocumentVersion: types.NewDocumentVersion(draftVersion),
-	}, nil
-}
-
-func (r *Resolver) UpdateDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateDocumentVersionInput) (*mcp.CallToolResult, types.UpdateDocumentVersionOutput, error) {
-	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionUpdate)
-
-	svc := r.ProboService(ctx, input.DocumentVersionID)
-
-	var content *string
-	if input.Content != nil {
-		c, err := markdownToProseMirrorJSON(*input.Content)
-		if err != nil {
-			panic(fmt.Errorf("cannot convert markdown to prosemirror: %w", err))
-		}
-		content = &c
-	}
-
-	documentVersion, err := svc.Documents.UpdateVersion(
-		ctx,
-		probo.UpdateDocumentVersionRequest{
-			ID:             input.DocumentVersionID,
-			Title:          input.Title,
-			Content:        content,
-			Classification: input.Classification,
-			DocumentType:   input.DocumentType,
-		},
-	)
-	if err != nil {
-		panic(fmt.Errorf("cannot update document version: %w", err))
-	}
-
-	return nil, types.UpdateDocumentVersionOutput{
-		DocumentVersion: types.NewDocumentVersion(documentVersion),
-	}, nil
-}
-
 func (r *Resolver) ListDocumentVersionSignaturesTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDocumentVersionSignaturesInput) (*mcp.CallToolResult, types.ListDocumentVersionSignaturesOutput, error) {
 	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionSignatureList)
 
@@ -2312,21 +2265,6 @@ func (r *Resolver) RequestDocumentVersionSignatureTool(ctx context.Context, req
 	}, nil
 }
 
-func (r *Resolver) DeleteDraftDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteDraftDocumentVersionInput) (*mcp.CallToolResult, types.DeleteDraftDocumentVersionOutput, error) {
-	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionDeleteDraft)
-
-	svc := r.ProboService(ctx, input.DocumentVersionID)
-
-	err := svc.Documents.DeleteDraft(ctx, input.DocumentVersionID)
-	if err != nil {
-		panic(fmt.Errorf("cannot delete draft document version: %w", err))
-	}
-
-	return nil, types.DeleteDraftDocumentVersionOutput{
-		DeletedDocumentVersionID: input.DocumentVersionID,
-	}, nil
-}
-
 func (r *Resolver) DeleteDocumentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteDocumentInput) (*mcp.CallToolResult, types.DeleteDocumentOutput, error) {
 	r.MustAuthorize(ctx, input.DocumentID, probo.ActionDocumentDelete)
 
@@ -3971,3 +3909,18 @@ func (r *Resolver) SendSigningNotificationsTool(ctx context.Context, req *mcp.Ca
 		Success: true,
 	}, nil
 }
+
+func (r *Resolver) DeleteDocumentDraftTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteDocumentDraftInput) (*mcp.CallToolResult, types.DeleteDocumentDraftOutput, error) {
+	r.MustAuthorize(ctx, input.ID, probo.ActionDocumentDeleteDraft)
+
+	svc := r.ProboService(ctx, input.ID)
+
+	document, err := svc.Documents.DeleteDraft(ctx, input.ID)
+	if err != nil {
+		return nil, types.DeleteDocumentDraftOutput{}, fmt.Errorf("cannot delete document draft: %w", err)
+	}
+
+	return nil, types.DeleteDocumentDraftOutput{
+		Document: types.NewDocument(document),
+	}, nil
+}

pkg/server/api/mcp/v1/specification.yaml

@@ -5517,6 +5517,18 @@ components:
         id:
           $ref: "#/components/schemas/GID"
           description: Document ID
+        title:
+          type: string
+          description: Document title
+        content:
+          type: string
+          description: Document content in markdown format
+        classification:
+          $ref: "#/components/schemas/DocumentClassification"
+          description: Document classification
+        document_type:
+          $ref: "#/components/schemas/DocumentType"
+          description: Document type
         trust_center_visibility:
           $ref: "#/components/schemas/TrustCenterVisibility"
           description: Trust center visibility
@@ -5527,6 +5539,25 @@ components:
           description: Default approver profile IDs
 
     UpdateDocumentOutput:
+      type: object
+      required:
+        - document
+      properties:
+        document:
+          $ref: "#/components/schemas/Document"
+        document_version:
+          $ref: "#/components/schemas/DocumentVersion"
+
+    DeleteDocumentDraftInput:
+      type: object
+      required:
+        - id
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Document ID
+
+    DeleteDocumentDraftOutput:
       type: object
       required:
         - document
@@ -5618,75 +5649,6 @@ components:
         document_version:
           $ref: "#/components/schemas/DocumentVersion"
 
-    CreateDraftDocumentVersionInput:
-      type: object
-      required:
-        - document_id
-      properties:
-        document_id:
-          $ref: "#/components/schemas/GID"
-          description: Document ID
-        content:
-          type: string
-          description: Document content in markdown format
-
-    CreateDraftDocumentVersionOutput:
-      type: object
-      description: Created draft; document_version.content is markdown
-      required:
-        - document_version
-      properties:
-        document_version:
-          $ref: "#/components/schemas/DocumentVersion"
-
-    UpdateDocumentVersionInput:
-      type: object
-      required:
-        - document_version_id
-      properties:
-        document_version_id:
-          $ref: "#/components/schemas/GID"
-          description: Document version ID
-        title:
-          type: string
-          description: Document version title
-        content:
-          type: string
-          description: Document content in markdown format
-        classification:
-          $ref: "#/components/schemas/DocumentClassification"
-          description: Document classification
-        document_type:
-          $ref: "#/components/schemas/DocumentType"
-          description: Document type
-
-    UpdateDocumentVersionOutput:
-      type: object
-      description: Updated draft; document_version.content is markdown
-      required:
-        - document_version
-      properties:
-        document_version:
-          $ref: "#/components/schemas/DocumentVersion"
-
-    DeleteDraftDocumentVersionInput:
-      type: object
-      required:
-        - document_version_id
-      properties:
-        document_version_id:
-          $ref: "#/components/schemas/GID"
-          description: Document version ID
-
-    DeleteDraftDocumentVersionOutput:
-      type: object
-      required:
-        - deleted_document_version_id
-      properties:
-        deleted_document_version_id:
-          $ref: "#/components/schemas/GID"
-          description: Deleted document version ID
-
     PublishMajorDocumentVersionInput:
       type: object
       required:
@@ -8398,6 +8360,15 @@ tools:
       $ref: "#/components/schemas/UpdateDocumentInput"
     outputSchema:
       $ref: "#/components/schemas/UpdateDocumentOutput"
+  - name: deleteDocumentDraft
+    description: Delete the latest draft version of a document, reverting to the last published version. Cannot delete the initial v0.1 draft.
+    hints:
+      readonly: false
+      destructive: true
+    inputSchema:
+      $ref: "#/components/schemas/DeleteDocumentDraftInput"
+    outputSchema:
+      $ref: "#/components/schemas/DeleteDocumentDraftOutput"
   - name: archiveDocument
     description: Archive a document to prevent further modifications
     hints:
@@ -8432,30 +8403,6 @@ tools:
       $ref: "#/components/schemas/GetDocumentVersionInput"
     outputSchema:
       $ref: "#/components/schemas/GetDocumentVersionOutput"
-  - name: createDraftDocumentVersion
-    description: Create a new draft version from the latest published version
-    hints:
-      readonly: false
-    inputSchema:
-      $ref: "#/components/schemas/CreateDraftDocumentVersionInput"
-    outputSchema:
-      $ref: "#/components/schemas/CreateDraftDocumentVersionOutput"
-  - name: updateDocumentVersion
-    description: Update an existing draft document version content
-    hints:
-      readonly: false
-    inputSchema:
-      $ref: "#/components/schemas/UpdateDocumentVersionInput"
-    outputSchema:
-      $ref: "#/components/schemas/UpdateDocumentVersionOutput"
-  - name: deleteDraftDocumentVersion
-    description: Delete a draft document version
-    hints:
-      readonly: false
-    inputSchema:
-      $ref: "#/components/schemas/DeleteDraftDocumentVersionInput"
-    outputSchema:
-      $ref: "#/components/schemas/DeleteDraftDocumentVersionOutput"
   - name: publishMajorDocumentVersion
     description: Publish a draft document version as a new major version
     hints: