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
55 type ModelSelection,
66 NodeId,
77 type OrchestrationV2Run,
8+ type OrchestrationV2Subagent,
89 type OrchestrationV2ThreadProjection,
910 type OrchestrationV2ThreadShell,
1011 type OrchestrationV2TurnItem,
@@ -279,6 +280,25 @@ function taskStatusForRun( · 268 unmodified lines
279280 }
280281 }
281282
283+function 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+
282302 function isTerminalTaskStatus(
283303 status: OrchestratorMcpDelegateTaskResult["status"],
284304 ): status is TerminalTaskStatus {
@@ -773,7 +793,7 @@ const make = Effect.gen(function* () { · 488 unmodified lines
773793 );
774794 }
775795 const childProjection = yield* loadProjection(task.childThreadId);
776− const childRun = latestRun(childProjection);
796+ const childRun = delegatedTaskRun(childProjection, task);
777797 const status = taskStatusForRun(childRun);
778798 const derivedResult =
779799 task.result !== null
@@ -1139,7 +1159,9 @@ const make = Effect.gen(function* () { · 359 unmodified lines
11391159 } satisfies OrchestratorMcpTaskCancelResult;
11401160 }
11411161 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+ );
11431165 if (activeRun === undefined) {
11441166 return yield* failure(
11451167 "task_not_cancellable",
apps/server/src/mcp/OrchestratorMcpToolkit.integration.test.ts+52−2
@@ -408,7 +408,7 @@ describe("orchestrator MCP toolkit", () => { · 407 unmodified lines
408408 driver: ProviderDriverKind.make("claudeAgent"),
409409 capabilities: ClaudeProviderCapabilitiesV2,
410410 capturedTurns,
411− shouldComplete: () => true,
411+ shouldComplete: (turn) => turn.message.text !== cancellationPrompt,
412412 response: (turn) =>
413413 turn.message.text === delegatedPrompt
414414 ? delegatedResult
@@ -845,10 +845,60 @@ describe("orchestrator MCP toolkit", () => { · 430 unmodified lines
845845 delegatedStatusAfterFollowupCall.structuredContent,
846846 ).pipe(Effect.orDie);
847847 expect(delegatedStatusAfterFollowup).toMatchObject({
848− childRunId: childFollowup.runId,
848+ childRunId: delegated.childRunId,
849849 status: "completed",
850+ summary: delegatedResult,
850851 });
851852
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+
852902 // A wait-mode child (completionWake settled_only) that completes
853903 // while the parent run is live does not offer a wake: the
854904 // 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
66
77 import * as NodeServices from "@effect/platform-node/NodeServices";
88 import { it } from "@effect/vitest";
9+import * as Deferred from "effect/Deferred";
910 import * as Effect from "effect/Effect";
1011 import * as Fiber from "effect/Fiber";
1112 import * as Option from "effect/Option";
@@ -785,6 +786,56 @@ describe("AcpSessionRuntime", () => { · 773 unmodified lines
785786 ),
786787 );
787788
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+
788839 it.effect("completes session/load after replay becomes idle while its RPC stays pending", () =>
789840 Effect.gen(function* () {
790841 const runtime = yield* AcpSessionRuntime.AcpSessionRuntime;
apps/server/src/provider/acp/AcpSessionRuntime.ts+7−1
@@ -1597,7 +1597,13 @@ export const make = ( · 1596 unmodified lines
15971597 yield* acp.handleSessionUpdate((notification) =>
15981598 Effect.gen(function* () {
15991599 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+ ) {
16011607 if (sessionUpdateCountsAsLoadReplayActivity(notification, gate.value.sessionId)) {
16021608 const lastActivityAtMillis = yield* Clock.currentTimeMillis;
16031609 yield* Ref.set(