feat: migrate from ai/rsc to ai/ui (useChat)

app/page.tsx

@@ -1,18 +1,17 @@
 import { Chat } from '@/components/chat'
 import { nanoid } from '@/lib/utils'
-import { AI } from './actions'
+import { ChatProvider } from '@/components/chat-provider'
+import { MapDataProvider } from '@/components/map/map-data-context'
 
 export const maxDuration = 60
 
-import { MapDataProvider } from '@/components/map/map-data-context'
-
 export default function Page() {
   const id = nanoid()
   return (
-    <AI initialAIState={{ chatId: id, messages: [] }}>
+    <ChatProvider chatId={id}>
       <MapDataProvider>
         <Chat id={id} />
       </MapDataProvider>
-    </AI>
+    </ChatProvider>
   )
 }

app/search/[id]/page.tsx

@@ -1,76 +1,58 @@
 import { notFound, redirect } from 'next/navigation';
 import { Chat } from '@/components/chat';
-import { getChat, getChatMessages } from '@/lib/actions/chat'; // Added getChatMessages
-import { AI } from '@/app/actions';
+import { ChatProvider } from '@/components/chat-provider';
+import { getChat, getChatMessages } from '@/lib/actions/chat';
 import { MapDataProvider } from '@/components/map/map-data-context';
-import { getCurrentUserIdOnServer } from '@/lib/auth/get-current-user'; // For server-side auth
-import type { AIMessage } from '@/lib/types'; // For AIMessage type
-import type { Message as DrizzleMessage } from '@/lib/actions/chat-db'; // For DrizzleMessage type
+import { getCurrentUserIdOnServer } from '@/lib/auth/get-current-user';
+import type { Message } from 'ai/react';
+import type { Message as DrizzleMessage } from '@/lib/actions/chat-db';
 
 export const maxDuration = 60;
 
 export interface SearchPageProps {
-  params: Promise<{ id: string }>; // Keep as is for now
+  params: Promise<{ id: string }>;
 }
 
 export async function generateMetadata({ params }: SearchPageProps) {
-  const { id } = await params; // Keep as is for now
-  // TODO: Metadata generation might need authenticated user if chats are private
-  // For now, assuming getChat can be called or it handles anon access for metadata appropriately
-  const userId = await getCurrentUserIdOnServer(); // Attempt to get user for metadata
-  const chat = await getChat(id, userId || 'anonymous'); // Pass userId or 'anonymous' if none
+  const { id } = await params;
+  const userId = await getCurrentUserIdOnServer();
+  const chat = await getChat(id, userId || 'anonymous');
   return {
     title: chat?.title?.toString().slice(0, 50) || 'Search',
   };
 }
 
 export default async function SearchPage({ params }: SearchPageProps) {
-  const { id } = await params; // Keep as is for now
+  const { id } = await params;
   const userId = await getCurrentUserIdOnServer();
 
   if (!userId) {
-    // If no user, redirect to login or show appropriate page
-    // For now, redirecting to home, but a login page would be better.
     redirect('/');
   }
 
   const chat = await getChat(id, userId);
 
   if (!chat) {
-    // If chat doesn't exist or user doesn't have access (handled by getChat)
     notFound();
   }
 
-  // Fetch messages for the chat
   const dbMessages: DrizzleMessage[] = await getChatMessages(chat.id);
 
-  // Transform DrizzleMessages to AIMessages
-  const initialMessages: AIMessage[] = dbMessages.map((dbMsg): AIMessage => {
-    return {
+  const validRoles = new Set(['user', 'assistant', 'system'])
+  const initialMessages: Message[] = dbMessages
+    .filter((dbMsg) => validRoles.has(dbMsg.role))
+    .map((dbMsg): Message => ({
       id: dbMsg.id,
-      role: dbMsg.role as AIMessage['role'], // Cast role, ensure AIMessage['role'] includes all dbMsg.role possibilities
+      role: dbMsg.role as Message['role'],
       content: dbMsg.content,
       createdAt: dbMsg.createdAt ? new Date(dbMsg.createdAt) : undefined,
-      // 'type' and 'name' are not in the basic Drizzle 'messages' schema.
-      // These would be undefined unless specific logic is added to derive them.
-      // For instance, if a message with role 'tool' should have a 'name',
-      // or if some messages have a specific 'type' based on content or other flags.
-      // This mapping assumes standard user/assistant messages primarily.
-    };
-  });
+    }));
 
   return (
-    <AI
-      initialAIState={{
-        chatId: chat.id,
-        messages: initialMessages, // Use the transformed messages from the database
-        // isSharePage: true, // This was in PR#533, but share functionality is removed.
-                             // If needed for styling or other logic, it can be set.
-      }}
-    >
+    <ChatProvider chatId={chat.id} initialMessages={initialMessages}>
       <MapDataProvider>
         <Chat id={id} />
       </MapDataProvider>
-    </AI>
+    </ChatProvider>
   );
-}
+}

components/chat-messages.tsx

@@ -1,70 +1,128 @@
 'use client'
 
-import { StreamableValue, useUIState } from 'ai/rsc'
-import type { AI, UIState } from '@/app/actions'
+import type { Message } from 'ai/react'
 import { CollapsibleMessage } from './collapsible-message'
+import { Section } from './section'
+import { BotMessage } from './message'
+import { UserMessage } from './user-message'
+import { ToolResultRenderer } from './tool-result-renderer'
+import { useChatContext, type Annotation } from './chat-provider'
+import { Copilot } from './copilot'
+import SearchRelated from './search-related'
 
 interface ChatMessagesProps {
-  messages: UIState
+  messages: Message[]
 }
 
 export function ChatMessages({ messages }: ChatMessagesProps) {
-  if (!messages.length) {
+  const { annotations, isLoading } = useChatContext()
+
+  if (!messages.length && !annotations.length) {
     return null
   }
 
-  // Group messages based on ID, and if there are multiple messages with the same ID, combine them into one message
-  const groupedMessages = messages.reduce(
-    (acc: { [key: string]: any }, message) => {
-      if (!acc[message.id]) {
-        acc[message.id] = {
+  const renderedMessages: {
+    id: string
+    component: React.ReactNode
+    isCollapsed?: boolean
+    isAssistant?: boolean
+  }[] = []
+
+  // Render tool result annotations first (they come before the text)
+  const toolAnnotations = annotations.filter((a: Annotation) => a.type === 'tool_result')
+  for (let i = 0; i < toolAnnotations.length; i++) {
+    const ann = toolAnnotations[i]
+    renderedMessages.push({
+      id: `tool-${ann.toolName}-${i}`,
+      component: <ToolResultRenderer toolName={ann.toolName} result={ann.result} />,
+      isCollapsed: true
+    })
+  }
+
+  // Render chat messages
+  for (const message of messages) {
+    if (message.role === 'user') {
+      renderedMessages.push({
+        id: message.id,
+        component: <UserMessage content={message.content} />
+      })
+    } else if (message.role === 'assistant') {
+      if (message.content) {
+        renderedMessages.push({
           id: message.id,
-          components: [],
-          isCollapsed: message.isCollapsed
+          component: (
+            <Section title="response">
+              <BotMessage content={message.content} />
+            </Section>
+          ),
+          isAssistant: true
+        })
+      }
+
+      // Render tool invocations
+      if (message.toolInvocations) {
+        for (const invocation of message.toolInvocations) {
+          if (invocation.state === 'result') {
+            renderedMessages.push({
+              id: `${message.id}-tool-${invocation.toolCallId}`,
+              component: (
+                <ToolResultRenderer
+                  toolName={invocation.toolName}
+                  result={invocation.result}
+                />
+              ),
+              isCollapsed: true
+            })
+          }
         }
       }
-      acc[message.id].components.push(message.component)
-      return acc
-    },
-    {}
-  )
+    }
+  }
 
-  // Convert grouped messages into an array with explicit type
-  const groupedMessagesArray = Object.values(groupedMessages).map(group => ({
-    ...group,
-    components: group.components as React.ReactNode[]
-  })) as {
-    id: string
-    components: React.ReactNode[]
-    isCollapsed?: StreamableValue<boolean>
-  }[]
+  // Render inquiry annotation if present
+  const inquiry = annotations.find((a: Annotation) => a.type === 'inquiry')
+  if (inquiry) {
+    renderedMessages.push({
+      id: 'inquiry',
+      component: <Copilot inquiry={{ value: inquiry.data }} />
+    })
+  }
+
+  // Render related queries annotation
+  const related = annotations.findLast?.((a: Annotation) => a.type === 'related')
+  if (related && related.relatedQueries?.items?.length > 0) {
+    renderedMessages.push({
+      id: 'related',
+      component: (
+        <Section title="Related" separator={true}>
+          <SearchRelated relatedQueries={related.relatedQueries} />
+        </Section>
+      )
+    })
+  }
+
+  // Find last assistant message index for isLastMessage prop
+  let lastAssistantIndex = -1
+  for (let i = renderedMessages.length - 1; i >= 0; i--) {
+    if (renderedMessages[i].isAssistant) {
+      lastAssistantIndex = i
+      break
+    }
+  }
 
   return (
     <>
-      {groupedMessagesArray.map(
-        (
-          groupedMessage: {
-            id: string
-            components: React.ReactNode[]
-            isCollapsed?: StreamableValue<boolean>
-          },
-          index
-        ) => (
-          <CollapsibleMessage
-            key={`${groupedMessage.id}`}
-            message={{
-              id: groupedMessage.id,
-              component: groupedMessage.components.map((component, i) => (
-                <div key={`${groupedMessage.id}-${i}`}>{component}</div>
-              )),
-              isCollapsed: groupedMessage.isCollapsed
-            }}
-            isLastMessage={
-              groupedMessage.id === messages[messages.length - 1].id
-            }
-          />
-        )
-      )}
+      {renderedMessages.map((msg, index) => (
+        <CollapsibleMessage
+          key={msg.id}
+          message={{
+            id: msg.id,
+            component: <div>{msg.component}</div>,
+            isCollapsed: msg.isCollapsed
+          }}
+          isLastMessage={index === lastAssistantIndex}
+        />
+      ))}
     </>
   )
 }