"use client" import { type ReactNode, useEffect, useMemo, useRef, useState } from "react" import type { UIMessage } from "@ai-sdk/react" import { useQuery } from "@tanstack/react-query" import { Streamdown } from "streamdown" import { BookOpenIcon, ChevronDownIcon, ChevronRightIcon, ClockIcon, GlobeIcon, ListIcon, Loader2, PlusIcon, SearchIcon, TerminalIcon, WrenchIcon, XCircleIcon, } from "lucide-react" import { cn } from "@lib/utils" import { isWebSearchToolName } from "@/lib/chat-web-search-tools" import { buildCitationIndex, fetchDocumentsByIds, getDocumentSourceUrl, isMemoryToolOutputReady, mapDocumentsByKnownIds, type CitationTarget, type DocumentWithMemories, extractMemoryToolOutputs, } from "@/lib/chat-memory-tools" import { parseSourceAnnotatedMarkdown, stripSourceMarkup, } from "@/lib/source-annotations" import { modelNames, type ModelId } from "@/lib/models" import { RelatedMemories } from "./related-memories" import { MessageActions } from "./message-actions" const TOOL_META: Record = { bash: { label: "Memory", icon: TerminalIcon }, recallContext: { label: "Recall Memories", icon: BookOpenIcon }, discoverSpaces: { label: "Discover Spaces", icon: SearchIcon }, web_search: { label: "Web search", icon: GlobeIcon }, google_search: { label: "Google search", icon: GlobeIcon }, // legacy tool names kept for existing persisted messages searchMemories: { label: "Search Memories", icon: SearchIcon }, addMemory: { label: "Add Memory", icon: PlusIcon }, fetchMemory: { label: "Fetch Memory", icon: BookOpenIcon }, forgetMemory: { label: "Forget Memory", icon: XCircleIcon }, updateMemory: { label: "Update Memory", icon: BookOpenIcon }, forgetDocument: { label: "Forget Document", icon: XCircleIcon }, scheduleTask: { label: "Schedule Task", icon: ClockIcon }, listSchedules: { label: "List Schedules", icon: ListIcon }, cancelSchedule: { label: "Cancel Schedule", icon: XCircleIcon }, } type ToolCallDisplayPart = { type: string state?: string input?: unknown output?: unknown toolCallId?: string errorText?: string } type SourceUrlPart = { type: "source-url" sourceId: string url: string title?: string } function sourceHost(url: string): string { try { return new URL(url).hostname.replace(/^www\./, "") } catch { return url } } function faviconUrl(host: string): string { return `https://www.google.com/s2/favicons?sz=64&domain=${host}` } function safeExternalUrl(url: string | null | undefined): string | null { if (!url) return null if (url.startsWith("/") && !url.startsWith("//")) return url if (url.startsWith("#") && !url.startsWith("#sm-source:")) return url try { const parsed = new URL(url) return parsed.protocol === "http:" || parsed.protocol === "https:" ? url : null } catch { return null } } function isWebSearchPart(part: { type: string; toolName?: string }): boolean { if (part.type === "dynamic-tool") { return isWebSearchToolName(part.toolName ?? "") } if (part.type.startsWith("tool-")) { return isWebSearchToolName(part.type.slice("tool-".length)) } return false } function isMemoryRetrievalToolName(toolName: string): boolean { return ( toolName === "searchMemories" || toolName === "recallContext" || toolName === "discoverSpaces" ) } export function isChatToolDisplayPartType(type: string): boolean { return ( type === "tool-searchMemories" || type === "tool-recallContext" || type === "tool-discoverSpaces" || type === "tool-forgetMemory" || type === "tool-updateMemory" || type === "tool-forgetDocument" ) } function CitationLink({ href, label, source, }: { href: string label: string source?: SourceUrlPart }) { const url = safeExternalUrl(source?.url ?? href) ?? "" if (!url) return <>{label} const host = sourceHost(url) const rawTitle = source?.title?.trim() const hasTitle = !!rawTitle && rawTitle !== host && !/^https?:\/\//i.test(rawTitle) let path = "" try { path = new URL(url).pathname.replace(/\/$/, "") } catch {} return ( {label} {host} {hasTitle ? ( {rawTitle} ) : path ? ( {path} ) : null} ) } function sourceTitle( target: CitationTarget, document?: DocumentWithMemories, ): string { return ( document?.title?.trim() || target.title?.trim() || document?.customId || target.customId || target.documentId || target.sourceId ) } function sourceSummary( target: CitationTarget, document?: DocumentWithMemories, ): string | null { const summary = document?.summary || target.summary || (document as { content?: string } | undefined)?.content || null return summary ? summary.trim() : null } function sourceKind( target: CitationTarget, document?: DocumentWithMemories, ): string { return (document?.type || target.type || "memory").replaceAll("_", " ") } function SourceCitationLink({ sourceId, children, citationIndex, documentByKnownId, }: { sourceId: string children: ReactNode citationIndex: Map documentByKnownId: Map }) { const target = citationIndex.get(sourceId) if (!target) return <>{children} const document = (target.documentId ? documentByKnownId.get(target.documentId) : undefined) ?? (target.customId ? documentByKnownId.get(target.customId) : undefined) const url = safeExternalUrl( document ? getDocumentSourceUrl(document) : target.url, ) const title = sourceTitle(target, document) const summary = sourceSummary(target, document) return ( {url ? ( {children} ) : ( )} {sourceId} {title} {sourceKind(target, document)} {summary ? ( {summary} ) : null} {url ? ( Open source ) : null} ) } function makeMarkdownComponents( sources: SourceUrlPart[], citationIndex: Map, documentByKnownId: Map, ) { return { a: ({ href, children }: { href?: string; children?: ReactNode }) => { if (href?.startsWith("#sm-source:")) { const sourceId = (() => { try { return decodeURIComponent(href.slice("#sm-source:".length)) } catch { return null } })() if (!sourceId) return <>{children} return ( {children} ) } const label = typeof children === "string" ? children : Array.isArray(children) ? children.join("") : "" const match = label.match(/^\[?(\d+)\]?$/) const safeHref = safeExternalUrl(href) if (match && safeHref) { const n = Number(match[1]) const source = sources.find((s) => s.url === safeHref) ?? sources[n - 1] return } if (!safeHref) return <>{children} return ( {children} ) }, } } function WebSourcesPill({ sources }: { sources: SourceUrlPart[] }) { const [expanded, setExpanded] = useState(false) const ref = useRef(null) useEffect(() => { if (!expanded) return const onDown = (e: MouseEvent) => { if (ref.current && !ref.current.contains(e.target as Node)) { setExpanded(false) } } document.addEventListener("mousedown", onDown) return () => document.removeEventListener("mousedown", onDown) }, [expanded]) if (sources.length === 0) return null const faviconHosts: string[] = [] for (const s of sources) { const host = sourceHost(s.url) if (!faviconHosts.includes(host)) faviconHosts.push(host) if (faviconHosts.length >= 3) break } const count = sources.length return (
{expanded && (
)}
) } function BashToolDisplay({ part }: { part: ToolCallDisplayPart }) { const [expanded, setExpanded] = useState(false) const isLoading = part.state === "input-streaming" || part.state === "input-available" const isDone = part.state === "output-available" const isError = part.state === "error" || part.state === "output-error" const cmd = part.input && typeof part.input === "object" && "cmd" in part.input ? String((part.input as { cmd: string }).cmd) : undefined const output = isDone && part.output && typeof part.output === "object" ? (part.output as { stdout?: string; stderr?: string; exitCode?: number }) : undefined const hasOutput = output && ((output.stdout && output.stdout.length > 0) || (output.stderr && output.stderr.length > 0)) const errorText = part.errorText const hasExpandable = hasOutput || (isError && errorText) return (
{expanded && (hasOutput || (isError && errorText)) && (
{output?.stdout && output.stdout.length > 0 && (
							{output.stdout}
						
)} {output?.stderr && output.stderr.length > 0 && (
							{output.stderr}
						
)} {isError && errorText && (
							{errorText}
						
)}
)}
) } function ToolCallDisplay({ part }: { part: ToolCallDisplayPart }) { const [expanded, setExpanded] = useState(false) const toolName = part.type.replace("tool-", "") if (toolName === "bash") { return } if (isWebSearchToolName(toolName)) { if (part.state === "output-available") return null if (part.state === "error" || part.state === "output-error") { return (
Web search failed
) } return (
Searching the web…
) } const meta = TOOL_META[toolName] ?? (isWebSearchToolName(toolName) ? { label: "Web search", icon: GlobeIcon } : undefined) const Icon = meta?.icon ?? WrenchIcon const label = meta?.label ?? toolName const isLoading = part.state === "input-streaming" || part.state === "input-available" const isDone = part.state === "output-available" const isError = part.state === "error" || part.state === "output-error" const errorText = part.errorText if (isMemoryRetrievalToolName(toolName) && isMemoryToolOutputReady(part)) { return null } return (
{expanded && (
{part.input !== undefined && (
Input
								{typeof part.input === "string"
									? part.input
									: JSON.stringify(part.input, null, 2)}
							
)} {isDone && part.output !== undefined && (
Output
								{typeof part.output === "string"
									? part.output
									: JSON.stringify(part.output, null, 2)}
							
)} {isError && errorText && (
Error
								{errorText}
							
)}
)}
) } interface AgentMessageProps { message: UIMessage index: number messagesLength: number hoveredMessageId: string | null copiedMessageId: string | null messageFeedback: Record expandedMemories: string | null responseModel: ModelId | null onCopy: (messageId: string, text: string) => void onLike: (messageId: string) => void onDislike: (messageId: string) => void onToggleMemories: (messageId: string) => void } export function AgentMessage({ message, index, messagesLength, hoveredMessageId, copiedMessageId, messageFeedback, expandedMemories, responseModel, onCopy, onLike, onDislike, onToggleMemories, }: AgentMessageProps) { const isLastAgentMessage = index === messagesLength - 1 && message.role === "assistant" const isHovered = hoveredMessageId === message.id const messageText = message.parts .filter((part) => part.type === "text") .map((part) => part.text) .join(" ") const copyText = stripSourceMarkup(messageText) const memoryOutputs = useMemo( () => extractMemoryToolOutputs(message), [message], ) const citationIndex = useMemo( () => buildCitationIndex(memoryOutputs), [memoryOutputs], ) const allowedSourceIds = useMemo( () => new Set(citationIndex.keys()), [citationIndex], ) const sourceDocumentIds = useMemo(() => { const ids = new Set() for (const target of citationIndex.values()) { if (target.documentId) ids.add(target.documentId) if (target.customId) ids.add(target.customId) } return [...ids].sort() }, [citationIndex]) const { data: sourceDocuments = [] } = useQuery({ queryKey: ["chat-source-documents", sourceDocumentIds], queryFn: () => fetchDocumentsByIds(sourceDocumentIds), enabled: sourceDocumentIds.length > 0, staleTime: 5 * 60 * 1000, }) const documentByKnownId = useMemo( () => mapDocumentsByKnownIds(sourceDocuments), [sourceDocuments], ) const webSources = useMemo(() => { const seen = new Set() const out: SourceUrlPart[] = [] for (const part of message.parts) { if (part.type !== "source-url") continue const source = part as SourceUrlPart if (seen.has(source.url)) continue seen.add(source.url) out.push(source) } return out }, [message.parts]) const hasAssistantText = message.parts.some( (p) => p.type === "text" && (p as { text?: string }).text?.trim(), ) const markdownComponents = useMemo( () => makeMarkdownComponents(webSources, citationIndex, documentByKnownId), [webSources, citationIndex, documentByKnownId], ) const responseModelLabel = responseModel ? `${modelNames[responseModel].name} ${modelNames[responseModel].version}` : null return (
{message.parts.map((part, partIndex) => { if (part.type === "source-url") { return null } if (isWebSearchPart(part)) { return null } if (part.type === "source-document") { const doc = part as { type: "source-document" sourceId: string title: string filename?: string } return (
Document
{doc.title}
{doc.filename && (
{doc.filename}
)}
) } if (part.type === "text") { // Skip fragments mid-run — source-url citations split one answer into // many text parts; rendering each separately tears markdown (lists etc.). let prev = partIndex - 1 while (prev >= 0 && message.parts[prev]?.type === "source-url") { prev-- } if (prev >= 0 && message.parts[prev]?.type === "text") { return null } let runText = "" for (let j = partIndex; j < message.parts.length; j++) { const p = message.parts[j] if (p?.type === "text") runText += p.text else if (p?.type === "source-url") continue else break } return (
{ parseSourceAnnotatedMarkdown(runText, allowedSourceIds) .markdown }
) } if (part.type === "dynamic-tool") { const dt = part as { type: "dynamic-tool" toolName: string toolCallId: string state?: string input?: unknown output?: unknown errorText?: string } const displayState = dt.state === "output-error" ? "error" : dt.state return ( ) } if (part.type.startsWith("tool-")) { return ( ) } return null })}
{hasAssistantText && (
{webSources.length > 0 && (
)} {responseModelLabel && ( {responseModelLabel} )}
)}
) }