feat(server): refactor mcp (#14579)

#### PR Dependency Tree


* **PR #14579** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Full JSON-RPC MCP endpoint with batch requests, per-message
validation, method dispatch (initialize, ping, tools/list, tools/call)
and request cancellation
* Tool listing and execution with input validation, standardized
results, and improved error responses

* **Chores**
  * Removed an external protocol dependency
  * Bumped MCP server version to 1.0.1
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/backend/server/package.json

@@ -34,7 +34,6 @@
     "@fal-ai/serverless-client": "^0.15.0",
     "@google-cloud/opentelemetry-cloud-trace-exporter": "^3.0.0",
     "@google-cloud/opentelemetry-resource-util": "^3.0.0",
-    "@modelcontextprotocol/sdk": "^1.26.0",
     "@nestjs-cls/transactional": "^2.7.0",
     "@nestjs-cls/transactional-adapter-prisma": "^1.2.24",
     "@nestjs/apollo": "^13.0.4",

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

@@ -1,4 +1,3 @@
-import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import {
   Controller,
   Delete,
@@ -13,28 +12,51 @@ import {
 } from '@nestjs/common';
 import type { Request, Response } from 'express';
 
+import { Throttle } from '../../../base';
 import { CurrentUser } from '../../../core/auth';
-import { WorkspaceMcpProvider } from './provider';
+import { WorkspaceMcpProvider, type WorkspaceMcpServer } from './provider';
+
+type JsonRpcId = string | number | null;
+
+type JsonRpcErrorResponse = {
+  jsonrpc: '2.0';
+  error: { code: number; message: string };
+  id: JsonRpcId;
+};
+
+type JsonRpcSuccessResponse = {
+  jsonrpc: '2.0';
+  result: Record<string, unknown>;
+  id: JsonRpcId;
+};
+
+type JsonRpcResponse = JsonRpcErrorResponse | JsonRpcSuccessResponse;
+
+const JSON_RPC_VERSION = '2.0';
+const MAX_BATCH_SIZE = 20;
+const DEFAULT_PROTOCOL_VERSION = '2025-03-26';
+const SUPPORTED_PROTOCOL_VERSIONS = new Set([
+  '2025-11-25',
+  '2025-06-18',
+  '2025-03-26',
+  '2024-11-05',
+  '2024-10-07',
+]);
 
 @Controller('/api/workspaces/:workspaceId/mcp')
 export class WorkspaceMcpController {
   private readonly logger = new Logger(WorkspaceMcpController.name);
+
   constructor(private readonly provider: WorkspaceMcpProvider) {}
 
   @Get('/')
   @Delete('/')
   @HttpCode(HttpStatus.METHOD_NOT_ALLOWED)
   async STATELESS_MCP_ENDPOINT() {
-    return {
+    return this.errorResponse(null, -32000, 'Method not allowed.');
-      jsonrpc: '2.0',
-      error: {
-        code: -32000,
-        message: 'Method not allowed.',
-      },
-      id: null,
-    };
   }
 
+  @Throttle('default')
   @Post('/')
   async mcp(
     @Req() req: Request,
@@ -42,28 +64,202 @@ export class WorkspaceMcpController {
     @CurrentUser() user: CurrentUser,
     @Param('workspaceId') workspaceId: string
   ) {
-    let server = await this.provider.for(user.id, workspaceId);
+    const abortController = new AbortController();
-
+    req.on('close', () => abortController.abort());
-    const transport: StreamableHTTPServerTransport =
-      new StreamableHTTPServerTransport({
-        sessionIdGenerator: undefined,
-      });
-
-    const cleanup = () => {
-      transport.close().catch(e => {
-        this.logger.error('Failed to close MCP transport', e);
-      });
-      server.close().catch(e => {
-        this.logger.error('Failed to close MCP server', e);
-      });
-    };
 
     try {
-      res.on('close', cleanup);
+      const server = await this.provider.for(user.id, workspaceId);
-      await server.connect(transport);
+      const body = req.body as unknown;
-      await transport.handleRequest(req, res, req.body);
+      const isBatch = Array.isArray(body);
-    } catch {
+      const messages = isBatch ? body : [body];
-      cleanup();
+
+      if (!messages.length) {
+        res
+          .status(HttpStatus.BAD_REQUEST)
+          .json(this.errorResponse(null, -32600, 'Invalid Request'));
+        return;
+      }
+      if (messages.length > MAX_BATCH_SIZE) {
+        res
+          .status(HttpStatus.BAD_REQUEST)
+          .json(
+            this.errorResponse(
+              null,
+              -32600,
+              `Batch size exceeds limit (${MAX_BATCH_SIZE}).`
+            )
+          );
+        return;
+      }
+
+      const responses: JsonRpcResponse[] = [];
+      for (const message of messages) {
+        const response = await this.handleMessage(
+          message,
+          server,
+          abortController.signal
+        );
+        if (response) {
+          responses.push(response);
+        }
+      }
+
+      if (!responses.length) {
+        res.status(HttpStatus.ACCEPTED).send();
+        return;
+      }
+
+      res.status(HttpStatus.OK).json(isBatch ? responses : responses[0]);
+    } catch (error) {
+      this.logger.error('Failed to handle MCP request', error);
+      res
+        .status(HttpStatus.INTERNAL_SERVER_ERROR)
+        .json(this.errorResponse(null, -32603, 'Internal error'));
     }
   }
+
+  private async handleMessage(
+    message: unknown,
+    server: WorkspaceMcpServer,
+    signal: AbortSignal
+  ): Promise<JsonRpcResponse | null> {
+    const rawRequest = this.asObject(message);
+    if (!rawRequest || rawRequest.jsonrpc !== JSON_RPC_VERSION) {
+      return this.errorResponse(null, -32600, 'Invalid Request');
+    }
+
+    const method = rawRequest.method;
+    if (typeof method !== 'string') {
+      return this.errorResponse(null, -32600, 'Invalid Request');
+    }
+
+    const id = this.parseRequestId(rawRequest.id);
+    if (id === 'invalid') {
+      return this.errorResponse(null, -32600, 'Invalid Request');
+    }
+
+    const isNotification = id === undefined;
+    const responseId = id ?? null;
+
+    switch (method) {
+      case 'initialize': {
+        const params = this.asObject(rawRequest.params);
+        const requestedVersion =
+          params && typeof params.protocolVersion === 'string'
+            ? params.protocolVersion
+            : DEFAULT_PROTOCOL_VERSION;
+        const protocolVersion = SUPPORTED_PROTOCOL_VERSIONS.has(
+          requestedVersion
+        )
+          ? requestedVersion
+          : DEFAULT_PROTOCOL_VERSION;
+
+        if (isNotification) return null;
+
+        return this.successResponse(responseId, {
+          protocolVersion,
+          capabilities: { tools: {} },
+          serverInfo: { name: server.name, version: server.version },
+        });
+      }
+
+      case 'notifications/initialized':
+      case 'ping': {
+        if (isNotification) {
+          return null;
+        }
+        return this.successResponse(responseId, {});
+      }
+
+      case 'tools/list': {
+        if (isNotification) {
+          return null;
+        }
+        return this.successResponse(responseId, {
+          tools: server.tools.map(tool => ({
+            name: tool.name,
+            title: tool.title,
+            description: tool.description,
+            inputSchema: tool.inputSchema,
+          })),
+        });
+      }
+
+      case 'tools/call': {
+        const params = this.asObject(rawRequest.params);
+        if (!params || typeof params.name !== 'string') {
+          return this.errorResponse(responseId, -32602, 'Invalid params');
+        }
+
+        const tool = server.tools.find(item => item.name === params.name);
+        if (!tool) {
+          return this.errorResponse(
+            responseId,
+            -32602,
+            `Tool not found: ${params.name}`
+          );
+        }
+
+        const args = this.asObject(params.arguments) ?? {};
+        try {
+          const result = await tool.execute(args, { signal });
+          if (isNotification) return null;
+
+          return this.successResponse(
+            responseId,
+            result as Record<string, unknown>
+          );
+        } catch (error) {
+          this.logger.error(
+            `Error executing tool in mcp ${tool.name}`,
+            error instanceof Error ? error.stack : String(error)
+          );
+          return this.errorResponse(
+            responseId,
+            -32001,
+            `Error executing tool: ${error instanceof Error ? error.message : String(error)}`
+          );
+        }
+      }
+
+      default: {
+        if (isNotification) return null;
+        return this.errorResponse(responseId, -32601, 'Method not found');
+      }
+    }
+  }
+
+  private asObject(value: unknown): Record<string, unknown> | null {
+    if (!value || typeof value !== 'object' || Array.isArray(value)) {
+      return null;
+    }
+    return value as Record<string, unknown>;
+  }
+
+  private parseRequestId(value: unknown): JsonRpcId | undefined | 'invalid' {
+    if (value === undefined) return undefined;
+    if (
+      value === null ||
+      typeof value === 'string' ||
+      typeof value === 'number'
+    ) {
+      return value;
+    }
+    return 'invalid';
+  }
+
+  private successResponse(
+    id: JsonRpcId,
+    result: Record<string, unknown>
+  ): JsonRpcSuccessResponse {
+    return { jsonrpc: JSON_RPC_VERSION, result, id };
+  }
+
+  private errorResponse(
+    id: JsonRpcId,
+    code: number,
+    message: string
+  ): JsonRpcErrorResponse {
+    return { jsonrpc: JSON_RPC_VERSION, error: { code, message }, id };
+  }
 }

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

@@ -1,5 +1,3 @@
-import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
-import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import { Injectable } from '@nestjs/common';
 import { pick } from 'lodash-es';
 import z from 'zod/v3';
@@ -10,6 +8,94 @@ import { clearEmbeddingChunk } from '../../../models';
 import { IndexerService } from '../../indexer';
 import { CopilotContextService } from '../context/service';
 
+type McpTextContent = {
+  type: 'text';
+  text: string;
+};
+
+export type WorkspaceMcpToolResult = {
+  content: McpTextContent[];
+  isError?: boolean;
+};
+
+export type WorkspaceMcpToolDefinition = {
+  name: string;
+  title: string;
+  description: string;
+  inputSchema: Record<string, unknown>;
+  execute: (
+    args: Record<string, unknown>,
+    options: { signal: AbortSignal }
+  ) => Promise<WorkspaceMcpToolResult>;
+};
+
+export type WorkspaceMcpServer = {
+  name: string;
+  version: string;
+  tools: WorkspaceMcpToolDefinition[];
+};
+
+type ToolExecutorInput<T extends z.ZodTypeAny> = {
+  name: string;
+  title: string;
+  description: string;
+  parser: T;
+  inputSchema: Record<string, unknown>;
+  execute: (
+    args: z.infer<T>,
+    options: { signal: AbortSignal }
+  ) => Promise<WorkspaceMcpToolResult>;
+};
+
+function toolText(text: string): WorkspaceMcpToolResult {
+  return {
+    content: [{ type: 'text', text }],
+  };
+}
+
+function toolError(message: string): WorkspaceMcpToolResult {
+  return {
+    isError: true,
+    content: [{ type: 'text', text: message }],
+  };
+}
+
+function toInputError(error: z.ZodError) {
+  const details = error.issues
+    .map(issue => {
+      const path = issue.path.join('.');
+      return path ? `${path}: ${issue.message}` : issue.message;
+    })
+    .join('; ');
+  return toolError(`Invalid arguments: ${details || 'Invalid input'}`);
+}
+
+function abortIfNeeded(
+  signal: AbortSignal
+): WorkspaceMcpToolResult | undefined {
+  if (signal.aborted) return toolError('Request aborted.');
+  return;
+}
+
+function defineTool<T extends z.ZodTypeAny>(
+  config: ToolExecutorInput<T>
+): WorkspaceMcpToolDefinition {
+  return {
+    name: config.name,
+    title: config.title,
+    description: config.description,
+    inputSchema: config.inputSchema,
+    execute: async (args, options) => {
+      const aborted = abortIfNeeded(options.signal);
+      if (aborted) return aborted;
+
+      const parsed = config.parser.safeParse(args ?? {});
+      if (!parsed.success) return toInputError(parsed.error);
+      return await config.execute(parsed.data, options);
+    },
+  };
+}
+
 @Injectable()
 export class WorkspaceMcpProvider {
   constructor(
@@ -20,190 +106,182 @@ export class WorkspaceMcpProvider {
     private readonly indexer: IndexerService
   ) {}
 
-  async for(userId: string, workspaceId: string) {
+  async for(userId: string, workspaceId: string): Promise<WorkspaceMcpServer> {
     await this.ac.user(userId).workspace(workspaceId).assert('Workspace.Read');
 
-    const server = new McpServer({
+    const readDocument = defineTool({
-      name: `AFFiNE MCP Server for Workspace ${workspaceId}`,
+      name: 'read_document',
-      version: '1.0.0',
+      title: 'Read Document',
-    });
+      description: 'Read a document with given ID',
-
+      parser: z.object({ docId: z.string() }),
-    server.registerTool(
+      inputSchema: {
-      'read_document',
+        type: 'object',
-      {
+        properties: {
-        title: 'Read Document',
+          docId: { type: 'string' },
-        description: 'Read a document with given ID',
+        },
-        inputSchema: z.object({
+        required: ['docId'],
-          docId: z.string(),
+        additionalProperties: false,
-        }),
       },
-      async ({ docId }) => {
+      execute: async ({ docId }, options) => {
-        const notFoundError: CallToolResult = {
+        const notFoundError = toolError(`Doc with id ${docId} not found.`);
-          isError: true,
-          content: [
-            {
-              type: 'text',
-              text: `Doc with id ${docId} not found.`,
-            },
-          ],
-        };
 
         const accessible = await this.ac
           .user(userId)
           .workspace(workspaceId)
           .doc(docId)
           .can('Doc.Read');
+        if (!accessible) return notFoundError;
 
-        if (!accessible) {
+        const abortedAfterPermission = abortIfNeeded(options.signal);
-          return notFoundError;
+        if (abortedAfterPermission) return abortedAfterPermission;
-        }
 
         const content = await this.reader.getDocMarkdown(
           workspaceId,
           docId,
           false
         );
+        if (!content) return notFoundError;
 
-        if (!content) {
+        const abortedAfterRead = abortIfNeeded(options.signal);
-          return notFoundError;
+        if (abortedAfterRead) return abortedAfterRead;
-        }
 
-        return {
+        return toolText(content.markdown);
-          content: [
-            {
-              type: 'text',
-              text: content.markdown,
-            },
-          ],
-        } as const;
-      }
-    );
-
-    server.registerTool(
-      'semantic_search',
-      {
-        title: 'Semantic Search',
-        description:
-          'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts, recent documents).',
-        inputSchema: z.object({
-          query: z.string(),
-        }),
       },
-      async ({ query }, req) => {
+    });
-        query = query.trim();
+
-        if (!query) {
+    const semanticSearch = defineTool({
-          return {
+      name: 'semantic_search',
-            isError: true,
+      title: 'Semantic Search',
-            content: [
+      description:
-              {
+        'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts, recent documents).',
-                type: 'text',
+      parser: z.object({ query: z.string() }),
-                text: 'Query is required for semantic search.',
+      inputSchema: {
-              },
+        type: 'object',
-            ],
+        properties: {
-          };
+          query: { type: 'string' },
+        },
+        required: ['query'],
+        additionalProperties: false,
+      },
+      execute: async ({ query }, options) => {
+        const trimmed = query.trim();
+        if (!trimmed) {
+          return toolError('Query is required for semantic search.');
         }
 
         const chunks = await this.context.matchWorkspaceDocs(
           workspaceId,
-          query,
+          trimmed,
           5,
-          req.signal
+          options.signal
         );
 
+        const abortedAfterMatch = abortIfNeeded(options.signal);
+        if (abortedAfterMatch) return abortedAfterMatch;
+
         const docs = await this.ac
           .user(userId)
           .workspace(workspaceId)
           .docs(
-            chunks.filter(c => 'docId' in c),
+            chunks.filter(chunk => 'docId' in chunk),
             'Doc.Read'
           );
 
+        const abortedAfterDocs = abortIfNeeded(options.signal);
+        if (abortedAfterDocs) return abortedAfterDocs;
+
         return {
           content: docs.map(doc => ({
             type: 'text',
             text: clearEmbeddingChunk(doc).content,
           })),
-        } as const;
+        };
-      }
-    );
-
-    server.registerTool(
-      'keyword_search',
-      {
-        title: 'Keyword Search',
-        description:
-          'Fuzzy search all workspace documents for the exact keyword or phrase supplied and return passages ranked by textual match. Use this tool by default whenever a straightforward term-based or keyword-base lookup is sufficient.',
-        inputSchema: z.object({
-          query: z.string(),
-        }),
       },
-      async ({ query }) => {
+    });
-        query = query.trim();
+
-        if (!query) {
+    const keywordSearch = defineTool({
-          return {
+      name: 'keyword_search',
-            isError: true,
+      title: 'Keyword Search',
-            content: [
+      description:
-              {
+        'Fuzzy search all workspace documents for the exact keyword or phrase supplied and return passages ranked by textual match. Use this tool by default whenever a straightforward term-based or keyword-base lookup is sufficient.',
-                type: 'text',
+      parser: z.object({ query: z.string() }),
-                text: 'Query is required for keyword search.',
+      inputSchema: {
-              },
+        type: 'object',
-            ],
+        properties: {
-          };
+          query: { type: 'string' },
-        }
+        },
+        required: ['query'],
+        additionalProperties: false,
+      },
+      execute: async ({ query }, options) => {
+        const trimmed = query.trim();
+        if (!trimmed) return toolError('Query is required for keyword search.');
+
+        let docs = await this.indexer.searchDocsByKeyword(workspaceId, trimmed);
+
+        const abortedAfterSearch = abortIfNeeded(options.signal);
+        if (abortedAfterSearch) return abortedAfterSearch;
 
-        let docs = await this.indexer.searchDocsByKeyword(workspaceId, query);
         docs = await this.ac
           .user(userId)
           .workspace(workspaceId)
           .docs(docs, 'Doc.Read');
 
+        const abortedAfterDocs = abortIfNeeded(options.signal);
+        if (abortedAfterDocs) return abortedAfterDocs;
+
         return {
           content: docs.map(doc => ({
             type: 'text',
             text: JSON.stringify(pick(doc, 'docId', 'title', 'createdAt')),
           })),
-        } as const;
+        };
-      }
+      },
-    );
+    });
+
+    const tools = [readDocument, semanticSearch, keywordSearch];
 
     if (env.dev || env.namespaces.canary) {
-      // Write tools - create and update documents
+      const createDocument = defineTool({
-      server.registerTool(
+        name: 'create_document',
-        'create_document',
+        title: 'Create Document',
-        {
+        description:
-          title: 'Create Document',
+          '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.',
-          description:
+        parser: z.object({
-            '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.',
+          title: z.string().min(1),
-          inputSchema: z.object({
+          content: z.string(),
-            title: z.string().min(1).describe('The title of the new document'),
+        }),
-            content: z
+        inputSchema: {
-              .string()
+          type: 'object',
-              .describe('The markdown content for the document body'),
+          properties: {
-          }),
+            title: {
+              type: 'string',
+              description: 'The title of the new document',
+            },
+            content: {
+              type: 'string',
+              description: 'The markdown content for the document body',
+            },
+          },
+          required: ['title', 'content'],
+          additionalProperties: false,
         },
-        async ({ title, content }) => {
+        execute: async ({ title, content }, options) => {
           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 abortedAfterPermission = abortIfNeeded(options.signal);
-            const sanitizedTitle = title.replace(/[\r\n]+/g, ' ').trim();
+            if (abortedAfterPermission) return abortedAfterPermission;
-            if (!sanitizedTitle) {
-              throw new Error('Title cannot be empty');
-            }
 
-            // Strip any leading H1 from content to prevent duplicates
+            const sanitizedTitle = title.replace(/[\r\n]+/g, ' ').trim();
-            // Per CommonMark spec, ATX headings allow only 0-3 spaces before the #
+            if (!sanitizedTitle) throw new Error('Title cannot be empty');
-            // 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,
@@ -211,173 +289,145 @@ export class WorkspaceMcpProvider {
               userId
             );
 
-            return {
+            return toolText(
-              content: [
+              JSON.stringify({
-                {
+                success: true,
-                  type: 'text',
+                docId: result.docId,
-                  text: JSON.stringify({
+                message: `Document "${title}" created successfully`,
-                    success: true,
+              })
-                    docId: result.docId,
+            );
-                    message: `Document "${title}" created successfully`,
-                  }),
-                },
-              ],
-            } as const;
           } catch (error) {
-            return {
+            return toolError(
-              isError: true,
+              `Failed to create document: ${error instanceof Error ? error.message : 'Unknown error'}`
-              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,
+      const updateDocument = defineTool({
-            content: [
+        name: 'update_document',
-              {
+        title: 'Update Document',
-                type: 'text',
+        description:
-                text: `Doc with id ${docId} not found.`,
+          '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.',
-              },
+        parser: z.object({
-            ],
+          docId: z.string(),
-          };
+          content: z.string(),
+        }),
+        inputSchema: {
+          type: 'object',
+          properties: {
+            docId: {
+              type: 'string',
+              description: 'The ID of the document to update',
+            },
+            content: {
+              type: 'string',
+              description:
+                'The complete new markdown content for the document body (do NOT include a title H1)',
+            },
+          },
+          required: ['docId', 'content'],
+          additionalProperties: false,
+        },
+        execute: async ({ docId, content }, options) => {
+          const notFoundError = toolError(`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');
+          if (!accessible) return notFoundError;
 
-          if (!accessible) {
+          const abortedBeforeWrite = abortIfNeeded(options.signal);
-            return notFoundError;
+          if (abortedBeforeWrite) return abortedBeforeWrite;
-          }
 
           try {
-            // Update the document
             await this.writer.updateDoc(workspaceId, docId, content, userId);
-
+            return toolText(
-            return {
+              JSON.stringify({
-              content: [
+                success: true,
-                {
+                docId,
-                  type: 'text',
+                message: 'Document updated successfully',
-                  text: JSON.stringify({
+              })
-                    success: true,
+            );
-                    docId,
-                    message: `Document updated successfully`,
-                  }),
-                },
-              ],
-            } as const;
           } catch (error) {
-            return {
+            return toolError(
-              isError: true,
+              `Failed to update document: ${error instanceof Error ? error.message : 'Unknown error'}`
-              content: [
+            );
-                {
-                  type: 'text',
-                  text: `Failed to update document: ${error instanceof Error ? error.message : 'Unknown error'}`,
-                },
-              ],
-            };
           }
-        }
-      );
-
-      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,
+      const updateDocumentMeta = defineTool({
-            content: [
+        name: 'update_document_meta',
-              {
+        title: 'Update Document Metadata',
-                type: 'text',
+        description: 'Update document metadata (currently title only).',
-                text: `Doc with id ${docId} not found.`,
+        parser: z.object({
-              },
+          docId: z.string(),
-            ],
+          title: z.string().min(1),
-          };
+        }),
+        inputSchema: {
+          type: 'object',
+          properties: {
+            docId: {
+              type: 'string',
+              description: 'The ID of the document to update',
+            },
+            title: {
+              type: 'string',
+              description: 'The new document title',
+            },
+          },
+          required: ['docId', 'title'],
+          additionalProperties: false,
+        },
+        execute: async ({ docId, title }, options) => {
+          const notFoundError = toolError(`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');
+          if (!accessible) return notFoundError;
 
-          if (!accessible) {
+          const abortedAfterPermission = abortIfNeeded(options.signal);
-            return notFoundError;
+          if (abortedAfterPermission) return abortedAfterPermission;
-          }
 
           try {
             const sanitizedTitle = title.replace(/[\r\n]+/g, ' ').trim();
-            if (!sanitizedTitle) {
+            if (!sanitizedTitle) throw new Error('Title cannot be empty');
-              throw new Error('Title cannot be empty');
-            }
 
             await this.writer.updateDocMeta(
               workspaceId,
               docId,
-              {
+              { title: sanitizedTitle },
-                title: sanitizedTitle,
-              },
               userId
             );
 
-            return {
+            return toolText(
-              content: [
+              JSON.stringify({
-                {
+                success: true,
-                  type: 'text',
+                docId,
-                  text: JSON.stringify({
+                message: 'Document title updated successfully',
-                    success: true,
+              })
-                    docId,
+            );
-                    message: `Document title updated successfully`,
-                  }),
-                },
-              ],
-            } as const;
           } catch (error) {
-            return {
+            return toolError(
-              isError: true,
+              `Failed to update document metadata: ${error instanceof Error ? error.message : 'Unknown error'}`
-              content: [
+            );
-                {
-                  type: 'text',
-                  text: `Failed to update document metadata: ${error instanceof Error ? error.message : 'Unknown error'}`,
-                },
-              ],
-            };
           }
-        }
+        },
-      );
+      });
+
+      tools.push(createDocument, updateDocument, updateDocumentMeta);
     }
 
-    return server;
+    return {
+      name: `AFFiNE MCP Server for Workspace ${workspaceId}`,
+      version: '1.0.1',
+      tools,
+    };
   }
 }

yarn.lock

@@ -969,7 +969,6 @@ __metadata:
     "@fal-ai/serverless-client": "npm:^0.15.0"
     "@google-cloud/opentelemetry-cloud-trace-exporter": "npm:^3.0.0"
     "@google-cloud/opentelemetry-resource-util": "npm:^3.0.0"
-    "@modelcontextprotocol/sdk": "npm:^1.26.0"
     "@nestjs-cls/transactional": "npm:^2.7.0"
     "@nestjs-cls/transactional-adapter-prisma": "npm:^1.2.24"
     "@nestjs/apollo": "npm:^13.0.4"
@@ -6413,15 +6412,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@hono/node-server@npm:^1.19.9":
-  version: 1.19.9
-  resolution: "@hono/node-server@npm:1.19.9"
-  peerDependencies:
-    hono: ^4
-  checksum: 10/d4915c2e736ee1e3934b5538cde92b19914dc71346340528a04e4c7219afc7367965080cd1a5291ac9cbda7b0780b89b6ca93472a9418aa105d6d1183033dc8a
-  languageName: node
-  linkType: hard
-
 "@html-validate/stylish@npm:^4.1.0":
   version: 4.2.0
   resolution: "@html-validate/stylish@npm:4.2.0"
@@ -7632,39 +7622,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@modelcontextprotocol/sdk@npm:^1.26.0":
-  version: 1.26.0
-  resolution: "@modelcontextprotocol/sdk@npm:1.26.0"
-  dependencies:
-    "@hono/node-server": "npm:^1.19.9"
-    ajv: "npm:^8.17.1"
-    ajv-formats: "npm:^3.0.1"
-    content-type: "npm:^1.0.5"
-    cors: "npm:^2.8.5"
-    cross-spawn: "npm:^7.0.5"
-    eventsource: "npm:^3.0.2"
-    eventsource-parser: "npm:^3.0.0"
-    express: "npm:^5.2.1"
-    express-rate-limit: "npm:^8.2.1"
-    hono: "npm:^4.11.4"
-    jose: "npm:^6.1.3"
-    json-schema-typed: "npm:^8.0.2"
-    pkce-challenge: "npm:^5.0.0"
-    raw-body: "npm:^3.0.0"
-    zod: "npm:^3.25 || ^4.0"
-    zod-to-json-schema: "npm:^3.25.1"
-  peerDependencies:
-    "@cfworker/json-schema": ^4.1.1
-    zod: ^3.25 || ^4.0
-  peerDependenciesMeta:
-    "@cfworker/json-schema":
-      optional: true
-    zod:
-      optional: false
-  checksum: 10/a206b2a4d61a23be8b8f4c886528dd9348d11b17ce36013b350edf5c082b1c1f07941d52ea098f721daf3828085b6f6276bb844c484a0e9913edbc028517a3d5
-  languageName: node
-  linkType: hard
-
 "@module-federation/error-codes@npm:0.22.0":
   version: 0.22.0
   resolution: "@module-federation/error-codes@npm:0.22.0"
@@ -18306,20 +18263,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"ajv-formats@npm:^3.0.1":
-  version: 3.0.1
-  resolution: "ajv-formats@npm:3.0.1"
-  dependencies:
-    ajv: "npm:^8.0.0"
-  peerDependencies:
-    ajv: ^8.0.0
-  peerDependenciesMeta:
-    ajv:
-      optional: true
-  checksum: 10/5679b9f9ced9d0213a202a37f3aa91efcffe59a6de1a6e3da5c873344d3c161820a1f11cc29899661fee36271fd2895dd3851b6461c902a752ad661d1c1e8722
-  languageName: node
-  linkType: hard
-
 "ajv-keywords@npm:^3.4.1":
   version: 3.5.2
   resolution: "ajv-keywords@npm:3.5.2"
@@ -18364,7 +18307,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"ajv@npm:^8.0.0, ajv@npm:^8.11.0, ajv@npm:^8.17.1, ajv@npm:^8.9.0":
+"ajv@npm:^8.0.0, ajv@npm:^8.11.0, ajv@npm:^8.9.0":
   version: 8.18.0
   resolution: "ajv@npm:8.18.0"
   dependencies:
@@ -20906,7 +20849,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"cross-spawn@npm:^7.0.1, cross-spawn@npm:^7.0.3, cross-spawn@npm:^7.0.5, cross-spawn@npm:^7.0.6":
+"cross-spawn@npm:^7.0.1, cross-spawn@npm:^7.0.3, cross-spawn@npm:^7.0.6":
   version: 7.0.6
   resolution: "cross-spawn@npm:7.0.6"
   dependencies:
@@ -23260,22 +23203,13 @@ __metadata:
   languageName: node
   linkType: hard
 
-"eventsource-parser@npm:^3.0.0, eventsource-parser@npm:^3.0.1, eventsource-parser@npm:^3.0.6":
+"eventsource-parser@npm:^3.0.6":
   version: 3.0.6
   resolution: "eventsource-parser@npm:3.0.6"
   checksum: 10/febf7058b9c2168ecbb33e92711a1646e06bd1568f60b6eb6a01a8bf9f8fcd29cc8320d57247059cacf657a296280159f21306d2e3ff33309a9552b2ef889387
   languageName: node
   linkType: hard
 
-"eventsource@npm:^3.0.2":
-  version: 3.0.7
-  resolution: "eventsource@npm:3.0.7"
-  dependencies:
-    eventsource-parser: "npm:^3.0.1"
-  checksum: 10/e034915bc97068d1d38617951afd798e6776d6a3a78e36a7569c235b177c7afc2625c9fe82656f7341ab72c7eeecb3fd507b7f88e9328f2448872ff9c4742bb6
-  languageName: node
-  linkType: hard
-
 "exa-js@npm:^2.4.0":
   version: 2.4.0
   resolution: "exa-js@npm:2.4.0"
@@ -23364,18 +23298,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"express-rate-limit@npm:^8.2.1":
+"express@npm:5.2.1, express@npm:^5.0.0, express@npm:^5.0.1":
-  version: 8.2.1
-  resolution: "express-rate-limit@npm:8.2.1"
-  dependencies:
-    ip-address: "npm:10.0.1"
-  peerDependencies:
-    express: ">= 4.11"
-  checksum: 10/7cbf70df2e88e590e463d2d8f93380775b2ea181d97f2c50c2ff9f2c666c247f83109a852b21d9c99ccc5762119101f281f54a27252a2f1a0a918be6d71f955b
-  languageName: node
-  linkType: hard
-
-"express@npm:5.2.1, express@npm:^5.0.0, express@npm:^5.0.1, express@npm:^5.2.1":
   version: 5.2.1
   resolution: "express@npm:5.2.1"
   dependencies:
@@ -25067,13 +24990,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"hono@npm:^4.11.4":
-  version: 4.12.0
-  resolution: "hono@npm:4.12.0"
-  checksum: 10/8a4de1ac4394816cbdbc87ef813ce3fb767953dff4cd88fa519bec66df8d7d801dcfddea815677440c5b0fbed636ee7fbdab834f589154550d83ea4de713e769
-  languageName: node
-  linkType: hard
-
 "hosted-git-info@npm:^2.1.4":
   version: 2.8.9
   resolution: "hosted-git-info@npm:2.8.9"
@@ -25833,13 +25749,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"ip-address@npm:10.0.1":
-  version: 10.0.1
-  resolution: "ip-address@npm:10.0.1"
-  checksum: 10/09731acda32cd8e14c46830c137e7e5940f47b36d63ffb87c737331270287d631cf25aa95570907a67d3f919fdb25f4470c404eda21e62f22e0a55927f4dd0fb
-  languageName: node
-  linkType: hard
-
 "ip-address@npm:^9.0.5":
   version: 9.0.5
   resolution: "ip-address@npm:9.0.5"
@@ -26641,13 +26550,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"json-schema-typed@npm:^8.0.2":
-  version: 8.0.2
-  resolution: "json-schema-typed@npm:8.0.2"
-  checksum: 10/fa866d1fe91e3a94aa4fe007861475cd03dcaf47b719861cab171ef2f8598478007c634d29ae45de94ee34ddff4e13414c63ea5ff06c5b868b613142c699d511
-  languageName: node
-  linkType: hard
-
 "json-schema@npm:^0.4.0":
   version: 0.4.0
   resolution: "json-schema@npm:0.4.0"
@@ -30744,13 +30646,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"pkce-challenge@npm:^5.0.0":
-  version: 5.0.0
-  resolution: "pkce-challenge@npm:5.0.0"
-  checksum: 10/e60c06a0e0481cb82f80072053d5c479a7490758541c4226460450285dd5d72a995c44b3c553731ca7c2f64cc34b35f1d2e5f9de08d276b59899298f9efe1ddf
-  languageName: node
-  linkType: hard
-
 "pkg-types@npm:^1.2.0, pkg-types@npm:^1.3.0, pkg-types@npm:^1.3.1":
   version: 1.3.1
   resolution: "pkg-types@npm:1.3.1"
@@ -31648,7 +31543,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"raw-body@npm:^3.0.0, raw-body@npm:^3.0.1":
+"raw-body@npm:^3.0.1":
   version: 3.0.2
   resolution: "raw-body@npm:3.0.2"
   dependencies:
@@ -36901,7 +36796,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"zod-to-json-schema@npm:^3.20.0, zod-to-json-schema@npm:^3.25.1":
+"zod-to-json-schema@npm:^3.20.0":
   version: 3.25.1
   resolution: "zod-to-json-schema@npm:3.25.1"
   peerDependencies:
@@ -36926,7 +36821,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"zod@npm:^3.25 || ^4.0, zod@npm:^3.25.0 || ^4.0.0":
+"zod@npm:^3.25.0 || ^4.0.0":
   version: 4.3.6
   resolution: "zod@npm:4.3.6"
   checksum: 10/25fc0f62e01b557b4644bf0b393bbaf47542ab30877c37837ea8caf314a8713d220c7d7fe51f68ffa72f0e1018ddfa34d96f1973d23033f5a2a5a9b6b9d9da01