Skip to content

refactor: clean up more of the context getters #1062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
OrKoN merged 1 commit into from
Feb 27, 2026
Merged

refactor: clean up more of the context getters #1062

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
141 changes: 25 additions & 116 deletions src/McpContext.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ import type {
BrowserContext,
ConsoleMessage,
Debugger,
ElementHandle,
HTTPRequest,
Page,
ScreenRecorder,
Expand All @@ -34,7 +33,6 @@ import type {
import {Locator} from './third_party/index.js';
import {PredefinedNetworkConditions} from './third_party/index.js';
import {listPages} from './tools/pages.js';
import {takeSnapshot} from './tools/snapshot.js';
import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
import type {
Context,
Expand Down Expand Up @@ -109,7 +107,7 @@ export class McpContext implements Context {
#extensionServiceWorkers: ExtensionServiceWorker[] = [];

#mcpPages = new Map<Page, McpPage>();
#selectedPage?: Page;
#selectedPage?: McpPage;
#networkCollector: NetworkCollector;
#consoleCollector: ConsoleCollector;
#devtoolsUniverseManager: UniverseManager;
Expand Down Expand Up @@ -290,7 +288,7 @@ export class McpContext implements Context {
page = await this.browser.newPage({background});
}
await this.createPagesSnapshot();
this.selectPage(page);
this.selectPage(this.#getMcpPage(page));
this.#networkCollector.addPage(page);
this.#consoleCollector.addPage(page);
return this.#getMcpPage(page);
Expand All @@ -300,16 +298,15 @@ export class McpContext implements Context {
throw new Error(CLOSE_PAGE_ERROR);
}
const page = this.getPageById(pageId);
const mcpPage = this.#mcpPages.get(page);
if (mcpPage) {
mcpPage.dispose();
this.#mcpPages.delete(page);
if (page) {
page.dispose();
this.#mcpPages.delete(page.pptrPage);
}
const ctx = page.browserContext();
if (this.#focusedPagePerContext.get(ctx) === page) {
const ctx = page.pptrPage.browserContext();
if (this.#focusedPagePerContext.get(ctx) === page.pptrPage) {
this.#focusedPagePerContext.delete(ctx);
}
await page.close({runBeforeUnload: false});
await page.pptrPage.close({runBeforeUnload: false});
}

getNetworkRequestById(page: McpPage, reqid: number): HTTPRequest {
Expand Down Expand Up @@ -461,29 +458,21 @@ export class McpContext implements Context {
if (!page) {
throw new Error('No page selected');
}
if (page.isClosed()) {
if (page.pptrPage.isClosed()) {
throw new Error(
`The selected page has been closed. Call ${listPages().name} to see open pages.`,
);
}
return page;
return page.pptrPage;
}

getSelectedMcpPage(): McpPage {
const page = this.getSelectedPptrPage();
return this.#getMcpPage(page);
}

resolvePageById(pageId?: number): McpPage {
if (pageId === undefined) {
return this.getSelectedMcpPage();
}
const page = this.getPageById(pageId);
return this.#getMcpPage(page);
}

getPageById(pageId: number): Page {
const page = this.#pages.find(p => this.#mcpPages.get(p)?.id === pageId);
getPageById(pageId: number): McpPage {
const page = this.#mcpPages.values().find(mcpPage => mcpPage.id === pageId);
if (!page) {
throw new Error('No page found');
}
Expand All @@ -507,7 +496,7 @@ export class McpContext implements Context {
}

isPageSelected(page: Page): boolean {
return this.#selectedPage === page;
return this.#selectedPage?.pptrPage === page;
}

assertPageIsFocused(pageToCheck: Page | ContextPage): void {
Expand All @@ -524,20 +513,22 @@ export class McpContext implements Context {
}
}

selectPage(pageToSelect: Page | ContextPage): void {
const newPage =
'pptrPage' in pageToSelect ? pageToSelect.pptrPage : pageToSelect;
const ctx = newPage.browserContext();
selectPage(newPage: McpPage): void {
const ctx = newPage.pptrPage.browserContext();
const oldFocused = this.#focusedPagePerContext.get(ctx);
if (oldFocused && oldFocused !== newPage && !oldFocused.isClosed()) {
if (
oldFocused &&
oldFocused !== newPage.pptrPage &&
!oldFocused.isClosed()
) {
void oldFocused.emulateFocusedPage(false).catch(error => {
this.logger('Error turning off focused page emulation', error);
});
}
this.#focusedPagePerContext.set(ctx, newPage);
this.#focusedPagePerContext.set(ctx, newPage.pptrPage);
this.#selectedPage = newPage;
this.#updateSelectedPageTimeouts();
void newPage.emulateFocusedPage(true).catch(error => {
void newPage.pptrPage.emulateFocusedPage(true).catch(error => {
this.logger('Error turning on focused page emulation', error);
});
}
Expand Down Expand Up @@ -572,63 +563,6 @@ export class McpContext implements Context {
return undefined;
}

async getElementByUid(
uid: string,
page?: Page,
): Promise<ElementHandle<Element>> {
if (page) {
// Scoped search: only look in the target page's snapshot.
const mcpPage = this.#mcpPages.get(page);
if (!mcpPage?.textSnapshot) {
throw new Error(
`No snapshot found for page ${mcpPage?.id ?? '?'}. Use ${takeSnapshot.name} to capture one.`,
);
}
const node = mcpPage.textSnapshot.idToNode.get(uid);
if (!node) {
throw new Error(
`Element uid "${uid}" not found on page ${mcpPage.id}.`,
);
}
return this.#resolveElementHandle(node, uid);
}

// Cross-page search with context-focus validation.
let anySnapshot = false;
for (const [searchPage, mcpPage] of this.#mcpPages.entries()) {
if (!mcpPage.textSnapshot) {
continue;
}
anySnapshot = true;
const node = mcpPage.textSnapshot.idToNode.get(uid);
if (node) {
const ctx = searchPage.browserContext();
const contextSelectedPage = this.#focusedPagePerContext.get(ctx);
if (contextSelectedPage !== searchPage) {
const targetId = mcpPage.id;
const selectedId = contextSelectedPage
? this.#mcpPages.get(contextSelectedPage)?.id
: this.#getSelectedMcpPage().id;
throw new Error(
`Element uid "${uid}" belongs to page ${targetId}, but page ${selectedId} is currently selected. ` +
`Call select_page with pageId ${targetId} first.`,
);
}
// Align global #selectedPage for waitForEventsAfterAction etc.
if (this.#selectedPage !== searchPage) {
this.#selectedPage = searchPage;
}
return this.#resolveElementHandle(node, uid);
}
}
if (!anySnapshot) {
throw new Error(
`No snapshot found. Use ${takeSnapshot.name} to capture one.`,
);
}
throw new Error('No such element found in any snapshot.');
}

/**
* Creates a snapshot of the extension service workers.
*/
Expand Down Expand Up @@ -664,24 +598,6 @@ export class McpContext implements Context {
return this.#extensionServiceWorkers;
}

async #resolveElementHandle(
node: TextSnapshotNode,
uid: string,
): Promise<ElementHandle<Element>> {
const message = `Element with uid ${uid} no longer exists on the page.`;
try {
const handle = await node.elementHandle();
if (!handle) {
throw new Error(message);
}
return handle;
} catch (error) {
throw new Error(message, {
cause: error,
});
}
}

async createPagesSnapshot(): Promise<Page[]> {
const {pages: allPages, isolatedContextNames} = await this.#getAllPages();

Expand Down Expand Up @@ -717,10 +633,11 @@ export class McpContext implements Context {
});

if (
(!this.#selectedPage || this.#pages.indexOf(this.#selectedPage) === -1) &&
(!this.#selectedPage ||
this.#pages.indexOf(this.#selectedPage.pptrPage) === -1) &&
this.#pages[0]
) {
this.selectPage(this.#pages[0]);
this.selectPage(this.#getMcpPage(this.#pages[0]));
}

await this.detectOpenDevToolsWindows();
Expand Down Expand Up @@ -943,14 +860,6 @@ export class McpContext implements Context {
}
}

getTextSnapshot(targetPage?: Page): TextSnapshot | null {
const page = targetPage ?? this.#selectedPage;
if (!page) {
return null;
}
return this.#mcpPages.get(page)?.textSnapshot ?? null;
}

async saveTemporaryFile(
data: Uint8Array<ArrayBufferLike>,
filename: string,
Expand Down
4 changes: 1 addition & 3 deletions src/McpResponse.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,9 +274,7 @@ export class McpResponse implements Response {
this.#snapshotParams.verbose,
this.#devToolsData,
);
const textSnapshot = context.getTextSnapshot(
this.#snapshotParams.page?.pptrPage,
);
const textSnapshot = this.#page.textSnapshot;
if (textSnapshot) {
const formatter = new SnapshotFormatter(textSnapshot);
if (this.#snapshotParams.filePath) {
Expand Down
2 changes: 1 addition & 1 deletion src/server.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ export async function createMcpServer(
serverArgs.experimentalPageIdRouting &&
params.pageId &&
!serverArgs.slim
? context.resolvePageById(params.pageId)
? context.getPageById(params.pageId)
: context.getSelectedMcpPage();
response.setPage(page);
await tool.handler(
Expand Down
7 changes: 2 additions & 5 deletions src/tools/ToolDefinition.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ export interface ImageContentData {
export interface SnapshotParams {
verbose?: boolean;
filePath?: string;
page?: ContextPage;
}

export interface LighthouseData {
Expand Down Expand Up @@ -136,16 +135,14 @@ export type Context = Readonly<{
isCruxEnabled(): boolean;
recordedTraces(): TraceResult[];
storeTraceRecording(result: TraceResult): void;
getPageById(pageId: number): Page;
getPageById(pageId: number): ContextPage;
newPage(
background?: boolean,
isolatedContextName?: string,
): Promise<ContextPage>;
closePage(pageId: number): Promise<void>;
selectPage(page: Page): void;
selectPage(page: ContextPage): void;
assertPageIsFocused(page: Page): void;
getElementByUid(uid: string, page?: Page): Promise<ElementHandle<Element>>;
getAXNodeByUid(uid: string): TextSnapshotNode | undefined;
restoreEmulation(page: ContextPage): Promise<void>;
emulate(
options: {
Expand Down
Loading