GearBox/src/server/mcp/index.ts

import { randomUUID } from "node:crypto";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
import { Hono } from "hono";
import { db as prodDb } from "../../db/index.ts";
import { getUserCount, verifyApiKey } from "../services/auth.service.ts";
import { getCollectionSummary } from "./resources/collection.ts";
import {
	categoryToolDefinitions,
	registerCategoryTools,
} from "./tools/categories.ts";
import { imageToolDefinitions, registerImageTools } from "./tools/images.ts";
import { itemToolDefinitions, registerItemTools } from "./tools/items.ts";
import { registerSetupTools, setupToolDefinitions } from "./tools/setups.ts";
import { registerThreadTools, threadToolDefinitions } from "./tools/threads.ts";

type Db = typeof prodDb;

function createMcpServer(db: Db): McpServer {
	const server = new McpServer({ name: "GearBox", version: "1.0.0" });

	// Register item tools
	const itemHandlers = registerItemTools(db);
	for (const def of itemToolDefinitions) {
		const handler = itemHandlers[def.name as keyof typeof itemHandlers];
		server.tool(def.name, def.description, def.inputSchema, handler);
	}

	// Register category tools
	const categoryHandlers = registerCategoryTools(db);
	for (const def of categoryToolDefinitions) {
		const handler = categoryHandlers[def.name as keyof typeof categoryHandlers];
		server.tool(def.name, def.description, def.inputSchema, handler);
	}

	// Register thread tools
	const threadHandlers = registerThreadTools(db);
	for (const def of threadToolDefinitions) {
		const handler = threadHandlers[def.name as keyof typeof threadHandlers];
		server.tool(def.name, def.description, def.inputSchema, handler);
	}

	// Register setup tools
	const setupHandlers = registerSetupTools(db);
	for (const def of setupToolDefinitions) {
		const handler = setupHandlers[def.name as keyof typeof setupHandlers];
		server.tool(def.name, def.description, def.inputSchema, handler);
	}

	// Register image tools
	const imageHandlers = registerImageTools();
	for (const def of imageToolDefinitions) {
		const handler = imageHandlers[def.name as keyof typeof imageHandlers];
		server.tool(def.name, def.description, def.inputSchema, handler);
	}

	// Register collection summary resource
	server.resource(
		"collection-summary",
		"gearbox://collection/summary",
		{
			description:
				"Overview of the entire gear collection including totals, categories, and active research threads.",
			mimeType: "application/json",
		},
		async () => {
			const summary = getCollectionSummary(db);
			return {
				contents: [
					{
						uri: "gearbox://collection/summary",
						mimeType: "application/json",
						text: JSON.stringify(summary, null, 2),
					},
				],
			};
		},
	);

	return server;
}

// Store active transports by session ID
const transports = new Map<string, WebStandardStreamableHTTPServerTransport>();

export const mcpRoutes = new Hono();

// Auth middleware for all MCP requests
mcpRoutes.use("/*", async (c, next) => {
	const db = c.get("db") ?? prodDb;
	const apiKey = c.req.header("X-API-Key");

	// Require API key when auth is configured (users exist)
	if (getUserCount(db) > 0) {
		if (!apiKey) {
			return c.json({ error: "API key required" }, 401);
		}
		const valid = await verifyApiKey(db, apiKey);
		if (!valid) {
			return c.json({ error: "Invalid API key" }, 401);
		}
	}

	return next();
});

mcpRoutes.post("/", async (c) => {
	const db = c.get("db") ?? prodDb;

	// Check for existing session
	const sessionId = c.req.header("mcp-session-id");

	if (sessionId) {
		const transport = transports.get(sessionId);
		if (!transport) {
			return c.json({ error: "Session not found" }, 404);
		}
		const response = await transport.handleRequest(c.req.raw);
		return response;
	}

	// New session: create transport and MCP server
	const transport = new WebStandardStreamableHTTPServerTransport({
		sessionIdGenerator: () => randomUUID(),
		onsessioninitialized: (newSessionId) => {
			transports.set(newSessionId, transport);
		},
	});

	// Clean up on close
	transport.onclose = () => {
		const sid = [...transports.entries()].find(
			([_, t]) => t === transport,
		)?.[0];
		if (sid) transports.delete(sid);
	};

	const server = createMcpServer(db);
	await server.connect(transport);

	const response = await transport.handleRequest(c.req.raw);
	return response;
});

mcpRoutes.get("/", async (c) => {
	const sessionId = c.req.header("mcp-session-id");
	if (!sessionId) {
		return c.json({ error: "Session ID required" }, 400);
	}

	const transport = transports.get(sessionId);
	if (!transport) {
		return c.json({ error: "Session not found" }, 404);
	}

	const response = await transport.handleRequest(c.req.raw);
	return response;
});

mcpRoutes.delete("/", async (c) => {
	const sessionId = c.req.header("mcp-session-id");
	if (!sessionId) {
		return c.json({ error: "Session ID required" }, 400);
	}

	const transport = transports.get(sessionId);
	if (!transport) {
		return c.json({ error: "Session not found" }, 404);
	}

	await transport.close();
	transports.delete(sessionId);
	return c.text("", 200);
});