Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
342 changes: 190 additions & 152 deletions README.md

Large diffs are not rendered by default.

221 changes: 220 additions & 1 deletion src/McpContext.ts

Large diffs are not rendered by default.

84 changes: 84 additions & 0 deletions src/McpResponse.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,11 @@ import {handleDialog} from './tools/pages.js';
import type {ImageContentData, Response} from './tools/ToolDefinition.js';
import {paginate, type PaginationOptions} from './utils/pagination.js';

/**
* Represents a response from an MCP tool, handling the collection and
* formatting of various data types like text, images, and network requests.
* @public
*/
export class McpResponse implements Response {
#includePages = false;
#includeSnapshot = false;
Expand All @@ -35,14 +40,31 @@ export class McpResponse implements Response {
resourceTypes?: ResourceType[];
};

/**
* Sets whether to include page information in the response.
*
* @param value - True to include page information, false otherwise.
*/
setIncludePages(value: boolean): void {
this.#includePages = value;
}

/**
* Sets whether to include a snapshot in the response.
*
* @param value - True to include a snapshot, false otherwise.
*/
setIncludeSnapshot(value: boolean): void {
this.#includeSnapshot = value;
}

/**
* Sets whether to include network requests in the response, with optional
* pagination and filtering.
*
* @param value - True to include network requests, false otherwise.
* @param options - Options for pagination and resource type filtering.
*/
setIncludeNetworkRequests(
value: boolean,
options?: {
Expand All @@ -69,52 +91,103 @@ export class McpResponse implements Response {
};
}

/**
* Sets whether to include console data in the response.
*
* @param value - True to include console data, false otherwise.
*/
setIncludeConsoleData(value: boolean): void {
this.#includeConsoleData = value;
}

/**
* Attaches a specific network request to the response by its URL.
*
* @param url - The URL of the network request to attach.
*/
attachNetworkRequest(url: string): void {
this.#attachedNetworkRequestUrl = url;
}

/**
* Gets whether page information is included in the response.
*/
get includePages(): boolean {
return this.#includePages;
}

/**
* Gets whether network requests are included in the response.
*/
get includeNetworkRequests(): boolean {
return this.#networkRequestsOptions?.include ?? false;
}

/**
* Gets whether console data is included in the response.
*/
get includeConsoleData(): boolean {
return this.#includeConsoleData;
}
/**
* Gets the URL of the attached network request.
*/
get attachedNetworkRequestUrl(): string | undefined {
return this.#attachedNetworkRequestUrl;
}
/**
* Gets the page index for network request pagination.
*/
get networkRequestsPageIdx(): number | undefined {
return this.#networkRequestsOptions?.pagination?.pageIdx;
}

/**
* Appends a line of text to the response.
*
* @param value - The line of text to append.
*/
appendResponseLine(value: string): void {
this.#textResponseLines.push(value);
}

/**
* Attaches an image to the response.
*
* @param value - The image data to attach.
*/
attachImage(value: ImageContentData): void {
this.#images.push(value);
}

/**
* Gets the lines of text in the response.
*/
get responseLines(): readonly string[] {
return this.#textResponseLines;
}

/**
* Gets the images attached to the response.
*/
get images(): ImageContentData[] {
return this.#images;
}

/**
* Gets whether a snapshot is included in the response.
*/
get includeSnapshot(): boolean {
return this.#includeSnapshot;
}

/**
* Handles the response by creating snapshots and formatting the data.
*
* @param toolName - The name of the tool that generated the response.
* @param context - The MCP context.
* @returns A promise that resolves to an array of text and image content.
*/
async handle(
toolName: string,
context: McpContext,
Expand All @@ -140,6 +213,13 @@ export class McpResponse implements Response {
return this.format(toolName, context);
}

/**
* Formats the response into an array of text and image content.
*
* @param toolName - The name of the tool that generated the response.
* @param context - The MCP context.
* @returns An array of text and image content.
*/
format(
toolName: string,
context: McpContext,
Expand Down Expand Up @@ -314,6 +394,10 @@ Call ${handleDialog.name} to handle it before continuing.`);
return response;
}

/**
* Resets the response lines for testing purposes.
* @internal
*/
resetResponseLineForTesting() {
this.#textResponseLines = [];
}
Expand Down
22 changes: 21 additions & 1 deletion src/Mutex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,23 @@
* SPDX-License-Identifier: Apache-2.0
*/

/**
* A simple asynchronous mutex implementation.
* @public
*/
export class Mutex {
/**
* A guard that releases the mutex when disposed.
* @public
*/
static Guard = class Guard {
#mutex: Mutex;
constructor(mutex: Mutex) {
this.#mutex = mutex;
}
/**
* Releases the mutex.
*/
dispose(): void {
return this.#mutex.release();
}
Expand All @@ -18,7 +29,12 @@ export class Mutex {
#locked = false;
#acquirers: Array<() => void> = [];

// This is FIFO.
/**
* Acquires the mutex, waiting if necessary. This is a FIFO queue.
*
* @returns A promise that resolves with a guard, which will release the
* mutex when disposed.
*/
async acquire(): Promise<InstanceType<typeof Mutex.Guard>> {
if (!this.#locked) {
this.#locked = true;
Expand All @@ -30,6 +46,10 @@ export class Mutex {
return new Mutex.Guard(this);
}

/**
* Releases the mutex.
* @internal
*/
release(): void {
const resolve = this.#acquirers.shift();
if (!resolve) {
Expand Down
46 changes: 46 additions & 0 deletions src/PageCollector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,29 @@

import type {Browser, HTTPRequest, Page} from 'puppeteer-core';

/**
* A generic class for collecting data from Puppeteer pages. It handles page
* creation and navigation to manage data collection lifecycle.
*
* @template T The type of data to collect.
* @public
*/
export class PageCollector<T> {
#browser: Browser;
#initializer: (page: Page, collector: (item: T) => void) => void;
/**
* The Array in this map should only be set once
* As we use the reference to it.
* Use methods that manipulate the array in place.
* @protected
*/
protected storage = new WeakMap<Page, T[]>();

/**
* @param browser - The Puppeteer browser instance.
* @param initializer - A function that sets up the data collection for a
* page.
*/
constructor(
browser: Browser,
initializer: (page: Page, collector: (item: T) => void) => void,
Expand All @@ -24,6 +37,10 @@ export class PageCollector<T> {
this.#initializer = initializer;
}

/**
* Initializes the collector by setting up data collection for all existing
* pages and listening for new pages.
*/
async init() {
const pages = await this.#browser.pages();
for (const page of pages) {
Expand All @@ -39,6 +56,11 @@ export class PageCollector<T> {
});
}

/**
* Adds a new page to the collector and initializes it.
*
* @param page - The page to add.
*/
public addPage(page: Page) {
this.#initializePage(page);
}
Expand All @@ -63,6 +85,13 @@ export class PageCollector<T> {
});
}

/**
* Cleans up the stored data for a page. By default, it clears the entire
* collection.
*
* @param page - The page to clean up.
* @protected
*/
protected cleanup(page: Page) {
const collection = this.storage.get(page);
if (collection) {
Expand All @@ -71,12 +100,29 @@ export class PageCollector<T> {
}
}

/**
* Gets the collected data for a specific page.
*
* @param page - The page to get data for.
* @returns The collected data, or an empty array if none.
*/
getData(page: Page): T[] {
return this.storage.get(page) ?? [];
}
}

/**
* A specific implementation of PageCollector for collecting network requests.
* @public
*/
export class NetworkCollector extends PageCollector<HTTPRequest> {
/**
* Cleans up network requests by removing all requests before the last
* navigation.
*
* @param page - The page to clean up.
* @override
*/
override cleanup(page: Page) {
const requests = this.storage.get(page) ?? [];
if (!requests) {
Expand Down
38 changes: 35 additions & 3 deletions src/WaitForHelper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,11 @@ import type {CdpPage} from 'puppeteer-core/internal/cdp/Page.js';

import {logger} from './logger.js';

/**
* A helper class for waiting for various page events, such as stable DOM and
* navigation, after performing an action.
* @public
*/
export class WaitForHelper {
#abortController = new AbortController();
#page: CdpPage;
Expand All @@ -16,6 +21,12 @@ export class WaitForHelper {
#expectNavigationIn: number;
#navigationTimeout: number;

/**
* @param page - The Puppeteer page to wait for events on.
* @param cpuTimeoutMultiplier - The multiplier for CPU-bound timeouts.
* @param networkTimeoutMultiplier - The multiplier for network-bound
* timeouts.
*/
constructor(
page: Page,
cpuTimeoutMultiplier: number,
Expand All @@ -29,9 +40,10 @@ export class WaitForHelper {
}

/**
* A wrapper that executes a action and waits for
* a potential navigation, after which it waits
* for the DOM to be stable before returning.
* Waits for the DOM to be stable (i.e., no mutations for a certain period).
*
* @returns A promise that resolves when the DOM is stable.
* @throws If the timeout is reached before the DOM becomes stable.
*/
async waitForStableDom(): Promise<void> {
const stableDomObserver = await this.#page.evaluateHandle(timeout => {
Expand Down Expand Up @@ -82,6 +94,12 @@ export class WaitForHelper {
]);
}

/**
* Waits for a navigation to start.
*
* @returns A promise that resolves to true if a navigation starts, and false
* otherwise.
*/
async waitForNavigationStarted() {
// Currently Puppeteer does not have API
// For when a navigation is about to start
Expand Down Expand Up @@ -114,6 +132,12 @@ export class WaitForHelper {
]);
}

/**
* Creates a timeout promise that can be aborted.
*
* @param time - The timeout in milliseconds.
* @returns A promise that resolves after the timeout.
*/
timeout(time: number): Promise<void> {
return new Promise<void>(res => {
const id = setTimeout(res, time);
Expand All @@ -124,6 +148,14 @@ export class WaitForHelper {
});
}

/**
* Executes an action and then waits for events to settle, such as navigation
* and stable DOM.
*
* @param action - The action to perform.
* @returns A promise that resolves when all events have settled.
* @throws If the action throws an error.
*/
async waitForEventsAfterAction(
action: () => Promise<unknown>,
): Promise<void> {
Expand Down
Loading