fix: only handle the query generation keydown event when not streaming

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.25.2",
+    "version": "3.25.3",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

querybook/webapp/components/AIAssistant/AutoFixButton.scss

@@ -0,0 +1,7 @@
+.AutoFixModal {
+    .Modal-box {
+        .Modal-content {
+            min-height: 40vh;
+        }
+    }
+}

querybook/webapp/components/AIAssistant/AutoFixButton.tsx

@@ -9,6 +9,8 @@ import { Message } from 'ui/Message/Message';
 import { Modal } from 'ui/Modal/Modal';
 import { AccentText } from 'ui/StyledText/StyledText';
 
+import './AutoFixButton.scss';
+
 interface IProps {
     query: string;
     queryExecutionId: number;
@@ -22,7 +24,7 @@ export const AutoFixButton = ({
 }: IProps) => {
     const [show, setShow] = useState(false);
 
-    const { streamStatus, startStream, streamData } = useStream(
+    const { streamStatus, startStream, streamData, cancelStream } = useStream(
         '/ds/ai/query_auto_fix/',
         {
             query_execution_id: queryExecutionId,
@@ -35,37 +37,53 @@ export const AutoFixButton = ({
         fixed_query: fixedQuery,
     } = streamData;
 
-    const bottomDOM = fixedQuery && (
-        <div className="right-align mb16">
-            <Button
-                title="Reject"
-                color="cancel"
-                onClick={() => {
-                    setShow(false);
-                }}
-            />
-            <Button
-                title="Apply"
-                color="confirm"
-                onClick={() => {
-                    onUpdateQuery?.(fixedQuery);
-                    trackClick({
-                        component: ComponentType.AI_ASSISTANT,
-                        element: ElementType.QUERY_ERROR_AUTO_FIX_APPLY_BUTTON,
-                    });
-                    setShow(false);
-                }}
-            />
-        </div>
-    );
+    const bottomDOM =
+        streamStatus === StreamStatus.STREAMING ? (
+            <div className="right-align mb16">
+                <Button
+                    title="Stop Generating"
+                    color="light"
+                    onClick={cancelStream}
+                    className="mr8"
+                />
+            </div>
+        ) : (
+            fixedQuery && (
+                <div className="right-align mb16">
+                    <Button
+                        title="Reject"
+                        color="cancel"
+                        onClick={() => {
+                            setShow(false);
+                        }}
+                    />
+                    <Button
+                        title="Apply"
+                        color="confirm"
+                        onClick={() => {
+                            onUpdateQuery?.(fixedQuery);
+                            trackClick({
+                                component: ComponentType.AI_ASSISTANT,
+                                element:
+                                    ElementType.QUERY_ERROR_AUTO_FIX_APPLY_BUTTON,
+                            });
+                            setShow(false);
+                        }}
+                    />
+                </div>
+            )
+        );
     return (
         <>
             <Button
                 icon="Bug"
                 title="Auto fix"
                 onClick={() => {
                     setShow(true);
-                    if (streamStatus === StreamStatus.NOT_STARTED) {
+                    if (
+                        streamStatus === StreamStatus.NOT_STARTED ||
+                        streamStatus === StreamStatus.CANCELLED
+                    ) {
                         startStream();
                         trackClick({
                             component: ComponentType.AI_ASSISTANT,
@@ -77,9 +95,11 @@ export const AutoFixButton = ({
             {show && (
                 <Modal
                     onHide={() => {
+                        cancelStream();
                         setShow(false);
                     }}
                     bottomDOM={bottomDOM}
+                    className="AutoFixModal"
                 >
                     <Message
                         message="Note: This AI-powered auto fix may not be 100% accurate. Please use your own judgement and verify the result."

querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx

@@ -76,7 +76,7 @@ export const QueryGenerationModal = ({
         setTables(uniq([...tablesInQuery, ...tables]));
     }, [tablesInQuery]);
 
-    const { streamStatus, startStream, streamData } = useStream(
+    const { streamStatus, startStream, streamData, cancelStream } = useStream(
         '/ds/ai/generate_query/',
         {
             query_engine_id: engineId,
@@ -91,7 +91,10 @@ export const QueryGenerationModal = ({
 
     const onKeyDown = useCallback(
         (event: React.KeyboardEvent) => {
-            if (matchKeyPress(event, 'Enter')) {
+            if (
+                streamStatus !== StreamStatus.STREAMING &&
+                matchKeyPress(event, 'Enter')
+            ) {
                 startStream();
                 inputRef.current.blur();
                 trackClick({
@@ -105,7 +108,7 @@ export const QueryGenerationModal = ({
                 });
             }
         },
-        [startStream]
+        [streamStatus, startStream]
     );
 
     const questionBarDOM = (
@@ -148,6 +151,14 @@ export const QueryGenerationModal = ({
                     autoFocus: true,
                 }}
             />
+            {streamStatus === StreamStatus.STREAMING && (
+                <Button
+                    title="Stop Generating"
+                    color="light"
+                    onClick={cancelStream}
+                    className="mr8"
+                />
+            )}
         </div>
     );
 
@@ -183,7 +194,10 @@ export const QueryGenerationModal = ({
 
     return (
         <Modal
-            onHide={onHide}
+            onHide={() => {
+                cancelStream();
+                onHide();
+            }}
             className="QueryGenerationModal"
             bottomDOM={bottomDOM}
         >

querybook/webapp/hooks/useStream.ts

@@ -1,11 +1,12 @@
-import { useCallback, useState } from 'react';
+import { useCallback, useRef, useState } from 'react';
 
 import ds from 'lib/datasource';
 
 export enum StreamStatus {
     NOT_STARTED,
     STREAMING,
     FINISHED,
+    CANCELLED,
 }
 
 export function useStream(
@@ -16,15 +17,17 @@ export function useStream(
     streamData: { [key: string]: string };
     startStream: () => void;
     resetStream: () => void;
+    cancelStream: () => void;
 } {
     const [streamStatus, setSteamStatus] = useState(StreamStatus.NOT_STARTED);
     const [data, setData] = useState<{ [key: string]: string }>({});
+    const streamRef = useRef<EventSource | null>(null);
 
     const startStream = useCallback(() => {
         setSteamStatus(StreamStatus.STREAMING);
         setData({});
 
-        ds.stream(url, params, setData, (data) => {
+        streamRef.current = ds.stream(url, params, setData, (data) => {
             setData(data);
             setSteamStatus(StreamStatus.FINISHED);
         });
@@ -35,10 +38,18 @@ export function useStream(
         setData({});
     }, []);
 
+    const cancelStream = useCallback(() => {
+        if (streamStatus === StreamStatus.STREAMING) {
+            streamRef.current?.close();
+            setSteamStatus(StreamStatus.CANCELLED);
+        }
+    }, [streamStatus]);
+
     return {
         streamStatus,
         streamData: data,
         startStream,
         resetStream,
+        cancelStream,
     };
 }

querybook/webapp/lib/datasource.ts

@@ -201,6 +201,8 @@ function streamDatasource(
         parser.close();
         onStreamingEnd?.(parser.result);
     });
+
+    return eventSource;
 }
 
 export default {