132 lines
4.1 KiB
TypeScript
132 lines
4.1 KiB
TypeScript
import { Request, Response } from 'express';
|
|
import { randomUUID, createHash } from 'crypto';
|
|
import jwt from 'jsonwebtoken';
|
|
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp';
|
|
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp';
|
|
import { JWT_SECRET } from '../config';
|
|
import { db } from '../db/database';
|
|
import { User } from '../types';
|
|
import { registerResources } from './resources';
|
|
import { registerTools } from './tools';
|
|
|
|
interface McpSession {
|
|
transport: StreamableHTTPServerTransport;
|
|
userId: number;
|
|
lastActivity: number;
|
|
}
|
|
|
|
const sessions = new Map<string, McpSession>();
|
|
|
|
const SESSION_TTL_MS = 60 * 60 * 1000; // 1 hour
|
|
|
|
const sessionSweepInterval = setInterval(() => {
|
|
const cutoff = Date.now() - SESSION_TTL_MS;
|
|
for (const [sid, session] of sessions) {
|
|
if (session.lastActivity < cutoff) {
|
|
try { session.transport.close(); } catch { /* ignore */ }
|
|
sessions.delete(sid);
|
|
}
|
|
}
|
|
}, 10 * 60 * 1000); // sweep every 10 minutes
|
|
|
|
// Prevent the interval from keeping the process alive if nothing else is running
|
|
sessionSweepInterval.unref();
|
|
|
|
function verifyToken(authHeader: string | undefined): User | null {
|
|
const token = authHeader && authHeader.split(' ')[1];
|
|
if (!token) return null;
|
|
|
|
// Long-lived MCP API token (trek_...)
|
|
if (token.startsWith('trek_')) {
|
|
const hash = createHash('sha256').update(token).digest('hex');
|
|
const row = db.prepare(`
|
|
SELECT u.id, u.username, u.email, u.role
|
|
FROM mcp_tokens mt
|
|
JOIN users u ON mt.user_id = u.id
|
|
WHERE mt.token_hash = ?
|
|
`).get(hash) as User | undefined;
|
|
if (row) {
|
|
// Update last_used_at (fire-and-forget, non-blocking)
|
|
db.prepare('UPDATE mcp_tokens SET last_used_at = CURRENT_TIMESTAMP WHERE token_hash = ?').run(hash);
|
|
return row;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
// Short-lived JWT
|
|
try {
|
|
const decoded = jwt.verify(token, JWT_SECRET) as { id: number };
|
|
const user = db.prepare(
|
|
'SELECT id, username, email, role FROM users WHERE id = ?'
|
|
).get(decoded.id) as User | undefined;
|
|
return user || null;
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
export async function mcpHandler(req: Request, res: Response): Promise<void> {
|
|
const mcpAddon = db.prepare("SELECT enabled FROM addons WHERE id = 'mcp'").get() as { enabled: number } | undefined;
|
|
if (!mcpAddon || !mcpAddon.enabled) {
|
|
res.status(403).json({ error: 'MCP is not enabled' });
|
|
return;
|
|
}
|
|
|
|
const user = verifyToken(req.headers['authorization']);
|
|
if (!user) {
|
|
res.status(401).json({ error: 'Access token required' });
|
|
return;
|
|
}
|
|
|
|
const sessionId = req.headers['mcp-session-id'] as string | undefined;
|
|
|
|
// Resume an existing session
|
|
if (sessionId) {
|
|
const session = sessions.get(sessionId);
|
|
if (!session) {
|
|
res.status(404).json({ error: 'Session not found' });
|
|
return;
|
|
}
|
|
if (session.userId !== user.id) {
|
|
res.status(403).json({ error: 'Session belongs to a different user' });
|
|
return;
|
|
}
|
|
session.lastActivity = Date.now();
|
|
await session.transport.handleRequest(req, res, req.body);
|
|
return;
|
|
}
|
|
|
|
// Only POST can initialize a new session
|
|
if (req.method !== 'POST') {
|
|
res.status(400).json({ error: 'Missing mcp-session-id header' });
|
|
return;
|
|
}
|
|
|
|
// Create a new per-user MCP server and session
|
|
const server = new McpServer({ name: 'trek', version: '1.0.0' });
|
|
registerResources(server, user.id);
|
|
registerTools(server, user.id);
|
|
|
|
const transport = new StreamableHTTPServerTransport({
|
|
sessionIdGenerator: () => randomUUID(),
|
|
onsessioninitialized: (sid) => {
|
|
sessions.set(sid, { transport, userId: user.id, lastActivity: Date.now() });
|
|
},
|
|
onsessionclosed: (sid) => {
|
|
sessions.delete(sid);
|
|
},
|
|
});
|
|
|
|
await server.connect(transport);
|
|
await transport.handleRequest(req, res, req.body);
|
|
}
|
|
|
|
/** Close all active MCP sessions (call during graceful shutdown). */
|
|
export function closeMcpSessions(): void {
|
|
clearInterval(sessionSweepInterval);
|
|
for (const [, session] of sessions) {
|
|
try { session.transport.close(); } catch { /* ignore */ }
|
|
}
|
|
sessions.clear();
|
|
}
|