AI subscription is removed on page change

src/app/layout.tsx

@@ -23,6 +23,7 @@ import {GraphQLFormattedError} from "graphql/error";
 import {addIslandErrorNotification} from "@code0-tech/pictor/dist/components/island/Island.hook";
 import {createConsumer} from "@rails/actioncable";
 import ActionCableLink from "graphql-ruby-client/subscriptions/ActionCableLink";
+import {AIGenerationWatcherComponent} from "@edition/ai/components/AIGenerationWatcherComponent";
 
 /**
  * Load the Inter font with Latin subset and swap display strategy
@@ -222,6 +223,7 @@ export default function RootLayout({children}: Readonly<{ children: React.ReactN
         <Suspense fallback={"loading..."}>
             <ApolloProvider client={client}>
                 <Toaster position={"top-right"}/>
+                <AIGenerationWatcherComponent/>
                 {children}
             </ApolloProvider>
         </Suspense>

src/packages/ce/src/ai/components/AIChatComponent.tsx

@@ -25,25 +25,15 @@ import {Select} from "@radix-ui/react-select";
 import {IconChevronDown, IconPlayerStop, IconSend, IconSparkles2Filled} from "@tabler/icons-react";
 import {AIService} from "@edition/ai/services/AI.service";
 import {motion} from "framer-motion";
-import {Flow, NamespaceProject, Subscription} from "@code0-tech/sagittarius-graphql-types";
-import {useSubscription} from "@apollo/client/react";
-import generateFlowSubscription from "@edition/ai/services/subscriptions/AI.generateFlow.subscription.graphql"
-
-const GENERATING_VARIANTS = [
-    "Generating...",
-    "Thinking...",
-    "Analyzing your prompt...",
-    "Composing flow...",
-    "Crafting nodes...",
-    "Wiring it up...",
-    "Almost there..."
-]
+import {AiGenerateFlowSubscriptionPayload, Flow, NamespaceProject} from "@code0-tech/sagittarius-graphql-types";
+import {useAIGenerationStore} from "@edition/ai/hooks/AI.generation.hook";
+import {AIGeneratingMessageComponent} from "@edition/ai/components/AIGeneratingMessageComponent";
 
 export interface AIChatComponentProps {
     projectId: NamespaceProject['id']
     flowId?: Flow['id']
     prompt?: string
-    onData?: (data: any) => string | void
+    onData?: (data: AiGenerateFlowSubscriptionPayload) => string | void
 
 }
 
@@ -53,70 +43,56 @@ export const AIChatComponent: React.FC<AIChatComponentProps> = (props) => {
 
     const aiService = useService(AIService)
     const aiStore = useStore(AIService)
+    const addGeneration = useAIGenerationStore(s => s.addGeneration)
+    const removeGeneration = useAIGenerationStore(s => s.removeGeneration)
 
     const [promptState, setPromptState] = React.useState<string>(prompt)
     const [model, setModel] = React.useState<string | undefined>(undefined)
     const [executionIdentifier, setExecutionIdentifier] = React.useState<string | null>(null)
-    const [isAIActive, setIsAIActive] = React.useState(false)
-    const [aiVariantIndex, setAiVariantIndex] = React.useState(0)
     const [aiErrorMessage, setAiErrorMessage] = React.useState<string | null>(null)
 
+    const isThisGenerating = useAIGenerationStore(s =>
+        !!executionIdentifier && s.hasGeneration(executionIdentifier)
+    )
+
     const models = React.useMemo(
         () => aiService.values(),
         [aiStore]
     )
 
-    const {data} = useSubscription<Subscription>(generateFlowSubscription, {
-        variables: {executionIdentifier: executionIdentifier},
-        skip: !executionIdentifier,
-        onData: (data) => {
-            setIsAIActive(true)
-            if (data.data.data?.aiGenerateFlow?.flow) {
-                const result = onData?.(data.data.data?.aiGenerateFlow)
-                if (typeof result === "string") {
-                    setAiErrorMessage(result)
-                } else {
-                    setPromptState("")
-                }
-                setExecutionIdentifier(null)
-            } else if (data.data.data?.aiGenerateFlow?.flow === null) {
-                setExecutionIdentifier(null)
-                setAiErrorMessage("Generation failed. Try another model.")
-            }
-        },
-        onComplete: () => setIsAIActive(false),
-        onError: () => {
-            setIsAIActive(false)
-            setAiErrorMessage("Generation failed. Try another model.")
+    React.useEffect(() => {
+        if (executionIdentifier && !isThisGenerating) {
+            if (!aiErrorMessage) setPromptState("")
             setExecutionIdentifier(null)
-        },
-    })
+        }
+    }, [executionIdentifier, isThisGenerating, aiErrorMessage])
 
-    const aiLoading = React.useMemo(
-        () => ((!data || Object.keys(data).length === 0) && executionIdentifier && isAIActive),
-        [executionIdentifier, data, isAIActive]
-    )
+    const aiLoading = !!executionIdentifier && isThisGenerating
 
     const onSend = React.useCallback(() => {
+        setAiErrorMessage(null)
         aiService.generateFlow({
             prompt: promptState,
             projectId: projectId!,
             modelIdentifier: model!,
             flowId: flowId,
         }).then(payload => {
-            if ((payload?.errors?.length ?? 0) <= 0) {
-                setExecutionIdentifier(payload?.executionIdentifier ?? null)
+            if ((payload?.errors?.length ?? 0) <= 0 && payload?.executionIdentifier) {
+                const id = payload.executionIdentifier
+                setExecutionIdentifier(id)
+                addGeneration({
+                    executionIdentifier: id,
+                    onData: (data) => onData?.(data),
+                    onError: (message) => setAiErrorMessage(message),
+                })
             }
         })
-    }, [aiService, model, promptState])
+    }, [aiService, model, promptState, projectId, flowId, onData, addGeneration])
 
-    React.useEffect(() => {
-        const id = setInterval(
-            () => setAiVariantIndex(i => (i + 1) % GENERATING_VARIANTS.length),
-            4000
-        )
-        return () => clearInterval(id)
-    }, [])
+    const onStop = React.useCallback(() => {
+        if (executionIdentifier) removeGeneration(executionIdentifier)
+        setExecutionIdentifier(null)
+    }, [executionIdentifier, removeGeneration])
 
     React.useEffect(() => {
         setModel(models.length > 0
@@ -186,30 +162,7 @@ export const AIChatComponent: React.FC<AIChatComponentProps> = (props) => {
                                 height: "2.3rem",
                             }}
                         >
-                            <motion.div
-                                key={GENERATING_VARIANTS[aiVariantIndex]}
-                                initial={{y: 18, opacity: 0}}
-                                animate={{y: 0, opacity: 1}}
-                                exit={{y: -18, opacity: 0}}
-                                transition={{duration: 0.45, ease: [0.4, 0, 0.4, 1]}}
-                            >
-                                <motion.div
-                                    style={{
-                                        display: "block",
-                                        backgroundImage:
-                                            "linear-gradient(90deg, rgba(255,255,255,0.25) 0%, rgba(255,255,255,0.25) 40%, rgba(255,255,255,1) 50%, rgba(255,255,255,0.25) 60%, rgba(255,255,255,0.25) 100%)",
-                                        backgroundSize: "200% 100%",
-                                        WebkitBackgroundClip: "text",
-                                        backgroundClip: "text",
-                                        WebkitTextFillColor: "transparent",
-                                        color: "transparent",
-                                    }}
-                                    animate={{backgroundPosition: ["200% 0%", "-200% 0%"]}}
-                                    transition={{duration: 3, ease: "linear", repeat: Infinity}}
-                                >
-                                    <Text>{GENERATING_VARIANTS[aiVariantIndex]}</Text>
-                                </motion.div>
-                            </motion.div>
+                            <AIGeneratingMessageComponent/>
                         </div>
                     ) : (
                         <EditorInput
@@ -296,7 +249,7 @@ export const AIChatComponent: React.FC<AIChatComponentProps> = (props) => {
                                 </Text>
                             ) : null}
                             {aiLoading ? (
-                                <Button onClick={() => setExecutionIdentifier(null)} color={"secondary"}>
+                                <Button onClick={onStop} color={"secondary"}>
                                     <IconPlayerStop size={13}/>
                                 </Button>
                             ) : (

src/packages/ce/src/ai/components/AIGeneratingMessageComponent.tsx

@@ -0,0 +1,49 @@
+"use client"
+
+import React from "react"
+import {Text} from "@code0-tech/pictor"
+import {motion} from "framer-motion"
+
+const VARIANTS = [
+    "Generating...",
+    "Thinking...",
+    "Analyzing your prompt...",
+    "Composing flow...",
+    "Crafting nodes...",
+    "Wiring it up...",
+    "Almost there...",
+]
+
+export const AIGeneratingMessageComponent: React.FC = () => {
+
+    const [index, setIndex] = React.useState(0)
+
+    React.useEffect(() => {
+        const id = setInterval(() => setIndex(i => (i + 1) % VARIANTS.length), 4000)
+        return () => clearInterval(id)
+    }, [])
+
+    return <motion.div
+        key={VARIANTS[index]}
+        initial={{y: 12, opacity: 0}}
+        animate={{y: 0, opacity: 1}}
+        exit={{y: -12, opacity: 0}}
+        transition={{duration: 0.4, ease: [0.4, 0, 0.4, 1]}}
+    >
+        <motion.div
+            style={{
+                backgroundImage:
+                    "linear-gradient(90deg, rgba(226,112,255,0.45) 0%, rgba(226,112,255,0.45) 40%, rgba(226,112,255,1) 50%, rgba(226,112,255,0.45) 60%, rgba(226,112,255,0.45) 100%)",
+                backgroundSize: "200% 100%",
+                WebkitBackgroundClip: "text",
+                backgroundClip: "text",
+                WebkitTextFillColor: "transparent",
+                color: "transparent",
+            }}
+            animate={{backgroundPosition: ["200% 0%", "-200% 0%"]}}
+            transition={{duration: 3, ease: "linear", repeat: Infinity}}
+        >
+            <Text>{VARIANTS[index]}</Text>
+        </motion.div>
+    </motion.div>
+}

src/packages/ce/src/ai/components/AIGenerationSubscriberComponent.tsx

@@ -0,0 +1,40 @@
+"use client"
+
+import React from "react"
+import {useSubscription} from "@apollo/client/react"
+import {AiGenerateFlowSubscriptionPayload, Subscription} from "@code0-tech/sagittarius-graphql-types"
+import {AIGeneration, useAIGenerationStore} from "@edition/ai/hooks/AI.generation.hook"
+import generateFlowSubscription from "@edition/ai/services/subscriptions/AI.generateFlow.subscription.graphql"
+
+export interface AIGenerationSubscriberComponentProps {
+    generation: AIGeneration
+}
+
+export const AIGenerationSubscriberComponent: React.FC<AIGenerationSubscriberComponentProps> = ({generation}) => {
+
+    const removeGeneration = useAIGenerationStore(s => s.removeGeneration)
+
+    useSubscription<Subscription>(generateFlowSubscription, {
+        variables: {executionIdentifier: generation.executionIdentifier},
+        onData: (data) => {
+            const payload = data.data.data?.aiGenerateFlow as AiGenerateFlowSubscriptionPayload | undefined
+            if (payload?.flow) {
+                const result = generation.onData(payload)
+                if (typeof result === "string") {
+                    generation.onError?.(result)
+                }
+                removeGeneration(generation.executionIdentifier)
+            } else if (payload?.flow === null) {
+                generation.onError?.("Generation failed. Try another model.")
+                removeGeneration(generation.executionIdentifier)
+            }
+        },
+        onComplete: () => removeGeneration(generation.executionIdentifier),
+        onError: () => {
+            generation.onError?.("Generation failed. Try another model.")
+            removeGeneration(generation.executionIdentifier)
+        },
+    })
+
+    return null
+}

src/packages/ce/src/ai/components/AIGenerationWatcherComponent.tsx

@@ -0,0 +1,51 @@
+"use client"
+
+import React from "react"
+import {useAIGenerationStore} from "@edition/ai/hooks/AI.generation.hook"
+import {useIsland} from "@code0-tech/pictor/dist/components/island/Island.hook"
+import {IconSparkles2Filled} from "@tabler/icons-react"
+import {AIGeneratingMessageComponent} from "@edition/ai/components/AIGeneratingMessageComponent"
+import {AIGenerationSubscriberComponent} from "@edition/ai/components/AIGenerationSubscriberComponent"
+
+export const AIGenerationWatcherComponent: React.FC = () => {
+
+    const generations = useAIGenerationStore(s => s.generations)
+    const islandToastIdRef = React.useRef<number | null>(null)
+    const addToast = useIsland(s => s.addToast)
+    const removeToast = useIsland(s => s.removeToast)
+
+    const isGenerating = generations.length > 0
+
+    React.useEffect(() => {
+        if (isGenerating && islandToastIdRef.current === null) {
+            const id = Date.now()
+            islandToastIdRef.current = id
+            addToast({
+                id,
+                duration: Infinity,
+                index: 2,
+                icon: <IconSparkles2Filled size={16} color={"#e270ff"}/>,
+                message: <AIGeneratingMessageComponent/>,
+            })
+        } else if (!isGenerating && islandToastIdRef.current !== null) {
+            removeToast(islandToastIdRef.current)
+            islandToastIdRef.current = null
+        }
+    }, [isGenerating, addToast, removeToast])
+
+    React.useEffect(() => () => {
+        if (islandToastIdRef.current !== null) {
+            removeToast(islandToastIdRef.current)
+            islandToastIdRef.current = null
+        }
+    }, [removeToast])
+
+    return <>
+        {generations.map(generation => (
+            <AIGenerationSubscriberComponent
+                key={generation.executionIdentifier}
+                generation={generation}
+            />
+        ))}
+    </>
+}

src/packages/ce/src/ai/hooks/AI.generation.hook.ts

@@ -0,0 +1,29 @@
+import {create} from "zustand"
+import {AiGenerateFlowSubscriptionPayload} from "@code0-tech/sagittarius-graphql-types"
+
+export interface AIGeneration {
+    executionIdentifier: string
+    onData: (payload: AiGenerateFlowSubscriptionPayload) => string | void
+    onError?: (message: string) => void
+}
+
+interface AIGenerationState {
+    generations: AIGeneration[]
+    addGeneration: (generation: AIGeneration) => void
+    removeGeneration: (executionIdentifier: string) => void
+    hasGeneration: (executionIdentifier: string) => boolean
+}
+
+export const useAIGenerationStore = create<AIGenerationState>((setState, getState) => ({
+    generations: [],
+    addGeneration: (generation) => setState((state) => ({
+        generations: state.generations.some(g => g.executionIdentifier === generation.executionIdentifier)
+            ? state.generations
+            : [...state.generations, generation]
+    })),
+    removeGeneration: (executionIdentifier) => setState((state) => ({
+        generations: state.generations.filter(g => g.executionIdentifier !== executionIdentifier)
+    })),
+    hasGeneration: (executionIdentifier) =>
+        getState().generations.some(g => g.executionIdentifier === executionIdentifier),
+}))