🤖 fix: let built-in forked explore tasks run in parallel

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<string, Tool> = {
+        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);
+    });
   });
 });

src/common/utils/ai/cacheStrategy.ts

@@ -199,6 +199,15 @@ export function applyCacheControlToTools<T extends Record<string, Tool>>(
           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 {

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);
+  });
+});

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<TParameters, TResult>(
+  taskTool: Tool<TParameters, TResult>
+): Tool<TParameters, TResult> {
+  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<symbol, unknown>) | 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<unknown> => {
@@ -663,4 +685,5 @@ export const createTaskTool: ToolFactory = (config: ToolConfiguration) => {
       throw new Error("Task foreground wait ended without a terminal result");
     },
   });
+  return markBuiltInTaskTool(taskTool);
 };

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<T> {
   promise: Promise<T>;
@@ -26,15 +27,16 @@ function createDeferred<T>(): Deferred<T> {
 
 function callWrappedExecute(
   toolRecord: Record<string, unknown>,
+  args: unknown,
   options: unknown
 ): Promise<unknown> {
   const execute = toolRecord.execute;
   if (typeof execute !== "function") {
     throw new Error("Expected wrapped tool execute handler");
   }
 
-  const invoke = execute as (args: Record<string, never>, options: unknown) => Promise<unknown>;
-  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<string, unknown>,
+      {},
       {} as never
     );
     await startedA.promise;
 
-    const secondPromise = callWrappedExecute(
-      wrappedTools!.b as Record<string, unknown>,
-      {
-        abortSignal: controller.signal,
-      } as never
-    );
+    const secondPromise = callWrappedExecute(wrappedTools!.b as Record<string, unknown>, {}, {
+      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<void>(),
+      b: createDeferred<void>(),
+      writer: createDeferred<void>(),
+    };
+    const release = {
+      a: createDeferred<void>(),
+      b: createDeferred<void>(),
+      writer: createDeferred<void>(),
+    };
+
+    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<string, unknown>,
+        { id: "a", agentId: "explore" },
+        {} as never
+      ),
+      callWrappedExecute(
+        wrappedTools.task as Record<string, unknown>,
+        { id: "b", agentId: "explore" },
+        {} as never
+      ),
+      callWrappedExecute(wrappedTools.bash as Record<string, unknown>, {}, {} 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<void>();
+    const releaseA = createDeferred<void>();
+    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<string, unknown>,
+      { agentId: "explore", isolation: "none" },
+      {} as never
+    );
+    await startedA.promise;
+
+    const secondPromise = callWrappedExecute(
+      wrappedTools.task as Record<string, unknown>,
+      { 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<void>(), b: createDeferred<void>() };
+    const release = { a: createDeferred<void>(), b: createDeferred<void>() };
+
+    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<string, unknown>,
+        { id: "a", agentId: " Explore " },
+        {} as never
+      ),
+      callWrappedExecute(
+        wrappedTools.task as Record<string, unknown>,
+        { 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" }]);
+  });
 });