mcp: Add draft and publish (#11488)

* mcp: Add draft and publish

* refactor

* touch

server/models/Document.ts

@@ -48,7 +48,7 @@ import { ProsemirrorHelper } from "@shared/utils/ProsemirrorHelper";
 import { UrlHelper } from "@shared/utils/UrlHelper";
 import slugify from "@shared/utils/slugify";
 import { DocumentValidation } from "@shared/validations";
-import { ValidationError } from "@server/errors";
+import { InvalidRequestError, ValidationError } from "@server/errors";
 import { generateUrlId } from "@server/utils/url";
 import { createContext } from "@server/context";
 import Collection from "./Collection";
@@ -1098,10 +1098,12 @@ class Document extends ArchivableModel<
   };
 
   /**
+   * Unpublishes a published document, converting it back to a draft.
    *
-   * @param user User who is performing the action
-   * @param options.detach Whether to detach the document from the containing collection
-   * @returns Updated document
+   * @param ctx the API context.
+   * @param options.detach whether to detach the document from the containing collection.
+   * @returns updated document.
+   * @throws if the document has child documents.
    */
   unpublishWithCtx = async (ctx: APIContext, options: { detach: boolean }) => {
     // If the document is already a draft then calling unpublish should act like save
@@ -1112,11 +1114,21 @@ class Document extends ArchivableModel<
     const { user } = ctx.state.auth;
     const { transaction } = ctx.state;
 
+    const childDocumentIds = await this.findAllChildDocumentIds(
+      { archivedAt: { [Op.eq]: null } },
+      { transaction }
+    );
+    if (childDocumentIds.length > 0) {
+      throw InvalidRequestError(
+        "Cannot unpublish document with child documents"
+      );
+    }
+
     const collection = this.collectionId
       ? await Collection.findByPk(this.collectionId, {
           includeDocumentStructure: true,
           transaction,
-          lock: transaction.LOCK.UPDATE,
+          lock: transaction?.LOCK.UPDATE,
         })
       : undefined;
 

server/routes/api/documents/documents.ts

@@ -1542,27 +1542,12 @@ router.post(
   async (ctx: APIContext<T.DocumentsUnpublishReq>) => {
     const { id, detach } = ctx.input.body;
     const { user } = ctx.state.auth;
-    const { transaction } = ctx.state;
 
     const document = await Document.findByPk(id, {
       userId: user.id,
     });
     authorize(user, "unpublish", document);
 
-    const childDocumentIds = await document.findAllChildDocumentIds(
-      {
-        archivedAt: {
-          [Op.eq]: null,
-        },
-      },
-      { transaction }
-    );
-    if (childDocumentIds.length > 0) {
-      throw InvalidRequestError(
-        "Cannot unpublish document with child documents"
-      );
-    }
-
     await document.unpublishWithCtx(ctx, { detach });
 
     ctx.body = {

server/routes/mcp/index.test.ts

@@ -335,6 +335,54 @@ describe("POST /mcp/", () => {
       expect(data.url).toMatch(/^https?:\/\//);
     });
 
+    it("update_document unpublishes a document", async () => {
+      const { user, accessToken } = await buildOAuthUser();
+      const collection = await buildCollection({
+        teamId: user.teamId,
+        userId: user.id,
+      });
+      const document = await buildDocument({
+        teamId: user.teamId,
+        userId: user.id,
+        collectionId: collection.id,
+      });
+
+      const res = await callMcpTool(server, accessToken, "update_document", {
+        id: document.id,
+        publish: false,
+      });
+      const data = JSON.parse(res?.result?.content?.[0]?.text ?? "{}");
+
+      expect(data.id).toEqual(document.id);
+      expect(res?.result?.isError).toBeUndefined();
+    });
+
+    it("update_document fails to unpublish a document with children", async () => {
+      const { user, accessToken } = await buildOAuthUser();
+      const collection = await buildCollection({
+        teamId: user.teamId,
+        userId: user.id,
+      });
+      const document = await buildDocument({
+        teamId: user.teamId,
+        userId: user.id,
+        collectionId: collection.id,
+      });
+      await buildDocument({
+        teamId: user.teamId,
+        userId: user.id,
+        collectionId: collection.id,
+        parentDocumentId: document.id,
+      });
+
+      const res = await callMcpTool(server, accessToken, "update_document", {
+        id: document.id,
+        publish: false,
+      });
+
+      expect(res?.result?.isError).toBe(true);
+    });
+
     it("get_document resource returns metadata and markdown", async () => {
       const { user, accessToken } = await buildOAuthUser();
       const collection = await buildCollection({

server/tools/documents.ts

@@ -202,7 +202,7 @@ export function documentTools(server: McpServer, scopes: string[]) {
       {
         title: "Create document",
         description:
-          "Creates a new document. Requires a collectionId to place the document in a collection, or parentDocumentId to nest it under an existing document. Set publish to true to make the document visible in the collection.",
+          "Creates a new document. Requires a collectionId to place the document in a collection, or parentDocumentId to nest it under an existing document.",
         annotations: {
           idempotentHint: false,
           readOnlyHint: false,
@@ -229,6 +229,12 @@ export function documentTools(server: McpServer, scopes: string[]) {
             .string()
             .optional()
             .describe("The hex color for the document icon, e.g. #FF0000."),
+          publish: z
+            .boolean()
+            .optional()
+            .describe(
+              "Whether to publish the document. Defaults to true. Set to false to create as a draft."
+            ),
         },
       },
       async (input, context) => {
@@ -267,7 +273,7 @@ export function documentTools(server: McpServer, scopes: string[]) {
             icon: input.icon,
             color: input.color,
             parentDocumentId: parentDocumentId,
-            publish: true,
+            publish: input.publish !== false,
             collectionId: collection?.id,
           });
 
@@ -346,6 +352,12 @@ export function documentTools(server: McpServer, scopes: string[]) {
             .describe(
               "The hex color for the document icon. Set to null to remove."
             ),
+          publish: z
+            .boolean()
+            .optional()
+            .describe(
+              "Set to true to publish a draft document, or false to convert a published document back to a draft."
+            ),
         },
       },
       async (input, context) => {
@@ -359,12 +371,22 @@ export function documentTools(server: McpServer, scopes: string[]) {
             rejectOnEmpty: true,
           });
 
-          authorize(user, "update", document);
+          let updated;
 
-          const updated = await documentUpdater(ctx, {
-            document,
-            ...input,
-          });
+          if (input.publish === false) {
+            authorize(user, "unpublish", document);
+
+            updated = await document.unpublishWithCtx(ctx, {
+              detach: false,
+            });
+          } else {
+            authorize(user, "update", document);
+
+            updated = await documentUpdater(ctx, {
+              document,
+              ...input,
+            });
+          }
 
           const { text, ...attributes } = await presentDocument(
             undefined,