Skip to content

Commit c769279

Browse files
forehaloforehalo
committed
feat(server): basic mcp server
1 parent feb42e3 commit c769279

File tree

10 files changed

+293
-46
lines changed

10 files changed

+293
-46
lines changed

‎packages/backend/server/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
"@fal-ai/serverless-client": "^0.15.0",
4141
"@google-cloud/opentelemetry-cloud-trace-exporter": "^2.4.1",
4242
"@google-cloud/opentelemetry-resource-util": "^2.4.0",
43+
"@modelcontextprotocol/sdk": "^1.16.0",
4344
"@nestjs-cls/transactional": "^2.6.1",
4445
"@nestjs-cls/transactional-adapter-prisma": "^1.2.19",
4546
"@nestjs/apollo": "^13.0.4",

‎packages/backend/server/src/plugins/copilot/context/resolver.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,7 @@ import { CopilotEmbeddingJob } from '../embedding';
5252
import { COPILOT_LOCKER, CopilotType } from '../resolver';
5353
import { ChatSessionService } from '../session';
5454
import { CopilotStorage } from '../storage';
55-
import { MAX_EMBEDDABLE_SIZE } from '../types';
56-
import { getSignal, readStream } from '../utils';
55+
import { getSignal, MAX_EMBEDDABLE_SIZE, readStream } from '../utils';
5756
import { CopilotContextService } from './service';
5857

5958
@InputType()

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ import {
1717
import { CopilotController } from './controller';
1818
import { CopilotCronJobs } from './cron';
1919
import { CopilotEmbeddingJob } from './embedding';
20+
import { WorkspaceMcpController } from './mcp/controller';
21+
import { WorkspaceMcpProvider } from './mcp/provider';
2022
import { ChatMessageCache } from './message';
2123
import { PromptService } from './prompt';
2224
import { CopilotProviderFactory, CopilotProviders } from './providers';
@@ -78,7 +80,9 @@ import {
7880
UserCopilotResolver,
7981
PromptsManagementResolver,
8082
CopilotContextRootResolver,
83+
// mcp
84+
WorkspaceMcpProvider,
8185
],
82-
controllers: [CopilotController],
86+
controllers: [CopilotController, WorkspaceMcpController],
8387
})
8488
export class CopilotModule {}

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
2+
import {
3+
Controller,
4+
Delete,
5+
Get,
6+
HttpCode,
7+
HttpStatus,
8+
Logger,
9+
Param,
10+
Post,
11+
Req,
12+
Res,
13+
} from '@nestjs/common';
14+
import type { Request, Response } from 'express';
15+
16+
import { CurrentUser } from '../../../core/auth';
17+
import { WorkspaceMcpProvider } from './provider';
18+
19+
@Controller('/api/workspaces/:workspaceId/mcp')
20+
export class WorkspaceMcpController {
21+
private readonly logger = new Logger(WorkspaceMcpController.name);
22+
constructor(private readonly provider: WorkspaceMcpProvider) {}
23+
24+
@Get('/')
25+
@Delete('/')
26+
@HttpCode(HttpStatus.METHOD_NOT_ALLOWED)
27+
async STATELESS_MCP_ENDPOINT() {
28+
return {
29+
jsonrpc: '2.0',
30+
error: {
31+
code: -32000,
32+
message: 'Method not allowed.',
33+
},
34+
id: null,
35+
};
36+
}
37+
38+
@Post('/')
39+
async mcp(
40+
@Req() req: Request,
41+
@Res() res: Response,
42+
@CurrentUser() user: CurrentUser,
43+
@Param('workspaceId') workspaceId: string
44+
) {
45+
let server = await this.provider.for(user.id, workspaceId);
46+
47+
const transport: StreamableHTTPServerTransport =
48+
new StreamableHTTPServerTransport({
49+
sessionIdGenerator: undefined,
50+
});
51+
52+
const cleanup = () => {
53+
transport.close().catch(e => {
54+
this.logger.error('Failed to close MCP transport', e);
55+
});
56+
server.close().catch(e => {
57+
this.logger.error('Failed to close MCP server', e);
58+
});
59+
};
60+
61+
try {
62+
res.on('close', cleanup);
63+
await server.connect(transport);
64+
await transport.handleRequest(req, res, req.body);
65+
} catch {
66+
cleanup();
67+
}
68+
}
69+
}

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

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
2+
import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
3+
import { Injectable } from '@nestjs/common';
4+
import { pick } from 'lodash-es';
5+
import z from 'zod';
6+
7+
import { DocReader } from '../../../core/doc';
8+
import { AccessController } from '../../../core/permission';
9+
import { IndexerService } from '../../indexer';
10+
import { CopilotContextService } from '../context';
11+
import { clearEmbeddingChunk } from '../utils';
12+
13+
@Injectable()
14+
export class WorkspaceMcpProvider {
15+
constructor(
16+
private readonly ac: AccessController,
17+
private readonly reader: DocReader,
18+
private readonly context: CopilotContextService,
19+
private readonly indexer: IndexerService
20+
) {}
21+
22+
async for(userId: string, workspaceId: string) {
23+
await this.ac.user(userId).workspace(workspaceId).assert('Workspace.Read');
24+
25+
const server = new McpServer({
26+
name: `AFFiNE MCP Server for Workspace ${workspaceId}`,
27+
version: '1.0.0',
28+
});
29+
30+
server.registerTool(
31+
'read_document',
32+
{
33+
title: 'Read Document',
34+
description: 'Read a document with given ID',
35+
inputSchema: {
36+
docId: z.string(),
37+
},
38+
},
39+
async ({ docId }) => {
40+
const notFoundError: CallToolResult = {
41+
isError: true,
42+
content: [
43+
{
44+
type: 'text',
45+
text: `Doc with id ${docId} not found.`,
46+
},
47+
],
48+
};
49+
50+
const accessible = await this.ac
51+
.user(userId)
52+
.workspace(workspaceId)
53+
.doc(docId)
54+
.can('Doc.Read');
55+
56+
if (!accessible) {
57+
return notFoundError;
58+
}
59+
60+
const content = await this.reader.getDocMarkdown(
61+
workspaceId,
62+
docId,
63+
false
64+
);
65+
66+
if (!content) {
67+
return notFoundError;
68+
}
69+
70+
return {
71+
content: [
72+
{
73+
type: 'text',
74+
text: content.markdown,
75+
},
76+
],
77+
};
78+
}
79+
);
80+
81+
server.registerTool(
82+
'semantic_search',
83+
{
84+
title: 'Semantic Search',
85+
description:
86+
'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).',
87+
inputSchema: {
88+
query: z.string(),
89+
},
90+
},
91+
async ({ query }, req) => {
92+
query = query.trim();
93+
if (!query) {
94+
return {
95+
isError: true,
96+
content: [
97+
{
98+
type: 'text',
99+
text: 'Query is required for semantic search.',
100+
},
101+
],
102+
};
103+
}
104+
105+
const chunks = await this.context.matchWorkspaceDocs(
106+
workspaceId,
107+
query,
108+
5,
109+
req.signal
110+
);
111+
112+
const docs = await this.ac
113+
.user(userId)
114+
.workspace(workspaceId)
115+
.docs(
116+
chunks.filter(c => 'docId' in c),
117+
'Doc.Read'
118+
);
119+
120+
return {
121+
content: docs.map(doc => ({
122+
type: 'text',
123+
text: clearEmbeddingChunk(doc).content,
124+
})),
125+
};
126+
}
127+
);
128+
129+
server.registerTool(
130+
'keyword_search',
131+
{
132+
title: 'Keyword Search',
133+
description:
134+
'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.',
135+
inputSchema: {
136+
query: z.string(),
137+
},
138+
},
139+
async ({ query }) => {
140+
query = query.trim();
141+
if (!query) {
142+
return {
143+
isError: true,
144+
content: [
145+
{
146+
type: 'text',
147+
text: 'Query is required for keyword search.',
148+
},
149+
],
150+
};
151+
}
152+
153+
let docs = await this.indexer.searchDocsByKeyword(workspaceId, query);
154+
docs = await this.ac
155+
.user(userId)
156+
.workspace(workspaceId)
157+
.docs(docs, 'Doc.Read');
158+
159+
return {
160+
content: docs.map(doc => ({
161+
type: 'text',
162+
text: JSON.stringify(pick(doc, 'docId', 'title', 'createdAt')),
163+
})),
164+
};
165+
}
166+
);
167+
168+
return server;
169+
}
170+
}

‎packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts

Lines changed: 1 addition & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,34 +7,9 @@ import type { ChunkSimilarity, Models } from '../../../models';
77
import type { CopilotContextService } from '../context';
88
import type { ContextSession } from '../context/session';
99
import type { CopilotChatOptions } from '../providers';
10+
import { clearEmbeddingChunk } from '../utils';
1011
import { toolError } from './error';
1112

12-
const FILTER_PREFIX = [
13-
'Title: ',
14-
'Created at: ',
15-
'Updated at: ',
16-
'Created by: ',
17-
'Updated by: ',
18-
];
19-
20-
function clearEmbeddingChunk(chunk: ChunkSimilarity): ChunkSimilarity {
21-
if (chunk.content) {
22-
const lines = chunk.content.split('\n');
23-
let maxLines = 5;
24-
while (maxLines > 0 && lines.length > 0) {
25-
if (FILTER_PREFIX.some(prefix => lines[0].startsWith(prefix))) {
26-
lines.shift();
27-
maxLines--;
28-
} else {
29-
// only process consecutive metadata rows
30-
break;
31-
}
32-
}
33-
return { ...chunk, content: lines.join('\n') };
34-
}
35-
return chunk;
36-
}
37-
3813
export const buildDocSearchGetter = (
3914
ac: AccessController,
4015
context: CopilotContextService,

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import { z } from 'zod';
22

3-
import { OneMB } from '../../base';
43
import type { ChatPrompt } from './prompt';
54
import { PromptMessageSchema, PureMessageSchema } from './providers';
65

@@ -130,5 +129,3 @@ export type CopilotContextFile = {
130129
// embedding status
131130
status: 'in_progress' | 'completed' | 'failed';
132131
};
133-
134-
export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;

‎packages/backend/server/src/plugins/copilot/utils.ts

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ import { Readable } from 'node:stream';
22

33
import type { Request } from 'express';
44

5-
import { readBufferWithLimit } from '../../base';
5+
import { OneMB, readBufferWithLimit } from '../../base';
6+
import { ChunkSimilarity } from '../../models';
67
import { PromptTools } from './providers';
7-
import { MAX_EMBEDDABLE_SIZE, ToolsConfig } from './types';
8+
import { ToolsConfig } from './types';
9+
10+
export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;
811

912
export function readStream(
1013
readable: Readable,
@@ -80,3 +83,29 @@ export function getTools(
8083
});
8184
return result;
8285
}
86+
87+
const FILTER_PREFIX = [
88+
'Title: ',
89+
'Created at: ',
90+
'Updated at: ',
91+
'Created by: ',
92+
'Updated by: ',
93+
];
94+
95+
export function clearEmbeddingChunk(chunk: ChunkSimilarity): ChunkSimilarity {
96+
if (chunk.content) {
97+
const lines = chunk.content.split('\n');
98+
let maxLines = 5;
99+
while (maxLines > 0 && lines.length > 0) {
100+
if (FILTER_PREFIX.some(prefix => lines[0].startsWith(prefix))) {
101+
lines.shift();
102+
maxLines--;
103+
} else {
104+
// only process consecutive metadata rows
105+
break;
106+
}
107+
}
108+
return { ...chunk, content: lines.join('\n') };
109+
}
110+
return chunk;
111+
}

‎packages/backend/server/src/plugins/copilot/workspace/resolver.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import { CurrentUser } from '../../../core/auth';
2727
import { AccessController } from '../../../core/permission';
2828
import { WorkspaceType } from '../../../core/workspaces';
2929
import { COPILOT_LOCKER } from '../resolver';
30-
import { MAX_EMBEDDABLE_SIZE } from '../types';
30+
import { MAX_EMBEDDABLE_SIZE } from '../utils';
3131
import { CopilotWorkspaceService } from './service';
3232
import {
3333
CopilotWorkspaceFileType,

0 commit comments

Comments
 (0)