Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,7 @@ import { ListTree } from "lucide-react";

import CellWrapper from "@/shared/DataTableCells/CellWrapper";
import {
useOutputLoadingByPromptDatasetItemId,
useOutputStaleStatusByPromptDatasetItemId,
useOutputValueByPromptDatasetItemId,
useSelectedRuleIdsByPromptDatasetItemId,
useTraceIdByPromptDatasetItemId,
useOutputByPromptDatasetItemId,
useDatasetType,
useExperimentIdByPromptId,
} from "@/store/PlaygroundStore";
Expand Down Expand Up @@ -46,30 +42,23 @@ const PlaygroundOutputCell: React.FunctionComponent<

const workspaceName = useAppStore((state) => state.activeWorkspaceName);

const value = useOutputValueByPromptDatasetItemId(
promptId,
originalRow.dataItemId,
);

const isLoading = useOutputLoadingByPromptDatasetItemId(
promptId,
originalRow.dataItemId,
);

const stale = useOutputStaleStatusByPromptDatasetItemId(
promptId,
originalRow.dataItemId,
);

const traceId = useTraceIdByPromptDatasetItemId(
promptId,
originalRow.dataItemId,
);

const selectedRuleIds = useSelectedRuleIdsByPromptDatasetItemId(
// One store subscription per cell, not five. Each of the fields below used to
// come from its own hook, and every one of those hooks runs the *same*
// selector, so a single streaming token re-ran it (dataset items x prompts x 5)
// times. Reading the output object once and picking the fields off it is
// equivalent — the defaults below are the ones those hooks applied — while
// cutting the per-token selector work by 5x. The selector returns the stored
// object reference, which only changes when this cell's own output changes, so
// unrelated updates still don't re-render this cell.
const output = useOutputByPromptDatasetItemId(
promptId,
originalRow.dataItemId,
);
const value = output?.value ?? null;
const isLoading = output?.isLoading ?? false;
const stale = output?.stale ?? false;
const traceId = output?.traceId ?? null;
const selectedRuleIds = output?.selectedRuleIds;

const datasetType = useDatasetType();
const experimentId = useExperimentIdByPromptId(promptId);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { useEffect, useRef, useState } from "react";
import { useEffect, useState } from "react";
import { DatasetItem } from "@/types/datasets";
import { useHydrateDatasetItemData } from "@/v2/pages/PlaygroundPage/useHydrateDatasetItemData";
import { containsTruncatedMedia } from "@/lib/media";

export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): {
hydratedItems: DatasetItem[];
Expand All @@ -9,10 +10,16 @@ export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): {
const hydrateDatasetItemData = useHydrateDatasetItemData();
const [hydratedItems, setHydratedItems] = useState<DatasetItem[]>([]);
const [isHydrating, setIsHydrating] = useState(false);
const cancelledRef = useRef(false);

useEffect(() => {
cancelledRef.current = false;
// Scoped to this effect run rather than a shared ref. React runs the previous
// run's cleanup before this body, so an in-flight hydration from an earlier
// dataset observes its own `cancelled === true` and stops — even when this run
// returns early below and installs no cleanup of its own. A shared ref instead
// got reset here on every run, un-cancelling the previous run's loop; its
// continuation would then write the old dataset's data into the new array at a
// positional index. Also covers unmount.
let cancelled = false;

if (datasetItems.length === 0) {
setHydratedItems([]);
Expand All @@ -21,19 +28,41 @@ export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): {
}

setHydratedItems(datasetItems);

// Only items carrying truncated media need a round-trip — that is the sole
// condition under which hydrateDatasetItemData fetches anything. Selecting
// them up front means the loop below runs once per *media* item rather than
// once per row, and each pass rebuilt the whole array, so a 1000-row text
// dataset was doing ~1M array copies to arrive back at the data it started
// with.
const indexesToHydrate = datasetItems.reduce<number[]>(
(acc, item, index) => {
if (containsTruncatedMedia(item.data)) {
acc.push(index);
}
return acc;
},
[],
);

if (indexesToHydrate.length === 0) {
setIsHydrating(false);
return;
}
Comment on lines +48 to +51

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale media fetch overwrites new dataset

The no-media early-return path resets the shared cancelledRef.current to false without installing a replacement cleanup, so an earlier hydrateItems continuation can pass its post-await cancellation check and call setHydratedItems on the new array, overwriting it with stale data or updating state after unmount — should we use a generation-specific cancellation token, or install cleanup on every effect path?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`apps/opik-frontend/src/v2/pages/PlaygroundPage/useIncrementalDatasetHydration.ts`
around lines 42-45, fix the `useEffect` early-return path for datasets with no media.
Resetting the shared `cancelledRef` allows an older `hydrateItems` continuation to
update the replacement dataset or call setters after unmount; use a generation-specific
cancellation token/version, or install cleanup for every effect path, so stale async
work is always invalidated.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit fb17f8e addressed this comment by replacing the shared cancellation ref with an effect-scoped cancellation token. Previous async hydrations are invalidated by cleanup before a new effect run, including no-media early returns and unmounts.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct catch, and worth being precise about why it matters: this was a latent flaw that the no-media early return turned into a dangerous one.

cancelledRef is one ref shared across every effect run, and the effect body reset it to false on entry. React runs the previous run cleanup before the next body, so the sequence is: run 1 starts hydrating dataset A and installs cleanup -> dataset changes -> cleanup sets cancelled = true -> run 2 enters and resets it to false, un-cancelling run 1 still-pending loop -> run 1 continuation passes its post-await check and calls setHydratedItems on the new array.

Before this PR the only early return sat after setHydratedItems([]), so that stale continuation prev.map() ran over an empty array and produced nothing. Harmless, which is why it went unnoticed. The early return I added sits after setHydratedItems(datasetItems) with a populated array, so the same continuation writes dataset A data into dataset B array at a positional index — silent cross-dataset corruption, reachable by switching dataset or page size while a media-bearing dataset is still hydrating.

Fixed in fb17f8e by taking your second suggestion rather than a generation token: the flag is now scoped to each effect run (let cancelled = false) instead of being a shared ref. Each run observes only its own flag, so early-returning without installing cleanup cannot resurrect an older run loop, and no path needs its own cleanup to stay correct. That also fixes the pre-existing empty-dataset path and the setState-after-unmount case, and drops the useRef import.

Preferred it over a generation counter because it removes the shared mutable state entirely rather than versioning it, so the whole class of bug goes away instead of being guarded against.

One thing I did not do: there is no regression test. Reproducing it needs an in-flight media fetch interleaved with a dataset swap, and the hook has no existing test harness. The fix is reasoned against React cleanup ordering, not proven by test, so it is worth a human eye.


setIsHydrating(true);

const hydrateItems = async () => {
for (let i = 0; i < datasetItems.length; i++) {
if (cancelledRef.current) return;
for (const index of indexesToHydrate) {
if (cancelled) return;

const hydratedData = await hydrateDatasetItemData(datasetItems[i]);
const hydratedData = await hydrateDatasetItemData(datasetItems[index]);

if (cancelledRef.current) return;
if (cancelled) return;

setHydratedItems((prev) =>
prev.map((item, idx) =>
idx === i ? { ...item, data: hydratedData } : item,
idx === index ? { ...item, data: hydratedData } : item,
),
);
}
Expand All @@ -44,7 +73,7 @@ export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): {
hydrateItems();

return () => {
cancelledRef.current = true;
cancelled = true;
};
}, [datasetItems, hydrateDatasetItemData]);

Expand Down
Loading