diff --git a/apps/dev-playground/client/src/routes/agent.route.tsx b/apps/dev-playground/client/src/routes/agent.route.tsx index 6762a1a38..09d5fce57 100644 --- a/apps/dev-playground/client/src/routes/agent.route.tsx +++ b/apps/dev-playground/client/src/routes/agent.route.tsx @@ -1,47 +1,14 @@ import { getPluginClientConfig } from "@databricks/appkit-ui/js"; import { Button } from "@databricks/appkit-ui/react"; +import { ChatInput, Conversation } from "@databricks/appkit-ui/react/chat"; import { createFileRoute } from "@tanstack/react-router"; -import { useCallback, useEffect, useRef, useState } from "react"; +import type { UIMessage, UIMessageChunk } from "ai"; +import { useCallback, useRef, useState } from "react"; export const Route = createFileRoute("/agent")({ component: AgentRoute, }); -interface SSEEvent { - type: string; - delta?: string; - item_id?: string; - item?: { - type?: string; - id?: string; - call_id?: string; - name?: string; - arguments?: string; - output?: string; - status?: string; - }; - content?: string; - data?: Record; - error?: string; - sequence_number?: number; - output_index?: number; - approval_id?: string; - stream_id?: string; - tool_name?: string; - args?: unknown; - annotations?: { - readOnly?: boolean; - destructive?: boolean; - idempotent?: boolean; - }; -} - -interface ChatMessage { - id: number; - role: "user" | "assistant"; - content: string; -} - interface PendingApproval { approvalId: string; streamId: string; @@ -136,43 +103,26 @@ function useAutocomplete(enabled: boolean) { }; } +// Joins all text parts; a message can have multiple if the agent +// reopens text after a tool call. +function messageBodyText(message: UIMessage): string { + let body = ""; + for (const part of message.parts) { + if (part.type === "text") body += part.text; + } + return body; +} + function AgentRoute() { - const [messages, setMessages] = useState([]); - const [events, setEvents] = useState([]); - const [input, setInput] = useState(""); - const [isLoading, setIsLoading] = useState(false); - const [threadId, setThreadId] = useState(null); const [pendingApprovals, setPendingApprovals] = useState( [], ); - - const decideApproval = useCallback( - async (approvalId: string, decision: "approve" | "deny") => { - const approval = pendingApprovals.find( - (a) => a.approvalId === approvalId, - ); - if (!approval) return; - try { - await fetch("/api/agents/approve", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - streamId: approval.streamId, - approvalId, - decision, - }), - }); - } finally { - setPendingApprovals((prev) => - prev.filter((a) => a.approvalId !== approvalId), - ); - } - }, - [pendingApprovals], - ); - const messagesEndRef = useRef(null); + // Raw chunk log for the debug panel; ids keep React keys stable. + const [streamLog, setStreamLog] = useState< + Array<{ id: number; chunk: UIMessageChunk }> + >([]); + const nextChunkIdRef = useRef(0); const inputRef = useRef(null); - const msgIdCounter = useRef(0); const agentConfig = getPluginClientConfig<{ agents?: string[]; @@ -187,381 +137,394 @@ function AgentRoute() { clear: clearSuggestion, } = useAutocomplete(hasAutocomplete); - // biome-ignore lint/correctness/useExhaustiveDependencies: scroll on new messages - useEffect(() => { - messagesEndRef.current?.scrollIntoView({ behavior: "smooth" }); - }, [messages]); - - const sendMessage = useCallback(async () => { - if (!input.trim() || isLoading) return; - - clearSuggestion(); - const userMessage = input.trim(); - setInput(""); - setMessages((prev) => [ - ...prev, - { id: ++msgIdCounter.current, role: "user", content: userMessage }, - ]); - setEvents([]); - setIsLoading(true); - - try { - const response = await fetch("/api/agents/chat", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - message: userMessage, - ...(threadId && { threadId }), - }), - }); - - if (!response.ok) { - const error = await response.json(); - setMessages((prev) => [ - ...prev, - { - id: ++msgIdCounter.current, - role: "assistant", - content: `Error: ${error.error}`, - }, - ]); - return; - } - - const reader = response.body?.getReader(); - if (!reader) return; - - const decoder = new TextDecoder(); - let assistantContent = ""; - let buffer = ""; - - while (true) { - const { done, value } = await reader.read(); - if (done) break; - - buffer += decoder.decode(value, { stream: true }); - const lines = buffer.split("\n"); - buffer = lines.pop() ?? ""; - - for (const line of lines) { - if (!line.startsWith("data: ")) continue; - const data = line.slice(6).trim(); - if (!data || data === "[DONE]") continue; - - try { - const event: SSEEvent = JSON.parse(data); - if (!event.type) continue; - setEvents((prev) => [...prev, event]); - - if ( - event.type === "appkit.approval_pending" && - event.approval_id && - event.stream_id && - event.tool_name - ) { - setPendingApprovals((prev) => [ - ...prev, - { - approvalId: event.approval_id as string, - streamId: event.stream_id as string, - toolName: event.tool_name as string, - args: event.args, - }, - ]); - } - if (event.type === "appkit.metadata" && event.data?.threadId) { - setThreadId(event.data.threadId as string); - } - - if (event.type === "response.output_text.delta" && event.delta) { - assistantContent += event.delta; - setMessages((prev) => { - const updated = [...prev]; - const last = updated[updated.length - 1]; - if (last?.role === "assistant") { - updated[updated.length - 1] = { - ...last, - content: assistantContent, - }; - } else { - updated.push({ - id: ++msgIdCounter.current, - role: "assistant", - content: assistantContent, - }); - } - return updated; - }); - } - } catch { - // skip malformed events - } - } + const decideApproval = useCallback( + async ( + approvalId: string, + streamId: string, + decision: "approve" | "deny", + ) => { + try { + await fetch("/api/agents/approve", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ streamId, approvalId, decision }), + }); + } finally { + setPendingApprovals((prev) => + prev.filter((a) => a.approvalId !== approvalId), + ); } - } catch (err) { - setMessages((prev) => [ - ...prev, - { - id: ++msgIdCounter.current, - role: "assistant", - content: `Error: ${err instanceof Error ? err.message : "Unknown error"}`, - }, - ]); - } finally { - setIsLoading(false); - } - }, [input, isLoading, threadId, clearSuggestion]); - - const handleInputChange = (value: string) => { - setInput(value); - requestSuggestion(value); - }; - - const acceptSuggestion = () => { - if (!suggestion) return; - const newValue = input + suggestion; - setInput(newValue); - clearSuggestion(); - inputRef.current?.focus(); - }; + }, + [], + ); return (
-
-
-

Agent Chat

-

- AI agent with auto-discovered tools from all AppKit plugins. - {threadId && ( - - Thread: {threadId.slice(0, 8)}... - - )} -

-
- {hasAutocomplete && ( - - Autocomplete enabled - - )} -
- -
-
-
- {messages.length === 0 && ( -
-

- Send a message to start a conversation -

-

- The agent can use analytics, files, genie, and lakebase - tools. - {hasAutocomplete && " Start typing for inline suggestions."} -

-
- )} - - {messages.map((msg) => ( -
-
-

{msg.content}

+ { + if (part.type === "data-approval-pending") { + const payload = part.data as PendingApproval; + setPendingApprovals((prev) => + prev.some((p) => p.approvalId === payload.approvalId) + ? prev + : [...prev, payload], + ); + } + }} + onStreamPart={(chunk) => + setStreamLog((prev) => [ + ...prev, + { id: nextChunkIdRef.current++, chunk }, + ]) + } + > + {({ + id: chatId, + messages, + status, + error, + sendMessage, + stop, + containerRef, + }) => { + const isLoading = status === "submitted" || status === "streaming"; + // Show "Thinking..." until the assistant has produced visible + // text — the assistant message stub appears immediately on `start`. + const lastMessage = messages[messages.length - 1]; + const lastAssistantHasText = + lastMessage?.role === "assistant" && + messageBodyText(lastMessage).length > 0; + const showThinking = + isLoading && + pendingApprovals.length === 0 && + !lastAssistantHasText; + + return ( + <> +
+
+

Agent Chat

+

+ AI agent with auto-discovered tools from all plugins. + + Chat: {chatId.slice(0, 8)}... + +

+ {hasAutocomplete && ( + + Autocomplete enabled + + )}
- ))} - - {pendingApprovals.map((approval) => ( -
-
-
- - Destructive tool — approval required - -
-
- {approval.toolName} -
-                        {JSON.stringify(approval.args, null, 2)}
-                      
+ +
+
+
+ {messages.length === 0 && ( +
+

+ Send a message to start a conversation +

+

+ The agent can use analytics, files, genie, and + lakebase tools. + {hasAutocomplete && + " Start typing for inline suggestions."} +

+
+ )} + + {messages.map((msg) => { + const body = messageBodyText(msg); + if (!body) return null; + return ( +
+
+

+ {body} +

+
+
+ ); + })} + + {pendingApprovals.map((approval) => ( +
+
+
+ + Destructive tool — approval required + +
+
+ {approval.toolName} +
+                                {JSON.stringify(approval.args, null, 2)}
+                              
+
+
+ + +
+
+
+ ))} + + {showThinking && ( +
+
+

+ Thinking... +

+
+
+ )} + + {error && ( +
+
+

Error: {error.message}

+
+
+ )}
-
- - + {({ + value, + onChange, + submit, + isStreaming, + canSubmit, + handleKeyDown, + }) => ( +
+
+
+ {value} + + {suggestion} + +
+