lemonspace_app/components/canvas/__tests__/use-canvas-flow-reconciliation.test.tsx

// @vitest-environment jsdom

import { act, useEffect, useRef, useState } from "react";
import { createRoot, type Root } from "react-dom/client";
import type { Edge as RFEdge, Node as RFNode } from "@xyflow/react";
import { afterEach, describe, expect, it, vi } from "vitest";

import type { Doc, Id } from "@/convex/_generated/dataModel";
import { useCanvasFlowReconciliation } from "@/components/canvas/use-canvas-flow-reconciliation";

vi.mock("@/components/canvas/canvas-helpers", async () => {
  const actual = await vi.importActual<typeof import("@/components/canvas/canvas-helpers")>(
    "@/components/canvas/canvas-helpers",
  );

  return {
    ...actual,
    getPendingMovePinsFromLocalOps: vi.fn(() => new Map()),
    getPendingRemovedEdgeIdsFromLocalOps: vi.fn(() => new Set()),
  };
});

const asCanvasId = (id: string): Id<"canvases"> => id as Id<"canvases">;
const asNodeId = (id: string): Id<"nodes"> => id as Id<"nodes">;
type HarnessProps = {
  canvasId: Id<"canvases">;
  initialNodes: RFNode[];
  initialEdges: RFEdge[];
  convexNodes?: Doc<"nodes">[];
  convexEdges?: Doc<"edges">[];
  storageUrlsById: Record<string, string>;
  themeMode: "light" | "dark";
  edgeSyncNonce: number;
  isDragging: boolean;
  isResizing: boolean;
  resolvedRealIdByClientRequest: Map<string, Id<"nodes">>;
  pendingConnectionCreateIds: Set<string>;
  previousConvexNodeIdsSnapshot: Set<string>;
  pendingLocalPositionPins?: Map<string, { x: number; y: number }>;
  preferLocalPositionNodeIds?: Set<string>;
};

const latestStateRef: {
  current: {
    nodes: RFNode[];
    edges: RFEdge[];
    resolvedRealIdByClientRequest: Map<string, Id<"nodes">>;
    pendingConnectionCreateIds: Set<string>;
    previousConvexNodeIdsSnapshot: Set<string>;
  } | null;
} = { current: null };

(globalThis as typeof globalThis & { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;

function HookHarness(props: HarnessProps) {
  const [nodes, setNodes] = useState<RFNode[]>(props.initialNodes);
  const [edges, setEdges] = useState<RFEdge[]>(props.initialEdges);
  const nodesRef = useRef(nodes);
  const deletingNodeIds = useRef(new Set<string>());
  const convexNodeIdsSnapshotForEdgeCarryRef = useRef(
    props.previousConvexNodeIdsSnapshot,
  );
  const resolvedRealIdByClientRequestRef = useRef(
    props.resolvedRealIdByClientRequest,
  );
  const pendingConnectionCreatesRef = useRef(props.pendingConnectionCreateIds);
  const pendingLocalPositionUntilConvexMatchesRef = useRef(
    props.pendingLocalPositionPins ?? new Map<string, { x: number; y: number }>(),
  );
  const preferLocalPositionNodeIdsRef = useRef(
    props.preferLocalPositionNodeIds ?? new Set<string>(),
  );
  const isDraggingRef = useRef(props.isDragging);
  const isResizingRef = useRef(props.isResizing);

  useEffect(() => {
    nodesRef.current = nodes;
  }, [nodes]);

  useEffect(() => {
    isDraggingRef.current = props.isDragging;
    isResizingRef.current = props.isResizing;
  }, [props.isDragging, props.isResizing]);

  useCanvasFlowReconciliation({
    canvasId: props.canvasId,
    convexNodes: props.convexNodes,
    convexEdges: props.convexEdges,
    storageUrlsById: props.storageUrlsById,
    themeMode: props.themeMode,
    edges,
    edgeSyncNonce: props.edgeSyncNonce,
    setNodes,
    setEdges,
    refs: {
      nodesRef,
      deletingNodeIds,
      convexNodeIdsSnapshotForEdgeCarryRef,
      resolvedRealIdByClientRequestRef,
      pendingConnectionCreatesRef,
      pendingLocalPositionUntilConvexMatchesRef,
      preferLocalPositionNodeIdsRef,
      isDragging: isDraggingRef,
      isResizing: isResizingRef,
    },
  });

  useEffect(() => {
    latestStateRef.current = {
      nodes,
      edges,
      resolvedRealIdByClientRequest: resolvedRealIdByClientRequestRef.current,
      pendingConnectionCreateIds: pendingConnectionCreatesRef.current,
      previousConvexNodeIdsSnapshot: convexNodeIdsSnapshotForEdgeCarryRef.current,
    };
  }, [edges, nodes]);

  return null;
}

describe("useCanvasFlowReconciliation", () => {
  let container: HTMLDivElement | null = null;
  let root: Root | null = null;

  afterEach(async () => {
    latestStateRef.current = null;
    if (root) {
      await act(async () => {
        root?.unmount();
      });
    }
    container?.remove();
    root = null;
    container = null;
  });

  it("carries an optimistic connection edge until convex publishes the real edge", async () => {
    container = document.createElement("div");
    document.body.appendChild(container);
    root = createRoot(container);

    await act(async () => {
      root?.render(
        <HookHarness
          canvasId={asCanvasId("canvas-1")}
          initialNodes={[
            {
              id: "node-source",
              type: "image",
              position: { x: 0, y: 0 },
              data: {},
            },
            {
              id: "optimistic_req-1",
              type: "prompt",
              position: { x: 120, y: 80 },
              data: {},
            },
          ]}
          initialEdges={[
            {
              id: "optimistic_edge_req-1",
              source: "node-source",
              target: "optimistic_req-1",
            },
          ]}
          convexNodes={[
            {
              _id: asNodeId("node-source"),
              _creationTime: 0,
              canvasId: asCanvasId("canvas-1"),
              type: "image",
              positionX: 0,
              positionY: 0,
              width: 280,
              height: 200,
              data: {},
            } as Doc<"nodes">,
            {
              _id: asNodeId("node-real"),
              _creationTime: 1,
              canvasId: asCanvasId("canvas-1"),
              type: "prompt",
              positionX: 120,
              positionY: 80,
              width: 288,
              height: 220,
              data: {},
            } as Doc<"nodes">,
          ]}
          convexEdges={[]}
          storageUrlsById={{}}
          themeMode="light"
          edgeSyncNonce={0}
          isDragging={false}
          isResizing={false}
          resolvedRealIdByClientRequest={new Map()}
          pendingConnectionCreateIds={new Set(["req-1"])}
          previousConvexNodeIdsSnapshot={new Set(["node-source"])}
        />,
      );
    });

    expect(latestStateRef.current?.edges).toEqual([
      {
        id: "optimistic_edge_req-1",
        source: "node-source",
        target: "node-real",
      },
    ]);
    expect(latestStateRef.current?.resolvedRealIdByClientRequest).toEqual(
      new Map([["req-1", asNodeId("node-real")]]),
    );
    expect(latestStateRef.current?.pendingConnectionCreateIds).toEqual(
      new Set(["req-1"]),
    );
    expect(latestStateRef.current?.previousConvexNodeIdsSnapshot).toEqual(
      new Set(["node-source", "node-real"]),
    );
  });

  it("preserves local dragging nodes instead of swapping in convex nodes mid-drag", async () => {
    container = document.createElement("div");
    document.body.appendChild(container);
    root = createRoot(container);

    await act(async () => {
      root?.render(
        <HookHarness
          canvasId={asCanvasId("canvas-1")}
          initialNodes={[
            {
              id: "optimistic_req-drag",
              type: "image",
              position: { x: 320, y: 180 },
              data: { label: "local" },
              dragging: true,
            },
          ]}
          initialEdges={[]}
          convexNodes={[
            {
              _id: asNodeId("node-real"),
              _creationTime: 1,
              canvasId: asCanvasId("canvas-1"),
              type: "image",
              positionX: 20,
              positionY: 40,
              width: 280,
              height: 200,
              data: { label: "server" },
            } as Doc<"nodes">,
          ]}
          convexEdges={[] as Doc<"edges">[]}
          storageUrlsById={{}}
          themeMode="light"
          edgeSyncNonce={0}
          isDragging={false}
          isResizing={false}
          resolvedRealIdByClientRequest={new Map([
            ["req-drag", asNodeId("node-real")],
          ])}
          pendingConnectionCreateIds={new Set()}
          previousConvexNodeIdsSnapshot={new Set(["node-real"])}
        />,
      );
    });

    expect(latestStateRef.current?.nodes).toEqual([
      {
        id: "optimistic_req-drag",
        type: "image",
        position: { x: 320, y: 180 },
        data: { label: "local" },
        dragging: true,
      },
    ]);
  });
});