inspect-ai 0.3.81__py3-none-any.whl → 0.3.83__py3-none-any.whl

inspect_ai/_view/www/src/state/scrolling.ts

@@ -0,0 +1,206 @@
+import { RefObject, useCallback, useEffect, useRef } from "react";
+import { StateCallback, StateSnapshot, VirtuosoHandle } from "react-virtuoso";
+import { createLogger } from "../utils/logger";
+import { debounce } from "../utils/sync";
+import { useStore } from "./store";
+
+const log = createLogger("scrolling");
+
+export function useStatefulScrollPosition<
+  T extends HTMLElement = HTMLDivElement,
+>(
+  elementRef: RefObject<T | null>,
+  elementKey: string,
+  delay = 500,
+  scrollable = true,
+) {
+  const getScrollPosition = useStore(
+    (state) => state.appActions.getScrollPosition,
+  );
+  const setScrollPosition = useStore(
+    (state) => state.appActions.setScrollPosition,
+  );
+
+  // Create debounced scroll handler
+  const handleScroll = useCallback(
+    debounce((e: Event) => {
+      const target = e.target as HTMLElement;
+      const position = target.scrollTop;
+      log.debug(`Storing scroll position`, elementKey, position);
+      setScrollPosition(elementKey, position);
+    }, delay),
+    [elementKey, setScrollPosition, delay],
+  );
+
+  // Function to manually restore scroll position
+  const restoreScrollPosition = useCallback(() => {
+    const element = elementRef.current;
+    const savedPosition = getScrollPosition(elementKey);
+
+    if (element && savedPosition !== undefined) {
+      requestAnimationFrame(() => {
+        element.scrollTop = savedPosition;
+
+        requestAnimationFrame(() => {
+          if (element.scrollTop !== savedPosition) {
+            element.scrollTop = savedPosition;
+          }
+        });
+      });
+    }
+  }, [elementKey, getScrollPosition, elementRef]);
+
+  // Set up scroll listener and restore position on mount
+  useEffect(() => {
+    const element = elementRef.current;
+    if (!element || !scrollable) {
+      return;
+    }
+    log.debug(`Restore Scroll Hook`, elementKey);
+
+    // Restore scroll position on mount
+    const savedPosition = getScrollPosition(elementKey);
+    if (savedPosition !== undefined) {
+      log.debug(`Restoring scroll position`, savedPosition);
+      // Ensure the element has fully rendered
+      requestAnimationFrame(() => {
+        if (element.scrollTop !== savedPosition) {
+          element.scrollTop = savedPosition;
+        }
+      });
+    }
+
+    // Set up scroll listener
+    if (element.addEventListener) {
+      element.addEventListener("scroll", handleScroll);
+    } else {
+      log.warn("Element has no way to add event listener", element);
+    }
+
+    // Clean up
+    return () => {
+      if (element.removeEventListener) {
+        element.removeEventListener("scroll", handleScroll);
+      } else {
+        log.warn("Element has no way to remove event listener", element);
+      }
+    };
+  }, [elementKey, elementRef, handleScroll]);
+
+  return { restoreScrollPosition };
+}
+
+// Define a type for the debounced function that includes the cancel method
+type DebouncedFunction<T extends (...args: any[]) => any> = T & {
+  cancel: () => void;
+  flush: () => void;
+};
+
+export const useVirtuosoState = (
+  virtuosoRef: RefObject<VirtuosoHandle | null>,
+  elementKey: string,
+  delay = 1000,
+) => {
+  // Use useCallback to stabilize the selectors
+  const restoreState = useStore(
+    useCallback((state) => state.app.listPositions[elementKey], [elementKey]),
+  );
+
+  const setListPosition = useStore(
+    useCallback((state) => state.appActions.setListPosition, []),
+  );
+
+  const clearListPosition = useStore(
+    useCallback((state) => state.appActions.clearListPosition, []),
+  );
+
+  // Properly type the debounced function ref
+  const debouncedFnRef = useRef<DebouncedFunction<
+    (isScrolling: boolean) => void
+  > | null>(null);
+
+  // Create the state change handler
+  const handleStateChange: StateCallback = useCallback(
+    (state: StateSnapshot) => {
+      log.debug(`Storing list state: [${elementKey}]`, state);
+      setListPosition(elementKey, state);
+    },
+    [elementKey, setListPosition],
+  );
+
+  // Setup the debounced function once
+  useEffect(() => {
+    debouncedFnRef.current = debounce((isScrolling: boolean) => {
+      log.debug("List scroll", isScrolling);
+      const element = virtuosoRef.current;
+      if (!element) {
+        return;
+      }
+      element.getState(handleStateChange);
+    }, delay) as DebouncedFunction<(isScrolling: boolean) => void>;
+
+    return () => {
+      // Clear the stored position when component unmounts
+      clearListPosition(elementKey);
+    };
+  }, [delay, elementKey, handleStateChange, clearListPosition, virtuosoRef]);
+
+  // Return a stable function reference that uses the ref internally
+  const isScrolling = useCallback((scrolling: boolean) => {
+    if (!scrolling) {
+      return;
+    }
+
+    if (debouncedFnRef.current) {
+      debouncedFnRef.current(scrolling);
+    }
+  }, []);
+
+  // Use a state to prevent re-rendering just because the restore
+  // state changes
+  const stateRef = useRef(restoreState);
+  useEffect(() => {
+    stateRef.current = restoreState;
+  }, [restoreState]);
+
+  const getRestoreState = useCallback(() => stateRef.current, []);
+
+  return { getRestoreState, isScrolling };
+};
+
+export function useRafThrottle<T extends (...args: any[]) => any>(
+  callback: T,
+  dependencies: any[] = [],
+): (...args: Parameters<T>) => void {
+  const rafRef = useRef<number | null>(null);
+  const callbackRef = useRef<T>(callback);
+
+  // Update the callback ref when the callback changes
+  useEffect(() => {
+    callbackRef.current = callback;
+  }, [callback, ...dependencies]);
+
+  const throttledCallback = useCallback((...args: Parameters<T>) => {
+    // Skip if we already have a frame queued
+    if (rafRef.current) {
+      return;
+    }
+
+    rafRef.current = requestAnimationFrame(() => {
+      callbackRef.current(...args);
+      rafRef.current = null;
+    });
+  }, []);
+
+  // Clean up any pending animation frame on unmount
+  useEffect(() => {
+    return () => {
+      if (rafRef.current) {
+        cancelAnimationFrame(rafRef.current);
+        rafRef.current = null;
+      }
+    };
+  }, []);
+
+  return throttledCallback;
+}

inspect_ai/_view/www/src/state/store.ts

@@ -0,0 +1,168 @@
+import { create, StoreApi, UseBoundStore } from "zustand";
+import { devtools, persist } from "zustand/middleware";
+import { immer } from "zustand/middleware/immer";
+import { Capabilities, ClientAPI, ClientStorage } from "../api/types";
+import { createLogger } from "../utils/logger";
+import { AppSlice, createAppSlice, initializeAppSlice } from "./appSlice";
+import { createLogSlice, initalializeLogSlice, LogSlice } from "./logSlice";
+import { createLogsSlice, initializeLogsSlice, LogsSlice } from "./logsSlice";
+import {
+  createSampleSlice,
+  initializeSampleSlice,
+  SampleSlice,
+} from "./sampleSlice";
+import { filterState } from "./store_filter";
+
+const log = createLogger("store");
+
+export interface StoreState extends AppSlice, LogsSlice, LogSlice, SampleSlice {
+  // The shared api
+  api?: ClientAPI | null;
+
+  // Global actions
+  initialize: (api: ClientAPI, capabilities: Capabilities) => void;
+  cleanup: () => void;
+}
+
+// The data that will actually be persisted
+export type PersistedState = {
+  app: AppSlice["app"];
+  log: LogSlice["log"];
+  logs: LogsSlice["logs"];
+  sample: SampleSlice["sample"];
+};
+
+// The store implementation (this will be set when the store is initialized)
+let storeImplementation: UseBoundStore<StoreApi<StoreState>> | null = null;
+
+// Create a proxy store that forwards calls to the real store once initialized
+export const useStore = ((selector?: any) => {
+  if (!storeImplementation) {
+    throw new Error(
+      "Store accessed before initialization. Call initializeStore first.",
+    );
+  }
+  return selector ? storeImplementation(selector) : storeImplementation();
+}) as UseBoundStore<StoreApi<StoreState>>;
+
+// Initialize the store
+export const initializeStore = (
+  api: ClientAPI,
+  capabilities: Capabilities,
+  storage?: ClientStorage,
+) => {
+  // Create the storage implementation
+  const storageImplementation = {
+    getItem: <T>(name: string): T | null => {
+      return storage ? (storage.getItem(name) as T) : null;
+    },
+    setItem: <T>(name: string, value: T): void => {
+      if (storage) {
+        storage.setItem(name, value);
+      }
+    },
+    removeItem: (name: string): void => {
+      if (storage) {
+        storage.removeItem(name);
+      }
+    },
+  };
+
+  // Create the actual store
+  const store = create<StoreState>()(
+    devtools(
+      persist(
+        immer((set, get, store) => {
+          const [appSlice, appCleanup] = createAppSlice(
+            set as (fn: (state: StoreState) => void) => void,
+            get,
+            store,
+          );
+          const [logsSlice, logsCleanup] = createLogsSlice(
+            set as (fn: (state: StoreState) => void) => void,
+            get,
+            store,
+          );
+          const [logSlice, logCleanup] = createLogSlice(
+            set as (fn: (state: StoreState) => void) => void,
+            get,
+            store,
+          );
+          const [sampleSlice, sampleCleanup] = createSampleSlice(
+            set as (fn: (state: StoreState) => void) => void,
+            get,
+            store,
+          );
+
+          return {
+            // Shared state
+            api: null,
+
+            // Initialize
+            initialize: (api, capabilities) => {
+              set((state) => {
+                state.api = api;
+              });
+
+              // Initialize application slices
+              initializeAppSlice(
+                set as (fn: (state: StoreState) => void) => void,
+                capabilities,
+              );
+              initializeLogsSlice(
+                set as (fn: (state: StoreState) => void) => void,
+              );
+              initalializeLogSlice(
+                set as (fn: (state: StoreState) => void) => void,
+              );
+              initializeSampleSlice(
+                set as (fn: (state: StoreState) => void) => void,
+              );
+            },
+
+            // Create the slices and merge them in
+            ...appSlice,
+            ...logsSlice,
+            ...logSlice,
+            ...sampleSlice,
+
+            cleanup: () => {
+              appCleanup();
+              logsCleanup();
+              logCleanup();
+              sampleCleanup();
+            },
+          };
+        }),
+        {
+          name: "app-storage",
+          storage: storageImplementation,
+          partialize: (state) => {
+            const persisted: PersistedState = filterState({
+              app: state.app,
+              log: state.log,
+              logs: state.logs,
+              sample: state.sample,
+            });
+            return persisted as unknown as StoreState;
+          },
+          version: 1,
+          onRehydrateStorage: (state: StoreState) => {
+            return (hydrationState, error) => {
+              log.debug("REHYDRATING STATE");
+              if (error) {
+                log.debug("ERROR", { error });
+              } else {
+                log.debug("STATE", { state, hydrationState });
+              }
+            };
+          },
+        },
+      ),
+    ),
+  );
+
+  // Set the implementation and initialize it
+  storeImplementation = store as UseBoundStore<StoreApi<StoreState>>;
+  store.getState().initialize(api, capabilities);
+};

inspect_ai/_view/www/src/state/store_filter.ts

@@ -0,0 +1,84 @@
+import { PersistedState } from "./store";
+
+export function filterState(state: PersistedState) {
+  if (!state) {
+    return state;
+  }
+
+  // When saving state, we can't store vast amounts of data (like a large sample)
+  const filters = [filterLargeSample, filterLargeLogSummary];
+  return filters.reduce(
+    (filteredState, filter) => filter(filteredState),
+    state,
+  );
+}
+
+// Filters the selected Sample if it is large
+function filterLargeSample(state: PersistedState): PersistedState {
+  if (!state || !state.sample || !state.sample.selectedSample) {
+    return state;
+  }
+
+  const estimatedTotalSize = estimateSize(state.sample.selectedSample.messages);
+  if (estimatedTotalSize > 250000) {
+    return {
+      ...state,
+      sample: {
+        ...state.sample,
+        selectedSample: undefined,
+      },
+    };
+  } else {
+    return state;
+  }
+}
+
+// Filters the selectedlog if it is too large
+function filterLargeLogSummary(state: PersistedState): PersistedState {
+  if (!state || !state.log || !state.log.selectedLogSummary) {
+    return state;
+  }
+
+  const estimatedSize = estimateSize(
+    state.log.selectedLogSummary.sampleSummaries,
+  );
+  if (estimatedSize > 250000) {
+    return {
+      ...state,
+      log: {
+        ...state.log,
+        selectedLogSummary: undefined,
+      },
+    };
+  } else {
+    return state;
+  }
+}
+
+function estimateSize(list: unknown[], frequency = 0.2) {
+  if (!list || list.length === 0) {
+    return 0;
+  }
+
+  // Total number of samples
+  const sampleSize = Math.ceil(list.length * frequency);
+
+  // Get a proper random sample without duplicates
+  const messageIndices = new Set<number>();
+  while (
+    messageIndices.size < sampleSize &&
+    messageIndices.size < list.length
+  ) {
+    const randomIndex = Math.floor(Math.random() * list.length);
+    messageIndices.add(randomIndex);
+  }
+
+  // Calculate size from sampled messages
+  const totalSize = Array.from(messageIndices).reduce((size, index) => {
+    return size + JSON.stringify(list[index]).length;
+  }, 0);
+
+  // Estimate total size based on sample
+  const estimatedTotalSize = (totalSize / sampleSize) * list.length;
+  return estimatedTotalSize;
+}

inspect_ai/_view/www/src/state/utils.ts

@@ -0,0 +1,23 @@
+import { SampleSummary } from "../api/types";
+
+// Function to merge log samples with pending samples
+export const mergeSampleSummaries = (
+  logSamples: SampleSummary[],
+  pendingSamples: SampleSummary[],
+) => {
+  // Create a map of existing sample IDs to avoid duplicates
+  const existingSampleIds = new Set(
+    logSamples.map((sample) => `${sample.id}-${sample.epoch}`),
+  );
+
+  // Filter out any pending samples that already exist in the log
+  const uniquePendingSamples = pendingSamples
+    .filter((sample) => !existingSampleIds.has(`${sample.id}-${sample.epoch}`))
+    .map((sample) => {
+      // Always mark pending items as incomplete to be sure we trigger polling
+      return { ...sample, completed: false };
+    });
+
+  // Combine and return all samples
+  return [...logSamples, ...uniquePendingSamples];
+};

inspect_ai/_view/www/src/storage/index.ts

@@ -0,0 +1,26 @@
+import { ClientStorage } from "../api/types";
+import { PersistedState } from "../state/store";
+import { getVscodeApi } from "../utils/vscode";
+
+const resolveStorage = (): ClientStorage | undefined => {
+  const vscodeApi = getVscodeApi();
+  if (vscodeApi) {
+    return {
+      getItem: (_name: string) => {
+        const state = vscodeApi.getState() as PersistedState;
+        return state;
+      },
+      setItem: (_name: string, value: unknown) => {
+        // TODO: This is pretty gnarly type hijinks
+        const valObj = value as { state: PersistedState; version: number };
+        vscodeApi.setState(valObj);
+      },
+      removeItem: (_name: string) => {
+        vscodeApi.setState(null);
+      },
+    };
+  }
+  return undefined;
+};
+
+export default resolveStorage();