Move document classification from document to document version

Classification now lives exclusively on DocumentVersion. The field is
removed from the Document model, all SQL queries, GraphQL Document
type, SignableDocument type, UpdateDocumentInput, and MCP Document
schema.

New documents still accept classification in CreateDocumentInput,
applied to the first version. New drafts inherit classification from
the previous version. PDF generation uses the version classification.

The drawer allows editing classification on draft versions via the
updateDocumentVersion mutation. Classification is read-only on
published versions.

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

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

@@ -16,19 +16,18 @@ import { documentClassifications, documentTypes, formatDate, getDocumentClassifi
 import { useTranslate } from "@probo/i18n";
 import { Badge, Button, Drawer, IconCheckmark1, IconCrossLargeX, IconPencil, PropertyRow } from "@probo/ui";
 import { useState } from "react";
-import { useFragment } from "react-relay";
+import { useFragment, useMutation } from "react-relay";
 import { graphql } from "relay-runtime";
 import { z } from "zod";
 
 import type { DocumentLayoutDrawer_documentFragment$key } from "#/__generated__/core/DocumentLayoutDrawer_documentFragment.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";
 import { DocumentClassificationOptions } from "#/components/form/DocumentClassificationOptions";
 import { DocumentTypeOptions } from "#/components/form/DocumentTypeOptions";
 import { useFormWithSchema } from "#/hooks/useFormWithSchema";
-import { useMutationWithToasts } from "#/hooks/useMutationWithToasts";
-
 const documentFragment = graphql`
   fragment DocumentLayoutDrawer_documentFragment on Document {
     id
@@ -57,6 +56,16 @@ const updateDocumentMutation = graphql`
       document {
         id
         documentType
+      }
+    }
+  }
+`;
+
+const updateClassificationMutation = graphql`
+  mutation DocumentLayoutDrawer_updateClassificationMutation($input: UpdateDocumentVersionInput!) {
+    updateDocumentVersion(input: $input) {
+      documentVersion {
+        id
         classification
       }
     }
@@ -65,6 +74,9 @@ const updateDocumentMutation = graphql`
 
 const schema = z.object({
   documentType: z.enum(documentTypes),
+});
+
+const classificationSchema = z.object({
   classification: z.enum(documentClassifications),
 });
 
@@ -90,47 +102,56 @@ export function DocumentLayoutDrawer(props: {
     {
       defaultValues: {
         documentType: document.documentType,
+      },
+    },
+  );
+
+  const {
+    control: classificationControl,
+    handleSubmit: handleClassificationSubmit,
+    reset: resetClassification,
+  } = useFormWithSchema(
+    classificationSchema,
+    {
+      defaultValues: {
         classification: version.classification,
       },
     },
   );
 
   const [updateDocument, isUpdatingDocument]
-    = useMutationWithToasts<DocumentLayoutDrawerMutation>(
+    = useMutation<DocumentLayoutDrawerMutation>(updateDocumentMutation);
-      updateDocumentMutation,
-      {
-        successMessage: __("Document updated successfully."),
-        errorMessage: __("Failed to update document"),
-      },
-    );
 
-  const handleUpdateDocumentType = async (data: {
+  const [updateClassification, isUpdatingClassification]
+    = useMutation<DocumentLayoutDrawer_updateClassificationMutation>(updateClassificationMutation);
+
+  const handleUpdateDocumentType = (data: {
     documentType: (typeof documentTypes)[number];
   }) => {
-    await updateDocument({
+    updateDocument({
       variables: {
         input: {
           id: document.id,
           documentType: data.documentType,
         },
       },
-      onSuccess: () => {
+      onCompleted: () => {
         setIsEditingType(false);
       },
     });
   };
 
-  const handleUpdateClassification = async (data: {
+  const handleUpdateClassification = (data: {
     classification: (typeof documentClassifications)[number];
   }) => {
-    await updateDocument({
+    updateClassification({
       variables: {
         input: {
-          id: document.id,
+          documentVersionId: version.id,
           classification: data.classification,
         },
       },
-      onSuccess: () => {
+      onCompleted: () => {
         setIsEditingClassification(false);
       },
     });
@@ -176,16 +197,16 @@ export function DocumentLayoutDrawer(props: {
         {isEditingClassification
           ? (
               <EditablePropertyContent
-                onSave={() => void handleSubmit(handleUpdateClassification)()}
+                onSave={() => void handleClassificationSubmit(handleUpdateClassification)()}
                 onCancel={() => {
                   setIsEditingClassification(false);
-                  reset();
+                  resetClassification();
                 }}
-                disabled={isUpdatingDocument}
+                disabled={isUpdatingClassification}
               >
                 <ControlledField
                   name="classification"
-                  control={control}
+                  control={classificationControl}
                   type="select"
                 >
                   <DocumentClassificationOptions />
@@ -195,13 +216,10 @@ export function DocumentLayoutDrawer(props: {
           : (
               <ReadOnlyPropertyContent
                 onEdit={() => setIsEditingClassification(true)}
-                canEdit={canEdit}
+                canEdit={canEdit && isDraft}
               >
                 <div className="text-sm text-txt-secondary">
-                  {getDocumentClassificationLabel(
+                  {getDocumentClassificationLabel(__, version.classification)}
-                    __,
-                    version.classification,
-                  )}
                 </div>
               </ReadOnlyPropertyContent>
             )}

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

@@ -15,12 +15,11 @@
 import { formatDate, getDocumentClassificationLabel, getDocumentTypeLabel, sprintf } from "@probo/helpers";
 import { useTranslate } from "@probo/i18n";
 import { ActionDropdown, Badge, Checkbox, DropdownItem, IconTrashCan, Td, Tr, useConfirm } from "@probo/ui";
-import { useFragment } from "react-relay";
+import { useFragment, useMutation } from "react-relay";
 import { type DataID, graphql } from "relay-runtime";
 
 import type { DocumentListItem_deleteMutation } from "#/__generated__/core/DocumentListItem_deleteMutation.graphql";
 import type { DocumentListItemFragment$key } from "#/__generated__/core/DocumentListItemFragment.graphql";
-import { useMutationWithToasts } from "#/hooks/useMutationWithToasts";
 import { useOrganizationId } from "#/hooks/useOrganizationId";
 
 const fragment = graphql`
@@ -28,7 +27,6 @@ const fragment = graphql`
     id
     title
     documentType
-    classification
     updatedAt
     canDelete: permission(action: "core:document:delete")
     recentVersions: versions(first: 2 orderBy: { field: CREATED_AT direction: DESC }) {
@@ -38,6 +36,7 @@ const fragment = graphql`
           status
           major
           minor
+          classification
           approvalQuorums(first: 1, orderBy: { field: CREATED_AT, direction: DESC }) {
             edges {
               node {
@@ -117,23 +116,21 @@ export function DocumentListItem(props: {
     PUBLISHED: __("Published"),
   } as const;
 
-  const [deleteDocument] = useMutationWithToasts<DocumentListItem_deleteMutation>(
+  const [deleteDocument] = useMutation<DocumentListItem_deleteMutation>(deleteDocumentMutation);
-    deleteDocumentMutation,
-    {
-      successMessage: __("Document deleted successfully."),
-      errorMessage: __("Failed to delete document"),
-    },
-  );
   const confirm = useConfirm();
 
   const handleDelete = () => {
     confirm(
       () =>
+        new Promise<void>((resolve, reject) => {
           deleteDocument({
             variables: {
               connections: [connectionId],
               input: { documentId: document.id },
             },
+            onCompleted: () => resolve(),
+            onError: err => reject(err),
+          });
         }),
       {
         message: sprintf(
@@ -171,7 +168,7 @@ export function DocumentListItem(props: {
         {getDocumentTypeLabel(__, document.documentType)}
       </Td>
       <Td className="w-32">
-        {getDocumentClassificationLabel(__, document.classification)}
+        {getDocumentClassificationLabel(__, lastVersion.classification)}
       </Td>
       <Td className="w-60">
         {(() => {

apps/console/src/pages/organizations/employee/_components/ApprovableDocumentRow.tsx

@@ -28,9 +28,15 @@ const fragment = graphql`
     id
     title
     documentType
-    classification
     approvalState
     updatedAt
+    lastVersion: versions(first: 1 orderBy: { field: CREATED_AT direction: DESC }) {
+      edges {
+        node {
+          classification
+        }
+      }
+    }
   }
 `;
 
@@ -42,6 +48,7 @@ export function ApprovableDocumentRow({
   organizationId: string;
 }) {
   const document = useFragment<ApprovableDocumentRowFragment$key>(fragment, fKey);
+  const lastVersion = document.lastVersion.edges[0].node;
   const { __ } = useTranslate();
 
   const stateVariant = document.approvalState === "APPROVED"
@@ -64,7 +71,7 @@ export function ApprovableDocumentRow({
       </Td>
       <Td className="w-36">
         <Badge variant="neutral">
-          {getDocumentClassificationLabel(__, document.classification)}
+          {getDocumentClassificationLabel(__, lastVersion.classification)}
         </Badge>
       </Td>
       <Td className="w-40">{formatDate(document.updatedAt)}</Td>

apps/console/src/pages/organizations/employee/_components/DocumentRow.tsx

@@ -28,9 +28,15 @@ const fragment = graphql`
     id
     title
     documentType
-    classification
     signed
     updatedAt
+    lastVersion: versions(first: 1 orderBy: { field: CREATED_AT direction: DESC }) {
+      edges {
+        node {
+          classification
+        }
+      }
+    }
   }
 `;
 
@@ -42,6 +48,7 @@ export function DocumentRow({
   organizationId: string;
 }) {
   const document = useFragment<DocumentRowFragment$key>(fragment, fKey);
+  const lastVersion = document.lastVersion.edges[0].node;
   const { __ } = useTranslate();
 
   return (
@@ -52,7 +59,7 @@ export function DocumentRow({
       </Td>
       <Td className="w-36">
         <Badge variant="neutral">
-          {getDocumentClassificationLabel(__, document.classification)}
+          {getDocumentClassificationLabel(__, lastVersion.classification)}
         </Badge>
       </Td>
       <Td className="w-40">{formatDate(document.updatedAt)}</Td>

e2e/console/document_test.go

@@ -104,7 +104,6 @@ func TestDocument_Create(t *testing.T) {
 								id
 								title
 								documentType
-								classification
 							}
 						}
 					}
@@ -123,7 +122,6 @@ func TestDocument_Create(t *testing.T) {
 							ID           string `json:"id"`
 							Title        string `json:"title"`
 							DocumentType string `json:"documentType"`
-							Classification string `json:"classification"`
 						} `json:"node"`
 					} `json:"documentEdge"`
 				} `json:"createDocument"`

pkg/coredata/document.go

@@ -34,7 +34,6 @@ type (
 		OrganizationID        gid.GID               `db:"organization_id"`
 		Title                 string                `db:"title"`
 		DocumentType          DocumentType          `db:"document_type"`
-		Classification        DocumentClassification `db:"classification"`
 		CurrentPublishedMajor *int                  `db:"current_published_major"`
 		CurrentPublishedMinor *int                  `db:"current_published_minor"`
 		TrustCenterVisibility TrustCenterVisibility `db:"trust_center_visibility"`
@@ -128,7 +127,6 @@ SELECT
     organization_id,
     title,
     document_type,
-    classification,
     current_published_major,
     current_published_minor,
     trust_center_visibility,
@@ -182,7 +180,6 @@ SELECT
     organization_id,
     title,
     document_type,
-    classification,
     current_published_major,
     current_published_minor,
     trust_center_visibility,
@@ -237,7 +234,6 @@ SELECT
     organization_id,
     title,
     document_type,
-    classification,
     current_published_major,
     current_published_minor,
     trust_center_visibility,
@@ -321,7 +317,6 @@ SELECT
     organization_id,
     title,
     document_type,
-    classification,
     current_published_major,
     current_published_minor,
     trust_center_visibility,
@@ -374,7 +369,6 @@ SELECT
     organization_id,
     title,
     document_type,
-    classification,
     current_published_major,
     current_published_minor,
     trust_center_visibility,
@@ -441,7 +435,6 @@ SELECT
 	organization_id,
 	COALESCE(published_title, title) AS title,
 	document_type,
-	classification,
 	current_published_major,
 	current_published_minor,
 	trust_center_visibility,
@@ -491,7 +484,6 @@ INSERT INTO
 		organization_id,
 		title,
 		document_type,
-		classification,
 		current_published_major,
 		current_published_minor,
 		trust_center_visibility,
@@ -506,7 +498,6 @@ VALUES (
     @organization_id,
     @title,
     @document_type,
-    @classification,
     @current_published_major,
     @current_published_minor,
     @trust_center_visibility,
@@ -523,7 +514,6 @@ VALUES (
 		"organization_id":         p.OrganizationID,
 		"title":                   p.Title,
 		"document_type":           p.DocumentType,
-		"classification":          p.Classification,
 		"current_published_major": p.CurrentPublishedMajor,
 		"current_published_minor": p.CurrentPublishedMinor,
 		"trust_center_visibility": p.TrustCenterVisibility,
@@ -586,7 +576,6 @@ SET
 	current_published_major = @current_published_major,
 	current_published_minor = @current_published_minor,
 	document_type = @document_type,
-	classification = @classification,
 	trust_center_visibility = @trust_center_visibility,
 	status = @status,
 	archived_at = @archived_at,
@@ -605,7 +594,6 @@ WHERE
 		"current_published_major": p.CurrentPublishedMajor,
 		"current_published_minor": p.CurrentPublishedMinor,
 		"document_type":           p.DocumentType,
-		"classification":          p.Classification,
 		"trust_center_visibility": p.TrustCenterVisibility,
 		"status":                  p.Status,
 		"archived_at":             p.ArchivedAt,
@@ -678,7 +666,6 @@ SELECT
 	scoped_documents.organization_id,
 	scoped_documents.title,
 	scoped_documents.document_type,
-	scoped_documents.classification,
 	scoped_documents.current_published_major,
 	scoped_documents.current_published_minor,
 	scoped_documents.trust_center_visibility,
@@ -770,7 +757,6 @@ SELECT
 	scoped_documents.organization_id,
 	scoped_documents.title,
 	scoped_documents.document_type,
-	scoped_documents.classification,
 	scoped_documents.current_published_major,
 	scoped_documents.current_published_minor,
 	scoped_documents.trust_center_visibility,

pkg/coredata/migrations/20260328T120000Z.sql

@@ -0,0 +1,2 @@
+-- TODO: drop the classification column from documents.
+ALTER TABLE documents ALTER COLUMN classification SET DEFAULT 'SECRET';

pkg/probo/document_service.go

@@ -85,14 +85,14 @@ type (
 	UpdateDocumentRequest struct {
 		DocumentID            gid.GID
 		Title                 *string
-		Classification        *coredata.DocumentClassification
 		DocumentType          *coredata.DocumentType
 		TrustCenterVisibility *coredata.TrustCenterVisibility
 	}
 
 	UpdateDocumentVersionRequest struct {
 		ID             gid.GID
-		Content string
+		Content        *string
+		Classification *coredata.DocumentClassification
 	}
 
 	RequestSignatureRequest struct {
@@ -133,7 +133,6 @@ func (udr *UpdateDocumentRequest) Validate() error {
 
 	v.Check(udr.DocumentID, "document_id", validator.Required(), validator.GID(coredata.DocumentEntityType))
 	v.Check(udr.Title, "title", validator.SafeTextNoNewLine(TitleMaxLength))
-	v.Check(udr.Classification, "classification", validator.OneOfSlice(coredata.DocumentClassifications()))
 	v.Check(udr.DocumentType, "document_type", validator.OneOfSlice(coredata.DocumentTypes()))
 	v.Check(udr.TrustCenterVisibility, "trust_center_visibility", validator.OneOfSlice(coredata.TrustCenterVisibilities()))
 
@@ -144,7 +143,8 @@ func (udvr *UpdateDocumentVersionRequest) Validate() error {
 	v := validator.New()
 
 	v.Check(udvr.ID, "id", validator.Required(), validator.GID(coredata.DocumentVersionEntityType))
-	v.Check(udvr.Content, "content", validator.Required(), validator.NotEmpty(), validator.MaxLen(documentMaxLength))
+	v.Check(udvr.Content, "content", validator.NotEmpty(), validator.MaxLen(documentMaxLength))
+	v.Check(udvr.Classification, "classification", validator.OneOfSlice(coredata.DocumentClassifications()))
 
 	return v.Error()
 }
@@ -509,7 +509,6 @@ func (s *DocumentService) Create(
 		Title:                 req.Title,
 		DocumentType:          req.DocumentType,
 		TrustCenterVisibility: coredata.TrustCenterVisibilityNone,
-		Classification:        req.Classification,
 		Status:                coredata.DocumentStatusActive,
 		CreatedAt:             now,
 		UpdatedAt:             now,
@@ -763,8 +762,12 @@ func (s *DocumentService) UpdateVersion(
 			}
 
 			documentVersion.Title = document.Title
-			documentVersion.Classification = document.Classification
+			if req.Content != nil {
-			documentVersion.Content = req.Content
+				documentVersion.Content = *req.Content
+			}
+			if req.Classification != nil {
+				documentVersion.Classification = *req.Classification
+			}
 			documentVersion.UpdatedAt = time.Now()
 
 			if err := documentVersion.Update(ctx, conn, s.svc.scope); err != nil {
@@ -1002,7 +1005,7 @@ func (s *DocumentService) CreateDraft(
 			draftVersion.Title = document.Title
 			draftVersion.Major = latestVersion.Major
 			draftVersion.Minor = latestVersion.Minor + 1
-			draftVersion.Classification = document.Classification
+			draftVersion.Classification = latestVersion.Classification
 			draftVersion.Content = latestVersion.Content
 			draftVersion.Status = coredata.DocumentVersionStatusDraft
 			draftVersion.CreatedAt = now
@@ -1533,10 +1536,6 @@ func (s *DocumentService) Update(
 				document.Title = *req.Title
 			}
 
-			if req.Classification != nil {
-				document.Classification = *req.Classification
-			}
-
 			if req.DocumentType != nil {
 				document.DocumentType = *req.DocumentType
 			}
@@ -1559,7 +1558,6 @@ func (s *DocumentService) Update(
 			err := draftVersion.LoadLatestVersion(ctx, tx, s.svc.scope, req.DocumentID)
 			if err == nil && draftVersion.Status == coredata.DocumentVersionStatusDraft {
 				draftVersion.Title = document.Title
-				draftVersion.Classification = document.Classification
 				draftVersion.UpdatedAt = now
 
 				if err := draftVersion.Update(ctx, tx, s.svc.scope); err != nil {
@@ -1934,7 +1932,7 @@ func exportDocumentPDF(
 	}
 
 	classification := docgen.ClassificationSecret
-	switch document.Classification {
+	switch version.Classification {
 	case coredata.DocumentClassificationPublic:
 		classification = docgen.ClassificationPublic
 	case coredata.DocumentClassificationInternal:

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

@@ -2440,7 +2440,6 @@ type Document implements Node {
     title: String!
     description: String
     documentType: DocumentType!
-    classification: DocumentClassification!
     currentPublishedMajor: Int
     currentPublishedMinor: Int
     trustCenterVisibility: TrustCenterVisibility!
@@ -2481,7 +2480,6 @@ type EmployeeDocument
     title: String!
     description: String
     documentType: DocumentType!
-    classification: DocumentClassification!
     signed: Boolean @goField(forceResolver: true)
     approvalState: DocumentVersionApprovalDecisionState @goField(forceResolver: true)
 
@@ -2505,6 +2503,7 @@ type EmployeeDocumentVersion
     major: Int!
     minor: Int!
     status: DocumentVersionStatus!
+    classification: DocumentClassification!
     signed: Boolean! @goField(forceResolver: true)
     approvalDecision: DocumentVersionApprovalDecision @goField(forceResolver: true)
     publishedAt: Datetime
@@ -4482,7 +4481,6 @@ input UpdateDocumentInput {
     title: String
     content: String
     documentType: DocumentType
-    classification: DocumentClassification
     trustCenterVisibility: TrustCenterVisibility
 }
 
@@ -5738,7 +5736,8 @@ input DeleteDraftDocumentVersionInput {
 
 input UpdateDocumentVersionInput {
     documentVersionId: ID!
-    content: String!
+    content: String
+    classification: DocumentClassification
 }
 
 input CancelSignatureRequestInput {

pkg/server/api/console/v1/types/document.go

@@ -81,7 +81,6 @@ func NewDocument(document *coredata.Document) *Document {
 			ID: document.OrganizationID,
 		},
 		DocumentType:          document.DocumentType,
-		Classification:        document.Classification,
 		CurrentPublishedMajor: document.CurrentPublishedMajor,
 		CurrentPublishedMinor: document.CurrentPublishedMinor,
 		TrustCenterVisibility: document.TrustCenterVisibility,

pkg/server/api/console/v1/types/employee_document.go

@@ -46,7 +46,6 @@ type (
 		Title        string
 		Description  *string
 		DocumentType coredata.DocumentType
-		Classification coredata.DocumentClassification
 		CreatedAt    time.Time
 		UpdatedAt    time.Time
 
@@ -69,6 +68,7 @@ type (
 		Major          int
 		Minor          int
 		Status         coredata.DocumentVersionStatus
+		Classification coredata.DocumentClassification
 		PublishedAt    *time.Time
 		CreatedAt      time.Time
 		UpdatedAt      time.Time

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

@@ -1664,6 +1664,7 @@ func (r *employeeDocumentResolver) Versions(ctx context.Context, obj *types.Empl
 			Major:          v.Major,
 			Minor:          v.Minor,
 			Status:         v.Status,
+			Classification: v.Classification,
 			PublishedAt:    v.PublishedAt,
 			CreatedAt:      v.CreatedAt,
 			UpdatedAt:      v.UpdatedAt,
@@ -4669,7 +4670,6 @@ func (r *mutationResolver) UpdateDocument(ctx context.Context, input types.Updat
 		probo.UpdateDocumentRequest{
 			DocumentID:            input.ID,
 			Title:                 input.Title,
-			Classification:        input.Classification,
 			DocumentType:          input.DocumentType,
 			TrustCenterVisibility: input.TrustCenterVisibility,
 		},
@@ -5412,6 +5412,7 @@ func (r *mutationResolver) UpdateDocumentVersion(ctx context.Context, input type
 		probo.UpdateDocumentVersionRequest{
 			ID:             input.DocumentVersionID,
 			Content:        input.Content,
+			Classification: input.Classification,
 		},
 	)
 	if err != nil {
@@ -10246,7 +10247,6 @@ func (r *viewerResolver) SignableDocuments(ctx context.Context, obj *types.Viewe
 			ID:           doc.ID,
 			Title:        doc.Title,
 			DocumentType: doc.DocumentType,
-			Classification: doc.Classification,
 			CreatedAt:    doc.CreatedAt,
 			UpdatedAt:    doc.UpdatedAt,
 			FilterMode:   types.EmployeeDocumentFilterModeSignature,
@@ -10283,7 +10283,6 @@ func (r *viewerResolver) SignableDocument(ctx context.Context, obj *types.Viewer
 		ID:           document.ID,
 		Title:        document.Title,
 		DocumentType: document.DocumentType,
-		Classification: document.Classification,
 		CreatedAt:    document.CreatedAt,
 		UpdatedAt:    document.UpdatedAt,
 		FilterMode:   types.EmployeeDocumentFilterModeSignature,
@@ -10327,7 +10326,6 @@ func (r *viewerResolver) ApprovableDocuments(ctx context.Context, obj *types.Vie
 			ID:           doc.ID,
 			Title:        doc.Title,
 			DocumentType: doc.DocumentType,
-			Classification: doc.Classification,
 			CreatedAt:    doc.CreatedAt,
 			UpdatedAt:    doc.UpdatedAt,
 			FilterMode:   types.EmployeeDocumentFilterModeApproval,
@@ -10364,7 +10362,6 @@ func (r *viewerResolver) ApprovableDocument(ctx context.Context, obj *types.View
 		ID:           document.ID,
 		Title:        document.Title,
 		DocumentType: document.DocumentType,
-		Classification: document.Classification,
 		CreatedAt:    document.CreatedAt,
 		UpdatedAt:    document.UpdatedAt,
 		FilterMode:   types.EmployeeDocumentFilterModeApproval,

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

@@ -2106,7 +2106,6 @@ func (r *Resolver) UpdateDocumentTool(ctx context.Context, req *mcp.CallToolRequ
 		probo.UpdateDocumentRequest{
 			DocumentID:            input.ID,
 			Title:                 input.Title,
-			Classification:        input.Classification,
 			DocumentType:          input.DocumentType,
 			TrustCenterVisibility: input.TrustCenterVisibility,
 		},
@@ -2185,6 +2184,7 @@ func (r *Resolver) UpdateDocumentVersionTool(ctx context.Context, req *mcp.CallT
 		probo.UpdateDocumentVersionRequest{
 			ID:             input.DocumentVersionID,
 			Content:        input.Content,
+			Classification: input.Classification,
 		},
 	)
 	if err != nil {

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

@@ -5186,7 +5186,6 @@ components:
         - organization_id
         - title
         - document_type
-        - classification
         - trust_center_visibility
         - status
         - created_at
@@ -5204,9 +5203,6 @@ components:
         document_type:
           $ref: "#/components/schemas/DocumentType"
           description: Document type
-        classification:
-          $ref: "#/components/schemas/DocumentClassification"
-          description: Document classification
         current_published_major:
           type:
             - integer
@@ -5459,9 +5455,6 @@ components:
         title:
           type: string
           description: Document title
-        classification:
-          $ref: "#/components/schemas/DocumentClassification"
-          description: Document classification
         document_type:
           $ref: "#/components/schemas/DocumentType"
           description: Document type
@@ -5580,7 +5573,6 @@ components:
       type: object
       required:
         - document_version_id
-        - content
       properties:
         document_version_id:
           $ref: "#/components/schemas/GID"
@@ -5588,6 +5580,9 @@ components:
         content:
           type: string
           description: Document content
+        classification:
+          $ref: "#/components/schemas/DocumentClassification"
+          description: Document classification
 
     UpdateDocumentVersionOutput:
       type: object

pkg/server/api/mcp/v1/types/document.go

@@ -25,7 +25,6 @@ func NewDocument(d *coredata.Document) *Document {
 		OrganizationID:        d.OrganizationID,
 		Title:                 d.Title,
 		DocumentType:          d.DocumentType,
-		Classification:        d.Classification,
 		CurrentPublishedMajor: d.CurrentPublishedMajor,
 		CurrentPublishedMinor: d.CurrentPublishedMinor,
 		TrustCenterVisibility: d.TrustCenterVisibility,

pkg/trust/document_service.go

@@ -219,12 +219,14 @@ func (s *DocumentService) exportPDFData(
 		return nil, err
 	}
 
-	classification := docgen.ClassificationInternal
+	classification := docgen.ClassificationSecret
-	switch document.DocumentType {
+	switch version.Classification {
-	case coredata.DocumentTypePolicy:
+	case coredata.DocumentClassificationPublic:
+		classification = docgen.ClassificationPublic
+	case coredata.DocumentClassificationInternal:
+		classification = docgen.ClassificationInternal
+	case coredata.DocumentClassificationConfidential:
 		classification = docgen.ClassificationConfidential
-	case coredata.DocumentTypeGovernance:
-		classification = docgen.ClassificationSecret
 	}
 
 	horizontalLogoBase64 := ""