feat(mcp): resume queued calculation jobs

src/mcp/native-call-identity.ts

@@ -0,0 +1,20 @@
+import { isDeepStrictEqual } from "node:util";
+import { queryMcpToolId } from "../runtime/product-mcp-tools.js";
+
+const record = (value: unknown): value is Record<string, unknown> => !!value && typeof value === "object" && !Array.isArray(value);
+
+// Use only host-read native history. Ambiguous simultaneous identical calls fail
+// closed rather than assigning a second submission to the first call's job.
+export function resolveMatchingNativeCall(messages: unknown, toolName: string, args: Record<string, unknown>): string | null {
+  if (!Array.isArray(messages)) return null;
+  const ids = new Set<string>();
+  for (const message of messages) {
+    if (!record(message) || !Array.isArray(message.parts)) continue;
+    for (const part of message.parts) {
+      if (!record(part) || part.type !== "tool" || part.tool !== queryMcpToolId(toolName) || !record(part.state)) continue;
+      if (!["running", "pending"].includes(String(part.state.status)) || !isDeepStrictEqual(part.state.input, args)) continue;
+      if (typeof part.callID === "string" && part.callID.length > 0 && part.callID.length <= 256) ids.add(part.callID);
+    }
+  }
+  return ids.size === 1 ? [...ids][0]! : null;
+}

src/mcp/query-server.ts

@@ -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); } };

src/mcp/query-timeouts.ts

@@ -3,9 +3,9 @@ import type { Tool } from "@modelcontextprotocol/client";
 // Absolute validation bound; the product catalog advertises the administrator's
 // effective maximum (one hour by default), never an identity or authorization grant.
 const ABSOLUTE_CALCULATION_SECONDS = 86_400;
-const CALCULATION_OVERHEAD_MS = 60_000;
+
 export const DEFAULT_QUERY_REQUEST_TIMEOUT_MS = 35_000;
-export const MAX_NATIVE_QUERY_TIMEOUT_MS = ABSOLUTE_CALCULATION_SECONDS * 1000 + CALCULATION_OVERHEAD_MS + 5000;
+export const MAX_NATIVE_QUERY_TIMEOUT_MS = DEFAULT_QUERY_REQUEST_TIMEOUT_MS;
 
 function timeoutSchema(tool: Tool) {
   const property = tool.inputSchema.properties?.timeout_seconds;
@@ -22,12 +22,10 @@ export function queryTimeouts(tool: Tool, args: Record<string, unknown>) {
     throw Object.assign(new Error("Calculation timeout must be an integer within the advertised server limit"), { status: 400 });
   }
   const executionMs = seconds * 1000;
-  return { executionMs, callMs: executionMs + 5000, requestMs: executionMs + CALCULATION_OVERHEAD_MS };
+  // timeout_seconds controls the queued worker, never this short submission.
+  return { executionMs, callMs: 30_000, requestMs: DEFAULT_QUERY_REQUEST_TIMEOUT_MS };
 }
 
-export function nativeQueryTimeout(tools: Tool[]) {
-  return Math.max(DEFAULT_QUERY_REQUEST_TIMEOUT_MS, ...tools.map((tool) => {
-    const schema = timeoutSchema(tool);
-    return schema ? queryTimeouts(tool, { timeout_seconds: schema.maximum }).requestMs + 5000 : DEFAULT_QUERY_REQUEST_TIMEOUT_MS;
-  }));
+export function nativeQueryTimeout(_tools: Tool[]) {
+  return DEFAULT_QUERY_REQUEST_TIMEOUT_MS;
 }

src/mcp/task-calculation-waiter.ts

@@ -0,0 +1,53 @@
+import { randomUUID } from "node:crypto";
+import { setTimeout as delay } from "node:timers/promises";
+import type { McpQueryClient, QueryOwner } from "../product/mcp-query-client.js";
+
+export type TaskCalculationWaiter = (sessionId: string, signal: AbortSignal) => Promise<string | null>;
+const record = (value: unknown): value is Record<string, unknown> => !!value && typeof value === "object" && !Array.isArray(value);
+
+// The host waits between short status requests while the task lease continues
+// heartbeating. No native model turn or long-lived MCP request performs polling.
+export function createTaskCalculationWaiter(
+  client: Pick<McpQueryClient, "query">,
+  resolveOwner: (sessionId: string) => QueryOwner | null,
+  pause: (ms: number, signal: AbortSignal) => Promise<unknown> = (ms, signal) => delay(ms, undefined, { signal }),
+): TaskCalculationWaiter {
+  const delivered = new Map<string, Set<string>>();
+  return async (sessionId, signal) => {
+    const owner = resolveOwner(sessionId);
+    if (!owner || owner.context.purpose !== "task_execution") throw new Error("Calculation task owner unavailable");
+    const seen = delivered.get(sessionId) ?? new Set<string>();
+    delivered.set(sessionId, seen);
+    const terminal = new Map<string, string>();
+    for (;;) {
+      signal.throwIfAborted();
+      let pending = false;
+      let pollMs = 5000;
+      let offset = 0;
+      for (;;) {
+        const result = await client.query(sessionId, owner, "list_calculation_jobs", { limit: 100, offset }, randomUUID(), signal);
+        const body = result.structuredContent;
+        if (result.isError || !record(body) || !Array.isArray(body.data)) throw new Error("Calculation status unavailable");
+        for (const job of body.data) {
+          if (!record(job) || typeof job.job_id !== "string" || typeof job.status !== "string") throw new Error("Calculation status invalid");
+          if (["queued", "running", "cancel_requested"].includes(job.status)) {
+            pending = true;
+            if (typeof job.poll_after_ms === "number" && Number.isFinite(job.poll_after_ms)) pollMs = Math.min(pollMs, Math.max(1500, job.poll_after_ms));
+          } else if (["succeeded", "failed", "timed_out", "cancelled", "interrupted"].includes(job.status)) {
+            if (!seen.has(job.job_id)) terminal.set(job.job_id, job.status);
+          } else throw new Error("Calculation status invalid");
+        }
+        if (!record(body.page) || typeof body.page.has_more !== "boolean") throw new Error("Calculation status paging invalid");
+        if (!body.page.has_more) break;
+        offset += 100;
+        if (offset > 10000) throw new Error("Calculation status paging limit exceeded");
+      }
+      if (!pending) {
+        if (!terminal.size) return null;
+        for (const id of terminal.keys()) seen.add(id);
+        return `后台计算已进入终态：${JSON.stringify([...terminal].map(([job_id, status]) => ({ job_id, status })))}。继续当前调查，使用 get_calculation_result 按 job_id、data_path、limit、offset 小页读取成功结果后形成判断；失败、取消、超时或中断应按执行异常处理，不得当作无异常或数据缺失。不要重复提交这些作业。`;
+      }
+      await pause(pollMs, signal);
+    }
+  };
+}