diff --git a/packages/runtime/src/index.test.ts b/packages/runtime/src/index.test.ts index 0b0c708..57eb9a9 100644 --- a/packages/runtime/src/index.test.ts +++ b/packages/runtime/src/index.test.ts @@ -329,4 +329,32 @@ describe("BrowserRuntime", () => { message: "other failure", }); }); + + it("rejects createSession after shutdown without relaunching the browser", async () => { + const { browser } = createMockBrowserTree(); + mockChromiumLaunch.mockResolvedValueOnce(browser); + + const rt = new BrowserRuntime({ headless: true }); + await rt.createSession(); + await rt.shutdown(); + + mockChromiumLaunch.mockClear(); + await expect(rt.createSession()).rejects.toMatchObject({ + code: "COMMAND_FAILED", + message: "Runtime is shutting down; cannot create session", + }); + expect(mockChromiumLaunch).not.toHaveBeenCalled(); + }); + + it("shutdown is idempotent", async () => { + const { browser, browserClose } = createMockBrowserTree(); + mockChromiumLaunch.mockResolvedValueOnce(browser); + + const rt = new BrowserRuntime({ headless: true }); + await rt.createSession(); + await rt.shutdown(); + await rt.shutdown(); + + expect(browserClose).toHaveBeenCalledTimes(1); + }); }); diff --git a/packages/runtime/src/index.ts b/packages/runtime/src/index.ts index 6e1c945..f21280a 100644 --- a/packages/runtime/src/index.ts +++ b/packages/runtime/src/index.ts @@ -33,9 +33,19 @@ type SessionRecord = { export class BrowserRuntime { private browserPromise: Promise | null = null; private sessions = new Map(); + private shuttingDown = false; constructor(private readonly options: BrowserRuntimeOptions = {}) {} + private assertActive(operation: string) { + if (this.shuttingDown) { + throw new WebchainRuntimeError( + "COMMAND_FAILED", + `Runtime is shutting down; cannot ${operation}`, + ); + } + } + private async getBrowser(): Promise { if (!this.browserPromise) { this.browserPromise = this.launchBrowser(); @@ -70,6 +80,7 @@ export class BrowserRuntime { } async createSession(): Promise { + this.assertActive("create session"); const browser = await this.getBrowser(); const context = await browser.newContext(); @@ -96,6 +107,7 @@ export class BrowserRuntime { } async navigate(command: NavigateCommand): Promise { + this.assertActive("navigate"); const session = this.getSession(command.sessionId); try { await session.page.goto(command.url, { waitUntil: "domcontentloaded" }); @@ -111,6 +123,7 @@ export class BrowserRuntime { } async snapshot(command: SnapshotCommand): Promise { + this.assertActive("snapshot"); const session = this.getSession(command.sessionId); try { const page = session.page; @@ -164,6 +177,7 @@ export class BrowserRuntime { } async click(command: ClickCommand): Promise { + this.assertActive("click"); const session = this.getSession(command.sessionId); try { await session.page.locator(command.selector).first().click(); @@ -179,6 +193,7 @@ export class BrowserRuntime { } async type(command: TypeCommand): Promise { + this.assertActive("type"); const session = this.getSession(command.sessionId); try { await session.page.locator(command.selector).first().fill(command.text); @@ -196,6 +211,7 @@ export class BrowserRuntime { async closeSession( command: CloseSessionCommand, ): Promise { + this.assertActive("close session"); const session = this.getSession(command.sessionId); try { await session.context.close(); @@ -211,6 +227,11 @@ export class BrowserRuntime { } async shutdown() { + if (this.shuttingDown) { + return; + } + this.shuttingDown = true; + await Promise.all( [...this.sessions.values()].map(async (session) => { await session.context.close(); diff --git a/services/companion/src/index.ts b/services/companion/src/index.ts index 1797d20..2f91ffa 100644 --- a/services/companion/src/index.ts +++ b/services/companion/src/index.ts @@ -9,9 +9,16 @@ const runtime = new BrowserRuntime({ const { app } = await createCompanionApp({ runtime }); +let closing = false; + const closeGracefully = async () => { - await runtime.shutdown(); + if (closing) { + return; + } + closing = true; + await app.close(); + await runtime.shutdown(); process.exit(0); }; diff --git a/services/mcp/src/test-support/bootstrap-mcp-stack.ts b/services/mcp/src/test-support/bootstrap-mcp-stack.ts index 951ad0d..98f00f8 100644 --- a/services/mcp/src/test-support/bootstrap-mcp-stack.ts +++ b/services/mcp/src/test-support/bootstrap-mcp-stack.ts @@ -80,8 +80,8 @@ export async function bootstrapMcpStack( shutdown: async () => { await client.close(); await stdioTransport.close(); - await runtime.shutdown(); await app.close(); + await runtime.shutdown(); }, }; }