next-tjwater-agent/src/mcp/query-server.ts

import { randomUUID, timingSafeEqual } from "node:crypto";
import { createMcpHandler, McpServer, fromJsonSchema, type Tool } from "@modelcontextprotocol/server";
import * as z from "zod/v4";
import { getRuntimeSessionContext, type RuntimeSessionContext } from "../runtime/sessionContext.js";
import { LOCAL_QUERY_TOOLS, PRODUCT_MCP_TOOLS, queryMcpToolId } from "../runtime/product-mcp-tools.js";
import type { McpQueryClient, QueryOwner } from "../product/mcp-query-client.js";
import { localQueryDescriptions, localQueryInputs } from "./local-query-tools.js";
import { allowedQueryTools } from "../runtime/query-mcp-policy.js";
import { localEvidenceResult } from "./evidence-result.js";
import { nativeQueryTimeout, queryTimeouts } from "./query-timeouts.js";

export type LocalQueryName = typeof LOCAL_QUERY_TOOLS[number];
const productNames = new Set<string>(PRODUCT_MCP_TOOLS);

export function createQueryMcpHandler(options: {
  internalToken: string; productClient: Pick<McpQueryClient, "query"> | null; productTools: Tool[];
  resolveReader: (id: string) => Promise<QueryOwner | null>;
  executeLocal: (name: LocalQueryName, args: Record<string, unknown>, context: RuntimeSessionContext, signal: AbortSignal) => Promise<unknown>;
}) {
  const definitions: Tool[] = [...options.productTools, ...LOCAL_QUERY_TOOLS.map((name) => ({ name, description: localQueryDescriptions[name],
    inputSchema: z.toJSONSchema(localQueryInputs[name], { io: "input" }) as Tool["inputSchema"], annotations: { readOnlyHint: true, destructiveHint: false, openWorldHint: name === "geocode" } }))];
  const handler = createMcpHandler(() => {
    const server = new McpServer({ name: "tjwater-agent-queries", version: "1.0.0" });
    for (const definition of definitions) {
      server.registerTool(definition.name, { ...definition, inputSchema: fromJsonSchema<Record<string, unknown>>(definition.inputSchema as Parameters<typeof fromJsonSchema>[0]), outputSchema: definition.outputSchema ? fromJsonSchema(definition.outputSchema) : undefined }, async (args, request) => {
        const id = request.mcpReq._meta?.["ai.opencode/sessionID"];
        if (typeof id !== "string" || !id || id.length > 128) return rejected("Authenticated native session metadata required");
        const direct = getRuntimeSessionContext(id);
        const owner = direct ? { context: direct, runtimeSessionId: id, taskId: direct.taskId } : await options.resolveReader(id);
        if (!owner || !allowedQueryTools(owner.context, !direct).includes(queryMcpToolId(definition.name))) return rejected("Query unavailable for this session");
        try {
          const budget = queryTimeouts(definition, args);
          const signal = AbortSignal.any([request.mcpReq.signal, AbortSignal.timeout(budget.requestMs)]);
          if (productNames.has(definition.name)) {
            if (!options.productClient) return rejected("Product MCP is disabled");
            const query = definition.name === "get_task_context" && args.report_revision === undefined && owner.context.reportRevision
              ? { ...args, report_revision: owner.context.reportRevision } : args;
            // JSON-RPC IDs may restart on a native MCP reconnect. Use a host
            // request key, stable across this client's auth/transport retries,
            // so a later query cannot collide with a saved evidence receipt.
            return localEvidenceResult(await options.productClient.query(id, owner, definition.name, query, randomUUID(), signal));
          }
          const value = await options.executeLocal(definition.name as LocalQueryName, args, owner.context, signal);
          signal.throwIfAborted();
          return { content: [{ type: "text", text: JSON.stringify(value) }] };
        } catch (cause) {
          // Never return upstream bodies, headers or credentials across MCP.
          const status = cause && typeof cause === "object" && "status" in cause ? cause.status : undefined;
          return rejected(`Query unavailable${typeof status === "number" ? ` (${status})` : ""}`);
        }
      });
    }
    return server;
  }, { legacy: "reject", maxRequestBodySize: 1024 * 1024 });
  return {
    executionTimeoutMs: nativeQueryTimeout(definitions),
    async fetch(request: Request) {
      const url = new URL(request.url);
      if (url.hostname !== "127.0.0.1" || url.pathname !== "/mcp" || request.headers.has("Origin") || request.headers.has("Sec-Fetch-Site")) return new Response(null, { status: 403 });
      const provided = Buffer.from(request.headers.get("x-agent-internal-token") ?? "");
      const expected = Buffer.from(options.internalToken);
      if (!expected.length || provided.length !== expected.length || !timingSafeEqual(provided, expected)) return new Response(null, { status: 403 });
      return handler.fetch(request);
    },
    close: handler.close,
  };
}
const rejected = (message: string) => ({ content: [{ type: "text" as const, text: message }], isError: true });

export function startQueryMcpServer(options: Parameters<typeof createQueryMcpHandler>[0]) {
  const handler = createQueryMcpHandler(options);
  const server = Bun.serve({ hostname: "127.0.0.1", port: 0, idleTimeout: 60,
    fetch(request, server) {
      // Native calls enforce a per-request deadline; an idle socket must not
      // disconnect a calculation before that deadline while no result is ready.
      server.timeout(request, 0);
      return handler.fetch(request);
    },
  });
  return { url: `http://127.0.0.1:${server.port}/mcp`, executionTimeoutMs: handler.executionTimeoutMs,
    async close() { await handler.close(); await server.stop(true); } };
}