diff --git a/apps/studio/src/components/AiChatPanel.tsx b/apps/studio/src/components/AiChatPanel.tsx index b69a68721..9be289680 100644 --- a/apps/studio/src/components/AiChatPanel.tsx +++ b/apps/studio/src/components/AiChatPanel.tsx @@ -7,6 +7,7 @@ import type { UIMessage } from 'ai'; import { Bot, X, Send, Trash2, Sparkles, Wrench, CheckCircle2, XCircle, Loader2, ShieldAlert, + ChevronDown, ChevronRight, Brain, Zap, } from 'lucide-react'; import { Button } from '@/components/ui/button'; import { ScrollArea } from '@/components/ui/scroll-area'; @@ -33,6 +34,15 @@ interface AgentSummary { role: string; } +/** + * Track active thinking/reasoning state during streaming. + */ +interface ThinkingState { + reasoning: string[]; + activeSteps: Map; + completedSteps: string[]; +} + /** * Extract the text content from a UIMessage's parts array. */ @@ -160,6 +170,88 @@ function useAgentList(baseUrl: string) { // ── Tool Invocation State Labels ──────────────────────────────────── +/** + * Display reasoning/thinking information in a collapsible section. + */ +interface ReasoningDisplayProps { + reasoning: string[]; +} + +function ReasoningDisplay({ reasoning }: ReasoningDisplayProps) { + const [isExpanded, setIsExpanded] = useState(false); + + if (reasoning.length === 0) return null; + + return ( +
+ + {isExpanded && ( +
+ {reasoning.map((step, idx) => ( +

+ {step} +

+ ))} +
+ )} +
+ ); +} + +/** + * Display active step progress indicators. + */ +interface StepProgressProps { + activeSteps: Map; + completedSteps: string[]; +} + +function StepProgress({ activeSteps, completedSteps }: StepProgressProps) { + if (activeSteps.size === 0) return null; + + const totalSteps = completedSteps.length + activeSteps.size; + const currentStep = completedSteps.length + 1; + + return ( +
+
+ + + Step {currentStep} of {totalSteps} + +
+ {Array.from(activeSteps.values()).map((step, idx) => ( +
+ + {step.stepName} +
+ ))} +
+ ); +} + +// ── Tool Invocation State Labels ──────────────────────────────────── + interface ToolInvocationDisplayProps { part: Extract; onApprove?: (approvalId: string) => void; @@ -175,6 +267,21 @@ function ToolInvocationDisplay({ part, onApprove, onDeny }: ToolInvocationDispla switch (part.state) { case 'input-streaming': + return ( +
+ +
+ Planning to call {toolLabel} + {argsText && ( +

{argsText}

+ )} +
+
+ ); + case 'input-available': return (
(loadSelectedAgent); + const [thinkingState, setThinkingState] = useState({ + reasoning: [], + activeSteps: new Map(), + completedSteps: [], + }); const scrollRef = useRef(null); const inputRef = useRef(null); const baseUrl = getApiBaseUrl(); @@ -316,10 +428,47 @@ export function AiChatPanel() { const { messages, sendMessage, setMessages, status, error, addToolApprovalResponse } = useChat({ transport, messages: initialMessages, + onFinish: () => { + // Reset thinking state when stream completes + setThinkingState({ + reasoning: [], + activeSteps: new Map(), + completedSteps: [], + }); + }, }); const isStreaming = status === 'streaming' || status === 'submitted'; + // Extract reasoning and step progress from the latest assistant message parts + useEffect(() => { + if (!isStreaming || messages.length === 0) return; + + // Get the latest message + const lastMessage = messages[messages.length - 1]; + if (lastMessage.role !== 'assistant') return; + + // Process message parts for reasoning and steps + const reasoning: string[] = []; + const activeSteps = new Map(); + const completedSteps: string[] = []; + + (lastMessage.parts || []).forEach((part: any) => { + if (part.type === 'reasoning-delta' || part.type === 'reasoning') { + reasoning.push(part.text); + } else if (part.type === 'step-start') { + activeSteps.set(part.stepId, { + stepName: part.stepName, + startedAt: Date.now(), + }); + } else if (part.type === 'step-finish') { + completedSteps.push(part.stepName); + } + }); + + setThinkingState({ reasoning, activeSteps, completedSteps }); + }, [messages, isStreaming]); + // Persist messages to localStorage whenever they change useEffect(() => { if (messages.length > 0) { @@ -513,10 +662,30 @@ export function AiChatPanel() { ); })} {isStreaming && ( -
- - Thinking… -
+ <> + {/* Show reasoning if available */} + {thinkingState.reasoning.length > 0 && ( +
+ +
+ )} + {/* Show step progress if available */} + {thinkingState.activeSteps.size > 0 && ( +
+ +
+ )} + {/* Default thinking indicator when no detailed state available */} + {thinkingState.reasoning.length === 0 && thinkingState.activeSteps.size === 0 && ( +
+ + Thinking… +
+ )} + )} {error && (
diff --git a/packages/services/service-ai/src/__tests__/vercel-stream-encoder.test.ts b/packages/services/service-ai/src/__tests__/vercel-stream-encoder.test.ts index 8c1fe83df..48e304d62 100644 --- a/packages/services/service-ai/src/__tests__/vercel-stream-encoder.test.ts +++ b/packages/services/service-ai/src/__tests__/vercel-stream-encoder.test.ts @@ -128,6 +128,53 @@ describe('encodeStreamPart', () => { const part = { type: 'unknown-internal' } as unknown as TextStreamPart; expect(encodeStreamPart(part)).toBe(''); }); + + it('should encode reasoning-start part with g: prefix', () => { + const part = { + type: 'reasoning-start', + id: 'r1', + } as unknown as TextStreamPart; + + const frame = encodeStreamPart(part); + expect(frame).toBe('g:{"text":""}\n'); + }); + + it('should encode reasoning-delta part with g: prefix', () => { + const part = { + type: 'reasoning-delta', + id: 'r1', + text: 'Let me think through this step by step...', + } as unknown as TextStreamPart; + + const frame = encodeStreamPart(part); + expect(frame).toBe('g:{"text":"Let me think through this step by step..."}\n'); + }); + + it('should encode reasoning-end part as empty (no specific end marker)', () => { + const part = { + type: 'reasoning-end', + id: 'r1', + } as unknown as TextStreamPart; + + const frame = encodeStreamPart(part); + expect(frame).toBe(''); + }); + + it('should pass through custom step events', () => { + const part = { + type: 'step-start', + stepId: 'step_1', + stepName: 'Query database', + } as unknown as TextStreamPart; + + const frame = encodeStreamPart(part); + const payload = parseSSE(frame); + expect(payload).toEqual({ + type: 'step-start', + stepId: 'step_1', + stepName: 'Query database', + }); + }); }); // ───────────────────────────────────────────────────────────────── diff --git a/packages/services/service-ai/src/stream/vercel-stream-encoder.ts b/packages/services/service-ai/src/stream/vercel-stream-encoder.ts index 72b7a1186..88b41f774 100644 --- a/packages/services/service-ai/src/stream/vercel-stream-encoder.ts +++ b/packages/services/service-ai/src/stream/vercel-stream-encoder.ts @@ -23,6 +23,14 @@ function sse(data: object): string { return `data: ${JSON.stringify(data)}\n\n`; } +/** + * Encode data using Vercel AI SDK Data Stream Protocol prefixes. + * @see https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol + */ +function dataStreamLine(prefix: string, data: object): string { + return `${prefix}:${JSON.stringify(data)}\n`; +} + // ── Public API ────────────────────────────────────────────────────── /** @@ -71,8 +79,24 @@ export function encodeStreamPart(part: TextStreamPart): string { errorText: String(part.error), }); + // Handle reasoning/thinking streams (DeepSeek R1, o1-style models) + // Use 'g:' prefix for reasoning content per Vercel AI SDK protocol + case 'reasoning-start': + return dataStreamLine('g', { text: '' }); + + case 'reasoning-delta': + return dataStreamLine('g', { text: part.text }); + + case 'reasoning-end': + return ''; // No specific end marker needed for reasoning + // finish-step and finish are handled by the generator, not here default: + // Pass through any unknown event types that might be custom + // (e.g., step-start, step-finish from custom providers) + if ((part as any).type?.startsWith('step-')) { + return sse(part as any); + } return ''; } }