Efficiently Restore Inspect View State (#754)

* Add baseline state snapshot and restore This doesn’t properly restore state for many of the properties (the state is restore but subsequently overwrtitten) * Don’t preserve context * Don’t throw if the date format is incorrect * Fully restore state - This fixes sorting, epochs, filtering, etc.. to properly restore complete view state. - In addition, this fixes an issue where filtering could cause bad behavior for the view by switching between none, single, and multe sample mode as well as incorrectly hiding tools. * Don’t overly select sample tabs * Don’t reload an already loaded sample * get vscode knowledge out of viewer * Add throttling to saving state
UKGovernmentBEIS · Oct 25, 2024 · 2a59a3e · 2a59a3e
1 parent 1ece196
commit 2a59a3e
Show file tree

Hide file tree

Showing 13 changed files with 787 additions and 535 deletions.
diff --git a/src/inspect_ai/_view/www/dist/assets/index.js b/src/inspect_ai/_view/www/dist/assets/index.js
diff --git a/src/inspect_ai/_view/www/src/App.mjs b/src/inspect_ai/_view/www/src/App.mjs
diff --git a/src/inspect_ai/_view/www/src/Types.mjs b/src/inspect_ai/_view/www/src/Types.mjs
@@ -26,10 +26,14 @@
 /**
  * @typedef {Object} ScoreFilter
  * @property {string} [value]
- * @property {(sample: import("./api/Types.mjs").SampleSummary, value: string) => boolean} [filterFn]
+ * @property {string} [type]
  */
 
 /**
  * @typedef {Object} RenderContext
  * @property {(el: import("preact").JSX.Element) => void} afterBody:
  */
+
+/**
+ * @typedef {"none" | "single" | "many"} SampleMode
+ */
diff --git a/src/inspect_ai/_view/www/src/index.js b/src/inspect_ai/_view/www/src/index.js
@@ -3,5 +3,27 @@ import { html } from "htm/preact";
 
 import { App } from "./App.mjs";
 import api from "./api/index.mjs";
+import { getVscodeApi } from "./utils/vscode.mjs";
+import { throttle } from "./utils/sync.mjs";
 
-render(html`<${App} api=${api} />`, document.getElementById("app"));
+// Read any state from the page itself
+const vscode = getVscodeApi();
+let initialState = undefined;
+if (vscode) {
+  initialState = vscode.getState();
+}
+
+render(
+  html`<${App}
+    api=${api}
+    initialState=${initialState}
+    saveInitialState=${throttle((state) => {
+      const vscode = getVscodeApi();
+      if (vscode) {
+        console.log({ state });
+        vscode.setState(state);
+      }
+    }, 1000)}
+  />`,
+  document.getElementById("app"),
+);
diff --git a/src/inspect_ai/_view/www/src/samples/SampleDialog.mjs b/src/inspect_ai/_view/www/src/samples/SampleDialog.mjs
@@ -103,7 +103,6 @@ export const SampleDialog = ({
                 id=${id}
                 sample=${sample}
                 sampleDescriptor=${sampleDescriptor}
-                visible=${showingSampleDialog}
                 selectedTab=${selectedTab}
                 setSelectedTab=${setSelectedTab}
                 context=${context}

diff --git a/src/inspect_ai/_view/www/src/samples/SampleDisplay.mjs b/src/inspect_ai/_view/www/src/samples/SampleDisplay.mjs
@@ -1,5 +1,4 @@
 import { html } from "htm/preact";
-import { useEffect } from "preact/hooks";
 
 import { ChatView } from "../components/ChatView.mjs";
 import { MetaDataView } from "../components/MetaDataView.mjs";
@@ -96,15 +95,13 @@ export const InlineSampleDisplay = ({
  * @param {import("../samples/SamplesDescriptor.mjs").SamplesDescriptor} props.sampleDescriptor - the sample descriptor
  * @param {string} props.selectedTab - The selected tab
  * @param {(tab: string) => void} props.setSelectedTab - function to set the selected tab
- * @param {boolean} props.visible - Determines whether the sample is visible.
  * @param {import("../Types.mjs").RenderContext} props.context - the app context
  * @returns {import("preact").JSX.Element} The TranscriptView component.
  */
 export const SampleDisplay = ({
   id,
   sample,
   sampleDescriptor,
-  visible,
   selectedTab,
   setSelectedTab,
   context,
@@ -117,21 +114,6 @@ export const SampleDisplay = ({
     return html`<${EmptyPanel} />`;
   }
 
-  // Upon new dialog
-  useEffect(() => {
-    if (!visible) {
-      setSelectedTab(undefined);
-    } else {
-      if (selectedTab === undefined) {
-        const defaultTab =
-          sample.events && sample.events.length > 0
-            ? kSampleTranscriptTabId
-            : kSampleMessagesTabId;
-        setSelectedTab(defaultTab);
-      }
-    }
-  }, [visible]);
-
   // Tab selection
   const onSelectedTab = (e) => {
     const id = e.currentTarget.id;

diff --git a/src/inspect_ai/_view/www/src/samples/SampleList.mjs b/src/inspect_ai/_view/www/src/samples/SampleList.mjs
@@ -1,4 +1,5 @@
 import { html } from "htm/preact";
+import { useCallback } from "preact/hooks";
 import { useEffect, useMemo } from "preact/hooks";
 
 import { ApplicationStyles } from "../appearance/Styles.mjs";
@@ -115,25 +116,28 @@ export const SampleList = (props) => {
     }
   };
 
-  const onkeydown = (e) => {
-    switch (e.key) {
-      case "ArrowUp":
-        prevSample();
-        e.preventDefault();
-        e.stopPropagation();
-        return false;
-      case "ArrowDown":
-        nextSample();
-        e.preventDefault();
-        e.stopPropagation();
-        return false;
-      case "Enter":
-        showSample();
-        e.preventDefault();
-        e.stopPropagation();
-        return false;
-    }
-  };
+  const onkeydown = useCallback(
+    (e) => {
+      switch (e.key) {
+        case "ArrowUp":
+          prevSample();
+          e.preventDefault();
+          e.stopPropagation();
+          return false;
+        case "ArrowDown":
+          nextSample();
+          e.preventDefault();
+          e.stopPropagation();
+          return false;
+        case "Enter":
+          showSample(selectedIndex);
+          e.preventDefault();
+          e.stopPropagation();
+          return false;
+      }
+    },
+    [selectedIndex],
+  );
 
   const listStyle = { ...style, flex: "1", overflowY: "auto", outline: "none" };
 
@@ -239,7 +243,6 @@ const SampleRow = ({
   sampleDescriptor,
   height,
   selected,
-  setSelected,
   showSample,
 }) => {
   const selectedStyle = selected
@@ -257,13 +260,7 @@ const SampleRow = ({
     <div
       id=${`sample-${id}`}
       onclick=${() => {
-        if (setSelected) {
-          setSelected(index);
-        }
-
-        if (showSample) {
-          showSample();
-        }
+        showSample(index);
       }}
       style=${{
         height: `${height}px`,

diff --git a/src/inspect_ai/_view/www/src/samples/SamplesTab.mjs b/src/inspect_ai/_view/www/src/samples/SamplesTab.mjs
@@ -5,6 +5,7 @@ import { useCallback, useEffect, useRef, useState } from "preact/hooks";
 import { SampleDialog } from "./SampleDialog.mjs";
 import { SampleList } from "./SampleList.mjs";
 import { InlineSampleDisplay } from "./SampleDisplay.mjs";
+import { EmptyPanel } from "../components/EmptyPanel.mjs";
 
 /**
  * Renders Samples Tab
@@ -13,6 +14,7 @@ import { InlineSampleDisplay } from "./SampleDisplay.mjs";
  * @param {import("../types/log").Sample} [props.sample] - The sample
  * @param {string} [props.task_id] - The task id
  * @param {import("../api/Types.mjs").SampleSummary[]} [props.samples] - the samples
+ * @param {import("../Types.mjs").SampleMode} props.sampleMode - the mode for displaying samples
  * @param {"epoch" | "sample" | "none" } props.groupBy - how to group items
  * @param {"asc" | "desc" } props.groupByOrder - whether grouping is ascending or descending
  * @param {import("../samples/SamplesDescriptor.mjs").SamplesDescriptor} [props.sampleDescriptor] - the sample descriptor
@@ -36,6 +38,7 @@ export const SamplesTab = ({
   task_id,
   sample,
   samples,
+  sampleMode,
   groupBy,
   groupByOrder,
   sampleDescriptor,
@@ -57,9 +60,13 @@ export const SamplesTab = ({
   const sampleDialogRef = useRef(/** @type {HTMLElement|null} */ (null));
 
   // Shows the sample dialog
-  const showSample = useCallback(() => {
-    setShowingSampleDialog(true);
-  }, [sampleDialogRef]);
+  const showSample = useCallback(
+    (index) => {
+      setSelectedSampleIndex(index);
+      setShowingSampleDialog(true);
+    },
+    [sampleDialogRef],
+  );
 
   useEffect(() => {
     if (showingSampleDialog) {
@@ -100,17 +107,8 @@ export const SamplesTab = ({
         return item.type === "sample";
       }),
     );
-    if (items.length) {
-      setSelectedSampleIndex(0);
-    }
   }, [samples, groupBy, groupByOrder, sampleDescriptor]);
 
-  // Focus the sample list
-  useEffect(() => {
-    // Hide a dialog, if it is displaying
-    setShowingSampleDialog(false);
-  }, [items]);
-
   const nextSampleIndex = useCallback(() => {
     if (selectedSampleIndex < sampleItems.length - 1) {
       return selectedSampleIndex + 1;
@@ -139,7 +137,7 @@ export const SamplesTab = ({
   }, [selectedSampleIndex, samples, sampleStatus, previousSampleIndex]);
 
   const elements = [];
-  if (samples?.length === 1) {
+  if (sampleMode === "single") {
     elements.push(
       html` <${InlineSampleDisplay}
         key=${`${task_id}-single-sample`}
@@ -153,7 +151,7 @@ export const SamplesTab = ({
         context=${context}
       />`,
     );
-  } else {
+  } else if (sampleMode === "many") {
     elements.push(
       html`<${SampleList}
         listRef=${sampleListRef}
@@ -167,6 +165,8 @@ export const SamplesTab = ({
         showSample=${showSample}
       />`,
     );
+  } else {
+    elements.push(html`<${EmptyPanel} />`);
   }
 
   const title =