feat(mcp): resume queued calculation jobs

This commit is contained in:
2026-10-29 18:25:22 +08:00
parent 52681a8451
commit 799b99f310
20 changed files with 348 additions and 42 deletions
+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");
});