refactor: preserve messages on refresh

Author: kevinkim-ogp

packages/frontend/src/hooks/useChatStream.ts

@@ -1,18 +1,26 @@
 import { useCallback, useMemo } from 'react'
+import { useLocation, useNavigate } from 'react-router-dom'
 import type { UIMessage } from '@ai-sdk/react'
 import { useChat } from '@ai-sdk/react'
 import { useToast } from '@opengovsg/design-system-react'
 import { DefaultChatTransport } from 'ai'
 
+import * as URLS from '@/config/urls'
+import {
+  deduplicateMessages,
+  extractTextContent,
+  transformMessages,
+} from '@/pages/AiBuilder/helpers'
+
 export interface Message {
+  id: string // this is auto-generated by the AI SDK
   text: string
-  traceId?: string
-  generationId?: string
+  traceId?: string // only assistant messages have this
   isUser: boolean
 }
 
 // Custom message type with metadata
-type CustomUIMessage = UIMessage<{
+export type CustomUIMessage = UIMessage<{
   traceId?: string
 }>
 
@@ -22,6 +30,8 @@ export interface UseChatStreamOptions {
 
 export function useChatStream(options?: UseChatStreamOptions) {
   const toast = useToast()
+  const navigate = useNavigate()
+  const location = useLocation()
 
   const {
     messages: aiMessages,
@@ -34,9 +44,19 @@ export function useChatStream(options?: UseChatStreamOptions) {
       api: '/api/chat',
       credentials: 'include',
       prepareSendMessagesRequest: ({ messages }) => {
-        // Send all messages to maintain conversation context
+        // Convert initialMessages to the format expected by the API
+        const initialMsgs = (options?.initialMessages || []).map((msg) => ({
+          id: msg.id,
+          role: msg.isUser ? 'user' : 'assistant',
+          parts: [{ type: 'text', text: msg.text }],
+          ...(msg.traceId && { metadata: { traceId: msg.traceId } }),
+        }))
+
+        // Prepend initial messages to maintain full conversation context
+        const allMessages = [...initialMsgs, ...messages]
+
         const body = {
-          messages: messages,
+          messages: allMessages,
           sessionId: '',
         }
         return { body }
@@ -51,16 +71,29 @@ export function useChatStream(options?: UseChatStreamOptions) {
         position: 'top',
       })
     },
+    onFinish: ({ messages }) => {
+      // transform the messages and save to location state
+      // so that user can still access it if they refresh the page
+      const transformedMessages = transformMessages(messages)
+
+      // Combine initial messages with new messages to preserve full history
+      const allMessages = deduplicateMessages([
+        ...(options?.initialMessages || []),
+        ...transformedMessages,
+      ])
+
+      navigate(`${URLS.EDITOR}/ai`, {
+        state: {
+          ...location.state,
+          isFormMode: false,
+          chatInput: allMessages[allMessages.length - 1].text,
+          chatMessages: allMessages,
+        },
+        replace: true,
+      })
+    },
   })
 
-  // Helper function to extract text content from UIMessage
-  const extractTextContent = useCallback((msg: CustomUIMessage): string => {
-    return msg.parts
-      .filter((part) => part.type === 'text')
-      .map((part) => (part as any).text)
-      .join('')
-  }, [])
-
   // Transform AI SDK messages to our Message format
   const messages = useMemo<Message[]>(() => {
     const isActivelyStreaming = status === 'streaming' || status === 'submitted'
@@ -81,20 +114,13 @@ export function useChatStream(options?: UseChatStreamOptions) {
       }
     }
 
-    const transformedMessages = messagesToTransform.map((msg) => {
-      // Extract traceId from message metadata
-      const traceId = msg.metadata?.traceId
-
-      return {
-        text: extractTextContent(msg),
-        isUser: msg.role === 'user',
-        traceId: traceId,
-      }
-    })
-
-    // Combine initial messages with new messages
-    return [...initialMsgs, ...transformedMessages]
-  }, [aiMessages, extractTextContent, status, options?.initialMessages])
+    const transformedMessages = transformMessages(messagesToTransform)
+    const allMessages = deduplicateMessages([
+      ...initialMsgs,
+      ...transformedMessages,
+    ])
+    return allMessages
+  }, [aiMessages, options?.initialMessages, status])
 
   // Get the current streaming response (last assistant message that's still being streamed)
   const currentResponse = useMemo(() => {
@@ -109,7 +135,7 @@ export function useChatStream(options?: UseChatStreamOptions) {
       return extractTextContent(lastMessage)
     }
     return ''
-  }, [aiMessages, status, extractTextContent])
+  }, [aiMessages, status])
 
   // Wrapper for sendMessage that matches the expected signature
   const sendMessageWrapper = useCallback(
@@ -122,16 +148,12 @@ export function useChatStream(options?: UseChatStreamOptions) {
     [sendMessage],
   )
 
-  const cancelStream = useCallback(() => {
-    stop()
-  }, [stop])
-
   return {
     messages,
     currentResponse,
     isStreaming: status === 'submitted' || status === 'streaming',
     error: aiError?.message || null,
     sendMessage: sendMessageWrapper,
-    cancelStream,
+    cancelStream: stop,
   }
 }

packages/frontend/src/pages/AiBuilder/helpers.ts

@@ -1,6 +1,39 @@
+import { CustomUIMessage, Message } from '@/hooks/useChatStream'
+
 export const getPromptFromFormInput = (formInput: {
   trigger: string
   actions: string
 }) => {
   return `#### Start the workflow\n${formInput.trigger}\n\n#### Actions\n${formInput.actions}`
 }
+
+// deduplicate messages by id
+// there may be duplicates when the messages are combined
+export const deduplicateMessages = (messages: Message[]) => {
+  const seen = new Set()
+  return messages.filter((msg) => {
+    if (!msg.id || !seen.has(msg.id)) {
+      if (msg.id) {
+        seen.add(msg.id)
+        return true
+      }
+    }
+  })
+}
+
+// Helper function to extract text content from UIMessage
+export const extractTextContent = (msg: CustomUIMessage): string => {
+  return msg.parts
+    .filter((part) => part.type === 'text')
+    .map((part) => (part as any).text)
+    .join('')
+}
+
+export const transformMessages = (messages: CustomUIMessage[]) => {
+  return messages.map((msg) => ({
+    id: msg.id,
+    text: extractTextContent(msg),
+    isUser: msg.role === 'user',
+    traceId: msg.metadata?.traceId,
+  }))
+}