Add fix code button (#586)

* Change `useAddCell` to support source with all changes in a single transaction

Author: markmiro

src/components/notebook/cell/ExecutableCell.tsx

@@ -613,6 +613,8 @@ export const ExecutableCell: React.FC<ExecutableCellProps> = ({
 
       <ErrorBoundary FallbackComponent={OutputsErrorBoundary}>
         <MaybeCellOutputs
+          cellId={cell.id}
+          cellType={cell.cellType}
           isLoading={cell.executionState === "running" && !hasOutputs}
           outputs={hasOutputs ? outputs : staleOutputs}
           showOutput={showOutput}

src/components/outputs/MaybeCellOutputs.tsx

@@ -1,6 +1,6 @@
 import { cn } from "@/lib/utils";
 import { outputsDeltasQuery, processDeltas } from "@/queries/outputDeltas";
-import { OutputData, SAFE_MIME_TYPES } from "@runtimed/schema";
+import { CellType, OutputData, SAFE_MIME_TYPES } from "@runtimed/schema";
 import { groupConsecutiveStreamOutputs } from "@/util/output-grouping";
 import { useQuery } from "@livestore/react";
 import { useMemo, useState } from "react";
@@ -9,17 +9,22 @@ import { SingleOutput } from "./shared-with-iframe/SingleOutput";
 import { useDebounce } from "react-use";
 import { OutputsContainer } from "./shared-with-iframe/OutputsContainer";
 import { SuspenseSpinner } from "./shared-with-iframe/SuspenseSpinner";
+import { MaybeFixCodeButton } from "./MaybeFixCodeButton";
 
 /**
  * TODO: consider renaming this component
  * By default, we want to be ready to render the outputs
  */
 export const MaybeCellOutputs = ({
+  cellId,
+  cellType,
   outputs,
   shouldAlwaysUseIframe = false,
   isLoading,
   showOutput,
 }: {
+  cellId: string;
+  cellType: CellType;
   outputs: readonly OutputData[];
   shouldAlwaysUseIframe?: boolean;
   isLoading: boolean;
@@ -54,6 +59,13 @@ export const MaybeCellOutputs = ({
           isLoading && outputs.length ? "opacity-50" : "opacity-100"
         )}
       >
+        {cellType === "code" && (
+          <MaybeFixCodeButton
+            isLoading={isLoading}
+            cellId={cellId}
+            outputs={outputs}
+          />
+        )}
         {/* TODO: consider rendering an empty iframewhen we have a safe output currently rendered but cell input has changed */}
         {shouldUseIframe ? (
           <IframeOutput

src/components/outputs/MaybeFixCodeButton.tsx

@@ -0,0 +1,75 @@
+import { useAddCell } from "@/hooks/useAddCell";
+import { OutputData } from "@runtimed/schema";
+import { Bug } from "lucide-react";
+import { useCallback } from "react";
+import { Button } from "../ui/button";
+import { OutputsContainer } from "./shared-with-iframe/OutputsContainer";
+
+export function MaybeFixCodeButton({
+  isLoading,
+  cellId,
+  outputs,
+}: {
+  isLoading: boolean;
+  cellId: string;
+  outputs: readonly OutputData[];
+}) {
+  const errorOutput = outputs.find((output) => output.outputType === "error");
+
+  if (
+    !errorOutput ||
+    !errorOutput.data ||
+    typeof errorOutput.data !== "string"
+  ) {
+    return null;
+  }
+
+  return (
+    <OutputsContainer>
+      <FixCodeButton
+        cellId={cellId}
+        errorOutputData={errorOutput.data}
+        isLoading={isLoading}
+      />
+    </OutputsContainer>
+  );
+}
+
+const FixCodeButton = ({
+  cellId,
+  errorOutputData,
+  isLoading,
+}: {
+  cellId: string;
+  errorOutputData: string;
+  isLoading: boolean;
+}) => {
+  const { addCell } = useAddCell();
+
+  const handleFixCode = useCallback(() => {
+    const formattedErrorOutputData = formatErrorOutputData(errorOutputData);
+    const message = formatAiMessage(formattedErrorOutputData);
+    addCell(cellId, "ai", "after", message);
+  }, [addCell, cellId, errorOutputData]);
+
+  return (
+    <Button
+      variant="destructive"
+      size="sm"
+      onClick={handleFixCode}
+      disabled={isLoading}
+    >
+      <Bug /> Fix Code with AI{" "}
+    </Button>
+  );
+};
+
+/** Add code block to the error message */
+function formatAiMessage(errorString: string) {
+  return `Fix the error:\n \`\`\`json\n${errorString}\n\`\`\`\n`;
+}
+
+/** Formats the error for AI model to fix */
+function formatErrorOutputData(errorOutputData: string) {
+  return JSON.stringify(JSON.parse(errorOutputData), null, 2);
+}

src/hooks/useAddCell.ts

@@ -18,7 +18,8 @@ export const useAddCell = () => {
     (
       cellId?: string,
       cellType: CellType = "code",
-      position: "before" | "after" = "after"
+      position: "before" | "after" = "after",
+      source?: string
     ) => {
       const cellReferences = store.query(queries.cellsWithIndices$);
       const newCellId = `cell-${Date.now()}-${Math.random()
@@ -99,6 +100,34 @@ export const useAddCell = () => {
         );
       }
 
+      if (source) {
+        const cellEvents = [
+          // set cell source
+          events.cellSourceChanged({
+            id: newCellId,
+            source,
+            modifiedBy: userId,
+          }),
+          // hide cell input
+          events.cellSourceVisibilityToggled({
+            id: newCellId,
+            sourceVisible: false,
+            actorId: userId,
+          }),
+          // run cell
+          events.executionRequested({
+            cellId: newCellId,
+            actorId: userId,
+            requestedBy: userId,
+            queueId: `exec-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+            executionCount:
+              (store.query(queries.cellQuery.byId(newCellId))?.executionCount ||
+                0) + 1,
+          }),
+        ];
+        store.commit(...cellEvents);
+      }
+
       // Focus the new cell after creation
       setTimeout(() => store.setSignal(focusedCellSignal$, newCellId), 0);