Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/backend/server/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
6 changes: 5 additions & 1 deletion packages/backend/server/src/plugins/copilot/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -78,7 +80,9 @@ import {
UserCopilotResolver,
PromptsManagementResolver,
CopilotContextRootResolver,
// mcp
WorkspaceMcpProvider,
],
controllers: [CopilotController],
controllers: [CopilotController, WorkspaceMcpController],
})
export class CopilotModule {}
69 changes: 69 additions & 0 deletions packages/backend/server/src/plugins/copilot/mcp/controller.ts
Original file line number Diff line number Diff line change
@@ -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();
}
}
}
170 changes: 170 additions & 0 deletions packages/backend/server/src/plugins/copilot/mcp/provider.ts
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 0 additions & 3 deletions packages/backend/server/src/plugins/copilot/types.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import { z } from 'zod';

import { OneMB } from '../../base';
import type { ChatPrompt } from './prompt';
import { PromptMessageSchema, PureMessageSchema } from './providers';

Expand Down Expand Up @@ -130,5 +129,3 @@ export type CopilotContextFile = {
// embedding status
status: 'in_progress' | 'completed' | 'failed';
};

export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;
35 changes: 32 additions & 3 deletions packages/backend/server/src/plugins/copilot/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading