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()`.
This commit is contained in:
Alex Rudenko authored and GitHub committed 2026-10-20 12:59:10 +00:00
1 parent 9a6b03acf7
commit 8a8d8771c2
20 files changed
+1406 -383

No files matched your search

+4 -3
View File
@@ -186,8 +186,9 @@ export class BrowserManager {
* McpContext initialization built on it — if ensureBrowser() already
* resolved and it's that later step hanging (e.g. a dead CDP transport),
* the token rotation alone does nothing, since #browser is already
* cached. Safe to forget unconditionally here: the tool mutex serializes
* every call, so there's no concurrent caller to disrupt.
* cached. Safe to forget unconditionally here: a tool call that already
* received a Browser keeps using it, so forgetting the cached value only
* makes the next ensureBrowser() reconnect.
*/
abandonPendingAttempt(): void {
this.#browserAttempt = new BrowserAttempt();
@@ -361,7 +362,7 @@ export class BrowserManager {
// periodically and terminates it on a missed pong, which rejects the
// in-flight CDP calls and fires `disconnected` so #evictIfCurrent()
// drops the stale handle. Without this, such a call would hang until
// protocolTimeout while holding the tool mutex.
// protocolTimeout inside the context queue.
wsOptions: {
keepAlive: true,
headers: wsHeaders,
+47 -5
View File
@@ -24,6 +24,7 @@ import {type UncaughtError} from './collectors/PageCollector.js';
import {ServiceWorkerConsoleCollector} from './collectors/ServiceWorkerCollector.js';
import {
Locator,
Mutex,
type Browser,
type BrowserContext,
type Connection,
@@ -57,6 +58,7 @@ import type {Logger} from './types.js';
import {McpWorker} from './McpWorker.js';
import {getTempFilePath, resolveCanonicalPath} from './utils/files.js';
import {escapeForLog} from './utils/logger.js';
import type {TaskQueue, TaskRunner} from './utils/TaskRunner.js';
import {isAllowedUrl} from './utils/url.js';
export interface McpContextOptions {
@@ -65,6 +67,8 @@ export interface McpContextOptions {
navigationTimeout?: number;
// Shared heap snapshot manager to preserve loaded snapshots across browser reconnects.
heapSnapshotManager: HeapSnapshotManager;
// Creates the context queue and the page queues.
taskRunner: TaskRunner;
// Returns the browser to use. Called by ensureBrowser(), so the browser
// connection is only established once a tool needs it.
getBrowser: () => Promise<Browser>;
@@ -121,6 +125,19 @@ export class McpContext implements Context, BrowserlessContext {
#workers = new Map<Target, McpWorker>();
#mcpPages = new Map<Target, McpPage>();
/**
* Queue for the browser-wide state (page list, selected page, trace
* state). Every page queue is a child of it, so a task on it waits for
* every page task and blocks later page tasks.
*/
readonly taskQueue: TaskQueue;
#taskRunner: TaskRunner;
#pageTaskQueues = new Map<number, TaskQueue>();
/**
* Serializes `createPagesSnapshot()`. Two tasks on different page queues
* can call it at the same time.
*/
#pagesSnapshotMutex = new Mutex();
#selectedPage?: McpPage;
#selectedPageFallback?: {wasClosed: boolean};
#selectedPageFallbackRead = false;
@@ -150,6 +167,8 @@ export class McpContext implements Context, BrowserlessContext {
this.logger = logger;
this.#locatorClass = locatorClass;
this.#taskRunner = options.taskRunner;
this.taskQueue = options.taskRunner.createQueue();
this.#config = options.config;
this.#options = options;
this.#heapSnapshotManager = options.heapSnapshotManager;
@@ -329,7 +348,7 @@ export class McpContext implements Context, BrowserlessContext {
this.#serviceWorkerConsoleCollector?.dispose();
for (const mcpPage of this.#allMcpPages()) {
mcpPage.dispose();
this.#disposeMcpPage(mcpPage);
}
this.#mcpPages.clear();
this.#workers.clear();
@@ -366,7 +385,7 @@ export class McpContext implements Context, BrowserlessContext {
try {
const mcpPage = this.#mcpPages.get(target);
if (mcpPage) {
mcpPage.dispose();
this.#disposeMcpPage(mcpPage);
this.#mcpPages.delete(target);
}
this.#workers.delete(target);
@@ -672,6 +691,28 @@ export class McpContext implements Context, BrowserlessContext {
return page;
}
/**
* The queue for tasks on page `pageId`. The queue exists before the page
* does, so a tool call can enter it before the browser connects. It is a
* child of {@link taskQueue}.
*/
pageTaskQueue(pageId: number): TaskQueue {
let queue = this.#pageTaskQueues.get(pageId);
if (!queue) {
queue = this.#taskRunner.createQueue(this.taskQueue);
this.#pageTaskQueues.set(pageId, queue);
}
return queue;
}
#disposeMcpPage(mcpPage: McpPage): void {
mcpPage.dispose();
this.#pageTaskQueues
.get(mcpPage.id)
?.dispose(new Error(`Page ${mcpPage.id} is closed.`));
this.#pageTaskQueues.delete(mcpPage.id);
}
isPageSelected(page: McpPage): boolean {
return this.#selectedPage === page;
}
@@ -688,7 +729,7 @@ export class McpContext implements Context, BrowserlessContext {
this.#selectedPage !== newPage &&
!this.#mcpPages.has(this.#selectedPage.target)
) {
this.#selectedPage.dispose();
this.#disposeMcpPage(this.#selectedPage);
}
this.#selectedPage = newPage;
newPage.updateTimeouts();
@@ -791,6 +832,7 @@ export class McpContext implements Context, BrowserlessContext {
if (!this.#browser) {
return [];
}
using guard = await this.#pagesSnapshotMutex.acquire();
const allTargets = this.#fetchPageTargets();
for (const target of allTargets) {
@@ -802,7 +844,7 @@ export class McpContext implements Context, BrowserlessContext {
for (const [target, mcpPage] of this.#mcpPages) {
if (!currentTargets.has(target)) {
if (mcpPage !== this.#selectedPage || mcpPage.isClosed()) {
mcpPage.dispose();
this.#disposeMcpPage(mcpPage);
}
this.#mcpPages.delete(target);
}
@@ -836,7 +878,7 @@ export class McpContext implements Context, BrowserlessContext {
} else if (this.#selectedPage) {
const wasClosed = Boolean(this.#selectedPage.isClosed());
if (!wasClosed) {
this.#selectedPage.dispose();
this.#disposeMcpPage(this.#selectedPage);
}
this.#selectedPage = undefined;
this.#selectedPageFallback = {wasClosed};
+113 -119
View File
@@ -11,7 +11,7 @@ import type {McpPage} from './McpPage.js';
import {McpResponse, type DataFormat} from './McpResponse.js';
import {SlimMcpResponse} from './SlimMcpResponse.js';
import {ClearcutLogger} from './telemetry/ClearcutLogger.js';
import type {Browser, CallToolResult, Mutex} from './third_party/index.js';
import type {Browser, CallToolResult} from './third_party/index.js';
import {ProtocolError, ProtocolErrorCode, zod} from './third_party/index.js';
import {labels} from './tools/categories.js';
import {categoryToFlagName} from './config/category-options.js';
@@ -29,18 +29,13 @@ import {
isPageScopedTool,
isSlimTool,
} from './tools/ToolDefinition.js';
import {getErrorMessage} from './utils/errorHandling.js';
import {logger} from './utils/logger.js';
import type {TaskQueue} from './utils/TaskRunner.js';
import {fileURLToPath, pathToFileURL} from 'node:url';
import {isDeepStrictEqual} from 'node:util';
import {isLocalhost} from './utils/url.js';
type ToolLockGuard = Awaited<ReturnType<Mutex['acquire']>>;
interface ToolLockControl {
releaseIfYielding: () => void;
reacquire: () => Promise<void>;
}
/**
* Upper bound on how long a tool call may wait to acquire the browser
* (connect/launch and McpContext initialization). A connect attempt on a
@@ -223,7 +218,7 @@ function errorResult(text: string): CallToolResult {
/**
* The page id a page tool call targets, or `undefined` when the call targets
* the selected page.
* the selected page. `args` can be the raw or the parsed arguments.
*/
function routedPageId(
tool: AnyToolDefinition,
@@ -250,7 +245,6 @@ export class ToolHandler {
tool: AnyToolDefinition,
config: ReadonlyServerConfig,
private readonly getContext: () => Promise<McpContext>,
private readonly toolMutex: Mutex,
private readonly abandonPendingBrowserAttemptOnTimeout: () => void,
) {
this.#tool = tool;
@@ -283,9 +277,9 @@ export class ToolHandler {
/**
* Replaces the tool definition and re-evaluates its status against the
* current server config. Callers must hold the tool mutex, so calls that
* are already waiting for it run with the new state. Returns true if the
* tool listing visible to the client changed.
* current server config. Callers must pause the task runner, so calls that
* are already waiting run with the new state. Returns true if the tool
* listing visible to the client changed.
*/
update(tool: AnyToolDefinition): boolean {
const previousTool = this.#tool;
@@ -363,16 +357,12 @@ export class ToolHandler {
params: Record<string, unknown>,
response: McpResponse,
context: McpContext,
lock: ToolLockControl,
): ReturnType<McpResponse['handle']> {
try {
await validateToolFiles(tool, params, context, true);
lock.releaseIfYielding();
await tool.handler({params}, response, context);
} catch (err) {
response.setError(err);
} finally {
await lock.reacquire();
}
return await response.handle(context, this.#dataFormat(serverArgs));
}
@@ -383,7 +373,6 @@ export class ToolHandler {
params: Record<string, unknown>,
response: McpResponse,
context: McpContext,
lock: ToolLockControl,
onTelemetryData: (
devToolsData: DevToolsData | undefined,
pageUrl: string | undefined,
@@ -416,7 +405,6 @@ export class ToolHandler {
if (tool.blockedByDialog) {
page.throwIfDialogOpen();
}
lock.releaseIfYielding();
await tool.handler(
{
params,
@@ -426,7 +414,6 @@ export class ToolHandler {
context,
);
} else {
lock.releaseIfYielding();
await tool.handler(
{
params,
@@ -437,8 +424,6 @@ export class ToolHandler {
}
} catch (err) {
response.setError(err);
} finally {
await lock.reacquire();
}
onTelemetryData(
await context.getDevToolsData(page),
@@ -447,113 +432,122 @@ export class ToolHandler {
return await response.handle(context, this.#dataFormat(serverArgs));
}
/**
* The queue the call runs on. `'none'` tools run outside the runner.
* `'page'` tools run on the queue of the page named by `rawArgs.pageId`
* when page id routing applies; every other call runs on the context
* queue.
*/
#queueFor(
context: McpContext,
rawArgs: Record<string, unknown>,
): TaskQueue | undefined {
const tool = this.#tool;
switch (tool.concurrency) {
case 'none':
return undefined;
case 'page': {
const pageId = routedPageId(tool, this.#config.resolved, rawArgs);
return pageId !== undefined
? context.pageTaskQueue(pageId)
: context.taskQueue;
}
default:
return context.taskQueue;
}
}
handle = async (
rawArgs: Record<string, unknown> = {},
options?: {viaMcp?: boolean},
): Promise<CallToolResult> => {
let guard: ToolLockGuard | undefined = await this.toolMutex.acquire();
const context = await this.getContext();
let queueGuard: Disposable | undefined;
try {
if (options?.viaMcp && this.#disabled) {
throw new ProtocolError(
ProtocolErrorCode.InvalidParams,
`Tool ${this.#tool.name} disabled`,
);
}
queueGuard = await this.#queueFor(context, rawArgs)?.scheduleTask();
} catch (err) {
// The queue was disposed while the call waited, e.g. the page closed.
return errorResult(getErrorMessage(err));
}
using guard = queueGuard;
if (options?.viaMcp && this.#disabled) {
throw new ProtocolError(
ProtocolErrorCode.InvalidParams,
`Tool ${this.#tool.name} disabled`,
);
}
if (this.#disabledReason) {
return errorResult(this.#disabledReason);
}
if (this.#disabledReason) {
return errorResult(this.#disabledReason);
}
const tool = this.#tool;
const serverArgs = this.#config.resolved;
const parseResult = this.#registeredInputSchema.safeParse(rawArgs);
if (!parseResult.success) {
return errorResult(formatValidationError(tool.name, parseResult.error));
}
const params = parseResult.data;
const tool = this.#tool;
const serverArgs = this.#config.resolved;
const parseResult = this.#registeredInputSchema.safeParse(rawArgs);
if (!parseResult.success) {
return errorResult(formatValidationError(tool.name, parseResult.error));
}
const params = parseResult.data;
const startTime = Date.now();
let success = false;
let devToolsData: DevToolsData | undefined;
let pageUrl: string | undefined;
try {
logger?.(`${tool.name} request: ${JSON.stringify(params, null, ' ')}`);
const context = await this.getContext();
const response = isSlimTool(tool)
? new SlimMcpResponse(serverArgs)
: new McpResponse(serverArgs);
response.setRedactNetworkHeaders(serverArgs.redactNetworkHeaders);
// When the tool yields the lock, the handler waits for external events.
// Release the lock so other tool calls proceed, and re-acquire it
// before the response is built.
const lock: ToolLockControl = {
releaseIfYielding: () => {
if (tool.yieldsToolLock && guard) {
guard[Symbol.dispose]();
guard = undefined;
}
},
reacquire: async () => {
guard ??= await this.toolMutex.acquire();
},
};
const {content, structuredContent} = isBrowserlessTool(tool)
? await this.#runBrowserlessTool(
tool,
serverArgs,
params,
response,
context,
lock,
)
: await this.#runBrowserTool(
tool,
serverArgs,
params,
response,
context,
lock,
(data, url) => {
devToolsData = data;
pageUrl = url;
},
);
const result: CallToolResult & {
structuredContent?: Record<string, unknown>;
} = {
content,
};
if (response.error) {
result.isError = true;
}
success = true;
if (serverArgs.experimentalStructuredContent) {
result.structuredContent = structuredContent as Record<
string,
unknown
>;
}
return result;
} catch (err) {
logger?.(`${tool.name} error:`, err, err?.stack);
let errorText = err && 'message' in err ? err.message : String(err);
if ('cause' in err && err.cause) {
errorText += `\nCause: ${err.cause.message}`;
}
return errorResult(errorText);
} finally {
void ClearcutLogger.get()?.logToolInvocation({
toolName: tool.name,
params,
schema: tool.schema,
success,
latencyMs: Date.now() - startTime,
devToolsData,
pageUrl,
});
const startTime = Date.now();
let success = false;
let devToolsData: DevToolsData | undefined;
let pageUrl: string | undefined;
try {
logger?.(`${tool.name} request: ${JSON.stringify(params, null, ' ')}`);
const response = isSlimTool(tool)
? new SlimMcpResponse(serverArgs)
: new McpResponse(serverArgs);
response.setRedactNetworkHeaders(serverArgs.redactNetworkHeaders);
const {content, structuredContent} = isBrowserlessTool(tool)
? await this.#runBrowserlessTool(
tool,
serverArgs,
params,
response,
context,
)
: await this.#runBrowserTool(
tool,
serverArgs,
params,
response,
context,
(data, url) => {
devToolsData = data;
pageUrl = url;
},
);
const result: CallToolResult & {
structuredContent?: Record<string, unknown>;
} = {
content,
};
if (response.error) {
result.isError = true;
}
success = true;
if (serverArgs.experimentalStructuredContent) {
result.structuredContent = structuredContent as Record<string, unknown>;
}
return result;
} catch (err) {
logger?.(`${tool.name} error:`, err, err?.stack);
let errorText = err && 'message' in err ? err.message : String(err);
if ('cause' in err && err.cause) {
errorText += `\nCause: ${err.cause.message}`;
}
return errorResult(errorText);
} finally {
guard?.[Symbol.dispose]();
void ClearcutLogger.get()?.logToolInvocation({
toolName: tool.name,
params,
schema: tool.schema,
success,
latencyMs: Date.now() - startTime,
devToolsData,
pageUrl,
});
}
};
}
+22 -18
View File
@@ -26,7 +26,6 @@ import {
type RegisteredTool,
type Root,
type Transport,
Mutex,
puppeteer,
} from './third_party/index.js';
import {ToolHandler} from './ToolHandler.js';
@@ -35,15 +34,16 @@ import {createTools} from './tools/tools.js';
import {logCruxDisclaimer} from './utils/disclaimers.js';
import {getErrorMessage} from './utils/errorHandling.js';
import {logger} from './utils/logger.js';
import {TaskRunner} from './utils/TaskRunner.js';
import {VERSION} from './version.js';
puppeteer.setFollowSymlinks(false);
/**
* Timeout for a `roots/list` that a tool call is waiting on, matching the 5s
* default used for page operations. `getContext()` awaits it while
* `ToolHandler` holds the tool mutex, so leaving it unbounded lets a client
* that negotiates `roots` but does not answer block every tool for the SDK's
* default used for page operations. `getContext()` awaits it before
* `ToolHandler` enters the task runner, so leaving it unbounded lets a client
* that negotiates `roots` but does not answer delay every tool for the SDK's
* default of 60s. Background refreshes are not bounded by this, so roots a
* slow client sends late still land.
*/
@@ -64,7 +64,9 @@ export class McpServer {
#config: ServerConfig;
#browserManager: BrowserManager;
#context: McpContext;
#heapSnapshotManager = new HeapSnapshotManager();
/** Runs every tool call and config reload. */
#taskRunner = new TaskRunner();
#heapSnapshotManager = new HeapSnapshotManager(this.#taskRunner);
/**
* Client roots stay valid across browser reconnects and only the client can
@@ -73,11 +75,11 @@ export class McpServer {
*/
#lastClientRoots?: Root[];
/**
* Whether a tool call already waited for the initial `roots/list`, so a
* client that negotiates roots but never answers only delays the first call.
* The bounded initial `roots/list` request. Every tool call that starts
* before the client answered waits on the same request, so a client that
* negotiates roots but never answers only delays the first calls.
*/
#awaitedInitialRoots = false;
#toolMutex = new Mutex();
#initialRoots?: Promise<void>;
#configWatcher?: ConfigWatcher;
#tools = new Map<string, ToolEntry>();
@@ -113,6 +115,7 @@ export class McpServer {
this.#context = McpContext.create(logger, {
config: this.#config,
heapSnapshotManager: this.#heapSnapshotManager,
taskRunner: this.#taskRunner,
getBrowser: () => this.#browserManager.ensureBrowser(),
});
this.#context.setRoots(this.#combinedRoots());
@@ -180,7 +183,9 @@ export class McpServer {
* notification.
*/
async applyConfig(config: ServerConfig): Promise<void> {
using guard = await this.#toolMutex.acquire();
// Goes ahead of waiting tool calls and waits for running ones, so every
// waiting call runs with the new config.
using guard = await this.#taskRunner.pause();
await this.#applyConfigLocked(this.#config.mergeReloadable(config));
}
@@ -253,7 +258,7 @@ export class McpServer {
configPath: string,
reload: () => ServerConfig,
): Promise<void> {
using guard = await this.#toolMutex.acquire();
using guard = await this.#taskRunner.pause();
if (this.#configWatcher !== watcher) {
return;
}
@@ -364,11 +369,11 @@ export class McpServer {
}
async #getContext(): Promise<McpContext> {
if (this.#lastClientRoots === undefined && !this.#awaitedInitialRoots) {
this.#awaitedInitialRoots = true;
// Nothing listed yet, so this call has to wait – bounded, since it is
// holding the tool mutex, and a later background refresh still lands
await this.#updateRoots(ROOTS_REQUEST_TIMEOUT);
if (this.#lastClientRoots === undefined) {
// Nothing listed yet, so this call has to wait – bounded, since every
// tool call waits on it, and a later background refresh still lands
this.#initialRoots ??= this.#updateRoots(ROOTS_REQUEST_TIMEOUT);
await this.#initialRoots;
}
return this.#context;
}
@@ -378,14 +383,13 @@ export class McpServer {
tool,
this.#config,
() => this.#getContext(),
this.#toolMutex,
() => this.#browserManager.abandonPendingAttempt(),
);
}
/**
* Overrides the SDK's `tools/call` handler so enabled and input-schema
* checks run inside `ToolHandler.handle()` after acquiring `#toolMutex`.
* checks run inside `ToolHandler.handle()` inside the task runner.
*/
#setToolCallRequestHandler(): void {
this.server.server.setRequestHandler('tools/call', async request => {
+49
View File
@@ -9,6 +9,7 @@ import path from 'node:path';
import {DevTools} from '../third_party/index.js';
import {createIdGenerator} from '../utils/id.js';
import type {TaskQueue, TaskRunner} from '../utils/TaskRunner.js';
export interface HeapSnapshotAggregate {
id: number;
@@ -129,6 +130,38 @@ function hasValidHeapSnapshotExtension(filePath: string): boolean {
export class HeapSnapshotManager {
#snapshotIdGenerator = createIdGenerator();
#snapshots = new Map<string, Promise<HeapSnapshot>>();
/**
* Queue for all heap snapshot work, with one child queue per absolute file
* path. A read of one file holds the file queue, so reads of different
* files run at the same time. A read of several files holds this queue.
* `disposeSnapshot()` holds the file queue, so it waits for in-flight
* reads of the file, and a read that starts after it loads the file again.
*/
#taskQueue: TaskQueue;
#taskRunner: TaskRunner;
#fileTaskQueues = new Map<string, TaskQueue>();
constructor(taskRunner: TaskRunner) {
this.#taskRunner = taskRunner;
this.#taskQueue = taskRunner.createQueue();
}
/**
* Holds the queue of `filePath`, or the heap queue without a path, until
* the guard is disposed.
*/
async #scheduleTask(filePath?: string): Promise<Disposable> {
if (filePath === undefined) {
return await this.#taskQueue.scheduleTask();
}
const absolutePath = path.resolve(filePath);
let queue = this.#fileTaskQueues.get(absolutePath);
if (!queue) {
queue = this.#taskRunner.createQueue(this.#taskQueue);
this.#fileTaskQueues.set(absolutePath, queue);
}
return await queue.scheduleTask();
}
async getSnapshot(filePath: string): Promise<HeapSnapshot> {
if (!hasValidHeapSnapshotExtension(filePath)) {
@@ -174,6 +207,7 @@ export class HeapSnapshotManager {
filterName?: string,
objectId?: number,
): Promise<HeapSnapshotAggregateData> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const filter =
new DevTools.HeapSnapshotModel.HeapSnapshotModel.NodeFilter();
@@ -211,6 +245,7 @@ export class HeapSnapshotManager {
async getStats(
filePath: string,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
return await snapshot.proxy.getStatistics();
}
@@ -218,6 +253,7 @@ export class HeapSnapshotManager {
async getStaticData(
filePath: string,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
return snapshot.proxy.staticData;
}
@@ -225,6 +261,7 @@ export class HeapSnapshotManager {
async getNativeContextSizes(
filePath: string,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.NativeContextSizes> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
return await snapshot.proxy.getNativeContextSizes();
}
@@ -232,6 +269,7 @@ export class HeapSnapshotManager {
async getRetainedByContextSummary(
filePath: string,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.RetainedByContextSummary> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
return await snapshot.proxy.getRetainedByContextSummary();
}
@@ -242,6 +280,7 @@ export class HeapSnapshotManager {
filterName?: string,
objectId?: number,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const filter =
new DevTools.HeapSnapshotModel.HeapSnapshotModel.NodeFilter();
@@ -262,6 +301,7 @@ export class HeapSnapshotManager {
filePath: string,
nodeId: number,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const nodeIndex = await snapshot.proxy.nodeIndexForId(nodeId);
if (nodeIndex === undefined) {
@@ -275,6 +315,7 @@ export class HeapSnapshotManager {
filePath: string,
nodeId: number,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ObjectInfo> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const nodeIndex = await snapshot.proxy.nodeIndexForId(nodeId);
if (nodeIndex === undefined) {
@@ -286,6 +327,7 @@ export class HeapSnapshotManager {
async analyzeContexts(
filePath: string,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ContextAnalysisResult> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
return await snapshot.proxy.analyzeContexts();
}
@@ -297,6 +339,7 @@ export class HeapSnapshotManager {
maxNodes?: number,
maxSiblings?: number,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.RetainingPaths> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const nodeIndex = await snapshot.proxy.nodeIndexForId(nodeId);
if (nodeIndex === undefined) {
@@ -314,6 +357,7 @@ export class HeapSnapshotManager {
filePath: string,
nodeId: number,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.DominatorChain> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const nodeIndex = await snapshot.proxy.nodeIndexForId(nodeId);
if (nodeIndex === undefined) {
@@ -327,6 +371,7 @@ export class HeapSnapshotManager {
nodeId: number,
options?: HeapEdgesQueryOptions,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const nodeIndex = await snapshot.proxy.nodeIndexForId(nodeId);
if (nodeIndex === undefined) {
@@ -401,6 +446,7 @@ export class HeapSnapshotManager {
baseFilePath: string,
currentFilePath: string,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Diff[]> {
using guard = await this.#scheduleTask();
const [baseSnapshot, currentSnapshot] = await Promise.all([
this.getSnapshot(baseFilePath),
this.getSnapshot(currentFilePath),
@@ -480,6 +526,7 @@ export class HeapSnapshotManager {
}
async getDuplicateStrings(filePath: string): Promise<DuplicateStringGroup[]> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
return await snapshot.proxy.getDuplicateStrings();
}
@@ -488,6 +535,7 @@ export class HeapSnapshotManager {
filePath: string,
options: HeapQueryOptions,
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
using guard = await this.#scheduleTask(filePath);
const snapshot = await this.getSnapshot(filePath);
const provider = snapshot.proxy.queryObjects(options);
return await provider.serializeItemsRange(0, Infinity);
@@ -498,6 +546,7 @@ export class HeapSnapshotManager {
}
async disposeSnapshot(filePath: string): Promise<boolean> {
using guard = await this.#scheduleTask(filePath);
const absolutePath = path.resolve(filePath);
const cached = this.#snapshots.get(absolutePath);
if (cached) {
+38 -5
View File
@@ -78,6 +78,24 @@ export type FileVerificationOption =
type AllKeys<T> = T extends unknown ? keyof T : never;
/**
* Which task queue `ToolHandler` runs a tool call on. A task on a queue
* waits for tasks on the same queue, on its parent, and on its children.
*
* - `'exclusive'`: the context queue, for the whole call. No other browser
* tool runs at the same time. This is the default.
* - `'page'`: the queue of the page named by `pageId`, a child of the
* context queue. Tools on different pages run at the same time; tools on
* one page and `'exclusive'` tools wait. Without page id routing the call
* runs on the context queue. Only for page-scoped tools whose handler
* touches the target page and no browser-global state.
* - `'none'`: no queue. The tool runs at the same time as any other tool.
* For browserless tools, whose data guards itself (see
* `HeapSnapshotManager`), and for page tools that only read page state
* and wait for external events.
*/
export type ToolConcurrency = 'exclusive' | 'page' | 'none';
type ExtractSchemaField<Schema, K extends PropertyKey> = Schema extends unknown
? K extends keyof Schema
? Exclude<Schema[K], undefined>
@@ -119,12 +137,10 @@ export interface BaseToolDefinition<
schema: Schema;
blockedByDialog: boolean;
/**
* When true, the handler runs without the server tool lock, so other tool
* calls proceed while this handler waits. Page lookup before the handler
* and response formatting after it still run under the lock. Only set
* this for tools that wait for external events.
* Which queues the server holds while the tool runs. Defaults to
* `'exclusive'`. See {@link ToolConcurrency}.
*/
yieldsToolLock?: boolean;
concurrency?: ToolConcurrency;
verifyFilesSchema: Partial<
Record<keyof MergeSchema<Schema>, FileVerificationOption>
>;
@@ -139,6 +155,11 @@ export interface ToolDefinition<
* non-page tool must handle dialogs itself.
*/
blockedByDialog: false;
/**
* A non-page tool has no target page, so it cannot take the page queue. It
* touches browser state, so it always runs in a queue.
*/
concurrency?: Exclude<ToolConcurrency, 'page' | 'none'>;
handler(
request: Request<Schema>,
response: Response,
@@ -525,6 +546,11 @@ export function defineTool<Schema extends zod.ZodRawShape>(
interface PageToolDefinition<
Schema extends zod.ZodRawShape = zod.ZodRawShape,
> extends BaseToolDefinition<Schema> {
/**
* A page tool with `'none'` must only read page state. Its page lookup and
* response formatting run at the same time as other tools on the page.
*/
concurrency?: ToolConcurrency;
handler(
request: Request<Schema> & {page: ContextPage},
response: Response,
@@ -563,6 +589,10 @@ export interface BrowserlessToolDefinition<
* dialogs.
*/
blockedByDialog: false;
/**
* A browserless tool has no page.
*/
concurrency?: Exclude<ToolConcurrency, 'page'>;
handler(
request: Request<Schema>,
response: Response,
@@ -573,6 +603,8 @@ export interface BrowserlessToolDefinition<
/**
* Defines a tool that does not need a browser connection, e.g. a tool that only
* analyzes local files. Its handler only receives a {@link BrowserlessContext}.
* Such a tool runs without the context queue (`concurrency: 'none'`) unless
* the definition sets another value.
*/
export function defineBrowserlessTool<Schema extends zod.ZodRawShape>(
definition: (
@@ -581,6 +613,7 @@ export function defineBrowserlessTool<Schema extends zod.ZodRawShape>(
): (args: ParsedArguments) => BrowserlessToolDefinition<Schema> {
return (args: ParsedArguments): BrowserlessToolDefinition<Schema> => {
return {
concurrency: 'none',
...definition(args),
browserless: true,
};
+3 -2
View File
@@ -69,8 +69,9 @@ export const getDevtoolsComments = definePageTool(() => ({
),
},
blockedByDialog: false,
// The handler waits for comment events. Other tools run during the wait.
yieldsToolLock: true,
// The handler only reads page state and waits for comment events, so it
// runs outside the task runner and other tools run during the wait.
concurrency: 'none',
verifyFilesSchema: {},
handler: async (request, response) => {
const page = request.page;
+176
View File
@@ -0,0 +1,176 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
type Priority = 'normal' | 'high';
interface Entry {
queue: TaskQueue;
priority: Priority;
running: boolean;
resolve: () => void;
reject: (error: Error) => void;
}
/** The runner methods a queue uses. Not part of the public API. */
interface QueueHooks {
scheduleTask(queue: TaskQueue): Promise<Disposable>;
cancel(queue: TaskQueue, reason: Error): void;
}
/**
* A queue that the {@link TaskRunner} created. Tasks on one queue run one
* at a time. Tasks on a queue and on its ancestors or descendants also run
* one at a time. Tasks on unrelated queues run at the same time.
*/
export class TaskQueue {
readonly parent?: TaskQueue;
readonly #hooks: QueueHooks;
#disposed?: Error;
constructor(hooks: QueueHooks, parent: TaskQueue | undefined) {
this.#hooks = hooks;
this.parent = parent;
}
/**
* Schedules a task on this queue. Resolves when no earlier task on a
* related queue runs, with a guard that holds the queue until it is
* disposed. Rejects when the queue is disposed while the task waits. A
* second dispose of the guard throws.
*/
async scheduleTask(): Promise<Disposable> {
if (this.#disposed) {
throw this.#disposed;
}
return await this.#hooks.scheduleTask(this);
}
/**
* Rejects the callers that wait on this queue with `reason`, and rejects
* every later `scheduleTask()` call. The current holder keeps the queue.
*/
dispose(reason: Error): void {
this.#disposed = reason;
this.#hooks.cancel(this, reason);
}
/** True when `other` is this queue, an ancestor, or a descendant. */
isRelatedTo(other: TaskQueue): boolean {
return this.#isAncestorOrSelf(other) || other.#isAncestorOrSelf(this);
}
#isAncestorOrSelf(other: TaskQueue): boolean {
for (
let queue: TaskQueue | undefined = other;
queue;
queue = queue.parent
) {
if (queue === this) {
return true;
}
}
return false;
}
}
/**
* Runs every task of the server from one ordered list. Each task names the
* queue it runs on. A task starts when every earlier task on a related queue
* has finished. Waiting tasks count, so a waiting task blocks later related
* tasks and no task starves. `pause()` enters the list ahead of waiting tasks
* and relates to every queue.
*/
export class TaskRunner {
readonly #hooks: QueueHooks = {
scheduleTask: queue => this.#scheduleTask(queue),
cancel: (queue, reason) => this.#cancel(queue, reason),
};
readonly #root = new TaskQueue(this.#hooks, undefined);
#entries: Entry[] = [];
/**
* Creates a queue. Tasks on the new queue and tasks on `parent` wait for
* each other. Without `parent`, the queue is a child of the root, so only
* `pause()` relates to it.
*/
createQueue(parent: TaskQueue = this.#root): TaskQueue {
return new TaskQueue(this.#hooks, parent);
}
/**
* Stops new tasks on every queue. The pause goes ahead of waiting tasks and
* resolves when every running task has finished. Disposing the guard
* resumes the runner.
*/
async pause(): Promise<Disposable> {
return await this.#scheduleTask(this.#root, 'high');
}
async #scheduleTask(
queue: TaskQueue,
priority: Priority = 'normal',
): Promise<Disposable> {
const {guard, started} = this.#enter(queue, priority);
await started;
return guard;
}
#cancel(queue: TaskQueue, reason: Error): void {
for (const entry of [...this.#entries]) {
if (entry.queue === queue && !entry.running) {
this.#leave(entry);
entry.reject(reason);
}
}
}
#enter(
queue: TaskQueue,
priority: Priority,
): {guard: Disposable; started: Promise<void>} {
const {promise, resolve, reject} = Promise.withResolvers<void>();
const entry: Entry = {queue, priority, running: false, resolve, reject};
// A high-priority entry goes before the first waiting normal entry.
const index =
priority === 'high'
? this.#entries.findIndex(e => !e.running && e.priority === 'normal')
: -1;
this.#entries.splice(index === -1 ? this.#entries.length : index, 0, entry);
this.#schedule();
return {
started: promise,
guard: {
[Symbol.dispose]: () => {
this.#leave(entry);
},
},
};
}
#leave(entry: Entry): void {
const index = this.#entries.indexOf(entry);
if (index === -1) {
throw new Error('The guard was already disposed.');
}
this.#entries.splice(index, 1);
this.#schedule();
}
#schedule(): void {
for (const [index, entry] of this.#entries.entries()) {
if (entry.running) {
continue;
}
const blocked = this.#entries
.slice(0, index)
.some(earlier => earlier.queue.isRelatedTo(entry.queue));
if (!blocked) {
entry.running = true;
entry.resolve();
}
}
}
}
+1 -1
View File
@@ -53,7 +53,7 @@ export class WaitForHelper {
async waitForStableDom(): Promise<void> {
// Bound the setup evaluation against the stable-DOM timeout. Without this
// cap a paused renderer (e.g. an open dialog) would make evaluateHandle
// hang until protocolTimeout (default 180s) while the tool mutex is held.
// hang until protocolTimeout (default 180s) inside the context queue.
using stableDomObserver = await Promise.race([
this.#page.evaluateHandle(timeout => {
let timeoutId: ReturnType<typeof setTimeout>;
+65 -8
View File
@@ -20,6 +20,7 @@ import {NetworkFormatter} from '../src/formatters/NetworkFormatter.js';
import {McpContext} from '../src/McpContext.js';
import {McpPage} from '../src/McpPage.js';
import {HeapSnapshotManager} from '../src/processors/HeapSnapshotManager.js';
import {TaskRunner} from '../src/utils/TaskRunner.js';
import {TextSnapshot} from '../src/TextSnapshot.js';
import {
type Browser,
@@ -117,7 +118,8 @@ describe('McpContext', () => {
{
config: options?.config ?? createMockServerConfig(),
heapSnapshotManager:
options?.heapSnapshotManager ?? new HeapSnapshotManager(),
options?.heapSnapshotManager ??
new HeapSnapshotManager(new TaskRunner()),
},
Locator,
);
@@ -133,7 +135,7 @@ describe('McpContext', () => {
config: createMockServerConfig({
performanceCrux: false,
}),
heapSnapshotManager: new HeapSnapshotManager(),
heapSnapshotManager: new HeapSnapshotManager(new TaskRunner()),
},
Locator,
);
@@ -180,7 +182,7 @@ describe('McpContext', () => {
config: createMockServerConfig({
performanceCrux: false,
}),
heapSnapshotManager: new HeapSnapshotManager(),
heapSnapshotManager: new HeapSnapshotManager(new TaskRunner()),
},
Locator,
);
@@ -767,7 +769,7 @@ describe('McpContext', () => {
config: createMockServerConfig({
performanceCrux: false,
}),
heapSnapshotManager: new HeapSnapshotManager(),
heapSnapshotManager: new HeapSnapshotManager(new TaskRunner()),
};
const first = await createConnectedMcpContext(browser, options, Locator);
const idBeforeReconnect = (await first.newPage()).id;
@@ -807,13 +809,15 @@ describe('McpContext', () => {
describe('lazy browser connection', () => {
function createLazyContext(getBrowser: () => Promise<Browser>) {
const taskRunner = new TaskRunner();
return McpContext.create(
undefined,
{
config: createMockServerConfig({
performanceCrux: false,
}),
heapSnapshotManager: new HeapSnapshotManager(),
heapSnapshotManager: new HeapSnapshotManager(taskRunner),
taskRunner,
getBrowser,
},
Locator,
@@ -826,6 +830,59 @@ describe('McpContext', () => {
return browser;
}
it('serializes concurrent createPagesSnapshot() calls', async () => {
const browser = createBrowser();
browser.targets.returns([createMockPuppeteerTarget()]);
const {promise: initDone, resolve: finishInit} =
Promise.withResolvers<void>();
const init = sinon.stub(McpPage.prototype, 'init').resolves();
init.onSecondCall().returns(initDone);
const context = createLazyContext(async () => browser);
await context.ensureBrowser();
sinon.assert.calledOnce(init);
const first = context.createPagesSnapshot();
const second = context.createPagesSnapshot();
await new Promise(resolve => setImmediate(resolve));
sinon.assert.calledTwice(init);
finishInit();
await Promise.all([first, second]);
sinon.assert.calledThrice(init);
context.dispose();
});
it('returns one page task queue per page id', async () => {
const context = createLazyContext(async () => createBrowser());
const queue = context.pageTaskQueue(1);
assert.strictEqual(context.pageTaskQueue(1), queue);
assert.notStrictEqual(context.pageTaskQueue(2), queue);
assert.strictEqual(queue.parent, context.taskQueue);
context.dispose();
});
it('disposes the page task queue with the page', async () => {
const browser = createBrowser();
browser.targets.returns([createMockPuppeteerTarget()]);
sinon.stub(McpPage.prototype, 'init').resolves();
const context = createLazyContext(async () => browser);
await context.ensureBrowser();
const page = context.getSelectedMcpPage();
const {promise: release, resolve: finishRunning} =
Promise.withResolvers<void>();
const running = (async () => {
using guard = await context.pageTaskQueue(page.id).scheduleTask();
await release;
})();
const waiting = context.pageTaskQueue(page.id).scheduleTask();
context.dispose();
await assert.rejects(waiting, new RegExp(`Page ${page.id} is closed\\.`));
finishRunning();
await running;
});
it('does not connect until ensureBrowser() is called', async () => {
const browser = createBrowser();
const getBrowser = sinon.stub<[], Promise<Browser>>().resolves(browser);
@@ -1338,7 +1395,7 @@ describe('McpContext', () => {
allowedUrlPattern: options.allowlist,
blockedUrlPattern: options.blocklist,
}),
heapSnapshotManager: new HeapSnapshotManager(),
heapSnapshotManager: new HeapSnapshotManager(new TaskRunner()),
},
Locator,
);
@@ -1799,7 +1856,7 @@ describe('McpContext', () => {
});
it('disposes loaded heap snapshots', async () => {
const heapSnapshotManager = new HeapSnapshotManager();
const heapSnapshotManager = new HeapSnapshotManager(new TaskRunner());
const context = await createMockContext({
heapSnapshotManager,
});
@@ -1813,7 +1870,7 @@ describe('McpContext', () => {
});
it('does not dispose heapSnapshotManager when no snapshots are loaded', async () => {
const heapSnapshotManager = new HeapSnapshotManager();
const heapSnapshotManager = new HeapSnapshotManager(new TaskRunner());
const disposeSpy = sinon.spy(heapSnapshotManager, 'dispose');
const context = await createMockContext({heapSnapshotManager});
+506 -205
View File
File diff suppressed because it is too large. Load diff
@@ -22,6 +22,7 @@ import {
} from '../../src/processors/HeapSnapshotManager.js';
import {DevTools} from '../../src/third_party/index.js';
import {parseByteSizeRange} from '../../src/utils/bytes.js';
import {TaskRunner} from '../../src/utils/TaskRunner.js';
import {createMockContextAnalysisResult} from '../mocks.js';
const {formatBytesToKb} = DevTools.I18n.ByteUtilities;
@@ -500,7 +501,7 @@ describe('HeapSnapshotFormatter', () => {
const heap3Path = join(process.cwd(), 'tests/fixtures/heap-3.heapsnapshot');
before(() => {
manager = new HeapSnapshotManager();
manager = new HeapSnapshotManager(new TaskRunner());
});
after(() => {
+44 -1
View File
@@ -239,11 +239,54 @@ describe('e2e', {concurrency: true}, () => {
);
});
it('makes concurrent first calls wait for the same roots/list', async () => {
const fixturesRoot = path.resolve('tests/fixtures');
await withClient(
async client => {
// Answers late enough that both calls start before the roots arrive.
// Browserless tools hold no lock, so without a shared initial roots
// request the second call validates its path against no roots.
client.setRequestHandler('roots/list', async () => {
await new Promise(resolve => setTimeout(resolve, 1_000));
return {
roots: [{uri: pathToFileURL(fixturesRoot).href, name: 'fixtures'}],
};
});
const results = await Promise.all([
client.callTool({
name: 'get_heapsnapshot_summary',
arguments: {
filePath: path.join(fixturesRoot, 'example.heapsnapshot'),
},
}),
client.callTool({
name: 'get_heapsnapshot_summary',
arguments: {
filePath: path.join(fixturesRoot, 'example.heapsnapshot'),
},
}),
]);
for (const result of results) {
const content = result.content as TextContent[];
assert.strictEqual(result.isError, undefined, content[0]?.text);
}
},
['--memoryDebugging'],
{
capabilities: {
roots: {listChanged: true},
},
},
);
});
it('does not block tools if the client never answers roots/list', async () => {
await withClient(
async client => {
// A client that negotiates roots but never responds. getContext()
// awaits updateRoots() while holding the tool mutex, so an unbounded
// awaits updateRoots() while holding the context lock, so an unbounded
// request would stall this call for the SDK default of 60s.
client.setRequestHandler('roots/list', () => {
return new Promise<never>(() => {
+19 -1
View File
@@ -35,6 +35,7 @@ import {McpContext} from '../src/McpContext.js';
import {McpPage} from '../src/McpPage.js';
import {McpResponse} from '../src/McpResponse.js';
import {McpWorker, type WorkerType} from '../src/McpWorker.js';
import {TaskRunner, type TaskQueue} from '../src/utils/TaskRunner.js';
import type {
DuplicateStringGroup,
HeapSnapshotAggregate,
@@ -435,12 +436,16 @@ export function createMockPuppeteerTarget(
export function createMockMcpPage(
options: {
id?: number;
pptrPage?: MockPuppeteerPage;
target?: sinon.SinonStubbedInstance<MockTarget>;
} = {},
): MockMcpPage {
const page = sinon.createStubInstance(McpPage);
page.emulationSettings = {};
if (options.id !== undefined) {
Object.defineProperty(page, 'id', {value: options.id});
}
const pptrPage = options.pptrPage ?? createMockPuppeteerPage();
const target = options.target ?? createMockPuppeteerTarget({page: pptrPage});
page.init.resolves();
@@ -539,12 +544,25 @@ export function createMockMcpWorker(
}
export function createMockMcpContext(
options: {selectedPage?: MockMcpPage} = {},
options: {selectedPage?: MockMcpPage; taskRunner?: TaskRunner} = {},
): MockMcpContext {
const context = sinon.createStubInstance(McpContext);
const page = options.selectedPage ?? createMockMcpPage();
context.getSelectedMcpPage.returns(page satisfies McpPage);
context.getTargetInfos.resolves(new Map());
// Real queues, so handler tests observe the real scheduling.
const taskRunner = options.taskRunner ?? new TaskRunner();
const taskQueue = taskRunner.createQueue();
Object.defineProperty(context, 'taskQueue', {value: taskQueue});
const pageTaskQueues = new Map<number, TaskQueue>();
context.pageTaskQueue.callsFake(pageId => {
let queue = pageTaskQueues.get(pageId);
if (!queue) {
queue = taskRunner.createQueue(taskQueue);
pageTaskQueues.set(pageId, queue);
}
return queue;
});
return context;
}
+31 -7
View File
@@ -10,11 +10,12 @@ import {describe, it, after, before} from 'node:test';
import sinon from 'sinon';
import {HeapSnapshotManager} from '../../src/processors/HeapSnapshotManager.js';
import {TaskRunner} from '../../src/utils/TaskRunner.js';
import {DevTools} from '../../src/third_party/index.js';
describe('HeapSnapshotManager', () => {
it('rejects when a file without .heapsnapshot or .heaptimeline extension is passed', async () => {
const manager = new HeapSnapshotManager();
const manager = new HeapSnapshotManager(new TaskRunner());
await assert.rejects(
manager.getSnapshot('tests/fixtures/snapshot_diffs.js'),
@@ -29,7 +30,7 @@ describe('HeapSnapshotManager', () => {
'dispose',
);
const manager = new HeapSnapshotManager();
const manager = new HeapSnapshotManager(new TaskRunner());
// A path that passes into #loadSnapshot but fails on read. The worker is
// created before the read, so a failed load must still dispose it,
@@ -51,7 +52,7 @@ describe('HeapSnapshotManager', () => {
let manager: HeapSnapshotManager;
before(() => {
manager = new HeapSnapshotManager();
manager = new HeapSnapshotManager(new TaskRunner());
});
after(() => {
@@ -108,7 +109,7 @@ describe('HeapSnapshotManager', () => {
let manager: HeapSnapshotManager;
before(() => {
manager = new HeapSnapshotManager();
manager = new HeapSnapshotManager(new TaskRunner());
});
after(() => {
@@ -124,7 +125,7 @@ describe('HeapSnapshotManager', () => {
});
it('resolves class IDs in getNodesById without calling getAggregates first', async () => {
const freshManager = new HeapSnapshotManager();
const freshManager = new HeapSnapshotManager(new TaskRunner());
try {
const {aggregates} = await manager.getAggregates(filePath);
const functionAggregate = aggregates.find(
@@ -150,7 +151,7 @@ describe('HeapSnapshotManager', () => {
});
it('assigns deterministic class IDs even when getAggregates is called with a filter first', async () => {
const filteredFirstManager = new HeapSnapshotManager();
const filteredFirstManager = new HeapSnapshotManager(new TaskRunner());
try {
const unfilteredFirst = await manager.getAggregates(filePath);
const unfilteredById = new Map(
@@ -182,7 +183,7 @@ describe('HeapSnapshotManager', () => {
});
it('gets stats and staticData, and disposes snapshot', async () => {
const manager = new HeapSnapshotManager();
const manager = new HeapSnapshotManager(new TaskRunner());
try {
const filePath = 'tests/fixtures/example.heapsnapshot';
const stats = await manager.getStats(filePath);
@@ -201,4 +202,27 @@ describe('HeapSnapshotManager', () => {
manager.dispose();
}
});
it('makes disposeSnapshot wait for an in-flight read of the same file', async () => {
const manager = new HeapSnapshotManager(new TaskRunner());
try {
const filePath = 'tests/fixtures/example.heapsnapshot';
const order: string[] = [];
const read = manager.getStats(filePath).then(stats => {
order.push('read');
return stats;
});
const dispose = manager.disposeSnapshot(filePath).then(disposed => {
order.push('dispose');
return disposed;
});
const [stats, disposed] = await Promise.all([read, dispose]);
assert.ok(stats.total > 0);
assert.strictEqual(disposed, true);
assert.deepStrictEqual(order, ['read', 'dispose']);
assert.strictEqual(manager.hasSnapshots(), false);
} finally {
manager.dispose();
}
});
});
+2 -1
View File
@@ -15,6 +15,7 @@ import sinon from 'sinon';
import type {McpContext} from '../src/McpContext.js';
import {HeapSnapshotManager} from '../src/processors/HeapSnapshotManager.js';
import {TaskRunner} from '../src/utils/TaskRunner.js';
import {resolveCanonicalPath} from '../src/utils/files.js';
import {escapeForLog} from '../src/utils/logger.js';
@@ -176,7 +177,7 @@ describe('McpContext path validation escaping', () => {
config: createMockServerConfig({
performanceCrux: false,
}),
heapSnapshotManager: new HeapSnapshotManager(),
heapSnapshotManager: new HeapSnapshotManager(new TaskRunner()),
});
}
+2 -2
View File
@@ -89,9 +89,9 @@ describe('comments tools', () => {
]);
});
it('yields the tool lock while waiting', () => {
it('runs the handler unlocked while waiting', () => {
const {args} = createHandlerMocks();
assert.strictEqual(getDevtoolsComments(args).yieldsToolLock, true);
assert.strictEqual(getDevtoolsComments(args).concurrency, 'none');
});
it('waits for a comment event and fetches again', async () => {
+10 -3
View File
@@ -34,6 +34,7 @@ import {isDaemonRunning} from '../src/daemon/utils.js';
import {McpContext, type McpContextOptions} from '../src/McpContext.js';
import {McpResponse} from '../src/McpResponse.js';
import {HeapSnapshotManager} from '../src/processors/HeapSnapshotManager.js';
import {TaskRunner} from '../src/utils/TaskRunner.js';
import {TextSnapshot} from '../src/TextSnapshot.js';
import {DevTools} from '../src/third_party/index.js';
import {stableIdSymbol} from '../src/utils/id.js';
@@ -177,12 +178,18 @@ export async function withBrowser(
*/
export async function createConnectedMcpContext(
browser: Browser,
options: Omit<McpContextOptions, 'getBrowser'>,
options: Omit<McpContextOptions, 'getBrowser' | 'taskRunner'> & {
taskRunner?: TaskRunner;
},
locatorClass?: typeof Locator,
): Promise<McpContext> {
const context = McpContext.create(
undefined,
{...options, getBrowser: async () => browser},
{
taskRunner: new TaskRunner(),
...options,
getBrowser: async () => browser,
},
locatorClass,
);
await context.ensureBrowser();
@@ -225,7 +232,7 @@ export async function withMcpContext(
}),
});
const response = new McpResponse(parsedArgs);
const heapSnapshotManager = new HeapSnapshotManager();
const heapSnapshotManager = new HeapSnapshotManager(new TaskRunner());
if (context) {
context.dispose();
}
+271
View File
@@ -0,0 +1,271 @@
/**
* @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;
});
});
});
+1 -1
View File
@@ -66,7 +66,7 @@ describe('WaitForHelper', () => {
// The action opens a dialog asynchronously and passes no handleDialog.
// The dialog leaves the renderer paused; without the fix,
// waitForStableDom's setup evaluation would hang until protocolTimeout
// (~180s) while the tool mutex is held, freezing the session.
// (~180s) while the context lock is held, freezing the session.
try {
const result = await Promise.race([
mcpPage.waitForEventsAfterAction(async () => {