From 52b0dba7b0702eb2e4337585f8c5ce53bbb0ebe0 Mon Sep 17 00:00:00 2001
From: Jayesh Betala <jayesh.betala7@gmail.com>
Date: Mon, 11 May 2026 00:38:59 +0530
Subject: [PATCH] fix(browse): preserve prettyscreenshot path after hide

---
 browse/src/write-commands.ts              | 88 +++++++++++++++--------
 browse/test/prettyscreenshot-args.test.ts | 33 +++++++++
 2 files changed, 92 insertions(+), 29 deletions(-)
 create mode 100644 browse/test/prettyscreenshot-args.test.ts

diff --git a/browse/src/write-commands.ts b/browse/src/write-commands.ts
index 61c84d839..a7e2d7b9e 100644
--- a/browse/src/write-commands.ts
+++ b/browse/src/write-commands.ts
@@ -127,6 +127,62 @@ const CLEANUP_SELECTORS = {
   ],
 };
 
+export interface PrettyScreenshotOptions {
+  scrollTo?: string;
+  doCleanup: boolean;
+  hideSelectors: string[];
+  viewportWidth?: number;
+  outputPath?: string;
+}
+
+function isLikelyScreenshotOutputPath(arg: string): boolean {
+  return arg.startsWith('/')
+    || arg.startsWith('./')
+    || arg.startsWith('../')
+    || arg.startsWith('~')
+    || /\.(png|jpe?g|webp)$/i.test(arg);
+}
+
+export function parsePrettyScreenshotArgs(args: string[]): PrettyScreenshotOptions {
+  const options: PrettyScreenshotOptions = {
+    doCleanup: false,
+    hideSelectors: [],
+  };
+
+  for (let i = 0; i < args.length; i++) {
+    if (args[i] === '--scroll-to' && i + 1 < args.length) {
+      options.scrollTo = args[++i];
+    } else if (args[i] === '--cleanup') {
+      options.doCleanup = true;
+    } else if (args[i] === '--hide' && i + 1 < args.length) {
+      const values: string[] = [];
+      i++;
+      while (i < args.length && !args[i].startsWith('--')) {
+        values.push(args[i]);
+        i++;
+      }
+      i--; // Back up since the for loop will increment.
+
+      if (!options.outputPath && values.length > 1) {
+        const last = values[values.length - 1];
+        if (isLikelyScreenshotOutputPath(last)) {
+          options.outputPath = values.pop();
+        }
+      }
+      options.hideSelectors.push(...values);
+    } else if (args[i] === '--width' && i + 1 < args.length) {
+      options.viewportWidth = parseInt(args[++i], 10);
+      if (isNaN(options.viewportWidth)) throw new Error('--width must be a number');
+    } else if (!args[i].startsWith('--')) {
+      options.outputPath = args[i];
+    } else {
+      throw new Error(`Unknown prettyscreenshot flag: ${args[i]}`);
+    }
+  }
+
+  return options;
+}
+
 export async function handleWriteCommand(
   command: string,
   args: string[],
@@ -992,35 +1048,9 @@ export async function handleWriteCommand(
     }
 
     case 'prettyscreenshot': {
-      // Parse flags
-      let scrollTo: string | undefined;
-      let doCleanup = false;
-      const hideSelectors: string[] = [];
-      let viewportWidth: number | undefined;
-      let outputPath: string | undefined;
-
-      for (let i = 0; i < args.length; i++) {
-        if (args[i] === '--scroll-to' && i + 1 < args.length) {
-          scrollTo = args[++i];
-        } else if (args[i] === '--cleanup') {
-          doCleanup = true;
-        } else if (args[i] === '--hide' && i + 1 < args.length) {
-          // Collect all following non-flag args as selectors to hide
-          i++;
-          while (i < args.length && !args[i].startsWith('--')) {
-            hideSelectors.push(args[i]);
-            i++;
-          }
-          i--; // Back up since the for loop will increment
-        } else if (args[i] === '--width' && i + 1 < args.length) {
-          viewportWidth = parseInt(args[++i], 10);
-          if (isNaN(viewportWidth)) throw new Error('--width must be a number');
-        } else if (!args[i].startsWith('--')) {
-          outputPath = args[i];
-        } else {
-          throw new Error(`Unknown prettyscreenshot flag: ${args[i]}`);
-        }
-      }
+      const prettyOptions = parsePrettyScreenshotArgs(args);
+      const { scrollTo, doCleanup, hideSelectors, viewportWidth } = prettyOptions;
+      let { outputPath } = prettyOptions;
 
       // Default output path
       if (!outputPath) {
diff --git a/browse/test/prettyscreenshot-args.test.ts b/browse/test/prettyscreenshot-args.test.ts
new file mode 100644
index 000000000..d577f6f8d
--- /dev/null
+++ b/browse/test/prettyscreenshot-args.test.ts
@@ -0,0 +1,33 @@
+import { describe, expect, test } from 'bun:test';
+import { parsePrettyScreenshotArgs } from '../src/write-commands';
+
+describe('parsePrettyScreenshotArgs', () => {
+  test('treats trailing path after --hide as output path', () => {
+    const parsed = parsePrettyScreenshotArgs(['--hide', 'header', '/tmp/case-a.png']);
+
+    expect(parsed.hideSelectors).toEqual(['header']);
+    expect(parsed.outputPath).toBe('/tmp/case-a.png');
+  });
+
+  test('keeps multiple trailing --hide selectors when no output path is present', () => {
+    const parsed = parsePrettyScreenshotArgs(['--hide', 'header', 'footer']);
+
+    expect(parsed.hideSelectors).toEqual(['header', 'footer']);
+    expect(parsed.outputPath).toBeUndefined();
+  });
+
+  test('handles multiple --hide flags with explicit trailing output path', () => {
+    const parsed = parsePrettyScreenshotArgs(['--hide', 'header', '--hide', 'footer', './shot.webp']);
+
+    expect(parsed.hideSelectors).toEqual(['header', 'footer']);
+    expect(parsed.outputPath).toBe('./shot.webp');
+  });
+
+  test('still honors paths after a later flag', () => {
+    const parsed = parsePrettyScreenshotArgs(['--hide', 'header', '--cleanup', '/tmp/case-b.png']);
+
+    expect(parsed.hideSelectors).toEqual(['header']);
+    expect(parsed.doCleanup).toBe(true);
+    expect(parsed.outputPath).toBe('/tmp/case-b.png');
+  });
+});