fix: initial chat redirect bugs

gregnr · gregnr · commit 69c1b2483d7d · 2024-08-11T15:59:46.000-05:00
diff --git a/apps/postgres-new/app/page.tsx b/apps/postgres-new/app/page.tsx
@@ -68,18 +68,51 @@ export default function Page() {
     <Workspace
       databaseId={nextDatabaseId}
       visibility="local"
-      onStart={async () => {
+      onMessage={async () => {
         // Make the DB no longer hidden
         await updateDatabase({ id: nextDatabaseId, name: null, isHidden: false })
 
-        // Navigate to this DB's path
-        router.push(`/db/${nextDatabaseId}`)
-
         // Pre-load the next DB
         const nextId = uniqueId()
         localStorage.setItem('next-db-id', JSON.stringify(nextId))
         preloadDb(nextId)
       }}
+      onReply={async (message, append) => {
+        if (!dbManager) {
+          throw new Error('dbManager is not available')
+        }
+
+        const messages = await dbManager.getMessages(nextDatabaseId)
+        const isFirstReplyComplete =
+          !messages.some((message) => message.role === 'assistant' && !message.toolInvocations) &&
+          message.role === 'assistant' &&
+          !message.toolInvocations
+
+        // The model might run multiple tool calls before ending with a message, so
+        // we only want to redirect after all of these back-to-back calls finish
+        if (isFirstReplyComplete) {
+          router.push(`/db/${nextDatabaseId}`)
+
+          append({
+            role: 'user',
+            content: 'Name this conversation. No need to reply.',
+            data: {
+              automated: true,
+            },
+          })
+        }
+      }}
+      onCancelReply={(append) => {
+        router.push(`/db/${nextDatabaseId}`)
+
+        append({
+          role: 'user',
+          content: 'Name this conversation. No need to reply.',
+          data: {
+            automated: true,
+          },
+        })
+      }}
     />
   )
 }
diff --git a/apps/postgres-new/components/workspace.tsx b/apps/postgres-new/components/workspace.tsx
@@ -1,14 +1,13 @@
 'use client'
 
-import { CreateMessage, Message, useChat } from 'ai/react'
+import { CreateMessage, Message, useChat, UseChatHelpers } from 'ai/react'
 import { createContext, useCallback, useContext, useMemo } from 'react'
 import { useMessageCreateMutation } from '~/data/messages/message-create-mutation'
 import { useMessagesQuery } from '~/data/messages/messages-query'
 import { useTablesQuery } from '~/data/tables/tables-query'
 import { useOnToolCall } from '~/lib/hooks'
 import { useBreakpoint } from '~/lib/use-breakpoint'
 import { ensureMessageId, ensureToolResult } from '~/lib/util'
-import { useApp } from './app-provider'
 import Chat, { getInitialMessages } from './chat'
 import IDE from './ide'
 
@@ -27,13 +26,31 @@ export type WorkspaceProps = {
   visibility: Visibility
 
   /**
-   * Callback called when the conversation has started.
+   * Callback called after the user sends a message.
    */
-  onStart?: () => void | Promise<void>
+  onMessage?: (
+    message: Message | CreateMessage,
+    append: UseChatHelpers['append']
+  ) => void | Promise<void>
+
+  /**
+   * Callback called after the LLM finishes a reply.
+   */
+  onReply?: (message: Message, append: UseChatHelpers['append']) => void | Promise<void>
+
+  /**
+   * Callback called when the user cancels the reply.
+   */
+  onCancelReply?: (append: UseChatHelpers['append']) => void | Promise<void>
 }
 
-export default function Workspace({ databaseId, visibility, onStart }: WorkspaceProps) {
-  const { dbManager } = useApp()
+export default function Workspace({
+  databaseId,
+  visibility,
+  onMessage,
+  onReply,
+  onCancelReply,
+}: WorkspaceProps) {
   const isSmallBreakpoint = useBreakpoint('lg')
   const onToolCall = useOnToolCall(databaseId)
   const { mutateAsync: saveMessage } = useMessageCreateMutation(databaseId)
@@ -46,40 +63,18 @@ export default function Workspace({ databaseId, visibility, onStart }: Workspace
 
   const initialMessages = useMemo(() => (tables ? getInitialMessages(tables) : undefined), [tables])
 
-  const {
-    messages,
-    setMessages,
-    append,
-    stop: stopReply,
-  } = useChat({
+  const { messages, setMessages, append, stop } = useChat({
     id: databaseId,
     api: '/api/chat',
     maxToolRoundtrips: 10,
+    keepLastMessageOnError: true,
     onToolCall: onToolCall as any, // our `OnToolCall` type is more specific than `ai` SDK's
     initialMessages:
       existingMessages && existingMessages.length > 0 ? existingMessages : initialMessages,
     async onFinish(message) {
-      if (!dbManager) {
-        throw new Error('dbManager is not available')
-      }
-
+      // Order is important here
+      await onReply?.(message, append)
       await saveMessage({ message })
-
-      const database = await dbManager.getDatabase(databaseId)
-      const isStartOfConversation = database.isHidden && !message.toolInvocations
-
-      if (isStartOfConversation) {
-        await onStart?.()
-
-        // Intentionally using `append` vs `appendMessage` so that this message isn't persisted in the meta DB
-        await append({
-          role: 'user',
-          content: 'Name this conversation. No need to reply.',
-          data: {
-            automated: true,
-          },
-        })
-      }
     },
   })
 
@@ -90,12 +85,20 @@ export default function Workspace({ databaseId, visibility, onStart }: Workspace
         return isModified ? [...messages] : messages
       })
       ensureMessageId(message)
+
+      // Intentionally don't await so that framer animations aren't affected
       append(message)
       saveMessage({ message })
+      onMessage?.(message, append)
     },
-    [setMessages, saveMessage, append]
+    [onMessage, setMessages, saveMessage, append]
   )
 
+  const stopReply = useCallback(async () => {
+    stop()
+    onCancelReply?.(append)
+  }, [onCancelReply, stop, append])
+
   const isConversationStarted =
     initialMessages !== undefined && messages.length > initialMessages.length
 
@@ -134,7 +137,7 @@ export type WorkspaceContextValues = {
   messages: Message[]
   visibility: Visibility
   appendMessage(message: Message | CreateMessage): Promise<void>
-  stopReply(): void
+  stopReply(): Promise<void>
 }
 
 export const WorkspaceContext = createContext<WorkspaceContextValues | undefined>(undefined)
