about summary refs log tree commit diff
diff options
context:
space:
mode:
authorArun Isaac2026-05-11 22:00:12 +0100
committerArun Isaac2026-05-17 23:10:11 +0100
commit03a0edf0e2f0df29afdded81824c7aa75745fa24 (patch)
tree72959dad76249cd73b4d05d3b4c878126fb55abc
parent02b5c4789e680b7f5c12601290c42adc34ee72d3 (diff)
downloadkaagum-03a0edf0e2f0df29afdded81824c7aa75745fa24.tar.gz
kaagum-03a0edf0e2f0df29afdded81824c7aa75745fa24.tar.lz
kaagum-03a0edf0e2f0df29afdded81824c7aa75745fa24.zip
Use the state monad.
With the state monad, we no longer have to explicitly thread state
through every function call and return it as one of two values. As a
result, the code now reads more naturally.
-rw-r--r--kaagum/monads.scm141
-rw-r--r--kaagum/tea.scm1307
-rw-r--r--kaagum/utils.scm22
3 files changed, 792 insertions, 678 deletions
diff --git a/kaagum/monads.scm b/kaagum/monads.scm
new file mode 100644
index 0000000..14b84fd
--- /dev/null
+++ b/kaagum/monads.scm
@@ -0,0 +1,141 @@
+;;; kaagum --- Tiny, security-focused AI agent in Guile
+;;; Copyright © 2026 Arun Isaac <arunisaac@systemreboot.net>
+;;;
+;;; This file is part of kaagum.
+;;;
+;;; kaagum is free software: you can redistribute it and/or modify it
+;;; under the terms of the GNU General Public License as published by
+;;; the Free Software Foundation, either version 3 of the License, or
+;;; (at your option) any later version.
+;;;
+;;; kaagum is distributed in the hope that it will be useful, but
+;;; WITHOUT ANY WARRANTY; without even the implied warranty of
+;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+;;; General Public License for more details.
+;;;
+;;; You should have received a copy of the GNU General Public License
+;;; along with kaagum.  If not, see <https://www.gnu.org/licenses/>.
+
+(define-module (kaagum monads)
+  #:use-module (rnrs records syntactic)
+  #:use-module (srfi srfi-1)
+  #:use-module (srfi srfi-26)
+  #:use-module (ice-9 match)
+  #:export (state-bind
+            state-return
+            state-let*
+            state-begin
+            state-when
+            state-sequence
+            state-map
+            state-append-map
+            current-state
+            set-current-state
+            run-with-state))
+
+(define-syntax mlet*
+  (syntax-rules ()
+    ((_ bind return () body ...)
+     (begin
+       body ...))
+    ((_ bind return ((var mvalue) other-bindings ...) body ...)
+     (bind mvalue
+           (lambda (var)
+             (mlet* bind return (other-bindings ...)
+                    body ...))))))
+
+(define-syntax mbegin
+  (syntax-rules ()
+    ((_ bind return expression)
+     expression)
+    ((_ bind return first-expression body ...)
+     (bind first-expression
+           (lambda _
+             (mbegin bind return
+                     body ...))))))
+
+(define-syntax-rule (mwhen bind return condition body ...)
+  "When @var{condition} is true, evaluate @var{body} in monad. The monadic return
+value must be ignored. @var{bind} and @var{return} describe the monad in
+question."
+  (if condition
+      (mbegin bind return body ...)
+      (return #f)))
+
+(define-syntax-rule (sequence bind return mvalues)
+  "Convert a list of monadic @var{mvalues} into a monadic list
+of values. @var{bind} and @var{return} describe the monad in question."
+  (mlet* bind return ((reverse-list
+                       (fold (lambda (mvalue mresult)
+                               (mlet* bind return ((result mresult)
+                                                   (value mvalue))
+                                      (return (cons value result))))
+                             (return (list))
+                             mvalues)))
+         (return (reverse reverse-list))))
+
+(define-syntax-rule (mmap bind return mproc lists ...)
+  "Map monadic funcion @var{mproc} over @var{lists} and return a monadic
+list. @var{bind} and @var{return} describe the monad in question."
+  (sequence bind return (map mproc lists ...)))
+
+(define-syntax-rule (mappend-map bind return mproc lists ...)
+  "Map monadic funcion @var{mproc} over @var{lists} like
+@code{mmap}, but return a monadic list of the results appended together.
+@var{bind} and @var{return} describe the monad in question."
+  (mlet* bind return ((mapped (mmap bind return mproc lists ...)))
+         (return (apply append mapped))))
+
+(define-record-type (<stateful-value> stateful-value stateful-value?)
+  (fields (immutable state stateful-value-state)
+          (immutable value stateful-value-value)))
+
+;; We force inlining so that the source location of state-bind calls is
+;; preserved correctly. FIXME: It would be preferable to capture the source
+;; location and somehow write it into the metadata of the lambda function.
+(define-inlinable (state-bind mvalue mproc)
+  (lambda (state)
+    (match (mvalue state)
+      (($ <stateful-value> next-state value)
+       ((mproc value) next-state)))))
+
+(define-inlinable (state-return value)
+  (cut stateful-value <> value))
+
+(define-syntax-rule (state-let* bindings body ...)
+  (mlet* state-bind state-return bindings
+    body ...))
+
+(define-syntax-rule (state-begin body ...)
+  (mbegin state-bind state-return
+    body ...))
+
+(define-syntax-rule (state-when condition body ...)
+  (mwhen state-bind state-return condition
+    body ...))
+
+(define-inlinable (state-sequence mvalues)
+  (sequence state-bind state-return mvalues))
+
+(define-syntax-rule (state-map mproc lists ...)
+  (mmap state-bind state-return mproc lists ...))
+
+(define-syntax-rule (state-append-map mproc lists ...)
+  (mappend-map state-bind state-return mproc lists ...))
+
+(define-inlinable (current-state)
+  "Return the current state as a state-monadic value."
+  (lambda (state)
+    (stateful-value state state)))
+
+(define-inlinable (set-current-state new-state)
+  "Set @var{new-state} as the state. The monadic return value must be ignored."
+  (lambda _
+    (stateful-value new-state #t)))
+
+(define (run-with-state mvalue initial-state)
+  "Run state-monadic value @var{mvalue} starting with @var{initial-state}. Return
+two values---the value encapsulated in @var{mvalue} and the final state."
+  (match (mvalue initial-state)
+    (($ <stateful-value> state value)
+     (values value state))))
diff --git a/kaagum/tea.scm b/kaagum/tea.scm
index 40fdd40..6e750c0 100644
--- a/kaagum/tea.scm
+++ b/kaagum/tea.scm
@@ -30,6 +30,7 @@
   #:use-module (lens)
   #:use-module (kaagum config)
   #:use-module (kaagum lens)
+  #:use-module (kaagum monads)
   #:use-module (kaagum openai)
   #:use-module (kaagum records)
   #:use-module (kaagum tools)
@@ -196,40 +197,42 @@ state."
            (key-ref session-id)
            state-sessions))
 
-(define (state->llm-requests session-id state)
-  "Return a list of @code{<llm-request>} objects for session with @var{session-id}
-in @var{state}."
-  (if (and (null? (focus (state-tool-calls session-id)
-                         state))
-           (not (focus (state-session-cancelling? session-id)
-                       state)))
-      ;; There are no more tool calls in flight and a cancellation is not in
-      ;; progress; dispatch to LLM.
-      (list (llm-request session-id
-                         (focus (state-model session-id)
-                                state)
-                         (map (lambda (message)
-                                ;; Strip unnecessary fields (such as reasoning
-                                ;; fields) based on role.
-                                (let* ((role (focus (key-ref "role") message))
-                                       (allowed-fields
-                                        (cond
-                                         ((string=? role "user")
-                                          '("role" "content"))
-                                         ((string=? role "assistant")
-                                          '("role" "content" "tool_calls"))
-                                         ((string=? role "tool")
-                                          '("role" "content" "tool_call_id")))))
-                                  (filter (match-lambda
-                                            ((key . _)
-                                             (member key allowed-fields)))
-                                          message)))
-                              ;; Reverse because we have been prepending new
-                              ;; messages onto the list.
-                              (reverse (focus (state-messages session-id)
-                                              state)))))
-      ;; There are tool calls or a cancellation in progress; do nothing.
-      (list)))
+(define (llm-requests session-id)
+  "Return the state-monadic list of @code{<llm-request>} objects for session with
+@var{session-id}."
+  (state-let* ((state (current-state)))
+    (state-return
+     (if (and (null? (focus (state-tool-calls session-id)
+                            state))
+              (not (focus (state-session-cancelling? session-id)
+                          state)))
+         ;; There are no more tool calls in flight and a cancellation is not in
+         ;; progress; dispatch to LLM.
+         (list (llm-request session-id
+                            (focus (state-model session-id)
+                                   state)
+                            (map (lambda (message)
+                                   ;; Strip unnecessary fields (such as
+                                   ;; reasoning fields) based on role.
+                                   (let* ((role (focus (key-ref "role") message))
+                                          (allowed-fields
+                                           (cond
+                                            ((string=? role "user")
+                                             '("role" "content"))
+                                            ((string=? role "assistant")
+                                             '("role" "content" "tool_calls"))
+                                            ((string=? role "tool")
+                                             '("role" "content" "tool_call_id")))))
+                                     (filter (match-lambda
+                                               ((key . _)
+                                                (member key allowed-fields)))
+                                             message)))
+                                 ;; Reverse because we have been prepending new
+                                 ;; messages onto the list.
+                                 (reverse (focus (state-messages session-id)
+                                                 state)))))
+         ;; There are tool calls or a cancellation in progress; do nothing.
+         (list)))))
 
 (define-record-type* (<llm-request> llm-request llm-request?)
   (fields (session-id llm-request-session-id)
@@ -245,7 +248,7 @@ in @var{state}."
 
 (define-record-type* (<command> command command?)
   (fields (description command-description)
-          (next-state command-next-state)))
+          (effects command-effects)))
 
 (define (command->spec name command)
   "Serialize @var{command} of @var{name} to ACP-compatible JSON spec."
@@ -267,11 +270,12 @@ in @var{state}."
 
 (define %cwd-command
   (command "Print current working directory of the session"
-           (lambda (state session-id tools argument)
-             (values state
-                     (list (agent-message-chunk session-id
-                                                (focus (state-cwd session-id)
-                                                       state)))))))
+           (lambda (session-id tools argument)
+             (state-let* ((state (current-state)))
+               (state-return
+                (list (agent-message-chunk session-id
+                                           (focus (state-cwd session-id)
+                                                  state))))))))
 
 (define (markdown-table lines)
   "Return a markdown table built from @var{lines}. Each line is a list of strings,
@@ -292,33 +296,33 @@ the table."
 
 (define %tools-command
   (command "List available tools and their permission status"
-           (lambda (state session-id tools argument)
-             (let* ((allowed-tools (focus (state-allowed-tools session-id)
-                                          state))
-                    (rejected-tools (focus (state-rejected-tools session-id)
-                                           state))
-                    (lines
-                     (cons (list "Tool" "Permission")
-                           (map (match-lambda
-                                  ((name . _)
-                                   (list name
-                                         (cond
-                                          ((member name allowed-tools) "allow")
-                                          ((member name rejected-tools) "reject")
-                                          (else "prompt user")))))
-                                tools))))
-               (values state
-                       (list (agent-message-chunk session-id
-                                                  (markdown-table lines))))))))
+           (lambda (session-id tools argument)
+             (state-let* ((state (current-state)))
+               (let* ((allowed-tools (focus (state-allowed-tools session-id)
+                                            state))
+                      (rejected-tools (focus (state-rejected-tools session-id)
+                                             state))
+                      (lines
+                       (cons (list "Tool" "Permission")
+                             (map (match-lambda
+                                    ((name . _)
+                                     (list name
+                                           (cond
+                                            ((member name allowed-tools) "allow")
+                                            ((member name rejected-tools) "reject")
+                                            (else "prompt user")))))
+                                  tools))))
+                 (state-return
+                  (list (agent-message-chunk session-id
+                                             (markdown-table lines)))))))))
 
 (define %commands
   `(("cwd" . ,%cwd-command)
     ("tools" . ,%tools-command)))
 
-(define (next-state-slash-command state session-id tools command-name argument)
-  "Given current @var{state} and an invocation of slash var{command-name} with
-@var{argument} for @var{session-id}, return the next state and a list of
-effects.
+(define (slash-command-effects session-id tools command-name argument)
+  "Return the state-monadic list of effects invoking slash @var{command-name} with
+@var{argument} for @var{session-id}.
 
 @var{tools} is the same as in @code{run-tea-loop}."
   (cond
@@ -326,12 +330,11 @@ effects.
    ((focus (key-ref command-name)
            %commands)
     => (lambda (command)
-         ((command-next-state command) state session-id tools argument)))
+         ((command-effects command) session-id tools argument)))
    ;; command not found
    (else
-    (values state
-            (list (agent-message-chunk session-id
-                                       "Error: Unknown command"))))))
+    (state-return (list (agent-message-chunk session-id
+                                             "Error: Unknown command"))))))
 
 (define (model->spec id model)
   "Serialize @var{model} of @var{id} to ACP-compatible JSON spec."
@@ -368,432 +371,450 @@ the association list of @var{available-models}."
                                     (model->spec id model)))
                                  available-models))))))
 
-(define (next-state-client-request state request models tools)
-  "Given current @var{state} and a new ACP @var{request}, return the next state and
-a list of effects.
+(define (client-request-effects request models tools)
+  "Return the state-monadic list of effects for new ACP @var{request}.
 
 @var{tools} is the same as in @code{run-tea-loop}. @var{models} is the same as
 in @code{tea-loop}."
-  (let ((request-id (focus state-client-request-id
-                           state)))
-    (cond
-     ;; There is a pending request from the client; process it.
-     ((focus (key-ref "method") request)
-      => (match-lambda
-           ("initialize"
-            (values state
-                    (list (acp-message `(("jsonrpc" . "2.0")
-                                         ("id" . ,request-id)
-                                         ("result"
-                                          ("protocolVersion" . 1)
-                                          ("agentCapabilities")
-                                          ("agentInfo"
-                                           ("name" . ,%project)
-                                           ("title" . ,%project-title)
-                                           ("version" . ,%version))
-                                          ("authMethods" . #())))))))
-           ("session/new"
-            (let ((session-id
-                   (string-append "session-"
-                                  (number->string
-                                   (focus state-next-session-id state))))
-                  (model (match models
-                           (((model-id . _) . _)
-                            model-id))))
-              (values (-> state
-                          ;; Push new session onto list.
-                          (push state-sessions
-                                (cons session-id
-                                      ;; TODO: Check if cwd is an absolute path.
-                                      (session (focus (in "params" "cwd")
-                                                      request)
-                                               model))
-                                <>)
-                          ;; Increment next session ID.
-                          (over state-next-session-id
-                                1+
-                                <>))
-                      (map acp-message
-                           `(;; Return new session.
-                             (("jsonrpc" . "2.0")
-                              ("id" . ,request-id)
-                              ("result"
-                               ("sessionId" . ,session-id)
-                               ("configOptions" . ,(config-options model models))))
-                             ;; Advertise available commands.
-                             (("jsonrpc" . "2.0")
-                              ("method" . "session/update")
-                              ("params"
-                               ("sessionId" . ,session-id)
-                               ("update"
-                                ("sessionUpdate" . "available_commands_update")
-                                ("availableCommands" .
-                                 ,(list->vector (map (match-lambda
-                                                       ((name . command)
-                                                        (command->spec name command)))
-                                                     %commands)))))))))))
-           ("session/set_config_option"
-            (let ((session-id (focus (in "params" "sessionId")
-                                     request))
-                  (config-id (focus (in "params" "configId")
-                                    request))
-                  (model (focus (in "params" "value")
-                                request)))
-              (if (string=? config-id "model")
-                  ;; The client is setting the model.
-                  (if (focus (key-ref model)
-                             models)
-                      ;; Set model in state and respond to the client.
-                      (values (put (state-model session-id)
-                                   (focus (in "params" "value")
-                                          request)
-                                   state)
-                              (list (acp-message `(("jsonrpc" . "2.0")
-                                                   ("id" . ,request-id)
-                                                   ("result"
-                                                    ("configOptions" .
-                                                     ,(config-options model models)))))))
-                      ;; The client specified a model that is not one of the
-                      ;; available models.
-                      (values state
-                              (list (jsonrpc-error request-id
-                                                   -32602
-                                                   "Invalid model"))))
-                  ;; The client specified an unknown configId parameter.
-                  (values state
-                          (list (jsonrpc-error request-id
-                                               -32602
-                                               "Unsupported configId parameter"))))))
-           ("session/prompt"
-            (let ((session-id (focus (in "params" "sessionId")
-                                     request))
-                  ;; TODO: Filter to only allow "text" type content blocks.
-                  (prompt (focus (in "params" "prompt")
-                                 request)))
-              (cond
-               ;; slash command
-               ((string-match "^/([a-z0-9]*)(.*)"
-                              (focus (in-json 0 "text")
-                                     prompt))
-                => (lambda (mtch)
-                     (let ((command-name (match:substring mtch 1))
-                           (argument (string-trim (match:substring mtch 2))))
-                       (let-values (((state slash-command-effects)
-                                     (next-state-slash-command
-                                      state session-id tools command-name argument)))
-                         ;; End prompt turn immediately. This means slash
-                         ;; commands cannot send LLM requests or initiate other
-                         ;; exchanges.
-                         (let-values (((state end-turn-effects)
-                                       (next-state-end-turn state session-id models)))
-                           (values state
-                                   (append slash-command-effects
-                                           end-turn-effects)))))))
-               ;; regular prompt
-               (else
-                (let ((state (push (state-messages session-id)
-                                   `(("role" . "user")
-                                     ("content" . ,prompt))
-                                   state)))
-                  (values state
-                          (state->llm-requests session-id state)))))))
-           ("session/cancel"
-            (let ((session-id (focus (in "params" "sessionId")
-                                     request)))
-              ;; Reset state and end the prompt turn.
-              (values (-> state
-                          (put (state-session-cancelling? session-id)
-                               #f
-                               <>)
-                          (put (state-tool-calls session-id)
-                               '()
-                               <>)
-                          (put state-requests-alist
-                               '()
-                               <>))
-                      (list (acp-message `(("jsonrpc" . "2.0")
-                                           ("id" . ,request-id)
-                                           ("result"
-                                            ("stopReason" . "cancelled")))))))))))))
-
-(define (next-state-client-response state response)
-  "Given current @var{state} and a new ACP @var{response}, return the next state
-and a list of effects."
-  (let ((request-id (focus (key-ref "id")
-                           response)))
-    (cond
-     ;; The tool call for this response has been found; process it.
-     ((focus (compose (key-ref request-id)
-                      state-requests-alist)
-             state)
-      => (lambda (tool-call)
-           (let* ((session-id (tool-call-session-id tool-call))
-                  (outcome (focus (in "result" "outcome" "outcome")
-                                  response))
-                  (option-id (focus (in "result" "outcome" "optionId")
+  (state-let* ((state (current-state)))
+    (let ((request-id (focus state-client-request-id
+                             state)))
+      (cond
+       ;; There is a pending request from the client; process it.
+       ((focus (key-ref "method") request)
+        => (match-lambda
+             ("initialize"
+              (state-return
+               (list (acp-message `(("jsonrpc" . "2.0")
+                                    ("id" . ,request-id)
+                                    ("result"
+                                     ("protocolVersion" . 1)
+                                     ("agentCapabilities")
+                                     ("agentInfo"
+                                      ("name" . ,%project)
+                                      ("title" . ,%project-title)
+                                      ("version" . ,%version))
+                                     ("authMethods" . #())))))))
+             ("session/new"
+              (let ((session-id
+                     (string-append "session-"
+                                    (number->string
+                                     (focus state-next-session-id state))))
+                    (model (match models
+                             (((model-id . _) . _)
+                              model-id))))
+                (state-begin
+                 (state-let* ((state (current-state)))
+                   (set-current-state (-> state
+                                          ;; Push new session onto list.
+                                          (push state-sessions
+                                                (cons session-id
+                                                      ;; TODO: Check if cwd is an absolute path.
+                                                      (session (focus (in "params" "cwd")
+                                                                      request)
+                                                               model))
+                                                <>)
+                                          ;; Increment next session ID.
+                                          (over state-next-session-id
+                                                1+
+                                                <>))))
+                 (state-return
+                  (map acp-message
+                       `( ;; Return new session.
+                         (("jsonrpc" . "2.0")
+                          ("id" . ,request-id)
+                          ("result"
+                           ("sessionId" . ,session-id)
+                           ("configOptions" . ,(config-options model models))))
+                         ;; Advertise available commands.
+                         (("jsonrpc" . "2.0")
+                          ("method" . "session/update")
+                          ("params"
+                           ("sessionId" . ,session-id)
+                           ("update"
+                            ("sessionUpdate" . "available_commands_update")
+                            ("availableCommands" .
+                             ,(list->vector (map (match-lambda
+                                                   ((name . command)
+                                                    (command->spec name command)))
+                                                 %commands))))))))))))
+             ("session/set_config_option"
+              (let ((session-id (focus (in "params" "sessionId")
+                                       request))
+                    (config-id (focus (in "params" "configId")
+                                      request))
+                    (model (focus (in "params" "value")
+                                  request)))
+                (if (string=? config-id "model")
+                    ;; The client is setting the model.
+                    (if (focus (key-ref model)
+                               models)
+                        ;; Set model in state and respond to the client.
+                        (state-begin
+                         (state-let* ((state (current-state)))
+                           (set-current-state (put (state-model session-id)
+                                                   (focus (in "params" "value")
+                                                          request)
+                                                   state)))
+                         (state-return
+                          (list (acp-message `(("jsonrpc" . "2.0")
+                                               ("id" . ,request-id)
+                                               ("result"
+                                                ("configOptions" .
+                                                 ,(config-options model models))))))))
+                        ;; The client specified a model that is not one of the
+                        ;; available models.
+                        (state-return
+                         (list (jsonrpc-error request-id
+                                              -32602
+                                              "Invalid model"))))
+                    ;; The client specified an unknown configId parameter.
+                    (state-return
+                     (list (jsonrpc-error request-id
+                                          -32602
+                                          "Unsupported configId parameter"))))))
+             ("session/prompt"
+              (let ((session-id (focus (in "params" "sessionId")
+                                       request))
+                    ;; TODO: Filter to only allow "text" type content blocks.
+                    (prompt (focus (in "params" "prompt")
+                                   request)))
+                (cond
+                 ;; slash command
+                 ((string-match "^/([a-z0-9]*)(.*)"
+                                (focus (in-json 0 "text")
+                                       prompt))
+                  => (lambda (mtch)
+                       (let ((command-name (match:substring mtch 1))
+                             (argument (string-trim (match:substring mtch 2))))
+                         (state-let* ((slash-command-effects
+                                       (slash-command-effects
+                                        session-id tools command-name argument))
+                                      ;; End prompt turn immediately. This means
+                                      ;; slash commands cannot send LLM requests
+                                      ;; or initiate other exchanges.
+                                      (end-turn-effects
+                                       (end-turn-effects session-id models)))
+                           (state-return
+                            (append slash-command-effects
+                                    end-turn-effects))))))
+                 ;; regular prompt
+                 (else
+                  (state-begin
+                   (state-let* ((state (current-state)))
+                     (set-current-state (push (state-messages session-id)
+                                              `(("role" . "user")
+                                                ("content" . ,prompt))
+                                              state)))
+                   (llm-requests session-id))))))
+             ("session/cancel"
+              (let ((session-id (focus (in "params" "sessionId")
+                                       request)))
+                (state-begin
+                 ;; Reset state and end the prompt turn.
+                 (state-let* ((state (current-state)))
+                   (set-current-state (-> state
+                                          (put (state-session-cancelling? session-id)
+                                               #f
+                                               <>)
+                                          (put (state-tool-calls session-id)
+                                               '()
+                                               <>)
+                                          (put state-requests-alist
+                                               '()
+                                               <>))))
+                 (state-return
+                  (list (acp-message `(("jsonrpc" . "2.0")
+                                       ("id" . ,request-id)
+                                       ("result"
+                                        ("stopReason" . "cancelled")))))))))))))))
+
+(define (client-response-effects response)
+  "Return the state-monadic list of effects for a new ACP @var{response}."
+  (state-let* ((state (current-state)))
+    (let ((request-id (focus (key-ref "id")
+                             response)))
+      (cond
+       ;; The tool call for this response has been found; process it.
+       ((focus (compose (key-ref request-id)
+                        state-requests-alist)
+               state)
+        => (lambda (tool-call)
+             (let* ((session-id (tool-call-session-id tool-call))
+                    (outcome (focus (in "result" "outcome" "outcome")
                                     response))
-                  ;; Decode permission selection to symbol.
-                  (selection
-                   (cond
-                    ((and (string=? outcome "selected")
-                          (string=? option-id
-                                    (focus (key-ref "optionId") %tool-allow-once)))
-                     'allow-once)
-                    ((and (string=? outcome "selected")
-                          (string=? option-id
-                                    (focus (key-ref "optionId") %tool-allow-always)))
-                     'allow-always)
-                    ((and (string=? outcome "selected")
-                          (string=? option-id
-                                    (focus (key-ref "optionId") %tool-reject-once)))
-                     'reject-once)
-                    ((and (string=? outcome "selected")
-                          (string=? option-id
-                                    (focus (key-ref "optionId") %tool-reject-always)))
-                     'reject-always)
-                    ;; We don't explicitly look for "cancelled". We defensively
-                    ;; assume anything other than "selected" is "cancelled".
-                    ;; This protects us from buggy clients.
-                    (else 'cancel)))
-                  (state
-                   (-> state
-                       ;; If the tool was "allowed always", add it to the list
-                       ;; of allowed tools.
-                       (over (state-allowed-tools session-id)
-                             (lambda (allowed-tools)
-                               (if (eq? selection 'allow-always)
-                                   (cons (tool-call-function tool-call)
-                                         allowed-tools)
-                                   allowed-tools))
-                             <>)
-                       ;; If the tool was "rejected always", add it to the list
-                       ;; of rejected tools.
-                       (over (state-rejected-tools session-id)
-                             (lambda (rejected-tools)
-                               (if (eq? selection 'reject-always)
-                                   (cons (tool-call-function tool-call)
-                                         rejected-tools)
-                                   rejected-tools))
+                    (option-id (focus (in "result" "outcome" "optionId")
+                                      response))
+                    ;; Decode permission selection to symbol.
+                    (selection
+                     (cond
+                      ((and (string=? outcome "selected")
+                            (string=? option-id
+                                      (focus (key-ref "optionId") %tool-allow-once)))
+                       'allow-once)
+                      ((and (string=? outcome "selected")
+                            (string=? option-id
+                                      (focus (key-ref "optionId") %tool-allow-always)))
+                       'allow-always)
+                      ((and (string=? outcome "selected")
+                            (string=? option-id
+                                      (focus (key-ref "optionId") %tool-reject-once)))
+                       'reject-once)
+                      ((and (string=? outcome "selected")
+                            (string=? option-id
+                                      (focus (key-ref "optionId") %tool-reject-always)))
+                       'reject-always)
+                      ;; We don't explicitly look for "cancelled". We
+                      ;; defensively assume anything other than "selected" is
+                      ;; "cancelled". This protects us from buggy clients.
+                      (else 'cancel))))
+               (state-begin
+                (set-current-state (-> state
+                                       ;; If the tool was "allowed always", add
+                                       ;; it to the list of allowed tools.
+                                       (over (state-allowed-tools session-id)
+                                             (lambda (allowed-tools)
+                                               (if (eq? selection 'allow-always)
+                                                   (cons (tool-call-function tool-call)
+                                                         allowed-tools)
+                                                   allowed-tools))
+                                             <>)
+                                       ;; If the tool was "rejected always", add
+                                       ;; it to the list of rejected tools.
+                                       (over (state-rejected-tools session-id)
+                                             (lambda (rejected-tools)
+                                               (if (eq? selection 'reject-always)
+                                                   (cons (tool-call-function tool-call)
+                                                         rejected-tools)
+                                                   rejected-tools))
+                                             <>)
+                                       ;; Update tool status.
+                                       (put (compose tool-call-status
+                                                     (state-tool-call (tool-call-id tool-call)
+                                                                      session-id))
+                                            (cond
+                                             ((eq? selection 'cancel) 'cancelled)
+                                             ((memq selection '(allow-once allow-always))
+                                              'approved)
+                                             ((memq selection '(reject-once reject-always))
+                                              'rejected)
+                                             ;; This branch should be
+                                             ;; unreachable.
+                                             (else
+                                              (assertion-violation selection
+                                                                   "Invalid selection")))
+                                            <>)
+                                       ;; If the tool call was cancelled, set
+                                       ;; the cancelling flag to indicate that a
+                                       ;; cancellation is in progress.
+                                       (put (state-session-cancelling? session-id)
+                                            (eq? selection 'cancel)
+                                            <>)
+                                       ;; Unregister request corresponding to
+                                       ;; this response.
+                                       (alist-delete-over state-requests-alist
+                                                          request-id
+                                                          <>)))
+                ;; Request tool call evaluation. eval-tool-call handles
+                ;; cancelled and rejected tool calls correctly. We don't have to
+                ;; worry about it here.
+                (state-let* ((state (current-state)))
+                  (state-return
+                   (list (focus (state-tool-call (tool-call-id tool-call)
+                                                 (tool-call-session-id tool-call))
+                                state))))))))
+       ;; Client response is stale (it pertains to a request not in
+       ;; requests-alist). Silently ignore it.
+       (else (state-return '()))))))
+
+(define (send-agent-request-effects request context)
+  "Return the state-monadic list of effects sending @var{request} from the agent to
+the client. Stash @var{context} against request ID in @code{requests-alist} for
+future recall."
+  (state-let* ((state (current-state)))
+    (let ((request-id (focus state-agent-request-id state)))
+      (state-begin
+       (set-current-state (-> state
+                              ;; Register this request so we can recall the
+                              ;; details later when the response comes in.
+                              (push state-requests-alist
+                                    (cons request-id
+                                          context)
+                                    <>)
+                              ;; Bump agent request ID.
+                              (over state-agent-request-id
+                                    1+
+                                    <>)))
+       (state-return
+        ;; Set request ID in request JSON.
+        (list (if (assoc "id" (focus acp-message-json request))
+                  ;; The request JSON already has an "id" field; overwrite it.
+                  (put (compose (key-ref "id") acp-message-json)
+                       request-id
+                       request)
+                  ;; The request JSON has no "id" field; cons it on;
+                  (over acp-message-json
+                        (cut cons
+                             (cons "id" request-id)
                              <>)
-                       ;; Update tool status.
-                       (put (compose tool-call-status
-                                     (state-tool-call (tool-call-id tool-call)
-                                                      session-id))
-                            (cond
-                             ((eq? selection 'cancel) 'cancelled)
-                             ((memq selection '(allow-once allow-always))
-                              'approved)
-                             ((memq selection '(reject-once reject-always))
-                              'rejected)
-                             ;; This branch should be unreachable.
-                             (else
-                              (assertion-violation selection
-                                                   "Invalid selection")))
-                            <>)
-                       ;; If the tool call was cancelled, set the cancelling flag
-                       ;; to indicate that a cancellation is in progress.
-                       (put (state-session-cancelling? session-id)
-                            (eq? selection 'cancel)
-                            <>)
-                       ;; Unregister request corresponding to this response.
-                       (alist-delete-over state-requests-alist
-                                          request-id
-                                          <>))))
-             (values state
-                     ;; Request tool call evaluation. eval-tool-call handles
-                     ;; cancelled and rejected tool calls correctly. We don't
-                     ;; have to worry about it here.
-                     (list (focus (state-tool-call (tool-call-id tool-call)
-                                                   (tool-call-session-id tool-call))
-                                  state))))))
-     ;; Client response is stale (it pertains to a request not in
-     ;; requests-alist). Silently ignore it.
-     (else (values state '())))))
+                        request))))))))
 
-(define (next-state-send-agent-request state request context)
-  "Return an updated @var{state} and a list of effects sending @var{request} from
-the agent to the client. Stash @var{context} against request ID in
-@code{requests-alist} for future recall."
-  (let ((request-id (focus state-agent-request-id state)))
-    (values (-> state
-                ;; Register this request so we can recall the details later when
-                ;; the response comes in.
-                (push state-requests-alist
-                      (cons request-id
-                            context)
-                      <>)
-                ;; Bump agent request ID.
-                (over state-agent-request-id
-                      1+
-                      <>))
-            ;; Set request ID in request JSON.
-            (list (if (assoc "id" (focus acp-message-json request))
-                      ;; The request JSON already has an "id" field; overwrite
-                      ;; it.
-                      (put (compose (key-ref "id") acp-message-json)
-                           request-id
-                           request)
-                      ;; The request JSON has no "id" field; cons it on;
-                      (over acp-message-json
-                            (cut cons
-                                 (cons "id" request-id)
-                                 <>)
-                            request))))))
-
-(define (next-state-tool-call state session-id call-json tools)
-  "Given current @var{state} and a new tool @var{call-json}, return the next state
-and a list of effects.
+(define (tool-call-effects session-id call-json tools)
+  "Return the state-monadic list of effects for a new tool @var{call-json} in
+session with @var{session-id}.
 
 @var{tools} is the same as in @code{run-tea-loop}."
-  (guard (c ((tool-call-parse-failure? c)
-             (let ((call-id (focus (key-ref "id") call-json)))
-               (values (-> state
-                           ;; Push tool call response onto messages.
-                           (push (state-messages session-id)
-                                 `(("role" . "tool")
-                                   ("tool_call_id" . ,call-id)
-                                   ("content" . ,(tool-call-parse-failure-message c)))
-                                 <>))
-                       ;; Notify client about invalid tool call.
-                       (list (acp-message `(("jsonrpc" . "2.0")
-                                            ("method" . "session/update")
-                                            ("params"
-                                             ("sessionId" . ,session-id)
-                                             ("update"
-                                              ("sessionUpdate" . "tool_call_update")
-                                              ("toolCallId" . ,call-id)
-                                              ("status" . "failed"))))))))))
-    (let* ((call (spec->tool-call session-id
-                                  (focus (state-cwd session-id)
-                                         state)
-                                  tools
-                                  (focus (state-allowed-tools session-id)
-                                         state)
-                                  (focus (state-rejected-tools session-id)
-                                         state)
-                                  call-json))
-           (state
-            ;; Queue tool call.
-            (push (state-tool-calls session-id)
-                  (cons (tool-call-id call)
-                        call)
-                  state)))
-      (let-values (((state effects)
-                    (if (eq? (focus tool-call-status call)
-                             'pending-approval)
-                        ;; Tool call requires permission from the client;
-                        ;; dispatch a request.
-                        (next-state-send-agent-request
-                         state
-                         (acp-message `(("jsonrpc" . "2.0")
-                                        ("method" . "session/request_permission")
-                                        ("params"
-                                         ("sessionId" . ,session-id)
-                                         ("toolCall"
-                                          ("toolCallId" . ,(tool-call-id call)))
-                                         ("options" .
-                                          ,(vector %tool-allow-once
-                                                   %tool-allow-always
-                                                   %tool-reject-once
-                                                   %tool-reject-always)))))
-                         call)
-                        ;; Tool call is already pre-approved or pre-rejected;
-                        ;; schedule it.
-                        (values state
-                                (list call)))))
-        (values state
-                ;; Notify client about new tool call before other effects.
-                (cons (let ((tool (focus (key-ref (tool-call-function call))
-                                         tools))
-                            (args (alist->plist (tool-call-arguments call))))
-                        (acp-message `(("jsonrpc" . "2.0")
-                                       ("method" . "session/update")
-                                       ("params"
-                                        ("sessionId" . ,session-id)
-                                        ("update"
-                                         ("sessionUpdate" . "tool_call")
-                                         ("toolCallId" . ,(tool-call-id call))
-                                         ("title" . ,(apply (tool-title tool) args))
-                                         ("kind" . ,(tool-kind tool))
-                                         ("rawInput" . ,(tool-call-arguments call))
-                                         ("status" . "pending"))))))
-                      effects))))))
-
-(define (next-state-end-turn state session-id models)
-  "Given current @var{state}, return the next state and a list of effects for
-ending the turn of session with @var{session-id}.
+  (state-let* ((state (current-state)))
+    (guard (c ((tool-call-parse-failure? c)
+               (let ((call-id (focus (key-ref "id") call-json)))
+                 (state-begin
+                  (state-let* ((state (current-state)))
+                    (set-current-state (-> state
+                                           ;; Push tool call response onto messages.
+                                           (push (state-messages session-id)
+                                                 `(("role" . "tool")
+                                                   ("tool_call_id" . ,call-id)
+                                                   ("content" . ,(tool-call-parse-failure-message c)))
+                                                 <>))))
+                  ;; Notify client about invalid tool call.
+                  (state-return (list (acp-message `(("jsonrpc" . "2.0")
+                                                     ("method" . "session/update")
+                                                     ("params"
+                                                      ("sessionId" . ,session-id)
+                                                      ("update"
+                                                       ("sessionUpdate" . "tool_call_update")
+                                                       ("toolCallId" . ,call-id)
+                                                       ("status" . "failed")))))))))))
+      (let ((call (spec->tool-call session-id
+                                   (focus (state-cwd session-id)
+                                          state)
+                                   tools
+                                   (focus (state-allowed-tools session-id)
+                                          state)
+                                   (focus (state-rejected-tools session-id)
+                                          state)
+                                   call-json)))
+        (state-begin
+         ;; Queue tool call.
+         (set-current-state (push (state-tool-calls session-id)
+                                  (cons (tool-call-id call)
+                                        call)
+                                  state))
+         (state-let* ((effects
+                       (if (eq? (focus tool-call-status call)
+                                'pending-approval)
+                           ;; Tool call requires permission from the client;
+                           ;; dispatch a request.
+                           (send-agent-request-effects
+                            (acp-message `(("jsonrpc" . "2.0")
+                                           ("method" . "session/request_permission")
+                                           ("params"
+                                            ("sessionId" . ,session-id)
+                                            ("toolCall"
+                                             ("toolCallId" . ,(tool-call-id call)))
+                                            ("options" .
+                                             ,(vector %tool-allow-once
+                                                      %tool-allow-always
+                                                      %tool-reject-once
+                                                      %tool-reject-always)))))
+                            call)
+                           ;; Tool call is already pre-approved or pre-rejected;
+                           ;; schedule it.
+                           (state-return (list call)))))
+           ;; Notify client about new tool call before other effects.
+           (state-return
+            (cons (let ((tool (focus (key-ref (tool-call-function call))
+                                     tools))
+                        (args (alist->plist (tool-call-arguments call))))
+                    (acp-message `(("jsonrpc" . "2.0")
+                                   ("method" . "session/update")
+                                   ("params"
+                                    ("sessionId" . ,session-id)
+                                    ("update"
+                                     ("sessionUpdate" . "tool_call")
+                                     ("toolCallId" . ,(tool-call-id call))
+                                     ("title" . ,(apply (tool-title tool) args))
+                                     ("kind" . ,(tool-kind tool))
+                                     ("rawInput" . ,(tool-call-arguments call))
+                                     ("status" . "pending"))))))
+                  effects))))))))
+
+(define (end-turn-effects session-id models)
+  "Return the state-monadic list of effects for ending the turn of session with
+@var{session-id}.
 
 @var{models} is the same as in @code{tea-loop}."
-  (let ((input-tokens (focus (state-session-input-tokens session-id)
-                             state))
-        (output-tokens (focus (state-session-output-tokens session-id)
-                              state))
-        (thought-tokens (focus (state-session-thought-tokens session-id)
+  (state-let* ((state (current-state)))
+    (let ((input-tokens (focus (state-session-input-tokens session-id)
                                state))
-        (cache-read-tokens (focus (state-session-cache-read-tokens session-id)
-                                  state))
-        (cache-write-tokens (focus (state-session-cache-write-tokens session-id)
-                                   state))
-        (model-lens (key-ref (focus (state-model session-id)
-                                    state))))
-    (values (-> state
-                ;; Reset per-turn token counters.
-                (put (state-session-input-tokens session-id)
-                     0
-                     <>)
-                (put (state-session-output-tokens session-id)
-                     0
-                     <>)
-                (put (state-session-thought-tokens session-id)
-                     0
-                     <>)
-                (put (state-session-cache-read-tokens session-id)
-                     0
-                     <>)
-                (put (state-session-cache-write-tokens session-id)
-                     0
-                     <>))
-            (list (acp-message `(("jsonrpc" . "2.0")
-                                 ("id" . ,(focus state-client-request-id
-                                                 state))
-                                 ("result"
-                                  ("stopReason" . "end_turn")
-                                  ;; Report usage if it is non-zero. Usage can
-                                  ;; be zero if turn was for a slash command.
-                                  ,@(if (any (negate zero?)
-                                             (list input-tokens
-                                                   output-tokens
-                                                   thought-tokens
-                                                   cache-read-tokens
-                                                   cache-write-tokens))
-                                        `(("usage"
-                                           ("totalTokens" . ,(+ input-tokens
-                                                                output-tokens))
-                                           ("inputTokens" . ,input-tokens)
-                                           ("outputTokens" . ,output-tokens)
-                                           ("thoughtTokens" . ,thought-tokens)
-                                           ("cachedReadTokens" . ,cache-read-tokens)
-                                           ("cachedWriteTokens" . ,cache-write-tokens)))
-                                        '()))))
-                  (acp-message `(("jsonrpc" . "2.0")
-                                 ("method" . "session/update")
-                                 ("params"
-                                  ("sessionId" . ,session-id)
-                                  ("update"
-                                   ("sessionUpdate" . "usage_update")
-                                   ("used" . ,(+ input-tokens
-                                                 output-tokens))
-                                   ("size" . ,(focus (compose model-context-length
-                                                              model-lens)
-                                                     models))
-                                   ("cost"
-                                    ("amount" . ,(focus (state-session-cost session-id)
-                                                        state))
-                                    ("currency" . "USD"))))))))))
-
-(define (next-state-llm-response state response tools models)
-  "Given current @var{state} and a new LLM @var{response}, return the next state
-and a list of effects.
+          (output-tokens (focus (state-session-output-tokens session-id)
+                                state))
+          (thought-tokens (focus (state-session-thought-tokens session-id)
+                                 state))
+          (cache-read-tokens (focus (state-session-cache-read-tokens session-id)
+                                    state))
+          (cache-write-tokens (focus (state-session-cache-write-tokens session-id)
+                                     state)))
+      (state-begin
+       (set-current-state (-> state
+                              ;; Reset per-turn token counters.
+                              (put (state-session-input-tokens session-id)
+                                   0
+                                   <>)
+                              (put (state-session-output-tokens session-id)
+                                   0
+                                   <>)
+                              (put (state-session-thought-tokens session-id)
+                                   0
+                                   <>)
+                              (put (state-session-cache-read-tokens session-id)
+                                   0
+                                   <>)
+                              (put (state-session-cache-write-tokens session-id)
+                                   0
+                                   <>)))
+       (state-return
+        (list (acp-message `(("jsonrpc" . "2.0")
+                             ("id" . ,(focus state-client-request-id
+                                             state))
+                             ("result"
+                              ("stopReason" . "end_turn")
+                              ;; Report usage if it is non-zero. Usage can
+                              ;; be zero if turn was for a slash command.
+                              ,@(if (any (negate zero?)
+                                         (list input-tokens
+                                               output-tokens
+                                               thought-tokens
+                                               cache-read-tokens
+                                               cache-write-tokens))
+                                    `(("usage"
+                                       ("totalTokens" . ,(+ input-tokens
+                                                            output-tokens))
+                                       ("inputTokens" . ,input-tokens)
+                                       ("outputTokens" . ,output-tokens)
+                                       ("thoughtTokens" . ,thought-tokens)
+                                       ("cachedReadTokens" . ,cache-read-tokens)
+                                       ("cachedWriteTokens" . ,cache-write-tokens)))
+                                    '()))))
+              (acp-message `(("jsonrpc" . "2.0")
+                             ("method" . "session/update")
+                             ("params"
+                              ("sessionId" . ,session-id)
+                              ("update"
+                               ("sessionUpdate" . "usage_update")
+                               ("used" . ,(+ input-tokens
+                                             output-tokens))
+                               ("size" . ,(focus (compose model-context-length
+                                                          (key-ref (focus (state-model session-id)
+                                                                          state)))
+                                                 models))
+                               ("cost"
+                                ("amount" . ,(focus (state-session-cost session-id)
+                                                    state))
+                                ("currency" . "USD"))))))))))))
+
+(define (llm-response-effects response tools models)
+  "Return the state-monadic list of effects for new LLM @var{response}.
 
 @var{tools} is the same as in @code{run-tea-loop}. @var{models} is the same as
 in @code{tea-loop}."
@@ -805,113 +826,110 @@ in @code{tea-loop}."
                                    llm-reply)
                             => vector->list)
                            (else '()))))
-    (let-values (((state tool-call-effects)
-                  (foldn (lambda (call-json state effects)
-                           (let-values (((state new-effects)
-                                         (next-state-tool-call
-                                          state session-id call-json tools)))
-                             (values state
-                                     (append new-effects effects))))
-                         tool-calls-json
-                         (-> state
-                             (add (state-session-input-tokens session-id)
-                                  (focus (in-json "usage" "prompt_tokens")
-                                         (llm-response-json response))
-                                  <>)
-                             (add (state-session-output-tokens session-id)
-                                  (focus (in-json "usage" "completion_tokens")
-                                         (llm-response-json response))
-                                  <>)
-                             (add (state-session-thought-tokens session-id)
-                                  (or (focus (in-json "usage"
-                                                      "completion_tokens_details"
-                                                      "reasoning_tokens")
-                                             (llm-response-json response))
-                                      0)
-                                  <>)
-                             (add (state-session-cache-read-tokens session-id)
-                                  (focus (in-json "usage"
-                                                  "prompt_tokens_details"
-                                                  "cached_tokens")
-                                         (llm-response-json response))
-                                  <>)
-                             (add (state-session-cache-write-tokens session-id)
-                                  (or (focus (in-json "usage"
-                                                      "prompt_tokens_details"
-                                                      "cache_write_tokens")
-                                             (llm-response-json response))
-                                      0)
-                                  <>)
-                             (add (state-session-cost session-id)
-                                  (or (focus (in-json "usage" "cost")
-                                             (llm-response-json response))
-                                      0)
-                                  <>)
-                             ;; Push LLM response onto messages.
-                             (push (state-messages session-id)
-                                   llm-reply
-                                   <>))
-                         '())))
-      (let ((effects (cons (agent-message-chunk session-id
-                                                ;; Send LLM's text response.
-                                                (focus (key-ref "content")
-                                                       llm-reply))
-                           tool-call-effects)))
-        (if (null? tool-calls-json)
-            ;; There are no further tool calls,
-            (if (not (focus (state-session-cancelling? session-id)
-                            state))
-                ;; … and a cancellation is not in progress; end turn.
-                (let-values (((state end-turn-effects)
-                              (next-state-end-turn state session-id models)))
-                  (values state
-                          (append effects end-turn-effects)))
-                ;; Else, return what we have so far.
-                (values state effects))
-            ;; Maybe dispatch LLM requests.
-            (values state
-                    (append effects
-                            (state->llm-requests session-id state))))))))
-
-(define (next-state-tool-call-result state result)
-  "Given current @var{state} and a new tool call @var{result}, return the next
-state and a list of effects."
-  (let* ((session-id (tool-call-result-session-id result))
-         (state (-> state
-                    ;; Push tool call result onto messages.
-                    (push (state-messages session-id)
-                          (tool-call-result-json result)
-                          <>)
-                    ;; Delete tool call from session tool call list.
-                    (alist-delete-over (state-tool-calls session-id)
-                                       (tool-call-result-call-id result)
-                                       <>))))
-    (values state
-            ;; Send a notification for each tool call evaluated.
-            (cons (acp-message `(("jsonrpc" . "2.0")
-                                 ("method" . "session/update")
-                                 ("params"
-                                  ("sessionId" . ,session-id)
-                                  ("update"
-                                   ;; TODO: Add locations and rawOutput.
-                                   ("sessionUpdate" . "tool_call_update")
-                                   ("toolCallId" . ,(tool-call-result-call-id result))
-                                   ("status" .
-                                    ,(if (tool-call-result-success? result)
-                                         "completed"
-                                         "failed"))
-                                   ("content" .
-                                    ,(vector `(("type" . "content")
-                                               ("content"
-                                                ("type" . "text")
-                                                ("text" .
-                                                 ,(focus (key-ref "content")
-                                                         (tool-call-result-json result)))))))))))
-                  (state->llm-requests session-id state)))))
-
-(define (next-state state message models tools)
-  "Given current @var{state} and a new @var{message}, return the next state and a
-list of effects.
+    (state-begin
+     (state-let* ((state (current-state)))
+       (set-current-state (-> state
+                              (add (state-session-input-tokens session-id)
+                                   (focus (in-json "usage" "prompt_tokens")
+                                          (llm-response-json response))
+                                   <>)
+                              (add (state-session-output-tokens session-id)
+                                   (focus (in-json "usage" "completion_tokens")
+                                          (llm-response-json response))
+                                   <>)
+                              (add (state-session-thought-tokens session-id)
+                                   (or (focus (in-json "usage"
+                                                       "completion_tokens_details"
+                                                       "reasoning_tokens")
+                                              (llm-response-json response))
+                                       0)
+                                   <>)
+                              (add (state-session-cache-read-tokens session-id)
+                                   (focus (in-json "usage"
+                                                   "prompt_tokens_details"
+                                                   "cached_tokens")
+                                          (llm-response-json response))
+                                   <>)
+                              (add (state-session-cache-write-tokens session-id)
+                                   (or (focus (in-json "usage"
+                                                       "prompt_tokens_details"
+                                                       "cache_write_tokens")
+                                              (llm-response-json response))
+                                       0)
+                                   <>)
+                              (add (state-session-cost session-id)
+                                   (or (focus (in-json "usage" "cost")
+                                              (llm-response-json response))
+                                       0)
+                                   <>)
+                              ;; Push LLM response onto messages.
+                              (push (state-messages session-id)
+                                    llm-reply
+                                    <>))))
+     (state-let* ((tool-call-effects
+                   (state-append-map (lambda (call-json)
+                                       (tool-call-effects session-id call-json tools))
+                                     tool-calls-json)))
+       (let ((effects (cons (agent-message-chunk session-id
+                                                 ;; Send LLM's text response.
+                                                 (focus (key-ref "content")
+                                                        llm-reply))
+                            tool-call-effects)))
+         (if (null? tool-calls-json)
+             ;; There are no further tool calls,
+             (state-let* ((state (current-state)))
+               (if (not (focus (state-session-cancelling? session-id)
+                               state))
+                   ;; … and a cancellation is not in progress; end turn.
+                   (state-let* ((end-turn-effects
+                                 (end-turn-effects session-id models)))
+                     (state-return (append effects end-turn-effects)))
+                   ;; Else, return what we have so far.
+                   (state-return effects)))
+             ;; Maybe dispatch LLM requests.
+             (state-let* ((llm-requests (llm-requests session-id)))
+               (state-return (append effects llm-requests)))))))))
+
+(define (tool-call-result-effects result)
+  "Return the state-monadic list of effects for new tool call @var{result}."
+  (let ((session-id (tool-call-result-session-id result)))
+    (state-begin
+     (state-let* ((state (current-state)))
+       (set-current-state (-> state
+                              ;; Push tool call result onto messages.
+                              (push (state-messages session-id)
+                                    (tool-call-result-json result)
+                                    <>)
+                              ;; Delete tool call from session tool call list.
+                              (alist-delete-over (state-tool-calls session-id)
+                                                 (tool-call-result-call-id result)
+                                                 <>))))
+     ;; Send a notification for each tool call evaluated.
+     (state-let* ((llm-requests (llm-requests session-id)))
+       (state-return
+        (cons (acp-message `(("jsonrpc" . "2.0")
+                             ("method" . "session/update")
+                             ("params"
+                              ("sessionId" . ,session-id)
+                              ("update"
+                               ;; TODO: Add locations and rawOutput.
+                               ("sessionUpdate" . "tool_call_update")
+                               ("toolCallId" . ,(tool-call-result-call-id result))
+                               ("status" .
+                                ,(if (tool-call-result-success? result)
+                                     "completed"
+                                     "failed"))
+                               ("content" .
+                                ,(vector `(("type" . "content")
+                                           ("content"
+                                            ("type" . "text")
+                                            ("text" .
+                                             ,(focus (key-ref "content")
+                                                     (tool-call-result-json result)))))))))))
+              llm-requests))))))
+
+(define (acp-effects message models tools)
+  "Return the state-monadic list of effects for @var{message}.
 
 @var{tools} is the same as in @code{run-tea-loop}. @var{models} is the same as
 in @code{tea-loop}."
@@ -920,96 +938,73 @@ in @code{tea-loop}."
     (let ((json-message (focus acp-message-json message)))
       (if (focus (key-ref "result") json-message)
           ;; message is a response from the client.
-          (next-state-client-response state json-message)
+          (client-response-effects json-message)
           ;; message is a request/notification from the client.
-          (next-state-client-request (cond
-                                      ;; message is a request from the client.
-                                      ((focus (key-ref "id") json-message)
-                                       => (cut put
-                                               state-client-request-id
-                                               <>
-                                               state))
-                                      ;; message is a notification from the
-                                      ;; client.
-                                      (else state))
-                                     json-message
-                                     models
-                                     tools))))
+          (state-begin
+           ;; When message is a request from the client, put it in the state.
+           (let ((request-id (focus (key-ref "id") json-message)))
+             (state-when request-id
+               (state-let* ((state (current-state)))
+                 (set-current-state (put state-client-request-id
+                                         request-id
+                                         state)))))
+           (client-request-effects json-message models tools)))))
    ((llm-response? message)
-    (next-state-llm-response state message tools models))
+    (llm-response-effects message tools models))
    ((tool-call-result? message)
-    (next-state-tool-call-result state message))))
+    (tool-call-result-effects message))))
 
-(define (tea-loop state llm-base-uri llm-api-key models tools)
+(define* (tea-loop state llm-base-uri llm-api-key models tools
+                   #:optional (events (let ((line (get-line (current-input-port))))
+                                        (and (not (eof-object? line))
+                                             (list (acp-message (json-string->scm line)))))))
   "Run a @acronym{TEA, The Elm Architecture} loop starting with @var{state}.
 
 @var{llm-base-uri}, @var{llm-api-key} and @var{tools} are the same as in
 @code{run-tea-loop}. @var{models} is an association list mapping model IDs to
-@code{<model>} objects, the first element of which is the default model."
+@code{<model>} objects, the first element of which is the default model.
+@var{events} is an internal recursion variable."
   ;; Read a JSON-RPC message, handle it, and loop.
-  (let ((line (get-line (current-input-port))))
-    (unless (eof-object? line)
-      (tea-loop (handle-event (acp-message (json-string->scm line))
-                              state
-                              llm-base-uri
-                              llm-api-key
-                              models
-                              tools)
-                llm-base-uri
-                llm-api-key
-                models
-                tools))))
-
-(define (handle-event event state llm-base-uri llm-api-key models tools)
-  "Handle @var{event} with @var{state} and return a new state.
-
-@var{llm-base-uri}, @var{llm-api-key} and @var{tools} are the same as in
-@code{run-tea-loop}. @var{models} is the same as in @code{tea-loop}."
-  (let-values (((state effects)
-                ;; Compute the next state and collect the effects.
-                (next-state state event models tools)))
-    ;; Do the effects.
-    (fold (cut do-effect <> <> llm-base-uri llm-api-key models tools)
-          state
-          effects)))
-
-(define (do-effect effect state llm-base-uri llm-api-key models tools)
-  "Do @var{effect} with @var{state} and return a new state.
+  (when events
+    (let-values (((effects next-state)
+                  (run-with-state
+                    (state-append-map (cut acp-effects <> models tools)
+                                      events)
+                    state)))
+      (match (append-map (cut do-effect <> llm-base-uri llm-api-key tools)
+                         effects)
+        (()
+         (tea-loop next-state llm-base-uri llm-api-key models tools))
+        ((events ...)
+         (tea-loop next-state llm-base-uri llm-api-key models tools events))))))
+
+(define (do-effect effect llm-base-uri llm-api-key tools)
+  "Do @var{effect}. Return list of events produced by the effect.
 
 @var{llm-base-uri}, @var{llm-api-key} and @var{tools} are the same as in
-@code{run-tea-loop}. @var{models} is the same as in @code{tea-loop}."
+@code{run-tea-loop}."
   (cond
    ;; Send message to client, and return the state unchanged.
    ((acp-message? effect)
     (display (scm->json-string (focus acp-message-json effect)))
     (newline)
     (flush-output-port (current-output-port))
-    state)
+    (list))
    ;; Send request to LLM, handle the response, and return the new state.
    ((llm-request? effect)
     ;; TODO: Handle network failures in OpenAI query.
-    (handle-event (llm-response (llm-request-session-id effect)
-                                (openai-query llm-base-uri
-                                              llm-api-key
-                                              (llm-request-model effect)
-                                              (llm-request-messages effect)
-                                              (map (match-lambda
-                                                     ((name . tool)
-                                                      (tool->spec name tool)))
-                                                   tools)))
-                  state
-                  llm-base-uri
-                  llm-api-key
-                  models
-                  tools))
+    (list (llm-response (llm-request-session-id effect)
+                        (openai-query llm-base-uri
+                                      llm-api-key
+                                      (llm-request-model effect)
+                                      (llm-request-messages effect)
+                                      (map (match-lambda
+                                             ((name . tool)
+                                              (tool->spec name tool)))
+                                           tools)))))
    ;; Evaluate tool, handle the result, and return the new state.
    ((tool-call? effect)
-    (handle-event (eval-tool-call effect tools)
-                  state
-                  llm-base-uri
-                  llm-api-key
-                  models
-                  tools))))
+    (list (eval-tool-call effect tools)))))
 
 (define (run-tea-loop llm-base-uri llm-api-key model tools)
   "Run a @acronym{TEA, The Elm Architecture} loop. @var{llm-base-uri} is the base
diff --git a/kaagum/utils.scm b/kaagum/utils.scm
index b6292c1..dff3b05 100644
--- a/kaagum/utils.scm
+++ b/kaagum/utils.scm
@@ -22,7 +22,6 @@
   #:use-module (ice-9 match)
   #:use-module (ice-9 popen)
   #:export (->
-            foldn
             alist->plist
             call-with-input-pipe))
 
@@ -44,27 +43,6 @@ For example:
 => 16"
   (->-helper x (cut proc ...) ...))
 
-(define (foldn proc lst . inits)
-  "Apply @var{proc} to the elements of @var{lst} to build a result, and return
-that result. @var{proc} may return multiple values, in which case, an equal
-number of values are returned. Each @var{proc} call is @code{(proc element
-previous ...)} where @code{element} is an element of @var{lst}, and
-@code{(previous ...)} is the return from the previous call to @var{proc} or the
-given @var{inits} for the first call. For example,
-
-(foldn (lambda (n sum sum-of-squares)
-         (values (+ sum n)
-                 (+ sum-of-squares (expt n 2))))
-       (iota 10)
-       0 0)
-=> 45
-=> 285"
-  (apply values
-         (fold (lambda (element results)
-                 (call-with-values (cut apply proc element results) list))
-               inits
-               lst)))
-
 (define (alist->plist alist)
   "Convert association list @var{alist} to a property list. Keys in
 @var{alist} are converted to keywords."