add test for stagehand. Target.attachToTarget handling and better DCP types

This commit is contained in:
Tommy D. Rossi
2025-12-08 15:13:51 +01:00
parent e5630d76a3
commit d20cea7c49
8 changed files with 2096 additions and 78 deletions
+1
View File
@@ -26,6 +26,7 @@
"author": "Tommaso De Rossi, morse <beats.by.morse@gmail.com>",
"license": "",
"devDependencies": {
"@browserbasehq/stagehand": "^3.0.5",
"@mizchi/selector-generator": "1.50.0-next",
"@types/chrome": "^0.0.315",
"@types/node": "^24.10.1",
+28 -3
View File
@@ -1,26 +1,51 @@
import type { Protocol } from 'devtools-protocol';
import type { ProtocolMapping } from 'devtools-protocol/types/protocol-mapping.js';
export type CDPCommand<T extends keyof ProtocolMapping.Commands = keyof ProtocolMapping.Commands> = {
export type CDPCommandFor<T extends keyof ProtocolMapping.Commands> = {
id: number;
sessionId?: string;
method: T;
params?: ProtocolMapping.Commands[T]['paramsType'][0];
};
export type CDPResponse<T extends keyof ProtocolMapping.Commands = keyof ProtocolMapping.Commands> = {
export type CDPCommand = {
[K in keyof ProtocolMapping.Commands]: CDPCommandFor<K>;
}[keyof ProtocolMapping.Commands];
export type CDPResponseFor<T extends keyof ProtocolMapping.Commands> = {
id: number;
sessionId?: string;
result?: ProtocolMapping.Commands[T]['returnType'];
error?: { code?: number; message: string };
};
export type CDPEvent<T extends keyof ProtocolMapping.Events = keyof ProtocolMapping.Events> = {
export type CDPResponse = {
[K in keyof ProtocolMapping.Commands]: CDPResponseFor<K>;
}[keyof ProtocolMapping.Commands];
export type CDPEventFor<T extends keyof ProtocolMapping.Events> = {
method: T;
sessionId?: string;
params?: ProtocolMapping.Events[T][0];
};
export type CDPEvent = {
[K in keyof ProtocolMapping.Events]: CDPEventFor<K>;
}[keyof ProtocolMapping.Events];
export type CDPResponseBase = {
id: number;
sessionId?: string;
result?: unknown;
error?: { code?: number; message: string };
};
export type CDPEventBase = {
method: string;
sessionId?: string;
params?: unknown;
};
export type CDPMessage = CDPCommand | CDPResponse | CDPEvent;
export { Protocol, ProtocolMapping };
+72 -32
View File
@@ -3,7 +3,7 @@ import { serve } from '@hono/node-server'
import { createNodeWebSocket } from '@hono/node-ws'
import type { WSContext } from 'hono/ws'
import type { Protocol } from '../cdp-types.js'
import type { CDPCommand, CDPResponse, CDPEvent } from '../cdp-types.js'
import type { CDPCommand, CDPResponseBase, CDPEventBase, CDPEventFor } from '../cdp-types.js'
import type { ExtensionMessage, ExtensionEventMessage } from './protocol.js'
import chalk from 'chalk'
@@ -20,7 +20,8 @@ type PlaywrightClient = {
ws: WSContext
}
export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.0.0.1', logger = console }: { port?: number; host?: string; logger?: { log(...args: any[]): void; error(...args: any[]): void } } = {}) {
export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.0.0.1', logger }: { port?: number; host?: string; logger?: { log(...args: any[]): void; error(...args: any[]): void } } = {}) {
const connectedTargets = new Map<string, ConnectedTarget>()
const playwrightClients = new Map<string, PlaywrightClient>()
@@ -88,14 +89,14 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
if (direction === 'from-playwright') {
const clientLabel = clientId ? chalk.blue(`[${clientId}]`) : ''
logger.log(chalk.cyan('← Playwright'), clientLabel + ':', method + detailsStr)
logger?.log(chalk.cyan('← Playwright'), clientLabel + ':', method + detailsStr)
} else if (direction === 'from-extension') {
logger.log(chalk.yellow('← Extension:'), method + detailsStr)
logger?.log(chalk.yellow('← Extension:'), method + detailsStr)
} else if (direction === 'to-playwright') {
const color = source === 'server' ? chalk.magenta : chalk.green
const sourceLabel = source === 'server' ? chalk.gray(' (server-generated)') : ''
const clientLabel = clientId ? chalk.blue(`[${clientId}]`) : chalk.blue('[ALL]')
logger.log(color('→ Playwright'), clientLabel + ':', method + detailsStr + sourceLabel)
logger?.log(color('→ Playwright'), clientLabel + ':', method + detailsStr + sourceLabel)
}
}
@@ -104,7 +105,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
clientId,
source = 'extension'
}: {
message: CDPResponse | CDPEvent
message: CDPResponseBase | CDPEventBase
clientId?: string
source?: 'extension' | 'server'
}) {
@@ -189,6 +190,25 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
return {}
}
case 'Target.setDiscoverTargets': {
return {}
}
case 'Target.attachToTarget': {
const targetId = params?.targetId
if (!targetId) {
throw new Error('targetId is required for Target.attachToTarget')
}
for (const target of connectedTargets.values()) {
if (target.targetId === targetId) {
return { sessionId: target.sessionId } satisfies Protocol.Target.AttachToTargetResponse
}
}
throw new Error(`Target ${targetId} not found in connected targets`)
}
case 'Target.getTargetInfo': {
const targetId = params?.targetId
@@ -251,7 +271,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
app.post('/mcp-log', async (c) => {
try {
const { level, args } = await c.req.json()
const logFn = (logger as any)[level] || logger.log
const logFn = (logger as any)[level] || logger?.log
const prefix = chalk.red(`[MCP] [${level.toUpperCase()}]`)
logFn(prefix, ...args)
return c.json({ ok: true })
@@ -266,13 +286,13 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
return {
onOpen(_event, ws) {
if (playwrightClients.has(clientId)) {
logger.log(chalk.red(`Rejecting duplicate client ID: ${clientId}`))
logger?.log(chalk.red(`Rejecting duplicate client ID: ${clientId}`))
ws.close(1000, 'Client ID already connected')
return
}
playwrightClients.set(clientId, { id: clientId, ws })
logger.log(chalk.green(`Playwright client connected: ${clientId} (${playwrightClients.size} total)`))
logger?.log(chalk.green(`Playwright client connected: ${clientId} (${playwrightClients.size} total)`))
},
async onMessage(event, ws) {
@@ -321,8 +341,8 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
},
waitingForDebugger: false
}
} satisfies CDPEvent
logger.log(chalk.magenta('[Server] Target.attachedToTarget full payload:'), JSON.stringify(attachedPayload))
} satisfies CDPEventFor<'Target.attachedToTarget'>
logger?.log(chalk.magenta('[Server] Target.attachedToTarget full payload:'), JSON.stringify(attachedPayload))
sendToPlaywright({
message: attachedPayload,
clientId,
@@ -331,12 +351,32 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
}
}
if (method === 'Target.setDiscoverTargets' && (params)?.discover) {
for (const target of connectedTargets.values()) {
const targetCreatedPayload = {
method: 'Target.targetCreated',
params: {
targetInfo: {
...target.targetInfo,
attached: true
}
}
} satisfies CDPEventFor<'Target.targetCreated'>
logger?.log(chalk.magenta('[Server] Target.targetCreated full payload:'), JSON.stringify(targetCreatedPayload))
sendToPlaywright({
message: targetCreatedPayload,
clientId,
source: 'server'
})
}
}
sendToPlaywright({
message: { id, sessionId, result },
clientId
})
} catch (e) {
logger.error('Error handling CDP command:', method, params, e)
logger?.error('Error handling CDP command:', method, params, e)
sendToPlaywright({
message: {
id,
@@ -350,11 +390,11 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
onClose() {
playwrightClients.delete(clientId)
logger.log(chalk.yellow(`Playwright client disconnected: ${clientId} (${playwrightClients.size} remaining)`))
logger?.log(chalk.yellow(`Playwright client disconnected: ${clientId} (${playwrightClients.size} remaining)`))
},
onError(event) {
logger.error(`Playwright WebSocket error [${clientId}]:`, event)
logger?.error(`Playwright WebSocket error [${clientId}]:`, event)
}
}
}))
@@ -363,9 +403,9 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
return {
onOpen(_event, ws) {
if (extensionWs) {
logger.log(chalk.yellow('Closing existing extension connection to replace with new one'))
logger?.log(chalk.yellow('Closing existing extension connection to replace with new one'))
extensionWs.close(4001, 'Extension Replaced')
// Clear state from the old connection to prevent leaks
connectedTargets.clear()
for (const pending of extensionPendingRequests.values()) {
@@ -380,7 +420,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
}
extensionWs = ws
logger.log('Extension connected with clean state')
logger?.log('Extension connected with clean state')
},
async onMessage(event, ws) {
@@ -396,7 +436,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
if ('id' in message) {
const pending = extensionPendingRequests.get(message.id)
if (!pending) {
logger.log('Unexpected response with id:', message.id)
logger?.log('Unexpected response with id:', message.id)
return
}
@@ -409,7 +449,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
}
} else if (message.method === 'log') {
const { level, args } = message.params
const logFn = (logger as any)[level] || logger.log
const logFn = (logger as any)[level] || logger?.log
const prefix = chalk.yellow(`[Extension] [${level.toUpperCase()}]`)
logFn(prefix, ...args)
} else {
@@ -431,7 +471,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
if (method === 'Target.attachedToTarget') {
const targetParams = params as Protocol.Target.AttachedToTargetEvent
logger.log(chalk.yellow('[Extension] Target.attachedToTarget full payload:'), JSON.stringify({ method, params: targetParams, sessionId }))
logger?.log(chalk.yellow('[Extension] Target.attachedToTarget full payload:'), JSON.stringify({ method, params: targetParams, sessionId }))
// Check if we already sent this target to clients (e.g., from Target.setAutoAttach response)
const alreadyConnected = connectedTargets.has(targetParams.sessionId)
@@ -449,7 +489,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
message: {
method: 'Target.attachedToTarget',
params: targetParams
} as CDPEvent,
} as CDPEventBase,
source: 'extension'
})
}
@@ -461,7 +501,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
message: {
method: 'Target.detachedFromTarget',
params: detachParams
} as CDPEvent,
} as CDPEventBase,
source: 'extension'
})
} else if (method === 'Target.targetInfoChanged') {
@@ -477,7 +517,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
message: {
method: 'Target.targetInfoChanged',
params: infoParams
} as CDPEvent,
} as CDPEventBase,
source: 'extension'
})
} else {
@@ -486,7 +526,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
sessionId,
method,
params
} as CDPEvent,
} as CDPEventBase,
source: 'extension'
})
}
@@ -494,12 +534,12 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
},
onClose(event, ws) {
logger.log('Extension disconnected')
logger?.log('Extension disconnected')
// If this is an old connection closing after we've already established a new one,
// don't clear the global state
if (extensionWs && extensionWs !== ws) {
logger.log('Old extension connection closed, keeping new one active')
logger?.log('Old extension connection closed, keeping new one active')
return
}
@@ -518,7 +558,7 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
},
onError(event) {
logger.error('Extension WebSocket error:', event)
logger?.error('Extension WebSocket error:', event)
}
}
}))
@@ -530,11 +570,11 @@ export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.
const cdpEndpoint = `${wsHost}/cdp`
const extensionEndpoint = `${wsHost}/extension`
logger.log('CDP relay server started')
logger.log('Host:', host)
logger.log('Port:', port)
logger.log('Extension endpoint:', extensionEndpoint)
logger.log('CDP endpoint:', cdpEndpoint)
logger?.log('CDP relay server started')
logger?.log('Host:', host)
logger?.log('Port:', port)
logger?.log('Extension endpoint:', extensionEndpoint)
logger?.log('CDP endpoint:', cdpEndpoint)
return {
close() {
+3 -3
View File
@@ -1,4 +1,4 @@
import { CDPCommand, CDPEvent, ProtocolMapping } from '../cdp-types.js'
import { CDPEventFor, ProtocolMapping } from '../cdp-types.js'
export const VERSION = 1
@@ -32,9 +32,9 @@ export type ExtensionEventMessage =
[K in keyof ProtocolMapping.Events]: {
method: 'forwardCDPEvent'
params: {
method: CDPEvent<K>['method']
method: CDPEventFor<K>['method']
sessionId?: string
params?: CDPEvent<K>['params']
params?: CDPEventFor<K>['params']
}
}
}[keyof ProtocolMapping.Events]
+75 -21
View File
@@ -610,8 +610,8 @@ describe('MCP Server Tests', () => {
expect(results).toMatchInlineSnapshot(`
[
{
"title": "",
"url": "about:blank",
"title": "🎄 Twelve Days of Shell 🎄",
"url": "https://12days.cmdchallenge.com/",
},
{
"title": "Example Domain",
@@ -696,8 +696,8 @@ describe('MCP Server Tests', () => {
expect(results).toMatchInlineSnapshot(`
[
{
"title": "",
"url": "about:blank",
"title": "🎄 Twelve Days of Shell 🎄",
"url": "https://12days.cmdchallenge.com/",
},
{
"title": "Example Domain",
@@ -762,18 +762,18 @@ describe('MCP Server Tests', () => {
const pages = await browserContext.pages()
expect(pages.length).toBeGreaterThan(0)
const page = pages[0]
await page.goto('https://example.com/disconnect-test')
await page.waitForLoadState('networkidle')
await page.bringToFront()
// Enable extension on this page
const initialEnable = await serviceWorker.evaluate(async () => {
return await globalThis.toggleExtensionForActiveTab()
})
console.log('Initial enable result:', initialEnable)
expect(initialEnable.isConnected).toBe(true)
// Wait for extension to fully connect
await new Promise(resolve => setTimeout(resolve, 500))
@@ -790,7 +790,7 @@ describe('MCP Server Tests', () => {
`,
},
})
const beforeOutput = (beforeDisconnect as any).content[0].text
expect(beforeOutput).toContain('foundTestPage')
console.log('Before disconnect:', beforeOutput)
@@ -800,7 +800,7 @@ describe('MCP Server Tests', () => {
await serviceWorker.evaluate(async () => {
await globalThis.disconnectEverything()
})
// Wait for disconnect to complete
await new Promise(resolve => setTimeout(resolve, 500))
@@ -815,7 +815,7 @@ describe('MCP Server Tests', () => {
`,
},
})
const afterDisconnectOutput = (afterDisconnect as any).content[0].text
console.log('After disconnect:', afterDisconnectOutput)
expect(afterDisconnectOutput).toContain('Pages after disconnect: 0')
@@ -829,10 +829,10 @@ describe('MCP Server Tests', () => {
console.log('toggleExtensionForActiveTab result:', result)
return result
})
console.log('Reconnect result:', reconnectResult)
expect(reconnectResult.isConnected).toBe(true)
// Wait for extension to fully reconnect and relay server to be ready
console.log('Waiting for reconnection to stabilize...')
await new Promise(resolve => setTimeout(resolve, 1000))
@@ -863,25 +863,25 @@ describe('MCP Server Tests', () => {
console.log('Checking pages after reconnect...');
const pages = context.pages();
console.log('Pages after reconnect:', pages.length);
if (pages.length === 0) {
console.log('No pages found!');
return { pagesCount: 0, foundTestPage: false };
}
const testPage = pages.find(p => p.url().includes('disconnect-test'));
console.log('Found test page after reconnect:', !!testPage);
if (testPage) {
console.log('Test page URL:', testPage.url());
return { pagesCount: pages.length, foundTestPage: true, url: testPage.url() };
}
return { pagesCount: pages.length, foundTestPage: false };
`,
},
})
const afterReconnectOutput = (afterReconnect as any).content[0].text
console.log('After reconnect:', afterReconnectOutput)
expect(afterReconnectOutput).toContain('foundTestPage')
@@ -900,11 +900,11 @@ describe('MCP Server Tests', () => {
// Clear any existing logs from previous tests
clearAllLogs();
console.log('Cleared all existing logs');
// Verify connection is working
const pages = context.pages();
console.log('Current pages count:', pages.length);
return { success: true, pagesCount: pages.length };
`,
},
@@ -996,7 +996,7 @@ describe('MCP Server Tests', () => {
`,
},
})
const beforeReloadOutput = (beforeReloadResult as any).content[0].text
expect(beforeReloadOutput).toContain('[log] Before reload 99999')
@@ -1025,7 +1025,7 @@ describe('MCP Server Tests', () => {
`,
},
})
const afterReloadOutput = (afterReloadResult as any).content[0].text
expect(afterReloadOutput).toContain('[log] After reload 88888')
expect(afterReloadOutput).not.toContain('[log] Before reload 99999')
@@ -1357,6 +1357,60 @@ describe('MCP Server Tests', () => {
await page.close()
}, 60000)
it('should work with stagehand', async () => {
if (!browserContext) throw new Error('Browser not initialized')
const serviceWorker = await getExtensionServiceWorker(browserContext)
await serviceWorker.evaluate(async () => {
await globalThis.disconnectEverything()
})
await new Promise(r => setTimeout(r, 500))
const targetUrl = 'https://example.com/'
const enableResult = await serviceWorker.evaluate(async (url) => {
const tab = await chrome.tabs.create({ url, active: true })
await new Promise(r => setTimeout(r, 1000))
return await globalThis.toggleExtensionForActiveTab()
}, targetUrl)
console.log('Extension enabled:', enableResult)
expect(enableResult.isConnected).toBe(true)
await new Promise(r => setTimeout(r, 1000))
const { Stagehand } = await import('@browserbasehq/stagehand')
const stagehand = new Stagehand({
env: 'LOCAL',
verbose: 1,
disablePino: true,
localBrowserLaunchOptions: {
cdpUrl: getCdpUrl(),
},
})
console.log('Initializing Stagehand...')
await stagehand.init()
console.log('Stagehand initialized')
const context = stagehand.context
console.log('Stagehand context:', context)
expect(context).toBeDefined()
const pages = context.pages()
console.log('Stagehand pages:', pages.length, pages.map(p => p.url()))
const stagehandPage = pages.find(p => p.url().includes('example.com'))
expect(stagehandPage).toBeDefined()
const url = stagehandPage!.url()
console.log('Stagehand page URL:', url)
expect(url).toContain('example.com')
await stagehand.close()
}, 60000)
})
@@ -28,10 +28,10 @@ Return value:
- generic [ref=e21]:
- generic: ⌘
- generic: K
- link "100.8k" [ref=e22] [cursor=pointer]:
- link "100.5k" [ref=e22] [cursor=pointer]:
- /url: https://github.com/shadcn-ui/ui
- img
- generic [ref=e23]: 100.8k
- generic [ref=e23]: 100.5k
- button "Toggle theme" [ref=e24]:
- img
- generic [ref=e25]: Toggle theme
+1913 -17
View File
File diff suppressed because it is too large Load Diff
+2
View File
@@ -1,2 +1,4 @@
packages:
- ./*
preferWorkspacePackages: true