feat: refactor doc write in native (#14272)

packages/backend/server/src/plugins/copilot/mcp/provider.ts

@@ -166,146 +166,217 @@ export class WorkspaceMcpProvider {
       }
     );
 
-    // Write tools - create and update documents
-    server.registerTool(
-      'create_document',
-      {
-        title: 'Create Document',
-        description:
-          'Create a new document in the workspace with the given title and markdown content. Returns the ID of the created document.',
-        inputSchema: z.object({
-          title: z.string().min(1).describe('The title of the new document'),
-          content: z
-            .string()
-            .describe(
-              'The markdown content for the document body (should NOT include a title H1 - the title parameter will be used)'
-            ),
-        }),
-      },
-      async ({ title, content }) => {
-        try {
-          // Check if user can create docs in this workspace
-          await this.ac
+    if (env.dev || env.namespaces.canary) {
+      // Write tools - create and update documents
+      server.registerTool(
+        'create_document',
+        {
+          title: 'Create Document',
+          description:
+            'Create a new document in the workspace with the given title and markdown content. Returns the ID of the created document. This tool not support insert or update database block and image yet.',
+          inputSchema: z.object({
+            title: z.string().min(1).describe('The title of the new document'),
+            content: z
+              .string()
+              .describe('The markdown content for the document body'),
+          }),
+        },
+        async ({ title, content }) => {
+          try {
+            // Check if user can create docs in this workspace
+            await this.ac
+              .user(userId)
+              .workspace(workspaceId)
+              .assert('Workspace.CreateDoc');
+
+            // Sanitize title by removing newlines and trimming
+            const sanitizedTitle = title.replace(/[\r\n]+/g, ' ').trim();
+            if (!sanitizedTitle) {
+              throw new Error('Title cannot be empty');
+            }
+
+            // Strip any leading H1 from content to prevent duplicates
+            // Per CommonMark spec, ATX headings allow only 0-3 spaces before the #
+            // Handles: "# Title", "  # Title", "# Title #"
+            const strippedContent = content.replace(
+              /^[ \t]{0,3}#\s+[^\n]*#*\s*\n*/,
+              ''
+            );
+
+            // Create the document
+            const result = await this.writer.createDoc(
+              workspaceId,
+              sanitizedTitle,
+              strippedContent,
+              userId
+            );
+
+            return {
+              content: [
+                {
+                  type: 'text',
+                  text: JSON.stringify({
+                    success: true,
+                    docId: result.docId,
+                    message: `Document "${title}" created successfully`,
+                  }),
+                },
+              ],
+            } as const;
+          } catch (error) {
+            return {
+              isError: true,
+              content: [
+                {
+                  type: 'text',
+                  text: `Failed to create document: ${error instanceof Error ? error.message : 'Unknown error'}`,
+                },
+              ],
+            };
+          }
+        }
+      );
+
+      server.registerTool(
+        'update_document',
+        {
+          title: 'Update Document',
+          description:
+            'Update an existing document with new markdown content (body only). Uses structural diffing to apply minimal changes, preserving document history and enabling real-time collaboration. This does NOT update the document title. This tool not support insert or update database block and image yet.',
+          inputSchema: z.object({
+            docId: z.string().describe('The ID of the document to update'),
+            content: z
+              .string()
+              .describe(
+                'The complete new markdown content for the document body (do NOT include a title H1)'
+              ),
+          }),
+        },
+        async ({ docId, content }) => {
+          const notFoundError: CallToolResult = {
+            isError: true,
+            content: [
+              {
+                type: 'text',
+                text: `Doc with id ${docId} not found.`,
+              },
+            ],
+          };
+
+          // Use can() instead of assert() to avoid leaking doc existence info
+          const accessible = await this.ac
             .user(userId)
             .workspace(workspaceId)
-            .assert('Workspace.CreateDoc');
+            .doc(docId)
+            .can('Doc.Update');
 
-          // Combine title and content into markdown
-          // Sanitize title by removing newlines and trimming
-          const sanitizedTitle = title.replace(/[\r\n]+/g, ' ').trim();
-          if (!sanitizedTitle) {
-            throw new Error('Title cannot be empty');
+          if (!accessible) {
+            return notFoundError;
           }
 
-          // Strip any leading H1 from content to prevent duplicates
-          // Per CommonMark spec, ATX headings allow only 0-3 spaces before the #
-          // Handles: "# Title", "  # Title", "# Title #"
-          const strippedContent = content.replace(
-            /^[ \t]{0,3}#\s+[^\n]*#*\s*\n*/,
-            ''
-          );
+          try {
+            // Update the document
+            await this.writer.updateDoc(workspaceId, docId, content, userId);
 
-          const markdown = `# ${sanitizedTitle}\n\n${strippedContent}`;
+            return {
+              content: [
+                {
+                  type: 'text',
+                  text: JSON.stringify({
+                    success: true,
+                    docId,
+                    message: `Document updated successfully`,
+                  }),
+                },
+              ],
+            } as const;
+          } catch (error) {
+            return {
+              isError: true,
+              content: [
+                {
+                  type: 'text',
+                  text: `Failed to update document: ${error instanceof Error ? error.message : 'Unknown error'}`,
+                },
+              ],
+            };
+          }
+        }
+      );
 
-          // Create the document
-          const result = await this.writer.createDoc(
-            workspaceId,
-            markdown,
-            userId
-          );
-
-          return {
-            content: [
-              {
-                type: 'text',
-                text: JSON.stringify({
-                  success: true,
-                  docId: result.docId,
-                  message: `Document "${title}" created successfully`,
-                }),
-              },
-            ],
-          } as const;
-        } catch (error) {
-          return {
+      server.registerTool(
+        'update_document_meta',
+        {
+          title: 'Update Document Metadata',
+          description: 'Update document metadata (currently title only).',
+          inputSchema: z.object({
+            docId: z.string().describe('The ID of the document to update'),
+            title: z.string().min(1).describe('The new document title'),
+          }),
+        },
+        async ({ docId, title }) => {
+          const notFoundError: CallToolResult = {
             isError: true,
             content: [
               {
                 type: 'text',
-                text: `Failed to create document: ${error instanceof Error ? error.message : 'Unknown error'}`,
+                text: `Doc with id ${docId} not found.`,
               },
             ],
           };
-        }
-      }
-    );
 
-    server.registerTool(
-      'update_document',
-      {
-        title: 'Update Document',
-        description:
-          'Update an existing document with new markdown content. Uses structural diffing to apply minimal changes, preserving document history and enabling real-time collaboration.',
-        inputSchema: z.object({
-          docId: z.string().describe('The ID of the document to update'),
-          content: z
-            .string()
-            .describe(
-              'The complete new markdown content for the document (including title as H1)'
-            ),
-        }),
-      },
-      async ({ docId, content }) => {
-        const notFoundError: CallToolResult = {
-          isError: true,
-          content: [
-            {
-              type: 'text',
-              text: `Doc with id ${docId} not found.`,
-            },
-          ],
-        };
+          // Use can() instead of assert() to avoid leaking doc existence info
+          const accessible = await this.ac
+            .user(userId)
+            .workspace(workspaceId)
+            .doc(docId)
+            .can('Doc.Update');
 
-        // Use can() instead of assert() to avoid leaking doc existence info
-        const accessible = await this.ac
-          .user(userId)
-          .workspace(workspaceId)
-          .doc(docId)
-          .can('Doc.Update');
+          if (!accessible) {
+            return notFoundError;
+          }
 
-        if (!accessible) {
-          return notFoundError;
-        }
+          try {
+            const sanitizedTitle = title.replace(/[\r\n]+/g, ' ').trim();
+            if (!sanitizedTitle) {
+              throw new Error('Title cannot be empty');
+            }
 
-        try {
-          // Update the document
-          await this.writer.updateDoc(workspaceId, docId, content, userId);
-
-          return {
-            content: [
+            await this.writer.updateDocMeta(
+              workspaceId,
+              docId,
               {
-                type: 'text',
-                text: JSON.stringify({
-                  success: true,
-                  docId,
-                  message: `Document updated successfully`,
-                }),
+                title: sanitizedTitle,
               },
-            ],
-          } as const;
-        } catch (error) {
-          return {
-            isError: true,
-            content: [
-              {
-                type: 'text',
-                text: `Failed to update document: ${error instanceof Error ? error.message : 'Unknown error'}`,
-              },
-            ],
-          };
+              userId
+            );
+
+            return {
+              content: [
+                {
+                  type: 'text',
+                  text: JSON.stringify({
+                    success: true,
+                    docId,
+                    message: `Document title updated successfully`,
+                  }),
+                },
+              ],
+            } as const;
+          } catch (error) {
+            return {
+              isError: true,
+              content: [
+                {
+                  type: 'text',
+                  text: `Failed to update document metadata: ${error instanceof Error ? error.message : 'Unknown error'}`,
+                },
+              ],
+            };
+          }
         }
-      }
-    );
+      );
+    }
 
     return server;
   }

packages/backend/server/src/plugins/copilot/prompt/chat-prompt.ts

@@ -127,6 +127,7 @@ export class ChatPrompt {
       selectedMarkdown,
       selectedSnapshot,
       html,
+      currentDocId,
     } = params;
     return {
       'affine::date': new Date().toLocaleDateString(),
@@ -135,6 +136,8 @@ export class ChatPrompt {
       'affine::hasDocsRef': Array.isArray(docs) && docs.length > 0,
       'affine::hasFilesRef': Array.isArray(files) && files.length > 0,
       'affine::hasSelected': !!selectedMarkdown || !!selectedSnapshot || !!html,
+      'affine::hasCurrentDoc':
+        typeof currentDocId === 'string' && currentDocId.trim().length > 0,
     };
   }
 

packages/backend/server/src/plugins/copilot/prompt/prompts.ts

@@ -1950,6 +1950,13 @@ User's preferred language is {{affine::language}}.
 User's timezone is {{affine::timezone}}.
 </real_world_info>
 
+{{#affine::hasCurrentDoc}}
+<current_document_context>
+The user is chatting within the current document: {{currentDocId}}.
+If the user's request relates to this document, call the doc_read tool with docId {{currentDocId}} to read it before answering.
+</current_document_context>
+{{/affine::hasCurrentDoc}}
+
 <content_analysis>
 - If documents are provided, analyze all documents based on the user's query
 - Identify key information relevant to the user's specific request
@@ -2086,7 +2093,10 @@ Below is the user's query. Please respond in the user's preferred language witho
   config: {
     tools: [
       'docRead',
-      'sectionEdit',
+      'docCreate',
+      'docUpdate',
+      'docUpdateMeta',
+      // 'sectionEdit',
       'docKeywordSearch',
       'docSemanticSearch',
       'webSearch',

packages/backend/server/src/plugins/copilot/providers/provider.ts

@@ -9,7 +9,7 @@ import {
   CopilotProviderNotSupported,
   OnEvent,
 } from '../../../base';
-import { DocReader } from '../../../core/doc';
+import { DocReader, DocWriter } from '../../../core/doc';
 import { AccessController } from '../../../core/permission';
 import { Models } from '../../../models';
 import { IndexerService } from '../../indexer';
@@ -19,16 +19,22 @@ import {
   buildBlobContentGetter,
   buildContentGetter,
   buildDocContentGetter,
+  buildDocCreateHandler,
   buildDocKeywordSearchGetter,
   buildDocSearchGetter,
+  buildDocUpdateHandler,
+  buildDocUpdateMetaHandler,
   createBlobReadTool,
   createCodeArtifactTool,
   createConversationSummaryTool,
   createDocComposeTool,
+  createDocCreateTool,
   createDocEditTool,
   createDocKeywordSearchTool,
   createDocReadTool,
   createDocSemanticSearchTool,
+  createDocUpdateMetaTool,
+  createDocUpdateTool,
   createExaCrawlTool,
   createExaSearchTool,
   createSectionEditTool,
@@ -163,6 +169,7 @@ export abstract class CopilotProvider<C = any> {
         strict: false,
       });
       const docReader = this.moduleRef.get(DocReader, { strict: false });
+      const docWriter = this.moduleRef.get(DocWriter, { strict: false });
       const models = this.moduleRef.get(Models, { strict: false });
       const prompt = this.moduleRef.get(PromptService, {
         strict: false,
@@ -177,6 +184,12 @@ export abstract class CopilotProvider<C = any> {
           }
           continue;
         }
+        if (
+          !(env.dev || env.namespaces.canary) &&
+          ['docCreate', 'docUpdate', 'docUpdateMeta'].includes(tool)
+        ) {
+          continue;
+        }
         switch (tool) {
           case 'blobRead': {
             const docContext = options.session
@@ -244,6 +257,27 @@ export abstract class CopilotProvider<C = any> {
             tools.doc_read = createDocReadTool(getDoc.bind(null, options));
             break;
           }
+          case 'docCreate': {
+            const createDoc = buildDocCreateHandler(ac, docWriter);
+            tools.doc_create = createDocCreateTool(
+              createDoc.bind(null, options)
+            );
+            break;
+          }
+          case 'docUpdate': {
+            const updateDoc = buildDocUpdateHandler(ac, docWriter);
+            tools.doc_update = createDocUpdateTool(
+              updateDoc.bind(null, options)
+            );
+            break;
+          }
+          case 'docUpdateMeta': {
+            const updateDocMeta = buildDocUpdateMetaHandler(ac, docWriter);
+            tools.doc_update_meta = createDocUpdateMetaTool(
+              updateDocMeta.bind(null, options)
+            );
+            break;
+          }
           case 'webSearch': {
             tools.web_search_exa = createExaSearchTool(this.AFFiNEConfig);
             tools.web_crawl_exa = createExaCrawlTool(this.AFFiNEConfig);

packages/backend/server/src/plugins/copilot/providers/types.ts

@@ -66,6 +66,9 @@ export const PromptToolsSchema = z
     'docEdit',
     // work with indexer
     'docRead',
+    'docCreate',
+    'docUpdate',
+    'docUpdateMeta',
     'docKeywordSearch',
     // work with embeddings
     'docSemanticSearch',

packages/backend/server/src/plugins/copilot/tools/doc-write.ts

@@ -0,0 +1,207 @@
+import { Logger } from '@nestjs/common';
+import { tool } from 'ai';
+import { z } from 'zod';
+
+import { DocWriter } from '../../../core/doc';
+import { AccessController } from '../../../core/permission';
+import type { CopilotChatOptions } from '../providers';
+import { toolError } from './error';
+
+const logger = new Logger('DocWriteTool');
+
+const stripLeadingH1 = (content: string) =>
+  content.replace(/^[ \t]{0,3}#\s+[^\n]*#*\s*\n*/, '');
+
+const sanitizeTitle = (title: string) => title.replace(/[\r\n]+/g, ' ').trim();
+
+export const buildDocCreateHandler = (
+  ac: AccessController,
+  writer: DocWriter
+) => {
+  return async (
+    options: CopilotChatOptions,
+    title: string,
+    content: string
+  ) => {
+    if (!options?.user || !options.workspace) {
+      return toolError(
+        'Doc Create Failed',
+        'Missing user or workspace context'
+      );
+    }
+
+    await ac
+      .user(options.user)
+      .workspace(options.workspace)
+      .assert('Workspace.CreateDoc');
+
+    const sanitizedTitle = sanitizeTitle(title);
+    if (!sanitizedTitle) {
+      return toolError('Doc Create Failed', 'Title cannot be empty');
+    }
+
+    const strippedContent = stripLeadingH1(content);
+    const result = await writer.createDoc(
+      options.workspace,
+      sanitizedTitle,
+      strippedContent,
+      options.user
+    );
+
+    return {
+      success: true,
+      docId: result.docId,
+      message: `Document "${sanitizedTitle}" created successfully`,
+    };
+  };
+};
+
+export const buildDocUpdateHandler = (
+  ac: AccessController,
+  writer: DocWriter
+) => {
+  return async (
+    options: CopilotChatOptions,
+    docId: string,
+    content: string
+  ) => {
+    const notFound = toolError(
+      'Doc Update Failed',
+      `Doc with id ${docId} not found.`
+    );
+
+    if (!options?.user || !options.workspace) {
+      return notFound;
+    }
+
+    const canAccess = await ac
+      .user(options.user)
+      .workspace(options.workspace)
+      .doc(docId)
+      .can('Doc.Update');
+
+    if (!canAccess) {
+      return notFound;
+    }
+
+    await writer.updateDoc(options.workspace, docId, content, options.user);
+
+    return {
+      success: true,
+      docId,
+      message: 'Document updated successfully',
+    };
+  };
+};
+
+export const buildDocUpdateMetaHandler = (
+  ac: AccessController,
+  writer: DocWriter
+) => {
+  return async (options: CopilotChatOptions, docId: string, title: string) => {
+    const notFound = toolError(
+      'Doc Meta Update Failed',
+      `Doc with id ${docId} not found.`
+    );
+
+    if (!options?.user || !options.workspace) {
+      return notFound;
+    }
+
+    const canAccess = await ac
+      .user(options.user)
+      .workspace(options.workspace)
+      .doc(docId)
+      .can('Doc.Update');
+
+    if (!canAccess) {
+      return notFound;
+    }
+
+    const sanitizedTitle = sanitizeTitle(title);
+    if (!sanitizedTitle) {
+      return toolError('Doc Meta Update Failed', 'Title cannot be empty');
+    }
+
+    await writer.updateDocMeta(
+      options.workspace,
+      docId,
+      { title: sanitizedTitle },
+      options.user
+    );
+
+    return {
+      success: true,
+      docId,
+      message: 'Document title updated successfully',
+    };
+  };
+};
+
+export const createDocCreateTool = (
+  createDoc: (title: string, content: string) => Promise<object>
+) => {
+  return tool({
+    description:
+      'Create a new document in the workspace with the given title and markdown content. Returns the ID of the created document. This tool not support insert or update database block and image yet.',
+    inputSchema: z.object({
+      title: z.string().min(1).describe('The title of the new document'),
+      content: z
+        .string()
+        .describe('The markdown content for the document body'),
+    }),
+    execute: async ({ title, content }) => {
+      try {
+        return await createDoc(title, content);
+      } catch (err: any) {
+        logger.error(`Failed to create document: ${title}`, err);
+        return toolError('Doc Create Failed', err.message);
+      }
+    },
+  });
+};
+
+export const createDocUpdateTool = (
+  updateDoc: (docId: string, content: string) => Promise<object>
+) => {
+  return tool({
+    description:
+      'Update an existing document with new markdown content (body only). Uses structural diffing to apply minimal changes. This does NOT update the document title. This tool not support insert or update database block and image yet.',
+    inputSchema: z.object({
+      doc_id: z.string().describe('The ID of the document to update'),
+      content: z
+        .string()
+        .describe(
+          'The complete new markdown content for the document body (do NOT include a title H1)'
+        ),
+    }),
+    execute: async ({ doc_id, content }) => {
+      try {
+        return await updateDoc(doc_id, content);
+      } catch (err: any) {
+        logger.error(`Failed to update document: ${doc_id}`, err);
+        return toolError('Doc Update Failed', err.message);
+      }
+    },
+  });
+};
+
+export const createDocUpdateMetaTool = (
+  updateDocMeta: (docId: string, title: string) => Promise<object>
+) => {
+  return tool({
+    description: 'Update document metadata (currently title only).',
+    inputSchema: z.object({
+      doc_id: z.string().describe('The ID of the document to update'),
+      title: z.string().min(1).describe('The new document title'),
+    }),
+    execute: async ({ doc_id, title }) => {
+      try {
+        return await updateDocMeta(doc_id, title);
+      } catch (err: any) {
+        logger.error(`Failed to update document meta: ${doc_id}`, err);
+        return toolError('Doc Meta Update Failed', err.message);
+      }
+    },
+  });
+};

packages/backend/server/src/plugins/copilot/tools/index.ts

@@ -8,6 +8,11 @@ import { createDocEditTool } from './doc-edit';
 import { createDocKeywordSearchTool } from './doc-keyword-search';
 import { createDocReadTool } from './doc-read';
 import { createDocSemanticSearchTool } from './doc-semantic-search';
+import {
+  createDocCreateTool,
+  createDocUpdateMetaTool,
+  createDocUpdateTool,
+} from './doc-write';
 import { createExaCrawlTool } from './exa-crawl';
 import { createExaSearchTool } from './exa-search';
 import { createSectionEditTool } from './section-edit';
@@ -20,6 +25,9 @@ export interface CustomAITools extends ToolSet {
   doc_semantic_search: ReturnType<typeof createDocSemanticSearchTool>;
   doc_keyword_search: ReturnType<typeof createDocKeywordSearchTool>;
   doc_read: ReturnType<typeof createDocReadTool>;
+  doc_create: ReturnType<typeof createDocCreateTool>;
+  doc_update: ReturnType<typeof createDocUpdateTool>;
+  doc_update_meta: ReturnType<typeof createDocUpdateMetaTool>;
   doc_compose: ReturnType<typeof createDocComposeTool>;
   section_edit: ReturnType<typeof createSectionEditTool>;
   web_search_exa: ReturnType<typeof createExaSearchTool>;
@@ -34,6 +42,7 @@ export * from './doc-edit';
 export * from './doc-keyword-search';
 export * from './doc-read';
 export * from './doc-semantic-search';
+export * from './doc-write';
 export * from './error';
 export * from './exa-crawl';
 export * from './exa-search';