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 { verifyApiKey } from "../services/auth.service.ts";
import { verifyAccessToken } from "../services/oauth.service.ts";
import { getCollectionSummary } from "./resources/collection.ts";
import {
	catalogToolDefinitions,
	registerCatalogTools,
} from "./tools/catalog.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, userId: number): McpServer {
	const server = new McpServer({ name: "GearBox", version: "1.0.0" });

	// Register item tools
	const itemHandlers = registerItemTools(db, userId);
	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, userId);
	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, userId);
	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, userId);
	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 catalog tools (no userId needed — catalog is global)
	const catalogHandlers = registerCatalogTools(db);
	for (const def of catalogToolDefinitions) {
		const handler = catalogHandlers[def.name as keyof typeof catalogHandlers];
		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 = await getCollectionSummary(db, userId);
			return {
				contents: [
					{
						uri: "gearbox://collection/summary",
						mimeType: "application/json",
						text: JSON.stringify(summary, null, 2),
					},
				],
			};
		},
	);

	return server;
}

// Store active transports by session ID (with userId for session reuse)
const transports = new Map<
	string,
	{ transport: WebStandardStreamableHTTPServerTransport; userId: number }
>();

export const mcpRoutes = new Hono();

// Auth middleware for all MCP requests — resolves userId
mcpRoutes.use("/*", async (c, next) => {
	const db = c.get("db") ?? prodDb;

	// Try Bearer token first (OAuth)
	const authHeader = c.req.header("Authorization");
	if (authHeader?.startsWith("Bearer ")) {
		const token = authHeader.slice(7);
		const result = await verifyAccessToken(db, token);
		if (result) {
			c.set("userId", result.userId);
			return next();
		}
		return c.json({ error: "invalid_token" }, 401);
	}

	// Try API key
	const apiKey = c.req.header("X-API-Key");
	if (apiKey) {
		const result = await verifyApiKey(db, apiKey);
		if (result) {
			c.set("userId", result.userId);
			return next();
		}
		return c.json({ error: "Invalid API key" }, 401);
	}

	// No auth provided — return 401 with WWW-Authenticate to trigger OAuth flow (RFC 9728)
	const baseUrl = (
		process.env.GEARBOX_URL || new URL(c.req.url).origin
	).replace(/\/$/, "");
	return c.text("Unauthorized", 401, {
		"WWW-Authenticate": `Bearer resource_metadata="${baseUrl}/.well-known/oauth-protected-resource"`,
	});
});

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

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

	if (sessionId) {
		const session = transports.get(sessionId);
		if (!session) {
			return c.json({ error: "Session not found" }, 404);
		}
		const response = await session.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, userId });
		},
	});

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

	const server = createMcpServer(db, userId);
	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 session = transports.get(sessionId);
	if (!session) {
		return c.json({ error: "Session not found" }, 404);
	}

	const response = await session.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 session = transports.get(sessionId);
	if (!session) {
		return c.json({ error: "Session not found" }, 404);
	}

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