fix: wire Edit Component to update the task in the v2 editor

src/routes/v2/pages/Editor/nodes/TaskNode/context/TaskDetails/TaskDetails.tsx

@@ -17,6 +17,7 @@ import { useEditorSession } from "@/routes/v2/pages/Editor/store/EditorSessionCo
 import { useSpec } from "@/routes/v2/shared/providers/SpecContext";
 import { useSharedStores } from "@/routes/v2/shared/store/SharedStoreContext";
 import { SYSTEM_ANNOTATIONS, ZINDEX_ANNOTATION } from "@/utils/annotations";
+import type { HydratedComponentReference } from "@/utils/componentSpec";
 import { tracking } from "@/utils/tracking";
 
 import { getTaskYamlText } from "./components/actions/getTaskYamlText";
@@ -38,7 +39,7 @@ export const TaskDetails = observer(function TaskDetails({
   const { track } = useAnalytics();
   const { editor } = useSharedStores();
   const { undo } = useEditorSession();
-  const { renameTask } = useTaskActions();
+  const { renameTask, replaceTask } = useTaskActions();
   const notify = useToastNotification();
   const spec = useSpec();
   const task = useTask(entityId);
@@ -74,6 +75,23 @@ export const TaskDetails = observer(function TaskDetails({
 
   const isSubgraphTask = task.subgraphSpec !== undefined;
 
+  const handleComponentSaved = (
+    hydratedComponent: HydratedComponentReference,
+  ) => {
+    const result = replaceTask(spec, task.$id, hydratedComponent);
+    const lostInputs = result.inputDiff?.lostEntities ?? [];
+
+    if (lostInputs.length > 0) {
+      const inputNames = lostInputs.map((input) => input.name).join(", ");
+      notify(
+        `Component updated. Removed ${lostInputs.length} input(s) no longer defined: ${inputNames}.`,
+        "warning",
+      );
+    } else {
+      notify("Component updated", "success");
+    }
+  };
+
   const handleZIndexChange = (newZIndex: number) => {
     undo.withGroup("Update task z-index", () => {
       task.annotations.set(ZINDEX_ANNOTATION, newZIndex);
@@ -163,6 +181,7 @@ export const TaskDetails = observer(function TaskDetails({
           yamlText={yamlText}
           taskName={task.name}
           pythonCode={pythonCode}
+          onComponentSaved={handleComponentSaved}
         />
       </BlockStack>
 

src/routes/v2/pages/Editor/nodes/TaskNode/context/TaskDetails/components/ComponentRefBar.tsx

@@ -22,6 +22,7 @@ import { Text } from "@/components/ui/typography";
 import useToastNotification from "@/hooks/useToastNotification";
 import type { ComponentReference as ModelComponentReference } from "@/models/componentSpec/entities/types";
 import { useAnalytics } from "@/providers/AnalyticsProvider";
+import type { HydratedComponentReference } from "@/utils/componentSpec";
 import { getComponentName } from "@/utils/getComponentName";
 import { isSubgraph } from "@/utils/subgraphUtils";
 import { tracking } from "@/utils/tracking";
@@ -35,13 +36,17 @@ interface ComponentRefBarProps {
   yamlText: string;
   taskName: string;
   pythonCode: string | undefined;
+  onComponentSaved?: (
+    hydratedComponent: HydratedComponentReference,
+  ) => void | Promise<void>;
 }
 
 export function ComponentRefBar({
   componentRef,
   yamlText,
   taskName,
   pythonCode,
+  onComponentSaved,
 }: ComponentRefBarProps) {
   const { track } = useAnalytics();
   const notify = useToastNotification();
@@ -195,6 +200,7 @@ export function ComponentRefBar({
         <ComponentEditorDialog
           text={yamlText}
           onClose={() => setIsEditDialogOpen(false)}
+          onComponentSaved={onComponentSaved}
         />
       )}
 

src/routes/v2/pages/Editor/store/actions/__tests__/task.actions.replaceTask.test.ts

@@ -0,0 +1,118 @@
+import { describe, expect, it, vi } from "vitest";
+
+import { ComponentSpec, Task } from "@/models/componentSpec";
+import { replaceTask } from "@/routes/v2/pages/Editor/store/actions/task.actions";
+import type { ComponentReference, InputSpec } from "@/utils/componentSpec";
+
+// task.actions imports the editor node registry barrel (used by other actions,
+// not by replaceTask). Loading the real barrel pulls in the node manifests and
+// trips a module cycle in the test environment, so stub it.
+vi.mock("@/routes/v2/pages/Editor/nodes", () => ({
+  editorRegistry: { getByNodeId: () => undefined },
+}));
+
+const noopUndo = {
+  withGroup: <T>(_label: string, fn: () => T): T => fn(),
+};
+
+const datasetRef = (
+  digest: string,
+  inputs: InputSpec[],
+): ComponentReference => ({
+  name: "Chicago Taxi Trips dataset",
+  digest,
+  spec: {
+    name: "Chicago Taxi Trips dataset",
+    inputs,
+    outputs: [{ name: "Table" }],
+    implementation: {
+      container: { image: "alpine/curl", command: ["sh", "-ec", "true"] },
+    },
+  },
+});
+
+const LIMIT_AND_SELECT: InputSpec[] = [
+  { name: "Limit", type: "Integer", default: "1000" },
+  { name: "Select", type: "String" },
+];
+
+const makeSpecWithDataset = () => {
+  const spec = new ComponentSpec({ $id: "spec_1", name: "Pipeline" });
+  spec.addTask(
+    new Task({
+      $id: "dataset",
+      name: "dataset",
+      componentRef: datasetRef("old-digest", LIMIT_AND_SELECT),
+      arguments: [
+        { name: "Limit", value: "86" },
+        { name: "Select", value: "tips,trip_seconds" },
+      ],
+    }),
+  );
+  return spec;
+};
+
+describe("replaceTask (edit component definition, v2)", () => {
+  it("swaps the componentRef in place, keeping task id/name and compatible arguments", () => {
+    const spec = makeSpecWithDataset();
+
+    const result = replaceTask(
+      noopUndo,
+      spec,
+      "dataset",
+      datasetRef("new-digest", LIMIT_AND_SELECT),
+    );
+
+    const task = spec.tasks.find((t) => t.$id === "dataset");
+
+    expect(spec.tasks).toHaveLength(1);
+    expect(task?.name).toBe("dataset");
+    expect(task?.componentRef.digest).toBe("new-digest");
+    expect(task?.arguments.map((a) => a.name).sort()).toEqual([
+      "Limit",
+      "Select",
+    ]);
+    expect(result.inputDiff?.lostEntities ?? []).toEqual([]);
+  });
+
+  it("drops arguments and reports lostInputs for inputs the edit removed", () => {
+    const spec = makeSpecWithDataset();
+
+    const result = replaceTask(
+      noopUndo,
+      spec,
+      "dataset",
+      datasetRef("new-digest", [{ name: "Limit", type: "Integer" }]),
+    );
+
+    const task = spec.tasks.find((t) => t.$id === "dataset");
+
+    expect((result.inputDiff?.lostEntities ?? []).map((i) => i.name)).toEqual([
+      "Select",
+    ]);
+    expect(task?.arguments.map((a) => a.name)).toEqual(["Limit"]);
+  });
+
+  it("adds arguments (with defaults) for inputs the edit introduced", () => {
+    const spec = makeSpecWithDataset();
+
+    replaceTask(noopUndo, spec, "dataset", {
+      name: "Chicago Taxi Trips dataset",
+      digest: "new-digest",
+      spec: {
+        name: "Chicago Taxi Trips dataset",
+        inputs: [
+          ...LIMIT_AND_SELECT,
+          { name: "Format", type: "String", default: "csv" },
+        ],
+        outputs: [{ name: "Table" }],
+        implementation: { container: { image: "alpine/curl" } },
+      },
+    });
+
+    const task = spec.tasks.find((t) => t.$id === "dataset");
+    const formatArg = task?.arguments.find((a) => a.name === "Format");
+
+    expect(formatArg?.value).toBe("csv");
+  });
+});

vitest-setup.js

@@ -1 +1,25 @@
 import "@testing-library/jest-dom/vitest";
+
+// Polyfill ES2024 Set methods used by app code (e.g. computeDiff in
+// task.utils.ts). They ship in every browser we target, but CI runs the test
+// suite on Node 20, which predates them. Without this, any test that exercises
+// that code throws "Set.prototype.difference is not a function".
+if (typeof Set.prototype.difference !== "function") {
+  Set.prototype.difference = function difference(other) {
+    const result = new Set();
+    for (const value of this) {
+      if (!other.has(value)) result.add(value);
+    }
+    return result;
+  };
+}
+
+if (typeof Set.prototype.intersection !== "function") {
+  Set.prototype.intersection = function intersection(other) {
+    const result = new Set();
+    for (const value of this) {
+      if (other.has(value)) result.add(value);
+    }
+    return result;
+  };
+}