import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { z } from 'zod'
import fs from 'fs'
import path from 'path'
import { getSession } from '../state.js'
import { getSessionClientCount, requestExport } from '../websocket.js'
import { log } from '../logger.js'
/**
* 导出图表为文件
*/
export function registerExport(server: McpServer): void {
server.registerTool(
'export_diagram',
{
description:
'Export the current diagram to a PNG, SVG, or JSON file.\n\n' +
'Notes:\n' +
'- JSON format is saved directly by the server.\n' +
'- PNG/SVG are generated by the browser and sent back to the server to save.\n' +
'- Requires start_session to be called first and the browser to be connected.\n\n' +
'Multi-session support: Specify sessionId to target a specific session.',
inputSchema: z.object({
sessionId: z
.string()
.optional()
.describe('Session ID. If not provided, uses default session.'),
path: z.string().describe('File path to save (including filename)'),
format: z.enum(['png', 'svg', 'json']).default('json').describe('Export format'),
}),
},
async ({ sessionId, path: filePath, format }) => {
try {
const session = getSession(sessionId)
// 确保路径有正确扩展名
const ext = `.${format}`
const finalPath = filePath.endsWith(ext) ? filePath : `${filePath}${ext}`
// 准备导出数据
const exportData = {
elements: session.elements,
appState: session.appState,
version: session.version,
}
// 保存为 JSON (服务器端保存)
if (format === 'json') {
// 确保目录存在
const dir = path.dirname(finalPath)
if (!fs.existsSync(dir)) {
fs.mkdirSync(dir, { recursive: true })
}
fs.writeFileSync(finalPath, JSON.stringify(exportData, null, 2))
log.info(`Exported diagram to ${finalPath}`)
return {
content: [
{
type: 'text',
text:
`✅ Diagram exported!\n\n` +
`Path: ${finalPath}\n` +
`Format: ${format.toUpperCase()}\n` +
`Session: ${session.id}\n` +
`Elements: ${session.elements.length}`,
},
],
}
}
if (getSessionClientCount(session.id) === 0) {
return {
content: [
{
type: 'text',
text: `❌ No browser session detected. Please use start_session to start a session first: ${session.id}`,
},
],
isError: true,
}
}
const exportFormat = format === 'png' ? 'png' : 'svg'
const data = await requestExport({
sessionId: session.id,
format: exportFormat,
elements: session.elements,
appState: session.appState,
})
const dir = path.dirname(finalPath)
if (!fs.existsSync(dir)) {
fs.mkdirSync(dir, { recursive: true })
}
const buffer = Buffer.from(data, 'base64')
fs.writeFileSync(finalPath, buffer)
log.info(`Exported diagram to ${finalPath}`)
return {
content: [
{
type: 'text',
text:
`✅ Diagram exported!\n\n` +
`Path: ${finalPath}\n` +
`Format: ${format.toUpperCase()}\n` +
`Session: ${session.id}\n` +
`Elements: ${session.elements.length}`,
},
],
}
} catch (error) {
const message = error instanceof Error ? error.message : String(error)
return {
content: [{ type: 'text', text: `Error: ${message}` }],
isError: true,
}
}
},
)
}