OrgTJWater/next-tjwater-agent
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mcp): resume queued calculation jobs

Browse Source
This commit is contained in:
jiang
2026-10-29 18:25:22 +08:00
parent 52681a8451
commit 799b99f310
20 changed files with 348 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mcp/query-server.ts
+20 -13
View File
@@ -1,8 +1,8 @@
import { randomUUID, timingSafeEqual } from "node:crypto";
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, queryMcpToolId } from "../runtime/product-mcp-tools.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";
@@ -15,6 +15,7 @@ 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>;
@@ -37,10 +38,16 @@ export function createQueryMcpHandler(options: {
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.
const result = await options.productClient.query(id, owner, definition.name, query, randomUUID(), signal);
// 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);
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 options.productClient.query(id, owner, definition.name, query, callId, signal);
if (["export_project_model", "export_sensor_placement"].includes(definition.name) && !result.isError) {
if (!options.artifactWorkspace) return rejected("Export workspace unavailable");
return await saveMcpArtifact(result, owner.context, await options.artifactWorkspace(owner.context));
@@ -53,7 +60,12 @@ export function createQueryMcpHandler(options: {
} 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})` : ""}`);
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})` : ""}`);
}
});
}
@@ -77,12 +89,7 @@ const rejected = (message: string) => ({ content: [{ type: "text" as const, text
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);
},
fetch: (request) => 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); } };