tui: collapsible <think> block for reasoning models

ober

343225e50bcd53c05a9c19e53ae4480a7588b9fb

diff --git a/src/jcode/core/message.ss b/src/jcode/core/message.ss
index dc526a1..7ae86aa 100644
--- a/src/jcode/core/message.ss
+++ b/src/jcode/core/message.ss
@@ -12,6 +12,9 @@
         message-content
         message-tool-calls
         message-tool-call-id
+        message-thinking
+        message-thinking-set!
+        extract-thinking
         message->json
         json->message
         tool-call-id
@@ -20,9 +23,10 @@
         tool-call?)
 
 (import :std/text/json
-        :std/misc/uuid)
+        :std/misc/uuid
+        :std/misc/string)
 
-(defstruct message (role content tool-calls tool-call-id))
+(defstruct message (role content tool-calls tool-call-id thinking))
 
 ;; Private struct — public API uses tool-call-* wrappers below
 (defstruct tool-call-data (id name arguments))
@@ -41,28 +45,44 @@
   (make-tool-call-data id name arguments))
 
 (def (make-user-message content)
-  (make-message "user" content #f #f))
+  (make-message "user" content #f #f #f))
 
 (def (make-assistant-message content . tool-calls)
   (make-message "assistant" content
                 (if (null? tool-calls) #f (car tool-calls))
+                #f
                 #f))
 
 (def (make-tool-result call-id content)
-  (make-message "tool" content #f call-id))
+  (make-message "tool" content #f call-id #f))
 
 (def (make-system-message content)
-  (make-message "system" content #f #f))
+  (make-message "system" content #f #f #f))
 
 ;; JSON null is (void) in this library — normalize to #f
 (def (json-null->false v)
   (if (eq? v (void)) #f v))
 
+(def (extract-thinking text)
+  ;; Pull <think>...</think> out of TEXT. Returns (cons thinking-or-#f reply).
+  ;; Only handles a single complete pair; partial tags are left untouched.
+  (let ((s-pos (string-contains text "<think>"))
+        (e-pos (string-contains text "</think>")))
+    (if (and s-pos e-pos (< s-pos e-pos))
+      (let* ((think (substring text (+ s-pos 7) e-pos))
+             (before (substring text 0 s-pos))
+             (after (substring text (+ e-pos 8) (string-length text)))
+             (reply (string-trim (string-append before after))))
+        (cons (string-trim think) reply))
+      (cons #f text))))
+
 (def (message->json msg)
   (let ((ht (make-hash-table)))
     (hash-put! ht "role" (message-role msg))
     (when (message-content msg)
-      (hash-put! ht "content" (message-content msg)))
+      ;; Strip <think>...</think> blocks before sending back to the model —
+      ;; reasoning is for display only, the LLM doesn't need its own thoughts echoed.
+      (hash-put! ht "content" (cdr (extract-thinking (message-content msg)))))
     (when (message-tool-calls msg)
       (hash-put! ht "tool_calls"
         (map tool-call->json (message-tool-calls msg))))
@@ -92,7 +112,8 @@
       role
       content
       (and tool-calls-json (pair? tool-calls-json) (map json->tool-call tool-calls-json))
-      tc-id)))
+      tc-id
+      #f)))
 
 (def (json->tool-call json)
   (let ((fn (hash-ref json "function" #f)))
diff --git a/src/jcode/ui/tui-message.ss b/src/jcode/ui/tui-message.ss
index ccb6f9f..03475d7 100644
--- a/src/jcode/ui/tui-message.ss
+++ b/src/jcode/ui/tui-message.ss
@@ -5,9 +5,11 @@
   make-msg-block msg-block? msg-block-role msg-block-content
   msg-block-lines msg-block-height msg-block-tool-name
   msg-block-tool-status msg-block-collapsed? msg-block-metadata
+  msg-block-thinking msg-block-thinking-collapsed?
   msg-block-lines-set! msg-block-height-set!
   msg-block-content-set! msg-block-tool-status-set!
   msg-block-collapsed?-set!
+  msg-block-thinking-set! msg-block-thinking-collapsed?-set!
   reflow-message!
   render-msg-block!
   msg-block-user msg-block-assistant msg-block-tool msg-block-error msg-block-system)
@@ -23,41 +25,44 @@
 ;; ---- Message block struct ----
 
 (defstruct msg-block
-  (role           ;; 'user | 'assistant | 'tool | 'error | 'system
-   content        ;; raw text
-   lines          ;; list of (list-of (text . face-name))  — rendered segments per line
-   height         ;; number of terminal rows
-   tool-name      ;; string or #f
-   tool-status    ;; 'running | 'done | 'error | #f
-   collapsed?     ;; #t = show summary only
-   metadata)      ;; alist
+  (role                 ;; 'user | 'assistant | 'tool | 'error | 'system
+   content              ;; raw text (assistant: visible reply only, with <think> stripped)
+   lines                ;; list of (list-of (text . face-name))  — rendered segments per line
+   height               ;; number of terminal rows
+   tool-name            ;; string or #f
+   tool-status          ;; 'running | 'done | 'error | #f
+   collapsed?           ;; #t = show summary only (tools)
+   metadata             ;; alist
+   thinking             ;; string or #f — extracted <think> content (assistant only)
+   thinking-collapsed?) ;; #t = show "Reasoning" header only; #f = expand thinking text
   transparent: #t)
 
 ;; ---- Constructors ----
 
 (def (msg-block-user text)
-  (let ((m (make-msg-block 'user text '() 0 #f #f #f '())))
+  (let ((m (make-msg-block 'user text '() 0 #f #f #f '() #f #t)))
     (reflow-message! m 80)
     m))
 
 (def (msg-block-assistant text)
-  (let ((m (make-msg-block 'assistant text '() 0 #f #f #f '())))
+  (let ((m (make-msg-block 'assistant text '() 0 #f #f #f '() #f #t)))
     (reflow-message! m 80)
     m))
 
 (def (msg-block-tool name status content . meta)
   (let ((m (make-msg-block 'tool content '() 0 name status #t
-             (if (null? meta) '() (car meta)))))
+             (if (null? meta) '() (car meta))
+             #f #t)))
     (reflow-message! m 80)
     m))
 
 (def (msg-block-error text)
-  (let ((m (make-msg-block 'error text '() 0 #f #f #f '())))
+  (let ((m (make-msg-block 'error text '() 0 #f #f #f '() #f #t)))
     (reflow-message! m 80)
     m))
 
 (def (msg-block-system text)
-  (let ((m (make-msg-block 'system text '() 0 #f #f #f '())))
+  (let ((m (make-msg-block 'system text '() 0 #f #f #f '() #f #t)))
     (reflow-message! m 80)
     m))
 
@@ -121,15 +126,17 @@
                                (msg-block-tool-status msg)
                                (msg-block-collapsed? msg)
                                (msg-block-metadata msg)
+                               (msg-block-thinking msg)
+                               (msg-block-thinking-collapsed? msg)
                                width)))
     (let ((wrapped (wrap-segment-lines lines width)))
       (msg-block-lines-set! msg wrapped)
       (msg-block-height-set! msg (length wrapped)))))
 
-(def (render-content role content tool-name tool-status collapsed? metadata width)
+(def (render-content role content tool-name tool-status collapsed? metadata thinking thinking-collapsed? width)
   (case role
     ((user)      (render-user-content content))
-    ((assistant) (render-assistant-content content))
+    ((assistant) (render-assistant-content content thinking thinking-collapsed? width))
     ((tool)      (render-tool-content tool-name tool-status content collapsed? metadata width))
     ((error)     (render-error-content content))
     ((system)    (render-system-content content))
@@ -148,10 +155,32 @@
                (list (cons (string-append "      " line) 'user-text)))
              (cdr lines))))))
 
-(def (render-assistant-content content)
-  (if (string-empty? content)
-    '()
-    (md-render-lines (string-split content #\newline))))
+(def (render-assistant-content content thinking thinking-collapsed? width)
+  (let* ((has-thinking? (and thinking (not (string-empty? thinking))))
+         (thinking-lines (if has-thinking?
+                           (render-thinking-block thinking thinking-collapsed? width)
+                           '()))
+         (reply-lines (if (string-empty? content)
+                        '()
+                        (md-render-lines (string-split content #\newline)))))
+    (append thinking-lines reply-lines)))
+
+(def (render-thinking-block thinking collapsed? width)
+  (let* ((chevron (if collapsed? "▸" "▾"))
+         (n (string-length thinking))
+         (header (list (cons (string-append chevron " Reasoning ("
+                                            (number->string n)
+                                            " chars, press t to toggle)")
+                             'thinking-header))))
+    (if collapsed?
+      (list header)
+      (let* ((bar (make-string (min (max 1 width) 40) #\─))
+             (border-row (list (cons bar 'thinking-border)))
+             (body (map (lambda (line) (list (cons line 'thinking-text)))
+                        (string-split thinking #\newline))))
+        (append (list header border-row)
+                body
+                (list border-row))))))
 
 (def (render-tool-content tool-name tool-status content collapsed? metadata width)
   (let* ((icon (case tool-status
diff --git a/src/jcode/ui/tui.ss b/src/jcode/ui/tui.ss
index 03b62bf..c4bf153 100644
--- a/src/jcode/ui/tui.ss
+++ b/src/jcode/ui/tui.ss
@@ -369,6 +369,11 @@
        (toast-add! 'success "Theme" "Theme cycled")
        (app-state-dirty?-set! state #t))
 
+      ;; Global: Ctrl-R toggle Reasoning (<think>) on latest assistant block
+      ((= key TB_KEY_CTRL_R)
+       (tui-log "  -> toggle-reasoning")
+       (toggle-last-thinking! state))
+
       ;; Global: Shift-Tab (BACK_TAB) toggle PLAN/BUILD mode (opencode-style)
       ((= key TB_KEY_BACK_TAB)
        (tui-log "  -> toggle-mode (back-tab)")
@@ -847,6 +852,7 @@
      (app-state-dirty?-set! state #t))
     ((list 'agent-done)
      (tui-log "apply-agent-event: agent-done")
+     (finalize-last-assistant! state)
      (app-state-agent-busy?-set! state #f)
      (app-state-scroll-offset-set! state 0)
      (app-state-dirty?-set! state #t))
@@ -1168,6 +1174,39 @@
     ((assistant) (make-assistant-message (msg-block-content blk)))
     (else        (make-system-message    (msg-block-content blk)))))
 
+(def (finalize-last-assistant! state)
+  ;; If the trailing assistant block contains a <think>...</think> pair, hoist
+  ;; the thinking text into its own slot and replace content with just the reply.
+  (let loop ((msgs (reverse (app-state-messages state))))
+    (when (pair? msgs)
+      (let ((m (car msgs)))
+        (cond
+          ((not (eq? (msg-block-role m) 'assistant))
+           (loop (cdr msgs)))
+          ((msg-block-thinking m) (void)) ;; already finalized
+          (else
+           (let* ((parsed (extract-thinking (msg-block-content m)))
+                  (think (car parsed))
+                  (reply (cdr parsed)))
+             (when think
+               (msg-block-thinking-set! m think)
+               (msg-block-content-set! m reply)
+               (reflow-message! m (msg-area-width state))))))))))
+
+(def (toggle-last-thinking! state)
+  ;; Flip the collapsed state of the latest assistant block that has thinking.
+  (let loop ((msgs (reverse (app-state-messages state))))
+    (when (pair? msgs)
+      (let ((m (car msgs)))
+        (cond
+          ((and (eq? (msg-block-role m) 'assistant)
+                (msg-block-thinking m))
+           (msg-block-thinking-collapsed?-set! m
+             (not (msg-block-thinking-collapsed? m)))
+           (reflow-message! m (msg-area-width state))
+           (app-state-dirty?-set! state #t))
+          (else (loop (cdr msgs))))))))
+
 (def (tui-stream-token! state token)
   "Handle a streaming token from the LLM — called from agent thread."
   (let ((buf (app-state-stream-buf state)))
@@ -1194,6 +1233,7 @@
   (case event
     ((start)
      ;; Reset stream buffer so next round creates a fresh assistant block
+     (finalize-last-assistant! state)
      (app-state-stream-buf-set! state "")
      ;; Track tool counts in sidebar only — no inline message block
      (let ((tc (app-state-tool-counts state)))