diff --git a/packages/codemode/README.md b/packages/codemode/README.md new file mode 100644 index 0000000000..e3e19f7bbf --- /dev/null +++ b/packages/codemode/README.md @@ -0,0 +1,83 @@ +# Stagehand V4 code-mode MCP spike + +This private workspace package is a trusted-code prototype for one MCP tool: +`code_execute`. Generated JavaScript runs against a Stagehand V4 browser that is +always hosted on Browserbase. + +## Lifecycle + +- Starting the MCP server and listing tools do not create a browser. +- The first `action: "run"` lazily creates a Stagehand instance and Browserbase + session. +- The response returns an opaque `code_session_id`. +- Later calls that pass the same ID reuse the browser even if the MCP transport + disconnects and reconnects. +- `action: "close"` closes one logical code session. Server shutdown closes all + remaining sessions. + +`page`, `context`, `stagehand`, `z`, and `console` are available inside each +code cell. + +## Run locally + +Build Stagehand first so its extension assets exist, then start the MCP: + +```bash +pnpm build +BROWSERBASE_API_KEY="" pnpm --filter @browserbasehq/stagehand-codemode build +BROWSERBASE_API_KEY="" node packages/codemode/dist/cli.mjs +``` + +The Streamable HTTP endpoint defaults to +`http://localhost:8932/mcp`. Use `--stdio` for a local stdio transport. +Binding to a non-loopback host is rejected unless +`CODEMODE_MCP_BEARER_TOKEN` is set. Unauthenticated loopback HTTP accepts only +loopback `Host` and `Origin` values to block browser-based DNS rebinding. +Non-loopback HTTP is plaintext and must sit behind trusted TLS termination. + +Optional environment variables: + +- `CODEMODE_MCP_BEARER_TOKEN` +- `CODEMODE_DEFAULT_TIMEOUT_MS` +- `STAGEHAND_MODEL_NAME` +- `STAGEHAND_MODEL_API_KEY` +- `STAGEHAND_MODEL_BASE_URL` + +## Tool shape + +```json +{ + "action": "run", + "code": "await page.goto(\"https://example.com\"); return await page.title();" +} +``` + +To reuse the same browser: + +```json +{ + "action": "run", + "code_session_id": "", + "code": "return { url: await page.url(), title: await page.title() };" +} +``` + +The tool returns the same JSON envelope as text and MCP +`structuredContent` for compatibility across agent harnesses. + +A cell timeout is not safe to replay automatically: the error sets +`may_have_side_effects: true` and `retryable: false`. The child runtime closes +the browser on a best-effort basis and then exits so timed-out JavaScript cannot +overlap a later cell. + +## Security status + +This is not a multi-tenant sandbox. Each logical code session gets +separate-process lifecycle containment, not a security boundary. Evaluated +JavaScript can access Node globals, the filesystem, network, inherited +credentials, and the child process environment. Do not expose this prototype +through an unauthenticated public tunnel. + +A hosted version needs a real sandbox/container boundary, scoped credentials, +authentication and authorization, quotas, idle expiry, output limits, and +independent Browserbase lease cleanup. diff --git a/packages/codemode/package.json b/packages/codemode/package.json new file mode 100644 index 0000000000..afb76b1b24 --- /dev/null +++ b/packages/codemode/package.json @@ -0,0 +1,39 @@ +{ + "name": "@browserbasehq/stagehand-codemode", + "version": "4.0.0", + "private": true, + "description": "Stagehand V4 code-mode runtime and MCP server", + "bin": { + "stagehand-codemode": "./dist/cli.mjs" + }, + "type": "module", + "exports": { + ".": { + "types": "./dist/index.d.mts", + "import": "./dist/index.mjs" + } + }, + "scripts": { + "build": "tsdown", + "test": "vitest run --config vitest.config.ts", + "test:live": "pnpm run build && node tests/live-smoke.mjs", + "test:unit": "pnpm run test", + "typecheck": "tsc --noEmit -p tsconfig.json" + }, + "dependencies": { + "@browserbasehq/sdk": "catalog:", + "@browserbasehq/stagehand": "workspace:*", + "@modelcontextprotocol/sdk": "catalog:", + "zod": "catalog:" + }, + "devDependencies": { + "@types/node": "catalog:", + "tsdown": "catalog:", + "tsx": "catalog:", + "typescript": "catalog:", + "vitest": "catalog:" + }, + "engines": { + "node": ">=22.18.0" + } +} diff --git a/packages/codemode/src/child-runtime.ts b/packages/codemode/src/child-runtime.ts new file mode 100644 index 0000000000..78aeba4a3e --- /dev/null +++ b/packages/codemode/src/child-runtime.ts @@ -0,0 +1,399 @@ +import { fork, type ChildProcess } from "node:child_process"; +import { createHash, randomUUID } from "node:crypto"; +import { fileURLToPath } from "node:url"; +import Browserbase from "@browserbasehq/sdk"; +import type { + CodeRuntime, + RuntimeRunResult, + RuntimeStatus, + StagehandCodeRuntimeConfig, +} from "./types.js"; +import { CodeModeRuntimeError } from "./types.js"; +import type { ChildRequest, ChildResponse } from "./runtime-protocol.js"; + +type ChildRequestWithoutId = ChildRequest extends infer Request + ? Request extends { id: string } + ? Omit + : never + : never; + +type PendingRequest = { + resolve: (value: unknown) => void; + reject: (error: Error) => void; +}; + +const CHILD_EXIT_GRACE_MS = 2_000; +const PARENT_WATCHDOG_GRACE_MS = 250; +const CONFIGURE_TIMEOUT_MS = 120_000; +const BROWSERBASE_RELEASE_SWEEP_DELAYS_MS = [0, 500, 1_500] as const; + +type RequestControl = { + hardTimeoutMs?: number; + terminateOnAbort?: boolean; + timeoutError?: () => CodeModeRuntimeError; +}; + +export type StagehandChildRuntimeOptions = { + childModuleUrl?: URL; +}; + +export class StagehandChildRuntime implements CodeRuntime { + private child?: ChildProcess; + private configurePromise?: Promise; + private readonly pending = new Map(); + private closePromise?: Promise; + private forceTerminationPromise?: Promise; + private browserbaseReleasePromise?: Promise; + private closed = false; + + constructor( + private readonly codeSessionId: string, + private readonly config: StagehandCodeRuntimeConfig, + private readonly options: StagehandChildRuntimeOptions = {}, + ) {} + + async run(code: string, timeoutMs: number, signal?: AbortSignal): Promise { + await this.ensureConfigured(signal); + return (await this.request({ type: "run", code, timeoutMs }, signal, false, { + hardTimeoutMs: timeoutMs + PARENT_WATCHDOG_GRACE_MS, + terminateOnAbort: true, + timeoutError: () => + new CodeModeRuntimeError("timeout", `Code execution exceeded ${timeoutMs}ms.`, false, { + mayHaveSideEffects: true, + }), + })) as RuntimeRunResult; + } + + async status(signal?: AbortSignal): Promise { + await this.ensureConfigured(signal); + return (await this.request({ type: "status" }, signal)) as RuntimeStatus; + } + + async reset(signal?: AbortSignal): Promise { + await this.ensureConfigured(signal); + await this.request({ type: "reset" }, signal); + } + + close(): Promise { + this.closePromise ??= this.closeInternal(); + return this.closePromise; + } + + private async closeInternal(): Promise { + if (this.closed) return; + this.closed = true; + const child = this.child; + if (!child) return; + let acknowledged = false; + + try { + if (child.connected) { + await this.request({ type: "close" }, undefined, true, { + hardTimeoutMs: CHILD_EXIT_GRACE_MS, + timeoutError: () => + new CodeModeRuntimeError( + "runtime", + "Stagehand code runtime did not acknowledge close.", + ), + }); + acknowledged = true; + } + } catch { + // The child may already be gone. The exit handler rejects pending work. + } finally { + if (child.exitCode === null && child.signalCode === null) { + child.kill("SIGTERM"); + const exited = new Promise((resolve) => child.once("exit", () => resolve())); + const grace = new Promise((resolve) => setTimeout(resolve, CHILD_EXIT_GRACE_MS)); + await Promise.race([exited, grace]); + if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL"); + } + if (!acknowledged) await this.releaseBrowserbaseSessions(); + if (this.child === child) this.child = undefined; + } + } + + private async ensureConfigured(signal?: AbortSignal): Promise { + await this.forceTerminationPromise; + if (this.closed) { + throw new CodeModeRuntimeError("closed", "Code session is closed."); + } + if (!this.configurePromise) { + const configurePromise = (async () => { + this.spawnChild(); + await this.request( + { + type: "configure", + codeSessionId: this.codeSessionId, + config: this.config, + }, + signal, + false, + { + hardTimeoutMs: this.config.defaultTimeoutMs ?? CONFIGURE_TIMEOUT_MS, + terminateOnAbort: true, + timeoutError: () => + new CodeModeRuntimeError( + "runtime", + "Stagehand code runtime configuration timed out.", + true, + ), + }, + ); + })(); + this.configurePromise = configurePromise; + void configurePromise.catch(() => { + if (this.configurePromise === configurePromise) this.configurePromise = undefined; + }); + } + await this.configurePromise; + } + + private spawnChild(): void { + if (this.child) return; + this.forceTerminationPromise = undefined; + const modulePath = fileURLToPath( + this.options.childModuleUrl ?? new URL("./runtime-child.mjs", import.meta.url), + ); + const child = fork(modulePath, [], { + stdio: ["ignore", "pipe", "pipe", "ipc"], + execArgv: [], + }); + this.child = child; + child.stdout?.on("data", (chunk) => process.stderr.write(`[codemode child] ${chunk}`)); + child.stderr?.on("data", (chunk) => process.stderr.write(`[codemode child] ${chunk}`)); + child.on("message", (message) => this.handleMessage(message)); + child.once("exit", (code, signal) => { + const error = new CodeModeRuntimeError( + "runtime", + `Stagehand code runtime exited${signal ? ` with signal ${signal}` : ` with code ${code}`}.`, + true, + ); + const pendingRequests = [...this.pending.values()]; + this.pending.clear(); + if (this.child === child) { + this.child = undefined; + if (!this.closed) this.configurePromise = undefined; + } + if (this.closed) { + for (const pending of pendingRequests) pending.reject(error); + return; + } + const recovery = this.forceTerminationPromise ?? this.forceTerminate(); + void recovery.then( + () => { + for (const pending of pendingRequests) pending.reject(error); + }, + () => { + for (const pending of pendingRequests) pending.reject(error); + }, + ); + }); + child.once("error", (error) => { + const pendingRequests = [...this.pending.values()]; + this.pending.clear(); + const recovery = this.forceTerminate(); + void recovery.then( + () => { + for (const pending of pendingRequests) pending.reject(error); + }, + () => { + for (const pending of pendingRequests) pending.reject(error); + }, + ); + }); + } + + private handleMessage(message: unknown): void { + if (!isChildResponse(message)) return; + const pending = this.pending.get(message.id); + if (!pending) return; + this.pending.delete(message.id); + if (message.ok) { + pending.resolve(message.result); + return; + } + const error = new CodeModeRuntimeError( + message.error.kind, + message.error.message, + message.error.retryable, + { + cause: message.error, + mayHaveSideEffects: message.error.mayHaveSideEffects, + }, + ); + if (message.error.kind !== "timeout") { + pending.reject(error); + return; + } + const recovery = this.forceTerminate(); + void recovery.then( + () => pending.reject(error), + () => pending.reject(error), + ); + } + + private request( + request: ChildRequestWithoutId, + signal?: AbortSignal, + allowClosed = false, + control: RequestControl = {}, + ): Promise { + if (this.closed && !allowClosed) { + return Promise.reject(new CodeModeRuntimeError("closed", "Code session is closed.")); + } + const child = this.child; + if (!child?.connected) { + return Promise.reject( + new CodeModeRuntimeError("runtime", "Stagehand code runtime is not connected.", true), + ); + } + if (signal?.aborted) { + if (control.terminateOnAbort) void this.forceTerminate(); + return Promise.reject( + new CodeModeRuntimeError("aborted", "Code execution was aborted.", true, { + cause: signal.reason, + }), + ); + } + + const id = randomUUID(); + return new Promise((resolve, reject) => { + let watchdog: NodeJS.Timeout | undefined; + const cleanup = () => { + signal?.removeEventListener("abort", onAbort); + if (watchdog) clearTimeout(watchdog); + }; + const onAbort = () => { + if (!this.pending.delete(id)) return; + cleanup(); + if (control.terminateOnAbort) void this.forceTerminate(); + reject( + new CodeModeRuntimeError("aborted", "Code execution was aborted.", false, { + cause: signal?.reason, + mayHaveSideEffects: true, + }), + ); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + this.pending.set(id, { + resolve: (value) => { + cleanup(); + resolve(value); + }, + reject: (error) => { + cleanup(); + reject(error); + }, + }); + if (control.hardTimeoutMs !== undefined) { + watchdog = setTimeout(() => { + if (!this.pending.delete(id)) return; + cleanup(); + void this.forceTerminate(); + reject( + control.timeoutError?.() ?? + new CodeModeRuntimeError("runtime", "Stagehand child request timed out."), + ); + }, control.hardTimeoutMs); + } + child.send({ ...request, id } as ChildRequest, (error) => { + if (!error) return; + this.pending.delete(id); + cleanup(); + void this.forceTerminate(); + reject(error); + }); + }); + } + + private forceTerminate(): Promise { + this.forceTerminationPromise ??= (async () => { + const child = this.child; + if (child && child.exitCode === null && child.signalCode === null) { + const exited = new Promise((resolve) => child.once("exit", () => resolve())); + child.kill("SIGKILL"); + const grace = new Promise((resolve) => setTimeout(resolve, CHILD_EXIT_GRACE_MS)); + await Promise.race([exited, grace]); + } + if (this.child === child) this.child = undefined; + if (!this.closed) this.configurePromise = undefined; + await this.releaseBrowserbaseSessions(); + })(); + return this.forceTerminationPromise; + } + + private releaseBrowserbaseSessions(): Promise { + if (this.browserbaseReleasePromise) return this.browserbaseReleasePromise; + const release = this.releaseBrowserbaseSessionsInternal(); + this.browserbaseReleasePromise = release; + void release.finally(() => { + if (this.browserbaseReleasePromise === release) this.browserbaseReleasePromise = undefined; + }); + return release; + } + + private async releaseBrowserbaseSessionsInternal(): Promise { + const apiKey = this.config.browserbaseApiKey; + if (!apiKey) return; + const browserbase = new Browserbase({ apiKey }); + let lastError: unknown; + for (const delayMs of BROWSERBASE_RELEASE_SWEEP_DELAYS_MS) { + if (delayMs > 0) { + await new Promise((resolve) => setTimeout(resolve, delayMs)); + } + try { + const sessionIds = ( + await browserbase.sessions.list({ + status: "RUNNING", + }) + ) + .filter( + (session) => + session.userMetadata?.integration === "stagehand-codemode-mcp" && + session.userMetadata?.codeSessionHash === this.codeSessionHash, + ) + .map((session) => session.id); + await Promise.all( + sessionIds.map((sessionId) => + browserbase.sessions.update(sessionId, { status: "REQUEST_RELEASE" }), + ), + ); + lastError = undefined; + } catch (error) { + lastError = error; + } + } + if (lastError !== undefined) { + const message = + lastError instanceof Error + ? lastError.message + : typeof lastError === "string" + ? lastError + : "Unknown Browserbase API error."; + process.stderr.write(`Failed to release a Stagehand code-mode browser: ${message}\n`); + } + } + + private get codeSessionHash(): string { + return createHash("sha256").update(this.codeSessionId).digest("hex").slice(0, 16); + } +} + +function isChildResponse(value: unknown): value is ChildResponse { + return ( + typeof value === "object" && + value !== null && + "id" in value && + typeof value.id === "string" && + "ok" in value && + typeof value.ok === "boolean" + ); +} + +export function createStagehandChildRuntime( + codeSessionId: string, + config: StagehandCodeRuntimeConfig, +): CodeRuntime { + return new StagehandChildRuntime(codeSessionId, config); +} diff --git a/packages/codemode/src/cli.ts b/packages/codemode/src/cli.ts new file mode 100644 index 0000000000..cd8e720b1c --- /dev/null +++ b/packages/codemode/src/cli.ts @@ -0,0 +1,94 @@ +#!/usr/bin/env node + +import { createStagehandChildRuntime } from "./child-runtime.js"; +import { runtimeConfigFromEnv } from "./config.js"; +import { connectCodeModeStdio, startCodeModeHttpServer } from "./mcp-server.js"; +import { CodeSessionManager } from "./session-manager.js"; + +const options = parseArgs(process.argv.slice(2)); +const runtimeConfig = runtimeConfigFromEnv(); +const manager = new CodeSessionManager({ + runtimeFactory: (codeSessionId) => createStagehandChildRuntime(codeSessionId, runtimeConfig), + defaultTimeoutMs: runtimeConfig.defaultTimeoutMs, +}); + +let closing = false; +const shutdown = async (exitCode: number) => { + if (closing) return; + closing = true; + await manager.closeAll().catch((error) => { + process.stderr.write(`Failed to close code sessions: ${String(error)}\n`); + }); + process.exit(exitCode); +}; +process.once("SIGINT", () => void shutdown(0)); +process.once("SIGTERM", () => void shutdown(0)); + +if (options.stdio) { + await connectCodeModeStdio(manager); + process.stdin.once("end", () => void shutdown(0)); + process.stdin.once("close", () => void shutdown(0)); + process.stderr.write("Stagehand V4 code-mode MCP listening on stdio\n"); +} else { + const running = await startCodeModeHttpServer({ + manager, + host: options.host, + port: options.port, + bearerToken: process.env.CODEMODE_MCP_BEARER_TOKEN, + }); + process.stderr.write(`Stagehand V4 code-mode MCP listening on ${running.url}\n`); +} + +type CliOptions = { + stdio: boolean; + host: string; + port: number; +}; + +function parseArgs(args: string[]): CliOptions { + let stdio = false; + let host = "127.0.0.1"; + let port = 8932; + for (let index = 0; index < args.length; index += 1) { + const argument = args[index]; + if (argument === "--stdio") { + stdio = true; + continue; + } + if (argument === "--host") { + host = requireValue(args, ++index, "--host"); + continue; + } + if (argument === "--port") { + const value = Number(requireValue(args, ++index, "--port")); + if (!Number.isSafeInteger(value) || value < 0 || value > 65_535) { + throw new Error("--port must be an integer between 0 and 65535."); + } + port = value; + continue; + } + if (argument === "--help" || argument === "-h") { + process.stdout.write( + [ + "Usage: stagehand-codemode [--stdio] [--host 127.0.0.1] [--port 8932]", + "", + "Environment:", + " BROWSERBASE_API_KEY Required on the first code_execute run", + " STAGEHAND_MODEL_NAME Optional provider/model name for AI methods", + " STAGEHAND_MODEL_API_KEY Optional model provider key", + " CODEMODE_MCP_BEARER_TOKEN Optional HTTP bearer token", + "", + ].join("\n"), + ); + process.exit(0); + } + throw new Error(`Unknown argument: ${argument}`); + } + return { stdio, host, port }; +} + +function requireValue(args: string[], index: number, flag: string): string { + const value = args[index]; + if (!value) throw new Error(`${flag} requires a value.`); + return value; +} diff --git a/packages/codemode/src/config.ts b/packages/codemode/src/config.ts new file mode 100644 index 0000000000..272e5b8477 --- /dev/null +++ b/packages/codemode/src/config.ts @@ -0,0 +1,44 @@ +import type { StagehandCodeRuntimeConfig } from "./types.js"; + +export function runtimeConfigFromEnv( + env: NodeJS.ProcessEnv = process.env, +): StagehandCodeRuntimeConfig { + const explicitModelName = nonEmpty(env.STAGEHAND_MODEL_NAME); + const explicitModelApiKey = nonEmpty(env.STAGEHAND_MODEL_API_KEY); + const inferredGoogleKey = + nonEmpty(env.GEMINI_API_KEY) ?? + nonEmpty(env.GOOGLE_API_KEY) ?? + nonEmpty(env.GOOGLE_GENERATIVE_AI_API_KEY); + const modelName = + explicitModelName ?? (inferredGoogleKey ? "google/gemini-2.5-flash-lite" : undefined); + const modelApiKey = explicitModelApiKey ?? inferredGoogleKey; + + return { + browserbaseApiKey: nonEmpty(env.BROWSERBASE_API_KEY), + ...(modelName + ? { + model: { + modelName, + ...(modelApiKey ? { apiKey: modelApiKey } : {}), + ...(nonEmpty(env.STAGEHAND_MODEL_BASE_URL) + ? { baseURL: nonEmpty(env.STAGEHAND_MODEL_BASE_URL)! } + : {}), + }, + } + : {}), + ...(positiveInt(env.CODEMODE_DEFAULT_TIMEOUT_MS) + ? { defaultTimeoutMs: positiveInt(env.CODEMODE_DEFAULT_TIMEOUT_MS) } + : {}), + }; +} + +function nonEmpty(value: string | undefined): string | undefined { + const trimmed = value?.trim(); + return trimmed ? trimmed : undefined; +} + +function positiveInt(value: string | undefined): number | undefined { + if (!value) return undefined; + const parsed = Number(value); + return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : undefined; +} diff --git a/packages/codemode/src/facades.ts b/packages/codemode/src/facades.ts new file mode 100644 index 0000000000..3ddabaf6dc --- /dev/null +++ b/packages/codemode/src/facades.ts @@ -0,0 +1,329 @@ +import type { z } from "zod/v4"; + +type UnknownRecord = Record; +type UnknownMethod = (...args: unknown[]) => unknown; + +export interface CodeLocatorFacade { + click(...args: unknown[]): Promise; + hover(...args: unknown[]): Promise; + fill(...args: unknown[]): Promise; + count(...args: unknown[]): Promise; + isChecked(...args: unknown[]): Promise; + inputValue(...args: unknown[]): Promise; + isVisible(...args: unknown[]): Promise; + innerText(...args: unknown[]): Promise; + innerHtml(...args: unknown[]): Promise; + textContent(...args: unknown[]): Promise; + scrollTo(...args: unknown[]): Promise; + centroid(...args: unknown[]): Promise; + highlight(...args: unknown[]): Promise; + sendClickEvent(...args: unknown[]): Promise; + type(...args: unknown[]): Promise; + selectOption(...args: unknown[]): Promise; + first(): CodeLocatorFacade; + nth(index: number): CodeLocatorFacade; +} + +export interface CodePageFacade { + readonly pageId: string; + goto(...args: unknown[]): Promise; + reload(...args: unknown[]): Promise; + goBack(...args: unknown[]): Promise; + goForward(...args: unknown[]): Promise; + click(...args: unknown[]): Promise; + hover(...args: unknown[]): Promise; + scroll(...args: unknown[]): Promise; + dragAndDrop(...args: unknown[]): Promise; + type(...args: unknown[]): Promise; + keyPress(...args: unknown[]): Promise; + evaluate(...args: unknown[]): Promise; + addInitScript(...args: unknown[]): Promise; + setExtraHTTPHeaders(...args: unknown[]): Promise; + setViewportSize(...args: unknown[]): Promise; + waitForLoadState(...args: unknown[]): Promise; + waitForTimeout(...args: unknown[]): Promise; + waitForSelector(...args: unknown[]): Promise; + screenshot(...args: unknown[]): Promise; + snapshot(...args: unknown[]): Promise; + url(): Promise; + title(): Promise; + close(): Promise; + locator(selector: string): CodeLocatorFacade; +} + +export interface CodeClipboardFacade { + readText(options?: unknown): Promise; + writeText(text: string, options?: unknown): Promise; + clear(options?: unknown): Promise; + paste(options?: unknown): Promise; + copy(options?: unknown): Promise; + cut(options?: unknown): Promise; +} + +export interface CodeContextFacade { + readonly clipboard: CodeClipboardFacade; + pages(): Promise; + newPage(...args: unknown[]): Promise; + activePage(): Promise; + setActivePage(page: CodePageFacade): Promise; + addInitScript(...args: unknown[]): Promise; + setExtraHTTPHeaders(...args: unknown[]): Promise; + getDomainPolicy(...args: unknown[]): Promise; + setDomainPolicy(...args: unknown[]): Promise; + cookies(...args: unknown[]): Promise; + addCookies(...args: unknown[]): Promise; + clearCookies(...args: unknown[]): Promise; +} + +export interface CodeStagehandFacade { + act(instruction: unknown, options?: Record): Promise; + observe(instruction?: string, options?: Record): Promise; + extract( + instruction: string, + schema: z.ZodType, + options?: Record, + ): Promise; +} + +export type CodeFacades = { + context: CodeContextFacade; + wrapPage: (rawPage: unknown) => CodePageFacade; + stagehand: CodeStagehandFacade; +}; + +export function createCodeFacades(rawStagehand: unknown, rawContext: unknown): CodeFacades { + const stagehandTarget = requireObject(rawStagehand, "Stagehand instance"); + const contextTarget = requireObject(rawContext, "Stagehand browser context"); + const pagesByRaw = new WeakMap(); + const rawPagesByFacade = new WeakMap(); + const locatorsByRaw = new WeakMap(); + const rawLocatorsByFacade = new WeakMap(); + + const wrapLocator = (rawLocator: unknown): CodeLocatorFacade => { + const locatorTarget = requireObject(rawLocator, "Stagehand locator"); + const existing = locatorsByRaw.get(locatorTarget); + if (existing) return existing; + + const facade: CodeLocatorFacade = { + click: (...args) => invoke(locatorTarget, "click", args), + hover: (...args) => invoke(locatorTarget, "hover", args), + fill: (...args) => invoke(locatorTarget, "fill", args), + count: (...args) => invoke(locatorTarget, "count", args), + isChecked: (...args) => invoke(locatorTarget, "isChecked", args), + inputValue: (...args) => invoke(locatorTarget, "inputValue", args), + isVisible: (...args) => invoke(locatorTarget, "isVisible", args), + innerText: (...args) => invoke(locatorTarget, "innerText", args), + innerHtml: (...args) => invoke(locatorTarget, "innerHtml", args), + textContent: (...args) => invoke(locatorTarget, "textContent", args), + scrollTo: (...args) => invoke(locatorTarget, "scrollTo", args), + centroid: (...args) => invoke(locatorTarget, "centroid", args), + highlight: (...args) => invoke(locatorTarget, "highlight", args), + sendClickEvent: (...args) => invoke(locatorTarget, "sendClickEvent", args), + type: (...args) => invoke(locatorTarget, "type", args), + selectOption: (...args) => invoke(locatorTarget, "selectOption", args), + first: () => wrapLocator(invokeSync(locatorTarget, "first", [])), + nth: (index) => wrapLocator(invokeSync(locatorTarget, "nth", [index])), + }; + + Object.freeze(facade); + locatorsByRaw.set(locatorTarget, facade); + rawLocatorsByFacade.set(facade, locatorTarget); + return facade; + }; + + const wrapPage = (rawPage: unknown): CodePageFacade => { + const pageTarget = requireObject(rawPage, "Stagehand page"); + const existing = pagesByRaw.get(pageTarget); + if (existing) return existing; + + const facade: CodePageFacade = { + pageId: readString(pageTarget, "pageId"), + goto: async (...args) => { + await invoke(pageTarget, "goto", args); + return facade; + }, + reload: async (...args) => { + await invoke(pageTarget, "reload", args); + return facade; + }, + goBack: async (...args) => { + await invoke(pageTarget, "goBack", args); + return facade; + }, + goForward: async (...args) => { + await invoke(pageTarget, "goForward", args); + return facade; + }, + click: (...args) => invoke(pageTarget, "click", args), + hover: (...args) => invoke(pageTarget, "hover", args), + scroll: (...args) => invoke(pageTarget, "scroll", args), + dragAndDrop: (...args) => invoke(pageTarget, "dragAndDrop", args), + type: (...args) => invoke(pageTarget, "type", args), + keyPress: (...args) => invoke(pageTarget, "keyPress", args), + evaluate: (...args) => invoke(pageTarget, "evaluate", args), + addInitScript: (...args) => invoke(pageTarget, "addInitScript", args), + setExtraHTTPHeaders: (...args) => invoke(pageTarget, "setExtraHTTPHeaders", args), + setViewportSize: (...args) => invoke(pageTarget, "setViewportSize", args), + waitForLoadState: (...args) => invoke(pageTarget, "waitForLoadState", args), + waitForTimeout: (...args) => invoke(pageTarget, "waitForTimeout", args), + waitForSelector: (...args) => invoke(pageTarget, "waitForSelector", args), + screenshot: (...args) => invoke(pageTarget, "screenshot", unwrapScreenshotArgs(args)), + snapshot: (...args) => invoke(pageTarget, "snapshot", args), + url: () => invoke(pageTarget, "url", []) as Promise, + title: () => invoke(pageTarget, "title", []) as Promise, + close: () => invoke(pageTarget, "close", []) as Promise, + locator: (selector) => wrapLocator(invokeSync(pageTarget, "locator", [selector])), + }; + + Object.freeze(facade); + pagesByRaw.set(pageTarget, facade); + rawPagesByFacade.set(facade, pageTarget); + return facade; + }; + + const clipboardTarget = requireObject(contextTarget.clipboard, "Stagehand browser clipboard"); + const clipboard: CodeClipboardFacade = Object.freeze({ + readText: (options?: unknown) => + invoke(clipboardTarget, "readText", [unwrapPageOption(options)]), + writeText: (text: string, options?: unknown) => + invoke(clipboardTarget, "writeText", [text, unwrapPageOption(options)]), + clear: (options?: unknown) => invoke(clipboardTarget, "clear", [unwrapPageOption(options)]), + paste: (options?: unknown) => invoke(clipboardTarget, "paste", [unwrapPageOption(options)]), + copy: (options?: unknown) => invoke(clipboardTarget, "copy", [unwrapPageOption(options)]), + cut: (options?: unknown) => invoke(clipboardTarget, "cut", [unwrapPageOption(options)]), + }); + + const context: CodeContextFacade = Object.freeze({ + clipboard, + pages: async () => requireArray(await invoke(contextTarget, "pages", [])).map(wrapPage), + newPage: async (...args: unknown[]) => wrapPage(await invoke(contextTarget, "newPage", args)), + activePage: async () => { + const active = await invoke(contextTarget, "activePage", []); + return active === undefined || active === null ? undefined : wrapPage(active); + }, + setActivePage: async (page: CodePageFacade) => { + await invoke(contextTarget, "setActivePage", [ + requireOwnedPage(page, rawPagesByFacade, "context.setActivePage"), + ]); + }, + addInitScript: (...args: unknown[]) => invoke(contextTarget, "addInitScript", args), + setExtraHTTPHeaders: (...args: unknown[]) => invoke(contextTarget, "setExtraHTTPHeaders", args), + getDomainPolicy: (...args: unknown[]) => invoke(contextTarget, "getDomainPolicy", args), + setDomainPolicy: (...args: unknown[]) => invoke(contextTarget, "setDomainPolicy", args), + cookies: (...args: unknown[]) => invoke(contextTarget, "cookies", args), + addCookies: (...args: unknown[]) => invoke(contextTarget, "addCookies", args), + clearCookies: (...args: unknown[]) => invoke(contextTarget, "clearCookies", args), + }); + + const stagehand: CodeStagehandFacade = Object.freeze({ + act: (instruction: unknown, options?: Record) => + invoke( + stagehandTarget, + "act", + options === undefined ? [instruction] : [instruction, unwrapPageOption(options)], + ), + observe: (instruction?: string, options?: Record) => + invoke( + stagehandTarget, + "observe", + options === undefined + ? instruction === undefined + ? [] + : [instruction] + : [instruction, unwrapPageOption(options)], + ), + extract: (instruction: string, schema: z.ZodType, options?: Record) => + invoke( + stagehandTarget, + "extract", + options === undefined + ? [instruction, schema] + : [instruction, schema, unwrapPageOption(options)], + ), + }); + + return { context, wrapPage, stagehand }; + + function unwrapPageOption(options: unknown): unknown { + if (!isRecord(options) || options.page === undefined) return options; + return { + ...options, + page: requireOwnedPage(options.page, rawPagesByFacade, "page option"), + }; + } + + function unwrapScreenshotArgs(args: unknown[]): unknown[] { + const [options, ...rest] = args; + if (!isRecord(options) || !Array.isArray(options.mask)) return args; + const mask = options.mask.map((locator) => { + if (!isObject(locator)) { + throw new Error("screenshot mask entries must be Stagehand locator facades."); + } + const rawLocator = rawLocatorsByFacade.get(locator); + if (!rawLocator) { + throw new Error("screenshot mask entries must come from this code session."); + } + return rawLocator; + }); + return [{ ...options, mask }, ...rest]; + } +} + +function invoke(target: UnknownRecord, methodName: string, args: unknown[]): Promise { + const method = target[methodName]; + if (typeof method !== "function") { + throw new Error(`Stagehand object does not expose ${methodName}().`); + } + return Promise.resolve((method as UnknownMethod).apply(target, args)); +} + +function invokeSync(target: UnknownRecord, methodName: string, args: unknown[]): unknown { + const method = target[methodName]; + if (typeof method !== "function") { + throw new Error(`Stagehand object does not expose ${methodName}().`); + } + return (method as UnknownMethod).apply(target, args); +} + +function requireObject(value: unknown, label: string): UnknownRecord { + if (!isObject(value)) throw new Error(`${label} is not an object.`); + return value as UnknownRecord; +} + +function requireArray(value: unknown): unknown[] { + if (!Array.isArray(value)) { + throw new Error("Stagehand context.pages() did not return an array."); + } + return value; +} + +function requireOwnedPage( + page: unknown, + rawPagesByFacade: WeakMap, + label: string, +): object { + if (!isObject(page)) { + throw new Error(`${label} requires a page facade from this code session.`); + } + const rawPage = rawPagesByFacade.get(page); + if (!rawPage) { + throw new Error(`${label} requires a page facade from this code session.`); + } + return rawPage; +} + +function readString(target: UnknownRecord, key: string): string { + const value = target[key]; + if (typeof value !== "string" || value.length === 0) { + throw new Error(`Stagehand object is missing string property ${key}.`); + } + return value; +} + +function isObject(value: unknown): value is object { + return typeof value === "object" && value !== null; +} + +function isRecord(value: unknown): value is UnknownRecord { + return isObject(value) && !Array.isArray(value); +} diff --git a/packages/codemode/src/index.ts b/packages/codemode/src/index.ts new file mode 100644 index 0000000000..8b7f2e6248 --- /dev/null +++ b/packages/codemode/src/index.ts @@ -0,0 +1,16 @@ +export { createStagehandChildRuntime, StagehandChildRuntime } from "./child-runtime.js"; +export { runtimeConfigFromEnv } from "./config.js"; +export { + connectCodeModeStdio, + createCodeModeMcpServer, + startCodeModeHttpServer, + type CodeModeHttpServerOptions, + type RunningCodeModeHttpServer, +} from "./mcp-server.js"; +export { + CodeSessionManager, + codeExecuteResultText, + type CodeRuntimeFactory, + type CodeSessionManagerOptions, +} from "./session-manager.js"; +export * from "./types.js"; diff --git a/packages/codemode/src/mcp-server.ts b/packages/codemode/src/mcp-server.ts new file mode 100644 index 0000000000..7754548d67 --- /dev/null +++ b/packages/codemode/src/mcp-server.ts @@ -0,0 +1,253 @@ +import { timingSafeEqual } from "node:crypto"; +import http from "node:http"; +import { randomUUID } from "node:crypto"; +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import { z } from "zod/v4"; +import { CodeSessionManager, codeExecuteResultText } from "./session-manager.js"; +import { CODE_EXECUTE_ACTIONS, type CodeExecuteInput } from "./types.js"; + +const DEFAULT_PORT = 8932; +const DEFAULT_HOST = "127.0.0.1"; + +export type CodeModeHttpServerOptions = { + manager: CodeSessionManager; + port?: number; + host?: string; + bearerToken?: string; +}; + +export type RunningCodeModeHttpServer = { + url: string; + close(): Promise; +}; + +export function createCodeModeMcpServer(manager: CodeSessionManager): McpServer { + const server = new McpServer({ + name: "stagehand-codemode", + version: "0.1.0", + }); + + server.registerTool( + "code_execute", + { + title: "Execute Stagehand V4 code", + description: [ + "Execute JavaScript against a long-lived Stagehand V4 browser running exclusively on Browserbase.", + "The first action=run lazily creates the remote browser. Later calls reuse it by passing the returned opaque code_session_id.", + "page, context, stagehand, z, and console are already in scope. Use await directly and return JSON-serializable values.", + "Do not call stagehand.close(); use action=close when the task is finished.", + ].join(" "), + inputSchema: { + action: z.enum(CODE_EXECUTE_ACTIONS).default("run"), + code_session_id: z + .string() + .optional() + .describe("Opaque code session ID returned by an earlier run."), + code: z + .string() + .optional() + .describe("Async JavaScript function body. Required for action=run."), + timeout_ms: z + .number() + .int() + .positive() + .max(300_000) + .optional() + .describe("Execution timeout in milliseconds."), + }, + outputSchema: z + .object({ + ok: z.boolean(), + action: z.enum(CODE_EXECUTE_ACTIONS), + }) + .loose(), + }, + async (arguments_, extra) => { + const result = await manager.execute(arguments_ as CodeExecuteInput, extra.signal); + return { + content: [{ type: "text" as const, text: codeExecuteResultText(result) }], + structuredContent: result, + isError: !result.ok, + }; + }, + ); + + return server; +} + +export async function connectCodeModeStdio(manager: CodeSessionManager): Promise { + const server = createCodeModeMcpServer(manager); + await server.connect(new StdioServerTransport()); + return server; +} + +export async function startCodeModeHttpServer( + options: CodeModeHttpServerOptions, +): Promise { + const host = options.host ?? DEFAULT_HOST; + const port = options.port ?? DEFAULT_PORT; + if (!options.bearerToken && !isLoopbackHost(host)) { + throw new Error( + "CODEMODE_MCP_BEARER_TOKEN is required when the HTTP server binds to a non-loopback host.", + ); + } + const sessions = new Map< + string, + { transport: StreamableHTTPServerTransport; server: McpServer } + >(); + let closing = false; + + const httpServer = http.createServer(async (request, response) => { + try { + if (closing) { + response.writeHead(503).end("Server is shutting down"); + return; + } + if (!options.bearerToken && !isSafeLoopbackRequest(request)) { + response.writeHead(403).end("Forbidden"); + return; + } + if (!authorized(request, options.bearerToken)) { + response.writeHead(401, { "www-authenticate": "Bearer" }).end("Unauthorized"); + return; + } + const url = new URL(request.url ?? "/", `http://${request.headers.host ?? "localhost"}`); + if (url.pathname === "/health") { + response.writeHead(200, { "content-type": "application/json" }); + response.end( + JSON.stringify({ + ok: true, + browserProvisioning: "lazy", + activeCodeSessions: options.manager.activeSessionCount, + }), + ); + return; + } + if (url.pathname !== "/mcp") { + response.writeHead(404).end("Not found"); + return; + } + + const mcpSessionId = request.headers["mcp-session-id"]; + if (typeof mcpSessionId === "string") { + const active = sessions.get(mcpSessionId); + if (!active) { + response.writeHead(404).end("MCP session not found"); + return; + } + await active.transport.handleRequest(request, response); + return; + } + + if (request.method !== "POST") { + response.writeHead(400).end("MCP initialization requires POST"); + return; + } + + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + enableJsonResponse: true, + onsessioninitialized: (sessionId) => { + sessions.set(sessionId, { transport, server }); + }, + }); + const server = createCodeModeMcpServer(options.manager); + transport.onclose = () => { + if (transport.sessionId) sessions.delete(transport.sessionId); + }; + await server.connect(transport); + await transport.handleRequest(request, response); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + if (!response.headersSent) { + response.writeHead(500, { "content-type": "application/json" }); + } + if (!response.writableEnded) response.end(JSON.stringify({ error: message })); + } + }); + + await new Promise((resolve, reject) => { + httpServer.once("error", reject); + httpServer.listen(port, host, () => { + httpServer.off("error", reject); + resolve(); + }); + }); + const address = httpServer.address(); + if (!address || typeof address === "string") throw new Error("HTTP server did not bind a port."); + const resolvedHost = + address.address === "0.0.0.0" || address.address === "::" + ? "127.0.0.1" + : address.family === "IPv6" + ? `[${address.address}]` + : address.address; + const serverUrl = `http://${resolvedHost}:${address.port}/mcp`; + let closePromise: Promise | undefined; + + return { + url: serverUrl, + close() { + closePromise ??= (async () => { + closing = true; + const stopped = new Promise((resolve, reject) => { + httpServer.close((error) => (error ? reject(error) : resolve())); + }); + httpServer.closeAllConnections(); + for (const { transport, server } of sessions.values()) { + await transport.close().catch(() => undefined); + await server.close().catch(() => undefined); + } + sessions.clear(); + await options.manager.closeAll(); + await stopped; + })(); + return closePromise; + }, + }; +} + +function authorized(request: http.IncomingMessage, expected?: string): boolean { + if (!expected) return true; + const authorization = request.headers.authorization; + if (!authorization?.startsWith("Bearer ")) return false; + const actual = Buffer.from(authorization.slice("Bearer ".length)); + const wanted = Buffer.from(expected); + return actual.length === wanted.length && timingSafeEqual(actual, wanted); +} + +function isLoopbackHost(host: string): boolean { + return host === "127.0.0.1" || host === "localhost" || host === "::1"; +} + +function isSafeLoopbackRequest(request: http.IncomingMessage): boolean { + if (!isLoopbackAuthority(request.headers.host)) return false; + const origin = request.headers.origin; + if (origin === undefined) return true; + try { + const parsed = new URL(origin); + return ( + (parsed.protocol === "http:" || parsed.protocol === "https:") && + isLoopbackHost(normalizeHostname(parsed.hostname)) + ); + } catch { + return false; + } +} + +function isLoopbackAuthority(authority: string | undefined): boolean { + if (!authority) return false; + try { + const parsed = new URL(`http://${authority}`); + return isLoopbackHost(normalizeHostname(parsed.hostname)); + } catch { + return false; + } +} + +function normalizeHostname(hostname: string): string { + return hostname.startsWith("[") && hostname.endsWith("]") + ? hostname.slice(1, -1) + : hostname.toLowerCase(); +} diff --git a/packages/codemode/src/runtime-child.ts b/packages/codemode/src/runtime-child.ts new file mode 100644 index 0000000000..4525291edb --- /dev/null +++ b/packages/codemode/src/runtime-child.ts @@ -0,0 +1,247 @@ +import { createHash } from "node:crypto"; +import { Stagehand, type StagehandClientInitParams } from "@browserbasehq/stagehand"; +import { z } from "zod/v4"; +import { createCodeFacades, type CodeFacades, type CodePageFacade } from "./facades.js"; +import type { ChildRequest, ChildResponse } from "./runtime-protocol.js"; +import type { + CodeLogEntry, + CodePageState, + RuntimeRunResult, + RuntimeStatus, + StagehandCodeRuntimeConfig, +} from "./types.js"; + +const AsyncFunction = Object.getPrototypeOf(async function () {}).constructor as new ( + ...args: string[] +) => (...values: unknown[]) => Promise; + +let config: StagehandCodeRuntimeConfig | undefined; +let codeSessionId: string | undefined; +let stagehand: Stagehand | undefined; +let facades: CodeFacades | undefined; +let closed = false; +let queue = Promise.resolve(); + +process.on("message", (message: unknown) => { + if (!isChildRequest(message)) return; + queue = queue.then(() => handle(message)); +}); + +process.once("SIGTERM", () => void shutdown(0)); +process.once("SIGINT", () => void shutdown(0)); +process.once("disconnect", () => void shutdown(0)); + +async function handle(request: ChildRequest): Promise { + try { + switch (request.type) { + case "configure": { + if (config) throw new Error("Stagehand code runtime is already configured."); + config = request.config; + codeSessionId = request.codeSessionId; + send({ id: request.id, ok: true }); + return; + } + case "run": { + requireConfigured(); + const result = await runSnippet(request.code, request.timeoutMs); + send({ id: request.id, ok: true, result }); + return; + } + case "status": { + requireConfigured(); + send({ id: request.id, ok: true, result: await status() }); + return; + } + case "reset": { + requireConfigured(); + await closeStagehand(); + send({ id: request.id, ok: true, result: { reset: true } }); + return; + } + case "close": { + closed = true; + await closeStagehand(); + send({ id: request.id, ok: true, result: { closed: true } }); + setImmediate(() => process.exit(0)); + return; + } + } + } catch (error) { + const normalized = normalizeError(error); + send({ + id: request.id, + ok: false, + error: normalized, + page: await readPageState().catch(() => undefined), + }); + if (normalized.kind === "timeout") { + closed = true; + void closeStagehand() + .catch(() => undefined) + .finally(() => process.exit(1)); + setTimeout(() => process.exit(1), 2_000).unref(); + } + } +} + +async function ensureStagehand(): Promise<{ + stagehand: Stagehand; + facades: CodeFacades; +}> { + requireConfigured(); + if (stagehand && facades) return { stagehand, facades }; + const apiKey = config!.browserbaseApiKey; + if (!apiKey) { + throw new Error("BROWSERBASE_API_KEY is required before the first code_execute run."); + } + + const initParams: StagehandClientInitParams = { + apiKey, + browser: { + type: "browserbase", + userMetadata: { + integration: "stagehand-codemode-mcp", + codeSessionHash: createHash("sha256").update(codeSessionId!).digest("hex").slice(0, 16), + }, + }, + logging: { level: "off" }, + ...(config!.model ? { model: config!.model as StagehandClientInitParams["model"] } : {}), + }; + const next = new Stagehand(initParams); + await next.init(); + stagehand = next; + facades = createCodeFacades(next, next.context); + return { stagehand: next, facades }; +} + +async function runSnippet(code: string, timeoutMs: number): Promise { + if (closed) throw new Error("Stagehand code runtime is closed."); + const runtime = await ensureStagehand(); + const rawPage = + (await runtime.stagehand.context.activePage()) ?? + (await runtime.stagehand.context.pages())[0] ?? + (await runtime.stagehand.context.newPage()); + const page = runtime.facades.wrapPage(rawPage); + const logs: CodeLogEntry[] = []; + const codeConsole = Object.freeze({ + log: (...values: unknown[]) => logs.push({ level: "log", text: formatLog(values) }), + warn: (...values: unknown[]) => logs.push({ level: "warn", text: formatLog(values) }), + error: (...values: unknown[]) => logs.push({ level: "error", text: formatLog(values) }), + }); + const fn = new AsyncFunction("page", "context", "stagehand", "z", "console", code); + let timeout: NodeJS.Timeout | undefined; + try { + const value = await Promise.race([ + fn(page, runtime.facades.context, runtime.facades.stagehand, z, codeConsole), + new Promise((_, reject) => { + timeout = setTimeout(() => { + const error = new Error(`Code execution exceeded ${timeoutMs}ms.`); + error.name = "CodeExecutionTimeoutError"; + reject(error); + }, timeoutMs); + }), + ]); + return { + value: jsonSafe(value), + logs, + page: await readRequiredPageState(page), + }; + } finally { + if (timeout) clearTimeout(timeout); + } +} + +async function status(): Promise { + if (!stagehand) return { state: "idle" }; + return { + state: "ready", + page: await readPageState(), + }; +} + +async function readPageState(): Promise { + if (!stagehand || !facades) return undefined; + const rawPage = (await stagehand.context.activePage()) ?? (await stagehand.context.pages())[0]; + if (!rawPage) return undefined; + return readRequiredPageState(facades.wrapPage(rawPage)); +} + +async function readRequiredPageState(page: CodePageFacade): Promise { + const [url, title] = await Promise.all([page.url(), page.title()]); + return { url, title }; +} + +async function closeStagehand(): Promise { + const current = stagehand; + stagehand = undefined; + facades = undefined; + await current?.close(); +} + +async function shutdown(exitCode: number): Promise { + closed = true; + await closeStagehand().catch(() => undefined); + process.exit(exitCode); +} + +function requireConfigured(): void { + if (!config || !codeSessionId) throw new Error("Stagehand code runtime is not configured."); +} + +function send(response: ChildResponse): void { + if (process.connected) process.send?.(response); +} + +function normalizeError(error: unknown): Extract["error"] { + const normalized = error instanceof Error ? error : new Error(String(error)); + const timeout = normalized.name === "CodeExecutionTimeoutError"; + return { + name: normalized.name, + message: normalized.message, + kind: timeout ? "timeout" : closed ? "closed" : "runtime", + retryable: false, + mayHaveSideEffects: timeout, + ...(normalized.stack ? { stack: normalized.stack } : {}), + }; +} + +function formatLog(values: unknown[]): string { + return values + .map((value) => { + if (typeof value === "string") return value; + try { + return JSON.stringify(jsonSafe(value)); + } catch { + return String(value); + } + }) + .join(" "); +} + +function jsonSafe(value: unknown): unknown { + if (value === undefined) return undefined; + return JSON.parse( + JSON.stringify(value, (_key, nested) => { + if (typeof nested === "bigint") return nested.toString(); + if (nested instanceof Uint8Array) { + return { + type: "bytes", + encoding: "base64", + data: Buffer.from(nested).toString("base64"), + }; + } + return nested; + }), + ); +} + +function isChildRequest(value: unknown): value is ChildRequest { + return ( + typeof value === "object" && + value !== null && + "id" in value && + typeof value.id === "string" && + "type" in value && + typeof value.type === "string" + ); +} diff --git a/packages/codemode/src/runtime-protocol.ts b/packages/codemode/src/runtime-protocol.ts new file mode 100644 index 0000000000..79608875a7 --- /dev/null +++ b/packages/codemode/src/runtime-protocol.ts @@ -0,0 +1,54 @@ +import type { + CodeLogEntry, + CodePageState, + RuntimeRunResult, + RuntimeStatus, + StagehandCodeRuntimeConfig, +} from "./types.js"; + +export type ChildRequest = + | { + id: string; + type: "configure"; + codeSessionId: string; + config: StagehandCodeRuntimeConfig; + } + | { + id: string; + type: "run"; + code: string; + timeoutMs: number; + } + | { + id: string; + type: "status"; + } + | { + id: string; + type: "reset"; + } + | { + id: string; + type: "close"; + }; + +export type ChildResponse = + | { + id: string; + ok: true; + result?: RuntimeRunResult | RuntimeStatus | { closed: true } | { reset: true }; + } + | { + id: string; + ok: false; + error: { + name: string; + message: string; + kind: "runtime" | "timeout" | "closed"; + retryable: boolean; + mayHaveSideEffects: boolean; + stack?: string; + }; + page?: CodePageState; + logs?: CodeLogEntry[]; + }; diff --git a/packages/codemode/src/session-manager.ts b/packages/codemode/src/session-manager.ts new file mode 100644 index 0000000000..7cecba4ec5 --- /dev/null +++ b/packages/codemode/src/session-manager.ts @@ -0,0 +1,310 @@ +import { randomUUID } from "node:crypto"; +import type { + CodeExecuteFailure, + CodeExecuteInput, + CodeExecuteResult, + CodeExecuteSuccess, + CodePageState, + CodeRuntime, + CodeSessionState, + RuntimeStatus, +} from "./types.js"; +import { CodeModeRuntimeError } from "./types.js"; + +type ManagedSession = { + id: string; + runtime?: CodeRuntime; + queue: Promise; + state: Exclude; + page?: CodePageState; +}; + +export type CodeRuntimeFactory = (codeSessionId: string) => CodeRuntime; + +export type CodeSessionManagerOptions = { + runtimeFactory: CodeRuntimeFactory; + defaultTimeoutMs?: number; + sessionIdFactory?: () => string; +}; + +const DEFAULT_TIMEOUT_MS = 120_000; + +export class CodeSessionManager { + private readonly sessions = new Map(); + private readonly defaultTimeoutMs: number; + private readonly sessionIdFactory: () => string; + private closeAllPromise?: Promise; + + constructor(private readonly options: CodeSessionManagerOptions) { + this.defaultTimeoutMs = requirePositiveTimeout( + options.defaultTimeoutMs ?? DEFAULT_TIMEOUT_MS, + "defaultTimeoutMs", + ); + this.sessionIdFactory = + options.sessionIdFactory ?? (() => `code_${randomUUID().replaceAll("-", "")}`); + } + + async execute(input: CodeExecuteInput, signal?: AbortSignal): Promise { + const validationError = validateInput(input); + if (validationError) return failure(input, "validation", validationError); + + if (input.action === "status" && input.code_session_id === undefined) { + return { + ok: true, + action: "status", + state: "idle", + active_code_sessions: this.sessions.size, + }; + } + + if (input.action === "run") { + const created = input.code_session_id === undefined; + const session = created ? this.createSession() : this.sessions.get(input.code_session_id!); + if (!session) return sessionNotFound(input); + const result = await this.enqueue(session, input, signal, () => + this.run(session, input, signal), + ); + if (created && !result.ok && result.error.kind === "aborted") { + await session.runtime?.close().catch(() => undefined); + this.sessions.delete(session.id); + } + return result; + } + + const session = this.sessions.get(input.code_session_id!); + if (!session) return sessionNotFound(input); + + return this.enqueue(session, input, signal, async () => { + switch (input.action) { + case "status": + return this.status(session, signal); + case "reset": + return this.reset(session, signal); + case "close": + return this.close(session); + case "run": + throw new Error("run is handled before session lookup"); + } + }); + } + + async closeAll(): Promise { + this.closeAllPromise ??= (async () => { + const sessions = [...this.sessions.values()]; + await Promise.all(sessions.map((session) => this.close(session))); + })(); + return this.closeAllPromise; + } + + get activeSessionCount(): number { + return this.sessions.size; + } + + private createSession(): ManagedSession { + const id = this.sessionIdFactory(); + if (this.sessions.has(id)) { + throw new Error(`Code session ID factory returned a duplicate ID: ${id}`); + } + const session: ManagedSession = { + id, + queue: Promise.resolve(), + state: "idle", + }; + this.sessions.set(id, session); + return session; + } + + private async run( + session: ManagedSession, + input: CodeExecuteInput, + signal?: AbortSignal, + ): Promise { + throwIfAborted(signal); + const runtime = (session.runtime ??= this.options.runtimeFactory(session.id)); + session.state = "running"; + try { + const result = await runtime.run( + input.code!, + requirePositiveTimeout(input.timeout_ms ?? this.defaultTimeoutMs, "timeout_ms"), + signal, + ); + session.state = "ready"; + session.page = result.page; + return { + ok: true, + action: "run", + code_session_id: session.id, + state: session.state, + page: result.page, + ...(result.value === undefined ? {} : { value: result.value }), + ...(result.logs.length === 0 ? {} : { logs: result.logs }), + }; + } catch (error) { + session.state = "idle"; + throw error; + } + } + + private async status(session: ManagedSession, signal?: AbortSignal): Promise { + throwIfAborted(signal); + let runtimeStatus: RuntimeStatus | undefined; + if (session.runtime) { + runtimeStatus = await session.runtime.status(signal); + session.state = runtimeStatus.state; + session.page = runtimeStatus.page; + } + return { + ok: true, + action: "status", + code_session_id: session.id, + state: session.state, + ...(session.page ? { page: session.page } : {}), + }; + } + + private async reset(session: ManagedSession, signal?: AbortSignal): Promise { + throwIfAborted(signal); + await session.runtime?.reset(signal); + session.state = "idle"; + session.page = undefined; + return { + ok: true, + action: "reset", + code_session_id: session.id, + state: "idle", + }; + } + + private async close(session: ManagedSession): Promise { + try { + await session.runtime?.close(); + } finally { + session.state = "idle"; + session.page = undefined; + this.sessions.delete(session.id); + } + return { + ok: true, + action: "close", + code_session_id: session.id, + state: "closed", + }; + } + + private async enqueue( + session: ManagedSession, + input: CodeExecuteInput, + signal: AbortSignal | undefined, + operation: () => Promise, + ): Promise { + const pending = session.queue.then(async () => { + try { + throwIfAborted(signal); + return await operation(); + } catch (error) { + return failureFromError(input, session, error); + } + }); + session.queue = pending.then( + () => undefined, + () => undefined, + ); + return pending; + } +} + +function validateInput(input: CodeExecuteInput): string | undefined { + if (input.action === "run" && (!input.code || input.code.trim().length === 0)) { + return "action=run requires a non-empty code string."; + } + if ( + (input.action === "reset" || input.action === "close") && + input.code_session_id === undefined + ) { + return `action=${input.action} requires code_session_id.`; + } + if (input.code_session_id !== undefined && input.code_session_id.trim().length === 0) { + return "code_session_id must be a non-empty opaque identifier."; + } + if ( + input.timeout_ms !== undefined && + (!Number.isSafeInteger(input.timeout_ms) || input.timeout_ms <= 0) + ) { + return "timeout_ms must be a positive integer."; + } + return undefined; +} + +function sessionNotFound(input: CodeExecuteInput): CodeExecuteFailure { + return failure( + input, + "session_not_found", + `Code session ${input.code_session_id} was not found or is already closed.`, + ); +} + +function failure( + input: Pick, + kind: CodeExecuteFailure["error"]["kind"], + message: string, +): CodeExecuteFailure { + return { + ok: false, + action: input.action, + ...(input.code_session_id ? { code_session_id: input.code_session_id } : {}), + error: { + kind, + name: "CodeModeRuntimeError", + message, + retryable: kind === "timeout" || kind === "runtime", + }, + }; +} + +function failureFromError( + input: CodeExecuteInput, + session: ManagedSession, + error: unknown, +): CodeExecuteFailure { + const normalized = + error instanceof Error ? error : new Error(typeof error === "string" ? error : String(error)); + const kind = error instanceof CodeModeRuntimeError ? error.kind : "runtime"; + return { + ok: false, + action: input.action, + code_session_id: session.id, + state: session.state, + ...(session.page ? { page: session.page } : {}), + error: { + kind, + name: normalized.name, + message: normalized.message, + retryable: + error instanceof CodeModeRuntimeError + ? error.retryable + : kind === "runtime" || kind === "timeout", + ...(error instanceof CodeModeRuntimeError && error.mayHaveSideEffects + ? { may_have_side_effects: true } + : {}), + }, + }; +} + +function requirePositiveTimeout(value: number, label: string): number { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new Error(`${label} must be a positive integer.`); + } + return value; +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) { + throw new CodeModeRuntimeError("aborted", "Code execution was aborted.", false, { + cause: signal.reason, + }); + } +} + +export function codeExecuteResultText(result: CodeExecuteResult): string { + return JSON.stringify(result, null, 2); +} diff --git a/packages/codemode/src/types.ts b/packages/codemode/src/types.ts new file mode 100644 index 0000000000..4770ce5e88 --- /dev/null +++ b/packages/codemode/src/types.ts @@ -0,0 +1,103 @@ +export const CODE_EXECUTE_ACTIONS = ["run", "status", "reset", "close"] as const; + +export type CodeExecuteAction = (typeof CODE_EXECUTE_ACTIONS)[number]; + +export type CodeExecuteInput = { + action: CodeExecuteAction; + code_session_id?: string; + code?: string; + timeout_ms?: number; +}; + +export type CodeSessionState = "idle" | "running" | "ready" | "closed"; + +export type CodePageState = { + url: string; + title: string; +}; + +export type CodeLogEntry = { + level: "log" | "warn" | "error"; + text: string; +}; + +export type CodeExecuteErrorKind = + | "validation" + | "session_not_found" + | "runtime" + | "timeout" + | "aborted" + | "closed"; + +export type CodeExecuteError = { + kind: CodeExecuteErrorKind; + name: string; + message: string; + retryable: boolean; + may_have_side_effects?: boolean; +}; + +export type CodeExecuteSuccess = { + ok: true; + action: CodeExecuteAction; + code_session_id?: string; + state: CodeSessionState; + page?: CodePageState; + value?: unknown; + logs?: CodeLogEntry[]; + active_code_sessions?: number; +}; + +export type CodeExecuteFailure = { + ok: false; + action: CodeExecuteAction; + code_session_id?: string; + state?: CodeSessionState; + page?: CodePageState; + error: CodeExecuteError; +}; + +export type CodeExecuteResult = CodeExecuteSuccess | CodeExecuteFailure; + +export type RuntimeRunResult = { + value?: unknown; + logs: CodeLogEntry[]; + page: CodePageState; +}; + +export type RuntimeStatus = { + state: Exclude; + page?: CodePageState; +}; + +export interface CodeRuntime { + run(code: string, timeoutMs: number, signal?: AbortSignal): Promise; + status(signal?: AbortSignal): Promise; + reset(signal?: AbortSignal): Promise; + close(): Promise; +} + +export type StagehandCodeRuntimeConfig = { + browserbaseApiKey?: string; + model?: { + modelName: string; + apiKey?: string; + baseURL?: string; + }; + defaultTimeoutMs?: number; +}; + +export class CodeModeRuntimeError extends Error { + readonly mayHaveSideEffects: boolean; + + constructor( + readonly kind: CodeExecuteErrorKind, + message: string, + readonly retryable = false, + options?: ErrorOptions & { mayHaveSideEffects?: boolean }, + ) { + super(message, options); + this.name = "CodeModeRuntimeError"; + this.mayHaveSideEffects = options?.mayHaveSideEffects ?? false; + } +} diff --git a/packages/codemode/tests/child-runtime.test.ts b/packages/codemode/tests/child-runtime.test.ts new file mode 100644 index 0000000000..59c5c35c90 --- /dev/null +++ b/packages/codemode/tests/child-runtime.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from "vitest"; +import { StagehandChildRuntime } from "../src/child-runtime.js"; + +describe("StagehandChildRuntime", () => { + it("hard-stops synchronous generated code from the parent and can recover", async () => { + const runtime = new StagehandChildRuntime( + "code_watchdog", + {}, + { + childModuleUrl: new URL("./fixtures/hung-runtime-child.mjs", import.meta.url), + }, + ); + const startedAt = Date.now(); + + await expect(runtime.run("hang", 25)).rejects.toMatchObject({ + kind: "timeout", + retryable: false, + mayHaveSideEffects: true, + }); + expect(Date.now() - startedAt).toBeLessThan(3_000); + + await expect(runtime.run("recover", 1_000)).resolves.toMatchObject({ + value: "recovered", + page: { url: "about:blank" }, + }); + await runtime.close(); + }); + + it("awaits child-reported timeout cleanup before same-handle recovery", async () => { + const runtime = fixtureRuntime("code_child_timeout"); + + await expect(runtime.run("child-timeout", 1_000)).rejects.toMatchObject({ + kind: "timeout", + retryable: false, + mayHaveSideEffects: true, + }); + await expect(runtime.run("recover", 1_000)).resolves.toMatchObject({ + value: "recovered", + }); + await runtime.close(); + }); + + it("awaits unexpected-exit cleanup before spawning a replacement child", async () => { + const runtime = fixtureRuntime("code_child_crash"); + + await expect(runtime.run("crash", 1_000)).rejects.toMatchObject({ + kind: "runtime", + }); + await expect(runtime.run("recover", 1_000)).resolves.toMatchObject({ + value: "recovered", + }); + await runtime.close(); + }); +}); + +function fixtureRuntime(codeSessionId: string): StagehandChildRuntime { + return new StagehandChildRuntime( + codeSessionId, + {}, + { + childModuleUrl: new URL("./fixtures/hung-runtime-child.mjs", import.meta.url), + }, + ); +} diff --git a/packages/codemode/tests/fixtures/hung-runtime-child.mjs b/packages/codemode/tests/fixtures/hung-runtime-child.mjs new file mode 100644 index 0000000000..49a35c6c69 --- /dev/null +++ b/packages/codemode/tests/fixtures/hung-runtime-child.mjs @@ -0,0 +1,46 @@ +process.on("message", (request) => { + if (request.type === "configure") { + process.send?.({ id: request.id, ok: true }); + return; + } + if (request.type === "run" && request.code === "hang") { + while (true) { + // Deliberately block this process's event loop. The parent watchdog must + // remain able to terminate it. + } + } + if (request.type === "run" && request.code === "child-timeout") { + process.send?.({ + id: request.id, + ok: false, + error: { + name: "CodeExecutionTimeoutError", + message: "child timer fired", + kind: "timeout", + retryable: false, + mayHaveSideEffects: true, + }, + }); + setTimeout(() => process.exit(1), 1_000); + return; + } + if (request.type === "run" && request.code === "crash") { + process.exit(1); + } + if (request.type === "run") { + process.send?.({ + id: request.id, + ok: true, + result: { + value: "recovered", + logs: [], + page: { url: "about:blank", title: "" }, + }, + }); + return; + } + if (request.type === "close") { + process.send?.({ id: request.id, ok: true, result: { closed: true } }); + setImmediate(() => process.exit(0)); + } +}); diff --git a/packages/codemode/tests/live-smoke.mjs b/packages/codemode/tests/live-smoke.mjs new file mode 100644 index 0000000000..c8ccda0d7d --- /dev/null +++ b/packages/codemode/tests/live-smoke.mjs @@ -0,0 +1,185 @@ +import assert from "node:assert/strict"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; +import { + CodeSessionManager, + createStagehandChildRuntime, + runtimeConfigFromEnv, + startCodeModeHttpServer, +} from "../dist/index.mjs"; + +if (!process.env.BROWSERBASE_API_KEY) { + throw new Error("BROWSERBASE_API_KEY is required for the live code-mode smoke."); +} + +const runtimeConfig = runtimeConfigFromEnv(); +const manager = new CodeSessionManager({ + runtimeFactory: (codeSessionId) => createStagehandChildRuntime(codeSessionId, runtimeConfig), +}); +const server = await startCodeModeHttpServer({ + manager, + host: "127.0.0.1", + port: 0, +}); + +const connect = async (name) => { + const client = new Client({ name, version: "0.1.0" }); + await client.connect(new StreamableHTTPClientTransport(new URL(server.url))); + return client; +}; + +const parseResult = (result) => { + if (result.structuredContent && typeof result.structuredContent === "object") { + return result.structuredContent; + } + const text = result.content?.find((item) => item.type === "text")?.text; + if (!text) throw new Error("code_execute returned no JSON result."); + return JSON.parse(text); +}; + +try { + const healthUrl = new URL("/health", server.url); + const initialHealth = await fetch(healthUrl).then((response) => response.json()); + assert.equal(initialHealth.browserProvisioning, "lazy"); + assert.equal(initialHealth.activeCodeSessions, 0); + + const firstClient = await connect("codemode-live-smoke-first"); + const tools = await firstClient.listTools(); + assert.deepEqual( + tools.tools.map((tool) => tool.name), + ["code_execute"], + ); + assert.equal(tools.tools[0]?.outputSchema?.type, "object"); + + const firstToolResult = await firstClient.callTool({ + name: "code_execute", + arguments: { + action: "run", + code: ` + await page.goto( + "https://browserbase.github.io/stagehand-eval-sites/sites/new-tab/", + { waitUntil: "load" } + ); + return { + phase: "opened", + title: await page.title(), + url: await page.url(), + pageCount: (await context.pages()).length + }; + `, + }, + }); + assert.equal(typeof firstToolResult.structuredContent, "object"); + const first = parseResult(firstToolResult); + assert.equal(first.ok, true); + assert.equal(first.value.phase, "opened"); + assert.match(first.value.url, /stagehand-eval-sites\/sites\/new-tab/); + + // Destroy the MCP transport. The logical code session and remote browser must + // outlive this connection. + await firstClient.close(); + const disconnectedHealth = await fetch(healthUrl).then((response) => response.json()); + assert.equal(disconnectedHealth.activeCodeSessions, 1); + + const secondClient = await connect("codemode-live-smoke-second"); + const second = parseResult( + await secondClient.callTool({ + name: "code_execute", + arguments: { + action: "run", + code_session_id: first.code_session_id, + code: ` + return { + phase: "reused", + title: await page.title(), + url: await page.url(), + bodyIncludesWelcome: (await page.locator("body").innerText()).includes("Welcome"), + pageCount: (await context.pages()).length + }; + `, + }, + }), + ); + assert.equal(second.ok, true); + assert.equal(second.code_session_id, first.code_session_id); + assert.equal(second.value.url, first.value.url); + assert.equal(second.value.bodyIncludesWelcome, true); + + const closed = parseResult( + await secondClient.callTool({ + name: "code_execute", + arguments: { + action: "close", + code_session_id: first.code_session_id, + }, + }), + ); + assert.equal(closed.state, "closed"); + + const finalHealth = await fetch(healthUrl).then((response) => response.json()); + assert.equal(finalHealth.activeCodeSessions, 0); + + const timedOut = parseResult( + await secondClient.callTool({ + name: "code_execute", + arguments: { + action: "run", + code: "while (true) {}", + timeout_ms: 10_000, + }, + }), + ); + assert.equal(timedOut.ok, false); + assert.equal(timedOut.error.kind, "timeout"); + assert.equal(timedOut.error.retryable, false); + assert.equal(timedOut.error.may_have_side_effects, true); + + const recovered = parseResult( + await secondClient.callTool({ + name: "code_execute", + arguments: { + action: "run", + code_session_id: timedOut.code_session_id, + code: ` + await page.goto("https://example.com", { waitUntil: "load" }); + return await page.title(); + `, + }, + }), + ); + assert.equal(recovered.ok, true); + assert.equal(recovered.code_session_id, timedOut.code_session_id); + assert.equal(recovered.value, "Example Domain"); + await secondClient.callTool({ + name: "code_execute", + arguments: { + action: "close", + code_session_id: timedOut.code_session_id, + }, + }); + await secondClient.close(); + + // oxlint-disable-next-line no-console + console.log( + JSON.stringify( + { + status: "PASS", + tools: tools.tools.map((tool) => tool.name), + browserProvisioning: initialHealth.browserProvisioning, + structuredOutput: true, + transportReconnect: true, + sameCodeSession: true, + preservedUrl: second.value.url, + preservedBodyState: second.value.bodyIncludesWelcome, + explicitClose: closed.state, + activeCodeSessionsAfterClose: finalHealth.activeCodeSessions, + synchronousLoopWatchdog: timedOut.error.kind, + sameHandleRecoveredAfterWatchdog: recovered.code_session_id === timedOut.code_session_id, + }, + null, + 2, + ), + ); +} finally { + await server.close(); +} diff --git a/packages/codemode/tests/mcp-server.test.ts b/packages/codemode/tests/mcp-server.test.ts new file mode 100644 index 0000000000..0d1fd89aa5 --- /dev/null +++ b/packages/codemode/tests/mcp-server.test.ts @@ -0,0 +1,140 @@ +import http from "node:http"; +import { describe, expect, it, vi } from "vitest"; +import { startCodeModeHttpServer } from "../src/mcp-server.js"; +import { CodeSessionManager } from "../src/session-manager.js"; +import type { CodeRuntime } from "../src/types.js"; + +function managerThatMustStayIdle(): CodeSessionManager { + return new CodeSessionManager({ + runtimeFactory: () => { + throw new Error("The HTTP safety test must not create a browser runtime."); + }, + }); +} + +describe("code-mode HTTP server", () => { + it("refuses an unauthenticated non-loopback bind", async () => { + await expect( + startCodeModeHttpServer({ + manager: managerThatMustStayIdle(), + host: "0.0.0.0", + port: 0, + }), + ).rejects.toThrow(/bearer_token is required/i); + }); + + it("allows an unauthenticated loopback server for local development", async () => { + const server = await startCodeModeHttpServer({ + manager: managerThatMustStayIdle(), + host: "127.0.0.1", + port: 0, + }); + try { + const health = await fetch(new URL("/health", server.url)).then((response) => + response.json(), + ); + expect(health).toMatchObject({ + ok: true, + browserProvisioning: "lazy", + activeCodeSessions: 0, + }); + } finally { + await server.close(); + } + }); + + it("rejects hostile Host and Origin headers on unauthenticated loopback", async () => { + const server = await startCodeModeHttpServer({ + manager: managerThatMustStayIdle(), + host: "127.0.0.1", + port: 0, + }); + try { + const hostileHost = await request(server.url, { host: "attacker.example" }); + expect(hostileHost.status).toBe(403); + + const hostileOrigin = await request(server.url, { + origin: "https://attacker.example", + }); + expect(hostileOrigin.status).toBe(403); + } finally { + await server.close(); + } + }); + + it("protects health metadata when bearer authentication is configured", async () => { + const server = await startCodeModeHttpServer({ + manager: managerThatMustStayIdle(), + host: "127.0.0.1", + port: 0, + bearerToken: "test-token", + }); + try { + const healthUrl = new URL("/health", server.url); + expect((await request(healthUrl)).status).toBe(401); + const authorized = await request(healthUrl, { + authorization: "Bearer test-token", + }); + expect(authorized.status).toBe(200); + } finally { + await server.close(); + } + }); + + it("closes every code session when the HTTP server stops", async () => { + const runtime = fakeRuntime(); + const close = vi.spyOn(runtime, "close"); + const manager = new CodeSessionManager({ + runtimeFactory: () => runtime, + sessionIdFactory: () => "code_server_shutdown", + }); + await manager.execute({ action: "run", code: "return true" }); + const server = await startCodeModeHttpServer({ + manager, + host: "127.0.0.1", + port: 0, + }); + + await server.close(); + + expect(close).toHaveBeenCalledOnce(); + expect(manager.activeSessionCount).toBe(0); + }); +}); + +function fakeRuntime(): CodeRuntime { + return { + async run() { + return { + value: true, + logs: [], + page: { url: "about:blank", title: "" }, + }; + }, + async status() { + return { state: "ready" }; + }, + async reset() {}, + async close() {}, + }; +} + +function request( + target: string | URL, + headers: Record = {}, +): Promise<{ status: number; body: string }> { + return new Promise((resolve, reject) => { + const outgoing = http.request(target, { headers }, (response) => { + const chunks: Buffer[] = []; + response.on("data", (chunk: Buffer) => chunks.push(chunk)); + response.once("end", () => { + resolve({ + status: response.statusCode ?? 0, + body: Buffer.concat(chunks).toString("utf8"), + }); + }); + }); + outgoing.once("error", reject); + outgoing.end(); + }); +} diff --git a/packages/codemode/tests/session-manager.test.ts b/packages/codemode/tests/session-manager.test.ts new file mode 100644 index 0000000000..9f5010cb1e --- /dev/null +++ b/packages/codemode/tests/session-manager.test.ts @@ -0,0 +1,149 @@ +import { describe, expect, it, vi } from "vitest"; +import { CodeSessionManager } from "../src/session-manager.js"; +import { CodeModeRuntimeError } from "../src/types.js"; +import type { CodeRuntime } from "../src/types.js"; + +function fakeRuntime(): CodeRuntime { + let url = "about:blank"; + return { + async run(code) { + if (code.includes("navigate")) url = "https://example.com/"; + return { + value: { code, url }, + logs: [], + page: { url, title: url === "about:blank" ? "" : "Example Domain" }, + }; + }, + async status() { + return { state: "ready", page: { url, title: "" } }; + }, + async reset() { + url = "about:blank"; + }, + async close() {}, + }; +} + +describe("CodeSessionManager", () => { + it("does not create a runtime until the first run", async () => { + const runtimeFactory = vi.fn(fakeRuntime); + const manager = new CodeSessionManager({ runtimeFactory }); + + await expect(manager.execute({ action: "status" })).resolves.toMatchObject({ + ok: true, + state: "idle", + active_code_sessions: 0, + }); + expect(runtimeFactory).not.toHaveBeenCalled(); + }); + + it("reuses one long-lived runtime for calls with the same opaque session ID", async () => { + const runtimeFactory = vi.fn(fakeRuntime); + const manager = new CodeSessionManager({ + runtimeFactory, + sessionIdFactory: () => "code_test", + }); + + const first = await manager.execute({ action: "run", code: "navigate" }); + expect(first).toMatchObject({ + ok: true, + code_session_id: "code_test", + value: { url: "https://example.com/" }, + }); + const second = await manager.execute({ + action: "run", + code_session_id: "code_test", + code: "read", + }); + expect(second).toMatchObject({ + ok: true, + code_session_id: "code_test", + value: { url: "https://example.com/" }, + }); + expect(runtimeFactory).toHaveBeenCalledTimes(1); + }); + + it("closes and removes a session only when explicitly requested", async () => { + const runtime = fakeRuntime(); + const close = vi.spyOn(runtime, "close"); + const manager = new CodeSessionManager({ + runtimeFactory: () => runtime, + sessionIdFactory: () => "code_test", + }); + await manager.execute({ action: "run", code: "navigate" }); + await manager.execute({ + action: "status", + code_session_id: "code_test", + }); + expect(close).not.toHaveBeenCalled(); + + await expect( + manager.execute({ action: "close", code_session_id: "code_test" }), + ).resolves.toMatchObject({ ok: true, state: "closed" }); + expect(close).toHaveBeenCalledOnce(); + expect(manager.activeSessionCount).toBe(0); + }); + + it("does not expose runtime stack traces in model-visible failures", async () => { + const manager = new CodeSessionManager({ + runtimeFactory: () => ({ + ...fakeRuntime(), + async run() { + throw new Error("synthetic failure"); + }, + }), + sessionIdFactory: () => "code_test", + }); + + const result = await manager.execute({ action: "run", code: "throw" }); + expect(result).toMatchObject({ + ok: false, + error: { message: "synthetic failure" }, + }); + if (result.ok) throw new Error("Expected a failure result."); + expect(result.error).not.toHaveProperty("stack"); + }); + + it("marks uncertain timeout failures as unsafe to retry", async () => { + const manager = new CodeSessionManager({ + runtimeFactory: () => ({ + ...fakeRuntime(), + async run() { + throw new CodeModeRuntimeError("timeout", "cell timed out", false, { + mayHaveSideEffects: true, + }); + }, + }), + sessionIdFactory: () => "code_test", + }); + + const result = await manager.execute({ action: "run", code: "slow mutation" }); + expect(result).toMatchObject({ + ok: false, + error: { + kind: "timeout", + retryable: false, + may_have_side_effects: true, + }, + }); + }); + + it("does not leak a new logical session when its request is already aborted", async () => { + const runtimeFactory = vi.fn(fakeRuntime); + const manager = new CodeSessionManager({ + runtimeFactory, + sessionIdFactory: () => "code_aborted", + }); + const controller = new AbortController(); + controller.abort("caller disconnected"); + + const result = await manager.execute({ action: "run", code: "return true" }, controller.signal); + + expect(result).toMatchObject({ + ok: false, + error: { kind: "aborted", retryable: false }, + }); + expect(runtimeFactory).not.toHaveBeenCalled(); + expect(manager.activeSessionCount).toBe(0); + }); +}); diff --git a/packages/codemode/tsconfig.json b/packages/codemode/tsconfig.json new file mode 100644 index 0000000000..80589e8c80 --- /dev/null +++ b/packages/codemode/tsconfig.json @@ -0,0 +1,12 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "noImplicitAny": true, + "skipLibCheck": true, + "strict": true, + "target": "ES2022", + "types": ["node"] + }, + "include": ["src/**/*.ts", "tests/**/*.ts"], + "exclude": ["dist", "node_modules"] +} diff --git a/packages/codemode/tsdown.config.ts b/packages/codemode/tsdown.config.ts new file mode 100644 index 0000000000..9f3a9f7007 --- /dev/null +++ b/packages/codemode/tsdown.config.ts @@ -0,0 +1,13 @@ +import { defineConfig } from "tsdown"; + +export default defineConfig({ + entry: ["src/index.ts", "src/cli.ts", "src/runtime-child.ts"], + format: ["esm"], + platform: "node", + target: "node22", + dts: { + sourcemap: true, + }, + sourcemap: true, + outDir: "dist", +}); diff --git a/packages/codemode/vitest.config.ts b/packages/codemode/vitest.config.ts new file mode 100644 index 0000000000..19384e80fd --- /dev/null +++ b/packages/codemode/vitest.config.ts @@ -0,0 +1,7 @@ +import { defineConfig } from "vitest/config"; + +export default defineConfig({ + test: { + include: ["tests/**/*.test.ts"], + }, +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b32293f560..65d86c3083 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -244,6 +244,9 @@ catalogs: '@mdx-js/mdx': specifier: 3.1.1 version: 3.1.1 + '@modelcontextprotocol/sdk': + specifier: 1.29.0 + version: 1.29.0 '@opentelemetry/api': specifier: 1.9.1 version: 1.9.1 @@ -409,6 +412,37 @@ importers: specifier: 'catalog:' version: 4.1.9(@opentelemetry/api@1.9.1)(@types/node@24.13.2)(vite@8.1.3(@types/node@24.13.2)(esbuild@0.28.1)(jiti@1.21.7)(tsx@4.23.1)(yaml@2.9.0)) + packages/codemode: + dependencies: + '@browserbasehq/sdk': + specifier: 'catalog:' + version: 2.16.0 + '@browserbasehq/stagehand': + specifier: workspace:* + version: link:../sdk-ts + '@modelcontextprotocol/sdk': + specifier: 'catalog:' + version: 1.29.0(zod@4.4.3) + zod: + specifier: 'catalog:' + version: 4.4.3 + devDependencies: + '@types/node': + specifier: 'catalog:' + version: 24.13.2 + tsdown: + specifier: 'catalog:' + version: 0.22.3(publint@0.3.21)(tsx@4.23.1)(typescript@5.9.3) + tsx: + specifier: 'catalog:' + version: 4.23.1 + typescript: + specifier: 'catalog:' + version: 5.9.3 + vitest: + specifier: 'catalog:' + version: 4.1.9(@opentelemetry/api@1.9.1)(@types/node@24.13.2)(vite@8.1.3(@types/node@24.13.2)(esbuild@0.28.1)(jiti@1.21.7)(tsx@4.23.1)(yaml@2.9.0)) + packages/docs: devDependencies: '@mdx-js/mdx': diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 90e5e925a0..1ffc7f1bd4 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -4,6 +4,7 @@ packages: catalogMode: prefer catalog: + "@modelcontextprotocol/sdk": 1.29.0 "@ast-grep/lang-go": 0.0.6 "@ast-grep/lang-python": 0.0.6 "@ast-grep/napi": 0.44.1 diff --git a/turbo.json b/turbo.json index 1841a2c24e..03e51917e7 100644 --- a/turbo.json +++ b/turbo.json @@ -31,6 +31,11 @@ "inputs": ["$TURBO_DEFAULT$", "!dist/**"], "outputs": ["dist/**"] }, + "@browserbasehq/stagehand-codemode#build": { + "dependsOn": ["^build"], + "inputs": ["$TURBO_DEFAULT$", "!dist/**"], + "outputs": ["dist/**"] + }, "//#generate:python": { "dependsOn": ["@browserbasehq/stagehand-protocol#build"], "inputs": ["$TURBO_DEFAULT$", "!packages/sdk-python/src/stagehand/_generated/**"], @@ -61,6 +66,9 @@ "@browserbasehq/stagehand-evals#typecheck": { "dependsOn": ["^build"] }, + "@browserbasehq/stagehand-codemode#typecheck": { + "dependsOn": ["^build"] + }, "@browserbasehq/stagehand-docs#typecheck": {}, "test:unit": { "dependsOn": ["^build"],