fix(grid): stable row IDs and live API handle

js/packages/ui/src/components/data/ScreenshotDataGrid.tsx

@@ -166,11 +166,16 @@ function _ScreenshotDataGrid<TData = DataGridRow>(
   const gridApiRef = useRef<GridReadyEvent["api"] | null>(null);
 
   // Expose both API and DOM element through ref
+  // Use getters to ensure handle always returns the current (live) api, not a stale snapshot
   useImperativeHandle(
     ref,
     () => ({
-      api: gridApiRef.current,
-      element: containerRef.current,
+      get api() {
+        return gridApiRef.current;
+      },
+      get element() {
+        return containerRef.current;
+      },
     }),
     [],
   );
@@ -205,17 +210,27 @@ function _ScreenshotDataGrid<TData = DataGridRow>(
     return () => renderers.noRowsFallback;
   }, [renderers?.noRowsFallback]);
 
-  // Generate row ID from __rowKey if available
+  // Generate row ID from __rowKey, _index, or rowIndex if available
   const resolvedGetRowId = useMemo(() => {
     if (getRowId) return getRowId;
     return (params: GetRowIdParams<TData>) => {
-      const data = params.data as DataGridRow;
+      const data = params.data as DataGridRow & {
+        _index?: number;
+        rowIndex?: number;
+      };
       if (data?.__rowKey !== undefined) {
         return String(data.__rowKey);
       }
-      // Use rowIndex from the data or generate a random ID
-      const index = (params.data as unknown as { rowIndex?: number })?.rowIndex;
-      return String(index ?? Math.random());
+      // Use _index (stable id for query results from dataFrameToRowObjects)
+      if (data?._index !== undefined) {
+        return String(data._index);
+      }
+      // Fall back to rowIndex from the data
+      if (data?.rowIndex !== undefined) {
+        return String(data.rowIndex);
+      }
+      // Last resort: generate a random ID
+      return String(Math.random());
     };
   }, [getRowId]);
 

js/packages/ui/src/components/data/__tests__/ScreenshotDataGrid.test.tsx

@@ -54,3 +54,50 @@ describe("ScreenshotDataGrid grid options", () => {
     expect(props.enableCellTextSelection).toBe(false);
   });
 });
+
+describe("ScreenshotDataGrid getRowId", () => {
+  beforeEach(() => {
+    agGridProps.mockClear();
+  });
+
+  it("uses _index when present (query result rows)", () => {
+    render(<ScreenshotDataGrid columnDefs={[]} rowData={[]} />);
+    const props = agGridProps.mock.calls[0][0];
+    const getRowId = props.getRowId;
+
+    expect(getRowId({ data: { _index: 42 } })).toBe("42");
+  });
+
+  it("uses __rowKey when present (takes precedence)", () => {
+    render(<ScreenshotDataGrid columnDefs={[]} rowData={[]} />);
+    const props = agGridProps.mock.calls[0][0];
+    const getRowId = props.getRowId;
+
+    expect(getRowId({ data: { __rowKey: "custom-key", _index: 42 } })).toBe(
+      "custom-key",
+    );
+  });
+
+  it("falls back to rowIndex when _index is absent", () => {
+    render(<ScreenshotDataGrid columnDefs={[]} rowData={[]} />);
+    const props = agGridProps.mock.calls[0][0];
+    const getRowId = props.getRowId;
+
+    expect(getRowId({ data: { rowIndex: 10 } })).toBe("10");
+  });
+
+  it("generates a random ID as last resort", () => {
+    render(<ScreenshotDataGrid columnDefs={[]} rowData={[]} />);
+    const props = agGridProps.mock.calls[0][0];
+    const getRowId = props.getRowId;
+
+    const id1 = getRowId({ data: {} });
+    const id2 = getRowId({ data: {} });
+
+    // Both should be valid strings (not undefined)
+    expect(id1).toBeTruthy();
+    expect(id2).toBeTruthy();
+    // Random IDs should be different
+    expect(id1).not.toBe(id2);
+  });
+});