fix: Edit Component Definition now updates the selected task

src/components/shared/ComponentEditor/ComponentEditorDialog.test.tsx

@@ -337,5 +337,54 @@ describe("<ComponentEditorDialog />", () => {
         expect(onCloseMock).toHaveBeenCalledTimes(1);
       });
     });
+
+    test("calls onComponentSaved instead of importing to the library when provided", async () => {
+      const onCloseMock = vi.fn();
+      const onComponentSavedMock = vi.fn();
+      const mockHydratedComponent = {
+        spec: {
+          implementation: { container: { image: "test" } },
+          metadata: { annotations: {} },
+        },
+        name: "test-component",
+        digest: "abc123",
+        text: "name: test-component",
+      };
+
+      vi.mocked(hydrateComponentReference).mockResolvedValue(
+        mockHydratedComponent,
+      );
+
+      renderWithProviders(
+        <ComponentEditorDialog
+          text="name: test-component"
+          onClose={onCloseMock}
+          onComponentSaved={onComponentSavedMock}
+        />,
+      );
+
+      await waitFor(() => {
+        expect(
+          screen.getByRole("button", { name: /Save/i }),
+        ).toBeInTheDocument();
+      });
+
+      fireEvent.click(screen.getByRole("button", { name: /Save/i }));
+
+      await waitFor(() => {
+        // The edited component is handed to the caller to apply.
+        expect(onComponentSavedMock).toHaveBeenCalledWith(
+          mockHydratedComponent,
+        );
+        expect(onCloseMock).toHaveBeenCalledTimes(1);
+      });
+
+      // It must NOT fall back to the library-import behavior.
+      expect(mockAddToComponentLibrary).not.toHaveBeenCalled();
+      expect(mockToast).not.toHaveBeenCalledWith(
+        expect.stringContaining("imported successfully"),
+        "success",
+      );
+    });
   });
 });

src/components/shared/ComponentEditor/ComponentEditorDialog.tsx

@@ -11,7 +11,10 @@ import useToastNotification from "@/hooks/useToastNotification";
 import { useAnalytics } from "@/providers/AnalyticsProvider";
 import { useComponentLibrary } from "@/providers/ComponentLibraryProvider";
 import { hydrateComponentReference } from "@/services/componentService";
-import { isContainerImplementation } from "@/utils/componentSpec";
+import {
+  type HydratedComponentReference,
+  isContainerImplementation,
+} from "@/utils/componentSpec";
 import { saveComponent } from "@/utils/localforage";
 
 import { FullscreenElement } from "../FullscreenElement";
@@ -72,10 +75,21 @@ export const ComponentEditorDialog = withSuspenseWrapper(
     text,
     templateName = "empty",
     onClose,
+    onComponentSaved,
   }: {
     text?: string;
     templateName?: SupportedTemplate;
     onClose: () => void;
+    /**
+     * When provided, the editor is being used to edit an existing target (e.g.
+     * a selected task's component) rather than to import a brand new component
+     * into the library. The callback receives the hydrated, edited component
+     * and is responsible for applying it (and any user feedback). When omitted,
+     * the editor falls back to importing the component into the library.
+     */
+    onComponentSaved?: (
+      hydratedComponent: HydratedComponentReference,
+    ) => void | Promise<void>;
   }) => {
     const notify = useToastNotification();
     const { track } = useAnalytics();
@@ -190,16 +204,23 @@ export const ComponentEditorDialog = withSuspenseWrapper(
 
       onClose();
 
-      await addToComponentLibrary(hydratedComponent, "editor_save");
+      if (onComponentSaved) {
+        // Editing an existing target (e.g. a selected task): apply the edit to
+        // that target instead of importing a new library component. The caller
+        // owns the success/feedback messaging.
+        await onComponentSaved(hydratedComponent);
+      } else {
+        await addToComponentLibrary(hydratedComponent, "editor_save");
+        notify(
+          `Component ${hydratedComponent.name} imported successfully`,
+          "success",
+        );
+      }
 
       track("component_editor.save.completed", {
         mode,
         selected_template: templateName,
       });
-      notify(
-        `Component ${hydratedComponent.name} imported successfully`,
-        "success",
-      );
     };
 
     const handleClose = () => {

src/components/shared/ReactFlow/FlowCanvas/utils/replaceTaskComponentRef.test.ts

@@ -0,0 +1,163 @@
+import { describe, expect, it } from "vitest";
+
+import type { ComponentReference, GraphSpec } from "@/utils/componentSpec";
+
+import { replaceTaskComponentRef } from "./replaceTaskComponentRef";
+
+const taxiRef = (command: string): ComponentReference => ({
+  name: "Chicago Taxi Trips dataset",
+  digest: "old-digest",
+  spec: {
+    name: "Chicago Taxi Trips dataset",
+    inputs: [
+      { name: "Limit", type: "Integer", default: "1000" },
+      { name: "Select", type: "String" },
+    ],
+    outputs: [{ name: "Table" }],
+    implementation: { container: { image: "alpine/curl", command: [command] } },
+  },
+});
+
+const baseGraphSpec = (): GraphSpec => ({
+  tasks: {
+    dataset: {
+      componentRef: taxiRef("sed -E 's/x/'\\1'/g'"),
+      annotations: { "editor.position": '{"x":10,"y":20}' },
+      arguments: {
+        Limit: { graphInput: { inputName: "Input" } },
+        Select: "tips,trip_seconds",
+      },
+    },
+    train: {
+      componentRef: { name: "Train" },
+      arguments: {
+        training_data: {
+          taskOutput: { taskId: "dataset", outputName: "Table" },
+        },
+      },
+    },
+  },
+});
+
+describe("replaceTaskComponentRef", () => {
+  it("swaps the componentRef in place without renaming the task", () => {
+    const graphSpec = baseGraphSpec();
+    const fixedRef: ComponentReference = {
+      ...taxiRef("sed -E 's/x/\\1/g'"),
+      digest: "new-digest",
+    };
+
+    const { updatedGraphSpec, lostInputs } = replaceTaskComponentRef(
+      "dataset",
+      fixedRef,
+      graphSpec,
+    );
+
+    // Same key, no "dataset 2".
+    expect(Object.keys(updatedGraphSpec.tasks)).toEqual(["dataset", "train"]);
+    expect(updatedGraphSpec.tasks.dataset.componentRef.digest).toBe(
+      "new-digest",
+    );
+    expect(lostInputs).toEqual([]);
+  });
+
+  it("preserves the task's arguments, annotations and downstream wiring", () => {
+    const graphSpec = baseGraphSpec();
+    const fixedRef = { ...taxiRef("fixed"), digest: "new-digest" };
+
+    const { updatedGraphSpec } = replaceTaskComponentRef(
+      "dataset",
+      fixedRef,
+      graphSpec,
+    );
+
+    expect(updatedGraphSpec.tasks.dataset.arguments).toEqual({
+      Limit: { graphInput: { inputName: "Input" } },
+      Select: "tips,trip_seconds",
+    });
+    expect(updatedGraphSpec.tasks.dataset.annotations).toEqual({
+      "editor.position": '{"x":10,"y":20}',
+    });
+    // Downstream consumer still wired to the (unchanged) output/taskId.
+    expect(updatedGraphSpec.tasks.train.arguments).toEqual({
+      training_data: {
+        taskOutput: { taskId: "dataset", outputName: "Table" },
+      },
+    });
+  });
+
+  it("does not mutate the original graph spec", () => {
+    const graphSpec = baseGraphSpec();
+    const fixedRef = { ...taxiRef("fixed"), digest: "new-digest" };
+
+    replaceTaskComponentRef("dataset", fixedRef, graphSpec);
+
+    expect(graphSpec.tasks.dataset.componentRef.digest).toBe("old-digest");
+  });
+
+  it("drops arguments for inputs the edited component no longer defines", () => {
+    const graphSpec = baseGraphSpec();
+    const refWithoutSelect: ComponentReference = {
+      name: "Chicago Taxi Trips dataset",
+      digest: "new-digest",
+      spec: {
+        name: "Chicago Taxi Trips dataset",
+        inputs: [{ name: "Limit", type: "Integer" }],
+        outputs: [{ name: "Table" }],
+        implementation: { container: { image: "alpine/curl" } },
+      },
+    };
+
+    const { updatedGraphSpec, lostInputs } = replaceTaskComponentRef(
+      "dataset",
+      refWithoutSelect,
+      graphSpec,
+    );
+
+    expect(lostInputs.map((input) => input.name)).toEqual(["Select"]);
+    expect(updatedGraphSpec.tasks.dataset.arguments).toEqual({
+      Limit: { graphInput: { inputName: "Input" } },
+    });
+  });
+
+  it("drops downstream bindings to outputs the edited component no longer produces", () => {
+    const graphSpec = baseGraphSpec();
+    const refWithoutTable: ComponentReference = {
+      name: "Chicago Taxi Trips dataset",
+      digest: "new-digest",
+      spec: {
+        name: "Chicago Taxi Trips dataset",
+        inputs: [
+          { name: "Limit", type: "Integer" },
+          { name: "Select", type: "String" },
+        ],
+        outputs: [{ name: "RenamedTable" }],
+        implementation: { container: { image: "alpine/curl" } },
+      },
+    };
+
+    const { updatedGraphSpec } = replaceTaskComponentRef(
+      "dataset",
+      refWithoutTable,
+      graphSpec,
+    );
+
+    expect(updatedGraphSpec.tasks.train.arguments).toEqual({});
+  });
+
+  it("returns the graph unchanged when the task does not exist", () => {
+    const graphSpec = baseGraphSpec();
+    const fixedRef = { ...taxiRef("fixed"), digest: "new-digest" };
+
+    const { updatedGraphSpec, lostInputs } = replaceTaskComponentRef(
+      "missing",
+      fixedRef,
+      graphSpec,
+    );
+
+    expect(lostInputs).toEqual([]);
+    expect(updatedGraphSpec.tasks.dataset.componentRef.digest).toBe(
+      "old-digest",
+    );
+  });
+});

src/components/shared/ReactFlow/FlowCanvas/utils/replaceTaskComponentRef.ts

@@ -0,0 +1,100 @@
+import type {
+  ArgumentType,
+  ComponentReference,
+  GraphSpec,
+  InputSpec,
+} from "@/utils/componentSpec";
+import { deepClone } from "@/utils/deepClone";
+
+/**
+ * Replaces a task's componentRef in place (keeping the same taskId, position,
+ * annotations and execution options) while preserving as much of the task's
+ * existing wiring as possible.
+ *
+ * This is the right primitive for "Edit Component Definition": the user has
+ * selected a task and edited its component, so their intent is to update that
+ * specific task and keep everything they already configured.
+ *
+ * Unlike `replaceTaskNode` (used by "Upgrade"), this does NOT rename the task:
+ * renaming would re-key the task to a new unique id derived from the component
+ * name and is surprising for an in-place edit.
+ *
+ * Wiring is preserved with two safe exceptions, mirroring the upgrade flow:
+ *  - argument bindings on this task for inputs that no longer exist in the
+ *    edited component are dropped (and reported via `lostInputs`);
+ *  - argument bindings on downstream tasks that consume an output this
+ *    component no longer produces are dropped, to avoid dangling references.
+ *
+ * In the common case (editing the command/image, no input/output changes)
+ * nothing is dropped and every argument and default is preserved.
+ */
+export const replaceTaskComponentRef = (
+  taskId: string,
+  newComponentRef: ComponentReference,
+  graphSpec: GraphSpec,
+) => {
+  const updatedGraphSpec = deepClone(graphSpec);
+  const task = updatedGraphSpec.tasks[taskId];
+
+  if (!task) {
+    return { updatedGraphSpec, lostInputs: [] as InputSpec[] };
+  }
+
+  const oldInputs = task.componentRef.spec?.inputs ?? [];
+  const newInputNames = new Set(
+    (newComponentRef.spec?.inputs ?? []).map((input) => input.name),
+  );
+  const newOutputNames = new Set(
+    (newComponentRef.spec?.outputs ?? []).map((output) => output.name),
+  );
+
+  // Inputs present on the old component but missing from the edited one.
+  const lostInputs = oldInputs.filter(
+    (input) => !newInputNames.has(input.name),
+  );
+
+  task.componentRef = newComponentRef;
+
+  // Drop this task's argument bindings for inputs that were removed.
+  if (task.arguments) {
+    task.arguments = Object.entries(task.arguments).reduce(
+      (acc, [inputName, argument]) => {
+        if (newInputNames.has(inputName)) {
+          acc[inputName] = argument;
+        }
+        return acc;
+      },
+      {} as Record<string, ArgumentType>,
+    );
+  }
+
+  // Drop downstream argument bindings that consume an output this component no
+  // longer produces. Outputs that still exist keep working because the taskId
+  // is unchanged.
+  Object.values(updatedGraphSpec.tasks).forEach((downstreamTask) => {
+    if (!downstreamTask.arguments) {
+      return;
+    }
+
+    downstreamTask.arguments = Object.entries(downstreamTask.arguments).reduce(
+      (acc, [inputName, argument]) => {
+        if (
+          typeof argument === "object" &&
+          argument !== null &&
+          "taskOutput" in argument &&
+          argument.taskOutput &&
+          argument.taskOutput.taskId === taskId &&
+          !newOutputNames.has(argument.taskOutput.outputName)
+        ) {
+          return acc;
+        }
+
+        acc[inputName] = argument;
+        return acc;
+      },
+      {} as Record<string, ArgumentType>,
+    );
+  });
+
+  return { updatedGraphSpec, lostInputs } as const;
+};

src/components/shared/TaskDetails/Actions.tsx

@@ -56,7 +56,7 @@ const TaskActions = ({
     />
   );
   const editComponent = !readOnly && (
-    <EditComponentButton componentRef={componentRef} />
+    <EditComponentButton componentRef={componentRef} taskId={taskId} />
   );
 
   // Canvas Actions