fix: verify and restore browser state after recording starts

src/browser/session.test.ts

@@ -1,5 +1,17 @@
-import { describe, expect, it } from 'vitest';
-import { buildOpenBrowserCommand } from './session.js';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { buildOpenBrowserCommand, getViewport, urlsMatch, verifyBrowserState } from './session.js';
+
+const { abMock } = vi.hoisted(() => ({
+  abMock: vi.fn(),
+}));
+
+vi.mock('../utils/exec.js', async () => {
+  const actual = await vi.importActual<typeof import('../utils/exec.js')>('../utils/exec.js');
+  return {
+    ...actual,
+    ab: abMock,
+  };
+});
 
 describe('buildOpenBrowserCommand', () => {
   it('builds a default open command without extra flags', () => {
@@ -23,3 +35,74 @@ describe('buildOpenBrowserCommand', () => {
     );
   });
 });
+
+describe('urlsMatch', () => {
+  it('treats trailing slashes as equivalent', () => {
+    expect(urlsMatch('http://localhost:3000', 'http://localhost:3000/')).toBe(true);
+  });
+
+  it('keeps query strings significant', () => {
+    expect(urlsMatch('http://localhost:3000/users?tab=open', 'http://localhost:3000/users?tab=closed')).toBe(
+      false,
+    );
+  });
+});
+
+describe('getViewport', () => {
+  beforeEach(() => {
+    abMock.mockReset();
+  });
+
+  it('returns viewport dimensions from the page context', () => {
+    abMock.mockReturnValueOnce(JSON.stringify({ width: 1280, height: 720 }));
+
+    expect(getViewport()).toEqual({ width: 1280, height: 720 });
+  });
+
+  it('parses double-encoded viewport payloads from agent-browser eval', () => {
+    abMock.mockReturnValueOnce(JSON.stringify(JSON.stringify({ width: 1920, height: 1080 })));
+
+    expect(getViewport()).toEqual({ width: 1920, height: 1080 });
+    expect(abMock).toHaveBeenCalledWith(
+      "eval 'JSON.stringify({width: window.innerWidth, height: window.innerHeight})'",
+      { session: undefined },
+    );
+  });
+});
+
+describe('verifyBrowserState', () => {
+  beforeEach(() => {
+    abMock.mockReset();
+  });
+
+  it('returns the observed state when URL and viewport match', () => {
+    abMock
+      .mockReturnValueOnce('http://localhost:3000/')
+      .mockReturnValueOnce(JSON.stringify({ width: 1280, height: 720 }));
+
+    expect(verifyBrowserState('http://localhost:3000', { width: 1280, height: 720 })).toEqual({
+      url: 'http://localhost:3000/',
+      viewport: { width: 1280, height: 720 },
+    });
+  });
+
+  it('throws a targeted error when the browser is on the wrong page', () => {
+    abMock
+      .mockReturnValueOnce('http://localhost:3000/about')
+      .mockReturnValueOnce(JSON.stringify({ width: 1280, height: 720 }));
+
+    expect(() => verifyBrowserState('http://localhost:3000', { width: 1280, height: 720 })).toThrow(
+      /attached to the wrong page or session/i,
+    );
+  });
+
+  it('throws a targeted error when the viewport does not match', () => {
+    abMock
+      .mockReturnValueOnce('http://localhost:3000/')
+      .mockReturnValueOnce(JSON.stringify({ width: 1600, height: 914 }));
+
+    expect(() => verifyBrowserState('http://localhost:3000', { width: 1280, height: 720 })).toThrow(
+      /viewport is 1600x914, expected 1280x720/i,
+    );
+  });
+});

src/browser/session.ts

@@ -16,10 +16,11 @@ export function buildOpenBrowserCommand(
   return `open ${url}${suffix}`;
 }
 
-/**
- * Initialize a browser session.
- * Opens the browser and sets viewport dimensions.
- */
+export interface BrowserState {
+  url: string;
+  viewport: ViewportConfig | null;
+}
+
 export function openBrowser(
   url: string,
   viewport: ViewportConfig,
@@ -28,12 +29,15 @@ export function openBrowser(
   browserConfig?: BrowserConfig,
 ): void {
   ab(buildOpenBrowserCommand(url, headless, browserConfig), { timeoutMs: 60000, session: sessionName });
-  ab(`set viewport ${viewport.width} ${viewport.height}`, { session: sessionName });
+  applyViewport(viewport, sessionName);
+}
+
+export function applyViewport(viewport: ViewportConfig, sessionName?: string): void {
+  const scale = viewport.deviceScaleFactor;
+  const suffix = typeof scale === 'number' && Number.isFinite(scale) ? ` ${scale}` : '';
+  ab(`set viewport ${viewport.width} ${viewport.height}${suffix}`, { session: sessionName });
 }
 
-/**
- * Close the browser session.
- */
 export function closeBrowser(sessionName?: string): void {
   try {
     ab('close', { session: sessionName });
@@ -42,9 +46,6 @@ export function closeBrowser(sessionName?: string): void {
   }
 }
 
-/**
- * Check if agent-browser is installed and accessible.
- */
 export function checkAgentBrowser(): boolean {
   try {
     ab('--version', 5000);
@@ -54,9 +55,6 @@ export function checkAgentBrowser(): boolean {
   }
 }
 
-/**
- * Get any console errors from the current page.
- */
 export function getConsoleErrors(sessionName?: string): string {
   try {
     return ab('errors', { session: sessionName });
@@ -65,9 +63,6 @@ export function getConsoleErrors(sessionName?: string): string {
   }
 }
 
-/**
- * Get console output from the current page.
- */
 export function getConsoleOutput(sessionName?: string): string {
   try {
     return ab('console', { session: sessionName });
@@ -78,28 +73,21 @@ export function getConsoleOutput(sessionName?: string): string {
 
 export interface ConsoleMessage {
   text: string;
-  timestamp: number; // epoch ms
-  type: string; // log, warn, error, etc.
+  timestamp: number;
+  type: string;
 }
 
-/**
- * Get console output as structured JSON with per-message timestamps.
- */
 export function getConsoleOutputJson(sessionName?: string): ConsoleMessage[] {
   try {
     const raw = ab('console --json', { session: sessionName });
     const parsed = JSON.parse(raw);
-    // agent-browser wraps JSON output: {success, data: {messages: [...]}, error}
     const messages = parsed?.data?.messages ?? parsed;
     return Array.isArray(messages) ? messages : [];
   } catch {
     return [];
   }
 }
 
-/**
- * Get the current page title.
- */
 export function getPageTitle(sessionName?: string): string {
   try {
     return ab('get title', { session: sessionName });
@@ -108,13 +96,87 @@ export function getPageTitle(sessionName?: string): string {
   }
 }
 
-/**
- * Get the current page URL.
- */
 export function getPageUrl(sessionName?: string): string {
   try {
     return ab('get url', { session: sessionName });
   } catch {
     return '';
   }
 }
+
+export function getViewport(sessionName?: string): ViewportConfig | null {
+  try {
+    const raw = ab("eval 'JSON.stringify({width: window.innerWidth, height: window.innerHeight})'", {
+      session: sessionName,
+    });
+    const parsed = parseViewportPayload(raw);
+    if (
+      typeof parsed?.width === 'number' &&
+      Number.isFinite(parsed.width) &&
+      typeof parsed?.height === 'number' &&
+      Number.isFinite(parsed.height)
+    ) {
+      return { width: parsed.width, height: parsed.height };
+    }
+    return null;
+  } catch {
+    return null;
+  }
+}
+
+function parseViewportPayload(raw: string): unknown {
+  const parsed = JSON.parse(raw);
+  if (typeof parsed === 'string') {
+    return JSON.parse(parsed);
+  }
+  return parsed;
+}
+
+function normalizeUrlForComparison(value: string): string {
+  try {
+    const url = new URL(value);
+    const pathname = url.pathname !== '/' ? url.pathname.replace(/\/+$/, '') : '/';
+    return `${url.origin}${pathname}${url.search}`;
+  } catch {
+    return value.trim();
+  }
+}
+
+export function urlsMatch(expectedUrl: string, actualUrl: string): boolean {
+  return normalizeUrlForComparison(expectedUrl) === normalizeUrlForComparison(actualUrl);
+}
+
+export function verifyBrowserState(
+  expectedUrl: string,
+  expectedViewport: ViewportConfig,
+  sessionName?: string,
+): BrowserState {
+  const url = getPageUrl(sessionName);
+  const viewport = getViewport(sessionName);
+
+  if (!url) {
+    throw new ProofShotError(
+      'Could not read the current browser URL after recording started. The browser session may not be attached correctly.',
+    );
+  }
+
+  if (!urlsMatch(expectedUrl, url)) {
+    throw new ProofShotError(
+      `Browser navigated to ${url}, expected ${expectedUrl}. Recording may be attached to the wrong page or session.`,
+    );
+  }
+
+  if (!viewport) {
+    throw new ProofShotError(
+      'Could not read the current viewport after recording started. The browser session may not be attached correctly.',
+    );
+  }
+
+  if (viewport.width !== expectedViewport.width || viewport.height !== expectedViewport.height) {
+    throw new ProofShotError(
+      `Browser viewport is ${viewport.width}x${viewport.height}, expected ${expectedViewport.width}x${expectedViewport.height}. Recording may be attached to the wrong page or session.`,
+    );
+  }
+
+  return { url, viewport };
+}

src/commands/start.test.ts

@@ -5,8 +5,11 @@ const mocks = vi.hoisted(() => ({
   loadConfig: vi.fn(),
   ensureDevServer: vi.fn(),
   openBrowser: vi.fn(),
+  applyViewport: vi.fn(),
   closeBrowser: vi.fn(),
+  verifyBrowserState: vi.fn(),
   startRecording: vi.fn(),
+  stopRecording: vi.fn(),
   ensureOutputDir: vi.fn(),
   generateTimestamp: vi.fn(),
   generateSessionDirName: vi.fn(),
@@ -28,11 +31,14 @@ vi.mock('../server/start.js', () => ({
 
 vi.mock('../browser/session.js', () => ({
   openBrowser: mocks.openBrowser,
+  applyViewport: mocks.applyViewport,
   closeBrowser: mocks.closeBrowser,
+  verifyBrowserState: mocks.verifyBrowserState,
 }));
 
 vi.mock('../browser/capture.js', () => ({
   startRecording: mocks.startRecording,
+  stopRecording: mocks.stopRecording,
 }));
 
 vi.mock('../artifacts/bundle.js', () => ({
@@ -92,20 +98,50 @@ describe('startCommand', () => {
     Object.values(mocks).forEach((mock) => mock.mockReset());
   });
 
-  it('closes the browser when recording never starts after all retries', async () => {
-    mocks.startRecording.mockImplementation(() => {
-      throw new Error('Recording session could not be initialized');
+  it('stops the active recording session and closes the browser when verification keeps failing', async () => {
+    mocks.verifyBrowserState.mockImplementation(() => {
+      throw new Error('Browser viewport is 800x535, expected 1280x720.');
     });
 
     const commandPromise = startCommand({}).catch((error) => error);
     await vi.runAllTimersAsync();
 
     await expect(commandPromise).resolves.toMatchObject({ message: 'process.exit:1' });
     expect(mocks.startRecording).toHaveBeenCalledTimes(3);
+    expect(mocks.stopRecording).toHaveBeenCalledTimes(3);
     expect(mocks.closeBrowser).toHaveBeenCalledTimes(1);
     expect(mocks.saveSession).not.toHaveBeenCalled();
   });
 
+  it('reapplies the configured viewport after recording starts before verification', async () => {
+    mocks.verifyBrowserState.mockReturnValue({
+      url: 'http://localhost:3000/',
+      viewport: { width: 1280, height: 720 },
+    });
+
+    await startCommand({});
+
+    expect(mocks.startRecording).toHaveBeenCalledWith(
+      expect.stringContaining('session.webm'),
+      'proofshot-2026-04-08_07-28-00',
+    );
+    expect(mocks.applyViewport).toHaveBeenCalledWith(
+      { width: 1280, height: 720 },
+      'proofshot-2026-04-08_07-28-00',
+    );
+    expect(mocks.verifyBrowserState).toHaveBeenCalledWith(
+      'http://localhost:3000',
+      { width: 1280, height: 720 },
+      'proofshot-2026-04-08_07-28-00',
+    );
+    expect(mocks.applyViewport.mock.invocationCallOrder[0]).toBeGreaterThan(
+      mocks.startRecording.mock.invocationCallOrder[0],
+    );
+    expect(mocks.verifyBrowserState.mock.invocationCallOrder[0]).toBeGreaterThan(
+      mocks.applyViewport.mock.invocationCallOrder[0],
+    );
+  });
+
   it('does not try to stop recording when recording never started', async () => {
     mocks.startRecording.mockImplementation(() => {
       throw new Error('Recording already active');
@@ -116,10 +152,11 @@ describe('startCommand', () => {
 
     await expect(commandPromise).resolves.toMatchObject({ message: 'process.exit:1' });
     expect(mocks.startRecording).toHaveBeenCalledTimes(3);
+    expect(mocks.stopRecording).not.toHaveBeenCalled();
     expect(mocks.closeBrowser).toHaveBeenCalledTimes(1);
   });
 
-  it('closes the session-scoped browser when browser open fails', async () => {
+  it('closes the browser when browser open fails', async () => {
     mocks.openBrowser.mockImplementation(() => {
       throw new Error('Chrome exited early without writing DevToolsActivePort');
     });