// mcp/src/tools/imageTools.ts
import type { Tool } from '@modelcontextprotocol/sdk/types.js';
import type { WhkerDBClient } from '../client/WhkerDBClient.js';
import type { SocketClient } from '../client/SocketClient.js';
import { EventType } from '@whkerdb/shared';
import { v4 as uuid } from 'uuid';
import * as fs from 'fs';
import * as path from 'path';
/**
* 图片工具定义
*/
export const imageTools: Tool[] = [
{
name: 'whkerdb_upload_image',
description: '上传图片到房间',
inputSchema: {
type: 'object',
properties: {
roomId: {
type: 'string',
description: '房间 ID',
},
filePath: {
type: 'string',
description: '图片文件的本地路径',
},
base64: {
type: 'string',
description: '图片的 Base64 编码（与 filePath 二选一）',
},
filename: {
type: 'string',
description: '文件名（使用 base64 时必填）',
},
mimeType: {
type: 'string',
description: 'MIME 类型（使用 base64 时可选，默认 image/png）',
},
},
required: ['roomId'],
},
},
{
name: 'whkerdb_list_images',
description: '列出房间内的所有图片',
inputSchema: {
type: 'object',
properties: {},
},
},
];
/**
* 根据文件扩展名获取 MIME 类型
*/
function getMimeType(filename: string): string {
const ext = path.extname(filename).toLowerCase();
const mimeTypes: Record<string, string> = {
'.png': 'image/png',
'.jpg': 'image/jpeg',
'.jpeg': 'image/jpeg',
'.gif': 'image/gif',
'.webp': 'image/webp',
'.svg': 'image/svg+xml',
};
return mimeTypes[ext] || 'image/png';
}
/**
* 图片工具处理器
*/
export function createImageToolHandlers(
httpClient: WhkerDBClient,
getSocketClient: () => SocketClient | null
) {
const ensureConnected = () => {
const client = getSocketClient();
if (!client?.isConnected()) {
throw new Error('未连接到任何房间，请先使用 whkerdb_join_room 加入房间');
}
return client;
};
return {
whkerdb_upload_image: async (args: {
roomId: string;
filePath?: string;
base64?: string;
filename?: string;
mimeType?: string;
}) => {
let fileBuffer: Buffer;
let filename: string;
let mimeType: 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);
mimeType = args.mimeType || getMimeType(filename);
} else if (args.base64) {
// 从 Base64 解码
if (!args.filename) {
throw new Error('使用 base64 上传时必须提供 filename');
}
fileBuffer = Buffer.from(args.base64, 'base64');
filename = args.filename;
mimeType = args.mimeType || getMimeType(filename);
} else {
throw new Error('必须提供 filePath 或 base64');
}
// 上传图片
const uploadResult = await httpClient.uploadImage(args.roomId, fileBuffer, filename, mimeType);
// 如果已连接到房间，发送图片上传事件
const socketClient = getSocketClient();
if (socketClient?.isConnected() && socketClient.getCurrentRoomId() === args.roomId) {
await socketClient.applyEvent({
id: uuid(),
type: EventType.IMAGE_UPLOADED,
roomId: args.roomId,
userId: socketClient.getUserId(),
clientId: socketClient.getClientId(),
timestamp: Date.now(),
version: socketClient.getNextVersion(),
aggregateId: args.roomId,
payload: {
imageId: uploadResult.imageId,
filename: uploadResult.filename,
mimeType: uploadResult.mimeType,
width: 0, // 服务器未返回尺寸
height: 0,
fileSize: uploadResult.size,
},
});
}
return {
content: [
{
type: 'text' as const,
text: JSON.stringify(
{
success: true,
imageId: uploadResult.imageId,
filename: uploadResult.filename,
mimeType: uploadResult.mimeType,
fileSize: uploadResult.size,
url: uploadResult.url,
},
null,
2
),
},
],
};
},
whkerdb_list_images: async () => {
const client = ensureConnected();
const snapshot = client.getSnapshot();
if (!snapshot) throw new Error('无法获取房间快照');
const images = snapshot.images.map(([id, info]) => ({
id,
filename: info.filename,
mimeType: info.mimeType,
width: info.width,
height: info.height,
fileSize: info.fileSize,
uploadedAt: info.uploadedAt,
url: httpClient.getImageUrl(id),
}));
return {
content: [
{
type: 'text' as const,
text: JSON.stringify(
{
count: images.length,
images,
},
null,
2
),
},
],
};
},
};
}
