Update AI message handling in useChat hook#109
Conversation
|
@sbansal1999 is attempting to deploy a commit to the Databuddy Team on Vercel. A member of the Team first needs to authorize it. |
WalkthroughThe hook now maintains a single in-flight assistant message during streaming. It inserts a placeholder assistant message, incrementally updates it on each SSE chunk via an internal updateAiMessage helper, and omits appending duplicates. Final content is persisted on completion. Unused import removed; error/rate-limit handling unchanged. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Hook as useChat hook
participant SSE as Stream (SSE)
participant IDB as IndexedDB
User->>Hook: sendMessage(userMessage)
Note over Hook: Create assistantMessage (working copy)
Hook->>Hook: Insert placeholder assistant message (empty content)
Hook->>SSE: Open stream with userMessage
loop For each SSE chunk
SSE-->>Hook: delta/content update
Hook->>Hook: Merge delta into assistantMessage
Hook->>Hook: updateAiMessage(assistantMessage)<br/>(update last message in array)
end
alt Stream complete
Hook->>IDB: Persist final assistantMessage
else Error / 429
Hook->>Hook: Handle error / cooldown logic
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.ts (3)
157-167: Give the assistant placeholder a stable temp id to make updates reliable.Without a stable id, in-flight updates can’t reliably locate the correct message. Assign a temporary id when inserting the placeholder and reconcile it when the server sends
metadata.messageId.Apply this diff:
- let assistantMessage: Message = { - id: '', + const placeholderId = `ai-${Date.now()}-${Math.random().toString(36).slice(2)}`; + let assistantMessage: Message = { + id: placeholderId, type: 'assistant', content: '', timestamp: new Date(), hasVisualization: false, thinkingSteps: [], }; - - setMessages((prev) => [...prev, assistantMessage]); + setMessages((prev) => [...prev, assistantMessage]);And keep the existing
metadatahandler (Lines 306-311) that setsassistantMessage.id = update.data.messageId; with the updatedupdateAiMessageit will swap the temp id entry in-place once metadata arrives.
191-216: Rate-limit and generic error branches never surface the assistant message to the UI.In both the 429 flow and the catch block you mutate
assistantMessagebut don’t callupdateAiMessage, leaving the placeholder empty on screen. Persisting to IndexedDB is optional, but the UI must be updated.Apply this diff:
if (response.status === 429) { const errorData = await response.json(); if (errorData.code === 'RATE_LIMIT_EXCEEDED') { assistantMessage = { ...assistantMessage, content: "⏱️ You've reached the rate limit. Please wait 60 seconds before sending another message.", }; setIsLoading(false); setIsRateLimited(true); + // Surface the message to the UI (and optionally persist) + updateAiMessage(assistantMessage); + try { + await chatDB.saveMessage(assistantMessage, websiteId || ''); + } catch (e) { + console.error('Failed to save rate-limited message:', e); + } // Clear any existing timeout} catch (error) { console.error('Failed to get AI response:', error); assistantMessage = { ...assistantMessage, content: "I apologize, but I'm having trouble processing your request right now. Please try again in a moment.", }; + // Ensure the placeholder shows the error + updateAiMessage(assistantMessage); + try { + await chatDB.saveMessage(assistantMessage, websiteId || ''); + } catch (e) { + console.error('Failed to save error message to IndexedDB:', e); + } } finally { setIsLoading(false); }Also applies to: 327-336
221-226: SSE parsing isn’t chunk-safe; updates can be dropped when JSON spans chunks.Splitting each
Uint8Arrayby'\n'assumes chunk boundaries align with SSE events, which isn’t guaranteed. If a JSON event crosses chunk boundaries, it gets lost. Use a rolling buffer and split on\n\n(event terminator) to parse complete events.Apply this diff to buffer and parse events safely:
- const reader = response.body?.getReader(); + const reader = response.body?.getReader(); if (!reader) { throw new Error('No response stream available'); } try { - while (true) { - const { done, value } = await reader.read(); - if (done) { - break; - } - - const chunk = new TextDecoder().decode(value); - const lines = chunk.split('\n'); - - for (const line of lines) { - if (line.startsWith('data: ')) { - try { - const update: StreamingUpdate = JSON.parse(line.slice(6)); - - switch (update.type) { + const decoder = new TextDecoder(); + let buffer = ''; + while (true) { + const { done, value } = await reader.read(); + if (done) { + // Optionally flush any trailing event without final \n\n + if (buffer.includes('data:')) { + const dataLines = buffer + .split('\n') + .filter((l) => l.startsWith('data:')); + const data = dataLines + .map((l) => l.replace(/^data:\s?/, '')) + .join(''); + if (data) { + try { + const update: StreamingUpdate = JSON.parse(data); + switch (update.type) { + case 'thinking': { + assistantMessage = { + ...assistantMessage, + thinkingSteps: [ + ...(assistantMessage.thinkingSteps || []), + update.content, + ], + }; + break; + } + case 'progress': { + assistantMessage = { + ...assistantMessage, + content: update.content, + hasVisualization: update.data?.hasVisualization, + chartType: update.data?.chartType as any, + data: update.data?.data, + responseType: update.data?.responseType, + metricValue: update.data?.metricValue, + metricLabel: update.data?.metricLabel, + }; + scrollToBottom(); + break; + } + case 'complete': { + assistantMessage = { + ...assistantMessage, + type: 'assistant' as const, + content: update.content, + timestamp: new Date(), + hasVisualization: update.data?.hasVisualization, + chartType: update.data?.chartType as any, + data: update.data?.data, + responseType: update.data?.responseType, + metricValue: update.data?.metricValue, + metricLabel: update.data?.metricLabel, + debugInfo: update.debugInfo, + }; + try { + await chatDB.saveMessage(assistantMessage, websiteId || ''); + } catch (error) { + console.error('Failed to save Databunny message to IndexedDB:', error); + } + scrollToBottom(); + break; + } + case 'error': { + assistantMessage = { + ...assistantMessage, + content: update.content, + debugInfo: update.debugInfo, + }; + break; + } + case 'metadata': { + assistantMessage = { + ...assistantMessage, + id: update.data.messageId, + }; + setConversationId(update.data.conversationId); + break; + } + default: { + break; + } + } + updateAiMessage(assistantMessage); + } catch (_parseError) { + console.warn('Failed to parse SSE data (flush):', buffer); + } + } + } + break; + } + + buffer += decoder.decode(value, { stream: true }); + let boundaryIdx; + while ((boundaryIdx = buffer.indexOf('\n\n')) !== -1) { + const eventBlock = buffer.slice(0, boundaryIdx); + buffer = buffer.slice(boundaryIdx + 2); + + const dataLines = eventBlock + .split('\n') + .filter((l) => l.startsWith('data:')); + const data = dataLines + .map((l) => l.replace(/^data:\s?/, '')) + .join(''); + if (!data) continue; + + try { + const update: StreamingUpdate = JSON.parse(data); + + switch (update.type) { case 'thinking': { assistantMessage = { ...assistantMessage, thinkingSteps: [ ...(assistantMessage.thinkingSteps || []), update.content, ], }; break; } case 'progress': { assistantMessage = { ...assistantMessage, content: update.content, hasVisualization: update.data?.hasVisualization, chartType: update.data?.chartType as any, data: update.data?.data, responseType: update.data?.responseType, metricValue: update.data?.metricValue, metricLabel: update.data?.metricLabel, }; scrollToBottom(); break; } case 'complete': { assistantMessage = { ...assistantMessage, type: 'assistant' as const, content: update.content, timestamp: new Date(), hasVisualization: update.data?.hasVisualization, chartType: update.data?.chartType as any, data: update.data?.data, responseType: update.data?.responseType, metricValue: update.data?.metricValue, metricLabel: update.data?.metricLabel, debugInfo: update.debugInfo, }; - // Save completed assistant message to IndexedDB + // Save completed assistant message to IndexedDB try { await chatDB.saveMessage( assistantMessage, websiteId || '' ); } catch (error) { console.error( 'Failed to save Databunny message to IndexedDB:', error ); } scrollToBottom(); break; } case 'error': { assistantMessage = { ...assistantMessage, content: update.content, debugInfo: update.debugInfo, }; break; } case 'metadata': { assistantMessage = { ...assistantMessage, id: update.data.messageId, }; setConversationId(update.data.conversationId); break; } default: { break; } - } - updateAiMessage(assistantMessage); - } catch (_parseError) { - console.warn('Failed to parse SSE data:', line); - } - } - } + } + updateAiMessage(assistantMessage); + } catch (_parseError) { + console.warn('Failed to parse SSE data (event):', eventBlock); + } + } }Also applies to: 227-323
🧹 Nitpick comments (1)
apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.ts (1)
317-317: Optional: micro-batch UI updates to reduce re-renders during streaming.Frequent
setMessagescan thrash rendering on long answers. Consider coalescingupdateAiMessageinto animation frames.Example helper (outside the hook) and usage:
// Simple rAF throttle const rafThrottle = <T extends (...args: any[]) => void>(fn: T) => { let queued = false; let lastArgs: Parameters<T>; return (...args: Parameters<T>) => { lastArgs = args; if (queued) return; queued = true; requestAnimationFrame(() => { queued = false; fn(...lastArgs); }); }; }; // inside useChat, wrap once: const updateAiMessageThrottled = useMemo(() => rafThrottle(updateAiMessage), [updateAiMessage]); // then call: updateAiMessageThrottled(assistantMessage);
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.ts(3 hunks)
| function updateAiMessage(message: Message) { | ||
| setMessages((prev) => { | ||
| //TODO: find a way to update the message with the correct id | ||
| const newMessages = [...prev]; | ||
| newMessages[newMessages.length - 1] = message; | ||
| return newMessages; | ||
| }); | ||
| } | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Don’t blindly replace the last entry; update by id with a safe fallback.
Replacing prev[prev.length - 1] risks overwriting a user message if the placeholder append hasn’t committed yet or if other state updates interleave. Also add a guard for empty arrays and merge fields to avoid accidental data loss.
Apply this diff:
-function updateAiMessage(message: Message) {
- setMessages((prev) => {
- //TODO: find a way to update the message with the correct id
- const newMessages = [...prev];
- newMessages[newMessages.length - 1] = message;
- return newMessages;
- });
-}
+function updateAiMessage(updated: Message) {
+ setMessages((prev) => {
+ if (!prev.length) return prev;
+ // Prefer id match when available
+ let idx = -1;
+ if (updated.id) {
+ // Avoid findLastIndex for broader runtime support
+ for (let i = prev.length - 1; i >= 0; i--) {
+ if (prev[i].id === updated.id) {
+ idx = i;
+ break;
+ }
+ }
+ }
+ // Fallback: last assistant message (covers placeholder before server id arrives)
+ if (idx === -1) {
+ for (let i = prev.length - 1; i >= 0; i--) {
+ if (prev[i].type === 'assistant') {
+ idx = i;
+ break;
+ }
+ }
+ }
+ // If still not found, append to avoid clobbering user messages
+ if (idx === -1) return [...prev, updated];
+
+ const newMessages = prev.slice();
+ newMessages[idx] = { ...newMessages[idx], ...updated };
+ return newMessages;
+ });
+}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| function updateAiMessage(message: Message) { | |
| setMessages((prev) => { | |
| //TODO: find a way to update the message with the correct id | |
| const newMessages = [...prev]; | |
| newMessages[newMessages.length - 1] = message; | |
| return newMessages; | |
| }); | |
| } | |
| function updateAiMessage(updated: Message) { | |
| setMessages((prev) => { | |
| if (!prev.length) return prev; | |
| // Prefer id match when available | |
| let idx = -1; | |
| if (updated.id) { | |
| // Avoid findLastIndex for broader runtime support | |
| for (let i = prev.length - 1; i >= 0; i--) { | |
| if (prev[i].id === updated.id) { | |
| idx = i; | |
| break; | |
| } | |
| } | |
| } | |
| // Fallback: last assistant message (covers placeholder before server id arrives) | |
| if (idx === -1) { | |
| for (let i = prev.length - 1; i >= 0; i--) { | |
| if (prev[i].type === 'assistant') { | |
| idx = i; | |
| break; | |
| } | |
| } | |
| } | |
| // If still not found, append to avoid clobbering user messages | |
| if (idx === -1) return [...prev, updated]; | |
| const newMessages = prev.slice(); | |
| newMessages[idx] = { ...newMessages[idx], ...updated }; | |
| return newMessages; | |
| }); | |
| } |
2 participants
This pull request refactors the message update logic in the
useChathook to ensure that AI-generated messages are properly updated during streaming responses, rather than being appended multiple times. The changes streamline how assistant messages are managed in state, improving consistency and preventing duplicate entries.Message update logic improvements:
updateAiMessagefunction to update the latest assistant message in themessagesstate, ensuring that streamed AI responses are reflected in place rather than appended as new messages. (apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.tsR123-R131)updateAiMessagefor each streamed chunk, replacing the previous approach of appending messages. (apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.tsR317)Code cleanup:
trpcimport fromuse-chat.tsto tidy up dependencies. (apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.tsL4)Summary by CodeRabbit