Add stop generating button to cody web (#53191)

Add stop generating button for in progress message for cody in web. <img width="760" alt="image" src="https://github.com/sourcegraph/sourcegraph/assets/22571395/8ea16443-a125-4a73-8859-eb815cdd89e3"> https://www.loom.com/share/811352d2fb6942ff9851c9dedda281e1 ## Test plan - visit /cody - ask a question - click on stop generating button - reload the page
sourcegraph · Jun 13, 2023 · 6b6008e · 6b6008e
1 parent 7b07d55
commit 6b6008e
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 3 deletions.
diff --git a/client/cody-shared/src/chat/useClient.ts b/client/cody-shared/src/chat/useClient.ts
@@ -68,6 +68,7 @@ export interface CodyClient {
     setEditorScope: (editor: Editor) => void
     toggleIncludeInferredRepository: () => void
     toggleIncludeInferredFile: () => void
+    abortMessageInProgress: () => void
 }
 
 interface CodyClientProps {
@@ -91,6 +92,7 @@ export const useClient = ({
     const [transcript, setTranscriptState] = useState<Transcript | null>(initialTranscript)
     const [chatMessages, setChatMessagesState] = useState<ChatMessage[]>([])
     const [isMessageInProgress, setIsMessageInProgressState] = useState<boolean>(false)
+    const [abortMessageInProgressInternal, setAbortMessageInProgress] = useState<() => void>(() => () => undefined)
 
     const messageInProgress: ChatMessage | null = useMemo(() => {
         if (isMessageInProgress) {
@@ -104,6 +106,18 @@ export const useClient = ({
         return null
     }, [chatMessages, isMessageInProgress])
 
+    const abortMessageInProgress = useCallback(() => {
+        abortMessageInProgressInternal()
+
+        transcript
+            ?.toChatPromise()
+            .then(messages => {
+                setChatMessagesState(messages)
+                setIsMessageInProgressState(false)
+            })
+            .catch(error => console.error(`aborting in progress message failed: ${error}`))
+    }, [abortMessageInProgressInternal, transcript, setChatMessagesState, setIsMessageInProgressState])
+
     const setTranscript = useCallback(async (transcript: Transcript): Promise<void> => {
         const messages = await transcript.toChatPromise()
 
@@ -256,8 +270,9 @@ export const useClient = ({
 
             const responsePrefix = interaction.getAssistantMessage().prefix ?? ''
             let rawText = ''
-            return new Promise(resolve => {
-                chatClient.chat(prompt, {
+
+            const updatedTranscript = await new Promise<Transcript | null>(resolve => {
+                const abort = chatClient.chat(prompt, {
                     onChange(_rawText) {
                         rawText = _rawText
 
@@ -297,7 +312,16 @@ export const useClient = ({
                         resolve(transcript)
                     },
                 })
+
+                setAbortMessageInProgress(() => () => {
+                    abort()
+                    resolve(transcript)
+                })
             })
+
+            setAbortMessageInProgress(() => () => undefined)
+
+            return updatedTranscript
         },
         [
             config,
@@ -310,6 +334,7 @@ export const useClient = ({
             chatClient,
             isMessageInProgress,
             onEvent,
+            setAbortMessageInProgress,
         ]
     )
 
@@ -363,6 +388,7 @@ export const useClient = ({
             editMessage,
             toggleIncludeInferredRepository,
             toggleIncludeInferredFile,
+            abortMessageInProgress,
         }),
         [
             transcript,
@@ -381,6 +407,7 @@ export const useClient = ({
             editMessage,
             toggleIncludeInferredRepository,
             toggleIncludeInferredFile,
+            abortMessageInProgress,
         ]
     )
 }
diff --git a/client/cody-shared/src/sourcegraph-api/completions/browserClient.ts b/client/cody-shared/src/sourcegraph-api/completions/browserClient.ts
@@ -58,6 +58,8 @@ export class SourcegraphBrowserCompletionsClient extends SourcegraphCompletionsC
             abort.abort()
             console.error(error)
         })
-        return () => abort.abort()
+        return () => {
+            abort.abort()
+        }
     }
 }
diff --git a/client/cody-ui/src/Chat.tsx b/client/cody-ui/src/Chat.tsx
@@ -41,6 +41,8 @@ interface ChatProps extends ChatClassNames {
     setSuggestions?: (suggestions: undefined | []) => void
     needsEmailVerification?: boolean
     needsEmailVerificationNotice?: React.FunctionComponent
+    abortMessageInProgressComponent?: React.FunctionComponent<{ onAbortMessageInProgress: () => void }>
+    onAbortMessageInProgress?: () => void
 }
 
 interface ChatClassNames extends TranscriptItemClassNames {
@@ -128,6 +130,8 @@ export const Chat: React.FunctionComponent<ChatProps> = ({
     needsEmailVerificationNotice: NeedsEmailVerificationNotice,
     contextStatusComponent: ContextStatusComponent,
     contextStatusComponentProps = {},
+    abortMessageInProgressComponent,
+    onAbortMessageInProgress,
 }) => {
     const [inputRows, setInputRows] = useState(5)
     const [historyIndex, setHistoryIndex] = useState(inputHistory.length)
@@ -254,6 +258,8 @@ export const Chat: React.FunctionComponent<ChatProps> = ({
                     copyButtonOnSubmit={copyButtonOnSubmit}
                     submitButtonComponent={SubmitButton}
                     chatInputClassName={chatInputClassName}
+                    abortMessageInProgressComponent={abortMessageInProgressComponent}
+                    onAbortMessageInProgress={onAbortMessageInProgress}
                 />
             )}
 

diff --git a/client/cody-ui/src/chat/Transcript.tsx b/client/cody-ui/src/chat/Transcript.tsx
@@ -32,6 +32,8 @@ export const Transcript: React.FunctionComponent<
         feedbackButtonsOnSubmit?: (text: string) => void
         copyButtonOnSubmit?: CopyButtonProps['copyButtonOnSubmit']
         submitButtonComponent?: React.FunctionComponent<ChatUISubmitButtonProps>
+        abortMessageInProgressComponent?: React.FunctionComponent<{ onAbortMessageInProgress: () => void }>
+        onAbortMessageInProgress?: () => void
     } & TranscriptItemClassNames
 > = React.memo(function TranscriptContent({
     transcript,
@@ -54,6 +56,8 @@ export const Transcript: React.FunctionComponent<
     copyButtonOnSubmit,
     submitButtonComponent,
     chatInputClassName,
+    abortMessageInProgressComponent,
+    onAbortMessageInProgress,
 }) {
     const transcriptContainerRef = useRef<HTMLDivElement>(null)
     useEffect(() => {
@@ -141,6 +145,8 @@ export const Transcript: React.FunctionComponent<
                     copyButtonOnSubmit={copyButtonOnSubmit}
                     submitButtonComponent={submitButtonComponent}
                     chatInputClassName={chatInputClassName}
+                    abortMessageInProgressComponent={abortMessageInProgressComponent}
+                    onAbortMessageInProgress={onAbortMessageInProgress}
                 />
             )}
         </div>

diff --git a/client/cody-ui/src/chat/TranscriptItem.tsx b/client/cody-ui/src/chat/TranscriptItem.tsx
@@ -51,6 +51,8 @@ export const TranscriptItem: React.FunctionComponent<
         showFeedbackButtons: boolean
         copyButtonOnSubmit?: CopyButtonProps['copyButtonOnSubmit']
         submitButtonComponent?: React.FunctionComponent<ChatUISubmitButtonProps>
+        abortMessageInProgressComponent?: React.FunctionComponent<{ onAbortMessageInProgress: () => void }>
+        onAbortMessageInProgress?: () => void
     } & TranscriptItemClassNames
 > = React.memo(function TranscriptItemContent({
     message,
@@ -74,6 +76,8 @@ export const TranscriptItem: React.FunctionComponent<
     copyButtonOnSubmit,
     submitButtonComponent: SubmitButton,
     chatInputClassName,
+    abortMessageInProgressComponent: AbortMessageInProgressButton,
+    onAbortMessageInProgress = () => {},
 }) {
     const [formInput, setFormInput] = useState<string>(message.displayText ?? '')
     const textarea =
@@ -178,6 +182,9 @@ export const TranscriptItem: React.FunctionComponent<
                 ) : inProgress ? (
                     <BlinkingCursor />
                 ) : null}
+                {inProgress && AbortMessageInProgressButton && (
+                    <AbortMessageInProgressButton onAbortMessageInProgress={onAbortMessageInProgress} />
+                )}
             </div>
         </div>
     )

diff --git a/client/web/src/cody/components/ChatUI/ChatUi.tsx b/client/web/src/cody/components/ChatUI/ChatUi.tsx
@@ -45,6 +45,7 @@ export const ChatUI: React.FC<IChatUIProps> = ({ codyChatStore }): JSX.Element =
         setScope,
         toggleIncludeInferredRepository,
         toggleIncludeInferredFile,
+        abortMessageInProgress,
     } = codyChatStore
 
     const [formInput, setFormInput] = useState('')
@@ -102,11 +103,29 @@ export const ChatUI: React.FC<IChatUIProps> = ({ codyChatStore }): JSX.Element =
                 needsEmailVerificationNotice={NeedsEmailVerificationNotice}
                 contextStatusComponent={ScopeSelector}
                 contextStatusComponentProps={scopeSelectorProps}
+                abortMessageInProgressComponent={AbortMessageInProgress}
+                onAbortMessageInProgress={abortMessageInProgress}
             />
         </>
     )
 }
 
+interface IAbortMessageInProgressProps {
+    onAbortMessageInProgress: () => void
+}
+
+const AbortMessageInProgress: React.FunctionComponent<IAbortMessageInProgressProps> = React.memo(
+    function AbortMessageInProgressButton({ onAbortMessageInProgress }) {
+        return (
+            <div className="d-flex justify-content-center w-100 mt-4 mb-2">
+                <Button onClick={onAbortMessageInProgress} variant="secondary" outline={true} size="sm">
+                    Stop generating
+                </Button>
+            </div>
+        )
+    }
+)
+
 export const ScrollDownButton = React.memo(function ScrollDownButtonContent({
     onClick,
 }: {

diff --git a/client/web/src/cody/useCodyChat.tsx b/client/web/src/cody/useCodyChat.tsx
@@ -31,6 +31,7 @@ export interface CodyChatStore
         | 'setEditorScope'
         | 'toggleIncludeInferredRepository'
         | 'toggleIncludeInferredFile'
+        | 'abortMessageInProgress'
     > {
     readonly transcriptHistory: TranscriptJSON[]
     readonly loaded: boolean
@@ -65,6 +66,7 @@ export const codyChatStoreMock: CodyChatStore = {
     loadTranscriptFromHistory: () => Promise.resolve(),
     toggleIncludeInferredRepository: () => {},
     toggleIncludeInferredFile: () => {},
+    abortMessageInProgress: () => {},
 }
 
 interface CodyChatProps {
@@ -106,6 +108,7 @@ export const useCodyChat = ({
         setScope: setScopeInternal,
         setEditorScope,
         setTranscript,
+        abortMessageInProgress,
         toggleIncludeInferredRepository: toggleIncludeInferredRepositoryInternal,
         toggleIncludeInferredFile: toggleIncludeInferredFileInternal,
         initializeNewChat: initializeNewChatInternal,
@@ -364,6 +367,7 @@ export const useCodyChat = ({
         setEditorScope,
         toggleIncludeInferredRepository,
         toggleIncludeInferredFile,
+        abortMessageInProgress,
     }
 }