Merge pull request #1098 from objectstack-ai/claude/improve-ai-chat-thinking-phase

Fix AI Chat Reasoning Display with Vercel AI SDK Data Stream Protocol

Author: hotlong

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<string, { stepName: string; startedAt: number }>;
+  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 (
+    <div
+      data-testid="reasoning-display"
+      className="flex flex-col gap-1 rounded-md border border-border/30 bg-muted/30 px-2.5 py-2 text-xs"
+    >
+      <button
+        onClick={() => setIsExpanded(!isExpanded)}
+        className="flex items-center gap-1.5 text-left text-muted-foreground hover:text-foreground transition-colors"
+      >
+        {isExpanded ? (
+          <ChevronDown className="h-3 w-3 shrink-0" />
+        ) : (
+          <ChevronRight className="h-3 w-3 shrink-0" />
+        )}
+        <Brain className="h-3 w-3 shrink-0" />
+        <span className="font-medium">Thinking</span>
+        <span className="text-[10px] opacity-60">
+          ({reasoning.length} step{reasoning.length !== 1 ? 's' : ''})
+        </span>
+      </button>
+      {isExpanded && (
+        <div className="mt-1 space-y-1 pl-5 text-muted-foreground italic border-l-2 border-border/30">
+          {reasoning.map((step, idx) => (
+            <p key={idx} className="text-[11px] leading-relaxed">
+              {step}
+            </p>
+          ))}
+        </div>
+      )}
+    </div>
+  );
+}
+
+/**
+ * Display active step progress indicators.
+ */
+interface StepProgressProps {
+  activeSteps: Map<string, { stepName: string; startedAt: number }>;
+  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 (
+    <div
+      data-testid="step-progress"
+      className="flex flex-col gap-1.5 rounded-md border border-blue-500/30 bg-blue-500/5 px-2.5 py-2 text-xs"
+    >
+      <div className="flex items-center gap-2">
+        <Zap className="h-3 w-3 shrink-0 text-blue-600 dark:text-blue-400" />
+        <span className="font-medium text-blue-700 dark:text-blue-300">
+          Step {currentStep} of {totalSteps}
+        </span>
+      </div>
+      {Array.from(activeSteps.values()).map((step, idx) => (
+        <div key={idx} className="flex items-center gap-2 pl-5">
+          <Loader2 className="h-3 w-3 shrink-0 animate-spin text-blue-600 dark:text-blue-400" />
+          <span className="text-blue-700 dark:text-blue-300">{step.stepName}</span>
+        </div>
+      ))}
+    </div>
+  );
+}
+
+// ── Tool Invocation State Labels ────────────────────────────────────
+
 interface ToolInvocationDisplayProps {
   part: Extract<UIMessage['parts'][number], { type: 'dynamic-tool' }>;
   onApprove?: (approvalId: string) => void;
@@ -175,6 +267,21 @@ function ToolInvocationDisplay({ part, onApprove, onDeny }: ToolInvocationDispla
 
   switch (part.state) {
     case 'input-streaming':
+      return (
+        <div
+          data-testid="tool-invocation-planning"
+          className="flex items-start gap-2 rounded-md border border-blue-500/40 bg-blue-500/10 px-2.5 py-2 text-xs"
+        >
+          <Loader2 className="mt-0.5 h-3.5 w-3.5 shrink-0 animate-spin text-blue-600 dark:text-blue-400" />
+          <div className="min-w-0">
+            <span className="font-medium text-blue-700 dark:text-blue-300">Planning to call {toolLabel}</span>
+            {argsText && (
+              <p className="mt-0.5 truncate text-blue-600/80 dark:text-blue-300/80">{argsText}</p>
+            )}
+          </div>
+        </div>
+      );
+
     case 'input-available':
       return (
         <div
@@ -289,6 +396,11 @@ export function AiChatPanel() {
   const { isOpen, setOpen, toggle } = useAiChatPanel();
   const [input, setInput] = useState('');
   const [selectedAgent, setSelectedAgent] = useState<string>(loadSelectedAgent);
+  const [thinkingState, setThinkingState] = useState<ThinkingState>({
+    reasoning: [],
+    activeSteps: new Map(),
+    completedSteps: [],
+  });
   const scrollRef = useRef<HTMLDivElement>(null);
   const inputRef = useRef<HTMLTextAreaElement>(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<string, { stepName: string; startedAt: number }>();
+    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 && (
-            <div className="mr-8 flex items-center gap-2 rounded-lg bg-muted px-3 py-2 text-sm text-muted-foreground">
-              <span className="inline-block h-2 w-2 animate-pulse rounded-full bg-primary" />
-              Thinking…
-            </div>
+            <>
+              {/* Show reasoning if available */}
+              {thinkingState.reasoning.length > 0 && (
+                <div className="mr-8">
+                  <ReasoningDisplay reasoning={thinkingState.reasoning} />
+                </div>
+              )}
+              {/* Show step progress if available */}
+              {thinkingState.activeSteps.size > 0 && (
+                <div className="mr-8">
+                  <StepProgress
+                    activeSteps={thinkingState.activeSteps}
+                    completedSteps={thinkingState.completedSteps}
+                  />
+                </div>
+              )}
+              {/* Default thinking indicator when no detailed state available */}
+              {thinkingState.reasoning.length === 0 && thinkingState.activeSteps.size === 0 && (
+                <div className="mr-8 flex items-center gap-2 rounded-lg bg-muted px-3 py-2 text-sm text-muted-foreground">
+                  <span className="inline-block h-2 w-2 animate-pulse rounded-full bg-primary" />
+                  Thinking…
+                </div>
+              )}
+            </>
           )}
           {error && (
             <div className="flex items-start gap-2 rounded-lg border border-destructive/50 bg-destructive/10 px-3 py-2 text-sm text-destructive">

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<ToolSet>;
     expect(encodeStreamPart(part)).toBe('');
   });
+
+  it('should encode reasoning-start part with g: prefix', () => {
+    const part = {
+      type: 'reasoning-start',
+      id: 'r1',
+    } as unknown as TextStreamPart<ToolSet>;
+
+    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<ToolSet>;
+
+    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<ToolSet>;
+
+    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<ToolSet>;
+
+    const frame = encodeStreamPart(part);
+    const payload = parseSSE(frame);
+    expect(payload).toEqual({
+      type: 'step-start',
+      stepId: 'step_1',
+      stepName: 'Query database',
+    });
+  });
 });
 
 // ─────────────────────────────────────────────────────────────────

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<ToolSet>): 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 '';
   }
 }