From e27b13ad37ac2b01bfe0ae5e001e454d5b825b46 Mon Sep 17 00:00:00 2001
From: Stef Lewandowski <stef@stef.io>
Date: Fri, 27 Sep 2024 10:27:46 +0100
Subject: [PATCH] fix: stop interaction except in moderation or idle (#165)

---
 .../AppComponents/Chat/chat-panel.tsx         | 13 +++-
 .../AppComponents/Chat/chat-quick-buttons.tsx | 73 ++++---------------
 .../AppComponents/Chat/prompt-form.tsx        | 57 +++++----------
 .../ContextProviders/ChatProvider.tsx         | 54 ++++++++++++++
 4 files changed, 97 insertions(+), 100 deletions(-)

diff --git a/apps/nextjs/src/components/AppComponents/Chat/chat-panel.tsx b/apps/nextjs/src/components/AppComponents/Chat/chat-panel.tsx
index bc7ca6739..fca32ed91 100644
--- a/apps/nextjs/src/components/AppComponents/Chat/chat-panel.tsx
+++ b/apps/nextjs/src/components/AppComponents/Chat/chat-panel.tsx
@@ -23,7 +23,16 @@ export function ChatPanel({
   isDemoLocked,
 }: Readonly<ChatPanelProps>) {
   const chat = useLessonChat();
-  const { id, isLoading, input, setInput, append, ailaStreamingStatus } = chat;
+  const {
+    id,
+    isLoading,
+    input,
+    setInput,
+    append,
+    ailaStreamingStatus,
+    queueUserAction,
+    queuedUserAction,
+  } = chat;
 
   const { trackEvent } = useAnalytics();
   const containerClass = `grid w-full grid-cols-1 ${isEmptyScreen ? "sm:grid-cols-1" : ""} peer-[[data-state=open]]:group-[]:lg:pl-[250px] peer-[[data-state=open]]:group-[]:xl:pl-[300px]`;
@@ -49,6 +58,8 @@ export function ChatPanel({
             setInput={setInput}
             ailaStreamingStatus={ailaStreamingStatus}
             isEmptyScreen={isEmptyScreen}
+            queueUserAction={queueUserAction}
+            queuedUserAction={queuedUserAction}
           />
         )}
         {isDemoLocked && <LockedPromptForm />}
diff --git a/apps/nextjs/src/components/AppComponents/Chat/chat-quick-buttons.tsx b/apps/nextjs/src/components/AppComponents/Chat/chat-quick-buttons.tsx
index f8cba9df9..06e9cada7 100644
--- a/apps/nextjs/src/components/AppComponents/Chat/chat-quick-buttons.tsx
+++ b/apps/nextjs/src/components/AppComponents/Chat/chat-quick-buttons.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useState, useRef } from "react";
+import { useCallback } from "react";
 
 import { useLessonChat } from "@/components/ContextProviders/ChatProvider";
 import { Icon } from "@/components/Icon";
@@ -13,75 +13,36 @@ interface QuickActionButtonsProps {
   isEmptyScreen: boolean;
 }
 
-type QueuedUserAction = "regenerate" | "continue";
-
 const QuickActionButtons = ({ isEmptyScreen }: QuickActionButtonsProps) => {
   const chat = useLessonChat();
   const { trackEvent } = useAnalytics();
   const lessonPlanTracking = useLessonPlanTracking();
   const { setDialogWindow } = useDialog();
-  const [queuedUserAction, setQueuedUserAction] =
-    useState<QueuedUserAction | null>(null);
-  const { messages, reload, append, id, stop, ailaStreamingStatus } = chat;
-  const isExecutingAction = useRef(false);
+  const {
+    messages,
+    id,
+    stop,
+    ailaStreamingStatus,
+    queueUserAction,
+    queuedUserAction,
+  } = chat;
 
   const shouldAllowUserAction =
     ["Idle", "Moderating"].includes(ailaStreamingStatus) && !queuedUserAction;
-  const shouldQueueUserAction = ailaStreamingStatus === "Moderating";
 
   const handleRegenerate = useCallback(() => {
     trackEvent("chat:regenerate", { id: id });
     const lastUserMessage =
       messages.findLast((m) => m.role === "user")?.content || "";
     lessonPlanTracking.onClickRetry(lastUserMessage);
-    reload();
-  }, [reload, lessonPlanTracking, messages, trackEvent, id]);
+    queueUserAction("regenerate");
+  }, [queueUserAction, lessonPlanTracking, messages, trackEvent, id]);
 
   const handleContinue = useCallback(async () => {
     trackEvent("chat:continue");
     lessonPlanTracking.onClickContinue();
-    await append({
-      content: "Continue",
-      role: "user",
-    });
-  }, [append, lessonPlanTracking, trackEvent]);
-
-  const queueUserAction = useCallback(
-    (action: QueuedUserAction) => {
-      setQueuedUserAction(action);
-    },
-    [setQueuedUserAction],
-  );
-
-  const executeQueuedAction = useCallback(async () => {
-    if (!queuedUserAction || shouldQueueUserAction || isExecutingAction.current)
-      return;
-
-    isExecutingAction.current = true;
-    const actionToExecute = queuedUserAction;
-    setQueuedUserAction(null);
-
-    try {
-      if (actionToExecute === "regenerate") {
-        handleRegenerate();
-      } else if (actionToExecute === "continue") {
-        await handleContinue();
-      }
-    } catch (error) {
-      console.error("Error handling queued action:", error);
-    } finally {
-      isExecutingAction.current = false;
-    }
-  }, [
-    queuedUserAction,
-    shouldQueueUserAction,
-    handleRegenerate,
-    handleContinue,
-  ]);
-
-  useEffect(() => {
-    executeQueuedAction();
-  }, [executeQueuedAction]);
+    queueUserAction("continue");
+  }, [queueUserAction, lessonPlanTracking, trackEvent]);
 
   return (
     <div className="-ml-7 flex justify-between space-x-7 rounded-bl rounded-br pt-8">
@@ -93,10 +54,6 @@ const QuickActionButtons = ({ isEmptyScreen }: QuickActionButtonsProps) => {
               size="sm"
               variant="text-link"
               onClick={() => {
-                if (shouldQueueUserAction) {
-                  queueUserAction("regenerate");
-                  return;
-                }
                 handleRegenerate();
               }}
             >
@@ -145,10 +102,6 @@ const QuickActionButtons = ({ isEmptyScreen }: QuickActionButtonsProps) => {
         variant="primary"
         disabled={!shouldAllowUserAction}
         onClick={async () => {
-          if (shouldQueueUserAction) {
-            queueUserAction("continue");
-            return;
-          }
           await handleContinue();
         }}
         testId="chat-continue"
diff --git a/apps/nextjs/src/components/AppComponents/Chat/prompt-form.tsx b/apps/nextjs/src/components/AppComponents/Chat/prompt-form.tsx
index 1bd8d3974..80837d9fa 100644
--- a/apps/nextjs/src/components/AppComponents/Chat/prompt-form.tsx
+++ b/apps/nextjs/src/components/AppComponents/Chat/prompt-form.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useState, useEffect, useRef } from "react";
+import { useCallback, useEffect, useRef } from "react";
 import Textarea from "react-textarea-autosize";
 
 import { UseChatHelpers } from "ai/react";
@@ -21,6 +21,8 @@ export interface PromptFormProps
   isEmptyScreen: boolean;
   placeholder?: string;
   ailaStreamingStatus: AilaStreamingStatus;
+  queuedUserAction?: string | null;
+  queueUserAction?: (action: string) => void;
 }
 
 export function PromptForm({
@@ -30,12 +32,12 @@ export function PromptForm({
   setInput,
   isEmptyScreen,
   placeholder,
+  queuedUserAction,
+  queueUserAction,
 }: Readonly<PromptFormProps>) {
   const { formRef, onKeyDown } = useEnterSubmit();
   const inputRef = useRef<HTMLTextAreaElement>(null);
   const lessonPlanTracking = useLessonPlanTracking();
-  const [queuedUserInput, setQueuedUserInput] = useState<string | null>(null);
-  const timeoutRef = useRef<number | null>(null);
 
   useEffect(() => {
     if (inputRef.current) {
@@ -45,22 +47,6 @@ export function PromptForm({
 
   const sidebar = useSidebar();
 
-  const queueSubmission = useCallback(
-    (value: string) => {
-      setQueuedUserInput(value);
-      timeoutRef.current = window.setTimeout(() => {
-        setQueuedUserInput(null);
-      }, 10000);
-
-      return () => {
-        if (timeoutRef.current !== null) {
-          window.clearTimeout(timeoutRef.current);
-        }
-      };
-    },
-    [timeoutRef],
-  );
-
   const handleSubmit = useCallback(
     async (value: string) => {
       setInput("");
@@ -69,24 +55,17 @@ export function PromptForm({
       }
 
       lessonPlanTracking.onSubmitText(value);
-      onSubmit(value);
+      if (queueUserAction) {
+        queueUserAction(value);
+      } else {
+        onSubmit(value);
+      }
     },
-    [lessonPlanTracking, onSubmit, setInput, sidebar],
+    [lessonPlanTracking, queueUserAction, onSubmit, setInput, sidebar],
   );
 
-  useEffect(() => {
-    if (ailaStreamingStatus === "Idle" && queuedUserInput) {
-      handleSubmit(queuedUserInput);
-      setQueuedUserInput(null);
-      if (timeoutRef.current !== null) {
-        window.clearTimeout(timeoutRef.current);
-      }
-    }
-  }, [ailaStreamingStatus, handleSubmit, queuedUserInput]);
-
   const shouldAllowUserInput =
-    ["Idle", "Moderating"].includes(ailaStreamingStatus) || queuedUserInput;
-  const shouldQueueUserInput = ailaStreamingStatus === "Moderating";
+    ["Idle", "Moderating"].includes(ailaStreamingStatus) && !queuedUserAction;
 
   return (
     <form
@@ -95,10 +74,6 @@ export function PromptForm({
         if (!input?.trim()) {
           return;
         }
-        if (shouldQueueUserInput) {
-          queueSubmission(input);
-          return;
-        }
         handleSubmit(input);
       }}
       ref={formRef}
@@ -111,13 +86,17 @@ export function PromptForm({
       >
         <Textarea
           data-testid="chat-input"
+          disabled={!shouldAllowUserInput}
           ref={inputRef}
           tabIndex={0}
           onKeyDown={onKeyDown}
           rows={1}
           value={input}
           onChange={(e) => setInput(e.target.value)}
-          placeholder={handlePlaceholder(isEmptyScreen, placeholder)}
+          placeholder={handlePlaceholder(
+            isEmptyScreen,
+            queuedUserAction ?? placeholder,
+          )}
           spellCheck={false}
           className="min-h-[60px] w-full resize-none bg-transparent px-10 py-[1.3rem] text-base focus-within:outline-none"
         />
@@ -142,7 +121,7 @@ export function PromptForm({
 }
 
 function handlePlaceholder(isEmptyScreen: boolean, placeholder?: string) {
-  if (placeholder) {
+  if (placeholder && !["continue", "regenerate"].includes(placeholder)) {
     return placeholder;
   }
   return !isEmptyScreen
diff --git a/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx b/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx
index cbfdd6f72..94d3cf30f 100644
--- a/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx
+++ b/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx
@@ -1,5 +1,6 @@
 import React, {
   createContext,
+  useCallback,
   useContext,
   useEffect,
   useMemo,
@@ -58,6 +59,9 @@ export type ChatContextProps = {
   input: string;
   setInput: React.Dispatch<React.SetStateAction<string>>;
   chatAreaRef: React.RefObject<HTMLDivElement>;
+  queuedUserAction: string | null;
+  queueUserAction: (action: string) => void;
+  executeQueuedAction: () => Promise<void>;
 };
 
 const ChatContext = createContext<ChatContextProps | null>(null);
@@ -245,6 +249,50 @@ export function ChatProvider({ id, children }: Readonly<ChatProviderProps>) {
       messageHashes,
     });
 
+  // Handle queued user actions and messages
+
+  const [queuedUserAction, setQueuedUserAction] = useState<string | null>(null);
+  const isExecutingAction = useRef(false);
+
+  const queueUserAction = useCallback((action: string) => {
+    setQueuedUserAction(action);
+  }, []);
+
+  const executeQueuedAction = useCallback(async () => {
+    if (!queuedUserAction || !hasFinished || isExecutingAction.current) return;
+
+    isExecutingAction.current = true;
+    const actionToExecute = queuedUserAction;
+    setQueuedUserAction(null);
+
+    try {
+      if (actionToExecute === "continue") {
+        await append({
+          content: "Continue",
+          role: "user",
+        });
+      } else if (actionToExecute === "regenerate") {
+        reload();
+      } else {
+        // Assume it's a user message
+        await append({
+          content: actionToExecute,
+          role: "user",
+        });
+      }
+    } catch (error) {
+      console.error("Error handling queued action:", error);
+    } finally {
+      isExecutingAction.current = false;
+    }
+  }, [queuedUserAction, hasFinished, append, reload]);
+
+  useEffect(() => {
+    if (hasFinished) {
+      executeQueuedAction();
+    }
+  }, [hasFinished, executeQueuedAction]);
+
   /**
    *  If the state is being restored from a previous lesson plan, set the lesson plan
    */
@@ -330,6 +378,9 @@ export function ChatProvider({ id, children }: Readonly<ChatProviderProps>) {
       setInput,
       partialPatches,
       validPatches,
+      queuedUserAction,
+      queueUserAction,
+      executeQueuedAction,
     }),
     [
       id,
@@ -352,6 +403,9 @@ export function ChatProvider({ id, children }: Readonly<ChatProviderProps>) {
       partialPatches,
       validPatches,
       overrideLessonPlan,
+      queuedUserAction,
+      queueUserAction,
+      executeQueuedAction,
     ],
   );