From 2d45a776eeb7be09933314e09decdcac67d4e045 Mon Sep 17 00:00:00 2001 From: ethan Date: Mon, 22 Jun 2026 15:50:39 +1000 Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20fix:=20let=20built-in=20forked?= =?UTF-8?q?=20explore=20tasks=20run=20in=20parallel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the per-stream AsyncMutex in withSequentialExecution with a fair reader-writer lock (SharedExecutionLock). Built-in `task` calls targeting the `explore` agent (without `isolation: "none"`) share the read side and overlap with each other; every other tool keeps the exclusive write lock and stays serialized. The built-in task tool is identified by a non-forgeable Symbol marker set in createTaskTool. Preserve that Symbol marker through Anthropic cache-control recreation (applyCacheControlToTools), which rebuilds the last function tool via createTool() and would otherwise drop the marker — silently serializing explore fan-out whenever tool-policy filtering leaves `task` as the last tool. --- src/common/utils/ai/cacheStrategy.test.ts | 24 +++ src/common/utils/ai/cacheStrategy.ts | 9 + src/node/services/tools/task.test.ts | 22 +- src/node/services/tools/task.ts | 27 ++- .../tools/withSequentialExecution.test.ts | 193 +++++++++++++++++- .../services/tools/withSequentialExecution.ts | 133 +++++++++++- 6 files changed, 390 insertions(+), 18 deletions(-) diff --git a/src/common/utils/ai/cacheStrategy.test.ts b/src/common/utils/ai/cacheStrategy.test.ts index 1bb76b24a4..510ff54401 100644 --- a/src/common/utils/ai/cacheStrategy.test.ts +++ b/src/common/utils/ai/cacheStrategy.test.ts @@ -10,6 +10,7 @@ import { createCachedSystemMessage, applyCacheControlToTools, } from "./cacheStrategy"; +import { markBuiltInTaskTool, isBuiltInTaskTool } from "@/node/services/tools/task"; describe("cacheStrategy", () => { describe("supportsAnthropicCache", () => { @@ -385,5 +386,28 @@ describe("cacheStrategy", () => { }); expect(result.readFile).toEqual(toolsWithDynamicTool.readFile); }); + + it("preserves the built-in task marker when recreating the last function tool", () => { + // Cache control recreates the last function tool via createTool(). Built-in explore-task + // parallelism depends on a symbol marker surviving that recreation; if it were dropped the + // task tool would silently fall back to serialized execution. + const taskTool = markBuiltInTaskTool( + tool({ + description: "task", + inputSchema: z.object({ prompt: z.string() }), + execute: () => Promise.resolve("ok"), + }) + ); + const tools: Record = { + readFile: mockTools.readFile, + task: taskTool, + }; + + const result = applyCacheControlToTools(tools, "anthropic:claude-3-5-sonnet"); + + expect(isBuiltInTaskTool(result.task)).toBe(true); + // A recreated tool is a different object; sanity-check the marker rode along on the copy. + expect(result.task).not.toBe(taskTool); + }); }); }); diff --git a/src/common/utils/ai/cacheStrategy.ts b/src/common/utils/ai/cacheStrategy.ts index 0b31d8fd87..40312361a2 100644 --- a/src/common/utils/ai/cacheStrategy.ts +++ b/src/common/utils/ai/cacheStrategy.ts @@ -199,6 +199,15 @@ export function applyCacheControlToTools>( execute: existingTool.execute, providerOptions: cacheOpts, }); + // createTool() returns a fresh object that drops any extra own symbol markers attached + // to the original (e.g. the built-in task-tool marker that lets sibling explore tasks run + // in parallel). Copy them over so downstream wrappers still recognize the recreated tool. + for (const marker of Object.getOwnPropertySymbols(existingTool)) { + const descriptor = Object.getOwnPropertyDescriptor(existingTool, marker); + if (descriptor) { + Object.defineProperty(cachedTool, marker, descriptor); + } + } cachedTools[key as keyof T] = cachedTool as unknown as T[keyof T]; } } else { diff --git a/src/node/services/tools/task.test.ts b/src/node/services/tools/task.test.ts index b86d387b87..6a93176d84 100644 --- a/src/node/services/tools/task.test.ts +++ b/src/node/services/tools/task.test.ts @@ -1,7 +1,10 @@ import { describe, it, expect, mock } from "bun:test"; import type { TaskCreatedEvent } from "@/common/types/stream"; -import { createTaskTool } from "./task"; +import { tool } from "ai"; +import { z } from "zod"; + +import { createTaskTool, markBuiltInTaskTool, isBuiltInTaskTool } from "./task"; import { createTestToolConfig, mockToolCallOptions, TestTempDir } from "./testHelpers"; import { Ok, Err } from "@/common/types/result"; import { ForegroundWaitBackgroundedError, type TaskService } from "@/node/services/taskService"; @@ -1117,3 +1120,20 @@ describe("task tool", () => { expect(waitForAgentReport).not.toHaveBeenCalled(); }); }); + +describe("built-in task marker", () => { + function makeTool() { + return tool({ + description: "task", + inputSchema: z.object({ prompt: z.string() }), + execute: () => Promise.resolve("ok"), + }); + } + + it("marks and recognizes the built-in task tool", () => { + const t = makeTool(); + expect(isBuiltInTaskTool(t)).toBe(false); + markBuiltInTaskTool(t); + expect(isBuiltInTaskTool(t)).toBe(true); + }); +}); diff --git a/src/node/services/tools/task.ts b/src/node/services/tools/task.ts index 5785daef1c..8c116bbf88 100644 --- a/src/node/services/tools/task.ts +++ b/src/node/services/tools/task.ts @@ -1,6 +1,6 @@ import { randomUUID } from "node:crypto"; -import { tool } from "ai"; +import { tool, type Tool } from "ai"; import type { z } from "zod"; import type { ToolConfiguration, ToolFactory } from "@/common/utils/tools/tools"; @@ -36,6 +36,28 @@ import { import { normalizeModelInput } from "@/common/utils/ai/normalizeModelInput"; import { coerceNonEmptyString } from "@/node/services/taskUtils"; +const BUILT_IN_TASK_TOOL_MARKER = Symbol("muxBuiltInTaskTool"); + +export function markBuiltInTaskTool( + taskTool: Tool +): Tool { + Object.defineProperty(taskTool, BUILT_IN_TASK_TOOL_MARKER, { + value: true, + // enumerable so object spread (wrapWithInitWait) and descriptor clones (withHooks, + // cloneToolPreservingDescriptors, cache control) carry the marker forward to every wrapper — + // that is what lets sibling explore task calls share the parallel reader lock downstream. + enumerable: true, + configurable: true, + }); + return taskTool; +} + +export function isBuiltInTaskTool(tool: Tool | undefined): boolean { + return Boolean( + (tool as (Tool & Record) | undefined)?.[BUILT_IN_TASK_TOOL_MARKER] === true + ); +} + /** Resolve the parent workspace's runtime mode from the injected MUX_RUNTIME env. */ function resolveRuntimeMode(config: ToolConfiguration): RuntimeMode | undefined { const runtimeValue = config.muxEnv?.MUX_RUNTIME; @@ -331,7 +353,7 @@ export const createTaskTool: ToolFactory = (config: ToolConfiguration) => { const inputSchema = buildTaskToolAgentArgsSchema({ includeIsolation: runtimeModeSupportsSharedTaskWorkspace(runtimeMode), }); - return tool({ + const taskTool = tool({ description: buildTaskDescription(config), inputSchema, execute: async (args, { abortSignal, toolCallId }): Promise => { @@ -663,4 +685,5 @@ export const createTaskTool: ToolFactory = (config: ToolConfiguration) => { throw new Error("Task foreground wait ended without a terminal result"); }, }); + return markBuiltInTaskTool(taskTool); }; diff --git a/src/node/services/tools/withSequentialExecution.test.ts b/src/node/services/tools/withSequentialExecution.test.ts index 20826eaa82..6e1c8e52d0 100644 --- a/src/node/services/tools/withSequentialExecution.test.ts +++ b/src/node/services/tools/withSequentialExecution.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from "bun:test"; import { tool } from "ai"; import { z } from "zod"; import { withSequentialExecution } from "./withSequentialExecution"; +import { markBuiltInTaskTool } from "./task"; interface Deferred { promise: Promise; @@ -26,6 +27,7 @@ function createDeferred(): Deferred { function callWrappedExecute( toolRecord: Record, + args: unknown, options: unknown ): Promise { const execute = toolRecord.execute; @@ -33,8 +35,8 @@ function callWrappedExecute( throw new Error("Expected wrapped tool execute handler"); } - const invoke = execute as (args: Record, options: unknown) => Promise; - return invoke({}, options); + const invoke = execute as (args: unknown, options: unknown) => unknown; + return Promise.resolve(invoke(args, options)); } describe("withSequentialExecution", () => { @@ -156,16 +158,14 @@ describe("withSequentialExecution", () => { const controller = new AbortController(); const firstPromise = callWrappedExecute( wrappedTools!.a as Record, + {}, {} as never ); await startedA.promise; - const secondPromise = callWrappedExecute( - wrappedTools!.b as Record, - { - abortSignal: controller.signal, - } as never - ); + const secondPromise = callWrappedExecute(wrappedTools!.b as Record, {}, { + abortSignal: controller.signal, + } as never); controller.abort(); try { @@ -184,4 +184,181 @@ describe("withSequentialExecution", () => { expect(startedB).toBe(false); expect(executionLog).toEqual(["start A", "end A"]); }); + + test("runs forked built-in explore tasks in parallel while writers wait", async () => { + const executionLog: string[] = []; + const started = { + a: createDeferred(), + b: createDeferred(), + writer: createDeferred(), + }; + const release = { + a: createDeferred(), + b: createDeferred(), + writer: createDeferred(), + }; + + const tools = { + task: markBuiltInTaskTool( + tool({ + description: "Task", + inputSchema: z.object({ + id: z.enum(["a", "b"]), + agentId: z.string(), + isolation: z.string().optional(), + }), + execute: async ({ id }: { id: "a" | "b" }) => { + executionLog.push(`start ${id}`); + started[id].resolve(); + await release[id].promise; + executionLog.push(`end ${id}`); + return { task: id }; + }, + }) + ), + bash: tool({ + description: "Writer", + inputSchema: z.object({}), + execute: async () => { + executionLog.push("start writer"); + started.writer.resolve(); + await release.writer.promise; + executionLog.push("end writer"); + return { tool: "writer" }; + }, + }), + }; + + const wrappedTools = withSequentialExecution(tools)!; + const resultsPromise = Promise.all([ + callWrappedExecute( + wrappedTools.task as Record, + { id: "a", agentId: "explore" }, + {} as never + ), + callWrappedExecute( + wrappedTools.task as Record, + { id: "b", agentId: "explore" }, + {} as never + ), + callWrappedExecute(wrappedTools.bash as Record, {}, {} as never), + ]); + + await started.a.promise; + await started.b.promise; + await Promise.resolve(); + expect(executionLog).toEqual(["start a", "start b"]); + + release.a.resolve(); + await Promise.resolve(); + expect(executionLog).toEqual(["start a", "start b", "end a"]); + + release.b.resolve(); + await started.writer.promise; + await Promise.resolve(); + expect(executionLog).toEqual(["start a", "start b", "end a", "end b", "start writer"]); + + release.writer.resolve(); + const results = await resultsPromise; + expect(results).toEqual([{ task: "a" }, { task: "b" }, { tool: "writer" }]); + }); + + test("keeps shared-workspace explore tasks serialized", async () => { + const executionLog: string[] = []; + const startedA = createDeferred(); + const releaseA = createDeferred(); + let startedB = false; + + const tools = { + task: markBuiltInTaskTool( + tool({ + description: "Task", + inputSchema: z.object({ agentId: z.string(), isolation: z.string().optional() }), + execute: async ({ isolation }: { isolation?: string }) => { + executionLog.push(`start ${isolation ?? "fork"}`); + if (isolation === "none") { + if (!startedB) { + startedA.resolve(); + await releaseA.promise; + executionLog.push("end first"); + } else { + executionLog.push("start second"); + } + } + return { ok: true }; + }, + }) + ), + }; + + const wrappedTools = withSequentialExecution(tools)!; + const firstPromise = callWrappedExecute( + wrappedTools.task as Record, + { agentId: "explore", isolation: "none" }, + {} as never + ); + await startedA.promise; + + const secondPromise = callWrappedExecute( + wrappedTools.task as Record, + { agentId: "explore", isolation: "none" }, + {} as never + ); + await Promise.resolve(); + expect(executionLog).toEqual(["start none"]); + + startedB = true; + releaseA.resolve(); + await firstPromise; + await secondPromise; + expect(executionLog).toEqual(["start none", "end first", "start none", "start second"]); + }); + + test("treats non-canonical explore agent ids as readers (trim + lowercase)", async () => { + // The task inputSchema passes the raw agentId through; classification must mirror the schema's + // trim()/toLowerCase() normalization so " Explore " / "EXPLORE" still share the reader lock. + const executionLog: string[] = []; + const started = { a: createDeferred(), b: createDeferred() }; + const release = { a: createDeferred(), b: createDeferred() }; + + const tools = { + task: markBuiltInTaskTool( + tool({ + description: "Task", + inputSchema: z.object({ id: z.enum(["a", "b"]), agentId: z.string() }), + execute: async ({ id }: { id: "a" | "b" }) => { + executionLog.push(`start ${id}`); + started[id].resolve(); + await release[id].promise; + return { task: id }; + }, + }) + ), + }; + + const wrappedTools = withSequentialExecution(tools)!; + const resultsPromise = Promise.all([ + callWrappedExecute( + wrappedTools.task as Record, + { id: "a", agentId: " Explore " }, + {} as never + ), + callWrappedExecute( + wrappedTools.task as Record, + { id: "b", agentId: "EXPLORE" }, + {} as never + ), + ]); + + // Both overlap before either finishes — proving they share the reader lock despite casing. + await started.a.promise; + await started.b.promise; + await Promise.resolve(); + expect(executionLog).toEqual(["start a", "start b"]); + + release.a.resolve(); + release.b.resolve(); + const results = await resultsPromise; + expect(results).toEqual([{ task: "a" }, { task: "b" }]); + }); }); diff --git a/src/node/services/tools/withSequentialExecution.ts b/src/node/services/tools/withSequentialExecution.ts index 972ce115a3..00f03c64c0 100644 --- a/src/node/services/tools/withSequentialExecution.ts +++ b/src/node/services/tools/withSequentialExecution.ts @@ -2,9 +2,16 @@ import type { Tool } from "ai"; import assert from "@/common/utils/assert"; import { cloneToolPreservingDescriptors } from "@/common/utils/tools/cloneToolPreservingDescriptors"; import { AsyncMutex } from "@/node/utils/concurrency/asyncMutex"; +import { isBuiltInTaskTool } from "@/node/services/tools/task"; type AsyncMutexGuard = Awaited>; +interface ParallelTaskArgs { + agentId?: unknown; + subagent_type?: unknown; + isolation?: unknown; +} + interface ToolExecutionContext { abortSignal?: AbortSignal; } @@ -18,6 +25,58 @@ function getAbortSignal(options: unknown): AbortSignal | undefined { return context.abortSignal; } +function getRequestedAgentId(args: unknown): string | undefined { + if (typeof args !== "object" || args === null) { + return undefined; + } + + // The task tool's inputSchema (buildTaskToolAgentArgsSchema) passes the raw agentId through; the + // trim()/toLowerCase() normalization only happens inside execute via TOOL_DEFINITIONS.task.schema. + // Mirror that normalization here so valid-but-non-canonical ids ("Explore", " explore ") are + // still recognized and share the reader lock instead of falling back to the writer lock. + const normalize = (value: unknown): string | undefined => { + if (typeof value !== "string") { + return undefined; + } + const normalized = value.trim().toLowerCase(); + return normalized.length > 0 ? normalized : undefined; + }; + + const taskArgs = args as ParallelTaskArgs; + return normalize(taskArgs.agentId) ?? normalize(taskArgs.subagent_type); +} + +// Decide whether a tool call may share the read side of the lock with sibling explore tasks. +// +// MAINTAINER DECISION (do not "harden" this by inspecting the resolved tool set or runtime fork +// isolation): we trust the `explore` agent id as a declaration of read-only intent. The `explore` +// contract is read-only *by prompt*, not by tool removal — the built-in explore agent still has +// `bash`, and a custom project/global override named `explore` may also enable bash/exec. Letting +// two such agents run in parallel is explicitly fine and NOT a race we are trying to prevent here, +// even when the runtime gives them a *shared* checkout. In local runtime the task tool does not +// expose `isolation`, so explore calls arrive with `isolation === undefined` and forks point at +// the same project directory; two parallel explore calls writing to that checkout via bash is +// accepted as best-effort. The only `isolation` value that opts OUT of parallelism is the explicit +// `"none"` (caller-declared shared workspace). What we DO keep serialized regardless: direct +// mutating tools (file_edit_*, bash, config writes) and non-explore forked tasks — those always +// take the exclusive write lock. +// +// This also covers parent-side tool hooks: when a repo configures `.mux/tool_pre`/`tool_post`/ +// `tool_hook`, those scripts wrap the task tool's execute and run in the parent checkout. Two +// sibling explore task calls therefore run those hook scripts concurrently. That is the same +// best-effort tradeoff as concurrent explore bash/exec, so the built-in task marker is preserved +// through hook wrapping (see wrapToolsWithHooks) rather than stripped — stripping it serialized +// every explore task in the common case of a repo that merely has a `tool_post` formatter. +function canRunWithSiblingExploreTasks(baseTool: Tool, args: unknown): boolean { + if (!isBuiltInTaskTool(baseTool)) { + return false; + } + if (getRequestedAgentId(args) !== "explore") { + return false; + } + return (args as ParallelTaskArgs | null)?.isolation !== "none"; +} + function releaseLockAfterAbort(acquirePromise: Promise): void { void acquirePromise .then(async (lock) => { @@ -74,13 +133,71 @@ async function acquireLockOrAbort( } } +class SharedExecutionReadGuard implements AsyncDisposable { + constructor(private readonly lock: SharedExecutionLock) {} + + async [Symbol.asyncDispose](): Promise { + await this.lock.releaseRead(); + } +} + +class SharedExecutionWriteGuard implements AsyncDisposable { + constructor( + private readonly turnstileLock: AsyncMutexGuard, + private readonly roomEmptyLock: AsyncMutexGuard + ) {} + + async [Symbol.asyncDispose](): Promise { + await this.roomEmptyLock[Symbol.asyncDispose](); + await this.turnstileLock[Symbol.asyncDispose](); + } +} + +class SharedExecutionLock { + private readonly turnstile = new AsyncMutex(); + private readonly roomEmpty = new AsyncMutex(); + private readonly readerCountLock = new AsyncMutex(); + private activeReaders = 0; + private roomEmptyGuard: AsyncMutexGuard | undefined; + + async acquireRead(abortSignal?: AbortSignal): Promise { + await using _turnstile = await acquireLockOrAbort(this.turnstile, abortSignal); + await using _readerCountLock = await acquireLockOrAbort(this.readerCountLock, abortSignal); + if (this.activeReaders === 0) { + this.roomEmptyGuard = await acquireLockOrAbort(this.roomEmpty, abortSignal); + } + this.activeReaders += 1; + return new SharedExecutionReadGuard(this); + } + + async acquireWrite(abortSignal?: AbortSignal): Promise { + const turnstileLock = await acquireLockOrAbort(this.turnstile, abortSignal); + try { + const roomEmptyLock = await acquireLockOrAbort(this.roomEmpty, abortSignal); + return new SharedExecutionWriteGuard(turnstileLock, roomEmptyLock); + } catch (error) { + await turnstileLock[Symbol.asyncDispose](); + throw error; + } + } + + async releaseRead(): Promise { + await using _readerCountLock = await this.readerCountLock.acquire(); + assert(this.activeReaders > 0, "SharedExecutionLock.releaseRead called with no active readers"); + this.activeReaders -= 1; + if (this.activeReaders === 0) { + const roomEmptyGuard = this.roomEmptyGuard; + this.roomEmptyGuard = undefined; + await roomEmptyGuard?.[Symbol.asyncDispose](); + } + } +} + /** * Serialize sibling tool execution for a single stream without changing the - * provider's parallel-tool-call planning behavior. - * - * We intentionally scope the mutex to the returned tool map so independent - * streams can still execute concurrently. Holding the lock across the full - * execute chain preserves ordering across all per-tool wrappers and side effects. + * provider's parallel-tool-call planning behavior. Built-in forked explore + * tasks share the read side so they can overlap with each other, while every + * other tool call stays exclusive. */ export function withSequentialExecution( tools: Record | undefined @@ -89,7 +206,7 @@ export function withSequentialExecution( return tools; } - const executionLock = new AsyncMutex(); + const executionLock = new SharedExecutionLock(); const wrappedTools: Record = { ...tools }; for (const [toolName, baseTool] of Object.entries(tools)) { @@ -111,7 +228,9 @@ export function withSequentialExecution( wrappedToolRecord.execute = async (args: unknown, options: unknown) => { const abortSignal = getAbortSignal(options); - await using _lock = await acquireLockOrAbort(executionLock, abortSignal); + await using _lock = canRunWithSiblingExploreTasks(baseTool, args) + ? await executionLock.acquireRead(abortSignal) + : await executionLock.acquireWrite(abortSignal); return await executeFn.call(baseTool, args, options); };