Mask non-deterministic chart data in visual (Argos) tests

.github/workflows/e2e.yml

@@ -91,6 +91,7 @@ jobs:
             "NEXT_PUBLIC_BASE_URL=http://localhost:3000" \
             "NEXT_PUBLIC_BACKEND_URL=http://localhost:8000" \
             "HF_TOKEN=${HF_TOKEN}" \
+            "NEXT_PUBLIC_VISUAL_TEST=true" \
             > .env
           cp .env workbench/_web/.env
           # Backend reads its own .env; mirror the secrets so CONFIG=e2e

workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { useState, useCallback, useRef, useEffect } from "react";
+import { useState, useCallback, useRef, useEffect, useMemo } from "react";
 import { useParams } from "next/navigation";
 import { useQuery, useIsMutating } from "@tanstack/react-query";
 import { getChartById, getConfigForChart } from "@/lib/queries/chartQueries";
@@ -16,6 +16,7 @@ import { useUpdateChartConfig } from "@/lib/api/configApi";
 import { NotebookExporter } from "@/components/NotebookExporter";
 import { ChartModelPill } from "@/components/charts/ChartModelPill";
 import { chartModelFromConfig, isChartStale } from "@/lib/configModelDiff";
+import { isVisualTestMode, maskActivationPatchingDataForVisualTest } from "@/lib/visualTest";
 import { useModelsQuery } from "@/lib/api/modelsApi";
 import { useWorkspace } from "@/stores/useWorkspace";
 
@@ -79,6 +80,14 @@ export function ActivationPatchingDisplay() {
     const hasData =
         patchingChart?.data && "lines" in patchingChart.data && patchingChart.data.lines.length > 0;
 
+    // In visual-test mode, flatten the noisy plotted series so Argos snapshots
+    // are deterministic against real NDIF (chart chrome stays, lines ignored).
+    const widgetData = useMemo(() => {
+        const raw = patchingChart?.data;
+        if (!raw || !hasData) return raw;
+        return isVisualTestMode() ? maskActivationPatchingDataForVisualTest(raw) : raw;
+    }, [patchingChart?.data, hasData]);
+
     // Get the chart's saved name (treat "Untitled Chart" default as empty)
     const rawChartName = patchingChart?.name || "";
     const chartName = rawChartName === "Untitled Chart" ? "" : rawChartName;
@@ -280,7 +289,7 @@ export function ActivationPatchingDisplay() {
             {/* Chart area */}
             <div className="flex-1 p-4 min-h-0">
                 <ActivationPatchingWidget
-                    data={patchingChart!.data!}
+                    data={widgetData!}
                     darkMode={isDarkMode}
                     transparentBackground
                     mode={

workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { useState, useCallback, useRef, useEffect } from "react";
+import { useState, useCallback, useRef, useEffect, useMemo } from "react";
 import { useParams } from "next/navigation";
 import { useQuery, useIsMutating } from "@tanstack/react-query";
 import { getChartById, getConfigForChart } from "@/lib/queries/chartQueries";
@@ -16,6 +16,7 @@ import { useUpdateChartName } from "@/lib/api/chartApi";
 import { useUpdateChartConfig } from "@/lib/api/configApi";
 import { ChartModelPill } from "@/components/charts/ChartModelPill";
 import { chartModelFromConfig, isChartStale } from "@/lib/configModelDiff";
+import { isVisualTestMode, maskLogitLensDataForVisualTest } from "@/lib/visualTest";
 
 interface Lens2Chart {
     id: string;
@@ -67,6 +68,14 @@ export function Lens2Display() {
     const lens2Config = config as Lens2Config | undefined;
     const hasData = lens2Chart?.data && "meta" in lens2Chart.data;
 
+    // In visual-test mode, blank the noisy heatmap values so Argos snapshots
+    // are deterministic against real NDIF (keeps only the final prediction).
+    const widgetData = useMemo(() => {
+        const raw = lens2Chart?.data as LogitLensData | undefined;
+        if (!raw || !hasData) return raw;
+        return isVisualTestMode() ? maskLogitLensDataForVisualTest(raw) : raw;
+    }, [lens2Chart?.data, hasData]);
+
     // ── Persist heatmap UI state (pins, selection, layer window, appearance)
     // into the chart config, mirroring ActivationPatchingDisplay. Debounced
     // because the widget emits on every interaction; restored on mount via
@@ -218,7 +227,7 @@ export function Lens2Display() {
                 {stale && chartModel && <ChartModelPill modelName={chartModel} />}
             </div>
             <LogitLensWidget
-                data={lens2Chart.data! as LogitLensData}
+                data={widgetData as LogitLensData}
                 darkMode={isDarkMode}
                 uiState={savedUiState}
                 onStateChange={handleStateChange}

workbench/_web/src/lib/visualTest.ts

@@ -0,0 +1,80 @@
+/**
+ * Visual-test (Argos) determinism helpers.
+ *
+ * The visual regression suite runs against the real NDIF service (see
+ * `.github/workflows/e2e.yml`). Real model runs are only *approximately*
+ * reproducible: floating-point results drift run-to-run across NDIF hardware,
+ * and that drift shows up as pixel noise in the heatmap cell colors and the
+ * plotted trajectory/line curves — which makes Argos snapshots flaky even
+ * though the UI itself is correct.
+ *
+ * Rather than mock NDIF away, we keep exercising the real request/response
+ * path and only neutralize the non-deterministic *values* right before they
+ * reach the visualization widgets, when `NEXT_PUBLIC_VISUAL_TEST` is set:
+ *
+ *  - Logit-lens heatmap: blank every cell except the final prediction (the
+ *    last token of the last layer), and flatten the tracked trajectories so
+ *    that one remaining cell's color (and the layer skyline) is stable.
+ *  - Activation-patching line plot: flatten the plotted series so the drawn
+ *    lines are constant. The chart chrome (axes, mode bar, token selector)
+ *    still renders, but the noisy curves are ignored.
+ *
+ * This flag is `NEXT_PUBLIC_` so it is inlined into the client bundle at build
+ * time; it must be present when `next build` runs.
+ */
+
+import type { LogitLensData } from "nnsightful";
+import type { ActivationPatchingData } from "@/types/activationPatching";
+
+export function isVisualTestMode(): boolean {
+    return process.env.NEXT_PUBLIC_VISUAL_TEST === "true";
+}
+
+/**
+ * Blank all heatmap content except the final-prediction cell (last token of
+ * the last layer). Token text at that cell stays real so the snapshot still
+ * verifies the model's actual next-token prediction rendered; everything else
+ * — including the values that drive cell colors — is zeroed so the image is
+ * deterministic.
+ */
+export function maskLogitLensDataForVisualTest(data: LogitLensData): LogitLensData {
+    const lastLayer = data.layers.length - 1;
+    const lastPos = data.input.length - 1;
+
+    // topk is indexed [layer][position]; keep only the final-prediction cell.
+    const topk = data.topk.map((layerRow, layerIdx) =>
+        layerRow.map((cell, pos) => (layerIdx === lastLayer && pos === lastPos ? cell : [""])),
+    );
+
+    // Flatten every tracked trajectory to a constant. The final cell reads its
+    // probability from here, so this pins its color; blanked cells reference
+    // tokens absent from `tracked` and render at probability 0.
+    const tracked = data.tracked.map((posMap) => {
+        const flat: Record<string, number[]> = {};
+        for (const [token, trajectory] of Object.entries(posMap)) {
+            flat[token] = trajectory.map(() => 1);
+        }
+        return flat;
+    });
+
+    const entropy = data.entropy?.map((layerRow) => layerRow.map(() => 0));
+
+    return { ...data, topk, tracked, ...(entropy ? { entropy } : {}) };
+}
+
+/**
+ * Flatten the activation-patching series so the plotted lines are constant and
+ * therefore deterministic. Layer count, token labels and chart chrome are left
+ * intact — only the curve values are ignored.
+ */
+export function maskActivationPatchingDataForVisualTest(
+    data: ActivationPatchingData,
+): ActivationPatchingData {
+    const flatten = (grid: number[][]) => grid.map((row) => row.map(() => 0));
+    return {
+        ...data,
+        lines: flatten(data.lines),
+        ranks: flatten(data.ranks),
+        prob_diffs: flatten(data.prob_diffs),
+    };
+}