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
.opencode/agents/instruction.md
+5 -1
View File
@@ -70,7 +70,11 @@ temperature: 0.2
| 给定调泵方案的影响 | `analyze_pressure_regulation` |
| 压力监测点选址 | `optimize_sensor_placement` |
调用计算工具时，主动按任务设置 `timeout_seconds`：结合模型规模、`duration`、水力步长对应的计算步数和已有耗时选择等待预算，并遵守工具 schema 公布的上限（默认 1 小时）。`duration` 是模拟时间，`timeout_seconds` 是实际等待时间，两者不能混用。默认等待 120 秒；更长任务需主动调整预算，也不要一律取最大值。超时后底层计算可能仍在运行，在线模拟可能已经写入部分基线；先核实状态，不自动重复整段计算。
长计算使用普通业务工具提交后台作业，提交和状态读取应快速返回。按规模、`duration` 和步数设置 `timeout_seconds`，它是后台作业的执行预算，不是 HTTP/MCP 等待时长；`duration` 是模拟时长。遵守 schema 上限，爆管定位默认 1800 秒，其他计算通常默认 120 秒。
返回 `semantics=task`、`data.job_id` 及 queued/running 状态仅表示已受理，不是计算完成或分析证据。保留并向用户说明作业句柄和当前状态；聊天工具结果保存该句柄，可在断线、重连或下一轮恢复。用 `get_calculation_job(job_id)` 读取状态，用 `list_calculation_jobs` 找回当前授权范围作业；完成后用 `get_calculation_result(job_id,data_path,limit,offset)` 小页读取结果。遵守 `poll_after_ms`，一次答复最多检查两次状态；仍在运行时返回作业状态和恢复方式，不连续忙轮询，不让一次 MCP 请求等待整个计算。提交超时或网络异常时先查已有作业，不能直接重复提交。
`cancel_calculation_job` 是独立状态变更，需要现有权限和审批流程，不能按只读查询自动放行。只读子代理仅能读取授权的作业状态和结果，不能提交或取消。保存方案、创建选址的预览和批准流程保持有效，实际执行才进入队列。失败、取消或执行超时不是无异常结论；在线模拟可能有部分基线写入，应保留状态并说明恢复边界。
爆管定位 `run_burst_location` 默认使用 monitoring 数据。`pressure_scada_ids` 可省略，由后端自动读取当前项目的 SCADA 压力测点，不因用户未提供节点 ID 而阻塞，也不猜测或把候选管段 ID 当压力节点。正常对照窗可成对省略，默认前一天同时间段；主动说明这一对照假设，并根据返回的实际测点、时间窗和数据覆盖解释结果，不能宣称默认时段已确认正常。数据不足时报告缺口，不自动换窗或补造数据。simulation 模式仍须已有分析运行 UUID，使用同窗模型基线；在线 `run_simulation` 没有生成该 UUID。
src/mcp/native-call-identity.ts
+20
View File
@@ -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
+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); } };
src/mcp/query-timeouts.ts
+6 -8
View File
@@ -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
+53
View File
@@ -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);
}
};
}
src/routes/chat.ts
+3
View File
@@ -1,3 +1,4 @@
import type { TaskCalculationWaiter } from "../mcp/task-calculation-waiter.js";
import { registerInteractionRequestRoutes } from "./interaction-request-routes.js";
import type { InteractionRequestStore } from "../sessions/interaction-request-store.js";
import { Router, type NextFunction, type Request, type Response } from "express";
@@ -33,6 +34,7 @@ import { registerArtifactRoutes } from "./artifact-routes.js";
import type { AgentMessageCoordinator } from "../chat/agent-message-coordinator.js";
export type BuildChatRouterOptions = {
waitTaskCalculations?: TaskCalculationWaiter;
interactionRequestStore?: InteractionRequestStore;
agentMessages?: AgentMessageCoordinator;
frontendActionCoordinator: FrontendActionCoordinator;
@@ -170,6 +172,7 @@ export const buildChatRouter = (options: BuildChatRouterOptions) => {
registerResponseInputRoutes(chatRouter, { runtime, sessionMetadataStore, runState });
registerChatStreamRoute(chatRouter, {
waitTaskCalculations: options.waitTaskCalculations,
interactionRequestStore: options.interactionRequestStore,
agentMessages,
frontendActionCoordinator,
src/routes/chatStream.ts
+6 -4
View File
@@ -77,7 +77,7 @@ type StreamPromptOptions = {
onInboxEvent?: (type: string, data: { inboxID?: string; delivery?: "steer" | "queue" }) => void;
waitForUserInput?: () => Promise<boolean>;
beforePrompt?: () => Promise<void>;
beforeIdle?: () => Promise<boolean>;
beforeIdle?: () => Promise<boolean | string>;
beforeRichAnswerClose?: () => Promise<boolean>;
runtime: OpencodeRuntimeAdapter;
sessionId: string;
@@ -904,7 +904,7 @@ export const streamPromptResponse = async ({
id: part.id,
phase: "tool",
status: normalizeToolStatus(part.state.status),
title: getToolProgressTitle(part.tool, part.state.status),
title: getToolProgressTitle(part.tool, part.state.status, part.state.status === "completed" ? part.state.output : undefined),
tool: { name: part.tool, call_id: part.callID, input: sanitizeToolProgressInput(toolParams) },
detail: buildToolProgressDetail(
part.tool,
@@ -912,6 +912,7 @@ export const streamPromptResponse = async ({
toolParams,
reason,
part.state.status === "error" ? part.state.error : undefined,
part.state.status === "completed" ? part.state.output : undefined,
),
});
if (
@@ -1022,12 +1023,13 @@ export const streamPromptResponse = async ({
});
continue;
}
if (beforeIdle && !await beforeIdle()) {
const idleDecision = beforeIdle ? await beforeIdle() : true;
if (idleDecision === false || typeof idleDecision === "string") {
// A durable message won the race with completion. Continue within
// the same managed session/lease, never create another task run.
finalAnswer = undefined;
sawResponseActivity = false;
await runtime.prompt(sessionId, "请处理收信箱中新到达的调查协作消息，按需回复，然后完成本轮工作。", toRuntimeModel(model));
await runtime.prompt(sessionId, typeof idleDecision === "string" ? idleDecision : "请处理收信箱中新到达的调查协作消息，按需回复，然后完成本轮工作。", toRuntimeModel(model));
continue;
}
if (waitForUserInput && await waitForUserInput()) {
src/routes/chatStreamEvents.ts
+24 -1
View File
@@ -468,6 +468,7 @@ export const buildToolProgressDetail = (
params: Record<string, unknown>,
reason: string,
error?: string,
output?: string,
) => {
const toolName = toolLabels[tool] ?? tool;
const reasonText = reason ? `；说明：${containsForbiddenPersistentContent(reason) ? "[redacted]" : reason}` : "";
@@ -478,6 +479,8 @@ export const buildToolProgressDetail = (
return `${toolName} 调用失败${reasonText}${paramsText}${errorText}`;
}
if (status === "completed") {
const calculation = calculationProgress(output);
if (calculation) return `${toolName}：${calculation.label}；计算任务：${calculation.jobId}${reasonText}${paramsText}`;
return `${toolName} 已执行完成${reasonText}${paramsText}`;
}
if (status === "pending") {
@@ -486,8 +489,28 @@ export const buildToolProgressDetail = (
return `${toolName} 正在执行${reasonText}${paramsText}`;
};
export const getToolProgressTitle = (tool: string, status: string) => {
const calculationProgress = (output?: string) => {
if (!output || output.length > 100_000) return null;
try {
const value = JSON.parse(output);
if (value?.semantics !== "task" || typeof value.data?.job_id !== "string" ||
!/^[0-9a-f-]{36}$/i.test(value.data.job_id)) return null;
const labels: Record<string, string> = {
queued: "已提交，等待计算", running: "后台计算中", cancel_requested: "已请求取消，等待计算停止",
succeeded: "计算已完成", failed: "计算失败", cancelled: "计算已取消",
interrupted: "计算已中断", timed_out: "计算执行超时",
};
const label = labels[value.data.status];
return label ? { label, jobId: value.data.job_id } : null;
} catch { return null; }
};
export const getToolProgressTitle = (tool: string, status: string, output?: string) => {
const toolName = toolLabels[tool] ?? tool;
if (status === "completed") {
const calculation = calculationProgress(output);
if (calculation) return `${toolName}：${calculation.label}`;
}
if (status === "completed") return `${toolName} 已完成`;
if (status === "error") return `${toolName} 执行失败`;
if (status === "pending") return `准备调用 ${toolName}`;
src/routes/chatStreamRoute.ts
+13 -6
View File
@@ -1,3 +1,4 @@
import type { TaskCalculationWaiter } from "../mcp/task-calculation-waiter.js";
import { createResponsesProjection, responseRequestSchema, responsePrompt } from "../chat/responses-protocol.js";
import { ResponseInputTracker, WaitingResponseInputQueue } from "../chat/response-inputs.js";
import { reconcileInteractionDelivery } from "../sessions/interaction-delivery-reconciliation.js";
@@ -252,6 +253,7 @@ export const buildFailedTaskExecutionPlan = (history: Record<string, unknown>[])
};
type RegisterChatStreamRouteOptions = {
waitTaskCalculations?: TaskCalculationWaiter;
interactionRequestStore?: InteractionRequestStore;
agentMessages?: AgentMessageCoordinator;
frontendActionCoordinator: FrontendActionCoordinator;
@@ -269,6 +271,7 @@ export const registerChatStreamRoute = (
chatRouter: Router,
{
interactionRequestStore,
waitTaskCalculations,
frontendActionCoordinator,
learningOrchestrator,
memoryStore,
@@ -1005,9 +1008,11 @@ export const registerChatStreamRoute = (
if (messaging && (messaging.hasInputThisTurn(binding.sessionId) || !await messaging.closeInteractive(activeSessionRecord))) return false;
return inputTracker.canClose();
},
beforeIdle: messaging ? async () => {
if (sessionPurpose !== "task_execution") return messaging.closeInteractive(activeSessionRecord);
if (await messaging.flush(binding.sessionId)) return false;
beforeIdle: messaging || (sessionPurpose === "task_execution" && waitTaskCalculations) ? async () => {
if (sessionPurpose !== "task_execution") return messaging ? messaging.closeInteractive(activeSessionRecord) : true;
if (await messaging?.flush(binding.sessionId)) return false;
const calculationResume = await waitTaskCalculations?.(binding.sessionId, abortController.signal);
if (calculationResume) return calculationResume;
await taskArtifactPublishQueue;
if (taskLeaseFailure || !activeSessionRecord.taskId || !activeSessionRecord.runId) return true;
taskCompleting = true;
@@ -1018,13 +1023,13 @@ export const registerChatStreamRoute = (
const task = await operationalTaskClient.completeTaskRun(binding.sessionId,
activeSessionRecord.taskId, activeSessionRecord.runId, clientSessionId);
taskRunFinalized = true;
messaging.stop(binding.sessionId);
messaging?.stop(binding.sessionId);
await stopTaskHeartbeat();
if (task.lifecycle_status === "failed") throw new Error("Server rejected incomplete investigation artifacts");
return true;
} catch (error) {
if (error && typeof error === "object" && "code" in error && error.code === "PENDING_AGENT_MESSAGES") {
await messaging.flush(binding.sessionId);
await messaging?.flush(binding.sessionId);
return false;
}
throw error;
@@ -1048,7 +1053,6 @@ export const registerChatStreamRoute = (
publish("execution_plan", buildFailedTaskExecutionPlan(executionPlanHistory));
}
await taskArtifactPublishQueue;
await stopTaskHeartbeat();
if (
sessionPurpose === "task_execution"
&& activeSessionRecord.taskId
@@ -1068,6 +1072,7 @@ export const registerChatStreamRoute = (
taskRunFinalized = true;
} else {
try {
if (await waitTaskCalculations?.(binding.sessionId, abortController.signal)) throw Object.assign(new Error("Calculation results require investigation continuation"), { code: "CALCULATION_CONTINUATION_REQUIRED" });
const completedTask = await operationalTaskClient.completeTaskRun(
binding.sessionId,
activeSessionRecord.taskId,
@@ -1096,6 +1101,7 @@ export const registerChatStreamRoute = (
}
}
}
await stopTaskHeartbeat();
await persistQueue.catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
});
@@ -1153,6 +1159,7 @@ export const registerChatStreamRoute = (
session_id: clientSessionId,
title: sessionTitle,
});
await stopTaskHeartbeat();
await persistQueue.catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
});
src/runtime/product-mcp-tools.ts
+11
View File
@@ -1,5 +1,8 @@
export const PRODUCT_MCP_READ_TOOLS = [
"get_mutation_status",
"get_calculation_job",
"get_calculation_result",
"list_calculation_jobs",
"export_project_model",
"export_sensor_placement",
"get_analysis_run",
@@ -58,6 +61,7 @@ export const PRODUCT_MCP_ANALYSIS_TOOLS = [
] as const;
export const PRODUCT_MCP_WRITE_TOOLS = [
"cancel_task_run",
"cancel_calculation_job",
"clean_scada_readings",
"create_operational_task",
"create_sensor_placement",
@@ -81,3 +85,10 @@ export const LOCAL_QUERY_TOOLS = ["geocode", "session_search", "list_investigati
export const LOCAL_QUERY_TOOL_IDS = LOCAL_QUERY_TOOLS.map(queryMcpToolId);
export const QUERY_MCP_TOOL_IDS = [...PRODUCT_MCP_TOOL_IDS, ...LOCAL_QUERY_TOOL_IDS];
export const QUERY_MCP_AUTOALLOW_TOOL_IDS = [...PRODUCT_MCP_READ_TOOL_IDS, ...PRODUCT_MCP_ANALYSIS_TOOLS.map(queryMcpToolId), ...LOCAL_QUERY_TOOL_IDS];
export const PRODUCT_MCP_QUEUED_TOOLS = [
"run_burst_location", "run_leakage_identification", "run_burst_detection",
"run_simulation", "run_burst_simulation", "run_contaminant_simulation",
"analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation",
"optimize_sensor_placement", "save_simulation_scenario", "create_sensor_placement",
] as const;
src/server.ts
+16
View File
@@ -2,6 +2,8 @@ import { InteractionRequestStore } from "./sessions/interaction-request-store.js
import { registerChatInteractionBridge } from "./routes/chat-interaction-bridge.js";
import { writeQueryMcpConnection } from "./runtime/query-mcp-connection.js";
import { startQueryMcpServer } from "./mcp/query-server.js";
import { resolveMatchingNativeCall } from "./mcp/native-call-identity.js";
import { createTaskCalculationWaiter } from "./mcp/task-calculation-waiter.js";
import { createLocalQueryExecutor } from "./mcp/local-query-tools.js";
import { createProductMcpClient } from "./product/mcp-query-client.js";
import { randomUUID } from "node:crypto";
@@ -225,6 +227,10 @@ app.get("/health", async (_req, res) => {
app.use(
"/api/v1",
buildChatRouter({
waitTaskCalculations: productMcpClient ? createTaskCalculationWaiter(productMcpClient, (id) => {
const context = getRuntimeSessionContext(id);
return context ? { context, runtimeSessionId: id, taskId: context.taskId } : null;
}) : undefined,
interactionRequestStore,
agentMessages,
frontendActionCoordinator,
@@ -284,6 +290,16 @@ const warmupOpencodeRuntime = async () => {
const sandboxProbe = await probeLandlockSandbox();
const queryMcp = startQueryMcpServer({
resolveNativeCall: async (sessionId, toolName, args) => {
// Native tool state can arrive just after the MCP request. Only use a
// unique current call from host-read history, never a model-supplied key.
for (let attempt = 0; attempt < 3; attempt++) {
const callId = resolveMatchingNativeCall(await opencodeRuntime.messages(sessionId, 20), toolName, args);
if (callId) return callId;
if (attempt < 2) await new Promise((resolve) => setTimeout(resolve, 100));
}
return null;
},
artifactWorkspace: (context) => opencodeRuntime.ensureSessionWorkspace(context.sessionId, context),
internalToken, productClient: productMcpClient,
productTools: productMcpClient ? await productMcpClient.prepare() : [],
tests/frontendAction/callIdPlugin.test.ts
+9
View File
@@ -1,13 +1,22 @@
import { afterAll, beforeAll, describe, expect, it } from "bun:test";
import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import plugin from "../../.opencode/plugins/tjwater-tools.js";
const originalConnectionPath = process.env.TJWATER_AGENT_MCP_CONNECTION_PATH;
const connectionRoot = mkdtempSync(join(tmpdir(), "call-id-plugin-"));
const originalInternalToken = process.env.AGENT_INTERNAL_TOKEN;
beforeAll(() => {
process.env.TJWATER_AGENT_MCP_CONNECTION_PATH = join(connectionRoot, "absent.json");
process.env.AGENT_INTERNAL_TOKEN = "test-internal-token";
});
afterAll(() => {
if (originalConnectionPath === undefined) delete process.env.TJWATER_AGENT_MCP_CONNECTION_PATH;
else process.env.TJWATER_AGENT_MCP_CONNECTION_PATH = originalConnectionPath;
rmSync(connectionRoot, { recursive: true, force: true });
if (originalInternalToken === undefined) delete process.env.AGENT_INTERNAL_TOKEN;
else process.env.AGENT_INTERNAL_TOKEN = originalInternalToken;
});
tests/mcp/calculationProgress.test.ts
+16
View File
@@ -0,0 +1,16 @@
import { expect, it } from "bun:test";
import { buildToolProgressDetail, getToolProgressTitle } from "../../src/routes/chatStreamEvents.js";
it("distinguishes completed submissions from completed calculations", () => {
const jobId = "f40ca28d-1ea4-4ed7-8e1a-4cc476e748e8";
const output = (status: string) => JSON.stringify({ semantics: "task", data: { job_id: jobId, status } });
const detail = buildToolProgressDetail("run_burst_location", "completed", {}, "", undefined, output("queued"));
expect(detail).toContain("等待计算");
expect(detail).toContain(jobId);
expect(detail).not.toContain("已执行完成");
expect(getToolProgressTitle("run_burst_location", "completed", output("running"))).toContain("后台计算中");
expect(getToolProgressTitle("cancel_calculation_job", "completed", output("cancel_requested"))).toContain("等待计算停止");
expect(getToolProgressTitle("get_calculation_job", "completed", output("failed"))).toContain("计算失败");
expect(getToolProgressTitle("query", "completed", "not-json")).toBe("query 已完成");
expect(buildToolProgressDetail("query", "completed", {}, "")).toContain("已执行完成");
});
tests/mcp/nativeCallIdentity.test.ts
+42
View File
@@ -0,0 +1,42 @@
import { fromJsonSchema } from "@modelcontextprotocol/server";
import { normalizeRuntimeMessage } from "../../src/runtime/opencode2-events.js";
import { expect, it } from "bun:test";
import { resolveMatchingNativeCall } from "../../src/mcp/native-call-identity.js";
import { allowedQueryTools } from "../../src/runtime/query-mcp-policy.js";
import { QUERY_MCP_AUTOALLOW_TOOL_IDS } from "../../src/runtime/product-mcp-tools.js";
const part = (callID: string, status = "running", input: unknown = { duration: 60 }) => ({ type: "tool", tool: "mcp_tjwater_run_simulation", callID, state: { status, input } });
it("resolves only an exact active persisted native call and fails closed on ambiguity", () => {
expect(resolveMatchingNativeCall([{ parts: [part("old", "completed"), part("current")] }], "run_simulation", { duration: 60 })).toBe("current");
expect(resolveMatchingNativeCall([{ parts: [part("current")] }], "run_simulation", { duration: 61 })).toBeNull();
expect(resolveMatchingNativeCall([{ parts: [part("a"), part("b")] }], "run_simulation", { duration: 60 })).toBeNull();
});
it("allows readers only job reads and never autoallows cancellation", () => {
const context = { sessionId: "s", clientSessionId: "c", actorKey: "a", projectKey: "p", projectId: "p", traceId: "t", purpose: "task_execution" as const };
const tools = allowedQueryTools(context, true);
for (const name of ["get_calculation_job", "get_calculation_result", "list_calculation_jobs"]) expect(tools).toContain(`mcp_tjwater_${name}`);
expect(tools).not.toContain("mcp_tjwater_cancel_calculation_job");
expect(tools).not.toContain("mcp_tjwater_run_simulation");
expect(QUERY_MCP_AUTOALLOW_TOOL_IDS).not.toContain("mcp_tjwater_cancel_calculation_job");
});
it("matches omitted MCP defaults against the actual native message-list normalizer", async () => {
const schema = fromJsonSchema<Record<string, unknown>>({
type: "object", additionalProperties: false,
properties: { timeout_seconds: { type: "integer", default: 120 }, dry_run: { type: "boolean", default: true } },
});
for (const input of [{}, { timeout_seconds: 900 }, { dry_run: false }]) {
const validated = await schema["~standard"].validate(input);
expect(validated.issues).toBeUndefined();
if (!("value" in validated)) throw new Error("MCP input validation failed");
expect(validated.value).toEqual(input);
const native = {
id: "assistant-message", type: "assistant", time: { created: 1 },
content: [{ type: "tool", id: "native-call", name: "mcp_tjwater_run_simulation", time: { created: 1 }, state: { status: "running", input } }],
} as unknown as Parameters<typeof normalizeRuntimeMessage>[0];
const messages = [normalizeRuntimeMessage(native, "native-session")];
expect(resolveMatchingNativeCall(messages, "run_simulation", validated.value!)).toBe("native-call");
expect(resolveMatchingNativeCall(messages, "run_simulation", { ...validated.value, arbitrary_parameter: true })).toBeNull();
}
});
tests/mcp/queryServer.test.ts
+26 -2
View File
@@ -23,7 +23,7 @@ it("advertises simulation budgets and derives the native connection ceiling from
}, additionalProperties: false } };
const { server, client, call, query } = await setup({ productTools: [simulation] });
setRuntimeSessionContext(owner);
expect(server.executionTimeoutMs).toBe(3_665_000);
expect(server.executionTimeoutMs).toBe(35_000);
expect((await client.listTools()).tools[0].inputSchema.properties?.timeout_seconds).toMatchObject({ maximum: 3600 });
expect((await call("run_simulation", { timeout_seconds: 900 })).isError).not.toBeTrue();
expect((query.mock.calls[0] as unknown as unknown[])[3]).toEqual({ timeout_seconds: 900 });
@@ -45,7 +45,7 @@ it("discovers calculations natively but denies verified readers before upstream
async function setup(extra: Partial<Parameters<typeof startQueryMcpServer>[0]> = {}) {
const executeLocal = mock(async () => ({ safe: true }));
const query = mock(async () => ({ content: [{ type: "text" as const, text: "evidence" }], structuredContent: { evidence_id: "E1", data: { value: 1 } } }));
const server = startQueryMcpServer({ internalToken: "MCP_TEST_ONLY", productClient: { query }, productTools: [productTool], resolveReader: async () => null, executeLocal, ...extra });
const server = startQueryMcpServer({ internalToken: "MCP_TEST_ONLY", productClient: { query }, productTools: [productTool], resolveNativeCall: async () => "native-call-1", resolveReader: async () => null, executeLocal, ...extra });
servers.push(server);
const client = new Client({ name: "query-test", version: "1" }, { versionNegotiation: { mode: { pin: "2026-07-28" } } });
clients.push(client);
@@ -159,3 +159,27 @@ it("replaces unavailable evidence resources with the authorized paging tool with
data: { evidence_read: { tool: "mcp_tjwater_get_task_evidence", arguments: { evidence_id: "E1", limit: 20, offset: 0 } } } });
expect(upstream.structuredContent.data.resource_uri).toBe(uri);
});
it("keeps the same host request identity across reconnect retries and rejects missing identity", async () => {
setRuntimeSessionContext(owner);
const { call, query } = await setup();
await call("get_task_context");
await call("get_task_context");
const ids = query.mock.calls.map((args) => (args as unknown as unknown[])[4]);
expect(ids[0]).toMatch(/^[a-f0-9]{64}$/);
expect(ids[1]).toBe(ids[0]);
const missing = await setup({ productTools: [{ name: "run_simulation", inputSchema: { type: "object", properties: {} } }], resolveNativeCall: async () => null });
expect((await missing.call("run_simulation")).isError).toBeTrue();
expect(missing.query).not.toHaveBeenCalled();
});
it("distinguishes safe network and timeout failures without upstream text", async () => {
setRuntimeSessionContext(owner);
for (const [cause, message] of [[new DOMException("SECRET", "TimeoutError"), "Query timed out"], [new TypeError("SECRET"), "Query network failure"]] as const) {
const { call } = await setup({ productClient: { query: async () => { throw cause; } } });
const result = await call("get_task_context");
expect(JSON.stringify(result)).toContain(message);
expect(JSON.stringify(result)).not.toContain("SECRET");
}
});
tests/mcp/queryTimeouts.test.ts
+12 -4
View File
@@ -7,8 +7,8 @@ const simulation: Tool = { name: "run_simulation", inputSchema: { type: "object"
} } };
it("budgets calculations independently from their simulated duration and from ordinary reads", () => {
expect(queryTimeouts(simulation, { duration: 86400, timeout_seconds: 900 })).toEqual({ executionMs: 900_000, callMs: 905_000, requestMs: 960_000 });
expect(queryTimeouts(simulation, {})).toEqual({ executionMs: 120_000, callMs: 125_000, requestMs: 180_000 });
expect(queryTimeouts(simulation, { duration: 86400, timeout_seconds: 900 })).toEqual({ executionMs: 900_000, callMs: 30_000, requestMs: 35_000 });
expect(queryTimeouts(simulation, {})).toEqual({ executionMs: 120_000, callMs: 30_000, requestMs: 35_000 });
expect(queryTimeouts({ name: "list_elements", inputSchema: { type: "object" } }, {})).toEqual({ executionMs: 30_000, callMs: 30_000, requestMs: 35_000 });
});
@@ -21,6 +21,14 @@ it("honors the effective advertised maximum and rejects invalid per-call budgets
} } };
expect(queryTimeouts(limited, { timeout_seconds: 120 }).executionMs).toBe(120_000);
expect(() => queryTimeouts(limited, { timeout_seconds: 121 })).toThrow();
expect(nativeQueryTimeout([simulation, limited])).toBe(3_665_000);
expect(nativeQueryTimeout([limited])).toBe(185_000);
expect(nativeQueryTimeout([simulation, limited])).toBe(35_000);
expect(nativeQueryTimeout([limited])).toBe(35_000);
});
it("uses the catalog's thirty-minute burst-location default as a worker execution budget", () => {
const location: Tool = { name: "run_burst_location", inputSchema: { type: "object", properties: {
timeout_seconds: { type: "integer", minimum: 1, maximum: 3600, default: 1800 },
} } };
expect(queryTimeouts(location, {})).toEqual({ executionMs: 1_800_000, callMs: 30_000, requestMs: 35_000 });
expect(nativeQueryTimeout([location])).toBe(35_000);
});
tests/mcp/taskCalculationWaiter.test.ts
+38
View File
@@ -0,0 +1,38 @@
import { expect, it, mock } from "bun:test";
import { createTaskCalculationWaiter } from "../../src/mcp/task-calculation-waiter.js";
const owner = { context: { sessionId: "s", clientSessionId: "c", actorKey: "a", projectKey: "p", projectId: "p", traceId: "t", purpose: "task_execution" as const }, runtimeSessionId: "s" };
const result = (status: string) => ({ content: [], structuredContent: { data: [{ job_id: "job-1", status, poll_after_ms: 1500 }], page: { has_more: false } } });
it("waits in the host across bounded status reads then resumes once for terminal results", async () => {
const outputs = [result("queued"), result("running"), result("succeeded"), result("succeeded")];
const query = mock(async () => outputs.shift()!);
const pause = mock(async (_ms: number, _signal: AbortSignal) => {});
const wait = createTaskCalculationWaiter({ query }, () => owner, pause);
const signal = new AbortController().signal;
expect(await wait("s", signal)).toContain("get_calculation_result");
expect(pause).toHaveBeenCalledTimes(2);
expect(pause.mock.calls[0]?.[0]).toBe(1500);
expect(await wait("s", signal)).toBeNull();
expect(query).toHaveBeenCalledTimes(4);
});
it("does not turn unavailable status into a completed task", async () => {
const query = mock(async () => ({ content: [], isError: true }));
await expect(createTaskCalculationWaiter({ query }, () => owner)("s", new AbortController().signal)).rejects.toThrow("unavailable");
});
it("aborts host waiting without issuing Server cancellation", async () => {
const controller = new AbortController();
const names: string[] = [];
const query = mock(async (_id: string, _owner: unknown, name: string) => { names.push(name); return result("running"); });
const wait = createTaskCalculationWaiter({ query }, () => owner, async () => { controller.abort(); });
await expect(wait("s", controller.signal)).rejects.toThrow();
expect(names).toEqual(["list_calculation_jobs"]);
});
it("resumes recovery for a worker-interrupted job instead of waiting or completing silently", async () => {
const query = mock(async () => result("interrupted"));
const pause = mock(async (_ms: number, _signal: AbortSignal) => {});
const wait = createTaskCalculationWaiter({ query }, () => owner, pause);
const resume = await wait("s", new AbortController().signal);
expect(resume).toContain("interrupted");
expect(resume).toContain("不要重复提交");
expect(pause).not.toHaveBeenCalled();
});
tests/opencode/queryMcp.test.ts
+3 -1
View File
@@ -19,7 +19,9 @@ it("configures a pinned native MCP connection using the host-owned endpoint and
writeQueryMcpConnection(join(root, "sessions"), "http://127.0.0.1:12345/mcp");
expect(queryMcpConnection()).toMatchObject({ type: "remote", protocol: "2026-07-28", oauth: false, codemode: false, headers: { "x-agent-internal-token": "MCP_TEST_ONLY" } });
writeQueryMcpConnection(join(root, "sessions"), "http://127.0.0.1:12345/mcp", 3_665_000);
expect(queryMcpConnection()?.timeout?.execution).toBe(3_665_000);
expect(() => queryMcpConnection()).toThrow("timeout");
writeQueryMcpConnection(join(root, "sessions"), "http://127.0.0.1:12345/mcp", 35_000);
expect(queryMcpConnection()?.timeout?.execution).toBe(35_000);
writeQueryMcpConnection(join(root, "sessions"), "http://127.0.0.1:12345/mcp", Infinity);
expect(() => queryMcpConnection()).toThrow("timeout");
writeQueryMcpConnection(join(root, "sessions"), "https://foreign.test/mcp");
tests/product/mcpQueryClient.test.ts
+2 -2
View File
@@ -45,8 +45,8 @@ it("forwards the selected simulation budget through HTTP and SDK while rejecting
await expect(client.query(id, owner(id), "run_simulation", { timeout_seconds: 3601 }, "excess")).rejects.toMatchObject({ status: 400 });
expect(requests).toBe(0);
const result = await client.query(id, owner(id), "run_simulation", { timeout_seconds: 900 }, "long");
expect(transportTimeout).toBe(905_000);
expect(sdkTimeout).toBe(905_000);
expect(transportTimeout).toBe(30_000);
expect(sdkTimeout).toBe(30_000);
expect(result.structuredContent).toMatchObject({ arguments: { timeout_seconds: 900, context_handle: "x".repeat(43) } });
});
tests/routes/chatStream.test.ts
+23
View File
@@ -45,6 +45,29 @@ describe("streamPromptResponse", () => {
expect(emitted.filter((event) => event === "done")).toHaveLength(1);
});
it("resumes the same investigation with host-provided calculation result guidance", async () => {
const prompts: string[] = [];
let idleChecks = 0;
const runtime = {
prompt: async (_sessionId: string, text: string) => { prompts.push(text); },
subscribeEvents: async () => createEventStream([1, 2].flatMap((round) => [
{ type: "message.part.updated", properties: { sessionID: "s", part: {
id: `calculation-text-${round}`, type: "text", text: `round ${round}`, sessionID: "s", messageID: `m-${round}`,
} } },
{ type: "session.idle", properties: { sessionID: "s" } },
])),
messages: async () => [],
} as unknown as OpencodeRuntimeAdapter;
const result = await streamPromptResponse({
runtime, sessionId: "s", clientSessionId: "client", message: "investigate", suppressAnswerUi: true,
beforeIdle: async () => ++idleChecks === 1 ? "Read get_calculation_result for job-1" : true,
write: () => {},
});
expect(result.failed).toBeFalse();
expect(prompts).toEqual(["investigate", "Read get_calculation_result for job-1"]);
expect(idleChecks).toBe(2);
});
it("keeps the first pending event when the V2 prompt acknowledgement arrives first", async () => {
const emitted: { event: string; data: Record<string, unknown> }[] = [];
const runtime = {