From 3ab3ca22ffba18336c73a92eb60564a8b6d5eb0c Mon Sep 17 00:00:00 2001
From: lucaronin <luca.rossi.917@gmail.com>
Date: Wed, 4 Mar 2026 12:09:02 +0100
Subject: [PATCH] feat: add reasoning streaming, markdown response, structured
 context snapshot

- Rust: add ThinkingDelta event to ClaudeStreamEvent for reasoning chunks
- ai-agent.ts: forward ThinkingDelta events via onThinking callback
- useAiAgent: stream reasoning live, accumulate response internally,
  reveal as complete block on done
- AiMessage: auto-collapse reasoning when done, use MarkdownContent
  for response rendering, update tests for new behavior
- ai-context: add buildContextSnapshot() for structured JSON context
  with activeNote, openTabs, noteListFilter, vault summary

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src-tauri/src/claude_cli.rs       |  9 +++
 src/components/AiMessage.test.tsx | 39 ++++++++-----
 src/components/AiMessage.tsx      | 29 ++++++++--
 src/hooks/useAiAgent.ts           | 32 +++++++++--
 src/utils/ai-agent.ts             |  5 ++
 src/utils/ai-context.ts           | 94 +++++++++++++++++++++++++++++--
 6 files changed, 177 insertions(+), 31 deletions(-)

diff --git a/src-tauri/src/claude_cli.rs b/src-tauri/src/claude_cli.rs
index 3876ccf2..726dfd5f 100644
--- a/src-tauri/src/claude_cli.rs
+++ b/src-tauri/src/claude_cli.rs
@@ -19,6 +19,8 @@ pub enum ClaudeStreamEvent {
     Init { session_id: String },
     /// Incremental text chunk.
     TextDelta { text: String },
+    /// Incremental thinking/reasoning chunk.
+    ThinkingDelta { text: String },
     /// A tool call started (agent mode only).
     ToolStart {
         tool_name: String,
@@ -402,6 +404,13 @@ where
                         });
                     }
                 }
+                Some("thinking_delta") => {
+                    if let Some(text) = delta["thinking"].as_str() {
+                        emit(ClaudeStreamEvent::ThinkingDelta {
+                            text: text.to_string(),
+                        });
+                    }
+                }
                 Some("input_json_delta") => {
                     if let (Some(partial), Some(ref tid)) =
                         (delta["partial_json"].as_str(), &state.current_tool_id)
diff --git a/src/components/AiMessage.test.tsx b/src/components/AiMessage.test.tsx
index aa8a040f..7137c34d 100644
--- a/src/components/AiMessage.test.tsx
+++ b/src/components/AiMessage.test.tsx
@@ -2,15 +2,20 @@ import { describe, it, expect, vi } from 'vitest'
 import { render, screen, fireEvent } from '@testing-library/react'
 import { AiMessage } from './AiMessage'
 
+vi.mock('./MarkdownContent', () => ({
+  MarkdownContent: ({ content }: { content: string }) => <div data-testid="markdown-content">{content}</div>,
+}))
+
 describe('AiMessage', () => {
   it('renders user message', () => {
     render(<AiMessage userMessage="Hello AI" actions={[]} />)
     expect(screen.getByText('Hello AI')).toBeTruthy()
   })
 
-  it('renders response text', () => {
-    render(<AiMessage userMessage="Ask" actions={[]} response="Here is the answer" />)
-    expect(screen.getByText('Here is the answer')).toBeTruthy()
+  it('renders response as markdown', () => {
+    render(<AiMessage userMessage="Ask" actions={[]} response="Here is the **answer**" />)
+    expect(screen.getByTestId('markdown-content')).toBeTruthy()
+    expect(screen.getByText('Here is the **answer**')).toBeTruthy()
   })
 
   it('shows undo button with response', () => {
@@ -18,23 +23,31 @@ describe('AiMessage', () => {
     expect(screen.getByTestId('undo-button')).toBeTruthy()
   })
 
-  it('renders reasoning toggle collapsed by default', () => {
-    render(<AiMessage userMessage="Ask" reasoning="Thinking about it..." actions={[]} />)
+  it('shows reasoning expanded while streaming (reasoningDone=false)', () => {
+    render(<AiMessage userMessage="Ask" reasoning="Thinking about it..." reasoningDone={false} actions={[]} />)
     expect(screen.getByTestId('reasoning-toggle')).toBeTruthy()
-    expect(screen.queryByTestId('reasoning-content')).toBeNull()
-  })
-
-  it('expands reasoning on toggle click', () => {
-    render(<AiMessage userMessage="Ask" reasoning="Thinking about it..." actions={[]} />)
-    fireEvent.click(screen.getByTestId('reasoning-toggle'))
     expect(screen.getByTestId('reasoning-content')).toBeTruthy()
     expect(screen.getByText('Thinking about it...')).toBeTruthy()
   })
 
-  it('collapses reasoning on second click', () => {
-    render(<AiMessage userMessage="Ask" reasoning="Thinking..." actions={[]} />)
+  it('auto-collapses reasoning when reasoningDone=true', () => {
+    render(<AiMessage userMessage="Ask" reasoning="Thinking..." reasoningDone actions={[]} />)
+    expect(screen.getByTestId('reasoning-toggle')).toBeTruthy()
+    expect(screen.queryByTestId('reasoning-content')).toBeNull()
+  })
+
+  it('expands collapsed reasoning on toggle click', () => {
+    render(<AiMessage userMessage="Ask" reasoning="Thinking..." reasoningDone actions={[]} />)
+    // Starts collapsed (reasoningDone=true)
+    expect(screen.queryByTestId('reasoning-content')).toBeNull()
     fireEvent.click(screen.getByTestId('reasoning-toggle'))
     expect(screen.getByTestId('reasoning-content')).toBeTruthy()
+  })
+
+  it('collapses expanded reasoning on toggle click', () => {
+    render(<AiMessage userMessage="Ask" reasoning="Thinking..." reasoningDone={false} actions={[]} />)
+    // Starts expanded (reasoningDone=false)
+    expect(screen.getByTestId('reasoning-content')).toBeTruthy()
     fireEvent.click(screen.getByTestId('reasoning-toggle'))
     expect(screen.queryByTestId('reasoning-content')).toBeNull()
   })
diff --git a/src/components/AiMessage.tsx b/src/components/AiMessage.tsx
index 3202cb12..a47c24c9 100644
--- a/src/components/AiMessage.tsx
+++ b/src/components/AiMessage.tsx
@@ -1,6 +1,7 @@
-import { useState, useCallback } from 'react'
+import { useState, useCallback, useEffect, useRef } from 'react'
 import { CaretRight, CaretDown, Brain, ArrowCounterClockwise } from '@phosphor-icons/react'
 import { AiActionCard, type AiActionStatus } from './AiActionCard'
+import { MarkdownContent } from './MarkdownContent'
 
 export interface AiAction {
   tool: string
@@ -15,6 +16,7 @@ export interface AiAction {
 export interface AiMessageProps {
   userMessage: string
   reasoning?: string
+  reasoningDone?: boolean
   actions: AiAction[]
   response?: string
   isStreaming?: boolean
@@ -44,6 +46,14 @@ function UserBubble({ content }: { content: string }) {
 function ReasoningBlock({ text, expanded, onToggle }: {
   text: string; expanded: boolean; onToggle: () => void
 }) {
+  const contentRef = useRef<HTMLDivElement>(null)
+
+  useEffect(() => {
+    if (expanded && contentRef.current) {
+      contentRef.current.scrollTop = contentRef.current.scrollHeight
+    }
+  }, [expanded, text])
+
   return (
     <div style={{ marginBottom: 8 }}>
       <button
@@ -58,8 +68,9 @@ function ReasoningBlock({ text, expanded, onToggle }: {
       </button>
       {expanded && (
         <div
+          ref={contentRef}
           className="text-muted-foreground"
-          style={{ fontSize: 12, lineHeight: 1.5, padding: '4px 0 4px 20px' }}
+          style={{ fontSize: 12, lineHeight: 1.5, padding: '4px 0 4px 20px', maxHeight: 200, overflowY: 'auto' }}
           data-testid="reasoning-content"
         >
           {text}
@@ -98,7 +109,7 @@ function ActionCardsList({ actions, onOpenNote, expandedIds, onToggleExpand }: {
 function ResponseBlock({ text }: { text: string }) {
   return (
     <div style={{ marginBottom: 4 }}>
-      <div style={{ fontSize: 13, lineHeight: 1.6 }}>{text}</div>
+      <MarkdownContent content={text} />
       <button
         className="flex items-center gap-1 border-none bg-transparent p-0 text-muted-foreground cursor-pointer hover:text-foreground transition-colors"
         style={{ fontSize: 11, marginTop: 4 }}
@@ -123,10 +134,16 @@ function StreamingIndicator() {
   )
 }
 
-export function AiMessage({ userMessage, reasoning, actions, response, isStreaming, onOpenNote }: AiMessageProps) {
-  const [reasoningExpanded, setReasoningExpanded] = useState(false)
+export function AiMessage({ userMessage, reasoning, reasoningDone, actions, response, isStreaming, onOpenNote }: AiMessageProps) {
+  // Manual override: null = follow auto behavior, true/false = user forced
+  const [userOverride, setUserOverride] = useState(false)
   const [expandedActions, setExpandedActions] = useState<Set<string>>(new Set())
 
+  // Auto: expanded while reasoning streams, collapsed once done
+  // User can manually toggle to override the auto state
+  const autoExpanded = !reasoningDone
+  const reasoningExpanded = userOverride ? !autoExpanded : autoExpanded
+
   const toggleAction = useCallback((toolId: string) => {
     setExpandedActions(prev => {
       const next = new Set(prev)
@@ -143,7 +160,7 @@ export function AiMessage({ userMessage, reasoning, actions, response, isStreami
         <ReasoningBlock
           text={reasoning}
           expanded={reasoningExpanded}
-          onToggle={() => setReasoningExpanded(!reasoningExpanded)}
+          onToggle={() => setUserOverride(prev => !prev)}
         />
       )}
       {actions.length > 0 && (
diff --git a/src/hooks/useAiAgent.ts b/src/hooks/useAiAgent.ts
index 793d31ce..9f182625 100644
--- a/src/hooks/useAiAgent.ts
+++ b/src/hooks/useAiAgent.ts
@@ -3,6 +3,9 @@
  * Uses Claude CLI subprocess with MCP tools via Tauri.
  *
  * States: idle -> thinking -> tool-executing -> done/error
+ *
+ * Reasoning streams live while Claude thinks, then auto-collapses.
+ * Response text accumulates internally and is revealed as a complete block on done.
  */
 import { useState, useCallback, useRef, useEffect } from 'react'
 import type { AiAction } from '../components/AiMessage'
@@ -14,6 +17,7 @@ export type AgentStatus = 'idle' | 'thinking' | 'tool-executing' | 'done' | 'err
 export interface AiAgentMessage {
   userMessage: string
   reasoning?: string
+  reasoningDone?: boolean
   actions: AiAction[]
   response?: string
   isStreaming?: boolean
@@ -25,6 +29,7 @@ export function useAiAgent(vaultPath: string, contextPrompt?: string) {
   const [status, setStatus] = useState<AgentStatus>('idle')
   const abortRef = useRef({ aborted: false })
   const contextRef = useRef(contextPrompt)
+  const responseAccRef = useRef('')
   useEffect(() => {
     contextRef.current = contextPrompt
   }, [contextPrompt])
@@ -42,6 +47,7 @@ export function useAiAgent(vaultPath: string, contextPrompt?: string) {
     }
 
     abortRef.current = { aborted: false }
+    responseAccRef.current = ''
 
     const messageId = nextMessageId()
     setMessages(prev => [...prev, {
@@ -53,23 +59,31 @@ export function useAiAgent(vaultPath: string, contextPrompt?: string) {
       setMessages(prev => prev.map(m => m.id === messageId ? fn(m) : m))
     }
 
-    // When a contextual prompt is provided (from buildContextualPrompt),
-    // use it directly — it already includes the system preamble.
+    const markReasoningDone = () => {
+      update(m => m.reasoningDone ? m : { ...m, reasoningDone: true })
+    }
+
     const systemPrompt = contextRef.current ?? buildAgentSystemPrompt()
 
     await streamClaudeAgent(text.trim(), systemPrompt, vaultPath, {
-      onText: (text) => {
+      onThinking: (chunk) => {
         if (abortRef.current.aborted) return
-        update(m => ({ ...m, response: (m.response ?? '') + text }))
+        update(m => ({ ...m, reasoning: (m.reasoning ?? '') + chunk }))
+      },
+
+      onText: (chunk) => {
+        if (abortRef.current.aborted) return
+        markReasoningDone()
+        responseAccRef.current += chunk
       },
 
       onToolStart: (toolName, toolId, input) => {
         if (abortRef.current.aborted) return
+        markReasoningDone()
         setStatus('tool-executing')
         update(m => {
           const existing = m.actions.find(a => a.toolId === toolId)
           if (existing) {
-            // Re-emitted with input data — update the existing action
             return {
               ...m,
               actions: m.actions.map(a =>
@@ -103,10 +117,12 @@ export function useAiAgent(vaultPath: string, contextPrompt?: string) {
       onError: (error) => {
         if (abortRef.current.aborted) return
         setStatus('error')
+        const partial = responseAccRef.current
         update(m => ({
           ...m,
           isStreaming: false,
-          response: (m.response ?? '') + `\nError: ${error}`,
+          reasoningDone: true,
+          response: partial ? `${partial}\n\nError: ${error}` : `Error: ${error}`,
           actions: m.actions.map(a =>
             a.status === 'pending' ? { ...a, status: 'error' as const } : a,
           ),
@@ -116,9 +132,12 @@ export function useAiAgent(vaultPath: string, contextPrompt?: string) {
       onDone: () => {
         if (abortRef.current.aborted) return
         setStatus('done')
+        const finalResponse = responseAccRef.current || undefined
         update(m => ({
           ...m,
           isStreaming: false,
+          reasoningDone: true,
+          response: finalResponse,
           actions: m.actions.map(a => a.status === 'pending' ? { ...a, status: 'done' as const } : a),
         }))
       },
@@ -127,6 +146,7 @@ export function useAiAgent(vaultPath: string, contextPrompt?: string) {
 
   const clearConversation = useCallback(() => {
     abortRef.current.aborted = true
+    responseAccRef.current = ''
     setMessages([])
     setStatus('idle')
   }, [])
diff --git a/src/utils/ai-agent.ts b/src/utils/ai-agent.ts
index 0553c957..035a4075 100644
--- a/src/utils/ai-agent.ts
+++ b/src/utils/ai-agent.ts
@@ -24,6 +24,7 @@ export function buildAgentSystemPrompt(vaultContext?: string): string {
 type ClaudeAgentStreamEvent =
   | { kind: 'Init'; session_id: string }
   | { kind: 'TextDelta'; text: string }
+  | { kind: 'ThinkingDelta'; text: string }
   | { kind: 'ToolStart'; tool_name: string; tool_id: string; input?: string }
   | { kind: 'ToolDone'; tool_id: string; output?: string }
   | { kind: 'Result'; text: string; session_id: string }
@@ -32,6 +33,7 @@ type ClaudeAgentStreamEvent =
 
 export interface AgentStreamCallbacks {
   onText: (text: string) => void
+  onThinking: (text: string) => void
   onToolStart: (toolName: string, toolId: string, input?: string) => void
   onToolDone: (toolId: string, output?: string) => void
   onError: (message: string) => void
@@ -65,6 +67,9 @@ export async function streamClaudeAgent(
       case 'TextDelta':
         callbacks.onText(data.text)
         break
+      case 'ThinkingDelta':
+        callbacks.onThinking(data.text)
+        break
       case 'ToolStart':
         callbacks.onToolStart(data.tool_name, data.tool_id, data.input)
         break
diff --git a/src/utils/ai-context.ts b/src/utils/ai-context.ts
index 0b0c4a99..55caebf3 100644
--- a/src/utils/ai-context.ts
+++ b/src/utils/ai-context.ts
@@ -1,6 +1,6 @@
 /**
- * AI contextual chat — builds the context note list from the active note
- * and its first-degree linked notes (outgoingLinks + relationships).
+ * AI contextual chat — builds a structured context snapshot from the active note,
+ * open tabs, vault metadata, and optional explicit note references.
  */
 
 import type { VaultEntry } from '../types'
@@ -34,19 +34,16 @@ export function collectLinkedEntries(
     }
   }
 
-  // outgoingLinks are raw targets (no [[ ]] wrapper)
   for (const target of active.outgoingLinks) {
     addTarget(target)
   }
 
-  // relationships values are wikilink references like [[target]]
   for (const refs of Object.values(active.relationships)) {
     for (const ref of refs) {
       addTarget(wikilinkTarget(ref))
     }
   }
 
-  // belongsTo and relatedTo are also wikilink references
   for (const ref of active.belongsTo) {
     addTarget(wikilinkTarget(ref))
   }
@@ -57,7 +54,92 @@ export function collectLinkedEntries(
   return linked
 }
 
-/** Build a contextual system prompt from the active note and its linked notes. */
+/** A note reference from the user's [[wikilink]] selection in the chat input. */
+export interface NoteReference {
+  title: string
+  path: string
+  type: string | null
+}
+
+/** Parameters for building the structured context snapshot. */
+export interface ContextSnapshotParams {
+  activeEntry: VaultEntry
+  allContent: Record<string, string>
+  openTabs?: VaultEntry[]
+  noteListFilter?: { type: string | null; query: string }
+  entries: VaultEntry[]
+  references?: NoteReference[]
+}
+
+function entryFrontmatter(e: VaultEntry): Record<string, unknown> {
+  const fm: Record<string, unknown> = {}
+  if (e.isA) fm.type = e.isA
+  if (e.status) fm.status = e.status
+  if (e.owner) fm.owner = e.owner
+  if (e.belongsTo.length > 0) fm.belongsTo = e.belongsTo
+  if (e.relatedTo.length > 0) fm.relatedTo = e.relatedTo
+  if (Object.keys(e.relationships).length > 0) fm.relationships = e.relationships
+  return fm
+}
+
+/** Build a structured context snapshot as a system prompt for Claude. */
+export function buildContextSnapshot(params: ContextSnapshotParams): string {
+  const { activeEntry, allContent, openTabs, noteListFilter, entries, references } = params
+
+  const snapshot: Record<string, unknown> = {
+    activeNote: {
+      path: activeEntry.path,
+      title: activeEntry.title,
+      type: activeEntry.isA ?? 'Note',
+      frontmatter: entryFrontmatter(activeEntry),
+      body: allContent[activeEntry.path] ?? '',
+    },
+  }
+
+  const otherTabs = openTabs?.filter(t => t.path !== activeEntry.path)
+  if (otherTabs && otherTabs.length > 0) {
+    snapshot.openTabs = otherTabs.map(t => ({
+      path: t.path,
+      title: t.title,
+      type: t.isA ?? 'Note',
+      frontmatter: entryFrontmatter(t),
+    }))
+  }
+
+  if (noteListFilter && (noteListFilter.type || noteListFilter.query)) {
+    snapshot.noteListFilter = noteListFilter
+  }
+
+  const types = new Set<string>()
+  for (const e of entries) {
+    if (e.isA) types.add(e.isA)
+  }
+  snapshot.vault = {
+    types: [...types].sort(),
+    totalNotes: entries.length,
+  }
+
+  if (references && references.length > 0) {
+    snapshot.referencedNotes = references
+      .filter(ref => allContent[ref.path] !== undefined)
+      .map(ref => ({
+        path: ref.path,
+        title: ref.title,
+        type: ref.type ?? 'Note',
+        body: allContent[ref.path] ?? '',
+      }))
+  }
+
+  const preamble = [
+    'You are an AI assistant integrated into Laputa, a personal knowledge management app.',
+    'The user is viewing a specific note. Use the structured context below to answer questions accurately.',
+    'You can also use MCP tools to search, read, create, or edit notes in the vault.',
+  ].join('\n')
+
+  return `${preamble}\n\n## Context Snapshot\n\`\`\`json\n${JSON.stringify(snapshot, null, 2)}\n\`\`\``
+}
+
+/** Legacy: Build a contextual system prompt (text-based). */
 export function buildContextualPrompt(
   active: VaultEntry,
   linkedEntries: VaultEntry[],