task-manager/server/mcp.ts

import { Request, Response } from "express";
import { storage } from "./storage";
import { User } from "../shared/schema";
import { randomBytes } from "crypto";

interface JsonRpcRequest {
    jsonrpc: "2.0";
    method: string;
    params?: any;
    id: number | string;
}

interface JsonRpcResponse {
    jsonrpc: "2.0";
    result?: any;
    error?: {
        code: number;
        message: string;
        data?: any;
    };
    id: number | string | null;
}

export class McpServer {
    private clients: Map<string, Response> = new Map();

    async authenticate(req: Request): Promise<User | null> {
        let key = req.headers["x-api-key"] as string;
        if (!key && req.query.apiKey) {
            key = req.query.apiKey as string;
        }
        // Bearer token support
        if (!key && req.headers["authorization"]) {
            const auth = req.headers["authorization"];
            if (auth.startsWith("Bearer ")) {
                key = auth.substring(7);
            }
        }

        if (!key) return null;
        return (await storage.getUserByApiKey(key)) || null;
    }

    async handleSse(req: Request, res: Response) {
        const user = await this.authenticate(req);
        if (!user) {
            res.status(401).send("Unauthorized: Invalid API Key");
            return;
        }

        res.writeHead(200, {
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*", // Allow local connections
        });

        const sessionId = randomBytes(8).toString("hex");
        this.clients.set(sessionId, res);

        const endpoint = `/api/mcp/messages`;
        res.write(`event: endpoint\ndata: ${endpoint}\n\n`);

        // Keep alive
        const interval = setInterval(() => {
            res.write(": keepalive\n\n");
        }, 15000);

        req.on("close", () => {
            clearInterval(interval);
            this.clients.delete(sessionId);
        });
    }

    async handleMessage(req: Request, res: Response) {
        const user = await this.authenticate(req);
        if (!user) {
            res.status(401).json({ error: "Unauthorized" });
            return;
        }

        const body = req.body as JsonRpcRequest;
        try {
            const result = await this.processRequest(body, user);
            res.json({
                jsonrpc: "2.0",
                result,
                id: body.id
            });
        } catch (err: any) {
            res.json({
                jsonrpc: "2.0",
                error: {
                    code: -32000,
                    message: err.message || "Internal Server Error"
                },
                id: body.id
            });
        }
    }

    private async processRequest(req: JsonRpcRequest, user: User): Promise<any> {
        switch (req.method) {
            case "initialize":
                return {
                    protocolVersion: "2024-11-05",
                    capabilities: {
                        tools: {},
                        resources: {}
                    },
                    serverInfo: {
                        name: "TaskFlow MCP",
                        version: "1.0.0"
                    }
                };
            case "tools/list": // MCP method
            case "listTools": // Legacy fallback
                return {
                    tools: [
                        {
                            name: "list_tasks",
                            description: "List all tasks for the user",
                            inputSchema: {
                                type: "object",
                                properties: {
                                    status: { type: "string", enum: ["todo", "inProgress", "done"], description: "Filter by status" },
                                    limit: { type: "number", description: "Limit number of tasks" }
                                }
                            }
                        },
                        {
                            name: "create_task",
                            description: "Create a new task",
                            inputSchema: {
                                type: "object",
                                properties: {
                                    title: { type: "string", description: "Title of the task" },
                                    description: { type: "string", description: "Description" },
                                    priority: { type: "string", enum: ["low", "medium", "high"] }
                                },
                                required: ["title"]
                            }
                        },
                        {
                            name: "complete_task",
                            description: "Mark a task as completed",
                            inputSchema: {
                                type: "object",
                                properties: {
                                    id: { type: "string", description: "Task ID" }
                                },
                                required: ["id"]
                            }
                        }
                    ]
                };
            case "tools/call": // MCP method
            case "callTool": // Legacy
                return await this.handleToolCall(req.params.name, req.params.arguments, user);

            case "notifications/initialized":
                return true;

            default:
                throw new Error(`Method ${req.method} not found`);
        }
    }

    private async handleToolCall(name: string, args: any, user: User) {
        switch (name) {
            case "list_tasks": {
                let tasks = await storage.getTasksForUser(user.id);
                if (args?.status) {
                    tasks = tasks.filter((t) => t.status === args.status);
                }
                if (args?.limit) {
                    tasks = tasks.slice(0, args.limit);
                }
                return { content: [{ type: "text", text: JSON.stringify(tasks, null, 2) }] };
            }
            case "create_task": {
                if (!args.title) throw new Error("Title is required");
                const task = await storage.createTask({
                    title: args.title,
                    description: args.description || "",
                    priority: args.priority || "medium",
                    status: "todo",
                    isTracking: false,
                    timeTracked: 0,
                    energyLevel: "medium",
                    estimatedDuration: 15,
                    dueDate: null,
                    notes: "",
                    labelId: null,
                    userId: user.id
                });
                return { content: [{ type: "text", text: JSON.stringify(task, null, 2) }] };
            }
            case "complete_task": {
                if (!args.id) throw new Error("ID is required");
                const task = await storage.getTask(args.id);
                if (!task || task.userId !== user.id) throw new Error("Task not found");
                const updated = await storage.updateTask(args.id, { status: "done" });
                return { content: [{ type: "text", text: JSON.stringify(updated, null, 2) }] };
            }
            default:
                throw new Error(`Tool ${name} not implemented`);
        }
    }
}

export const mcpServer = new McpServer();