Add context pruning to one-shot paths; fix TUI word-wrap

ober

e73f4073761c6f787b361d15b32a4feecb3d52ea

diff --git a/lib/jcode/core/agent.sls b/lib/jcode/core/agent.sls
index 017ad64..47ab5a3 100644
--- a/lib/jcode/core/agent.sls
+++ b/lib/jcode/core/agent.sls
@@ -25,6 +25,54 @@
   (def current-tool-cb (make-parameter #f))
   (def current-usage-cb (make-parameter #f))
   (def *max-tool-rounds* 8)
+  (def *prune-protect-chars* 16000)
+  (def *tool-result-stub* "[Old tool result cleared]")
+  (def (trim-messages messages)
+       "Prune old tool results: walk backwards, protect recent ones, stub the rest."
+       (let* ([reversed (reverse messages)]
+              [tool-chars 0]
+              [pruned 0]
+              [total-before (apply
+                              +
+                              (map (lambda (m)
+                                     (string-length
+                                       (json-object->string
+                                         (message->json m))))
+                                   messages))])
+         (let ([result (reverse
+                         (map (lambda (m)
+                                (if (and (equal? (message-role m) "tool")
+                                         (message-content m))
+                                    (let ([size (string-length
+                                                  (message-content m))])
+                                      (set! tool-chars (+ tool-chars size))
+                                      (if (> tool-chars
+                                             *prune-protect-chars*)
+                                          (begin
+                                            (set! pruned (+ pruned 1))
+                                            (make-tool-result
+                                              (or (message-tool-call-id m)
+                                                  "")
+                                              *tool-result-stub*))
+                                          m))
+                                    m))
+                              reversed))])
+           (let ([total-after (apply
+                                +
+                                (map (lambda (m)
+                                       (string-length
+                                         (json-object->string
+                                           (message->json m))))
+                                     result))])
+             (log-info
+               logger
+               "trim-messages"
+               `((msgs . ,(length messages))
+                  (before-bytes . ,total-before)
+                  (after-bytes . ,total-after)
+                  (tool-chars . ,tool-chars)
+                  (pruned . ,pruned))))
+           result)))
   (def (agent-run session-id user-input)
        (log-info logger "agent-run" `((session . ,session-id)))
        (let ([existing (session-get-messages session-id)])
@@ -47,7 +95,8 @@
   (def (agent-loop session-id messages round)
        (let* ([provider (get-current-provider)]
               [tools (get-tool-schemas)]
-              [response (provider-chat provider messages tools)])
+              [msgs (trim-messages messages)]
+              [response (provider-chat provider msgs tools)])
          (log-debug
            logger
            "got-response"
@@ -57,7 +106,16 @@
            [(not (message-tool-calls response)) response]
            [(>= round *max-tool-rounds*)
             (log-warn logger "max-rounds" `((round . ,round)))
-            response]
+            (let ([results (execute-tool-calls
+                             (message-tool-calls response))])
+              (for-each
+                (lambda (r) (session-add-message session-id r))
+                results)
+              (let* ([final-msgs (trim-messages
+                                   (session-get-messages session-id))]
+                     [final (provider-chat provider final-msgs '())])
+                (session-add-message session-id final)
+                final))]
            [else
             (let ([results (execute-tool-calls
                              (message-tool-calls response))])
@@ -66,15 +124,16 @@
                 results)
               (agent-loop
                 session-id
-                (session-get-messages session-id)
+                (trim-messages (session-get-messages session-id))
                 (+ round 1)))])))
   (def (agent-loop-stream session-id messages round)
        (let* ([provider (get-current-provider)]
-              [tools (get-tool-schemas)])
+              [tools (get-tool-schemas)]
+              [msgs (trim-messages messages)])
          (let-values ([(content tool-calls usage)
                        (provider-stream-chat
                          provider
-                         messages
+                         msgs
                          tools
                          (current-stream-cb))])
            (when (and usage (current-usage-cb))
@@ -87,7 +146,22 @@
                [(null? tool-calls) response]
                [(>= round *max-tool-rounds*)
                 (log-warn logger "max-rounds" `((round . ,round)))
-                response]
+                (let ([results (execute-tool-calls tool-calls)])
+                  (for-each
+                    (lambda (r) (session-add-message session-id r))
+                    results)
+                  (let-values ([(fc _tc _u)
+                                (provider-stream-chat
+                                  provider
+                                  (trim-messages
+                                    (session-get-messages session-id))
+                                  '()
+                                  (current-stream-cb))])
+                    (let ([final (make-assistant-message
+                                   (if (string=? fc "") #f fc)
+                                   #f)])
+                      (session-add-message session-id final)
+                      final)))]
                [else
                 (let ([results (execute-tool-calls tool-calls)])
                   (for-each
@@ -95,7 +169,7 @@
                     results)
                   (agent-loop-stream
                     session-id
-                    (session-get-messages session-id)
+                    (trim-messages (session-get-messages session-id))
                     (+ round 1)))])))))
   (def (execute-tool-calls tool-calls)
        (log-info
@@ -139,45 +213,62 @@
              (agent-chat-loop-stream provider messages tools 0)
              (agent-chat-loop provider messages tools 0))))
   (def (agent-chat-loop provider messages tools round)
-       (let ([response (provider-chat provider messages tools)])
+       (let* ([msgs (trim-messages messages)]
+              [response (provider-chat provider msgs tools)])
          (cond
            [(not (message-tool-calls response))
             (message-content response)]
            [(>= round *max-tool-rounds*)
-            (or (message-content response) "")]
+            (let* ([results (execute-tool-calls
+                              (message-tool-calls response))]
+                   [new-messages (trim-messages
+                                   (append msgs (list response) results))]
+                   [final (provider-chat provider new-messages '())])
+              (or (message-content final) ""))]
            [else
             (let* ([results (execute-tool-calls
                               (message-tool-calls response))]
-                   [new-messages (append
-                                   messages
-                                   (list response)
-                                   results)])
+                   [new-messages (append msgs (list response) results)])
               (agent-chat-loop
                 provider
                 new-messages
                 tools
                 (+ round 1)))])))
   (def (agent-chat-loop-stream provider messages tools round)
-       (let-values ([(content tool-calls usage)
-                     (provider-stream-chat
-                       provider
-                       messages
-                       tools
-                       (current-stream-cb))])
-         (cond
-           [(null? tool-calls) content]
-           [(>= round *max-tool-rounds*) content]
-           [else
-            (let* ([response (make-assistant-message
-                               (if (string=? content "") #f content)
-                               tool-calls)]
-                   [results (execute-tool-calls tool-calls)]
-                   [new-msgs (append messages (list response) results)])
-              (agent-chat-loop-stream
-                provider
-                new-msgs
-                tools
-                (+ round 1)))])))
+       (let* ([msgs (trim-messages messages)])
+         (let-values ([(content tool-calls usage)
+                       (provider-stream-chat
+                         provider
+                         msgs
+                         tools
+                         (current-stream-cb))])
+           (cond
+             [(null? tool-calls) content]
+             [(>= round *max-tool-rounds*)
+              (let* ([response (make-assistant-message
+                                 (if (string=? content "") #f content)
+                                 tool-calls)]
+                     [results (execute-tool-calls tool-calls)]
+                     [new-msgs (trim-messages
+                                 (append msgs (list response) results))])
+                (let-values ([(fc _tc _u)
+                              (provider-stream-chat
+                                provider
+                                new-msgs
+                                '()
+                                (current-stream-cb))])
+                  fc))]
+             [else
+              (let* ([response (make-assistant-message
+                                 (if (string=? content "") #f content)
+                                 tool-calls)]
+                     [results (execute-tool-calls tool-calls)]
+                     [new-msgs (append msgs (list response) results)])
+                (agent-chat-loop-stream
+                  provider
+                  new-msgs
+                  tools
+                  (+ round 1)))]))))
   (def (agent-step messages)
        (let* ([provider (get-current-provider)]
               [tools (get-tool-schemas)])
diff --git a/lib/jcode/ui/tui-message.sls b/lib/jcode/ui/tui-message.sls
index 07ba06d..fd8fcbb 100644
--- a/lib/jcode/ui/tui-message.sls
+++ b/lib/jcode/ui/tui-message.sls
@@ -47,14 +47,59 @@
        (let ([m (make-msg-block 'system text '() 0 #f #f #f '())])
          (reflow-message! m 80)
          m))
+  (def (wrap-segment-lines lines width)
+       "Wrap each line of segments to fit within width columns."
+       (apply
+         append
+         (map (lambda (segs) (wrap-one-line segs width)) lines)))
+  (def (wrap-one-line segs width)
+       "Wrap a single line of segments into multiple lines if needed."
+       (let loop ([segs segs] [cur-line '()] [col 0] [out '()])
+         (cond
+           [(null? segs)
+            (reverse
+              (if (null? cur-line) out (cons (reverse cur-line) out)))]
+           [else
+            (let* ([seg (car segs)]
+                   [text (car seg)]
+                   [face (cdr seg)]
+                   [len (string-length text)])
+              (if (<= (+ col len) width)
+                  (loop (cdr segs) (cons seg cur-line) (+ col len) out)
+                  (let break ([pos 0]
+                              [cur-line cur-line]
+                              [col col]
+                              [out out])
+                    (let ([remaining (- len pos)] [avail (- width col)])
+                      (cond
+                        [(<= remaining 0)
+                         (loop (cdr segs) cur-line col out)]
+                        [(<= remaining avail)
+                         (loop
+                           (cdr segs)
+                           (cons
+                             (cons (substring text pos len) face)
+                             cur-line)
+                           (+ col remaining)
+                           out)]
+                        [else
+                         (let ([chunk (substring text pos (+ pos avail))])
+                           (break
+                             (+ pos avail)
+                             '()
+                             0
+                             (cons
+                               (reverse (cons (cons chunk face) cur-line))
+                               out)))])))))])))
   (def (reflow-message! msg width)
        "Re-render message content into lines for given terminal width."
        (let ([lines (render-content (msg-block-role msg) (msg-block-content msg)
                       (msg-block-tool-name msg) (msg-block-tool-status msg)
                       (msg-block-collapsed? msg) (msg-block-metadata msg)
                       width)])
-         (msg-block-lines-set! msg lines)
-         (msg-block-height-set! msg (length lines))))
+         (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)
        (case role
diff --git a/src/jcode/core/agent.ss b/src/jcode/core/agent.ss
index 559645d..adbb8bb 100644
--- a/src/jcode/core/agent.ss
+++ b/src/jcode/core/agent.ss
@@ -48,6 +48,45 @@ Prefer using the edit tool over write for modifying existing files." (current-di
 (def current-usage-cb (make-parameter #f))
 
 (def *max-tool-rounds* 8)
+(def *prune-protect-chars* 16000)  ;; ~4k tokens of recent tool results to keep
+(def *tool-result-stub* "[Old tool result cleared]")
+
+(def (trim-messages messages)
+  "Prune old tool results: walk backwards, protect recent ones, stub the rest."
+  (let* ((reversed (reverse messages))
+         (tool-chars 0)
+         (pruned 0)
+         (total-before (apply + (map (lambda (m)
+                                       (string-length
+                                         (json-object->string (message->json m))))
+                                     messages))))
+    (let ((result
+            (reverse
+              (map (lambda (m)
+                     (if (and (equal? (message-role m) "tool")
+                              (message-content m))
+                       (let ((size (string-length (message-content m))))
+                         (set! tool-chars (+ tool-chars size))
+                         (if (> tool-chars *prune-protect-chars*)
+                           (begin
+                             (set! pruned (+ pruned 1))
+                             (make-tool-result
+                               (or (message-tool-call-id m) "")
+                               *tool-result-stub*))
+                           m))
+                       m))
+                   reversed))))
+      (let ((total-after (apply + (map (lambda (m)
+                                         (string-length
+                                           (json-object->string (message->json m))))
+                                       result))))
+        (log-info logger "trim-messages"
+          `((msgs . ,(length messages))
+            (before-bytes . ,total-before)
+            (after-bytes . ,total-after)
+            (tool-chars . ,tool-chars)
+            (pruned . ,pruned))))
+      result)))
 
 (def (agent-run session-id user-input)
   (log-info logger "agent-run" `((session . ,session-id)))
@@ -62,27 +101,34 @@ Prefer using the edit tool over write for modifying existing files." (current-di
 (def (agent-loop session-id messages round)
   (let* ((provider (get-current-provider))
          (tools (get-tool-schemas))
-         (response (provider-chat provider messages tools)))
+         (msgs (trim-messages messages))
+         (response (provider-chat provider msgs tools)))
     (log-debug logger "got-response" `((role . ,(message-role response))))
     (session-add-message session-id response)
     (cond
       ((not (message-tool-calls response)) response)
       ((>= round *max-tool-rounds*)
        (log-warn logger "max-rounds" `((round . ,round)))
-       response)
+       (let ((results (execute-tool-calls (message-tool-calls response))))
+         (for-each (lambda (r) (session-add-message session-id r)) results)
+         (let* ((final-msgs (trim-messages (session-get-messages session-id)))
+                (final (provider-chat provider final-msgs '())))
+           (session-add-message session-id final)
+           final)))
       (else
        (let ((results (execute-tool-calls (message-tool-calls response))))
          (for-each
            (lambda (result) (session-add-message session-id result))
            results)
-         (agent-loop session-id (session-get-messages session-id) (+ round 1)))))))
+         (agent-loop session-id (trim-messages (session-get-messages session-id)) (+ round 1)))))))
 
 (def (agent-loop-stream session-id messages round)
   ;; Streaming version: calls (current-stream-cb) for each text token.
   (let* ((provider (get-current-provider))
-         (tools    (get-tool-schemas)))
+         (tools    (get-tool-schemas))
+         (msgs     (trim-messages messages)))
     (let-values (((content tool-calls usage)
-                  (provider-stream-chat provider messages tools (current-stream-cb))))
+                  (provider-stream-chat provider msgs tools (current-stream-cb))))
       (when (and usage (current-usage-cb))
         ((current-usage-cb) usage))
       (let ((response (make-assistant-message
@@ -93,11 +139,20 @@ Prefer using the edit tool over write for modifying existing files." (current-di
           ((null? tool-calls) response)
           ((>= round *max-tool-rounds*)
            (log-warn logger "max-rounds" `((round . ,round)))
-           response)
+           (let ((results (execute-tool-calls tool-calls)))
+             (for-each (lambda (r) (session-add-message session-id r)) results)
+             (let-values (((fc _tc _u)
+                           (provider-stream-chat provider
+                             (trim-messages (session-get-messages session-id))
+                             '() (current-stream-cb))))
+               (let ((final (make-assistant-message
+                              (if (string=? fc "") #f fc) #f)))
+                 (session-add-message session-id final)
+                 final))))
           (else
            (let ((results (execute-tool-calls tool-calls)))
              (for-each (lambda (r) (session-add-message session-id r)) results)
-             (agent-loop-stream session-id (session-get-messages session-id) (+ round 1)))))))))
+             (agent-loop-stream session-id (trim-messages (session-get-messages session-id)) (+ round 1)))))))))
 
 (def (execute-tool-calls tool-calls)
   (log-info logger "executing-tools" `((count . ,(length tool-calls))))
@@ -138,28 +193,42 @@ Prefer using the edit tool over write for modifying existing files." (current-di
       (agent-chat-loop provider messages tools 0))))
 
 (def (agent-chat-loop provider messages tools round)
-  (let ((response (provider-chat provider messages tools)))
+  (let* ((msgs (trim-messages messages))
+         (response (provider-chat provider msgs tools)))
     (cond
       ((not (message-tool-calls response)) (message-content response))
-      ((>= round *max-tool-rounds*) (or (message-content response) ""))
+      ((>= round *max-tool-rounds*)
+       (let* ((results (execute-tool-calls (message-tool-calls response)))
+              (new-messages (trim-messages (append msgs (list response) results)))
+              (final (provider-chat provider new-messages '())))
+         (or (message-content final) "")))
       (else
        (let* ((results (execute-tool-calls (message-tool-calls response)))
-              (new-messages (append messages (list response) results)))
+              (new-messages (append msgs (list response) results)))
          (agent-chat-loop provider new-messages tools (+ round 1)))))))
 
 (def (agent-chat-loop-stream provider messages tools round)
-  (let-values (((content tool-calls usage)
-                (provider-stream-chat provider messages tools (current-stream-cb))))
-    (cond
-      ((null? tool-calls) content)
-      ((>= round *max-tool-rounds*) content)
-      (else
-       (let* ((response  (make-assistant-message
-                           (if (string=? content "") #f content)
-                           tool-calls))
-              (results   (execute-tool-calls tool-calls))
-              (new-msgs  (append messages (list response) results)))
-         (agent-chat-loop-stream provider new-msgs tools (+ round 1)))))))
+  (let* ((msgs (trim-messages messages)))
+    (let-values (((content tool-calls usage)
+                  (provider-stream-chat provider msgs tools (current-stream-cb))))
+      (cond
+        ((null? tool-calls) content)
+        ((>= round *max-tool-rounds*)
+         (let* ((response  (make-assistant-message
+                             (if (string=? content "") #f content)
+                             tool-calls))
+                (results   (execute-tool-calls tool-calls))
+                (new-msgs  (trim-messages (append msgs (list response) results))))
+           (let-values (((fc _tc _u)
+                         (provider-stream-chat provider new-msgs '() (current-stream-cb))))
+             fc)))
+        (else
+         (let* ((response  (make-assistant-message
+                             (if (string=? content "") #f content)
+                             tool-calls))
+                (results   (execute-tool-calls tool-calls))
+                (new-msgs  (append msgs (list response) results)))
+           (agent-chat-loop-stream provider new-msgs tools (+ round 1))))))))
 
 (def (agent-step messages)
   (let* ((provider (get-current-provider))
diff --git a/src/jcode/ui/tui-message.ss b/src/jcode/ui/tui-message.ss
index c1dfa18..0ad8992 100644
--- a/src/jcode/ui/tui-message.ss
+++ b/src/jcode/ui/tui-message.ss
@@ -59,6 +59,45 @@
     (reflow-message! m 80)
     m))
 
+;; ---- Word-wrap segments to fit terminal width ----
+
+(def (wrap-segment-lines lines width)
+  "Wrap each line of segments to fit within width columns."
+  (apply append (map (lambda (segs) (wrap-one-line segs width)) lines)))
+
+(def (wrap-one-line segs width)
+  "Wrap a single line of segments into multiple lines if needed."
+  (let loop ((segs segs) (cur-line '()) (col 0) (out '()))
+    (cond
+      ((null? segs)
+       (reverse (if (null? cur-line) out (cons (reverse cur-line) out))))
+      (else
+       (let* ((seg (car segs))
+              (text (car seg))
+              (face (cdr seg))
+              (len (string-length text)))
+         (if (<= (+ col len) width)
+           ;; Fits on current line
+           (loop (cdr segs) (cons seg cur-line) (+ col len) out)
+           ;; Need to break this segment
+           (let break ((pos 0) (cur-line cur-line) (col col) (out out))
+             (let ((remaining (- len pos))
+                   (avail (- width col)))
+               (cond
+                 ((<= remaining 0)
+                  (loop (cdr segs) cur-line col out))
+                 ((<= remaining avail)
+                  (loop (cdr segs)
+                        (cons (cons (substring text pos len) face) cur-line)
+                        (+ col remaining) out))
+                 (else
+                  ;; Fill current line, start new one
+                  (let ((chunk (substring text pos (+ pos avail))))
+                    (break (+ pos avail)
+                           '()
+                           0
+                           (cons (reverse (cons (cons chunk face) cur-line)) out)))))))))))))
+
 ;; ---- Reflow: re-render message for given width ----
 
 (def (reflow-message! msg width)
@@ -70,8 +109,9 @@
                                (msg-block-collapsed? msg)
                                (msg-block-metadata msg)
                                width)))
-    (msg-block-lines-set! msg lines)
-    (msg-block-height-set! msg (length lines))))
+    (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)
   (case role