Add workflow reconnection logic for page reload scenarios

frontend/apps/app/components/SessionDetailPage/SessionDetailPageClient.tsx

@@ -1,7 +1,6 @@
 'use client'
 
 import {
-  isHumanMessage,
   mapStoredMessagesToChatMessages,
   type StoredMessage,
 } from '@langchain/core/messages'
@@ -18,6 +17,8 @@ import { useStream } from './hooks/useStream'
 import { SQL_REVIEW_COMMENTS } from './mock'
 import styles from './SessionDetailPage.module.css'
 import type { Version } from './types'
+import { determineWorkflowAction } from './utils/determineWorkflowAction'
+import { getWorkflowInProgress } from './utils/workflowStorage'
 
 type Props = {
   buildingSchemaId: string
@@ -116,7 +117,7 @@ export const SessionDetailPageClient: FC<Props> = ({
     (artifact !== null || selectedVersion !== null) && activeTab
 
   const chatMessages = mapStoredMessagesToChatMessages(initialMessages)
-  const { isStreaming, messages, start, error } = useStream({
+  const { isStreaming, messages, start, replay, error } = useStream({
     initialMessages: chatMessages,
     designSessionId,
     senderName,
@@ -130,27 +131,37 @@ export const SessionDetailPageClient: FC<Props> = ({
   // Auto-trigger workflow on page load if there's an unanswered user message
   useEffect(() => {
     const triggerInitialWorkflow = async () => {
-      // Skip if already triggered
-      if (hasTriggeredInitialWorkflow.current) return
+      const isWorkflowInProgress = getWorkflowInProgress(designSessionId)
 
-      if (messages.length !== 1) return
+      const action = determineWorkflowAction(
+        messages,
+        isWorkflowInProgress,
+        hasTriggeredInitialWorkflow.current,
+      )
 
-      const firstItem = messages[0]
-      if (!firstItem || !isHumanMessage(firstItem)) return
+      if (action.type === 'none') return
 
       // Mark as triggered before the async call
       hasTriggeredInitialWorkflow.current = true
 
-      // Trigger the workflow for the initial user message
-      await start({
-        designSessionId,
-        userInput: firstItem.text,
-        isDeepModelingEnabled,
-      })
+      if (action.type === 'replay') {
+        // Trigger replay for interrupted workflow
+        await replay({
+          designSessionId,
+          isDeepModelingEnabled,
+        })
+      } else if (action.type === 'start') {
+        // Trigger the workflow for the initial user message
+        await start({
+          designSessionId,
+          userInput: action.userInput,
+          isDeepModelingEnabled,
+        })
+      }
     }
 
     triggerInitialWorkflow()
-  }, [messages, designSessionId, isDeepModelingEnabled, start])
+  }, [messages, designSessionId, isDeepModelingEnabled, start, replay])
 
   return (
     <div className={styles.container}>

frontend/apps/app/components/SessionDetailPage/hooks/useStream/useStream.ts

@@ -11,6 +11,10 @@ import { useCallback, useMemo, useRef, useState } from 'react'
 import { object, safeParse, string } from 'valibot'
 import { useNavigationGuard } from '../../../../hooks/useNavigationGuard'
 import { ERROR_MESSAGES } from '../../components/Chat/constants/chatConstants'
+import {
+  clearWorkflowInProgress,
+  setWorkflowInProgress,
+} from '../../utils/workflowStorage'
 import { parseSse } from './parseSse'
 import { useSessionStorageOnce } from './useSessionStorageOnce'
 
@@ -86,15 +90,19 @@ export const useStream = ({
   const abortRef = useRef<AbortController | null>(null)
   const retryCountRef = useRef(0)
 
-  const finalizeStream = useCallback(() => {
+  const completeWorkflow = useCallback((sessionId: string) => {
     setIsStreaming(false)
     abortRef.current = null
     retryCountRef.current = 0
+    clearWorkflowInProgress(sessionId)
   }, [])
 
-  const stop = useCallback(() => {
+  const abortWorkflow = useCallback(() => {
     abortRef.current?.abort()
+    setIsStreaming(false)
     abortRef.current = null
+    retryCountRef.current = 0
+    // Do NOT clear workflow flag - allow reconnection
   }, [])
 
   const clearError = useCallback(() => {
@@ -103,14 +111,7 @@ export const useStream = ({
 
   useNavigationGuard((_event) => {
     if (isStreaming) {
-      const shouldContinue = window.confirm(
-        "Design session is currently running. If you leave now, you'll lose your session progress. Continue anyway?",
-      )
-      if (shouldContinue) {
-        abortRef.current?.abort()
-        return true
-      }
-      return false
+      abortWorkflow()
     }
     return true
   })
@@ -171,7 +172,7 @@ export const useStream = ({
   const runStreamAttempt = useCallback(
     async (
       endpoint: string,
-      payload: unknown,
+      params: StartParams | ReplayParams,
     ): Promise<Result<StreamAttemptStatus, StreamError>> => {
       abortRef.current?.abort()
 
@@ -184,12 +185,12 @@ export const useStream = ({
         const res = await fetch(endpoint, {
           method: 'POST',
           headers: { 'Content-Type': 'application/json' },
-          body: JSON.stringify(payload),
+          body: JSON.stringify(params),
           signal: controller.signal,
         })
 
         if (!res.body) {
-          finalizeStream()
+          abortWorkflow()
           return err({
             type: 'network',
             message: ERROR_MESSAGES.FETCH_FAILED,
@@ -201,7 +202,7 @@ export const useStream = ({
 
         if (!endEventReceived) {
           if (controller.signal.aborted) {
-            finalizeStream()
+            abortWorkflow()
             return err({
               type: 'abort',
               message: 'Request was aborted',
@@ -213,10 +214,10 @@ export const useStream = ({
           return ok('shouldRetry')
         }
 
-        finalizeStream()
+        completeWorkflow(params.designSessionId)
         return ok('complete')
       } catch (unknownError) {
-        finalizeStream()
+        abortWorkflow()
 
         if (
           unknownError instanceof Error &&
@@ -234,7 +235,7 @@ export const useStream = ({
         })
       }
     },
-    [finalizeStream, processStreamEvents],
+    [completeWorkflow, abortWorkflow, processStreamEvents],
   )
 
   const replay = useCallback(
@@ -254,14 +255,14 @@ export const useStream = ({
       }
 
       const timeoutMessage = ERROR_MESSAGES.CONNECTION_TIMEOUT
-      finalizeStream()
+      abortWorkflow()
       setError(timeoutMessage)
       return err({
         type: 'timeout',
         message: timeoutMessage,
       })
     },
-    [finalizeStream, runStreamAttempt],
+    [abortWorkflow, runStreamAttempt],
   )
 
   const start = useCallback(
@@ -284,6 +285,9 @@ export const useStream = ({
         isFirstMessage.current = false
       }
 
+      // Set workflow in progress flag
+      setWorkflowInProgress(params.designSessionId)
+
       const result = await runStreamAttempt('/api/chat/stream', params)
 
       if (result.isErr()) {
@@ -309,8 +313,8 @@ export const useStream = ({
     messages,
     isStreaming,
     error,
-    stop,
     start,
+    replay,
     clearError,
   }
 }

frontend/apps/app/components/SessionDetailPage/utils/determineWorkflowAction.test.ts

@@ -0,0 +1,90 @@
+import { AIMessage, HumanMessage } from '@langchain/core/messages'
+import { describe, expect, it } from 'vitest'
+import { determineWorkflowAction } from './determineWorkflowAction'
+
+describe('determineWorkflowAction', () => {
+  describe('Priority 1: Already triggered', () => {
+    it('returns none when hasTriggered is true', () => {
+      const result = determineWorkflowAction(
+        [new HumanMessage('test')],
+        false, // isWorkflowInProgress
+        true, // hasTriggered
+      )
+      expect(result).toEqual({ type: 'none' })
+    })
+
+    it('returns none even when both hasTriggered and isWorkflowInProgress are true', () => {
+      const result = determineWorkflowAction(
+        [new HumanMessage('test')],
+        true, // isWorkflowInProgress
+        true, // hasTriggered
+      )
+      expect(result).toEqual({ type: 'none' })
+    })
+  })
+
+  describe('Priority 2: Workflow in progress flag exists', () => {
+    it('returns replay when isWorkflowInProgress is true', () => {
+      const result = determineWorkflowAction(
+        [],
+        true, // isWorkflowInProgress
+        false, // hasTriggered
+      )
+      expect(result).toEqual({ type: 'replay' })
+    })
+
+    it('returns replay even with multiple messages when isWorkflowInProgress is true', () => {
+      const result = determineWorkflowAction(
+        [new HumanMessage('first'), new AIMessage('response')],
+        true, // isWorkflowInProgress
+        false, // hasTriggered
+      )
+      expect(result).toEqual({ type: 'replay' })
+    })
+  })
+
+  describe('Priority 3: Single unanswered user message', () => {
+    it('returns start when there is a single HumanMessage', () => {
+      const result = determineWorkflowAction(
+        [new HumanMessage('hello')],
+        false, // isWorkflowInProgress
+        false, // hasTriggered
+      )
+      expect(result).toEqual({
+        type: 'start',
+        userInput: 'hello',
+      })
+    })
+
+    it('handles when content is not a string', () => {
+      const message = new HumanMessage({ content: 'test' })
+      const result = determineWorkflowAction([message], false, false)
+      expect(result.type).toBe('start')
+    })
+  })
+
+  describe('Priority 4: Do nothing', () => {
+    it('returns none when messages is empty', () => {
+      const result = determineWorkflowAction([], false, false)
+      expect(result).toEqual({ type: 'none' })
+    })
+
+    it('returns none when there are multiple messages', () => {
+      const result = determineWorkflowAction(
+        [new HumanMessage('first'), new AIMessage('response')],
+        false,
+        false,
+      )
+      expect(result).toEqual({ type: 'none' })
+    })
+
+    it('returns none when single message is AIMessage', () => {
+      const result = determineWorkflowAction(
+        [new AIMessage('ai message')],
+        false,
+        false,
+      )
+      expect(result).toEqual({ type: 'none' })
+    })
+  })
+})

frontend/apps/app/components/SessionDetailPage/utils/determineWorkflowAction.ts

@@ -0,0 +1,45 @@
+import type { BaseMessage } from '@langchain/core/messages'
+import { isHumanMessage } from '@langchain/core/messages'
+
+type WorkflowAction =
+  | { type: 'start'; userInput: string }
+  | { type: 'replay' }
+  | { type: 'none' }
+
+/**
+ * Pure function to determine workflow action
+ *
+ * @param messages - Message history
+ * @param isWorkflowInProgress - Whether workflow is in progress (from sessionStorage)
+ * @param hasTriggered - Whether workflow has already been triggered
+ * @returns Action to execute
+ */
+export const determineWorkflowAction = (
+  messages: BaseMessage[],
+  isWorkflowInProgress: boolean,
+  hasTriggered: boolean,
+): WorkflowAction => {
+  // 1. Already triggered - do nothing
+  if (hasTriggered) {
+    return { type: 'none' }
+  }
+
+  // 2. Workflow in progress flag exists - resume interrupted workflow
+  if (isWorkflowInProgress) {
+    return { type: 'replay' }
+  }
+
+  // 3. Single unanswered user message - start new workflow
+  if (messages.length === 1) {
+    const firstMessage = messages[0]
+    if (firstMessage && isHumanMessage(firstMessage)) {
+      return {
+        type: 'start',
+        userInput: firstMessage.text,
+      }
+    }
+  }
+
+  // 4. Otherwise - do nothing
+  return { type: 'none' }
+}

frontend/apps/app/components/SessionDetailPage/utils/workflowStorage.ts

@@ -0,0 +1,29 @@
+const WORKFLOW_KEY_PREFIX = 'liam:workflow:'
+
+/**
+ * Get whether workflow is in progress
+ */
+export const getWorkflowInProgress = (designSessionId: string): boolean => {
+  if (typeof window === 'undefined') return false
+  const key = `${WORKFLOW_KEY_PREFIX}${designSessionId}`
+  const value = sessionStorage.getItem(key)
+  return value === 'in_progress'
+}
+
+/**
+ * Set workflow in progress flag
+ */
+export const setWorkflowInProgress = (designSessionId: string): void => {
+  if (typeof window === 'undefined') return
+  const key = `${WORKFLOW_KEY_PREFIX}${designSessionId}`
+  sessionStorage.setItem(key, 'in_progress')
+}
+
+/**
+ * Clear workflow in progress flag
+ */
+export const clearWorkflowInProgress = (designSessionId: string): void => {
+  if (typeof window === 'undefined') return
+  const key = `${WORKFLOW_KEY_PREFIX}${designSessionId}`
+  sessionStorage.removeItem(key)
+}