Files
Alex Rudenko 8a8d8771c2 feat: run independent tool calls in parallel with a TaskRunner (#3243)
## Summary

`ToolHandler` held one Puppeteer `Mutex` for every tool call. This PR
replaces it with a central `TaskRunner` that creates `TaskQueue` objects
and runs every task from one ordered list. A queue has an optional
parent queue. A task on a queue starts when every earlier task on the
same queue, an ancestor, or a descendant has finished. Tasks on
unrelated queues run at the same time. Waiting tasks count, so a waiting
task blocks later related tasks and no task starves.

```
TaskRunner (root)
├── context queue (McpContext)
│   ├── page 1 queue
│   └── page 2 queue
└── heap queue (HeapSnapshotManager)
    ├── file a queue
    └── file b queue
```

## User-visible changes

- Heap snapshot tools run in parallel with browser tools.
- `close_heapsnapshot` waits for running reads of the same file.
- With `--pageIdRouting`, page tools that declare `concurrency: 'page'`
and target different pages run in parallel. No tool declares it in this
PR; PR 3 of the stack does.
- `get_devtools_comments` runs at the same time as any other tool.
Before, only its handler ran unlocked.
- A tool call that waits on a page that closes fails with `Page N is
closed.`
- Two first tool calls that wait for the initial `roots/list` share one
request. Before, the second call validated paths against no roots and
failed with "Access denied".

## API

`src/utils/TaskRunner.ts` (new, ~170 lines):

| API | Behavior |
| --- | --- |
| `runner.createQueue(parent?)` | Creates a queue. Without `parent`,
only `pause()` relates to it. |
| `runner.pause()` | Enters the list ahead of waiting tasks and relates
to every queue. Resolves with a guard when running tasks have finished.
Disposing the guard resumes. |
| `queue.scheduleTask()` | Resolves with a guard when no earlier task on
a related queue runs. Callers write `using guard = await
queue.scheduleTask();`. Rejects when the queue is disposed while the
task waits. |
| `queue.dispose(reason)` | Rejects waiting tasks with `reason` and
rejects later `scheduleTask()` calls. The current holder keeps the
queue. |

Every hold is a `Disposable`, released by `using`. A second dispose of a
guard throws. There are no lock modes: holding a queue is exclusive on
that queue and implicitly shared on its ancestors.

Tool definitions declare a `concurrency` mode (`ToolConcurrency`):

| Mode | Queue | Tools |
| --- | --- | --- |
| `'exclusive'` (default) | context queue | every tool that touches
browser-global state |
| `'page'` | `pageTaskQueue(pageId)`; the context queue without page id
routing | page tools that touch only the target page |
| `'none'` | none | browserless tools; `get_devtools_comments` |

`yieldsToolLock` is gone. `get_devtools_comments` only reads page state
and waits for comment events, so it runs outside the runner with
`concurrency: 'none'`. The type system limits the modes per definition
kind: `defineTool` tools cannot take `'page'` or `'none'`, browserless
tools cannot take `'page'`.

## Changes

- `ToolHandler.handle()` resolves the context, picks the queue from
`rawArgs.pageId` and `tool.concurrency`, and holds it with `using guard
= await queue.scheduleTask()` for the rest of the upstream body, which
is unchanged. The disabled check, input parsing, `ensureBrowser()`, page
lookup, handler, and response formatting all run while the queue is
held. A rejected `scheduleTask()` (disposed queue) returns a tool error
result.
- `McpContext` takes a `taskRunner` option, owns `taskQueue`, and
manages page queues keyed by page id in `pageTaskQueue(pageId)`.
`#disposeMcpPage()` disposes the page and its queue in one place.
- `McpContext.createPagesSnapshot()` runs under an internal `Mutex`
because two page tasks can call it at the same time.
- `HeapSnapshotManager` takes the `TaskRunner`, owns a heap queue, and
creates one child queue per absolute file path in
`#scheduleTask(filePath?)`. Every read method starts with `using guard =
await this.#scheduleTask(filePath)`; the two-file diff methods hold the
heap queue; `disposeSnapshot()` holds the file queue.
- `McpServer` owns the `TaskRunner`, passes it to `McpContext` and
`HeapSnapshotManager`, and holds `using guard = await
taskRunner.pause()` in `applyConfig()` and in the `--watchConfig` reload
instead of the mutex.
- `McpServer.#getContext()` stores the initial `roots/list` request as a
promise.
- `defineBrowserlessTool` sets `concurrency: 'none'`.
- `McpPage` is unchanged.

No tool schema or description changes, so `npm run gen` is not needed.

## Stack

1. `refactor: extract errorResult() and routedPageId() in ToolHandler`
(#3251).
2. **This PR**: the runner. No tool sets `concurrency: 'page'` yet, so
page tools still serialize on the context queue after this PR alone.
3. `feat: run page-scoped tools on different pages in parallel` (#3252):
adds `concurrency: 'page'` to the page tools.

## Related issues

No open issue is closed by this PR. Two are related:

- #2934 (a stuck page-side call blocks every other tool): once PR 3
lands, a stuck `'page'` tool (for example `execute_webmcp_tool`) no
longer blocks tools on other pages or heap snapshot tools. It still
blocks `'exclusive'` tools such as `list_pages`, and `evaluate_script`
is still `'exclusive'`. Cancellation is not addressed.
- #2939 (`take_snapshot` stalls for 180 s): same partial effect. The
stall itself is not addressed.

## Tests

- `tests/utils/TaskRunner.test.ts` (new): unrelated queues run together,
different roots run together, FIFO on one queue, parent waits for child,
child waits for parent, child waits behind a waiting parent, second
dispose of a guard throws, release on throw; pause (stops new tasks,
resolves after running tasks finish, goes ahead of waiting tasks);
dispose (rejects waiting tasks with the reason, rejects later runs,
parent runs after a rejected child).
- `tests/ToolHandler.test.ts`: `concurrency: 'exclusive'` suite
(serializes exclusive tools; a `'none'` page tool runs while an
exclusive tool holds the context queue), `concurrency: 'page'` suite
(different pages in parallel, same page serialized, exclusive tool
queued FIFO behind page tools, page queue held until the response is
built, call runs on the page queue, page tool without `pageId` runs on
the context queue, page queue released on throw, disposed page queue
fails the waiting call with `Page 1 is closed.`), browserless suite
(does not enter the context queue, runs while an exclusive tool runs),
`update` suite uses `taskRunner.pause()`.
- `tests/McpContext.test.ts`: one page queue per id with the context
queue as parent; disposing a page rejects a waiting task with `Page N is
closed.`; `createPagesSnapshot()` serializes concurrent calls.
- `tests/index.test.ts`: two concurrent first calls share one
`roots/list` request.
- `tests/processors/HeapSnapshotManager.test.ts`: `disposeSnapshot()`
waits for an in-flight read; a read after the dispose loads the file
again.
- `tests/tools/comments.test.ts`: `get_devtools_comments` has
`concurrency: 'none'`.
- `tests/mocks.ts`: `createMockMcpContext()` carries a real
`TaskRunner`, `taskQueue`, and `pageTaskQueue()`.
2026-10-20 12:59:10 +00:00

272 lines
8.0 KiB
TypeScript

/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import {describe, it} from 'node:test';
import sinon from 'sinon';
import {TaskRunner, type TaskQueue} from '../../src/utils/TaskRunner.js';
async function settle(): Promise<void> {
await new Promise(resolve => setImmediate(resolve));
}
interface ControlledTask {
/** True once the queue is held. */
started: boolean;
/** Settles when the queue is released. */
done: Promise<void>;
/** Releases the queue. */
finish: () => void;
}
/** Acquires `queue` and holds it until `finish()` is called. */
function start(queue: TaskQueue): ControlledTask {
const {promise: release, resolve: finish} = Promise.withResolvers<void>();
const controlled: ControlledTask = {
started: false,
done: Promise.resolve(),
finish,
};
controlled.done = (async () => {
using guard = await queue.scheduleTask();
controlled.started = true;
await release;
})();
return controlled;
}
function createTree() {
const runner = new TaskRunner();
const parent = runner.createQueue();
const childA = runner.createQueue(parent);
const childB = runner.createQueue(parent);
return {runner, parent, childA, childB};
}
describe('TaskRunner', () => {
it('runs tasks on unrelated queues at the same time', async () => {
const {childA, childB} = createTree();
const a = start(childA);
const b = start(childB);
await settle();
assert.strictEqual(a.started, true);
assert.strictEqual(b.started, true);
a.finish();
b.finish();
await Promise.all([a.done, b.done]);
});
it('runs tasks on queues of different roots at the same time', async () => {
const runner = new TaskRunner();
const a = start(runner.createQueue());
const b = start(runner.createQueue());
await settle();
assert.strictEqual(a.started, true);
assert.strictEqual(b.started, true);
a.finish();
b.finish();
await Promise.all([a.done, b.done]);
});
it('serializes holders of one queue in FIFO order', async () => {
const queue = new TaskRunner().createQueue();
const order: number[] = [];
await Promise.all(
[1, 2, 3].map(async id => {
using guard = await queue.scheduleTask();
await settle();
order.push(id);
}),
);
assert.deepStrictEqual(order, [1, 2, 3]);
});
it('orders by queue, not by task duration', async () => {
const clock = sinon.useFakeTimers();
try {
const runner = new TaskRunner();
const order: string[] = [];
const run = async (queue: TaskQueue, name: string, ms: number) => {
using guard = await queue.scheduleTask();
await new Promise(resolve => setTimeout(resolve, ms));
order.push(name);
};
const shared = runner.createQueue();
const sameQueue = Promise.all([
run(shared, 'slow', 300),
run(shared, 'fast', 100),
]);
await clock.tickAsync(1_000);
await sameQueue;
assert.deepStrictEqual(order, ['slow', 'fast']);
order.length = 0;
const unrelated = Promise.all([
run(runner.createQueue(), 'slow', 300),
run(runner.createQueue(), 'fast', 100),
]);
await clock.tickAsync(1_000);
await unrelated;
assert.deepStrictEqual(order, ['fast', 'slow']);
} finally {
clock.restore();
}
});
it('makes a parent task wait for a running child task', async () => {
const {parent, childA} = createTree();
const child = start(childA);
const parentTask = start(parent);
await settle();
assert.strictEqual(parentTask.started, false);
child.finish();
await settle();
assert.strictEqual(parentTask.started, true);
parentTask.finish();
await Promise.all([child.done, parentTask.done]);
});
it('makes a child task wait for a running parent task', async () => {
const {parent, childA} = createTree();
const parentTask = start(parent);
const child = start(childA);
await settle();
assert.strictEqual(child.started, false);
parentTask.finish();
await settle();
assert.strictEqual(child.started, true);
child.finish();
await Promise.all([child.done, parentTask.done]);
});
it('blocks a later child task behind a waiting parent task', async () => {
const {parent, childA, childB} = createTree();
const a = start(childA);
const parentTask = start(parent);
const b = start(childB);
await settle();
assert.strictEqual(b.started, false);
a.finish();
await settle();
assert.strictEqual(parentTask.started, true);
assert.strictEqual(b.started, false);
parentTask.finish();
await settle();
assert.strictEqual(b.started, true);
b.finish();
await Promise.all([a.done, parentTask.done, b.done]);
});
it('throws when a guard is disposed twice', async () => {
const queue = new TaskRunner().createQueue();
const guard = await queue.scheduleTask();
guard[Symbol.dispose]();
assert.throws(
() => guard[Symbol.dispose](),
/The guard was already disposed\./,
);
});
it('releases the queue when the holder throws', async () => {
const queue = new TaskRunner().createQueue();
await assert.rejects(
(async () => {
using guard = await queue.scheduleTask();
throw new Error('task failed');
})(),
/task failed/,
);
const next = start(queue);
await settle();
assert.strictEqual(next.started, true);
next.finish();
await next.done;
});
describe('pause', () => {
it('stops new tasks and resolves after running tasks finish', async () => {
const {runner, childA, childB} = createTree();
const running = start(childA);
await settle();
let paused = false;
const pause = runner.pause().then(guard => {
paused = true;
return guard;
});
const late = start(childB);
await settle();
assert.strictEqual(paused, false);
assert.strictEqual(late.started, false);
running.finish();
await running.done;
const guard = await pause;
await settle();
assert.strictEqual(late.started, false);
guard[Symbol.dispose]();
await settle();
assert.strictEqual(late.started, true);
late.finish();
await late.done;
});
it('goes ahead of waiting tasks', async () => {
const {runner, childA} = createTree();
const running = start(childA);
const waiting = start(childA);
await settle();
const pause = runner.pause();
running.finish();
const guard = await pause;
await settle();
assert.strictEqual(waiting.started, false);
guard[Symbol.dispose]();
await settle();
assert.strictEqual(waiting.started, true);
waiting.finish();
await Promise.all([running.done, waiting.done]);
});
});
describe('dispose', () => {
it('rejects waiting tasks with the reason and keeps the running task', async () => {
const queue = new TaskRunner().createQueue();
const running = start(queue);
const waiting = start(queue);
await settle();
queue.dispose(new Error('Page 1 is closed.'));
await assert.rejects(waiting.done, /Page 1 is closed\./);
assert.strictEqual(running.started, true);
running.finish();
await running.done;
});
it('rejects a later scheduleTask()', async () => {
const queue = new TaskRunner().createQueue();
queue.dispose(new Error('Page 1 is closed.'));
await assert.rejects(queue.scheduleTask(), /Page 1 is closed\./);
});
it('lets a parent task run after a waiting child task is rejected', async () => {
const {parent, childA} = createTree();
const running = start(childA);
const waiting = start(childA);
await settle();
childA.dispose(new Error('closed'));
await assert.rejects(waiting.done, /closed/);
running.finish();
await running.done;
const parentTask = start(parent);
await settle();
assert.strictEqual(parentTask.started, true);
parentTask.finish();
await parentTask.done;
});
});
});