diff --git a/packages/backend/server/package.json b/packages/backend/server/package.json index 2efd0ebfa8..3608195e70 100644 --- a/packages/backend/server/package.json +++ b/packages/backend/server/package.json @@ -40,6 +40,7 @@ "@fal-ai/serverless-client": "^0.15.0", "@google-cloud/opentelemetry-cloud-trace-exporter": "^2.4.1", "@google-cloud/opentelemetry-resource-util": "^2.4.0", + "@modelcontextprotocol/sdk": "^1.16.0", "@nestjs-cls/transactional": "^2.6.1", "@nestjs-cls/transactional-adapter-prisma": "^1.2.19", "@nestjs/apollo": "^13.0.4", diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index b3420cf4ae..9177caf073 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -52,8 +52,7 @@ import { CopilotEmbeddingJob } from '../embedding'; import { COPILOT_LOCKER, CopilotType } from '../resolver'; import { ChatSessionService } from '../session'; import { CopilotStorage } from '../storage'; -import { MAX_EMBEDDABLE_SIZE } from '../types'; -import { getSignal, readStream } from '../utils'; +import { getSignal, MAX_EMBEDDABLE_SIZE, readStream } from '../utils'; import { CopilotContextService } from './service'; @InputType() diff --git a/packages/backend/server/src/plugins/copilot/index.ts b/packages/backend/server/src/plugins/copilot/index.ts index 02982bbbb8..1c575ade46 100644 --- a/packages/backend/server/src/plugins/copilot/index.ts +++ b/packages/backend/server/src/plugins/copilot/index.ts @@ -17,6 +17,8 @@ import { import { CopilotController } from './controller'; import { CopilotCronJobs } from './cron'; import { CopilotEmbeddingJob } from './embedding'; +import { WorkspaceMcpController } from './mcp/controller'; +import { WorkspaceMcpProvider } from './mcp/provider'; import { ChatMessageCache } from './message'; import { PromptService } from './prompt'; import { CopilotProviderFactory, CopilotProviders } from './providers'; @@ -78,7 +80,9 @@ import { UserCopilotResolver, PromptsManagementResolver, CopilotContextRootResolver, + // mcp + WorkspaceMcpProvider, ], - controllers: [CopilotController], + controllers: [CopilotController, WorkspaceMcpController], }) export class CopilotModule {} diff --git a/packages/backend/server/src/plugins/copilot/mcp/controller.ts b/packages/backend/server/src/plugins/copilot/mcp/controller.ts new file mode 100644 index 0000000000..e738d9d31f --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/mcp/controller.ts @@ -0,0 +1,69 @@ +import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +import { + Controller, + Delete, + Get, + HttpCode, + HttpStatus, + Logger, + Param, + Post, + Req, + Res, +} from '@nestjs/common'; +import type { Request, Response } from 'express'; + +import { CurrentUser } from '../../../core/auth'; +import { WorkspaceMcpProvider } from './provider'; + +@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 { + jsonrpc: '2.0', + error: { + code: -32000, + message: 'Method not allowed.', + }, + id: null, + }; + } + + @Post('/') + async mcp( + @Req() req: Request, + @Res() res: Response, + @CurrentUser() user: CurrentUser, + @Param('workspaceId') workspaceId: string + ) { + let server = await this.provider.for(user.id, workspaceId); + + 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); + await server.connect(transport); + await transport.handleRequest(req, res, req.body); + } catch { + cleanup(); + } + } +} diff --git a/packages/backend/server/src/plugins/copilot/mcp/provider.ts b/packages/backend/server/src/plugins/copilot/mcp/provider.ts new file mode 100644 index 0000000000..c926ec8146 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/mcp/provider.ts @@ -0,0 +1,170 @@ +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'; + +import { DocReader } from '../../../core/doc'; +import { AccessController } from '../../../core/permission'; +import { IndexerService } from '../../indexer'; +import { CopilotContextService } from '../context'; +import { clearEmbeddingChunk } from '../utils'; + +@Injectable() +export class WorkspaceMcpProvider { + constructor( + private readonly ac: AccessController, + private readonly reader: DocReader, + private readonly context: CopilotContextService, + private readonly indexer: IndexerService + ) {} + + async for(userId: string, workspaceId: string) { + await this.ac.user(userId).workspace(workspaceId).assert('Workspace.Read'); + + const server = new McpServer({ + name: `AFFiNE MCP Server for Workspace ${workspaceId}`, + version: '1.0.0', + }); + + server.registerTool( + 'read_document', + { + title: 'Read Document', + description: 'Read a document with given ID', + inputSchema: { + docId: z.string(), + }, + }, + async ({ docId }) => { + const notFoundError: CallToolResult = { + 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; + } + + const content = await this.reader.getDocMarkdown( + workspaceId, + docId, + false + ); + + if (!content) { + return notFoundError; + } + + return { + content: [ + { + type: 'text', + text: content.markdown, + }, + ], + }; + } + ); + + 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: { + query: z.string(), + }, + }, + async ({ query }, req) => { + query = query.trim(); + if (!query) { + return { + isError: true, + content: [ + { + type: 'text', + text: 'Query is required for semantic search.', + }, + ], + }; + } + + const chunks = await this.context.matchWorkspaceDocs( + workspaceId, + query, + 5, + req.signal + ); + + const docs = await this.ac + .user(userId) + .workspace(workspaceId) + .docs( + chunks.filter(c => 'docId' in c), + 'Doc.Read' + ); + + return { + content: docs.map(doc => ({ + type: 'text', + text: clearEmbeddingChunk(doc).content, + })), + }; + } + ); + + 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: { + query: z.string(), + }, + }, + async ({ query }) => { + query = query.trim(); + if (!query) { + return { + isError: true, + content: [ + { + type: 'text', + text: 'Query is required for keyword search.', + }, + ], + }; + } + + let docs = await this.indexer.searchDocsByKeyword(workspaceId, query); + docs = await this.ac + .user(userId) + .workspace(workspaceId) + .docs(docs, 'Doc.Read'); + + return { + content: docs.map(doc => ({ + type: 'text', + text: JSON.stringify(pick(doc, 'docId', 'title', 'createdAt')), + })), + }; + } + ); + + return server; + } +} diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts index 83660be18e..9736ba3d8e 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts @@ -7,34 +7,9 @@ import type { ChunkSimilarity, Models } from '../../../models'; import type { CopilotContextService } from '../context'; import type { ContextSession } from '../context/session'; import type { CopilotChatOptions } from '../providers'; +import { clearEmbeddingChunk } from '../utils'; import { toolError } from './error'; -const FILTER_PREFIX = [ - 'Title: ', - 'Created at: ', - 'Updated at: ', - 'Created by: ', - 'Updated by: ', -]; - -function clearEmbeddingChunk(chunk: ChunkSimilarity): ChunkSimilarity { - if (chunk.content) { - const lines = chunk.content.split('\n'); - let maxLines = 5; - while (maxLines > 0 && lines.length > 0) { - if (FILTER_PREFIX.some(prefix => lines[0].startsWith(prefix))) { - lines.shift(); - maxLines--; - } else { - // only process consecutive metadata rows - break; - } - } - return { ...chunk, content: lines.join('\n') }; - } - return chunk; -} - export const buildDocSearchGetter = ( ac: AccessController, context: CopilotContextService, diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts index af6da6c7e4..71487152f9 100644 --- a/packages/backend/server/src/plugins/copilot/types.ts +++ b/packages/backend/server/src/plugins/copilot/types.ts @@ -1,6 +1,5 @@ import { z } from 'zod'; -import { OneMB } from '../../base'; import type { ChatPrompt } from './prompt'; import { PromptMessageSchema, PureMessageSchema } from './providers'; @@ -130,5 +129,3 @@ export type CopilotContextFile = { // embedding status status: 'in_progress' | 'completed' | 'failed'; }; - -export const MAX_EMBEDDABLE_SIZE = 50 * OneMB; diff --git a/packages/backend/server/src/plugins/copilot/utils.ts b/packages/backend/server/src/plugins/copilot/utils.ts index e22f4b7584..9e8ee77f43 100644 --- a/packages/backend/server/src/plugins/copilot/utils.ts +++ b/packages/backend/server/src/plugins/copilot/utils.ts @@ -2,9 +2,12 @@ import { Readable } from 'node:stream'; import type { Request } from 'express'; -import { readBufferWithLimit } from '../../base'; -import { PromptTools } from './providers'; -import { MAX_EMBEDDABLE_SIZE, ToolsConfig } from './types'; +import { OneMB, readBufferWithLimit } from '../../base'; +import type { ChunkSimilarity } from '../../models'; +import type { PromptTools } from './providers'; +import type { ToolsConfig } from './types'; + +export const MAX_EMBEDDABLE_SIZE = 50 * OneMB; export function readStream( readable: Readable, @@ -80,3 +83,29 @@ export function getTools( }); return result; } + +const FILTER_PREFIX = [ + 'Title: ', + 'Created at: ', + 'Updated at: ', + 'Created by: ', + 'Updated by: ', +]; + +export function clearEmbeddingChunk(chunk: ChunkSimilarity): ChunkSimilarity { + if (chunk.content) { + const lines = chunk.content.split('\n'); + let maxLines = 5; + while (maxLines > 0 && lines.length > 0) { + if (FILTER_PREFIX.some(prefix => lines[0].startsWith(prefix))) { + lines.shift(); + maxLines--; + } else { + // only process consecutive metadata rows + break; + } + } + return { ...chunk, content: lines.join('\n') }; + } + return chunk; +} diff --git a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts index 9d33c92ab9..936d931386 100644 --- a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts @@ -27,7 +27,7 @@ import { CurrentUser } from '../../../core/auth'; import { AccessController } from '../../../core/permission'; import { WorkspaceType } from '../../../core/workspaces'; import { COPILOT_LOCKER } from '../resolver'; -import { MAX_EMBEDDABLE_SIZE } from '../types'; +import { MAX_EMBEDDABLE_SIZE } from '../utils'; import { CopilotWorkspaceService } from './service'; import { CopilotWorkspaceFileType, diff --git a/yarn.lock b/yarn.lock index ab0db81ab3..abd6f2c6c3 100644 --- a/yarn.lock +++ b/yarn.lock @@ -931,6 +931,7 @@ __metadata: "@fal-ai/serverless-client": "npm:^0.15.0" "@google-cloud/opentelemetry-cloud-trace-exporter": "npm:^2.4.1" "@google-cloud/opentelemetry-resource-util": "npm:^2.4.0" + "@modelcontextprotocol/sdk": "npm:^1.16.0" "@nestjs-cls/transactional": "npm:^2.6.1" "@nestjs-cls/transactional-adapter-prisma": "npm:^1.2.19" "@nestjs/apollo": "npm:^13.0.4" @@ -7927,21 +7928,23 @@ __metadata: languageName: node linkType: hard -"@modelcontextprotocol/sdk@npm:^1.8.0": - version: 1.11.1 - resolution: "@modelcontextprotocol/sdk@npm:1.11.1" +"@modelcontextprotocol/sdk@npm:^1.16.0, @modelcontextprotocol/sdk@npm:^1.8.0": + version: 1.16.0 + resolution: "@modelcontextprotocol/sdk@npm:1.16.0" dependencies: + ajv: "npm:^6.12.6" content-type: "npm:^1.0.5" cors: "npm:^2.8.5" - cross-spawn: "npm:^7.0.3" + cross-spawn: "npm:^7.0.5" eventsource: "npm:^3.0.2" + eventsource-parser: "npm:^3.0.0" express: "npm:^5.0.1" express-rate-limit: "npm:^7.5.0" pkce-challenge: "npm:^5.0.0" raw-body: "npm:^3.0.0" zod: "npm:^3.23.8" zod-to-json-schema: "npm:^3.24.1" - checksum: 10/bf388e3f5082839ccf32eb4f16e086ead71310f30c3103ff99f337d7bcfc6da6b3377dc9bd64ac9b862969487081c36889faea02c5c30805695e8addac96b9a8 + checksum: 10/2f462fc1fc1f06c6c4b1e01e403f01e41ba5aaaaebd298381ccf70242d1c73423c9c4a2b7717b19c190e65a098979ccebf2f50b159c55870f5f3977d5b24e269 languageName: node linkType: hard @@ -16649,7 +16652,7 @@ __metadata: languageName: node linkType: hard -"ajv@npm:^6.12.0, ajv@npm:^6.12.4": +"ajv@npm:^6.12.0, ajv@npm:^6.12.4, ajv@npm:^6.12.6": version: 6.12.6 resolution: "ajv@npm:6.12.6" dependencies: @@ -19291,7 +19294,7 @@ __metadata: languageName: node linkType: hard -"cross-spawn@npm:^7.0.1, cross-spawn@npm:^7.0.3, cross-spawn@npm:^7.0.6": +"cross-spawn@npm:^7.0.1, cross-spawn@npm:^7.0.3, cross-spawn@npm:^7.0.5, cross-spawn@npm:^7.0.6": version: 7.0.6 resolution: "cross-spawn@npm:7.0.6" dependencies: @@ -21149,10 +21152,10 @@ __metadata: languageName: node linkType: hard -"eventsource-parser@npm:^3.0.1": - version: 3.0.1 - resolution: "eventsource-parser@npm:3.0.1" - checksum: 10/2730c54c3cb47d55d2967f2ece843f9fc95d8a11c2fef6fece8d17d9080193cbe3cd9ac7b04a325977f63cbf8c1664fdd0512dec1aec601666a5c5bd8564b61f +"eventsource-parser@npm:^3.0.0, eventsource-parser@npm:^3.0.1": + version: 3.0.3 + resolution: "eventsource-parser@npm:3.0.3" + checksum: 10/b8f8e79333441ad0eb9299e3fa693ab506892ffc53f0cc1d23134090351cf2d71c8e405a2e879f6acfbd2e17f41d5a00dafba05ff25c82141fc07078ad992187 languageName: node linkType: hard