History

fix: preserve orchestration task identity

Julius Marminge committed 3 months ago

Commit b4bc3291 parent b65e09c4 files · +134 −5Browse files

Unified · Split

4 files
  1. apps/server/src/mcp/OrchestratorMcpService.ts +24 −2
  2. apps/server/src/mcp/OrchestratorMcpToolkit.integration.test.ts +52 −2
  3. apps/server/src/provider/acp/AcpJsonRpcConnection.test.ts +51 −0
  4. apps/server/src/provider/acp/AcpSessionRuntime.ts +7 −1
apps/server/src/mcp/OrchestratorMcpService.ts+24−2
@@ -5,6 +5,7 @@ import { · 4 unmodified lines
5 type ModelSelection,5 type ModelSelection,
6 NodeId,6 NodeId,
7 type OrchestrationV2Run,7 type OrchestrationV2Run,
8 type OrchestrationV2Subagent,
8 type OrchestrationV2ThreadProjection,9 type OrchestrationV2ThreadProjection,
9 type OrchestrationV2ThreadShell,10 type OrchestrationV2ThreadShell,
10 type OrchestrationV2TurnItem,11 type OrchestrationV2TurnItem,
@@ -279,6 +280,25 @@ function taskStatusForRun( · 268 unmodified lines
279 }280 }
280}281}
281 282
283function delegatedTaskRun(
284 childProjection: OrchestrationV2ThreadProjection,
285 task: OrchestrationV2Subagent,
286): OrchestrationV2Run | undefined {
287 const spawnTransfer = childProjection.contextTransfers.find(
288 (transfer) =>
289 transfer.type === "subagent_spawn" &&
290 transfer.sourceThreadId === task.threadId &&
291 transfer.targetThreadId === task.childThreadId,
292 );
293 if (spawnTransfer === undefined) {
294 // Legacy delegated-task projections predate the durable spawn transfer.
295 return latestRun(childProjection);
296 }
297 return spawnTransfer.targetRunId === null
298 ? undefined
299 : childProjection.runs.find((run) => run.id === spawnTransfer.targetRunId);
300}
301
282function isTerminalTaskStatus(302function isTerminalTaskStatus(
283 status: OrchestratorMcpDelegateTaskResult["status"],303 status: OrchestratorMcpDelegateTaskResult["status"],
284): status is TerminalTaskStatus {304): status is TerminalTaskStatus {
@@ -773,7 +793,7 @@ const make = Effect.gen(function* () { · 488 unmodified lines
773 );793 );
774 }794 }
775 const childProjection = yield* loadProjection(task.childThreadId);795 const childProjection = yield* loadProjection(task.childThreadId);
776 const childRun = latestRun(childProjection);796 const childRun = delegatedTaskRun(childProjection, task);
777 const status = taskStatusForRun(childRun);797 const status = taskStatusForRun(childRun);
778 const derivedResult =798 const derivedResult =
779 task.result !== null799 task.result !== null
@@ -1139,7 +1159,9 @@ const make = Effect.gen(function* () { · 359 unmodified lines
1139 } satisfies OrchestratorMcpTaskCancelResult;1159 } satisfies OrchestratorMcpTaskCancelResult;
1140 }1160 }
1141 const child = yield* loadProjection(current.childThreadId);1161 const child = yield* loadProjection(current.childThreadId);
1142 const activeRun = child.runs.find(isActiveRun);1162 const activeRun = child.runs.find(
1163 (run) => run.id === current.childRunId && isActiveRun(run),
1164 );
1143 if (activeRun === undefined) {1165 if (activeRun === undefined) {
1144 return yield* failure(1166 return yield* failure(
1145 "task_not_cancellable",1167 "task_not_cancellable",
apps/server/src/mcp/OrchestratorMcpToolkit.integration.test.ts+52−2
@@ -408,7 +408,7 @@ describe("orchestrator MCP toolkit", () => { · 407 unmodified lines
408 driver: ProviderDriverKind.make("claudeAgent"),408 driver: ProviderDriverKind.make("claudeAgent"),
409 capabilities: ClaudeProviderCapabilitiesV2,409 capabilities: ClaudeProviderCapabilitiesV2,
410 capturedTurns,410 capturedTurns,
411 shouldComplete: () => true,411 shouldComplete: (turn) => turn.message.text !== cancellationPrompt,
412 response: (turn) =>412 response: (turn) =>
413 turn.message.text === delegatedPrompt413 turn.message.text === delegatedPrompt
414 ? delegatedResult414 ? delegatedResult
@@ -845,10 +845,60 @@ describe("orchestrator MCP toolkit", () => { · 430 unmodified lines
845 delegatedStatusAfterFollowupCall.structuredContent,845 delegatedStatusAfterFollowupCall.structuredContent,
846 ).pipe(Effect.orDie);846 ).pipe(Effect.orDie);
847 expect(delegatedStatusAfterFollowup).toMatchObject({847 expect(delegatedStatusAfterFollowup).toMatchObject({
848 childRunId: childFollowup.runId,848 childRunId: delegated.childRunId,
849 status: "completed",849 status: "completed",
850 summary: delegatedResult,
850 });851 });
851 852
853 const activeChildFollowupCall = yield* invoke("t3_thread_send", {
854 threadId: delegated.childThreadId,
855 message: cancellationPrompt,
856 clientRequestId: "delegated-child-active-followup-1",
857 });
858 const activeChildFollowup = yield* decodeThreadSendResult(
859 activeChildFollowupCall.structuredContent,
860 ).pipe(Effect.orDie);
861 yield* waitForProjection(orchestrator, delegated.childThreadId, (projection) =>
862 projection.runs.some(
863 (run) => run.id === activeChildFollowup.runId && run.status === "running",
864 ),
865 );
866 const completedTaskCancelCall = yield* invoke("task_cancel", {
867 taskId: delegated.taskId,
868 reason: "Must not interrupt a later unrelated child run.",
869 clientRequestId: "cancel-completed-delegated-task-1",
870 });
871 const completedTaskCancel = yield* decodeTaskCancelResult(
872 completedTaskCancelCall.structuredContent,
873 ).pipe(Effect.orDie);
874 expect(completedTaskCancel).toEqual({
875 taskId: delegated.taskId,
876 status: "completed",
877 });
878 expect(
879 (yield* orchestrator.getThreadProjection(delegated.childThreadId)).runs.find(
880 (run) => run.id === activeChildFollowup.runId,
881 )?.status,
882 ).toBe("running");
883 const activeChildCleanupCall = yield* invoke("t3_thread_interrupt", {
884 threadId: delegated.childThreadId,
885 runId: activeChildFollowup.runId,
886 reason: "Clean up the active follow-up after verifying task cancellation isolation.",
887 clientRequestId: "interrupt-delegated-child-followup-1",
888 });
889 const activeChildCleanup = yield* decodeThreadInterruptResult(
890 activeChildCleanupCall.structuredContent,
891 ).pipe(Effect.orDie);
892 expect(activeChildCleanup).toMatchObject({
893 runId: activeChildFollowup.runId,
894 status: "interrupt_requested",
895 });
896 yield* waitForProjection(orchestrator, delegated.childThreadId, (projection) =>
897 projection.runs.some(
898 (run) => run.id === activeChildFollowup.runId && run.status === "interrupted",
899 ),
900 );
901
852 // A wait-mode child (completionWake settled_only) that completes902 // A wait-mode child (completionWake settled_only) that completes
853 // while the parent run is live does not offer a wake: the903 // while the parent run is live does not offer a wake: the
854 // blocking delegate_task call above already returned the result.904 // blocking delegate_task call above already returned the result.
apps/server/src/provider/acp/AcpJsonRpcConnection.test.ts+51−0
@@ -6,6 +6,7 @@ import * as NodeFS from "node:fs"; · 5 unmodified lines
6 6
7import * as NodeServices from "@effect/platform-node/NodeServices";7import * as NodeServices from "@effect/platform-node/NodeServices";
8import { it } from "@effect/vitest";8import { it } from "@effect/vitest";
9import * as Deferred from "effect/Deferred";
9import * as Effect from "effect/Effect";10import * as Effect from "effect/Effect";
10import * as Fiber from "effect/Fiber";11import * as Fiber from "effect/Fiber";
11import * as Option from "effect/Option";12import * as Option from "effect/Option";
@@ -785,6 +786,56 @@ describe("AcpSessionRuntime", () => { · 773 unmodified lines
785 ),786 ),
786 );787 );
787 788
789 it.effect("keeps active-session updates while loading a different session", () =>
790 Effect.gen(function* () {
791 const loadStarted = yield* Deferred.make<void>();
792 const runtime = yield* AcpSessionRuntime.make({
793 authMethodId: "test",
794 spawn: {
795 command: mockAgentCommand,
796 args: mockAgentArgs,
797 env: {
798 T3_ACP_DELAY_LOAD_SESSION_AFTER_REPLAY: "1",
799 T3_ACP_LOAD_SESSION_DELAY_MS: "250",
800 },
801 },
802 cwd: process.cwd(),
803 sessionLoadTimeout: "2 seconds",
804 clientInfo: { name: "t3-test", version: "0.0.0" },
805 requestLogger: (event) =>
806 event.method === "session/load" && event.status === "started"
807 ? Deferred.succeed(loadStarted, undefined).pipe(Effect.asVoid)
808 : Effect.void,
809 });
810 yield* runtime.start();
811
812 const eventsFiber = yield* Stream.runCollect(Stream.take(runtime.getEvents(), 4)).pipe(
813 Effect.forkChild({ startImmediately: true }),
814 );
815 const loadFiber = yield* runtime
816 .loadSession("mock-session-2")
817 .pipe(Effect.forkChild({ startImmediately: true }));
818 yield* Deferred.await(loadStarted);
819
820 yield* runtime.prompt({
821 prompt: [{ type: "text", text: "active session prompt" }],
822 });
823 const events = Array.from(yield* Fiber.join(eventsFiber).pipe(Effect.timeout("2 seconds")));
824 expect(events.map((event) => event._tag)).toEqual([
825 "PlanUpdated",
826 "AssistantItemStarted",
827 "ContentDelta",
828 "AssistantItemCompleted",
829 ]);
830 expect(
831 events.some((event) => event._tag === "ContentDelta" && event.text === "hello from mock"),
832 ).toBe(true);
833
834 const loaded = yield* Fiber.join(loadFiber).pipe(Effect.timeout("2 seconds"));
835 expect(loaded.sessionId).toBe("mock-session-2");
836 }).pipe(Effect.scoped, Effect.provide(NodeServices.layer), TestClock.withLive),
837 );
838
788 it.effect("completes session/load after replay becomes idle while its RPC stays pending", () =>839 it.effect("completes session/load after replay becomes idle while its RPC stays pending", () =>
789 Effect.gen(function* () {840 Effect.gen(function* () {
790 const runtime = yield* AcpSessionRuntime.AcpSessionRuntime;841 const runtime = yield* AcpSessionRuntime.AcpSessionRuntime;
apps/server/src/provider/acp/AcpSessionRuntime.ts+7−1
@@ -1597,7 +1597,13 @@ export const make = ( · 1596 unmodified lines
1597 yield* acp.handleSessionUpdate((notification) =>1597 yield* acp.handleSessionUpdate((notification) =>
1598 Effect.gen(function* () {1598 Effect.gen(function* () {
1599 const gate = yield* Ref.get(sessionLoadGateRef);1599 const gate = yield* Ref.get(sessionLoadGateRef);
1600 if (Option.isSome(gate) && gate.value.active) {1600 // A different session can still have an in-flight prompt while this
1601 // load replays history, so quarantine only the loading session.
1602 if (
1603 Option.isSome(gate) &&
1604 gate.value.active &&
1605 notification.sessionId === gate.value.sessionId
1606 ) {
1601 if (sessionUpdateCountsAsLoadReplayActivity(notification, gate.value.sessionId)) {1607 if (sessionUpdateCountsAsLoadReplayActivity(notification, gate.value.sessionId)) {
1602 const lastActivityAtMillis = yield* Clock.currentTimeMillis;1608 const lastActivityAtMillis = yield* Clock.currentTimeMillis;
1603 yield* Ref.set(1609 yield* Ref.set(