refactor(utils): unify captureSnapshot + trace into lib/utils/trace.js

DavertMik · claude · DavertMik · commit eaca22dfb08a · 2026-04-26T17:37:29.000+03:00
Both files belonged to the same flow (capture artifacts -&gt; render trace.md).
Splitting them across two files made the import sites awkward, so they're
now consolidated:

- `lib/utils/captureSnapshot.js` deleted; its `captureSnapshot`,
  `normalizeBrowserLogs`, and `captureStorageState` moved into
  `lib/utils/trace.js`.
- `writeTraceMarkdown` (was inlined in `bin/mcp-server.js`) moved into
  `lib/utils/trace.js` too — it's the same trace-record concern.
- All three call sites (`bin/mcp-server.js`, `lib/plugin/aiTrace.js`,
  `lib/plugin/pageInfo.js`) now import from a single module.
- Stale `writeFileSync` import dropped from `bin/mcp-server.js`.

72 unit tests still passing.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/bin/mcp-server.js b/bin/mcp-server.js
@@ -5,21 +5,21 @@ import Codecept from '../lib/codecept.js'
 import container from '../lib/container.js'
 import { getParamsToString } from '../lib/parser.js'
 import { methodsOfObject, safeStringify, truncateString } from '../lib/utils.js'
-import { captureSnapshot } from '../lib/utils/captureSnapshot.js'
 import {
+  captureSnapshot,
   pickActingHelper,
   traceDirFor,
   snapshotDirFor,
-  artifactLinks,
   artifactsToFileUrls,
+  writeTraceMarkdown,
 } from '../lib/utils/trace.js'
 import event from '../lib/event.js'
 import { fileURLToPath } from 'url'
 import { dirname, resolve as resolvePath } from 'path'
 import path from 'path'
 import { spawn } from 'child_process'
 import { createRequire } from 'module'
-import { existsSync, readdirSync, writeFileSync } from 'fs'
+import { existsSync, readdirSync } from 'fs'
 import { mkdirp } from 'mkdirp'
 
 const require = createRequire(import.meta.url)
@@ -235,30 +235,6 @@ function outputBaseDir() {
   return global.output_dir || resolvePath(process.cwd(), 'output')
 }
 
-function writeTraceMarkdown({ dir, title, file, durationMs, commands, captured, error }) {
-  let md = `file: ${file || 'mcp'}\n`
-  md += `name: ${title}\n`
-  md += `time: ${(durationMs / 1000).toFixed(2)}s\n`
-  md += `---\n\n`
-
-  if (error) md += `Error: ${error}\n\n---\n\n`
-
-  if (commands && commands.length) {
-    md += `### Commands\n`
-    for (const c of commands) md += `- ${c}\n`
-    md += `\n`
-  }
-
-  md += `### Final State\n`
-  if (captured.url) md += `  > URL: ${captured.url}\n`
-  const links = artifactLinks(captured)
-  if (links) md += links + '\n'
-
-  const traceFile = path.join(dir, 'trace.md')
-  writeFileSync(traceFile, md)
-  return traceFile
-}
-
 async function initCodecept(configPath) {
   if (containerInitialized) return
 
diff --git a/lib/plugin/aiTrace.js b/lib/plugin/aiTrace.js
@@ -8,8 +8,7 @@ import recorder from '../recorder.js'
 import event from '../event.js'
 import output from '../output.js'
 import { deleteDir, clearString } from '../utils.js'
-import { captureSnapshot } from '../utils/captureSnapshot.js'
-import { pickActingHelper, traceDirFor, artifactLinks } from '../utils/trace.js'
+import { captureSnapshot, pickActingHelper, traceDirFor, artifactLinks } from '../utils/trace.js'
 
 const defaultConfig = {
   deleteSuccessful: false,
diff --git a/lib/plugin/pageInfo.js b/lib/plugin/pageInfo.js
@@ -4,8 +4,7 @@ import Container from '../container.js'
 import recorder from '../recorder.js'
 import event from '../event.js'
 import { scanForErrorMessages } from '../html.js'
-import { captureSnapshot } from '../utils/captureSnapshot.js'
-import { pickActingHelper } from '../utils/trace.js'
+import { captureSnapshot, pickActingHelper } from '../utils/trace.js'
 import { output } from '../index.js'
 import store from '../store.js'
 import { humanizeString, ucfirst } from '../utils.js'
diff --git a/lib/utils/captureSnapshot.js b/lib/utils/captureSnapshot.js
diff --git a/lib/utils/trace.js b/lib/utils/trace.js
@@ -1,7 +1,13 @@
 import crypto from 'crypto'
+import fs from 'fs'
 import path from 'path'
 import Container from '../container.js'
 import { clearString } from '../utils.js'
+import { formatHtml } from '../html.js'
+
+// ---------------------------------------------------------------------------
+// Helper / directory naming
+// ---------------------------------------------------------------------------
 
 export function pickActingHelper(helpers) {
   for (const name of Container.STANDARD_ACTING_HELPERS) {
@@ -21,6 +27,10 @@ export function snapshotDirFor(baseDir) {
   return path.resolve(baseDir, `snapshot_${Date.now()}_${hash}`)
 }
 
+// ---------------------------------------------------------------------------
+// Artifact link rendering (trace.md)
+// ---------------------------------------------------------------------------
+
 const ARTIFACT_LABELS = {
   html: 'HTML',
   aria: 'ARIA',
@@ -56,6 +66,30 @@ export function fileToUrl(dir, basename) {
   return `file://${path.join(dir, basename)}`
 }
 
+export function writeTraceMarkdown({ dir, title, file, durationMs, commands, captured, error }) {
+  let md = `file: ${file || 'mcp'}\n`
+  md += `name: ${title}\n`
+  md += `time: ${(durationMs / 1000).toFixed(2)}s\n`
+  md += `---\n\n`
+
+  if (error) md += `Error: ${error}\n\n---\n\n`
+
+  if (commands && commands.length) {
+    md += `### Commands\n`
+    for (const c of commands) md += `- ${c}\n`
+    md += `\n`
+  }
+
+  md += `### Final State\n`
+  if (captured.url) md += `  > URL: ${captured.url}\n`
+  const links = artifactLinks(captured)
+  if (links) md += links + '\n'
+
+  const traceFile = path.join(dir, 'trace.md')
+  fs.writeFileSync(traceFile, md)
+  return traceFile
+}
+
 export function artifactsToFileUrls(captured, dir) {
   const out = {}
   if (captured.url) out.url = captured.url
@@ -69,3 +103,134 @@ export function artifactsToFileUrls(captured, dir) {
   if (typeof captured.localStorageCount === 'number') out.localStorageCount = captured.localStorageCount
   return out
 }
+
+// ---------------------------------------------------------------------------
+// Snapshot capture (HTML / ARIA / screenshot / console / storage)
+// ---------------------------------------------------------------------------
+
+function normalizeBrowserLogs(logs) {
+  return (logs || []).map(l => {
+    if (typeof l === 'string') return l
+    if (l && typeof l.type === 'function' && typeof l.text === 'function') {
+      return { type: l.type(), text: l.text() }
+    }
+    return l
+  })
+}
+
+async function captureStorageState(helper) {
+  if (typeof helper.grabStorageState === 'function') {
+    try {
+      const state = await helper.grabStorageState()
+      if (state) return state
+    } catch {}
+  }
+
+  const state = { cookies: [], origins: [] }
+
+  if (typeof helper.grabCookie === 'function') {
+    try {
+      const cookies = await helper.grabCookie()
+      if (Array.isArray(cookies)) state.cookies = cookies
+    } catch {}
+  }
+
+  if (typeof helper.executeScript === 'function') {
+    try {
+      const result = await helper.executeScript(() => {
+        const out = { origin: location.origin, items: [] }
+        for (let i = 0; i < localStorage.length; i++) {
+          const name = localStorage.key(i)
+          out.items.push({ name, value: localStorage.getItem(name) })
+        }
+        return out
+      })
+      if (result?.items?.length) {
+        state.origins.push({ origin: result.origin, localStorage: result.items })
+      }
+    } catch {}
+  }
+
+  return state
+}
+
+export async function captureSnapshot(helper, {
+  dir,
+  prefix = 'snapshot',
+  fullPage = false,
+  captureURL = true,
+  captureScreenshot = true,
+  captureHTML = true,
+  captureARIA = true,
+  captureBrowserLogs = true,
+  captureStorage = true,
+} = {}) {
+  if (!helper) return {}
+  const out = {}
+
+  if (captureURL) {
+    try {
+      if (helper.grabCurrentUrl) out.url = await helper.grabCurrentUrl()
+    } catch {}
+  }
+
+  if (captureScreenshot && helper.saveScreenshot) {
+    try {
+      const file = `${prefix}_screenshot.png`
+      await helper.saveScreenshot(path.join(dir, file), fullPage)
+      out.screenshot = file
+    } catch {}
+  }
+
+  if (captureHTML && helper.grabSource) {
+    try {
+      const html = await helper.grabSource()
+      // Universal funnel: every captured HTML snapshot flows through formatHtml
+      // (minify -> cleanHtml -> beautify). Don't add direct grabSource->writeFile
+      // paths elsewhere; route through this util so trash-class cleanup stays
+      // consistent across aiTrace, pageInfo, and MCP tools.
+      const formatted = await formatHtml(html)
+      const file = `${prefix}_page.html`
+      fs.writeFileSync(path.join(dir, file), formatted)
+      out.html = file
+    } catch {}
+  }
+
+  if (captureARIA && helper.grabAriaSnapshot) {
+    try {
+      const aria = await helper.grabAriaSnapshot()
+      const file = `${prefix}_aria.txt`
+      fs.writeFileSync(path.join(dir, file), aria)
+      out.aria = file
+    } catch {}
+  }
+
+  if (captureBrowserLogs && helper.grabBrowserLogs) {
+    try {
+      const logs = await helper.grabBrowserLogs()
+      const normalized = normalizeBrowserLogs(logs)
+      const file = `${prefix}_console.json`
+      fs.writeFileSync(path.join(dir, file), JSON.stringify(normalized, null, 2))
+      out.console = file
+      out.consoleCount = normalized.length
+    } catch {}
+  }
+
+  if (captureStorage) {
+    try {
+      const state = await captureStorageState(helper)
+      const cookieCount = state.cookies?.length || 0
+      const localStorageCount = (state.origins || [])
+        .reduce((sum, o) => sum + (o.localStorage?.length || 0), 0)
+      if (cookieCount || localStorageCount) {
+        const file = `${prefix}_storage.json`
+        fs.writeFileSync(path.join(dir, file), JSON.stringify(state, null, 2))
+        out.storage = file
+        out.cookieCount = cookieCount
+        out.localStorageCount = localStorageCount
+      }
+    } catch {}
+  }
+
+  return out
+}
