feat: enhance chat components and token usage tracking

- Refactor AgentReasoning component to remove autoClose prop.
- Update ChatHeader to include detailed output and input token information.
- Add SelectedAttachments component in ChatInput for better file management.
- Improve ChatMessages by extracting source documents from message parts.
- Extend TokenUsage interface to include inputTokenDetails for better tracking.
- Modify ChatProvider to handle detailed token usage data including cache details.

Author: ssdeanx

app/chat/components/agent-reasoning.tsx

@@ -22,7 +22,6 @@ export function AgentReasoning({
         <Reasoning
             isStreaming={isStreaming}
             duration={duration}
-            autoClose={false}
             className={className}
         >
             <ReasoningTrigger />

app/chat/components/chat-header.tsx

@@ -199,6 +199,18 @@ export function ChatHeader() {
                             inputTokens: usage.inputTokens,
                             outputTokens: usage.outputTokens,
                             totalTokens: usage.totalTokens,
+                            outputTokenDetails:
+                                typeof usage.outputTokens === 'number'
+                                    ? { textTokens: usage.outputTokens, reasoningTokens: 0 }
+                                    : {
+                                        textTokens: (usage.outputTokens as unknown as { textTokens?: number }).textTokens ?? 0,
+                                        reasoningTokens: (usage.outputTokens as unknown as { reasoningTokens?: number }).reasoningTokens ?? 0,
+                                    },
+                            inputTokenDetails: {
+                                cacheReadTokens: (usage.inputTokenDetails as unknown as { cacheReadTokens?: number }).cacheReadTokens ?? 0,
+                                cacheWriteTokens: (usage.inputTokenDetails as unknown as { cacheWriteTokens?: number }).cacheWriteTokens ?? 0,
+                                noCacheTokens: (usage.inputTokenDetails as unknown as { noCacheTokens?: number }).noCacheTokens ?? 0,
+                            }
                         }}
                     >
                         <ContextTrigger />

app/chat/components/chat-input.tsx

@@ -7,15 +7,13 @@ import {
     PromptInputTools,
     PromptInputButton,
     PromptInputSubmit,
-    PromptInputAttachments,
-    PromptInputAttachment,
     PromptInputHeader,
     PromptInputBody,
     PromptInputActionMenu,
     PromptInputActionMenuTrigger,
     PromptInputActionMenuContent,
     PromptInputActionAddAttachments,
-    PromptInputSpeechButton,
+    usePromptInputAttachments,
 } from '@/src/components/ai-elements/prompt-input'
 import {
     ModelSelector,
@@ -48,6 +46,7 @@ import {
     MicIcon,
     SparklesIcon,
     ListTodoIcon,
+    XIcon,
 } from 'lucide-react'
 import { useMemo, useState, useRef } from 'react'
 import { MODEL_CONFIGS } from '../config/models'
@@ -58,6 +57,35 @@ import {
 } from '../config/agents'
 import { cn } from '@/lib/utils'
 
+function SelectedAttachments() {
+    const { attachments, removeAttachment } = usePromptInputAttachments()
+
+    if (attachments.length === 0) {return null}
+
+    return (
+        <div className="flex flex-wrap gap-2 p-2 border-b border-border/50">
+            {attachments.map((file, index) => (
+                <Badge
+                    key={`${file.name}-${index}`}
+                    variant="secondary"
+                    className="h-6 gap-1 pr-1"
+                >
+                    <span className="max-w-30 truncate text-[10px]">
+                        {file.name}
+                    </span>
+                    <button
+                        type="button"
+                        onClick={() => removeAttachment(index)}
+                        className="rounded-full hover:bg-foreground/10 transition-colors"
+                    >
+                        <XIcon className="size-3" />
+                    </button>
+                </Badge>
+            ))}
+        </div>
+    )
+}
+
 export function ChatInput() {
     const {
         sendMessage,
@@ -90,16 +118,9 @@ export function ChatInput() {
 
     const agentsByCategory = useMemo(() => getAgentsByCategory(), [])
 
-    /* Agent Selector - compact dropdown in input toolbar */
-
-    // Model Selector
-
-    const handleSubmit = async (message: {
-        text: string
-        files: unknown[]
-    }) => {
+    const handleSubmit = (message: { text: string; files: File[] }) => {
         if (message.text.trim()) {
-            sendMessage(message.text, message.files as File[])
+            sendMessage(message.text, message.files)
             setInput('')
         }
     }
@@ -125,7 +146,6 @@ export function ChatInput() {
                         <span className="flex items-center gap-1.5">
                             <BotIcon className="size-3" />
                             {agentConfig?.name ?? selectedAgent}
-
                         </span>
                         <span className="flex items-center gap-1.5">
                             <CpuIcon className="size-3" />
@@ -154,16 +174,7 @@ export function ChatInput() {
                     globalDrop
                 >
                     <PromptInputHeader>
-                        {supportsFiles && (
-                            <PromptInputAttachments>
-                                {(file) => (
-                                    <PromptInputAttachment
-                                        key={file.id}
-                                        data={file}
-                                    />
-                                )}
-                            </PromptInputAttachments>
-                        )}
+                        {supportsFiles && <SelectedAttachments />}
                     </PromptInputHeader>
 
                     <PromptInputBody>
@@ -187,18 +198,17 @@ export function ChatInput() {
                                 </PromptInputActionMenuContent>
                             </PromptInputActionMenu>
 
-                            <PromptInputSpeechButton
-                                onTranscriptionChange={setInput}
-                                textareaRef={textareaRef}
+                            <PromptInputButton
                                 className={cn(
                                     'magnetic transition-all duration-300',
                                     isSpeaking &&
                                         'text-primary glow-primary animate-ambient-pulse scale-110'
                                 )}
                                 onClick={() => setIsSpeaking(!isSpeaking)}
+                                title="Speech to text"
                             >
                                 <MicIcon className="size-4" />
-                            </PromptInputSpeechButton>
+                            </PromptInputButton>
 
                             <PromptInputButton
                                 onClick={() => {

app/chat/components/chat-messages.tsx

@@ -15,8 +15,6 @@ import {
     MessageToolbar,
     MessageActions,
     MessageAction,
-    MessageAttachment,
-    MessageAttachments,
 } from '@/src/components/ai-elements/message'
 import { Loader } from '@/src/components/ai-elements/loader'
 import {
@@ -54,13 +52,57 @@ import {
     ActivityIcon,
     NetworkIcon,
 } from 'lucide-react'
-import { useState, useCallback, useMemo, Fragment } from 'react'
-import type { UIMessage, FileUIPart } from 'ai'
+import { useState, useCallback, useMemo, Fragment, memo  } from 'react'
+import type {
+    UIMessage,
+    DynamicToolUIPart,
+    TextUIPart,
+    ReasoningUIPart,
+    ToolUIPart,
+    TextStreamPart,
+    TextPart,
+    ToolResultPart,
+    ReasoningOutput,
+    UIDataPartSchemas,
+    UIMessageChunk,
+    UIMessagePart,
+    DataContent,
+    FinishReason,
+    FileUIPart,
+    Tool,
+    DataUIPart,
+    SourceDocumentUIPart,
+    SourceUrlUIPart,
+    StepResult,
+    PrepareStepResult,
+    StepStartUIPart,
+    InferSchema,
+    InferUIDataParts,
+    InferAgentUIMessage,
+    InferToolInput,
+    InferUIMessageChunk,
+    InferToolOutput,
+    InferUITool,
+    InferUITools,
+    InferGenerateOutput,
+    InferStreamOutput,
+} from 'ai'
 import {
-    isTextUIPart,
-    isReasoningUIPart,
-    isToolOrDynamicToolUIPart,
+    safeValidateUIMessages,
+    getToolName,
+    getStaticToolName,
+    getTextFromDataUrl,
+    isDataUIPart,
     isFileUIPart,
+    isReasoningUIPart,
+    isTextUIPart,
+    isToolUIPart,
+    isStaticToolUIPart,
+    isDeepEqualData,
+    InvalidResponseDataError,
+    InvalidMessageRoleError,
+    InvalidArgumentError,
+    generateId
 } from 'ai'
 import type { BundledLanguage } from 'shiki'
 import { Button } from '@/ui/button'
@@ -84,6 +126,58 @@ type MastraDataPart =
     | NetworkDataPart
     | { type: `data-${string}`; id?: string; data: unknown }
 
+interface ChatMessagesProps {
+    messages: UIMessage[]
+    status: string
+    error: Error | undefined
+    onSuggestionClick: (suggestion: string) => void
+    onCopyMessage?: (messageId: string, content: string) => void
+    onRegenerate?: (messageId: string) => void
+}
+
+interface SourceDocument {
+    title?: string
+    url?: string
+    description?: string
+    sourceDocument?: string
+}
+
+
+function isSourceUrlPart(part: UIMessage['parts'][number]): part is SourceUrlUIPart {
+    return part.type === 'source-url'
+}
+
+function isSourceDocumentPart(
+    part: UIMessage['parts'][number]
+): part is SourceDocumentUIPart {
+    return part.type === 'source-document'
+}
+
+
+    // Extract sources from message parts
+const getSourcesFromParts = (parts: UIMessage['parts']): SourceDocument[] => {
+    const sources: SourceDocument[] = []
+    for (const part of parts) {
+        if (isSourceUrlPart(part)) {
+            sources.push({
+                title: part.title,
+                url: part.url,
+                description: part.url,
+            })
+            continue
+        }
+
+        if (isSourceDocumentPart(part)) {
+            sources.push({
+                title: part.title,
+                description: part.filename ?? part.mediaType,
+            })
+        }
+    }
+    return sources
+}
+
+
 /**
  * Type guard to check for type property
  */
@@ -207,8 +301,8 @@ function WorkflowDataSection({ part }: { part: WorkflowDataPart }) {
                             workflowData.status === 'success' || (workflowData.status as string) === 'completed'
                                 ? 'default'
                                 : workflowData.status === 'failed'
-                                  ? 'destructive'
-                                  : 'secondary'
+                                    ? 'destructive'
+                                    : 'secondary'
                         }
                         className="text-xs"
                     >
@@ -243,8 +337,8 @@ function WorkflowDataSection({ part }: { part: WorkflowDataPart }) {
                                                     ? 'default'
                                                     : stepData.status ===
                                                         'failed'
-                                                      ? 'destructive'
-                                                      : 'secondary'
+                                                        ? 'destructive'
+                                                        : 'secondary'
                                             }
                                             className="text-xs"
                                         >
@@ -364,8 +458,8 @@ function NetworkDataSection({ part }: { part: NetworkDataPart }) {
                                                 step.status === 'success' || (step.status as string) === 'completed'
                                                     ? 'default'
                                                     : step.status === 'failed'
-                                                      ? 'destructive'
-                                                      : 'secondary'
+                                                        ? 'destructive'
+                                                        : 'secondary'
                                             }
                                             className="text-xs"
                                         >
@@ -726,7 +820,7 @@ function MessageItem({
         const tools: ToolInvocationState[] = []
 
         for (const p of parts) {
-            if (isToolOrDynamicToolUIPart(p)) {
+            if (isToolUIPart(p)) {
                 tools.push(p as ToolInvocationState)
             }
         }
@@ -778,8 +872,8 @@ function MessageItem({
                     statusText === 'in-progress'
                         ? 'in-progress'
                         : statusText.trim().length > 0
-                          ? 'done'
-                          : ''
+                            ? 'done'
+                            : ''
                 return {
                     message: messageText,
                     status: normalizedStatus,
@@ -904,14 +998,20 @@ function MessageItem({
                     <MessageContent>
                         {/* User file attachments */}
                         {isUser && fileParts && fileParts.length > 0 && (
-                            <MessageAttachments>
+                            <div className="my-2 flex flex-wrap gap-2">
                                 {fileParts.map((file, idx) => (
-                                    <MessageAttachment
+                                    <div
                                         key={`file-${idx}`}
-                                        data={file}
-                                    />
+                                        className="flex items-center gap-2 rounded-lg border bg-muted/20 px-3 py-1.5 text-xs font-medium"
+                                    >
+                                        <span className="truncate max-w-50">
+                                            {file.filename ??
+                                                file.mediaType ??
+                                                `File ${idx + 1}`}
+                                        </span>
+                                    </div>
                                 ))}
-                            </MessageAttachments>
+                            </div>
                         )}
 
                         {/* Non-image files with inline preview controls */}

app/chat/providers/chat-context-types.ts

@@ -29,6 +29,7 @@ export interface Source {
 }
 
 export interface TokenUsage {
+    inputTokenDetails: any
     inputTokens: number
     outputTokens: number
     totalTokens: number