// mcp/src/tools/pdfTools.ts
import type { Tool } from '@modelcontextprotocol/sdk/types.js';
import type { WhkerDBClient } from '../client/WhkerDBClient.js';
import type { SocketClient } from '../client/SocketClient.js';
import type { RoomSnapshot } from '@whkerdb/shared';
import { EventType, createFileNode, createPageNode } from '@whkerdb/shared';
import { v4 as uuid } from 'uuid';
import * as fs from 'fs';
import * as path from 'path';
/**
* PDF 工具定义
*/
export const pdfTools: Tool[] = [
{
name: 'whkerdb_upload_pdf',
description: '上传 PDF 文件到房间',
inputSchema: {
type: 'object',
properties: {
roomId: {
type: 'string',
description: '房间 ID',
},
filePath: {
type: 'string',
description: 'PDF 文件的本地路径',
},
base64: {
type: 'string',
description: 'PDF 文件的 Base64 编码（与 filePath 二选一）',
},
filename: {
type: 'string',
description: '文件名（使用 base64 时必填）',
},
},
required: ['roomId'],
},
},
{
name: 'whkerdb_list_pdfs',
description: '列出房间内的所有 PDF',
inputSchema: {
type: 'object',
properties: {},
},
},
{
name: 'whkerdb_get_pdf_info',
description: '获取 PDF 的详细信息',
inputSchema: {
type: 'object',
properties: {
pdfId: {
type: 'string',
description: 'PDF ID',
},
},
required: ['pdfId'],
},
},
];
/**
* PDF 工具处理器
*/
export function createPdfToolHandlers(
httpClient: WhkerDBClient,
getSocketClient: () => SocketClient | null
) {
const ensureConnected = () => {
const client = getSocketClient();
if (!client?.isConnected()) {
throw new Error('未连接到任何房间，请先使用 whkerdb_join_room 加入房间');
}
return client;
};
return {
whkerdb_upload_pdf: async (args: {
roomId: string;
filePath?: string;
base64?: string;
filename?: string;
}) => {
let fileBuffer: Buffer;
let filename: string;
if (args.filePath) {
// 从文件路径读取
if (!fs.existsSync(args.filePath)) {
throw new Error(`文件不存在: ${args.filePath}`);
}
fileBuffer = fs.readFileSync(args.filePath);
filename = args.filename || path.basename(args.filePath);
} else if (args.base64) {
// 从 Base64 解码
if (!args.filename) {
throw new Error('使用 base64 上传时必须提供 filename');
}
fileBuffer = Buffer.from(args.base64, 'base64');
filename = args.filename;
} else {
throw new Error('必须提供 filePath 或 base64');
}
// 上传 PDF
const uploadResult = await httpClient.uploadPdf(args.roomId, fileBuffer, filename);
// 如果已连接到房间，同步 PDF 元信息，并创建文件节点和页面节点
const socketClient = getSocketClient();
if (socketClient?.isConnected() && socketClient.getCurrentRoomId() === args.roomId) {
await socketClient.applyEvent({
id: uuid(),
type: EventType.PDF_UPLOADED,
roomId: args.roomId,
userId: socketClient.getUserId(),
clientId: socketClient.getClientId(),
timestamp: Date.now(),
version: socketClient.getNextVersion(),
aggregateId: args.roomId,
payload: {
pdfId: uploadResult.pdfId,
filename: uploadResult.filename,
totalPages: uploadResult.totalPages,
fileSize: uploadResult.size,
},
});
const snapshot = socketClient.getSnapshot();
// 创建文件节点
const fileNode = createFileNode(
uploadResult.filename,
uploadResult.pdfId,
uploadResult.totalPages
);
const filePosition = snapshot ? snapshot.noteTree.metadata.rootIds.length : 0;
await socketClient.applyEvent({
id: uuid(),
type: EventType.NODE_ADDED,
roomId: args.roomId,
userId: socketClient.getUserId(),
clientId: socketClient.getClientId(),
timestamp: Date.now(),
version: socketClient.getNextVersion(),
aggregateId: args.roomId,
payload: {
node: fileNode,
parentId: null,
position: filePosition,
},
});
// 为每页创建页面节点
for (let i = 0; i < uploadResult.totalPages; i++) {
const pageSize = uploadResult.pageSizes[i] || { width: 595, height: 842 };
const pageNode = createPageNode(
`第 ${i + 1} 页`,
fileNode.id,
i + 1,
pageSize.width,
pageSize.height
);
await socketClient.applyEvent({
id: uuid(),
type: EventType.NODE_ADDED,
roomId: args.roomId,
userId: socketClient.getUserId(),
clientId: socketClient.getClientId(),
timestamp: Date.now(),
version: socketClient.getNextVersion(),
aggregateId: args.roomId,
payload: {
node: pageNode,
parentId: fileNode.id,
position: i,
},
});
}
return {
content: [
{
type: 'text' as const,
text: JSON.stringify(
{
success: true,
pdfId: uploadResult.pdfId,
filename: uploadResult.filename,
totalPages: uploadResult.totalPages,
fileSize: uploadResult.size,
fileNodeId: fileNode.id,
message: 'PDF 已上传并创建笔记树节点',
},
null,
2
),
},
],
};
}
return {
content: [
{
type: 'text' as const,
text: JSON.stringify(
{
success: true,
pdfId: uploadResult.pdfId,
filename: uploadResult.filename,
totalPages: uploadResult.totalPages,
fileSize: uploadResult.size,
message: 'PDF 已上传（未创建笔记树节点，请先加入房间）',
},
null,
2
),
},
],
};
},
whkerdb_list_pdfs: async () => {
const client = ensureConnected();
const snapshot = client.getSnapshot();
if (!snapshot) throw new Error('无法获取房间快照');
const pdfs = snapshot.pdfs.map(([id, info]) => ({
id,
filename: info.filename,
totalPages: info.totalPages,
fileSize: info.fileSize,
uploadedAt: info.uploadedAt,
fileNodeId: info.fileNodeId,
}));
return {
content: [
{
type: 'text' as const,
text: JSON.stringify(
{
count: pdfs.length,
pdfs,
},
null,
2
),
},
],
};
},
whkerdb_get_pdf_info: async (args: { pdfId: string }) => {
const info = await httpClient.getPdfInfo(args.pdfId);
const fileUrl = httpClient.getPdfFileUrl(args.pdfId);
return {
content: [
{
type: 'text' as const,
text: JSON.stringify(
{
...info,
fileUrl,
},
null,
2
),
},
],
};
},
};
}
