From 4498a3f66682aa9bace9d1b9fcbd5b0bbaf195e6 Mon Sep 17 00:00:00 2001 From: Rodrigo Verdiani Date: Thu, 11 Jun 2026 21:46:53 -0300 Subject: [PATCH] feat(chat): enhance streaming status and tool call events --- src/features/binary/AIChatPanel.tsx | 52 +++++--- src/features/binary/hooks/useAIChat.ts | 173 +++++++++++++++++++------ src/index.css | 1 + 3 files changed, 171 insertions(+), 55 deletions(-) diff --git a/src/features/binary/AIChatPanel.tsx b/src/features/binary/AIChatPanel.tsx index 8d90dc2..22af22e 100644 --- a/src/features/binary/AIChatPanel.tsx +++ b/src/features/binary/AIChatPanel.tsx @@ -29,6 +29,8 @@ export function AIChatPanel() { isLoadingMessages, isStreaming, streamingContent, + toolStatus, + isToolRunning, streamError, isCreatingSession, createSession, @@ -215,8 +217,6 @@ export function AIChatPanel() { <> {messages.map((message, idx) => { const isUser = message.role === "USER"; - const isLastMessage = idx === messages.length - 1; - const isStreamingMessage = isLastMessage && isStreaming && !message.id; return (
{isUser ? (
{message.content}
- ) : isStreamingMessage ? ( -
- {streamingContent ? ( - <> - {streamingContent} - - - ) : ( -
- - Analyzing... -
- )} -
) : (
)} + {isStreaming && ( + <> +
+
+ +
+
+
+
+ {streamingContent ? ( + <> + {streamingContent} + + + ) : ( +
+ + Analyzing... +
+ )} +
+
+
+
+ {toolStatus && ( +
+ {isToolRunning && ( + + )} + {toolStatus} +
+ )} + + )}
)} diff --git a/src/features/binary/hooks/useAIChat.ts b/src/features/binary/hooks/useAIChat.ts index 418ce87..2194b2a 100644 --- a/src/features/binary/hooks/useAIChat.ts +++ b/src/features/binary/hooks/useAIChat.ts @@ -10,6 +10,46 @@ import { } from "@/api/generated/chat/chat.ts"; import type { SessionResponse, MessageResponse } from "@/api/generated/api.schemas.ts"; +interface ChunkEvent { + type: "chunk"; + content: string; +} + +interface ToolCallStartEvent { + type: "tool_call_start"; + tool: string; +} + +interface ToolCallEndEvent { + type: "tool_call_end"; + tool: string; + durationMs: number; + resultSummary: string; +} + +interface ThinkingEvent { + type: "thinking"; + message: string; +} + +interface DoneEvent { + type: "done"; + sessionId: string; +} + +interface ErrorEvent { + type: "error"; + message: string; +} + +type SSEEvent = + | ChunkEvent + | ToolCallStartEvent + | ToolCallEndEvent + | ThinkingEvent + | DoneEvent + | ErrorEvent; + export function useAIChat() { const params = useParams(); const binaryId = params.binaryId!; @@ -20,6 +60,8 @@ export function useAIChat() { const [streamingContent, setStreamingContent] = useState(""); const [isStreaming, setIsStreaming] = useState(false); const [streamError, setStreamError] = useState(null); + const [toolStatus, setToolStatus] = useState(null); + const [isToolRunning, setIsToolRunning] = useState(false); const [pendingUserMessage, setPendingUserMessage] = useState(null); const { data: binary } = useGetBinary(binaryId, { @@ -89,6 +131,8 @@ export function useAIChat() { setIsStreaming(true); setStreamingContent(""); + setToolStatus(null); + setIsToolRunning(false); setStreamError(null); setPendingUserMessage({ role: "USER", @@ -137,23 +181,48 @@ export function useAIChat() { if (!line.startsWith("data:")) continue; const event = JSON.parse(line.slice(5).trim()); - if (event.type === "chunk") { - fullContent += event.content; - setStreamingContent(fullContent); - } else if (event.type === "done") { - setIsStreaming(false); - setStreamingContent(""); - setPendingUserMessage(null); - queryClient.invalidateQueries({ - queryKey: [`/chat/sessions/${activeSessionId}/messages`], - }); - queryClient.invalidateQueries({ - queryKey: [`/binaries/${binaryId}/chat/sessions`], - }); - } else if (event.type === "error") { - setIsStreaming(false); - setPendingUserMessage(null); - setStreamError(event.message); + const sseEvent = event as SSEEvent; + switch (sseEvent.type) { + case "chunk": + fullContent += sseEvent.content; + setStreamingContent(fullContent); + setToolStatus(null); + setIsToolRunning(false); + break; + case "tool_call_start": + setToolStatus(`Calling ${sseEvent.tool}...`); + setIsToolRunning(true); + break; + case "tool_call_end": + setToolStatus( + `${sseEvent.tool} completed (${sseEvent.durationMs}ms): ${sseEvent.resultSummary}` + ); + setIsToolRunning(false); + break; + case "thinking": + setToolStatus(sseEvent.message); + setIsToolRunning(true); + break; + case "done": + setIsStreaming(false); + setStreamingContent(""); + setToolStatus(null); + setIsToolRunning(false); + setPendingUserMessage(null); + queryClient.invalidateQueries({ + queryKey: [`/chat/sessions/${activeSessionId}/messages`], + }); + queryClient.invalidateQueries({ + queryKey: [`/binaries/${binaryId}/chat/sessions`], + }); + break; + case "error": + setIsStreaming(false); + setToolStatus(null); + setIsToolRunning(false); + setPendingUserMessage(null); + setStreamError(sseEvent.message); + break; } } } @@ -164,23 +233,48 @@ export function useAIChat() { if (!line.startsWith("data:")) continue; const event = JSON.parse(line.slice(5).trim()); - if (event.type === "chunk") { - fullContent += event.content; - setStreamingContent(fullContent); - } else if (event.type === "done") { - setIsStreaming(false); - setStreamingContent(""); - setPendingUserMessage(null); - queryClient.invalidateQueries({ - queryKey: [`/chat/sessions/${activeSessionId}/messages`], - }); - queryClient.invalidateQueries({ - queryKey: [`/binaries/${binaryId}/chat/sessions`], - }); - } else if (event.type === "error") { - setIsStreaming(false); - setPendingUserMessage(null); - setStreamError(event.message); + const sseEvent = event as SSEEvent; + switch (sseEvent.type) { + case "chunk": + fullContent += sseEvent.content; + setStreamingContent(fullContent); + setToolStatus(null); + setIsToolRunning(false); + break; + case "tool_call_start": + setToolStatus(`Calling ${sseEvent.tool}...`); + setIsToolRunning(true); + break; + case "tool_call_end": + setToolStatus( + `${sseEvent.tool} completed (${sseEvent.durationMs}ms): ${sseEvent.resultSummary}` + ); + setIsToolRunning(false); + break; + case "thinking": + setToolStatus(sseEvent.message); + setIsToolRunning(true); + break; + case "done": + setIsStreaming(false); + setStreamingContent(""); + setToolStatus(null); + setIsToolRunning(false); + setPendingUserMessage(null); + queryClient.invalidateQueries({ + queryKey: [`/chat/sessions/${activeSessionId}/messages`], + }); + queryClient.invalidateQueries({ + queryKey: [`/binaries/${binaryId}/chat/sessions`], + }); + break; + case "error": + setIsStreaming(false); + setToolStatus(null); + setIsToolRunning(false); + setPendingUserMessage(null); + setStreamError(sseEvent.message); + break; } } } @@ -190,6 +284,8 @@ export function useAIChat() { } catch (err: unknown) { if (err instanceof Error && err.name !== "AbortError") { setIsStreaming(false); + setToolStatus(null); + setIsToolRunning(false); setPendingUserMessage(null); setStreamError(err.message || "An unexpected error occurred"); } @@ -201,17 +297,14 @@ export function useAIChat() { const cancelStream = useCallback(() => { abortRef.current?.abort(); setIsStreaming(false); + setToolStatus(null); + setIsToolRunning(false); setPendingUserMessage(null); }, []); const messages: MessageResponse[] = [...historicalMessages]; if (isStreaming && pendingUserMessage) { messages.push(pendingUserMessage); - messages.push({ - role: "ASSISTANT", - content: streamingContent, - createdAt: new Date().toISOString(), - }); } const userFriendlyError = streamError @@ -231,6 +324,8 @@ export function useAIChat() { isLoadingMessages: messagesQuery.isLoading, isStreaming, streamingContent, + toolStatus, + isToolRunning, streamError: userFriendlyError, isCreatingSession: createSessionMutation.isPending, createSession, diff --git a/src/index.css b/src/index.css index 2266220..7fb8492 100644 --- a/src/index.css +++ b/src/index.css @@ -28,6 +28,7 @@ --color-input: var(--input); --color-border: var(--border); --color-destructive: var(--destructive); + --color-success: var(--success); --color-accent-foreground: var(--accent-foreground); --color-accent: var(--accent); --color-muted-foreground: var(--muted-foreground);