fix useEffect, and have switch

Author: Light2Dark

frontend/src/components/editor/ai/ai-completion-editor.tsx

@@ -22,6 +22,7 @@ import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import { selectAllText } from "@/core/codemirror/utils";
 import { useRuntimeManager } from "@/core/runtime/config";
+import { useEvent } from "@/hooks/useEvent";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { prettyError } from "@/utils/errors";
@@ -45,7 +46,7 @@ interface Props {
   declineChange: () => void;
   acceptChange: (rightHandCode: string) => void;
   enabled: boolean;
-  initialTrigger?: boolean;
+  triggerImmediately?: boolean;
   /**
    * Children shown when there is no completion
    */
@@ -68,13 +69,10 @@ export const AiCompletionEditor: React.FC<Props> = ({
   declineChange,
   acceptChange,
   enabled,
-  initialTrigger,
+  triggerImmediately,
   children,
 }) => {
-  const [hasTriggered, setHasTriggered] = useState(false);
-  const [completionBody, setCompletionBody] = useState<object>(
-    initialPrompt ? getAICompletionBody({ input: initialPrompt }) : {},
-  );
+  const [completionBody, setCompletionBody] = useState<object>({});
 
   const [includeOtherCells, setIncludeOtherCells] = useAtom(
     includeOtherCellsAtom,
@@ -99,7 +97,11 @@ export const AiCompletionEditor: React.FC<Props> = ({
     // Throttle the messages and data updates to 100ms
     experimental_throttle: 100,
     body: {
-      ...completionBody,
+      ...(Object.keys(completionBody).length > 0
+        ? completionBody
+        : initialPrompt
+          ? getAICompletionBody({ input: initialPrompt })
+          : {}),
       includeOtherCode: includeOtherCells ? getCodes(currentCode) : "",
       code: currentCode,
       language: currentLanguageAdapter,
@@ -119,6 +121,12 @@ export const AiCompletionEditor: React.FC<Props> = ({
   const inputRef = React.useRef<ReactCodeMirrorRef>(null);
   const completion = untrimmedCompletion.trimEnd();
 
+  const initialSubmit = useEvent(() => {
+    if (triggerImmediately && !isLoading && initialPrompt) {
+      handleSubmit();
+    }
+  });
+
   // Focus the input
   useEffect(() => {
     if (enabled) {
@@ -127,6 +135,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
           const input = inputRef.current;
           if (input?.view) {
             input.view.focus();
+            initialSubmit();
             return true;
           }
           return false;
@@ -136,7 +145,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
 
       selectAllText(inputRef.current?.view);
     }
-  }, [enabled]);
+  }, [enabled, initialSubmit]);
 
   // Reset the input when the prompt changes
   useEffect(() => {
@@ -145,18 +154,6 @@ export const AiCompletionEditor: React.FC<Props> = ({
     }
   }, [enabled, initialPrompt, setInput]);
 
-  // TODO: Does not work properly
-  if (!hasTriggered && initialTrigger) {
-    setHasTriggered(true);
-    // Use requestAnimationFrame for better timing
-    requestAnimationFrame(() => {
-      if (inputRef.current?.view) {
-        storePrompt(inputRef.current.view);
-      }
-      handleSubmit();
-    });
-  }
-
   const { theme } = useTheme();
 
   const handleAcceptCompletion = () => {

frontend/src/components/editor/cell/code/cell-editor.tsx

@@ -406,7 +406,7 @@ const CellEditorInternal = ({
     <AiCompletionEditor
       enabled={aiCompletionCell?.cellId === cellId}
       initialPrompt={aiCompletionCell?.initialPrompt}
-      initialTrigger={aiCompletionCell?.triggerImmediately}
+      triggerImmediately={aiCompletionCell?.triggerImmediately}
       currentCode={editorViewRef.current?.state.doc.toString() ?? code}
       currentLanguageAdapter={languageAdapter}
       declineChange={useEvent(() => {

frontend/src/components/editor/errors/auto-fix-atom.ts

@@ -0,0 +1,18 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useAtom } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+import { store } from "@/core/state/jotai";
+
+const BASE_KEY = "marimo:instant-ai-fix";
+
+const instantAIFixAtom = atomWithStorage<boolean>(BASE_KEY, false);
+
+export function useInstantAIFix() {
+  const [instantAIFix, setInstantAIFix] = useAtom(instantAIFixAtom);
+  return { instantAIFix, setInstantAIFix };
+}
+
+export function getInstantAIFix() {
+  return store.get(instantAIFixAtom);
+}

frontend/src/components/editor/errors/auto-fix.tsx

@@ -1,8 +1,9 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { useAtomValue, useSetAtom } from "jotai";
-import { WrenchIcon, ZapIcon } from "lucide-react";
+import { WrenchIcon, ZapIcon, ZapOffIcon } from "lucide-react";
 import { Button } from "@/components/ui/button";
+import { Switch } from "@/components/ui/switch";
 import { Tooltip } from "@/components/ui/tooltip";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
@@ -12,6 +13,7 @@ import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
 import { cn } from "@/utils/cn";
+import { useInstantAIFix } from "./auto-fix-atom";
 
 export const AutoFixButton = ({
   errors,
@@ -22,6 +24,7 @@ export const AutoFixButton = ({
   cellId: CellId;
   className?: string;
 }) => {
+  const { instantAIFix, setInstantAIFix } = useInstantAIFix();
   const { createNewCell } = useCellActions();
   const aiEnabled = useAtomValue(aiEnabledAtom);
   const autoFixes = errors.flatMap((error) =>
@@ -37,7 +40,7 @@ export const AutoFixButton = ({
   // multiple fixes.
   const firstFix = autoFixes[0];
 
-  const handleFix = (aiInstantFix = false) => {
+  const handleFix = () => {
     const editorView =
       store.get(notebookAtom).cellHandles[cellId].current?.editorView;
     firstFix.onFix({
@@ -53,33 +56,48 @@ export const AutoFixButton = ({
       cellId: cellId,
       aiFix: {
         setAiCompletionCell,
-        instantFix: aiInstantFix,
+        instantFix: instantAIFix,
       },
     });
     // Focus the editor
     editorView?.focus();
   };
 
   return (
-    <div className={cn("flex gap-2 my-2", className)}>
+    <div className={cn("flex gap-2 my-2 items-center", className)}>
       <Tooltip content={firstFix.description} align="start">
         <Button
           size="xs"
           variant="outline"
           className="font-normal"
-          onClick={() => handleFix(false)}
+          onClick={handleFix}
         >
           <WrenchIcon className="h-3 w-3 mr-2" />
           {firstFix.title}
         </Button>
       </Tooltip>
 
       {firstFix.fixType === "ai" && (
-        <Tooltip content="Instant fix" align="start">
-          <Button size="xs" variant="ghost" onClick={() => handleFix(true)}>
-            <ZapIcon className="h-3 w-3" />
-          </Button>
-        </Tooltip>
+        <div className="flex items-center gap-2">
+          <Switch
+            checked={instantAIFix}
+            onCheckedChange={() => setInstantAIFix(!instantAIFix)}
+            size="sm"
+            className="h-4 w-8"
+            title="Toggle instant AI fix mode"
+          />
+          <Tooltip
+            content={
+              instantAIFix ? "Instant fix enabled" : "Instant fix disabled"
+            }
+          >
+            {instantAIFix ? (
+              <ZapIcon className="h-3 w-3 text-amber-500" />
+            ) : (
+              <ZapOffIcon className="h-3 w-3 text-muted-foreground" />
+            )}
+          </Tooltip>
+        </div>
       )}
     </div>
   );