|
| 1 | +import type { UIMessage } from "@ai-sdk/react"; |
| 2 | +import { memo } from "react"; |
| 3 | +import { |
| 4 | + AssistantResponse, |
| 5 | + ChatBubble, |
| 6 | + ToolUseRow, |
| 7 | +} from "~/components/runs/v3/ai/AIChatMessages"; |
| 8 | +import { Popover, PopoverContent, PopoverTrigger } from "~/components/primitives/Popover"; |
| 9 | + |
| 10 | +// --------------------------------------------------------------------------- |
| 11 | +// AgentMessageView — renders an AI SDK UIMessage[] conversation. |
| 12 | +// |
| 13 | +// Extracted from the playground route so it can be reused on the run details |
| 14 | +// page when the user picks the Agent view. |
| 15 | +// |
| 16 | +// UIMessage part types (AI SDK): |
| 17 | +// text — markdown text content |
| 18 | +// reasoning — model reasoning/thinking |
| 19 | +// tool-{name} — tool call with input/output/state |
| 20 | +// source-url — citation link |
| 21 | +// source-document — citation document reference |
| 22 | +// file — file attachment (image, etc.) |
| 23 | +// step-start — visual separator between steps |
| 24 | +// data-{name} — custom data parts (rendered as a small popover) |
| 25 | +// --------------------------------------------------------------------------- |
| 26 | + |
| 27 | +export function AgentMessageView({ messages }: { messages: UIMessage[] }) { |
| 28 | + return ( |
| 29 | + <div className="mx-auto flex max-w-[800px] flex-col gap-2"> |
| 30 | + {messages.map((msg) => ( |
| 31 | + <MessageBubble key={msg.id} message={msg} /> |
| 32 | + ))} |
| 33 | + </div> |
| 34 | + ); |
| 35 | +} |
| 36 | + |
| 37 | +// Memoized so stable messages (anything older than the one currently |
| 38 | +// streaming) don't re-render on every chunk. This matters a lot during |
| 39 | +// `resumeStream()` history replay, where each re-render would otherwise |
| 40 | +// re-run Prism highlighting on every tool-call CodeBlock in the list. |
| 41 | +// |
| 42 | +// Default shallow prop comparison is fine: AI SDK's useChat keeps stable |
| 43 | +// references for messages that haven't changed, so only the last message |
| 44 | +// (the one receiving new chunks) re-renders. |
| 45 | +export const MessageBubble = memo(function MessageBubble({ |
| 46 | + message, |
| 47 | +}: { |
| 48 | + message: UIMessage; |
| 49 | +}) { |
| 50 | + if (message.role === "user") { |
| 51 | + const text = |
| 52 | + message.parts |
| 53 | + ?.filter((p) => p.type === "text") |
| 54 | + .map((p) => (p as { type: "text"; text: string }).text) |
| 55 | + .join("") ?? ""; |
| 56 | + |
| 57 | + return ( |
| 58 | + <div className="flex justify-end"> |
| 59 | + <div className="max-w-[80%] rounded-lg bg-indigo-600 px-4 py-2.5 text-sm text-white"> |
| 60 | + <div className="whitespace-pre-wrap">{text}</div> |
| 61 | + </div> |
| 62 | + </div> |
| 63 | + ); |
| 64 | + } |
| 65 | + |
| 66 | + if (message.role === "assistant") { |
| 67 | + const hasContent = message.parts && message.parts.length > 0; |
| 68 | + if (!hasContent) return null; |
| 69 | + |
| 70 | + return ( |
| 71 | + <div className="space-y-2"> |
| 72 | + {message.parts?.map((part, i) => renderPart(part, i))} |
| 73 | + </div> |
| 74 | + ); |
| 75 | + } |
| 76 | + |
| 77 | + return null; |
| 78 | +}); |
| 79 | + |
| 80 | +export function renderPart(part: UIMessage["parts"][number], i: number) { |
| 81 | + const p = part as any; |
| 82 | + const type = part.type as string; |
| 83 | + |
| 84 | + // Text — markdown rendered via AssistantResponse |
| 85 | + if (type === "text") { |
| 86 | + return p.text ? <AssistantResponse key={i} text={p.text} headerLabel="" /> : null; |
| 87 | + } |
| 88 | + |
| 89 | + // Reasoning — amber-bordered italic block |
| 90 | + if (type === "reasoning") { |
| 91 | + return ( |
| 92 | + <div key={i} className="border-l-2 border-amber-500/40 pl-2"> |
| 93 | + <ChatBubble> |
| 94 | + <div className="whitespace-pre-wrap text-xs italic text-amber-200/70"> |
| 95 | + {p.text ?? ""} |
| 96 | + </div> |
| 97 | + </ChatBubble> |
| 98 | + </div> |
| 99 | + ); |
| 100 | + } |
| 101 | + |
| 102 | + // Tool call — type: "tool-{name}" with toolCallId, input, output, state |
| 103 | + if (type.startsWith("tool-")) { |
| 104 | + const toolName = type.slice(5); |
| 105 | + |
| 106 | + // Sub-agent tool: output is a UIMessage with parts |
| 107 | + const isSubAgent = |
| 108 | + p.output != null && typeof p.output === "object" && Array.isArray(p.output.parts); |
| 109 | + |
| 110 | + // For sub-agent tools, show the last text part as the "output" tab |
| 111 | + // (mirrors what toModelOutput typically sends to the parent LLM) |
| 112 | + // instead of dumping the full UIMessage JSON. |
| 113 | + let resultOutput: string | undefined; |
| 114 | + if (isSubAgent) { |
| 115 | + const lastText = (p.output.parts as any[]) |
| 116 | + .filter((part: any) => part.type === "text" && part.text) |
| 117 | + .pop(); |
| 118 | + resultOutput = lastText?.text ?? undefined; |
| 119 | + } else if (p.output != null) { |
| 120 | + resultOutput = |
| 121 | + typeof p.output === "string" ? p.output : JSON.stringify(p.output, null, 2); |
| 122 | + } |
| 123 | + |
| 124 | + return ( |
| 125 | + <ToolUseRow |
| 126 | + key={i} |
| 127 | + tool={{ |
| 128 | + toolCallId: p.toolCallId ?? `tool-${i}`, |
| 129 | + toolName, |
| 130 | + inputJson: JSON.stringify(p.input ?? {}, null, 2), |
| 131 | + resultOutput, |
| 132 | + resultSummary: |
| 133 | + p.state === "input-streaming" || p.state === "input-available" |
| 134 | + ? "calling..." |
| 135 | + : p.state === "output-error" |
| 136 | + ? `error: ${p.errorText ?? "unknown"}` |
| 137 | + : undefined, |
| 138 | + subAgent: isSubAgent |
| 139 | + ? { |
| 140 | + parts: p.output.parts, |
| 141 | + isStreaming: p.state === "output-available" && p.preliminary === true, |
| 142 | + } |
| 143 | + : undefined, |
| 144 | + }} |
| 145 | + /> |
| 146 | + ); |
| 147 | + } |
| 148 | + |
| 149 | + // Source URL — clickable citation link |
| 150 | + if (type === "source-url") { |
| 151 | + return ( |
| 152 | + <div key={i} className="text-xs"> |
| 153 | + <a |
| 154 | + href={p.url} |
| 155 | + target="_blank" |
| 156 | + rel="noopener noreferrer" |
| 157 | + className="text-indigo-400 underline hover:text-indigo-300" |
| 158 | + > |
| 159 | + {p.title || p.url} |
| 160 | + </a> |
| 161 | + </div> |
| 162 | + ); |
| 163 | + } |
| 164 | + |
| 165 | + // Source document — citation label |
| 166 | + if (type === "source-document") { |
| 167 | + return ( |
| 168 | + <div key={i} className="text-xs text-text-dimmed"> |
| 169 | + {p.title} |
| 170 | + {p.mediaType ? ` (${p.mediaType})` : ""} |
| 171 | + </div> |
| 172 | + ); |
| 173 | + } |
| 174 | + |
| 175 | + // File — render as image if image type, otherwise as download link |
| 176 | + if (type === "file") { |
| 177 | + const isImage = typeof p.mediaType === "string" && p.mediaType.startsWith("image/"); |
| 178 | + if (isImage) { |
| 179 | + return ( |
| 180 | + <img |
| 181 | + key={i} |
| 182 | + src={p.url} |
| 183 | + alt={p.filename ?? "file"} |
| 184 | + className="max-h-64 rounded border border-charcoal-650" |
| 185 | + /> |
| 186 | + ); |
| 187 | + } |
| 188 | + return ( |
| 189 | + <div key={i} className="text-xs"> |
| 190 | + <a |
| 191 | + href={p.url} |
| 192 | + target="_blank" |
| 193 | + rel="noopener noreferrer" |
| 194 | + className="text-indigo-400 underline hover:text-indigo-300" |
| 195 | + > |
| 196 | + {p.filename ?? "Download file"} |
| 197 | + </a> |
| 198 | + </div> |
| 199 | + ); |
| 200 | + } |
| 201 | + |
| 202 | + // Step start — subtle dashed separator with centered label |
| 203 | + if (type === "step-start") { |
| 204 | + return ( |
| 205 | + <div key={i} className="flex items-center gap-2 py-0.5"> |
| 206 | + <div className="flex-1 border-t border-dashed border-charcoal-650" /> |
| 207 | + <span className="text-[10px] text-charcoal-500">step</span> |
| 208 | + <div className="flex-1 border-t border-dashed border-charcoal-650" /> |
| 209 | + </div> |
| 210 | + ); |
| 211 | + } |
| 212 | + |
| 213 | + // Data parts — type: "data-{name}", show as labeled JSON popover |
| 214 | + if (type.startsWith("data-")) { |
| 215 | + const dataName = type.slice(5); |
| 216 | + return <DataPartPopover key={i} name={dataName} data={p.data} />; |
| 217 | + } |
| 218 | + |
| 219 | + return null; |
| 220 | +} |
| 221 | + |
| 222 | +function DataPartPopover({ name, data }: { name: string; data: unknown }) { |
| 223 | + const formatted = JSON.stringify(data, null, 2); |
| 224 | + |
| 225 | + return ( |
| 226 | + <Popover> |
| 227 | + <PopoverTrigger asChild> |
| 228 | + <button |
| 229 | + type="button" |
| 230 | + className="inline-flex items-center gap-1 rounded border border-charcoal-650 bg-charcoal-800 px-1.5 py-0.5 font-mono text-[10px] text-text-dimmed transition-colors hover:border-charcoal-500 hover:text-text-bright" |
| 231 | + > |
| 232 | + <span className="text-purple-400">{name}</span> |
| 233 | + <span className="text-charcoal-500">{"{}"}</span> |
| 234 | + </button> |
| 235 | + </PopoverTrigger> |
| 236 | + <PopoverContent className="w-auto max-w-md p-0" align="start" sideOffset={4}> |
| 237 | + <div className="flex items-center justify-between border-b border-charcoal-650 px-2.5 py-1.5"> |
| 238 | + <span className="text-[10px] font-medium text-text-dimmed">data-{name}</span> |
| 239 | + </div> |
| 240 | + <div className="max-h-60 overflow-auto scrollbar-thin scrollbar-track-transparent scrollbar-thumb-charcoal-600"> |
| 241 | + <pre className="p-2.5 text-[11px] leading-relaxed text-text-bright">{formatted}</pre> |
| 242 | + </div> |
| 243 | + </PopoverContent> |
| 244 | + </Popover> |
| 245 | + ); |
| 246 | +} |
0 commit comments