feat(mcp): resume queued calculation jobs
This commit is contained in:
@@ -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");
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user