From c256de50aad943f142ae70bd0ff45c1a916126d6 Mon Sep 17 00:00:00 2001 From: luulam Date: Tue, 25 Aug 2026 15:54:11 +0700 Subject: [PATCH] feat(mcp): expand MCP server toolset to Notion-level parity - Ungate write tools (create/update/update_meta) from dev/canary env flag; now available in all environments with READ_WRITE credentials - Add list_documents: root-doc snapshot based listing with permission filter, pagination and title fallback parsing for unmerged docs - Add get_workspace_info: workspace name and member count - Add get_users: searchable member list (id/name/email/avatar/role) - Add get_comments / create_comment: comment roundtrip via CommentModel gated by Doc.Comments.Create permission - Enable READ_WRITE credential creation in all environments (resolver) - Bump MCP server version to 1.1.0 - Extend copilot e2e tests: full toolset assertions, doc/comment roundtrip over HTTP, read-only isolation --- .../src/__tests__/copilot/copilot.e2e.ts | 137 +++++++- .../src/plugins/copilot/mcp/provider.ts | 304 +++++++++++++++++- .../src/plugins/copilot/mcp/resolver.ts | 10 +- 3 files changed, 433 insertions(+), 18 deletions(-) diff --git a/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts index 695f0889a..c99cec842 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts @@ -14,7 +14,7 @@ import { Models } from '../../models'; import { CopilotFeatureService } from '../../plugins/copilot/feature'; import { McpCredentialService } from '../../plugins/copilot/mcp/credential'; import { WorkspaceMcpProvider } from '../../plugins/copilot/mcp/provider'; -import { installMockCopilotRuntime } from '../mocks'; +import { installMockCopilotRuntime, Mockers } from '../mocks'; import { createTestingApp, createWorkspace, type TestingApp } from '../utils'; import { chatWithImages, @@ -250,7 +250,14 @@ test('MCP credentials remain endpoint-bound through rotate, revoke and expiry', (await provider.for(user.id, target.id, McpAccessMode.READ_ONLY)).tools.map( tool => tool.name ), - ['read_document', 'doc_search'] + [ + 'read_document', + 'doc_search', + 'list_documents', + 'get_workspace_info', + 'get_users', + 'get_comments', + ] ); const rotated = await credentials.rotate( @@ -283,3 +290,129 @@ test('MCP credentials remain endpoint-bound through rotate, revoke and expiry', }); await t.throwsAsync(credentials.authenticate(disabled.token, target.id)); }); + +test('MCP server exposes full toolset with read-write credential and exercises doc/comment tools', async t => { + const { app } = t.context; + const auth = app.get(AuthService); + const credentials = app.get(McpCredentialService); + const provider = app.get(WorkspaceMcpProvider); + const user = await auth.signUp(`mcp-rw-${randomUUID()}@affine.pro`, '123456'); + const workspace = await app.create(Mockers.Workspace, { + owner: user, + snapshot: true, + }); + const issued = await credentials.create({ + userId: user.id, + workspaceId: workspace.id, + name: 'RW client', + accessMode: McpAccessMode.READ_WRITE, + expirationDays: 90, + }); + + // READ_WRITE no longer gated behind dev/canary + const toolNames = ( + await provider.for(user.id, workspace.id, McpAccessMode.READ_WRITE) + ).tools.map(tool => tool.name); + t.deepEqual(toolNames, [ + 'read_document', + 'doc_search', + 'create_document', + 'update_document', + 'update_document_meta', + 'list_documents', + 'get_workspace_info', + 'create_comment', + 'get_users', + 'get_comments', + ]); + const callTool = async (name: string, args: Record) => { + const response = await app + .POST(`/api/workspaces/${workspace.id}/mcp`) + .set('Authorization', `Bearer ${issued.token}`) + .send({ + jsonrpc: '2.0', + id: 1, + method: 'tools/call', + params: { name, arguments: args }, + }) + .expect(200); + const result = response.body.result as { + content?: { text: string }[]; + isError?: boolean; + }; + t.like( + { error: result.isError ? result.content?.[0].text : null }, + { error: null } + ); + t.truthy(result.content); + return result; + }; + + // create_document + const created = await callTool('create_document', { + title: 'MCP Test Doc', + content: '# MCP Test Doc\n\nHello from MCP.', + }); + const docId = JSON.parse(created.content![0].text).docId as string; + const listed = await callTool('list_documents', {}); + const docs = JSON.parse(listed.content![0].text).docs as { + doc_id: string; + title: string; + }[]; + t.log('list_documents result:', listed.content?.[0].text); + t.true( + docs.some(doc => doc.doc_id === docId && doc.title === 'MCP Test Doc') + ); + + // get_workspace_info + const wsInfo = await callTool('get_workspace_info', {}); + t.like(JSON.parse(wsInfo.content![0].text), { + workspace_id: workspace.id, + member_count: 1, + }); + + // get_users contains owner + const users = await callTool('get_users', {}); + const userList = JSON.parse(users.content![0].text).users as { id: string }[]; + t.true(userList.some(u => u.id === user.id)); + + // create_comment + get_comments roundtrip + const commented = await callTool('create_comment', { + docId, + content: 'A comment from MCP', + }); + const commentId = JSON.parse(commented.content![0].text).comment_id as string; + t.truthy(commentId); + + const comments = await callTool('get_comments', { docId }); + const parsed = JSON.parse(comments.content![0].text) as { + comments: { id: string; content: unknown }[]; + }; + t.is(parsed.comments.length, 1); + t.is(parsed.comments[0].id, commentId); + t.like(parsed.comments[0].content, { + type: 'paragraph', + content: [{ type: 'text', text: 'A comment from MCP' }], + }); +}); + +test('MCP read-only credential cannot create documents or comments', async t => { + const { app } = t.context; + const auth = app.get(AuthService); + const models = app.get(Models); + const provider = app.get(WorkspaceMcpProvider); + const user = await auth.signUp(`mcp-ro-${randomUUID()}@affine.pro`, '123456'); + const workspace = await models.workspace.create(user.id); + + const readOnlyTools = ( + await provider.for(user.id, workspace.id, McpAccessMode.READ_ONLY) + ).tools.map(tool => tool.name); + t.deepEqual(readOnlyTools.sort(), [ + 'doc_search', + 'get_comments', + 'get_users', + 'get_workspace_info', + 'list_documents', + 'read_document', + ]); +}); diff --git a/packages/backend/server/src/plugins/copilot/mcp/provider.ts b/packages/backend/server/src/plugins/copilot/mcp/provider.ts index b721c65fd..399225063 100644 --- a/packages/backend/server/src/plugins/copilot/mcp/provider.ts +++ b/packages/backend/server/src/plugins/copilot/mcp/provider.ts @@ -2,8 +2,11 @@ import { Injectable } from '@nestjs/common'; import { McpAccessMode } from '@prisma/client'; import z from 'zod/v3'; +import { PaginationInput } from '../../../base/graphql'; import { DocReader, DocWriter } from '../../../core/doc'; -import { PermissionAccess } from '../../../core/permission'; +import { PermissionAccess, PermissionService } from '../../../core/permission'; +import { readAllDocIdsFromWorkspaceSnapshot } from '../../../core/utils/blocksuite'; +import { Models, WorkspaceRole } from '../../../models'; import { DocumentRetrievalService } from '../retrieval/document'; type McpTextContent = { @@ -100,7 +103,9 @@ export class WorkspaceMcpProvider { private readonly ac: PermissionAccess, private readonly reader: DocReader, private readonly writer: DocWriter, - private readonly retrieval: DocumentRetrievalService + private readonly retrieval: DocumentRetrievalService, + private readonly models: Models, + private readonly permission: PermissionService ) {} async for( @@ -202,10 +207,7 @@ export class WorkspaceMcpProvider { const tools = [readDocument, docSearch]; - if ( - accessMode === McpAccessMode.READ_WRITE && - (env.dev || env.namespaces.canary) - ) { + if (accessMode === McpAccessMode.READ_WRITE) { const createDocument = defineTool({ name: 'create_document', title: 'Create Document', @@ -388,9 +390,297 @@ export class WorkspaceMcpProvider { tools.push(createDocument, updateDocument, updateDocumentMeta); } + const listDocuments = defineTool({ + name: 'list_documents', + title: 'List Documents', + description: + 'List documents in the workspace the credential owner can read, ordered by last update time (newest first). Returns doc IDs, titles and timestamps for pagination.', + parser: z.object({ + limit: z.number().int().min(1).max(100).optional(), + offset: z.number().int().min(0).optional(), + }), + inputSchema: { + type: 'object', + properties: { + limit: { type: 'integer', minimum: 1, maximum: 100 }, + offset: { type: 'integer', minimum: 0 }, + }, + additionalProperties: false, + }, + execute: async ({ limit, offset }, options) => { + await this.ac + .user(userId) + .workspace(workspaceId) + .assert('Workspace.Read'); + + const abortedAfterPermission = abortIfNeeded(options.signal); + if (abortedAfterPermission) return abortedAfterPermission; + + const pagination: PaginationInput = { + first: Math.min(limit ?? 20, 100), + offset: offset ?? 0, + }; + const rootDoc = await this.reader.getDoc(workspaceId, workspaceId); + if (!rootDoc) { + return toolText( + JSON.stringify({ total: 0, offset: pagination.offset, docs: [] }) + ); + } + const docIds = readAllDocIdsFromWorkspaceSnapshot(rootDoc.bin); + const readable = await this.permission.filterReadableDocs({ + userId, + workspaceId, + docs: docIds.map(docId => ({ docId })), + }); + const infos = ( + await Promise.all( + readable.map(async ({ docId }) => { + const info = await this.models.doc.getDocInfo(workspaceId, docId); + if (!info || !info.title) { + // Doc created but its updates have not been merged into a + // snapshot yet, so `workspace_pages` has no title. Parse the + // title from the pending yjs binary instead. + const markdown = await this.reader.getDocMarkdown( + workspaceId, + docId, + false + ); + if (!markdown) return null; + return { + ...info, + docId, + title: markdown.title, + createdAt: info?.createdAt ?? new Date(), + updatedAt: info?.updatedAt ?? new Date(), + }; + } + return info; + }) + ) + ).filter( + (info): info is NonNullable => + info !== null && !!info.title + ); + infos.sort( + (a, b) => + new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime() + ); + const page = infos.slice( + pagination.offset, + pagination.offset + pagination.first + ); + return toolText( + JSON.stringify({ + total: infos.length, + offset: pagination.offset, + docs: page.map(info => ({ + doc_id: info.docId, + title: info.title, + created_at: info.createdAt, + updated_at: info.updatedAt, + })), + }) + ); + }, + }); + + tools.push(listDocuments); + + const getWorkspaceInfo = defineTool({ + name: 'get_workspace_info', + title: 'Get Workspace Info', + description: 'Get the name and member count of the workspace.', + parser: z.object({}), + inputSchema: { + type: 'object', + properties: {}, + additionalProperties: false, + }, + execute: async (_args, options) => { + await this.ac + .user(userId) + .workspace(workspaceId) + .assert('Workspace.Read'); + + const aborted = abortIfNeeded(options.signal); + if (aborted) return aborted; + + const [workspace, memberCount] = await Promise.all([ + this.models.workspace.get(workspaceId), + this.models.workspaceUser.count(workspaceId), + ]); + if (!workspace) return toolError(`Workspace ${workspaceId} not found.`); + + return toolText( + JSON.stringify({ + workspace_id: workspaceId, + name: workspace.name, + member_count: memberCount, + }) + ); + }, + }); + + tools.push(getWorkspaceInfo); + + const getUsers = defineTool({ + name: 'get_users', + title: 'Get Workspace Users', + description: + 'List workspace members (id, name, email, avatar and role) so callers can mention or reference them.', + parser: z.object({ + query: z.string().trim().min(1).max(255).optional(), + limit: z.number().int().min(1).max(100).optional(), + offset: z.number().int().min(0).optional(), + }), + inputSchema: { + type: 'object', + properties: { + query: { type: 'string', description: 'Filter by name or email.' }, + limit: { type: 'integer', minimum: 1, maximum: 100 }, + offset: { type: 'integer', minimum: 0 }, + }, + additionalProperties: false, + }, + execute: async ({ query, limit, offset }, options) => { + await this.ac + .user(userId) + .workspace(workspaceId) + .assert('Workspace.Users.Read'); + + const aborted = abortIfNeeded(options.signal); + if (aborted) return aborted; + + const pagination: PaginationInput = { + first: Math.min(limit ?? 20, 100), + offset: offset ?? 0, + }; + const rows = query + ? await this.models.workspaceUser.search( + workspaceId, + query, + pagination + ) + : ( + await this.models.workspaceUser.paginate(workspaceId, pagination) + )[0]; + + return toolText( + JSON.stringify({ + users: rows.flatMap(row => + row.status === 'Accepted' && row.user + ? [ + { + id: row.user.id, + name: row.user.name, + email: row.user.email, + avatar_url: row.user.avatarUrl, + role: WorkspaceRole[row.type] ?? 'Unknown', + }, + ] + : [] + ), + }) + ); + }, + }); + + const getComments = defineTool({ + name: 'get_comments', + title: 'Get Document Comments', + description: + 'List comments (with replies) on a document, newest first. Requires the credential owner to be able to read comments on that document.', + parser: z.object({ + docId: z.string(), + limit: z.number().int().min(1).max(100).optional(), + }), + inputSchema: { + type: 'object', + properties: { + docId: { type: 'string', description: 'The document ID.' }, + limit: { type: 'integer', minimum: 1, maximum: 100 }, + }, + required: ['docId'], + additionalProperties: false, + }, + execute: async ({ docId, limit }, options) => { + await this.ac.user(userId).doc({ workspaceId, docId }).can('Doc.Read'); + + const aborted = abortIfNeeded(options.signal); + if (aborted) return aborted; + + const comments = await this.models.comment.list(workspaceId, docId, { + take: limit ?? 50, + }); + + return toolText(JSON.stringify({ comments })); + }, + }); + + const createComment = defineTool({ + name: 'create_comment', + title: 'Create Document Comment', + description: + 'Add a plain-text comment on a document as the credential owner. Returns the created comment ID.', + parser: z.object({ + docId: z.string(), + content: z.string().trim().min(1).max(5000), + }), + inputSchema: { + type: 'object', + properties: { + docId: { type: 'string', description: 'The document ID.' }, + content: { + type: 'string', + description: 'Plain text comment body.', + }, + }, + required: ['docId', 'content'], + additionalProperties: false, + }, + execute: async ({ docId, content }, options) => { + const accessible = await this.ac + .user(userId) + .doc({ workspaceId, docId }) + .can('Doc.Comments.Create'); + if (!accessible) return toolError(`Doc with id ${docId} not found.`); + + const aborted = abortIfNeeded(options.signal); + if (aborted) return aborted; + + try { + const comment = await this.models.comment.create({ + workspaceId, + docId, + userId, + content: { + type: 'paragraph', + content: [{ type: 'text', text: content }], + }, + }); + + return toolText( + JSON.stringify({ + success: true, + comment_id: comment.id, + }) + ); + } catch (error) { + return toolError( + `Failed to create comment: ${error instanceof Error ? error.message : 'Unknown error'}` + ); + } + }, + }); + + if (accessMode === McpAccessMode.READ_WRITE) { + tools.push(createComment); + } + tools.push(getUsers, getComments); + return { name: `AFFiNE MCP Server for Workspace ${workspaceId}`, - version: '1.0.1', + version: '1.1.0', tools, }; } diff --git a/packages/backend/server/src/plugins/copilot/mcp/resolver.ts b/packages/backend/server/src/plugins/copilot/mcp/resolver.ts index 0b2565db6..6c2a24708 100644 --- a/packages/backend/server/src/plugins/copilot/mcp/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/mcp/resolver.ts @@ -1,4 +1,3 @@ -import { BadRequestException } from '@nestjs/common'; import { Args, Field, @@ -109,7 +108,7 @@ export class McpCredentialResolver { @Query(() => Boolean) mcpCredentialReadWriteAvailable() { - return env.dev || env.namespaces.canary; + return true; } @Mutation(() => RevealedMcpCredentialType) @@ -117,13 +116,6 @@ export class McpCredentialResolver { @CurrentUser() user: CurrentUser, @Args('input') input: CreateMcpCredentialInput ) { - if ( - input.accessMode === McpAccessMode.READ_WRITE && - !env.dev && - !env.namespaces.canary - ) { - throw new BadRequestException('MCP write tools are not available'); - } await this.ac .user(user.id) .workspace(input.workspaceId)