CR fixes

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Author: hubertzub-db

integrations/appkit-agent/client/src/components/ChatPanel.tsx

@@ -1,5 +1,4 @@
 import { useEffect, useState } from 'react';
-import { useSearchParams } from 'react-router-dom';
 import type { ReactNode } from 'react';
 import { useChatStream, type UseChatStreamOptions } from '@/hooks/use-chat-stream';
 import { ChatHeader, type ChatHeaderProps } from './chat-header';
@@ -43,14 +42,11 @@ export function ChatPanel({
   const chat = useChatStream(streamOptions);
 
   // Handle query param auto-send (for ?query=... deep-links)
-  let searchParams: URLSearchParams | undefined;
-  try {
-    [searchParams] = useSearchParams();
-  } catch {
-    // Not inside a router — skip query param handling
-  }
-
-  const query = searchParams?.get('query');
+  // Avoid router hooks here so ChatPanel can render outside a Router.
+  const query =
+    typeof window === 'undefined'
+      ? null
+      : new URLSearchParams(window.location.search).get('query');
   const [hasAppendedQuery, setHasAppendedQuery] = useState(false);
 
   useEffect(() => {

integrations/appkit-agent/client/src/components/chat.tsx

@@ -1,38 +1,29 @@
-import type { LanguageModelUsage, UIMessageChunk } from 'ai';
-import { useChat } from '@ai-sdk/react';
-import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { useSWRConfig } from 'swr';
+import type { LanguageModelUsage } from 'ai';
+import { useEffect, useState } from 'react';
+import { useSearchParams } from 'react-router-dom';
 import { ChatHeader } from '@/components/chat-header';
-import { fetchWithErrorHandlers, generateUUID } from '@/lib/utils';
-import { MultimodalInput } from './multimodal-input';
-import { Messages } from './messages';
+import { useChatStream } from '@/hooks/use-chat-stream';
+import { softNavigateToChatId } from '@/lib/navigation';
+import { useAppConfig } from '@/contexts/AppConfigContext';
 import type {
   Attachment,
   ChatMessage,
   FeedbackMap,
   VisibilityType,
   ClientSession,
 } from '@/types';
-import { ChatSDKError } from '@/types';
-import { apiUrl } from '@/lib/config';
-import { unstable_serialize } from 'swr/infinite';
-import { getChatHistoryPaginationKey } from './sidebar-history';
-import { toast } from './toast';
-import { useSearchParams } from 'react-router-dom';
-import { useChatVisibility } from '@/hooks/use-chat-visibility';
-import { isCredentialErrorMessage } from '@/lib/oauth-error-utils';
-import { ChatTransport } from '../lib/ChatTransport';
-import { softNavigateToChatId } from '@/lib/navigation';
-import { useAppConfig } from '@/contexts/AppConfigContext';
 import { Greeting } from './greeting';
+import { Messages } from './messages';
+import { MultimodalInput } from './multimodal-input';
 
 export function Chat({
   id,
   initialMessages,
   initialChatModel,
   initialVisibilityType,
   isReadonly,
-  initialLastContext,
+  session: _session,
+  initialLastContext: _initialLastContext,
   feedback = {},
   title,
 }: {
@@ -46,218 +37,16 @@ export function Chat({
   feedback?: FeedbackMap;
   title?: string;
 }) {
-  const { visibilityType } = useChatVisibility({
-    chatId: id,
-    initialVisibilityType,
-  });
-
-  const { mutate } = useSWRConfig();
   const { chatHistoryEnabled } = useAppConfig();
 
-  const [input, setInput] = useState<string>('');
-  const [_usage, setUsage] = useState<LanguageModelUsage | undefined>(
-    initialLastContext,
-  );
-
-  const [lastPart, setLastPart] = useState<UIMessageChunk | undefined>();
-  const lastPartRef = useRef<UIMessageChunk | undefined>(lastPart);
-  lastPartRef.current = lastPart;
-
-  // Single counter for resume attempts - reset when stream parts are received
-  const resumeAttemptCountRef = useRef(0);
-  const maxResumeAttempts = 3;
-
-  const abortController = useRef<AbortController | null>(new AbortController());
-  useEffect(() => {
-    return () => {
-      abortController.current?.abort('ABORT_SIGNAL');
-    };
-  }, []);
-
-  const fetchWithAbort = useMemo(() => {
-    return async (input: RequestInfo | URL, init?: RequestInit) => {
-      // useChat does not cancel /stream requests when the component is unmounted
-      const signal = abortController.current?.signal;
-      return fetchWithErrorHandlers(input, { ...init, signal });
-    };
-  }, []);
-
-  const stop = useCallback(() => {
-    abortController.current?.abort('USER_ABORT_SIGNAL');
-  }, []);
-
-  const isNewChat = initialMessages.length === 0;
-  const didFetchHistoryOnNewChat = useRef(false);
-  const fetchChatHistory = useCallback(() => {
-    mutate(unstable_serialize(getChatHistoryPaginationKey));
-  }, [mutate]);
-
-  // For new chats, the title arrives via a `data-title` stream part
-  // once backend title generation completes — no separate fetch needed.
-  const [streamTitle, setStreamTitle] = useState<string | undefined>();
-  const [titlePending, setTitlePending] = useState(false);
-  const displayTitle = title ?? streamTitle;
-
-  const {
-    messages,
-    setMessages,
-    sendMessage,
-    status,
-    resumeStream,
-    clearError,
-    addToolApprovalResponse,
-    regenerate,
-  } = useChat<ChatMessage>({
+  const chat = useChatStream({
     id,
-    messages: initialMessages,
-    experimental_throttle: 100,
-    generateId: generateUUID,
-    resume: id !== undefined && initialMessages.length > 0, // Enable automatic stream resumption
-    transport: new ChatTransport({
-      onStreamPart: (part) => {
-        if (isNewChat && !didFetchHistoryOnNewChat.current) {
-          fetchChatHistory();
-          if (chatHistoryEnabled) {
-            setTitlePending(true);
-          }
-          didFetchHistoryOnNewChat.current = true;
-        }
-        // Reset resume attempts when we successfully receive stream parts
-        resumeAttemptCountRef.current = 0;
-        setLastPart(part);
-      },
-      api: apiUrl('/'),
-      fetch: fetchWithAbort,
-      prepareSendMessagesRequest({ messages, id, body }) {
-        const lastMessage = messages.at(-1);
-        const isUserMessage = lastMessage?.role === 'user';
-
-        // For continuations (non-user messages like tool results), we must always
-        // send previousMessages because the tool result only exists client-side
-        // and hasn't been saved to the database yet.
-        const needsPreviousMessages = !chatHistoryEnabled || !isUserMessage;
-
-        return {
-          body: {
-            id,
-            // Only include message field for user messages (new messages)
-            // For continuation (assistant messages with tool results), omit message field
-            ...(isUserMessage ? { message: lastMessage } : {}),
-            selectedChatModel: initialChatModel,
-            selectedVisibilityType: visibilityType,
-            nextMessageId: generateUUID(),
-            // Send previous messages when:
-            // 1. Database is disabled (ephemeral mode) - always need client-side messages
-            // 2. Continuation request (tool results) - tool result only exists client-side
-            ...(needsPreviousMessages
-              ? {
-                  previousMessages: isUserMessage
-                    ? messages.slice(0, -1)
-                    : messages,
-                }
-              : {}),
-            ...body,
-          },
-        };
-      },
-      prepareReconnectToStreamRequest({ id }) {
-        return {
-          api: apiUrl(`/${id}/stream`),
-          credentials: 'include',
-        };
-      },
-    }),
-    onData: (dataPart) => {
-      if (dataPart.type === 'data-usage') {
-        setUsage(dataPart.data as LanguageModelUsage);
-      }
-      if (dataPart.type === 'data-title') {
-        setStreamTitle(dataPart.data as string);
-        setTitlePending(false);
-        fetchChatHistory();
-      }
-    },
-    onFinish: ({
-      isAbort,
-      isDisconnect,
-      isError,
-      messages: finishedMessages,
-    }) => {
-      didFetchHistoryOnNewChat.current = false;
-      setTitlePending(false);
-
-      // If user aborted, don't try to resume
-      if (isAbort) {
-        console.log('[Chat onFinish] Stream was aborted by user, not resuming');
-        fetchChatHistory();
-        return;
-      }
-
-      // Check if the last message contains an OAuth credential error
-      // If so, don't try to resume - the user needs to authenticate first
-      const lastMessage = finishedMessages?.at(-1);
-      const hasOAuthError = lastMessage?.parts?.some(
-        (part) =>
-          part.type === 'data-error' &&
-          typeof part.data === 'string' &&
-          isCredentialErrorMessage(part.data),
-      );
-
-      if (hasOAuthError) {
-        console.log(
-          '[Chat onFinish] OAuth credential error detected, not resuming',
-        );
-        fetchChatHistory();
-        clearError();
-        return;
-      }
-
-      // Determine if we should attempt to resume:
-      // 1. Stream didn't end with a 'finish' part (incomplete)
-      // 2. It was a disconnect/error that terminated the stream
-      // 3. We haven't exceeded max resume attempts
-      const streamIncomplete = lastPartRef.current?.type !== 'finish';
-      const shouldResume =
-        streamIncomplete &&
-        (isDisconnect || isError || lastPartRef.current === undefined);
-
-      if (shouldResume && resumeAttemptCountRef.current < maxResumeAttempts) {
-        console.log(
-          '[Chat onFinish] Resuming stream. Attempt:',
-          resumeAttemptCountRef.current + 1,
-        );
-        resumeAttemptCountRef.current++;
-        // Ref: https://github.com/vercel/ai/issues/8477#issuecomment-3603209884
-        queueMicrotask(() => {
-          resumeStream();
-        })
-      } else {
-        // Stream completed normally or we've exhausted resume attempts
-        if (resumeAttemptCountRef.current >= maxResumeAttempts) {
-          console.warn('[Chat onFinish] Max resume attempts reached');
-        }
-        fetchChatHistory();
-      }
-    },
-    onError: (error) => {
-      console.log('[Chat onError] Error occurred:', error);
-
-      // Only show toast for explicit ChatSDKError (backend validation errors)
-      // Other errors (network, schema validation) are handled silently or in message parts
-      if (error instanceof ChatSDKError) {
-        toast({
-          type: 'error',
-          description: error.message,
-        });
-      } else {
-        // Non-ChatSDKError: Could be network error or in-stream error
-        // Log but don't toast - errors during streaming may be informational
-        console.warn('[Chat onError] Error during streaming:', error.message);
-      }
-      // Note: We don't call resumeStream here because onError can be called
-      // while the stream is still active (e.g., for data-error parts).
-      // Resume logic is handled exclusively in onFinish.
-    },
+    initialMessages,
+    model: initialChatModel,
+    initialVisibility: initialVisibilityType,
+    isReadonly,
+    feedback,
+    title,
   });
 
   const [searchParams] = useSearchParams();
@@ -267,33 +56,34 @@ export function Chat({
 
   useEffect(() => {
     if (query && !hasAppendedQuery) {
-      sendMessage({
+      chat.sendMessage({
         role: 'user' as const,
         parts: [{ type: 'text', text: query }],
       });
 
       setHasAppendedQuery(true);
-      softNavigateToChatId(id, chatHistoryEnabled);
+      softNavigateToChatId(chat.id, chatHistoryEnabled);
     }
-  }, [query, sendMessage, hasAppendedQuery, id, chatHistoryEnabled]);
+  }, [query, chat.sendMessage, hasAppendedQuery, chat.id, chatHistoryEnabled]);
 
+  const [input, setInput] = useState<string>('');
   const [attachments, setAttachments] = useState<Array<Attachment>>([]);
 
   const inputElement = <MultimodalInput
-    chatId={id}
+    chatId={chat.id}
     input={input}
     setInput={setInput}
-    status={status}
-    stop={stop}
+    status={chat.status}
+    stop={chat.stop}
     attachments={attachments}
     setAttachments={setAttachments}
-    messages={messages}
-    setMessages={setMessages}
-    sendMessage={sendMessage}
-    selectedVisibilityType={visibilityType}
+    messages={chat.messages}
+    setMessages={chat.setMessages}
+    sendMessage={chat.sendMessage}
+    selectedVisibilityType={chat.visibilityType}
   />
 
-  if (messages.length === 0) {
+  if (chat.messages.length === 0) {
     return (
       <div className="flex h-dvh min-w-0 flex-col bg-background">
         <ChatHeader empty />
@@ -308,30 +98,24 @@ export function Chat({
   }
 
   return (
-    <>
-      <div className="overscroll-behavior-contain flex h-dvh min-w-0 touch-pan-y flex-col bg-background">
-        <ChatHeader title={displayTitle} isLoadingTitle={titlePending && !displayTitle} />
-
-        <Messages
-          status={status}
-          messages={messages}
-          setMessages={setMessages}
-          addToolApprovalResponse={addToolApprovalResponse}
-          regenerate={regenerate}
-          sendMessage={sendMessage}
-          isReadonly={isReadonly}
-          selectedModelId={initialChatModel}
-          feedback={feedback}
-        />
-
-
-
-        <div className="sticky bottom-0 z-1 mx-auto flex w-full max-w-4xl gap-2 border-t-0 bg-background px-2 pb-3 md:px-4 md:pb-4">
-          {!isReadonly && (
-            inputElement
-          )}
-        </div>
+    <div className="overscroll-behavior-contain flex h-dvh min-w-0 touch-pan-y flex-col bg-background">
+      <ChatHeader title={chat.title} isLoadingTitle={chat.isTitleLoading} />
+
+      <Messages
+        status={chat.status}
+        messages={chat.messages}
+        setMessages={chat.setMessages}
+        addToolApprovalResponse={chat.addToolApprovalResponse}
+        regenerate={chat.regenerate}
+        sendMessage={chat.sendMessage}
+        isReadonly={chat.isReadonly}
+        selectedModelId={chat.model}
+        feedback={chat.feedback}
+      />
+
+      <div className="sticky bottom-0 z-1 mx-auto flex w-full max-w-4xl gap-2 border-t-0 bg-background px-2 pb-3 md:px-4 md:pb-4">
+        {!chat.isReadonly && inputElement}
       </div>
-    </>
+    </div>
   );
 }