Merge pull request #58 from Demonstrandum/cursor/run-selection-persistence-348f

Run selection persistence

Author: Demonstrandum

AGENTS_DEV.md

@@ -249,7 +249,8 @@ The frontend persists state to browser localStorage. This is the core mechanism
 | `_tb_profile.*`        | Saved profile data            | JSON `ProfileData`                                                 | Profile effects |
 | `_tb_profiles_index`   | List of profile names         | JSON string array                                                  | Profile effects |
 | `_tb_active_profile`   | Currently active profile name | Plain string                                                       | Profile effects |
-| `_tb_run_selection.v1` | Run visibility states         | `{version: 1, runSelection: [[id, bool], ...]}`                    | Runs effects    |
+| `_tb_run_selection.v1` | Run visibility states (NgRx)  | `{version: 1, runSelection: [[id, bool], ...]}`                    | Runs effects    |
+| `runSelectionState`    | Run visibility states (Polymer) | Base64-encoded JSON `{runName: bool, ...}`                       | tf-runs-selector |
 | `_tb_run_colors.v1`    | Color overrides               | `{version: 1, runColorOverrides: [...], groupKeyToColorId: [...]}` | Runs effects    |
 | `_tb_tag_filter.v1`    | Tag filter regex              | `{value: string, timestamp: number}`                               | Metrics effects |
 | `_tb_axis_scales.v1`   | Axis scales                   | `{version: 1, yAxisScale?: string, xAxisScale?: string}`           | Metrics effects |
@@ -260,6 +261,7 @@ The frontend persists state to browser localStorage. This is the core mechanism
 Important behaviors:
 
 - When loading run selection from localStorage, if **all** runs would be hidden, the selection is discarded and all runs default to visible.
+- Run selection is stored in two formats: `_tb_run_selection.v1` (NgRx, used by time-series dashboard) and `runSelectionState` (Polymer, used by old-style plugin dashboards). The NgRx effects write both formats to keep them in sync, and fall back to the Polymer format when the NgRx format is empty.
 - Tag filter persistence uses timestamps: user-set values override profile defaults.
 - Pins are synced to localStorage both when saving profiles and when pinning/unpinning cards directly.
 - Section expansion state is restored from localStorage before tag metadata loads, so the auto-expand-first-2-groups default only applies on truly fresh sessions with no persisted state.
@@ -488,6 +490,7 @@ If no persisted state exists and no profile specifies `expandedTagGroups`, the d
 | Issue | Status      | Description                                                                            |
 | ----- | ----------- | -------------------------------------------------------------------------------------- |
 | #25   | Implemented | Shift-select runs to toggle a range (shift+click to select a contiguous range of runs) |
+| #53   | Implemented | Run selection persistence across plugin dashboards and page reloads                    |
 
 ---
 

tensorbored/components/tf_color_scale/colorScale.ts

@@ -16,40 +16,42 @@ import * as d3 from 'd3';
 import {BaseStore} from '../tf_backend/baseStore';
 import {experimentsStore} from '../tf_backend/experimentsStore';
 import {runsStore} from '../tf_backend/runsStore';
-import {standard} from './palettes';
 
-// Example usage:
-// runs = ["train", "test", "test1", "test2"]
-// ccs = new ColorScale();
-// ccs.domain(runs);
-// ccs.getColor("train");
-// ccs.getColor("test1");
+/**
+ * Read the run-name → hex-color map from NgRx (exposed on window by
+ * RunsEffects). This is the single source of truth used by time-series.
+ */
+function readColorMap(): Record<string, string> {
+  const live = (window as any).__tbRunColorMap as
+    | Record<string, string>
+    | undefined;
+  if (!live) {
+    throw new Error('Missing run color map on window.__tbRunColorMap');
+  }
+  return live;
+}
+
 export class ColorScale {
   private identifiers = d3.map();
-  /**
-   * Creates a color scale with optional custom palette.
-   * @param {Array<string>} palette The color palette to use, as an
-   *   Array of hex strings. Defaults to the standard palette.
-   */
-  constructor(private readonly palette: string[] = standard) {}
-  /**
-   * Set the domain of strings.
-   * @param {Array<string>} strings - An array of possible strings to use as the
-   *     domain for your scale.
-   */
+
   public setDomain(strings: string[]): this {
     this.identifiers = d3.map();
-    strings.forEach((s, i) => {
-      this.identifiers.set(s, this.palette[i % this.palette.length]);
+    // Module-level initialization runs before the NgRx bridge seeds
+    // window.__tbRunColorMap. During that phase the run domain is empty.
+    // Keep strict behavior for non-empty domains only.
+    if (strings.length === 0) {
+      return this;
+    }
+    const stored = readColorMap();
+    strings.forEach((s) => {
+      if (stored[s] === undefined) {
+        throw new Error(`Missing run color for "${s}" in shared color map`);
+      }
+      this.identifiers.set(s, stored[s]);
     });
     return this;
   }
-  /**
-   * Use the color scale to transform an element in the domain into a color.
-   * @param {string} The input string to map to a color.
-   * @return {string} The color corresponding to that input string.
-   * @throws Will error if input string is not in the scale's domain.
-   */
+
   public getColor(s: string): string {
     if (!this.identifiers.has(s)) {
       throw new Error(`String ${s} was not in the domain.`);
@@ -58,21 +60,19 @@ export class ColorScale {
   }
 }
 
-/**
- * A color scale of a domain from a store. Automatically updated when the store
- * emits a change.
- */
 function createAutoUpdateColorScale(
   store: BaseStore,
   getDomain: () => string[]
 ): (runName: string) => string {
   const colorScale = new ColorScale();
-  function updateRunsColorScale(): void {
+  function update(): void {
     colorScale.setDomain(getDomain());
   }
-  store.addListener(updateRunsColorScale);
-  updateRunsColorScale();
-  return (domain) => colorScale.getColor(domain);
+  store.addListener(update);
+  // Re-read colors when the NgRx store subscription updates them.
+  window.addEventListener('tb-run-color-map-changed', update);
+  update();
+  return (runName) => colorScale.getColor(runName);
 }
 
 export const runsColorScale = createAutoUpdateColorScale(runsStore, () =>

tensorbored/components/tf_paginated_view/tf-category-paginated-view.ts

@@ -27,6 +27,37 @@ import {
 } from './paginatedViewStore';
 import {TfDomRepeat} from './tf-dom-repeat';
 
+const TAG_GROUP_EXPANSION_KEY = '_tb_tag_group_expansion.v1';
+
+/**
+ * Read the persisted expansion map (shared with the NgRx time-series
+ * dashboard).  Returns null when nothing is stored.
+ */
+function readExpansionMap(): Map<string, boolean> | null {
+  const raw = window.localStorage.getItem(TAG_GROUP_EXPANSION_KEY);
+  if (!raw) return null;
+  try {
+    const parsed = JSON.parse(raw) as {
+      version?: number;
+      groups?: Array<[string, boolean]>;
+    };
+    if (parsed.version !== 1 || !Array.isArray(parsed.groups)) return null;
+    return new Map(parsed.groups);
+  } catch {
+    return null;
+  }
+}
+
+function writeExpansionEntry(name: string, opened: boolean): void {
+  const map = readExpansionMap() ?? new Map<string, boolean>();
+  map.set(name, opened);
+  window.localStorage.setItem(
+    TAG_GROUP_EXPANSION_KEY,
+    JSON.stringify({version: 1, groups: Array.from(map.entries())})
+  );
+  window.dispatchEvent(new CustomEvent('tb-tag-group-expansion-changed'));
+}
+
 @customElement('tf-category-paginated-view')
 class TfCategoryPaginatedView<
   CategoryItem extends {}
@@ -375,6 +406,9 @@ class TfCategoryPaginatedView<
   }
   _togglePane() {
     this.opened = !this.opened;
+    if (this.category?.name) {
+      writeExpansionEntry(this.category.name, this.opened);
+    }
   }
 
   @observe('opened')
@@ -418,17 +452,53 @@ class TfCategoryPaginatedView<
     const {type, compositeSearch} = this.category.metadata as any;
     return compositeSearch && type === CategoryType.SEARCH_RESULTS;
   }
+  private _expansionChangedListener: (() => void) | null = null;
+
   ready() {
     super.ready();
-    this.opened = this.initialOpened == null ? true : this.initialOpened;
+    const stored = this.category?.name ? readExpansionMap() : null;
+    if (
+      stored !== null &&
+      this.category?.name &&
+      stored.has(this.category.name)
+    ) {
+      this.opened = stored.get(this.category.name)!;
+    } else {
+      this.opened = this.initialOpened == null ? true : this.initialOpened;
+    }
     this._limitListener = () => {
       this.set('_limit', getLimit());
     };
     addLimitListener(this._limitListener);
     this._limitListener();
+
+    this._expansionChangedListener = () => this._applyStoredExpansion();
+    window.addEventListener(
+      'tb-tag-group-expansion-changed',
+      this._expansionChangedListener
+    );
   }
+
   detached() {
     removeLimitListener(this._limitListener);
+    if (this._expansionChangedListener) {
+      window.removeEventListener(
+        'tb-tag-group-expansion-changed',
+        this._expansionChangedListener
+      );
+      this._expansionChangedListener = null;
+    }
+  }
+
+  _applyStoredExpansion() {
+    const stored = this.category?.name ? readExpansionMap() : null;
+    if (
+      stored !== null &&
+      this.category?.name &&
+      stored.has(this.category.name)
+    ) {
+      this.opened = stored.get(this.category.name)!;
+    }
   }
 
   @observe(

tensorbored/components/tf_runs_selector/tf-runs-selector.ts

@@ -24,6 +24,90 @@ import {runsColorScale} from '../tf_color_scale/colorScale';
 import '../tf_dashboard_common/tf-multi-checkbox';
 import '../tf_wbr_string/tf-wbr-string';
 
+const RUN_SELECTION_KEY = '_tb_run_selection.v1';
+
+/**
+ * Read the NgRx run-selection localStorage entry and return it as a
+ * bare-run-name → boolean map suitable for tf-multi-checkbox.
+ */
+function readSelectionFromLocalStorage(): Record<string, boolean> {
+  const raw = window.localStorage.getItem(RUN_SELECTION_KEY);
+  if (!raw) return {};
+  try {
+    const parsed = JSON.parse(raw) as {
+      version?: number;
+      runSelection?: Array<[string, boolean]>;
+    };
+    if (parsed.version !== 1 || !Array.isArray(parsed.runSelection)) return {};
+    const out: Record<string, boolean> = {};
+    for (const [runId, selected] of parsed.runSelection) {
+      const slashIdx = runId.indexOf('/');
+      const name = slashIdx >= 0 ? runId.substring(slashIdx + 1) : runId;
+      out[name] = selected;
+    }
+    return out;
+  } catch {
+    return {};
+  }
+}
+
+/**
+ * Merge a bare-run-name selection map back into the NgRx localStorage
+ * entry, preserving any run-IDs that we don't know about.
+ */
+function writeSelectionToLocalStorage(state: Record<string, boolean>): void {
+  const raw = window.localStorage.getItem(RUN_SELECTION_KEY);
+  let existing: Array<[string, boolean]> = [];
+  if (raw) {
+    try {
+      const parsed = JSON.parse(raw) as {
+        version?: number;
+        runSelection?: Array<[string, boolean]>;
+      };
+      if (parsed.version === 1 && Array.isArray(parsed.runSelection)) {
+        existing = parsed.runSelection;
+      }
+    } catch {
+      // ignore
+    }
+  }
+
+  // Build a set of bare names we're about to write so we can detect
+  // which existing entries to update vs. keep as-is.
+  const updatedIds = new Set<string>();
+  const result: Array<[string, boolean]> = [];
+
+  for (const [runId, _] of existing) {
+    const slashIdx = runId.indexOf('/');
+    const name = slashIdx >= 0 ? runId.substring(slashIdx + 1) : runId;
+    if (name in state) {
+      result.push([runId, state[name]]);
+      updatedIds.add(runId);
+    } else {
+      result.push([runId, _]);
+      updatedIds.add(runId);
+    }
+  }
+
+  // Add entries from `state` that weren't in `existing` (bare names).
+  for (const [name, selected] of Object.entries(state)) {
+    const alreadyCovered = existing.some(([runId]) => {
+      const slashIdx = runId.indexOf('/');
+      const n = slashIdx >= 0 ? runId.substring(slashIdx + 1) : runId;
+      return n === name;
+    });
+    if (!alreadyCovered) {
+      result.push([name, selected]);
+    }
+  }
+
+  window.localStorage.setItem(
+    RUN_SELECTION_KEY,
+    JSON.stringify({version: 1, runSelection: result})
+  );
+  window.dispatchEvent(new CustomEvent('tb-run-selection-changed'));
+}
+
 @customElement('tf-runs-selector')
 class TfRunsSelector extends LegacyElementMixin(PolymerElement) {
   static readonly template = html`
@@ -115,8 +199,9 @@ class TfRunsSelector extends LegacyElementMixin(PolymerElement) {
 
   @property({
     type: Object,
+    observer: '_storeRunSelectionState',
   })
-  runSelectionState: object = {};
+  runSelectionState: object = readSelectionFromLocalStorage();
 
   @property({
     type: String,
@@ -159,20 +244,44 @@ class TfRunsSelector extends LegacyElementMixin(PolymerElement) {
 
   _envStoreListener: baseStore.ListenKey;
 
+  private _selectionChangedListener: (() => void) | null = null;
+  private _syncingFromStorage = false;
+
   override attached() {
+    this._syncFromStorage();
+
     this._runStoreListener = runsStore.addListener(() => {
-      this.set('runs', runsStore.getRuns());
+      this.set('runs', runsStore.getRuns().slice().sort());
     });
-    this.set('runs', runsStore.getRuns());
+    this.set('runs', runsStore.getRuns().slice().sort());
     this._envStoreListener = environmentStore.addListener(() => {
       this.set('dataLocation', environmentStore.getDataLocation());
     });
     this.set('dataLocation', environmentStore.getDataLocation());
+
+    this._selectionChangedListener = () => this._syncFromStorage();
+    window.addEventListener(
+      'tb-run-selection-changed',
+      this._selectionChangedListener
+    );
+  }
+
+  private _syncFromStorage() {
+    this._syncingFromStorage = true;
+    this.set('runSelectionState', readSelectionFromLocalStorage());
+    this._syncingFromStorage = false;
   }
 
   override detached() {
     runsStore.removeListenerByKey(this._runStoreListener);
     environmentStore.removeListenerByKey(this._envStoreListener);
+    if (this._selectionChangedListener) {
+      window.removeEventListener(
+        'tb-run-selection-changed',
+        this._selectionChangedListener
+      );
+      this._selectionChangedListener = null;
+    }
   }
 
   _toggleAll() {
@@ -205,7 +314,10 @@ class TfRunsSelector extends LegacyElementMixin(PolymerElement) {
     return dataLocation && dataLocation.length > _dataLocationClipLength;
   }
 
-  // Run selection state and regex are no longer persisted to the URL hash.
-  // Run selection is managed via localStorage (runs_effects.ts).
-  // Regex filter is managed via Angular query params ('runFilter').
+  _storeRunSelectionState() {
+    if (this._syncingFromStorage) return;
+    writeSelectionToLocalStorage(
+      this.runSelectionState as Record<string, boolean>
+    );
+  }
 }