TUI: ESC during pre-stream thinking gives immediate feedback

ober

1e8655c5571c19a738cca4ef3a24a5f58ed0d750

diff --git a/lib/jcode/ui/tui.sls b/lib/jcode/ui/tui.sls
index ee6992a..03a00d5 100644
--- a/lib/jcode/ui/tui.sls
+++ b/lib/jcode/ui/tui.sls
@@ -256,6 +256,9 @@
                  (app-state-agent-busy? state))
             (tui-log "  -> cancel-agent (key=~a)" key)
             (set-car! *tui-stream-abort* #t)
+            (bump-tui-run-gen!)
+            (app-state-agent-busy?-set! state #f)
+            (add-message! state (msg-block-system "(interrupted)"))
             (app-state-dirty?-set! state #t)]
            [#t
             (if (app-state-agent-busy? state)
@@ -655,9 +658,15 @@
        (let ([servers (try (mcp-active-servers) (catch (e) '()))])
          (sidebar-state-mcp-set! (app-state-sidebar state) servers)))
   (def *tui-stream-abort* (cons #f #f))
+  (def *tui-run-gen* (cons 0 #f))
+  (def (tui-run-gen) (car *tui-run-gen*))
+  (def (bump-tui-run-gen!)
+       (set-car! *tui-run-gen* (+ 1 (car *tui-run-gen*))))
   (def (send-agent-event! ev)
        (let ([main (*main-thread*)])
          (when main (thread-send main ev))))
+  (def (send-worker-event! gen ev)
+       (send-agent-event! (cons gen ev)))
   (def (drain-agent-events! state)
        "Drain all pending agent events from the main thread's mailbox.\n   Apply them to state synchronously. Returns the number drained."
        (let loop ([n 0])
@@ -666,6 +675,15 @@
                n
                (begin (apply-agent-event! state ev) (loop (+ n 1)))))))
   (def (apply-agent-event! state ev)
+       (when (and (pair? ev) (number? (car ev)))
+         (let ([gen (car ev)] [body (cdr ev)])
+           (if (= gen (tui-run-gen))
+               (apply-agent-event-body! state body)
+               (tui-log
+                 "apply-agent-event: discarded stale gen=~a (current=~a)"
+                 gen
+                 (tui-run-gen))))))
+  (def (apply-agent-event-body! state ev)
        (match ev
          [(list 'stream-token token) (tui-stream-token! state token)]
          [(list 'tool-event op name args)
@@ -690,7 +708,7 @@
   (def (run-agent! state text)
        (app-state-agent-busy?-set! state #t)
        (app-state-stream-buf-set! state "")
-       (set-car! *tui-stream-abort* #f)
+       (set-car! *tui-stream-abort* #f) (bump-tui-run-gen!)
        (add-message! state (msg-block-assistant ""))
        (app-state-dirty?-set! state #t)
        (let ([p-name (or (current-provider-override)
@@ -703,47 +721,68 @@
            "run-agent: spawning worker thread, session=~a provider=~a model=~a key?=~a"
            s-id p-name m-name
            (and (config-get-provider-key p-name) #t))
-         (spawn
-           (lambda ()
-             (tui-log "worker: entered, main-thread=~a" (*main-thread*))
-             (try (parameterize ([current-provider-override p-override]
-                                 [current-model-override m-override]
-                                 [current-stream-cb
-                                  (lambda (token)
-                                    (when (car *tui-stream-abort*)
-                                      (error 'stream-aborted
-                                        "interrupted by user"))
-                                    (send-agent-event!
-                                      (list 'stream-token token)))]
-                                 [current-tool-cb
-                                  (lambda (event name args)
-                                    (tui-log
-                                      "worker: tool-cb ~a ~a"
-                                      event
-                                      name)
-                                    (send-agent-event!
-                                      (list 'tool-event event name args)))]
-                                 [current-usage-cb
-                                  (lambda (usage)
-                                    (send-agent-event!
-                                      (list 'usage-update usage)))])
-                    (tui-log "worker: calling agent-run")
-                    (agent-run s-id text)
-                    (tui-log
-                      "worker: agent-run returned, sending agent-done")
-                    (send-agent-event! (list 'agent-done))
-                    (tui-log
-                      "worker: agent-done sent, worker exiting normally"))
-                  (catch
-                    (e)
-                    (let ([msg (err->string e)])
-                      (tui-log "worker: CAUGHT exception: ~a" msg)
-                      (cond
-                        [(string-contains msg "stream-aborted")
-                         (send-agent-event! (list 'agent-cancelled))]
-                        [else
-                         (send-agent-event!
-                           (list 'agent-error msg))]))))))))
+         (let ([gen (tui-run-gen)])
+           (spawn
+             (lambda ()
+               (tui-log
+                 "worker: entered gen=~a main-thread=~a"
+                 gen
+                 (*main-thread*))
+               (try (parameterize ([current-provider-override p-override]
+                                   [current-model-override m-override]
+                                   [current-stream-cb
+                                    (lambda (token)
+                                      (when (or (car *tui-stream-abort*)
+                                                (not (= gen
+                                                        (tui-run-gen))))
+                                        (error 'stream-aborted
+                                          "interrupted by user"))
+                                      (send-worker-event!
+                                        gen
+                                        (list 'stream-token token)))]
+                                   [current-tool-cb
+                                    (lambda (event name args)
+                                      (when (or (car *tui-stream-abort*)
+                                                (not (= gen
+                                                        (tui-run-gen))))
+                                        (error 'stream-aborted
+                                          "interrupted by user"))
+                                      (tui-log
+                                        "worker: tool-cb ~a ~a"
+                                        event
+                                        name)
+                                      (send-worker-event!
+                                        gen
+                                        (list
+                                          'tool-event
+                                          event
+                                          name
+                                          args)))]
+                                   [current-usage-cb
+                                    (lambda (usage)
+                                      (send-worker-event!
+                                        gen
+                                        (list 'usage-update usage)))])
+                      (tui-log "worker: calling agent-run")
+                      (agent-run s-id text)
+                      (tui-log
+                        "worker: agent-run returned, sending agent-done")
+                      (send-worker-event! gen (list 'agent-done))
+                      (tui-log
+                        "worker: agent-done sent, worker exiting normally"))
+                    (catch
+                      (e)
+                      (let ([msg (err->string e)])
+                        (tui-log "worker: CAUGHT exception: ~a" msg)
+                        (cond
+                          [(string-contains msg "stream-aborted")
+                           (send-worker-event!
+                             gen
+                             (list 'agent-cancelled))]
+                          [else
+                           (send-worker-event!
+                             gen
+                             (list 'agent-error msg))])))))))))
   (def (tui-stream-token! state token)
        "Handle a streaming token from the LLM — called from agent thread."
        (let ([buf (app-state-stream-buf state)])
diff --git a/src/jcode/ui/tui.ss b/src/jcode/ui/tui.ss
index 49b40e5..c32e44d 100644
--- a/src/jcode/ui/tui.ss
+++ b/src/jcode/ui/tui.ss
@@ -352,11 +352,17 @@
          (max 0 (- (app-state-scroll-offset state) (quotient (msg-area-height state) 2))))
        (app-state-dirty?-set! state #t))
 
-      ;; Ctrl-C or ESC during agent: cancel
+      ;; Ctrl-C or ESC during agent: cancel. Orphan the worker (bump gen
+      ;; so its future events are discarded) and give immediate UI
+      ;; feedback. The worker will notice *tui-stream-abort* on its next
+      ;; callback and exit.
       ((and (or (= key TB_KEY_CTRL_C) (= key TB_KEY_ESC))
             (app-state-agent-busy? state))
        (tui-log "  -> cancel-agent (key=~a)" key)
        (set-car! *tui-stream-abort* #t)
+       (bump-tui-run-gen!)
+       (app-state-agent-busy?-set! state #f)
+       (add-message! state (msg-block-system "(interrupted)"))
        (app-state-dirty?-set! state #t))
 
       ;; Input handling
@@ -697,13 +703,26 @@
 ;; msg-block lists (which are not thread-safe under preemptive pthreads).
 
 ;; Abort flag: ESC/Ctrl-C during a busy agent sets the car to #t. The
-;; stream callback raises, aborting the in-flight HTTP stream.
+;; stream callback raises, aborting the in-flight HTTP stream. During
+;; the pre-stream "thinking" phase no callback fires, so the key
+;; handler also bumps *tui-run-gen* — the worker tags events with the
+;; gen it started in, and the main loop discards events from stale
+;; generations. This makes ESC responsive even before tokens arrive;
+;; the orphaned worker aborts whenever it finally unblocks.
 (def *tui-stream-abort* (cons #f #f))
+(def *tui-run-gen* (cons 0 #f))
+
+(def (tui-run-gen) (car *tui-run-gen*))
+(def (bump-tui-run-gen!) (set-car! *tui-run-gen* (+ 1 (car *tui-run-gen*))))
 
 (def (send-agent-event! ev)
   (let ((main (*main-thread*)))
     (when main (thread-send main ev))))
 
+(def (send-worker-event! gen ev)
+  ;; Events from an orphaned generation are dropped at apply time.
+  (send-agent-event! (cons gen ev)))
+
 (def (drain-agent-events! state)
   "Drain all pending agent events from the main thread's mailbox.
    Apply them to state synchronously. Returns the number drained."
@@ -716,6 +735,15 @@
           (loop (+ n 1)))))))
 
 (def (apply-agent-event! state ev)
+  ;; New-style events are (gen . body). Drop stale generations silently.
+  (when (and (pair? ev) (number? (car ev)))
+    (let ((gen (car ev)) (body (cdr ev)))
+      (if (= gen (tui-run-gen))
+        (apply-agent-event-body! state body)
+        (tui-log "apply-agent-event: discarded stale gen=~a (current=~a)"
+                 gen (tui-run-gen))))))
+
+(def (apply-agent-event-body! state ev)
   (match ev
     ((list 'stream-token token)
      (tui-stream-token! state token))
@@ -746,6 +774,7 @@
   (app-state-agent-busy?-set! state #t)
   (app-state-stream-buf-set! state "")
   (set-car! *tui-stream-abort* #f)
+  (bump-tui-run-gen!)
   ;; Add empty assistant message that will be filled by streaming
   (add-message! state (msg-block-assistant ""))
   (app-state-dirty?-set! state #t)
@@ -759,38 +788,46 @@
              s-id p-name m-name
              (and (config-get-provider-key p-name) #t))
     ;; Agent worker: sends events to main thread mailbox; never touches state.
-    (spawn
-      (lambda ()
-        (tui-log "worker: entered, main-thread=~a" (*main-thread*))
-        (try
-          (parameterize
-            ((current-provider-override p-override)
-             (current-model-override m-override)
-             (current-stream-cb
-               (lambda (token)
-                 (when (car *tui-stream-abort*)
-                   (error 'stream-aborted "interrupted by user"))
-                 (send-agent-event! (list 'stream-token token))))
-             (current-tool-cb
-               (lambda (event name args)
-                 (tui-log "worker: tool-cb ~a ~a" event name)
-                 (send-agent-event! (list 'tool-event event name args))))
-             (current-usage-cb
-               (lambda (usage)
-                 (send-agent-event! (list 'usage-update usage)))))
-            (tui-log "worker: calling agent-run")
-            (agent-run s-id text)
-            (tui-log "worker: agent-run returned, sending agent-done")
-            (send-agent-event! (list 'agent-done))
-            (tui-log "worker: agent-done sent, worker exiting normally"))
-          (catch (e)
-            (let ((msg (err->string e)))
-              (tui-log "worker: CAUGHT exception: ~a" msg)
-              (cond
-                ((string-contains msg "stream-aborted")
-                 (send-agent-event! (list 'agent-cancelled)))
-                (else
-                 (send-agent-event! (list 'agent-error msg)))))))))))
+    ;; Capture current generation; if the user ESCs while we're blocked in
+    ;; the pre-stream "thinking" phase, main bumps gen and ignores our
+    ;; future events. We still abort cleanly on the next callback.
+    (let ((gen (tui-run-gen)))
+      (spawn
+        (lambda ()
+          (tui-log "worker: entered gen=~a main-thread=~a" gen (*main-thread*))
+          (try
+            (parameterize
+              ((current-provider-override p-override)
+               (current-model-override m-override)
+               (current-stream-cb
+                 (lambda (token)
+                   (when (or (car *tui-stream-abort*)
+                             (not (= gen (tui-run-gen))))
+                     (error 'stream-aborted "interrupted by user"))
+                   (send-worker-event! gen (list 'stream-token token))))
+               (current-tool-cb
+                 (lambda (event name args)
+                   (when (or (car *tui-stream-abort*)
+                             (not (= gen (tui-run-gen))))
+                     (error 'stream-aborted "interrupted by user"))
+                   (tui-log "worker: tool-cb ~a ~a" event name)
+                   (send-worker-event! gen (list 'tool-event event name args))))
+               (current-usage-cb
+                 (lambda (usage)
+                   (send-worker-event! gen (list 'usage-update usage)))))
+              (tui-log "worker: calling agent-run")
+              (agent-run s-id text)
+              (tui-log "worker: agent-run returned, sending agent-done")
+              (send-worker-event! gen (list 'agent-done))
+              (tui-log "worker: agent-done sent, worker exiting normally"))
+            (catch (e)
+              (let ((msg (err->string e)))
+                (tui-log "worker: CAUGHT exception: ~a" msg)
+                (cond
+                  ((string-contains msg "stream-aborted")
+                   (send-worker-event! gen (list 'agent-cancelled)))
+                  (else
+                   (send-worker-event! gen (list 'agent-error msg))))))))))))
 
 (def (tui-stream-token! state token)
   "Handle a streaming token from the LLM — called from agent thread."