diff --git a/src/core/OriginHeaders.ts b/src/core/OriginHeaders.ts index 36d912e..e9bc563 100644 --- a/src/core/OriginHeaders.ts +++ b/src/core/OriginHeaders.ts @@ -14,6 +14,11 @@ export interface OriginHeaderConfig { type RouteHandler = (route: Route) => Promise; +interface SessionRouteRegistration { + handler: RouteHandler; + onClose: () => void; +} + /** * Manages per-origin HTTP headers and installs request interception on a Playwright page. * @@ -31,7 +36,7 @@ export class OriginHeaders { private installedPage: Page | null = null; private routeHandler: RouteHandler | null = null; private rollbackRoutes: Array<{ page: Page; handler: RouteHandler }> = []; - private readonly sessionRoutes = new Map(); + private readonly sessionRoutes = new Map(); constructor(config?: OriginHeaderConfig) { if (config) { @@ -101,9 +106,28 @@ export class OriginHeaders { if (this.sessionRoutes.has(page)) return; const handler = this.createRouteHandler(); - const routePromise = page.route("**/*", handler); - if (routePromise) await routePromise; - this.sessionRoutes.set(page, handler); + const onClose = () => { + this.sessionRoutes.delete(page); + }; + page.once("close", onClose); + + try { + const routePromise = page.route("**/*", handler); + if (routePromise) await routePromise; + } catch (error) { + page.off("close", onClose); + throw error; + } + + this.sessionRoutes.set(page, { handler, onClose }); + + // The page can close while route registration is awaiting Playwright. + // A close event that fires before ownership is recorded deletes nothing, + // so reconcile against the authoritative page state before returning. + if (page.isClosed()) { + page.off("close", onClose); + this.sessionRoutes.delete(page); + } } private createRouteHandler(): RouteHandler { @@ -170,7 +194,8 @@ export class OriginHeaders { const sessionRoutes = [...this.sessionRoutes]; this.sessionRoutes.clear(); - for (const [page, handler] of sessionRoutes) { + for (const [page, { handler, onClose }] of sessionRoutes) { + page.off("close", onClose); try { await page.unroute("**/*", handler); } catch { diff --git a/tests/unit/OriginHeadersSessionLifecycle.test.ts b/tests/unit/OriginHeadersSessionLifecycle.test.ts index 4035a73..d5dd4c9 100644 --- a/tests/unit/OriginHeadersSessionLifecycle.test.ts +++ b/tests/unit/OriginHeadersSessionLifecycle.test.ts @@ -2,9 +2,24 @@ import { describe, expect, it, vi } from "vitest"; import { OriginHeaders } from "../../src/core/OriginHeaders.js"; function createPage() { + let closed = false; + let closeHandler: (() => void) | null = null; return { route: vi.fn().mockResolvedValue(undefined), unroute: vi.fn().mockResolvedValue(undefined), + once: vi.fn((event: string, handler: () => void) => { + if (event === "close") closeHandler = handler; + }), + off: vi.fn((event: string, handler: () => void) => { + if (event === "close" && closeHandler === handler) closeHandler = null; + }), + isClosed: vi.fn(() => closed), + emitClose: () => { + closed = true; + const handler = closeHandler; + closeHandler = null; + handler?.(); + }, }; } @@ -22,15 +37,19 @@ describe("OriginHeaders session lifecycle", () => { expect(firstPage.route).toHaveBeenCalledTimes(1); expect(secondPage.route).toHaveBeenCalledTimes(1); + expect(firstPage.once).toHaveBeenCalledTimes(1); + expect(secondPage.once).toHaveBeenCalledTimes(1); expect(firstPage.unroute).not.toHaveBeenCalled(); expect(secondPage.unroute).not.toHaveBeenCalled(); await headers.dispose(); expect(firstPage.unroute).toHaveBeenCalledTimes(1); expect(secondPage.unroute).toHaveBeenCalledTimes(1); + expect(firstPage.off).toHaveBeenCalledTimes(1); + expect(secondPage.off).toHaveBeenCalledTimes(1); }); - it("does not claim ownership when a session route registration fails", async () => { + it("does not retain ownership or a close listener when route registration fails", async () => { const headers = new OriginHeaders({ "https://api.example.com": { Authorization: "Bearer secret" }, }); @@ -38,12 +57,49 @@ describe("OriginHeaders session lifecycle", () => { page.route.mockRejectedValue(new Error("route unavailable")); await expect(headers.installSessionPage(page as any)).rejects.toThrow("route unavailable"); + expect(page.once).toHaveBeenCalledTimes(1); + expect(page.off).toHaveBeenCalledTimes(1); await headers.dispose(); expect(page.unroute).not.toHaveBeenCalled(); }); - it("disposes every session route best-effort even when one page is already gone", async () => { + it("releases a naturally closed page without touching a surviving sibling", async () => { + const headers = new OriginHeaders({ + "https://api.example.com": { Authorization: "Bearer secret" }, + }); + const closedPage = createPage(); + const livePage = createPage(); + + await headers.installSessionPage(closedPage as any); + await headers.installSessionPage(livePage as any); + closedPage.emitClose(); + + expect(closedPage.unroute).not.toHaveBeenCalled(); + expect(livePage.unroute).not.toHaveBeenCalled(); + + await headers.dispose(); + expect(closedPage.unroute).not.toHaveBeenCalled(); + expect(livePage.unroute).toHaveBeenCalledTimes(1); + }); + + it("reconciles a page that closes while route registration is pending", async () => { + const headers = new OriginHeaders({ + "https://api.example.com": { Authorization: "Bearer secret" }, + }); + const page = createPage(); + page.route.mockImplementation(async () => { + page.emitClose(); + }); + + await headers.installSessionPage(page as any); + await headers.dispose(); + + expect(page.isClosed).toHaveBeenCalled(); + expect(page.unroute).not.toHaveBeenCalled(); + }); + + it("disposes every still-owned session route best-effort when a close event was missed", async () => { const headers = new OriginHeaders({ "https://api.example.com": { Authorization: "Bearer secret" }, }); diff --git a/tests/unit/OriginHeadersSessionRouting.test.ts b/tests/unit/OriginHeadersSessionRouting.test.ts index aada858..ef22067 100644 --- a/tests/unit/OriginHeadersSessionRouting.test.ts +++ b/tests/unit/OriginHeadersSessionRouting.test.ts @@ -35,17 +35,30 @@ const emptyState = { }; function createPage(order: string[]) { + let closed = false; + let closeHandler: (() => void) | null = null; return { route: vi.fn(async () => { order.push("route"); }), unroute: vi.fn().mockResolvedValue(undefined), on: vi.fn(), + once: vi.fn((event: string, handler: () => void) => { + if (event === "close") closeHandler = handler; + }), + off: vi.fn((event: string, handler: () => void) => { + if (event === "close" && closeHandler === handler) closeHandler = null; + }), goto: vi.fn(async () => { order.push("goto"); }), - close: vi.fn().mockResolvedValue(undefined), - isClosed: vi.fn(() => false), + close: vi.fn(async () => { + closed = true; + const handler = closeHandler; + closeHandler = null; + handler?.(); + }), + isClosed: vi.fn(() => closed), }; } @@ -123,9 +136,11 @@ describe("OriginHeaders session routing", () => { await controller.closePage(1); expect(firstPage.unroute).not.toHaveBeenCalled(); + expect(secondPage.unroute).not.toHaveBeenCalled(); expect(firstPage.route).toHaveBeenCalledTimes(1); await (controller as any).disposeOriginHeaders(); expect(firstPage.unroute).toHaveBeenCalledTimes(1); + expect(secondPage.unroute).not.toHaveBeenCalled(); }); it("installs origin headers on the recreated page before snapshot restoration navigation", async () => {