youling-task/lib/server.js

基于POMDP的研发任务分解工具，支持MCP协议

/**
 * MCP服务器实现 - 提供POMDP任务分解服务
 */
import express from 'express';
import dotenv from 'dotenv';
import { randomUUID } from 'crypto';
import { pomdpTaskDecomposer } from './pomdp-decomposer.js';
import { z } from 'zod';
import logger from './utils/logger.js';
// 按照SDK实际的导出路径导入
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
// 加载环境变量
dotenv.config();
// 默认端口号和密钥
const DEFAULT_PORT = 3000;
const DEFAULT_SECRET_KEY = 'youling-task-secret';
/**
 * 创建并启动YoulingTask MCP服务器
 * @param options 服务器配置选项
 * @returns 返回Express应用实例和服务器实例
 */
export async function startYoulingTaskServer(options) {
    const { port = DEFAULT_PORT, secretKey = process.env.MCP_SECRET_KEY || DEFAULT_SECRET_KEY, serverName = 'youling-task-decomposer', serverVersion = '1.0.0' } = options || {};
    // 创建Express应用
    const app = express();
    // 基本中间件
    app.use(express.json());
    // 健康检查端点
    app.get('/health', (req, res) => {
        res.status(200).json({ status: 'ok', timestamp: new Date().toISOString() });
    });
    // 创建MCP服务器实例
    const mcpServer = new McpServer({
        name: serverName,
        version: serverVersion
    });
    // 添加简单的认证中间件
    app.use((req, res, next) => {
        const authHeader = req.headers.authorization;
        // 如果请求路径是健康检查，则跳过认证
        if (req.path === '/health') {
            return next();
        }
        if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== secretKey) {
            return res.status(401).json({ error: 'Unauthorized' });
        }
        next();
    });
    // 注册POMDP任务分解资源
    mcpServer.tool('decompose-task', '使用POMDP方法将研发任务分解为详细计划', {
        taskDescription: z.string().describe('需要分解的研发任务描述'),
    }, 
    // 工具处理函数
    async (params) => {
        try {
            logger.info('收到分解任务请求', params);
            // 提取任务描述
            const { taskDescription } = params;
            // 执行POMDP任务分解
            const plan = await pomdpTaskDecomposer(taskDescription);
            // 返回分解结果
            return {
                content: [
                    {
                        type: 'text',
                        text: plan,
                    },
                ],
            };
        }
        catch (error) {
            logger.error('任务分解处理失败', { error: error instanceof Error ? error.message : String(error) });
            throw error;
        }
    });
    // 每个客户端请求创建一个新的ServerTransport和McpServer实例
    // 这样可以避免"Server already initialized"的问题
    // 存储为每个会话创建的服务器和传输层的映射
    const sessions = new Map();
    // 创建主HTTP传输层，仅用于初始会话分配
    const masterTransport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => randomUUID(), // 生成会话ID
        enableJsonResponse: true, // 使用JSON响应而不是SSE流
        onsessioninitialized: (sessionId) => {
            logger.info('MCP主会话已初始化', { sessionId });
            // 为这个会话创建一个新的服务器实例和传输层
            const sessionServer = new McpServer({
                name: serverName,
                version: serverVersion
            });
            // 为新服务器实例注册相同的工具
            sessionServer.tool('decompose-task', '使用POMDP方法将研发任务分解为详细计划', {
                taskDescription: z.string().describe('需要分解的研发任务描述'),
            }, async (params) => {
                try {
                    logger.info('会话处理分解任务请求', { sessionId, params });
                    const { taskDescription } = params;
                    const plan = await pomdpTaskDecomposer(taskDescription);
                    return {
                        content: [
                            {
                                type: 'text',
                                text: plan,
                            },
                        ],
                    };
                }
                catch (error) {
                    logger.error('会话任务分解失败', {
                        sessionId,
                        error: error instanceof Error ? error.message : String(error)
                    });
                    throw error;
                }
            });
            // 保存会话信息
            sessions.set(sessionId, {
                server: sessionServer,
                created: new Date()
            });
        }
    });
    // 连接主服务器和传输层
    await mcpServer.connect(masterTransport);
    // 定义MCP请求处理函数
    const mcpHandler = (req, res) => {
        // 检查会话ID
        const sessionId = req.headers['mcp-session-id'];
        if (sessionId && sessions.has(sessionId)) {
            // 会话已存在，使用对应的服务器实例处理
            logger.info('收到已存在会话的请求', { sessionId });
            // 目前不需要单独为每个会话创建传输层
            // 先使用主传输层处理所有请求
        }
        // 使用主传输层处理请求
        masterTransport.handleRequest(req, res, req.body);
    };
    // 注册MCP路由处理器
    app.post('/mcp', mcpHandler); // 处理POST请求
    app.get('/mcp', mcpHandler); // 处理GET请求的SSE流
    // 启动HTTP服务器
    const server = app.listen(port, () => {
        logger.info('服务器启动成功', {
            port: port,
            endpoint: `http://localhost:${port}/mcp`
        });
    });
    return { app, server };
}
/**
 * 如果直接运行此文件，则启动服务器
 */
if (process.argv[1] === process.argv[1]) {
    startYoulingTaskServer();
}
//# sourceMappingURL=server.js.map