Improve chat navigation, model selection, and streaming

components/chat/ChatAgent.tsx

@@ -1,20 +1,12 @@
 'use client';
 
-import { useState, useCallback, useRef, useEffect } from 'react';
-import type { ChatStatus } from 'ai';
-import { ChatInput, type PromptInputMessage } from './ChatInput';
+import { ChatInput } from './ChatInput';
 import { ChatMessages } from './ChatMessages';
-import { ChatPrompt, PromptSuggestion } from './ChatPrompt';
-import { ChatSidebar, type Chat } from './ChatSidebar';
-import { useChat, UIMessage } from '@ai-sdk/react';
-import { DefaultChatTransport } from 'ai';
-import {
-  createChat,
-  deleteChat,
-  getMessages,
-  listChats,
-  saveMessages,
-} from '@/services/indexed-db/chat-store';
+import { ChatPrompt } from './ChatPrompt';
+import { ChatSidebar } from './ChatSidebar';
+import { useChatSession } from '@/hooks/use-chat-session';
+import type { PromptSuggestion } from '@/config/suggestions';
+import { DEFAULT_OPENROUTER_MODEL_ID, type OpenRouterModelId } from '@/lib/openrouter';
 import {
   SidebarProvider,
   SidebarInset,
@@ -24,228 +16,32 @@ import {
 interface ChatAgentProps {
   suggestions?: PromptSuggestion[];
   webSearchEnabled?: boolean;
-}
-
-async function convertToDataUrl(url: string): Promise<string> {
-  const response = await fetch(url);
-  if (!response.ok) throw new Error('Could not read the selected file');
-
-  const blob = await response.blob();
-
-  return new Promise((resolve, reject) => {
-    const reader = new FileReader();
-    reader.onload = () => {
-      if (typeof reader.result === 'string') {
-        resolve(reader.result);
-      } else {
-        reject(new Error('Could not encode the selected file'));
-      }
-    };
-    reader.onerror = () => reject(reader.error);
-    reader.readAsDataURL(blob);
-  });
+  defaultModelId?: OpenRouterModelId;
 }
 
 export function ChatAgent({
   suggestions,
   webSearchEnabled = false,
+  defaultModelId = DEFAULT_OPENROUTER_MODEL_ID,
 }: ChatAgentProps) {
-  const [inputValue, setInputValue] = useState('');
-  const [chatId, setChatId] = useState<string | null>(null);
-  const [chats, setChats] = useState<Chat[]>([]);
-  const chatIdRef = useRef(chatId);
-  const pendingChatIdRef = useRef<string | null>(null);
-  chatIdRef.current = chatId;
-
-  const fetchChats = useCallback(async () => {
-    try {
-      setChats(await listChats());
-    } catch (error) {
-      console.error('Failed to fetch chats:', error);
-    }
-  }, []);
-
-  const transportRef = useRef(
-    new DefaultChatTransport({
-      api: '/api/chat',
-    })
-  );
-
   const {
+    inputValue,
+    setInputValue,
+    modelId,
+    setModelId,
+    chatId,
+    chats,
+    chatMessages,
     status,
-    messages: chatMessages,
-    setMessages: setChatMessages,
-    sendMessage,
-    stop,
-    error: aiError,
-    addToolApprovalResponse,
-  } = useChat({
-    transport: transportRef.current,
-    onFinish: async ({ messages }) => {
-      const activeChatId = pendingChatIdRef.current ?? chatIdRef.current;
-      if (!activeChatId) return;
-
-      try {
-        await saveMessages(activeChatId, messages);
-        await fetchChats();
-      } catch (error) {
-        console.error('Failed to save messages:', error);
-      } finally {
-        if (pendingChatIdRef.current === activeChatId) {
-          pendingChatIdRef.current = null;
-        }
-      }
-    },
-    onError: (error) => {
-      console.error('Chat error:', error);
-    },
-  });
-
-  const [isPreparingFiles, setIsPreparingFiles] = useState(false);
-
-  const chatStatus: ChatStatus = isPreparingFiles
-    ? 'submitted'
-    : status === 'submitted' || status === 'streaming'
-      ? status
-      : 'ready';
-
-  useEffect(() => {
-    fetchChats();
-  }, [fetchChats]);
-
-  const handleSelectChat = useCallback(
-    async (selectedChatId: string) => {
-      try {
-        const messages = await getMessages(selectedChatId);
-        chatIdRef.current = selectedChatId;
-        setChatId(selectedChatId);
-
-        setChatMessages(messages);
-      } catch (error) {
-        console.error('Failed to load chat:', error);
-      }
-    },
-    [setChatMessages]
-  );
-
-  const handleNewChat = useCallback(() => {
-    chatIdRef.current = null;
-    setChatId(null);
-    setChatMessages([]);
-    setInputValue('');
-  }, [setChatMessages]);
-
-  const handleDeleteChat = useCallback(
-    async (deletedChatId: string) => {
-      try {
-        await deleteChat(deletedChatId);
-        setChats((prev) => prev.filter((c) => c.id !== deletedChatId));
-        if (chatId === deletedChatId) {
-          chatIdRef.current = null;
-          setChatId(null);
-          setChatMessages([]);
-        }
-      } catch (error) {
-        console.error('Failed to delete chat:', error);
-      }
-    },
-    [chatId, setChatMessages]
-  );
-
-  const handleSendMessage = useCallback(
-    async (message: PromptInputMessage) => {
-      const trimmedValue = message.text?.trim();
-      if (!trimmedValue && !message.files?.length) return;
-      if (
-        status === 'submitted' ||
-        status === 'streaming' ||
-        isPreparingFiles
-      )
-        return;
-
-      let activeChatId = chatIdRef.current;
-
-      if (!chatIdRef.current) {
-        const newChatId = crypto.randomUUID();
-        const title = trimmedValue ? trimmedValue.slice(0, 100) : 'New Chat';
-
-        activeChatId = newChatId;
-        chatIdRef.current = newChatId;
-        setChatId(newChatId);
-
-        try {
-          await createChat(newChatId, title);
-          await fetchChats();
-        } catch (error) {
-          console.error('Failed to create chat:', error);
-        }
-      }
-
-      const parts: Array<
-        | { type: 'text'; text: string }
-        | { type: 'file'; url: string; mediaType: string; filename: string }
-      > = [];
-
-      if (trimmedValue) {
-        parts.push({ type: 'text', text: trimmedValue });
-      }
-
-      // PromptInput supplies files as base64 data URLs. Keep them inline so the
-      // demo never needs a remote storage bucket.
-      if (message.files?.length) {
-        setIsPreparingFiles(true);
-        for (const file of message.files) {
-          try {
-            const dataUrl = file.url.startsWith('data:')
-              ? file.url
-              : await convertToDataUrl(file.url);
-
-            parts.push({
-              type: 'file',
-              url: dataUrl,
-              mediaType: file.mediaType || 'application/octet-stream',
-              filename: file.filename || 'file',
-            });
-          } catch (err) {
-            console.error('File encoding error:', err);
-          }
-        }
-        setIsPreparingFiles(false);
-      }
-
-      if (parts.length === 0) return;
-
-      const userMessage: UIMessage = {
-        id: crypto.randomUUID(),
-        role: 'user',
-        parts,
-      };
-
-      setInputValue('');
-
-      if (activeChatId) {
-        try {
-          await saveMessages(activeChatId, [...chatMessages, userMessage]);
-        } catch (error) {
-          console.error('Failed to save user message:', error);
-        }
-      }
-
-      pendingChatIdRef.current = activeChatId;
-      await sendMessage(userMessage);
-    },
-    [
-      status,
-      sendMessage,
-      chatMessages,
-      isPreparingFiles,
-      fetchChats,
-    ]
-  );
-
-  const handleStop = useCallback(() => {
-    stop();
-  }, [stop]);
+    chatStatus,
+    aiError,
+    isPreparingFiles,
+    handleSelectChat,
+    handleNewChat,
+    handleDeleteChat,
+    handleSendMessage,
+    handleStop,
+  } = useChatSession(defaultModelId);
 
   const hasMessages = chatMessages.length > 0;
 
@@ -268,6 +64,9 @@ export function ChatAgent({
               value={inputValue}
               onChange={setInputValue}
               onSubmit={handleSendMessage}
+              modelId={modelId}
+              onModelChange={setModelId}
+              onStop={isPreparingFiles ? undefined : handleStop}
               status={chatStatus}
               suggestions={suggestions}
               webSearchEnabled={webSearchEnabled}
@@ -281,15 +80,23 @@ export function ChatAgent({
             <div className="flex-1 overflow-hidden">
               <ChatMessages
                 messages={chatMessages}
-                isLoading={status === 'submitted' || status === 'streaming'}
+                status={status}
               />
             </div>
             <div className="sticky bottom-0 bg-muted p-4">
               <div className="mx-auto max-w-3xl">
+                {aiError && (
+                  <p role="alert" className="mb-3 text-sm text-destructive">
+                    The response could not be completed. Please try again.
+                  </p>
+                )}
                 <ChatInput
                   value={inputValue}
                   onChange={setInputValue}
                   onSubmit={handleSendMessage}
+                  modelId={modelId}
+                  onModelChange={setModelId}
+                  onStop={isPreparingFiles ? undefined : handleStop}
                   status={chatStatus}
                 />
               </div>
@@ -301,4 +108,4 @@ export function ChatAgent({
   );
 }
 
-export type { UIMessage as Message };
+export type { UIMessage as Message } from 'ai';