diff --git a/src/node/services/taskService.replacementMcpOrdering.test.ts b/src/node/services/taskService.replacementMcpOrdering.test.ts new file mode 100644 index 00000000000..a22707e1f9a --- /dev/null +++ b/src/node/services/taskService.replacementMcpOrdering.test.ts @@ -0,0 +1,394 @@ +import assert from "node:assert/strict"; +import { describe, expect, spyOn, test } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as path from "node:path"; +import { MockLanguageModelV3, simulateReadableStream } from "ai/test"; +import { Ok } from "@/common/types/result"; +import { shellQuote } from "@/common/utils/shell"; +import type { + ProviderModelFactory, + ResolveAndCreateModelResult, +} from "@/node/services/providerModelFactory"; +import { + cleanupTestEnvironment, + createTestEnvironment, + setupProviders, +} from "../../../tests/ipc/setup"; +import { + HAIKU_MODEL, + cleanupTempGitRepo, + createTempGitRepo, + createWorkspace, + generateBranchName, + resolveOrpcClient, +} from "../../../tests/ipc/helpers"; + +// Runs under `bun test` with the real ServiceContainer (TaskService, WorkspaceService, +// AIService, MCPServerManager) and a real stdio MCP process; only the language model is +// substituted (as in mcpIdentity.assembly.test.ts). +const RECORDING_SERVER = path.resolve( + import.meta.dir, + "../../../tests/fixtures/mcp/recording-server.ts" +); +const ATTEMPT = "att_00000000000000d4"; +const RUN = { runId: "wfr_mcp_order", stepId: "summarize", inputHash: "hash-mcp" }; + +interface RecordedEvent { + event: string; + cwd: string; + at: number; +} + +async function readRecord(file: string): Promise { + // Absent until the stub process first starts. + const raw = await fs.readFile(file, "utf8").catch((error: NodeJS.ErrnoException) => { + if (error.code === "ENOENT") return ""; + throw error; + }); + return raw + .split("\n") + .filter((line) => line.length > 0) + .map((line) => JSON.parse(line) as RecordedEvent); +} + +/** A text-only answer: the child's turn ends normally without calling any tool. */ +function answeringModel(): MockLanguageModelV3 { + const usage = { + inputTokens: { total: 1, noCache: 1, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 1, text: 1, reasoning: 0 }, + }; + return new MockLanguageModelV3({ + doStream: () => + Promise.resolve({ + stream: simulateReadableStream({ + chunks: [ + { type: "stream-start", warnings: [] }, + { type: "text-start", id: "a" }, + { type: "text-delta", id: "a", delta: "done" }, + { type: "text-end", id: "a" }, + { type: "finish", finishReason: { unified: "stop", raw: undefined }, usage }, + ], + initialDelayInMs: null, + chunkDelayInMs: null, + }), + }), + }); +} + +type TestEnvironment = Awaited>; + +/** + * Parent workspace with the recording MCP server enabled (project config + trust), a model + * that answers in text, the workflow step's retired child, and pass-through spies that ledger + * every MCP entry point by workspace. `launchReplacement` claims the retired child and + * launches its replacement through createMany's claim, returning the replacement's id. + */ +async function setUpReplacement(env: TestEnvironment, repoPath: string, recordFile: string) { + await setupProviders(env, { anthropic: { apiKey: "mock-model-key" } }); + // Project MCP config, consented by project trust (createWorkspace trusts the project): + // the server is enabled for every workspace of the project, the replacement included. + await fs.mkdir(path.join(repoPath, ".xum"), { recursive: true }); + await fs.writeFile( + path.join(repoPath, ".xum", "mcp.jsonc"), + JSON.stringify({ + servers: { + recorder: [process.execPath, RECORDING_SERVER, recordFile].map(shellQuote).join(" "), + }, + }) + ); + const parent = await createWorkspace(env, repoPath, generateBranchName("mcp-order")); + if (!parent.success) throw new Error(parent.error); + const parentId = parent.metadata.id; + // The parent's checkout materializes (and is sanitized) after create returns. + await env.services.initStateManager.waitForInit(parentId); + + const factory = ( + env.services.aiService as unknown as { providerModelFactory: ProviderModelFactory } + ).providerModelFactory; + spyOn(factory, "resolveAndCreateModel").mockImplementation(() => + Promise.resolve( + Ok({ + model: answeringModel(), + effectiveModelString: HAIKU_MODEL, + canonicalModelString: HAIKU_MODEL, + canonicalProviderName: "anthropic", + canonicalModelId: HAIKU_MODEL.slice(HAIKU_MODEL.indexOf(":") + 1), + wireProviderName: "anthropic", + routedThroughGateway: false, + } satisfies ResolveAndCreateModelResult) + ) + ); + + // The workflow step's previous child: ended without a report (interrupted, no receipt + // needed in the owning process), so the runner may claim and replace it. + await env.config.editConfig((cfg) => { + const project = cfg.projects.get(repoPath); + assert(project, "parent project must be registered"); + project.workspaces.push({ + id: "retiredmcp", + name: "retired-mcp", + path: path.join(env.config.srcDir, "retired-mcp"), + createdAt: new Date().toISOString(), + parentWorkspaceId: parentId, + agentType: "explore", + agentId: "explore", + runtimeConfig: parent.metadata.runtimeConfig, + taskStatus: "interrupted", + taskAttemptId: ATTEMPT, + workflowTask: { runId: RUN.runId, stepId: RUN.stepId }, + }); + return cfg; + }); + + // In-process order ledger. Every MCP entry point is a pass-through spy on the real + // manager, so a start that the stub process has not yet recorded is still caught. + const ledger: Array<{ step: string; workspaceId: string }> = []; + const mcp = env.services.mcpServerManager; + for (const method of ["getToolsForWorkspace", "getPromptsForWorkspace"] as const) { + const real = mcp[method].bind(mcp) as (...args: unknown[]) => Promise; + spyOn(mcp, method).mockImplementation(((...args: unknown[]) => { + ledger.push({ + step: method, + workspaceId: (args[0] as { workspaceId: string }).workspaceId, + }); + return real(...args); + }) as never); + } + + const launchReplacement = async (): Promise => { + const taskService = env.services.taskService; + const claim = await taskService.claimRetiredAttempt("retiredmcp", ATTEMPT, RUN); + assert(claim.success, `claim must succeed: ${claim.success ? "" : claim.error}`); + const created = await taskService.createMany( + [ + { + parentWorkspaceId: parentId, + kind: "agent", + agentId: "explore", + prompt: "Summarize durable workflows", + title: "Replacement", + workflowTask: { runId: RUN.runId, stepId: RUN.stepId }, + }, + ], + { retires: [{ taskId: "retiredmcp", attemptId: ATTEMPT, nonce: claim.data.nonce }] } + ); + assert(created.success, `createMany must succeed: ${created.success ? "" : created.error}`); + const replacementId = created.data[0]?.taskId; + assert(replacementId, "createMany must return the replacement's task id"); + return replacementId; + }; + return { parentId, ledger, launchReplacement }; +} + +/** + * Pass-through spies on init state: `parked` resolves once a request for `target.workspaceId` + * waits on init, and `completed` lists every workspace whose init was completed (endInit), the + * point that releases such waiters. + */ +function observeInit(env: TestEnvironment) { + const initStateManager = env.services.initStateManager; + const realWaitForInit = initStateManager.waitForInit.bind(initStateManager); + const realEndInit = initStateManager.endInit.bind(initStateManager); + const parked = Promise.withResolvers(); + const target: { workspaceId?: string } = {}; + const completed: string[] = []; + spyOn(initStateManager, "waitForInit").mockImplementation((workspaceId, signal) => { + if (workspaceId === target.workspaceId) parked.resolve(); + return realWaitForInit(workspaceId, signal); + }); + spyOn(initStateManager, "endInit").mockImplementation((workspaceId, exitCode) => { + completed.push(workspaceId); + return realEndInit(workspaceId, exitCode); + }); + return { target, parked: parked.promise, completed }; +} + +/** + * Gate 4 of G2 end to end (#4576): a workflow replacement's launch must not start MCP servers or + * run prompt discovery before its checkout's plugin overrides are sanitized. Two discovery + * routes reach a new child: the launch's own first send (turn assembly lists MCP tools and + * prompts) and a client's prompt-catalog request (`workspace.mcp.prompts.list`), which a + * renderer can issue as soon as the published row appears. + */ +describe("workflow replacement launch: MCP after sanitize", () => { + test("no MCP process, tool list or prompt discovery reaches the replacement until its sanitize returns", async () => { + const env = await createTestEnvironment(); + const repoPath = await createTempGitRepo(); + const recordFile = path.join(env.tempDir, "mcp-record.jsonl"); + try { + const { parentId, ledger, launchReplacement } = await setUpReplacement( + env, + repoPath, + recordFile + ); + + // Hold the replacement's sanitize (the real call runs once released). + const workspaceService = env.services.workspaceService; + const realSanitize = + workspaceService.sanitizeMaterializedTaskWorkspace.bind(workspaceService); + const sanitizeEntered = Promise.withResolvers<{ taskId: string; checkout: string }>(); + const releaseSanitize = Promise.withResolvers(); + let sanitizedAt = Number.POSITIVE_INFINITY; + let recordAtSanitize: RecordedEvent[] | undefined; + spyOn(workspaceService, "sanitizeMaterializedTaskWorkspace").mockImplementation( + async (...args) => { + if (args[0] === parentId) return realSanitize(...args); + sanitizeEntered.resolve({ taskId: args[0], checkout: args[1] }); + await releaseSanitize.promise; + const result = await realSanitize(...args); + sanitizedAt = Date.now(); + recordAtSanitize = await readRecord(recordFile); + ledger.push({ step: "sanitized", workspaceId: args[0] }); + return result; + } + ); + // Witness that a client's prompt-catalog request has parked on the launch's init state. + const initWaits = observeInit(env); + const initStateManager = env.services.initStateManager; + + const replacementId = await launchReplacement(); + + const held = await sanitizeEntered.promise; + expect(held.taskId).toBe(replacementId); + // Sanitize is held: the launch has neither started a server nor listed anything. + expect(ledger.filter((entry) => entry.workspaceId === replacementId)).toEqual([]); + expect(await readRecord(recordFile)).toEqual([]); + // The row is published and its checkout materialized: a renderer may ask for the new + // workspace's prompt catalog now. The request must wait for sanitize, not race it. + initWaits.target.workspaceId = replacementId; + const probe = resolveOrpcClient(env).workspace.mcp.prompts.list({ + workspaceId: replacementId, + }); + await initWaits.parked; + // Parked, not passed through: the launch's init is still running, and it completes only + // after sanitize (waitForInit returns at once for a completed or absent init). + expect(initStateManager.getInitState(replacementId)?.status).toBe("running"); + expect(initWaits.completed).not.toContain(replacementId); + + releaseSanitize.resolve(); + + // Positive control: discovery genuinely reaches this child's stub once sanitized. + const prompts = await probe; + expect(prompts.map((p) => `${p.serverName}/${p.promptName}`)).toContain("recorder/recorded"); + const checkout = await fs.realpath(held.checkout); + // The launch's own first send assembles the turn: it lists this child's MCP tools. + const deadline = Date.now() + 30_000; + const sendListedTools = () => + ledger.some((e) => e.workspaceId === replacementId && e.step === "getToolsForWorkspace"); + while (!sendListedTools()) { + assert(Date.now() < deadline, "the replacement's first send never listed MCP tools"); + await new Promise((resolve) => setTimeout(resolve, 25)); + } + const childEvents = (await readRecord(recordFile)).filter((e) => e.cwd === checkout); + // Let the child's turn finish before teardown: it reports, and a completed task's row + // may then be auto-deleted. + const status = () => + env.config + .loadConfigOrDefault() + .projects.get(repoPath) + ?.workspaces.find((w) => w.id === replacementId)?.taskStatus; + while (status() !== "reported" && status() !== undefined) { + assert( + Date.now() < deadline, + `the replacement never reported (status ${String(status())})` + ); + await new Promise((resolve) => setTimeout(resolve, 25)); + } + + // Nothing ran before sanitize returned, in-process or in the stub process... + expect(recordAtSanitize).toEqual([]); + const replacementSteps = ledger + .filter((entry) => entry.workspaceId === replacementId) + .map((entry) => entry.step); + expect(replacementSteps[0]).toBe("sanitized"); + // ...and afterwards both routes reached it: the probe's catalog and the send's tool list. + expect(replacementSteps).toContain("getPromptsForWorkspace"); + expect(replacementSteps).toContain("getToolsForWorkspace"); + expect(childEvents[0]?.event).toBe("start"); + const childMethods = childEvents.map((e) => e.event); + expect(childMethods).toContain("tools/list"); + expect(childMethods).toContain("prompts/list"); + for (const event of childEvents) expect(event.at).toBeGreaterThanOrEqual(sanitizedAt); + } finally { + await cleanupTestEnvironment(env); + await cleanupTempGitRepo(repoPath); + } + }, 120_000); + + test("a failed sanitize releases no parked discovery into the checkout before its reclaim", async () => { + const env = await createTestEnvironment(); + const repoPath = await createTempGitRepo(); + const recordFile = path.join(env.tempDir, "mcp-record.jsonl"); + try { + const { parentId, ledger, launchReplacement } = await setUpReplacement( + env, + repoPath, + recordFile + ); + // The replacement's sanitize fails; the launch must then unpublish the row and delete + // the checkout. The first config edit after the failure (the reclaim's unpublish) is held. + const reclaimHeld = Promise.withResolvers(); + const releaseReclaim = Promise.withResolvers(); + let holdNextEdit = false; + const realEditConfig = env.config.editConfig.bind(env.config); + spyOn(env.config, "editConfig").mockImplementation((async (...args: unknown[]) => { + if (holdNextEdit) { + holdNextEdit = false; + reclaimHeld.resolve(); + await releaseReclaim.promise; + } + return (realEditConfig as (...a: unknown[]) => Promise)(...args); + }) as never); + let checkout: string | undefined; + spyOn(env.services.workspaceService, "sanitizeMaterializedTaskWorkspace").mockImplementation( + async (id: string, workspacePath: string) => { + if (id === parentId) return undefined; + checkout = await fs.realpath(workspacePath); + holdNextEdit = true; + return "fixture: sanitize failed"; + } + ); + const initWaits = observeInit(env); + + const replacementId = await launchReplacement(); + await reclaimHeld.promise; + assert(checkout, "the replacement's sanitize must have run"); + // The unsanitized row is still published; a renderer asks for its prompt catalog. + initWaits.target.workspaceId = replacementId; + const probe = resolveOrpcClient(env) + .workspace.mcp.prompts.list({ workspaceId: replacementId }) + .then( + (prompts) => ({ prompts }), + (error: unknown) => ({ error }) + ); + await initWaits.parked; + // Parked: init completes only after the reclaim attempt (completing it is what releases + // the request), so nothing reaches the manager or starts a server while it is held. + expect(initWaits.completed).not.toContain(replacementId); + expect(ledger.filter((entry) => entry.workspaceId === replacementId)).toEqual([]); + expect(await readRecord(recordFile)).toEqual([]); + + releaseReclaim.resolve(); + const outcome = await probe; + // The reclaim removed the row, checkout and session dir before the request resumed: it + // finds no workspace and fails, and no server ever ran in the reclaimed checkout. + expect("error" in outcome).toBe(true); + expect(ledger.filter((entry) => entry.workspaceId === replacementId)).toEqual([]); + expect((await readRecord(recordFile)).filter((e) => e.cwd === checkout)).toEqual([]); + expect( + env.config + .loadConfigOrDefault() + .projects.get(repoPath) + ?.workspaces.some((w) => w.id === replacementId) + ).toBe(false); + expect(await fs.stat(checkout).catch(() => null)).toBeNull(); + // Dropping (not completing) the reclaimed task's init writes no init-status.json back. + expect( + await fs.stat(path.join(env.config.sessionsDir, replacementId)).catch(() => null) + ).toBeNull(); + } finally { + await cleanupTestEnvironment(env); + await cleanupTempGitRepo(repoPath); + } + }, 120_000); +}); diff --git a/src/node/services/taskService.testHarness.ts b/src/node/services/taskService.testHarness.ts index 97915d3a517..a9c398eb2f0 100644 --- a/src/node/services/taskService.testHarness.ts +++ b/src/node/services/taskService.testHarness.ts @@ -61,6 +61,7 @@ export function createMockInitStateManager(): InitStateManager { appendOutput: mock(() => undefined), reportProgress: mock(() => undefined), endInit: mock(() => Promise.resolve()), + clearInMemoryState: mock(() => undefined), getInitState: mock(() => undefined), readInitStatus: mock(() => Promise.resolve(null)), } as unknown as InitStateManager; diff --git a/src/node/services/taskService.ts b/src/node/services/taskService.ts index 0a093d9686d..a1f57ca71ee 100644 --- a/src/node/services/taskService.ts +++ b/src/node/services/taskService.ts @@ -6460,6 +6460,7 @@ export class TaskService implements AgentTaskIntegration { * and session dir are named after the task id, so no other task can reuse them once the row is * gone. Anything short of a confirmed unpublication — no or another owner, a moved row, a lost * or unverifiable write — retains everything; the failure is then recorded on the row. + * Returns whether the row was unpublished and the checkout reclaimed. */ private async reclaimUnsanitizedTaskCheckout( runtime: Runtime, @@ -6467,7 +6468,7 @@ export class TaskService implements AgentTaskIntegration { workspaceName: string, taskId: string, expectedAttemptId: string | undefined - ): Promise { + ): Promise { let unpublished = false; if ( expectedAttemptId != null && @@ -6501,11 +6502,12 @@ export class TaskService implements AgentTaskIntegration { log.warn("Task launch: unsanitized checkout retained (its row is still published)", { taskId, }); - return; + return false; } await this.rollbackFailedTaskCreate(runtime, projectPath, workspaceName, taskId, { rowUnpublished: true, }); + return true; } private async getExistingMaterializedTaskLaunch( @@ -6980,18 +6982,32 @@ export class TaskService implements AgentTaskIntegration { forkedRuntimeConfig ); if (sanitizeError !== undefined) { - initLogger.logComplete(-1); // Reclaim the just-materialized worktree/session before failing the // launch: the throw reaches scheduleReservedTaskLaunch, which only // marks the task interrupted — without this cleanup the physical // checkout would accumulate and collide with later same-name forks. - await this.reclaimUnsanitizedTaskCheckout( - runtimeForTaskWorkspace, - plan.parentMeta.projectPath, - plan.workspaceName, - plan.taskId, - plan.attemptId - ); + let reclaimed = false; + try { + reclaimed = await this.reclaimUnsanitizedTaskCheckout( + runtimeForTaskWorkspace, + plan.parentMeta.projectPath, + plan.workspaceName, + plan.taskId, + plan.attemptId + ); + } finally { + // SECURITY: init ends only after the reclaim attempt. Ending it releases every + // request parked in waitForInit (MCP prompt discovery among them); released while + // the row is still published, one would start MCP servers inside this unsanitized + // checkout. A reclaimed task is gone (row, checkout, session dir): drop its init state + // as workspace removal does, since completing it would recreate the session dir to + // persist init-status.json. A retained checkout completes init so waiters don't hang. + if (reclaimed) { + this.initStateManager.clearInMemoryState(plan.taskId); + } else { + initLogger.logComplete(-1); + } + } throw new Error(sanitizeError); } } diff --git a/tests/fixtures/mcp/recording-server.ts b/tests/fixtures/mcp/recording-server.ts new file mode 100644 index 00000000000..cceb562e5bf --- /dev/null +++ b/tests/fixtures/mcp/recording-server.ts @@ -0,0 +1,34 @@ +import { appendFileSync } from "node:fs"; +import { createInterface } from "node:readline"; + +// Stdio MCP stub that appends one JSON line per event to the file named by argv[2]: its own +// process start and every JSON-RPC message it receives, each with the process cwd (the checkout +// the host launched it in) and a timestamp. Tests read the file to prove whether, where and +// when a real server process ran. Serves tools and prompts so both discovery lists are exercised. +const recordFile = process.argv[2]; +if (!recordFile) throw new Error("recording-server: pass the record file path as argv[2]"); + +const record = (event: string) => + appendFileSync(recordFile, JSON.stringify({ event, cwd: process.cwd(), at: Date.now() }) + "\n"); +record("start"); + +const results: Record = { + initialize: { + protocolVersion: "2025-11-25", + capabilities: { tools: {}, prompts: {} }, + serverInfo: { name: "recorder", version: "1" }, + }, + "tools/list": { tools: [{ name: "probe", inputSchema: { type: "object", properties: {} } }] }, + "prompts/list": { prompts: [{ name: "recorded" }] }, + ping: {}, +}; + +createInterface({ input: process.stdin, crlfDelay: Infinity }).on("line", (line) => { + const request = JSON.parse(line) as { id?: string | number; method: string }; + record(request.method); + if (request.id === undefined) return; + const result = results[request.method]; + const envelope = + result === undefined ? { error: { code: -32601, message: "Method not found" } } : { result }; + process.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: request.id, ...envelope }) + "\n"); +});