Skip to content
Merged
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
135 changes: 130 additions & 5 deletions apps/web/components/graph-studio/graph-studio-restore.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import type { GraphCanvas } from "./graph-canvas";
import type { GraphConsole } from "./graph-console";
import type { GraphStudioDialogs } from "./graph-studio-dialogs";
import type { GraphToolbar } from "./graph-toolbar";
import type { GraphNodeDefinition, GraphWorkflowPayload, GraphWorkflowRecord } from "./types";
import type { GraphNodeDefinition, GraphRun, GraphWorkflowPayload, GraphWorkflowRecord } from "./types";
import { __resetGraphStartupReadsForTests } from "./utils/graph-startup-reads";
import {
graphWorkflowSnapshotSignature,
Expand All @@ -18,21 +18,29 @@ import {
} from "./utils/graph-tabs";

let loadSavedWorkflow: ComponentProps<typeof GraphStudioDialogs>["onLoadWorkflow"];
let restoreHistoryRun: ComponentProps<typeof GraphStudioDialogs>["onRestoreRun"];

// Keep the real restore, tab, hydration and version coordination; replace visual surfaces only.
vi.mock("./graph-canvas", () => ({
GraphCanvas: ({ nodes }: ComponentProps<typeof GraphCanvas>) => (
GraphCanvas: ({ nodes, edges, groups = [] }: ComponentProps<typeof GraphCanvas>) => (
<>
<output data-testid="canvas">{nodes.map((node) => node.id).join(",")}</output>
<output data-testid="canvas-state">{JSON.stringify({ nodes: nodes.map((node) => ({ id: node.id, fields: node.data.fields })), edges: edges.map(({ id, source, target }) => ({ id, source, target })), groups: groups.map(({ id, node_ids }) => ({ id, node_ids })) })}</output>
<button onClick={() => nodes[0]?.data.onFieldChange(nodes[0].id, "text", "Edited note")}>Edit node</button>
</>
),
}));
vi.mock("./graph-toolbar", () => ({
GraphToolbar: ({ workflowName, onNewTab }: ComponentProps<typeof GraphToolbar>) => (
GraphToolbar: ({ workflowName, activeTabId, tabs, run, onNewTab, onSwitchTab, onCloseTab, onUndo, canUndo }: ComponentProps<typeof GraphToolbar>) => (
<>
<h1>{workflowName}</h1>
<output data-testid="identity">{JSON.stringify({ activeTabId, workflowName, runId: run?.run_id ?? null })}</output>
<button onClick={onNewTab}>New tab</button>
<button onClick={onUndo} disabled={!canUndo}>Undo</button>
{tabs?.map((tab) => <span key={tab.tab_id}>
<button onClick={() => onSwitchTab?.(tab.tab_id)}>Switch {tab.tab_id}</button>
<button onClick={() => onCloseTab?.(tab.tab_id)}>Close {tab.tab_id}</button>
</span>)}
</>
),
}));
Expand All @@ -41,8 +49,9 @@ vi.mock("./graph-console", () => ({
}));
vi.mock("./graph-left-rail", () => ({ GraphLeftRail: () => null }));
vi.mock("./graph-studio-dialogs", () => ({
GraphStudioDialogs: ({ onLoadWorkflow }: ComponentProps<typeof GraphStudioDialogs>) => {
GraphStudioDialogs: ({ onLoadWorkflow, onRestoreRun }: ComponentProps<typeof GraphStudioDialogs>) => {
loadSavedWorkflow = onLoadWorkflow;
restoreHistoryRun = onRestoreRun;
return null;
},
}));
Expand All @@ -53,6 +62,10 @@ const definition: GraphNodeDefinition = {
type: "utility.note", title: "Note", category: "Utility", fields: [],
ports: { inputs: [], outputs: [] },
};
const historyDefinitions: GraphNodeDefinition[] = [
{ type: "prompt.text", title: "Text", category: "Prompt", fields: [{ id: "text", label: "Text", type: "textarea" }], ports: { inputs: [], outputs: [{ id: "text", label: "Text", type: "text" }] } },
{ type: "display.any", title: "Display", category: "Preview", fields: [], ports: { inputs: [{ id: "value", label: "Value", type: "any" }], outputs: [] } },
];
const localWorkflow: GraphWorkflowPayload = {
schema_version: 1, workflow_id: "saved-a", name: "Local snapshot",
nodes: [{ id: "local-node", type: definition.type, position: { x: 10, y: 20 }, fields: {} }],
Expand All @@ -64,10 +77,12 @@ const savedRecord = {
} satisfies GraphWorkflowRecord;
const recordUrl = "/api/control/media/graph/workflows/saved-a";
const reloadRecord = vi.fn<() => Promise<Response>>();
const historyRequests: { runId: string; resolve: (response: Response) => void }[] = [];
const fetcher = vi.fn(async (url: string) => {
if (/\/graph\/runs\/history-[^/]+$/.test(url)) return new Promise<Response>((resolve) => historyRequests.push({ runId: url.split("/").at(-1)!, resolve }));
if (url === recordUrl) return reloadRecord();
if (url === "/api/control/health") return Response.json({ install_id: "restore-test" });
if (url === "/api/control/media/graph/node-definitions") return Response.json({ items: [definition] });
if (url === "/api/control/media/graph/node-definitions") return Response.json({ items: [definition, ...historyDefinitions] });
return Response.json({ items: [], available_credits: 0 });
});

Expand All @@ -87,6 +102,7 @@ beforeEach(() => {
vi.stubGlobal("ResizeObserver", vi.fn(() => ({ observe: vi.fn(), disconnect: vi.fn() })));
vi.stubGlobal("fetch", fetcher);
fetcher.mockClear();
historyRequests.length = 0;
reloadRecord.mockReset();
reloadRecord.mockResolvedValue(Response.json(savedRecord));
storeSnapshot();
Expand Down Expand Up @@ -236,3 +252,112 @@ it("does not accept a partially hydrated saved graph as clean", async () => {
fireEvent.click(screen.getByRole("button", { name: "New tab" }));
expect(readGraphTabSession("restore-test")?.tabs.find((tab) => tab.workflow_id === "saved-a")?.dirty).toBe(true);
});

const historicalRun: GraphRun = {
run_id: "history-a", workflow_id: "saved-a", status: "completed",
workflow_json: {
schema_version: 1, workflow_id: "saved-a", name: "Historical A",
nodes: [
{ id: "history-text", type: "prompt.text", position: { x: 100, y: 100 }, fields: { text: "Historical text" } },
{ id: "history-result", type: "display.any", position: { x: 600, y: 100 }, fields: {} },
],
edges: [{ id: "history-edge", source: "history-text", source_port: "text", target: "history-result", target_port: "value" }],
metadata: { groups: [{ id: "history-group", title: "History group", color: "default", node_ids: ["history-text", "history-result"], bounds: { x: 0, y: 0, width: 1000, height: 600 } }] },
},
};

function draftState() {
const session = readGraphTabSession("restore-test")!;
return { active: session.active_tab_id, tabs: session.tabs.map(({ updated_at: _timestamp, ...tab }) => tab) };
}

async function readyForHistory() {
const source = readGraphTabSession("restore-test")!.tabs[0];
writeGraphTabSession("restore-test", "tab-a", [source, {
tab_id: "tab-b", workflow_id: null, workflow_name: "Draft B", dirty: true,
workflow_json: { ...localWorkflow, workflow_id: null, name: "Draft B", nodes: [{ ...localWorkflow.nodes[0], id: "node-b", fields: { text: "Keep B" } }] },
console_lines: ["B console"],
}]);
render(<GraphStudio />);
await waitFor(() => expect(screen.getByTestId("canvas").textContent).toBe("saved-node"));
}

function beginHistoryRestore(runId = "history-a") {
const completion = Promise.resolve(restoreHistoryRun({ run_id: runId, workflow_id: "saved-a", status: "completed" }));
expect(historyRequests.at(-1)?.runId).toBe(runId);
return completion;
}

it.each(["switch to B", "close source", "new tab", "away and back"])("ignores a late history restore after %s", async (change) => {
await readyForHistory();
const completion = beginHistoryRestore();
if (change === "new tab") fireEvent.click(screen.getByRole("button", { name: "New tab", exact: true }));
else if (change === "close source") fireEvent.click(screen.getByRole("button", { name: "Close tab-a", exact: true }));
else {
fireEvent.click(screen.getByRole("button", { name: "Switch tab-b", exact: true }));
if (change === "away and back") fireEvent.click(screen.getByRole("button", { name: "Switch tab-a", exact: true }));
}
const before = { identity: screen.getByTestId("identity").textContent, canvas: screen.getByTestId("canvas-state").textContent, drafts: draftState(), console: screen.getByTestId("console").textContent };
await act(async () => { historyRequests[0].resolve(Response.json(historicalRun)); await completion; });
expect(screen.getByTestId("identity").textContent).toBe(before.identity);
expect(screen.getByTestId("canvas-state").textContent).toBe(before.canvas);
expect(draftState()).toEqual(before.drafts);
expect(screen.getByTestId("console").textContent).toBe(before.console);
});

it("keeps intervening same-tab edits and their undo history instead of restoring an obsolete snapshot", async () => {
await readyForHistory();
const completion = beginHistoryRestore();
fireEvent.click(screen.getByRole("button", { name: "Edit node", exact: true }));
const before = draftState();
await act(async () => { historyRequests[0].resolve(Response.json(historicalRun)); await completion; });
expect(screen.getByTestId("canvas").textContent).toBe("saved-node");
expect(draftState()).toEqual(before);
expect(before.tabs.find((tab) => tab.tab_id === "tab-a")?.workflow_json?.nodes[0].fields.text).toBe("Edited note");
fireEvent.click(screen.getByRole("button", { name: "Undo", exact: true }));
expect(screen.getByTestId("canvas").textContent).toBe("saved-node");
expect(readGraphTabSession("restore-test")!.tabs[0].workflow_json?.nodes[0].fields.text).not.toBe("Edited note");
});

it("keeps a newer history restore when an older reply arrives afterward", async () => {
await readyForHistory();
const older = beginHistoryRestore();
const newer = beginHistoryRestore("history-newer");
await act(async () => { historyRequests[1].resolve(Response.json({ ...historicalRun, run_id: "history-newer", workflow_json: { ...historicalRun.workflow_json, name: "Newer history" } })); await newer; });
const before = draftState();
await act(async () => { historyRequests[0].resolve(Response.json(historicalRun)); await older; });
expect(screen.getByRole("heading").textContent).toBe("Newer history");
expect(draftState()).toEqual(before);
expect(screen.getByTestId("console").textContent).not.toContain("Restored graph run history-a.");
});

it("restores the same owner's historical nodes, edges, groups, run and console, with normal undo", async () => {
await readyForHistory();
const completion = beginHistoryRestore();
await act(async () => { historyRequests[0].resolve(Response.json(historicalRun)); await completion; });
expect(JSON.parse(screen.getByTestId("identity").textContent!)).toEqual({ activeTabId: "tab-a", workflowName: "Historical A", runId: "history-a" });
expect(JSON.parse(screen.getByTestId("canvas-state").textContent!)).toMatchObject({ nodes: [{ id: "history-text" }, { id: "history-result" }], edges: [{ id: "history-edge", source: "history-text", target: "history-result" }], groups: [{ id: "history-group", node_ids: ["history-text", "history-result"] }] });
const source = readGraphTabSession("restore-test")!.tabs.find((tab) => tab.tab_id === "tab-a")!;
expect(source).toMatchObject({ workflow_id: "saved-a", workflow_name: "Historical A", run_id: "history-a" });
expect(source.workflow_json?.edges).toEqual(historicalRun.workflow_json?.edges);
expect(source.workflow_json?.metadata?.groups).toMatchObject([{ id: "history-group", node_ids: ["history-text", "history-result"] }]);
expect(source.console_lines).toContain("Restored graph run history-a.");
fireEvent.click(screen.getByRole("button", { name: "Undo", exact: true }));
expect(screen.getByTestId("canvas").textContent).toBe("saved-node");
});

it.each([false, true])("reports a failed history read only while its source still owns the workspace (stale=%s)", async (stale) => {
await readyForHistory();
const completion = beginHistoryRestore();
if (stale) fireEvent.click(screen.getByRole("button", { name: "Switch tab-b", exact: true }));
const before = { drafts: draftState(), canvas: screen.getByTestId("canvas-state").textContent, console: screen.getByTestId("console").textContent };
await act(async () => { historyRequests[0].resolve(Response.json({ detail: "History unavailable" }, { status: 503 })); await completion; });
expect(screen.getByTestId("canvas-state").textContent).toBe(before.canvas);
if (stale) {
expect(draftState()).toEqual(before.drafts);
expect(screen.getByTestId("console").textContent).toBe(before.console);
} else {
expect(screen.getByTestId("console").textContent).toContain("Run history-a could not be restored:");
expect(readGraphTabSession("restore-test")!.tabs[0].workflow_json).toEqual(before.drafts.tabs[0].workflow_json);
}
});
14 changes: 12 additions & 2 deletions apps/web/components/graph-studio/graph-studio.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -1789,14 +1789,22 @@ function GraphStudioClient() {
const restoreRunFromHistory = useCallback(
async (historyRun: GraphRunHistoryItem) => {
markWorkspaceChanged();
const sourceTabId = activeTabIdRef.current;
const sourceVersion = workspaceRestoreVersionRef.current;
const sourceSignature = graphWorkflowSnapshotSignature(currentWorkflowPayload);
const ownsSource = () =>
activeTabIdRef.current === sourceTabId &&
workspaceRestoreVersionRef.current === sourceVersion &&
graphWorkflowSnapshotSignature(currentHistorySnapshotRef.current?.workflow) === sourceSignature;
try {
const fullRun = historyRun.workflow_json?.nodes?.length
const fullRun = historyRun.workflow_json?.nodes.length
? (historyRun as GraphRun)
: await jsonFetch<GraphRun>(
`/api/control/media/graph/runs/${historyRun.run_id}`,
);
if (!ownsSource()) return;
const workflow = fullRun.workflow_json;
if (!workflow?.nodes?.length) {
if (!workflow?.nodes.length) {
appendConsole(
`Run ${historyRun.run_id} does not include a restorable workflow snapshot.`,
);
Expand All @@ -1811,13 +1819,15 @@ function GraphStudioClient() {
setSidebarDialog(null);
appendConsole(`Restored graph run ${historyRun.run_id}.`);
} catch (error) {
if (!ownsSource()) return;
appendConsole(
`Run ${historyRun.run_id} could not be restored: ${(error as Error).message}`,
);
}
},
[
appendConsole,
currentWorkflowPayload,
hydrateWorkflowPayload,
markWorkspaceChanged,
workflowUpdatedAt,
Expand Down
Loading