feat: refactor and fixes client streaming flow, user can send messages at the same time seamlessly

Author: NamesMT

apps/backend-convex/convex/_generated/api.d.ts

@@ -10,7 +10,7 @@
 
 import type * as authInfo from "../authInfo.js";
 import type * as crons from "../crons.js";
-import type * as http_ai from "../http/ai.js";
+import type * as http_chat from "../http/chat.js";
 import type * as http from "../http.js";
 import type * as messages from "../messages.js";
 import type * as tasks from "../tasks.js";
@@ -33,7 +33,7 @@ import type {
 declare const fullApi: ApiFromModules<{
   authInfo: typeof authInfo;
   crons: typeof crons;
-  "http/ai": typeof http_ai;
+  "http/chat": typeof http_chat;
   http: typeof http;
   messages: typeof messages;
   tasks: typeof tasks;

apps/backend-convex/convex/http.ts

@@ -2,9 +2,9 @@ import type { HonoWithConvex } from 'convex-helpers/server/hono'
 import type { ActionCtx } from './_generated/server'
 import { HttpRouterWithHono } from 'convex-helpers/server/hono'
 import { Hono } from 'hono'
-import { aiApp } from './http/ai'
+import { chatApp } from './http/chat'
 
 const app: HonoWithConvex<ActionCtx> = new Hono()
-app.route('/api/ai', aiApp)
+app.route('/api/chat', chatApp)
 
 export default new HttpRouterWithHono(app)

apps/backend-convex/convex/http/ai.ts apps/backend-convex/convex/http/chat.tsapps/backend-convex/convex/http/ai.ts renamed to apps/backend-convex/convex/http/chat.ts

@@ -19,20 +19,23 @@ const rateLimiter = new RateLimiter(components.rateLimiter, {
   aiChat: { kind: 'token bucket', rate: 10, period: MINUTE, capacity: 3 },
 })
 
-export const aiApp: HonoWithConvex<ActionCtx> = new Hono()
-aiApp
+export const chatApp: HonoWithConvex<ActionCtx> = new Hono()
+chatApp
   .use(cors())
   .post(
-    '/chat/stream',
+    '/stream',
     zValidator('json', z.object({
       threadId: z.string(),
       provider: z.string(),
       model: z.string(),
       apiKey: z.optional(z.string()),
       content: z.optional(z.string()),
+      // Optionally set the stream id when creating a new stream message for identification
+      streamId: z.optional(z.string()),
       context: z.optional(z.object({
         from: z.optional(z.string()),
       })),
+      // Will bypass `content` and `streamId`, resumes the target streamId.
       resumeStreamId: z.optional(z.string()),
       finishOnly: z.optional(z.boolean()),
       lockerKey: z.optional(z.string()),
@@ -52,6 +55,7 @@ aiApp
         finishOnly,
         lockerKey,
       } = c.req.valid('json')
+      let { streamId } = c.req.valid('json')
 
       // getUserIdentity on HTTP Action will throw if not authenticated 🤦‍♂️
       const userIdentity = await c.env.auth.getUserIdentity().catch(() => null)
@@ -65,21 +69,19 @@ aiApp
       const threadId = _threadId as Id<'threads'>
       const thread = await c.env.runQuery(api.threads.get, { threadId, lockerKey })
 
-      let streamId: string
       let streamingMessageId: Id<'messages'>
       let existingMessage: Doc<'messages'> | null = null
 
-      // Disable SSE resume for now until we have pub sub
+      // Disable SSE resume, if you want SSE resume, implement a pub-sub.
       if (resumeStreamId)
         throw new ConvexError('SSE stream resume is disabled')
 
+      // On SSE resume
       if (resumeStreamId) {
         streamId = resumeStreamId
 
         // Check if there's an existing streaming message to resume
-        existingMessage = await c.env.runQuery(internal.messages.getStreamingMessage, {
-          streamId,
-        })
+        existingMessage = await c.env.runQuery(internal.messages.getStreamingMessage, { streamId })
 
         if (!existingMessage) {
           // If no streaming message found, just return success
@@ -96,12 +98,17 @@ aiApp
           c.text('OK')
         }
       }
+      // On new stream
       else if (content) {
         if (thread.frozen)
           throw new ConvexError(`Can't send new messages to frozen thread`)
 
-        // Create new streaming session
-        streamId = `${Date.now()}_${randomStr(10)}`
+        // If user provides a streamId, check if its properly unused
+        if (streamId) {
+          if (await c.env.runQuery(internal.messages.getStreamingMessage, { streamId }))
+            throw new ConvexError('streamId is already in use')
+        }
+        streamId = streamId ?? `stream-${Date.now()}_${randomStr(4)}`
 
         // Add user message to thread
         await c.env.runMutation(internal.messages.internalAdd, {
@@ -114,7 +121,7 @@ aiApp
           lockerKey,
         })
 
-        // Add assistant message to thread (initial streaming)
+        // Add assistant message to thread
         streamingMessageId = await c.env.runMutation(internal.messages.internalAdd, {
           threadId,
           role: 'assistant',

apps/frontend/app/components/chat/ChatInterface.vue

@@ -2,7 +2,7 @@
 <script setup lang="ts">
 import type { Doc, Id } from 'backend-convex/convex/_generated/dataModel'
 import type Lenis from 'lenis'
-import { keyBy, randomStr, sleep, uniquePromise } from '@namesmt/utils'
+import { keyBy, objectPick, randomStr, sleep, uniquePromise } from '@namesmt/utils'
 import { api } from 'backend-convex/convex/_generated/api'
 import { useConvexClient } from 'convex-vue'
 import { countdown, debounce, getInstance, throttle } from 'kontroll'
@@ -43,7 +43,7 @@ const cachedThreadsMessages: {
 } = {}
 const messages = ref<Array<CustomMessage>>([])
 const messagesKeyed = computed(() => keyBy(messages.value, 'id'))
-const streamingMessages = ref(0)
+const streamingMessagesMap = reactive<Record<string, true>>({ })
 const isFetching = ref(false)
 const chatInput = ref('')
 
@@ -56,14 +56,25 @@ const { ignoreUpdates: ignorePathUpdate } = watchIgnorable(
 
     messages.value = cachedThreadsMessages[threadId as string] ?? []
 
-    doScrollBottom({ smooth: false })
+    nextTick(() => { doScrollBottom({ smooth: false, maybe: true }) })
 
     if (threadId) {
       isFetching.value = true
       await convex.query(api.messages.listByThread, { threadId: threadId as Doc<'threads'>['_id'], lockerKey: getLockerKey(threadId) })
-        .then((existingMessages) => {
+        .then((messagesFromConvex) => {
           if (threadIdRef.value === threadId) {
-            messages.value = existingMessages.map(customMessageTransform)
+            if (threadId === oldThreadId) {
+              for (const index in messagesFromConvex) {
+                const m = messagesFromConvex[index]!
+                if (m.streamId && !streamingMessagesMap[m.streamId]) {
+                  messages.value.push(customMessageTransform(messagesFromConvex[+index - 1]!))
+                  messages.value.push(customMessageTransform(m))
+                }
+              }
+            }
+            else {
+              messages.value = messagesFromConvex.map(customMessageTransform)
+            }
           }
         })
         .catch((e) => {
@@ -93,11 +104,12 @@ const { ignoreUpdates: ignorePathUpdate } = watchIgnorable(
           && message.streamId
         ) {
           console.log('Attempting to resume stream for session:', message.streamId)
-          uniquePromise(message.streamId, () => resumeStreamProcess(message.streamId!, message.id))
+          nextTick(() => { uniquePromise(message.streamId!, () => resumeStreamProcess(message.streamId!, message.id)) })
         }
       }
 
-      nextTick(() => doScrollBottom({ tries: 6 }))
+      if (threadId !== oldThreadId)
+        nextTick(() => doScrollBottom({ tries: 6 }))
     }
   },
   { immediate: true },
@@ -140,6 +152,8 @@ async function handleSubmit({ input }: HandleSubmitArgs) {
       .then(() => { sleep(500).then(() => handleSubmit({ input })) })
   }
 
+  const streamId = `stream-${Date.now()}_${randomStr(4)}`
+
   // Optimistically add the messages
   messages.value.push({
     id: `user-${Date.now()}_${randomStr(4)}`,
@@ -153,7 +167,7 @@ async function handleSubmit({ input }: HandleSubmitArgs) {
     model: chatContext.activeAgent.value.model,
     content: '',
     isStreaming: true,
-    streamId: undefined,
+    streamId,
   } as any as CustomMessage)
 
   // For some reason creating object reference first does not work, so we push and then get last message
@@ -187,8 +201,8 @@ async function handleSubmit({ input }: HandleSubmitArgs) {
   // Wraps in a kontroller to make sure there is only one stream on the same message
   throttle(
     1,
-    () => streamToMessage({ message: targetMessage, content: userInput }),
-    { key: `messageStream-${targetMessage.id}` },
+    () => streamToMessage({ message: targetMessage, content: userInput, streamId }),
+    { key: `messageStream-${streamId}` },
   )
 }
 
@@ -200,10 +214,10 @@ async function resumeStreamProcess(streamSessionId: string, messageId: string) {
   if (getInstance(threadIdRef.value))
     return console.warn('Trying to resume stream for message that is currently streaming:', messageId)
 
-  // Currently we doesn't support SSE resume yet
+  // Currently SSE resume not implemented yet
   // await streamToMessage({ message, resumeStreamId: streamSessionId })
 
-  // Using custom convex polling instead
+  // Using custom convex polling resume instead
   await pollToMessage({ message, resumeStreamId: streamSessionId })
 }
 
@@ -218,22 +232,26 @@ async function pollToMessage({ message, resumeStreamId, threadId = threadIdRef.v
     return
   }
 
+  streamingMessagesMap[resumeStreamId] = true
+  console.log(`Polling: ${message.id}`)
+
   const messageFromConvex = await convex.query(api.messages.get, {
     messageId: message._id,
     lockerKey: getLockerKey(threadId),
   })
-  Object.assign(message, customMessageTransform(messageFromConvex))
+  Object.assign(message, objectPick(messageFromConvex, ['content', 'context', 'isStreaming']))
 
   if (message.isStreaming) {
     // Wraps in a kontroller to make sure there is only one stream on the same message
     countdown(
       500,
-      () => pollToMessage({ message, resumeStreamId, threadId }),
-      { key: `messageStream-${message.id}` },
+      () => { nextTick(() => { pollToMessage({ message, resumeStreamId, threadId }) }) },
+      { key: `messageStream-${resumeStreamId}` },
     )
   }
   else {
-    console.log('Poll completed')
+    console.log(`Poll completed: ${message.id}`)
+    delete streamingMessagesMap[resumeStreamId]
   }
 
   nextTick(() => { doScrollBottom({ maybe: true }) })
@@ -242,17 +260,19 @@ async function pollToMessage({ message, resumeStreamId, threadId = threadIdRef.v
 interface StreamToMessageArgs {
   message: CustomMessage
   content?: string
+  streamId?: string
   resumeStreamId?: string
 }
-async function streamToMessage({ message, content, resumeStreamId }: StreamToMessageArgs) {
+async function streamToMessage({ message, content, streamId, resumeStreamId }: StreamToMessageArgs) {
   try {
-    ++streamingMessages.value
+    streamingMessagesMap[(streamId ?? resumeStreamId)!] = true
 
     const currentThreadId = threadIdRef.value
     const { response, abortController } = await postChatStream({
       threadId: currentThreadId as Id<'threads'>,
       ...chatContext.activeAgent.value,
       content,
+      streamId,
       resumeStreamId,
     })
 
@@ -315,6 +335,8 @@ async function streamToMessage({ message, content, resumeStreamId }: StreamToMes
       if (state.content)
         message.content += state.content
 
+      console.log({ chunk, write: state.content })
+
       nextTick(() => { doScrollBottom({ maybe: true }) })
     }
 
@@ -325,14 +347,14 @@ async function streamToMessage({ message, content, resumeStreamId }: StreamToMes
     message!.content += `\nError: ${(error as Error).message}`
   }
   finally {
-    --streamingMessages.value
+    delete streamingMessagesMap[(streamId ?? resumeStreamId)!]
   }
 
   console.log('Stream completed')
 }
 
 async function _branchThreadFromMessage({ messageId, lockerKey }: BranchThreadFromMessageArgs) {
-  if (streamingMessages.value > 0)
+  if (Object.keys(streamingMessagesMap).length > 0)
     throw new Error('Can not branch while streaming')
 
   const messagesLte = messages.value.slice(0, messages.value.findIndex(m => m._id === messageId) + 1)
@@ -417,7 +439,7 @@ function doScrollBottom({ smooth = true, maybe = false, tries = 0, lastScrollTop
     </VueLenis>
 
     <PrompterArea
-      v-bind="{ nearTopBottom, lenisRef, streamingMessages }"
+      v-bind="{ nearTopBottom, lenisRef, streamingMessagesMap }"
       v-model:chat-input="chatInput"
       @submit="(input) => handleSubmit({ input })"
     />

apps/frontend/app/components/chat/interface/PrompterArea.vue

@@ -4,7 +4,7 @@ import type Lenis from 'lenis'
 const props = defineProps<{
   nearTopBottom: Array<null | boolean | number>
   lenisRef: undefined | { $el: HTMLElement, lenis: Lenis }
-  streamingMessages: number
+  streamingMessagesMap: Record<string, true>
 }>()
 
 const emit = defineEmits<{
@@ -28,7 +28,7 @@ function handleSubmit({ confirmMultiStream = false }) {
   if (!userInput)
     return
 
-  if (!confirmMultiStream && props.streamingMessages > 0) {
+  if (!confirmMultiStream && Object.keys(props.streamingMessagesMap).length > 0) {
     multiStreamConfirmDialogOpen.value = true
     return
   }

apps/frontend/app/utils/chat.ts

@@ -81,6 +81,7 @@ export interface PostChatStreamArgs {
   model: string
   apiKey?: string
   content?: string
+  streamId?: string
   resumeStreamId?: string
   finishOnly?: boolean
   abortController?: AbortController
@@ -93,7 +94,7 @@ export async function postChatStream(args: PostChatStreamArgs) {
   const { convexApiUrl } = useRuntimeConfig().public
   const { $auth } = useNuxtApp()
 
-  const response = await fetch(`${convexApiUrl}/api/ai/chat/stream`, {
+  const response = await fetch(`${convexApiUrl}/api/chat/stream`, {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',