diff --git a/apps/server/src/mcp/OrchestratorMcpService.test.ts b/apps/server/src/mcp/OrchestratorMcpService.test.ts index 6450f9acac2cb..d1d6108b7fc9d 100644 --- a/apps/server/src/mcp/OrchestratorMcpService.test.ts +++ b/apps/server/src/mcp/OrchestratorMcpService.test.ts @@ -24,6 +24,7 @@ import { } from "@t3tools/contracts"; import { checkpointWorkspace } from "@t3tools/provider-testing/replayWorkspace"; import * as DateTime from "effect/DateTime"; +import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; import * as Fiber from "effect/Fiber"; import * as Layer from "effect/Layer"; @@ -1912,6 +1913,157 @@ describe("OrchestratorMcpService provider resolution", () => { ), ); + it.effect.each(["local", "remote"] as const)( + "a blocking %s delegation cancelled mid-wait wakes its parent like an async task", + (placement) => + Effect.gen(function* () { + const driver = ProviderDriverKind.make("codex"); + const layerRegistry = ProviderAdapterRegistry.layerFromAdapters([ + { + instanceId: codexInstanceId, + driver, + getCapabilities: () => Effect.succeed(CodexProviderCapabilitiesV2), + planSelectionTransition: () => Effect.succeed({ type: "apply_on_next_turn" as const }), + openSession: () => Effect.die("The child never starts: the effect worker is off."), + }, + ]); + const layerDatabase = SqlitePersistence.layerMemory; + const layerOrchestrator = ProviderReplayHarness.layerWithRegistry( + { name: "mcp-wait-cancelled" }, + layerRegistry, + { databaseLayer: layerDatabase, runEffectWorker: false }, + ); + const layerSettings = ServerSettings.layerTest({ + providerInstances: { [codexInstanceId]: { driver: "codex", enabled: true } }, + }); + const layerThreads = ThreadManagementService.layer.pipe( + Layer.provideMerge(layerOrchestrator), + Layer.provide( + LegacyV1ThreadImporter.layer.pipe( + Layer.provide(layerOrchestrator), + Layer.provide(layerDatabase), + ), + ), + Layer.provide(CommandReceiptStore.layer.pipe(Layer.provide(layerDatabase))), + Layer.provideMerge(layerSettings), + ); + const waiting = yield* Deferred.make(); + const layerWatchedThreads = Layer.effect( + ThreadManagementService.ThreadManagementService, + Effect.gen(function* () { + const threads = yield* ThreadManagementService.ThreadManagementService; + return { + ...threads, + streamStoredEventsFrom: (input) => + Stream.unwrap( + Deferred.succeed(waiting, undefined).pipe( + Effect.as(threads.streamStoredEventsFrom(input)), + ), + ), + } satisfies ThreadManagementService.ThreadManagementService["Service"]; + }), + ).pipe(Layer.provideMerge(layerThreads)); + const remoteChild = { + environmentId: EnvironmentId.make("environment:mcp-providers-box"), + threadId: ThreadId.make("thread:mcp-providers-on-the-box"), + label: "Box", + }; + // Records the task as the real delegation does once the linked environment launched it. + const layerRemoteDelegation = Layer.unwrap( + Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + return Layer.mock(RemoteDelegation.RemoteDelegation)({ + delegate: (_scope, input) => + Effect.gen(function* () { + const run = (yield* orchestrator.getThreadProjection(parentThreadId)).runs[0]!; + yield* orchestrator.dispatch({ + type: "delegated_task.remote.request", + commandId: CommandId.make("command:remote-delegate:wait-cancelled"), + parentThreadId, + parentRunId: run.id, + parentNodeId: run.rootNodeId!, + task: input.task, + driver: ProviderDriverKind.make("remote"), + modelSelection: run.modelSelection, + remoteChild, + completionWake: "settled_only", + }); + const parent = yield* orchestrator.getThreadProjection(parentThreadId); + return { taskId: parent.subagents[0]!.id }; + }).pipe(Effect.orDie), + }); + }), + ).pipe(Layer.provide(layerWatchedThreads)); + const layerTest = layerOrchestratorService.pipe( + Layer.provideMerge(layerWatchedThreads), + Layer.provide( + Layer.mergeAll( + NodeServices.layer, + layerRegistry, + layerSettings, + providerRegistryLayer([ + providerSnapshot({ instanceId: codexInstanceId, driver, model: "gpt-5.4" }), + ]), + Layer.mock(ProjectService.ProjectService)({}), + Layer.mock(SecretRequests.SecretRequests)({}), + layerRemoteDelegation, + Layer.mock(ScheduledTaskService.ScheduledTaskService)({}), + ), + ), + ); + yield* Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const service = yield* OrchestratorMcpService.OrchestratorMcpService; + yield* orchestrator.dispatch({ + type: "thread.create", + createdBy: "user", + creationSource: "web", + commandId: CommandId.make("parent:create"), + threadId: parentThreadId, + projectId, + title: "MCP parent", + modelSelection: { instanceId: codexInstanceId, model: "gpt-5.4" }, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: null, + }); + yield* orchestrator.dispatch({ + type: "message.dispatch", + createdBy: "user", + creationSource: "web", + commandId: CommandId.make("parent:start"), + threadId: parentThreadId, + messageId: MessageId.make("parent:message"), + text: "Delegate a review and wait for it.", + attachments: [], + dispatchMode: { type: "start_immediately" }, + }); + const call = yield* service + .delegateTask(scope, { + task: "Review the change.", + ...(placement === "remote" + ? { target: { environmentId: remoteChild.environmentId } } + : {}), + mode: "wait", + clientRequestId: "delegate-wait-cancelled", + }) + .pipe(Effect.forkChild); + yield* Deferred.await(waiting); + // A steered Claude parent cancels the tools/call, which interrupts this fiber. + yield* Fiber.interrupt(call); + const parent = yield* orchestrator.getThreadProjection(parentThreadId); + const task = parent.subagents.find((subagent) => subagent.origin === "app_owned"); + assert.equal(task?.status, "running"); + assert.equal( + task?.completionWake, + "always", + "a cancelled blocking delegation left the task settled_only, so its result never reaches the still-live parent run", + ); + }).pipe(Effect.provide(layerTest)); + }), + ); + it.effect( "inherits an available parent instance for driver-only targets and otherwise selects a healthy peer", () => diff --git a/apps/server/src/mcp/OrchestratorMcpService.ts b/apps/server/src/mcp/OrchestratorMcpService.ts index d8dacb35a0617..e606edef8caea 100644 --- a/apps/server/src/mcp/OrchestratorMcpService.ts +++ b/apps/server/src/mcp/OrchestratorMcpService.ts @@ -1600,33 +1600,59 @@ const make = Effect.gen(function* () { return response; }); - /** `mode=wait` on a task already recorded, as `delegateTask` does for its own. */ + /** + * Blocks for `mode=wait`. A wait that ends without a terminal result, including an + * interrupted tools/call, upgrades the wake so a later result still reaches the parent. + */ const awaitTask = Effect.fn("OrchestratorMcpService.awaitTask")(function* ( scope: McpThreadInvocationScope, taskId: NodeId, + wakePolicyRequestKey: string, requestedTimeoutMs: number | undefined, ): Effect.fn.Return { const timeoutMs = Math.min( MAX_WAIT_TIMEOUT_MS, Math.max(1, requestedTimeoutMs ?? DEFAULT_WAIT_TIMEOUT_MS), ); - const waited = yield* waitForTask(scope, taskId, timeoutMs); - if (Option.isSome(waited)) return waited.value; - // As for a local task: once the blocking wait ends, a later result wakes the parent. - yield* threadManagement + const upgradeWake = threadManagement .dispatch({ type: "delegated_task.wake-policy", commandId: stableCommandId({ scope, - requestKey: taskId, + requestKey: wakePolicyRequestKey, operation: "delegate-task-wake-policy", }), parentThreadId: scope.thread.threadId, taskId, completionWake: "always", }) - .pipe(Effect.ignoreCause({ log: true })); - return yield* readTask(scope, taskId, true, true); + .pipe( + Effect.catch((error) => + Effect.logWarning("orchestrator-mcp.delegate-task.wake-policy-failed", { + taskId, + outcome: + error._tag === "OrchestratorCommandPreviouslyRejectedError" + ? "previously_rejected" + : "dispatch_failed", + error, + }), + ), + Effect.catchCause((cause) => + Effect.logWarning("orchestrator-mcp.delegate-task.wake-policy-failed", { + taskId, + outcome: "defect", + cause, + }), + ), + ); + const waited = yield* waitForTask(scope, taskId, timeoutMs).pipe( + Effect.onExit((exit) => (Exit.isSuccess(exit) ? Effect.void : upgradeWake)), + ); + if (Option.isSome(waited) && isTerminalTaskStatus(waited.value.status)) { + return waited.value; + } + yield* upgradeWake; + return yield* readTask(scope, taskId, Option.isNone(waited), true); }); /** @@ -2190,7 +2216,7 @@ const make = Effect.gen(function* () { target: { ...input.target, environmentId }, }); return yield* input.mode === "wait" - ? awaitTask(scope, taskId, input.timeoutMs) + ? awaitTask(scope, taskId, taskId, input.timeoutMs) : readTask(scope, taskId, false, true); } const { scope, parent } = yield* loadThreadCaller(callerScope, "delegate_task"); @@ -2288,8 +2314,8 @@ const make = Effect.gen(function* () { interactionMode, ...(workspaceStrategy === undefined ? {} : { workspaceStrategy }), // Async delegations wake the parent on every child terminal; wait - // delegations deliver through the blocking tool call, so a wake is - // only needed if the parent settled first (timeout, disconnect). + // delegations deliver through the blocking tool call until `awaitTask` + // gives up the wait (timeout, cancelled call) and upgrades the wake. completionWake: input.mode === "wait" ? "settled_only" : "always", }) .pipe( @@ -2329,51 +2355,9 @@ const make = Effect.gen(function* () { }); } - if (input.mode !== "wait") { - return yield* readTask(scope, taskId, false, true); - } - const timeoutMs = Math.min( - MAX_WAIT_TIMEOUT_MS, - Math.max(1, input.timeoutMs ?? DEFAULT_WAIT_TIMEOUT_MS), - ); - const waited = yield* waitForTask(scope, taskId, timeoutMs); - if (Option.isSome(waited) && isTerminalTaskStatus(waited.value.status)) { - return waited.value; - } - // An incomplete wait gives up delivery: upgrade so a later terminal - // wakes the parent even mid-turn. - yield* threadManagement - .dispatch({ - type: "delegated_task.wake-policy", - commandId: stableCommandId({ - scope, - requestKey: key, - operation: "delegate-task-wake-policy", - }), - parentThreadId: scope.thread.threadId, - taskId, - completionWake: "always", - }) - .pipe( - Effect.catch((error) => - Effect.logWarning("orchestrator-mcp.delegate-task.wake-policy-failed", { - taskId, - outcome: - error._tag === "OrchestratorCommandPreviouslyRejectedError" - ? "previously_rejected" - : "dispatch_failed", - error, - }), - ), - Effect.catchCause((cause) => - Effect.logWarning("orchestrator-mcp.delegate-task.wake-policy-failed", { - taskId, - outcome: "defect", - cause, - }), - ), - ); - return yield* readTask(scope, taskId, Option.isNone(waited), true); + return yield* input.mode === "wait" + ? awaitTask(scope, taskId, key, input.timeoutMs) + : readTask(scope, taskId, false, true); }), taskStatus: Effect.fn("OrchestratorMcpService.taskStatus")(function* (callerScope, taskId) { const scope = yield* requireThreadScope(callerScope, "task_status");