refactor: unify message rendering pattern for tools and annotations

components/chat-messages.tsx

@@ -1,5 +1,4 @@
 import { Message } from 'ai'
-import RelatedQuestions from './related-questions'
 import { RenderMessage } from './render-message'
 import { Spinner } from './ui/spinner'
 import { useState, useEffect } from 'react'
@@ -30,37 +29,32 @@ export function ChatMessages({
     messages.length -
     1 -
     [...messages].reverse().findIndex(msg => msg.role === 'user')
+
   const showSpinner = isLoading && messages[messages.length - 1].role === 'user'
 
+  const getIsOpen = (id: string) => {
+    const index = messages.findIndex(msg => msg.id === id.split('-')[0])
+    return openStates[id] ?? index >= lastUserIndex
+  }
+
+  const handleOpenChange = (id: string, open: boolean) => {
+    setOpenStates(prev => ({
+      ...prev,
+      [id]: open
+    }))
+  }
+
   return (
     <div className="relative mx-auto px-4 w-full">
-      {messages.map((message, index) => (
+      {messages.map(message => (
         <div key={message.id} className="mb-4 flex flex-col gap-4">
           <RenderMessage
             message={message}
-            isOpen={openStates[message.id] ?? index >= lastUserIndex}
-            onOpenChange={(open: boolean) => {
-              setOpenStates(prev => ({
-                ...prev,
-                [message.id]: open
-              }))
-            }}
+            messageId={message.id}
+            getIsOpen={getIsOpen}
+            onOpenChange={handleOpenChange}
+            onQuerySelect={onQuerySelect}
           />
-          {!message.toolInvocations && message.annotations && (
-            <RelatedQuestions
-              annotations={message.annotations}
-              onQuerySelect={onQuerySelect}
-              isOpen={
-                openStates[`${message.id}-related`] ?? index >= lastUserIndex
-              }
-              onOpenChange={(open: boolean) => {
-                setOpenStates(prev => ({
-                  ...prev,
-                  [`${message.id}-related`]: open
-                }))
-              }}
-            />
-          )}
         </div>
       ))}
       {showSpinner && <Spinner />}

components/render-message.tsx

@@ -2,17 +2,22 @@ import { Message } from 'ai'
 import { UserMessage } from './user-message'
 import { ToolSection } from './tool-section'
 import { AnswerSection } from './answer-section'
+import RelatedQuestions from './related-questions'
 
 interface RenderMessageProps {
   message: Message
-  isOpen: boolean
-  onOpenChange: (open: boolean) => void
+  messageId: string
+  getIsOpen: (id: string) => boolean
+  onOpenChange: (id: string, open: boolean) => void
+  onQuerySelect: (query: string) => void
 }
 
 export function RenderMessage({
   message,
-  isOpen,
-  onOpenChange
+  messageId,
+  getIsOpen,
+  onOpenChange,
+  onQuerySelect
 }: RenderMessageProps) {
   if (message.role === 'user') {
     return <UserMessage message={message.content} />
@@ -25,19 +30,29 @@ export function RenderMessage({
           <ToolSection
             key={tool.toolCallId}
             tool={tool}
-            isOpen={isOpen}
-            onOpenChange={onOpenChange}
+            isOpen={getIsOpen(messageId)}
+            onOpenChange={open => onOpenChange(messageId, open)}
           />
         ))}
       </>
     )
   }
 
   return (
-    <AnswerSection
-      content={message.content}
-      isOpen={isOpen}
-      onOpenChange={onOpenChange}
-    />
+    <>
+      <AnswerSection
+        content={message.content}
+        isOpen={getIsOpen(messageId)}
+        onOpenChange={open => onOpenChange(messageId, open)}
+      />
+      {!message.toolInvocations && message.annotations && (
+        <RelatedQuestions
+          annotations={message.annotations}
+          onQuerySelect={onQuerySelect}
+          isOpen={getIsOpen(`${messageId}-related`)}
+          onOpenChange={open => onOpenChange(`${messageId}-related`, open)}
+        />
+      )}
+    </>
   )
 }