112 lines
8.0 KiB
TypeScript
112 lines
8.0 KiB
TypeScript
import { createHash, 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, PRODUCT_MCP_QUEUED_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 { boundedQuery, MAX_NATIVE_QUERY_TIMEOUT_MS, modelQueryTool, nativeQueryTimeout, queryTimeouts } from "./query-timeouts.js";
|
|
import { saveMcpArtifact } from "./artifact-files.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[];
|
|
resolveNativeCall?: (sessionId: string, toolName: string, args: Record<string, unknown>) => Promise<string | null>;
|
|
resolveReader: (id: string) => Promise<QueryOwner | null>;
|
|
executeLocal: (name: LocalQueryName, args: Record<string, unknown>, context: RuntimeSessionContext, signal: AbortSignal) => Promise<unknown>;
|
|
artifactWorkspace?: (context: RuntimeSessionContext) => Promise<string>;
|
|
}) {
|
|
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 originalDefinition of definitions) {
|
|
const definition = modelQueryTool(originalDefinition);
|
|
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");
|
|
try {
|
|
const budget = queryTimeouts(originalDefinition);
|
|
const signal = AbortSignal.any([request.mcpReq.signal, AbortSignal.timeout(budget.requestMs)]);
|
|
return await boundedQuery(async () => {
|
|
const direct = getRuntimeSessionContext(id);
|
|
const owner = direct ? { context: direct, runtimeSessionId: id, taskId: direct.taskId } : await options.resolveReader(id);
|
|
signal.throwIfAborted();
|
|
if (!owner || !allowedQueryTools(owner.context, !direct).includes(queryMcpToolId(definition.name))) return rejected("Query unavailable for this session");
|
|
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;
|
|
// Native JSON-RPC IDs restart on reconnect. Resolve the persisted
|
|
// native tool call from host history, never from model arguments.
|
|
const nativeCallId = await options.resolveNativeCall?.(id, definition.name, args);
|
|
signal.throwIfAborted();
|
|
const requiresIdentity = (PRODUCT_MCP_QUEUED_TOOLS as readonly string[]).includes(definition.name)
|
|
&& (!["save_simulation_scenario", "create_sensor_placement"].includes(definition.name) || args.dry_run === false);
|
|
if (requiresIdentity && !nativeCallId) return rejected("Native tool call identity unavailable; check calculation jobs before retrying");
|
|
const callId = createHash("sha256").update(JSON.stringify([
|
|
owner.context.actorKey, owner.context.projectKey, id, definition.name, nativeCallId ?? randomUUID(),
|
|
])).digest("hex");
|
|
const result = await boundedQuery((boundedSignal) => options.productClient!.query(id, owner, definition.name, query, callId, boundedSignal), signal);
|
|
if (["export_project_model", "export_sensor_placement"].includes(definition.name) && !result.isError) {
|
|
if (!options.artifactWorkspace) return rejected("Export workspace unavailable");
|
|
const workspace = await options.artifactWorkspace(owner.context);
|
|
signal.throwIfAborted();
|
|
return await saveMcpArtifact(result, owner.context, workspace);
|
|
}
|
|
return localEvidenceResult(result);
|
|
}
|
|
const value = await boundedQuery((boundedSignal) => options.executeLocal(definition.name as LocalQueryName, args, owner.context, boundedSignal), signal);
|
|
signal.throwIfAborted();
|
|
return { content: [{ type: "text", text: JSON.stringify(value) }] };
|
|
}, signal);
|
|
} catch (cause) {
|
|
// Never return upstream bodies, headers or credentials across MCP.
|
|
const status = cause && typeof cause === "object" && "status" in cause ? cause.status : undefined;
|
|
const name = cause instanceof Error ? cause.name : undefined;
|
|
const code = cause && typeof cause === "object" && "code" in cause ? cause.code : undefined;
|
|
const category = name === "TimeoutError" || code === -32001 || status === 408 || status === 504 ? "Query timed out"
|
|
: name === "AbortError" ? "Query cancelled"
|
|
: cause instanceof TypeError || ["ECONNRESET", "ECONNREFUSED", "ENOTFOUND", "EPIPE"].includes(String(code)) ? "Query network failure" : "Query unavailable";
|
|
return rejected(`${category}${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, // Bun caps socket idleTimeout at 255 seconds. Enforce a finite five-minute
|
|
// HTTP response deadline explicitly so the advertised 300s budget can finish.
|
|
idleTimeout: 0,
|
|
fetch: async (request) => {
|
|
try {
|
|
return await boundedQuery((signal) => handler.fetch(new Request(request, { signal })), AbortSignal.any([request.signal, AbortSignal.timeout(MAX_NATIVE_QUERY_TIMEOUT_MS)]));
|
|
} catch {
|
|
return new Response("Query timed out", { status: 504 });
|
|
}
|
|
},
|
|
});
|
|
return { url: `http://127.0.0.1:${server.port}/mcp`, executionTimeoutMs: handler.executionTimeoutMs,
|
|
async close() { await handler.close(); await server.stop(true); } };
|
|
}
|