Abort hidden streamed tools during forced repair

ober

d399dd6476b5d5e54a0eadb74c54a4f5ba12e7a6

diff --git a/src/jcode/provider/provider.ss b/src/jcode/provider/provider.ss
index a042804..14d702f 100644
--- a/src/jcode/provider/provider.ss
+++ b/src/jcode/provider/provider.ss
@@ -1012,6 +1012,26 @@
     (hash-put! choice "function" fn)
     choice))
 
+(def (hashish-get ht key)
+  (guard (_ [else #f])
+    (hash-get ht key)))
+
+(def (openai-tool-choice-name choice)
+  (let ((fn (hashish-get choice "function")))
+    (hashish-get fn "name")))
+
+(def forced-tool-choice-retry-marker
+  "JCODE_INTERNAL_FORCE_FIRST_EDIT_RETRY")
+
+(def (forced-tool-choice-retry-content expected actual)
+  (string-append
+    forced-tool-choice-retry-marker
+    " provider streamed hidden tool "
+    actual
+    " while this turn required "
+    expected
+    "."))
+
 (def (openai-max-tokens provider)
   (let* ((configured
            (or (positive-int-value (getenv "JCODE_MAX_TOKENS"))
@@ -2367,7 +2387,10 @@
          (logprob-box (box '()))
          (entropy-box (box '()))
          (finish-reason-box (box #f))
-         (saw-done-box (box #f)))
+         (saw-done-box (box #f))
+         (forced-tool-name
+           (openai-tool-choice-name (current-tool-choice-override)))
+         (forced-tool-mismatch-box (box #f)))
     (let* ((body-json (json-object->string body))
            (dummy (begin
                     (log-info logger "stream-request"
@@ -2378,6 +2401,10 @@
                           (headers . ,(redact-headers headers))
                           (body . ,body-json))))))
            (http-status
+             (guard (e [#t
+                         (if (unbox forced-tool-mismatch-box)
+                           200
+                           (raise e))])
            (jcode-http-post-stream url headers body-json
       (lambda (event-str)
         (when event-str
@@ -2443,7 +2470,18 @@
                                  (when fn
                                    (let ((name (hash-get fn "name"))
                                          (args (hash-get fn "arguments")))
-                                     (when name (hash-put! acc "name" name))
+                                     (when name
+                                       (when (and forced-tool-name
+                                                  (string? name)
+                                                  (not (string=? name forced-tool-name)))
+                                         (set-box! forced-tool-mismatch-box
+                                                   (cons forced-tool-name name))
+                                         (log-warn logger "forced-tool-stream-mismatch"
+                                           `((expected . ,forced-tool-name)
+                                             (actual . ,name)))
+                                         (error 'openai-stream-chat
+                                           "provider streamed hidden tool during forced tool turn"))
+                                       (hash-put! acc "name" name))
                                      (when args
                                        (hash-put! acc "args"
                                          (string-append
@@ -2452,7 +2490,7 @@
                              tcs))))))))))))
              (lambda (status response-body)
                (raise-provider-api-error (provider-name provider)
-                                         status response-body)))))
+                                         status response-body))))))
       (unless (= http-status 200)
         (log-error logger "stream-http-error"
           `((status . ,http-status) (url . ,url)))))
@@ -2460,7 +2498,9 @@
     ;; sent either the OpenAI [DONE] sentinel or a choice finish_reason. Local
     ;; gateways can close mid-generation while still returning HTTP 200; do
     ;; not turn that partial reasoning into an assistant turn.
-    (unless (or (unbox saw-done-box) (unbox finish-reason-box))
+    (unless (or (unbox forced-tool-mismatch-box)
+                (unbox saw-done-box)
+                (unbox finish-reason-box))
       (error 'openai-stream-chat
         "stream closed before a terminal event (finish_reason or [DONE])"))
     ;; Build result
@@ -2471,11 +2511,18 @@
            ;; otherwise trip the auto-escalator and produce an invalid
            ;; assistant message for downstream providers).
            (initial-content (cond
+                              ((unbox forced-tool-mismatch-box)
+                               (let ((mismatch (unbox forced-tool-mismatch-box)))
+                                 (forced-tool-choice-retry-content
+                                   (car mismatch)
+                                   (cdr mismatch))))
                               ((> (string-length raw-content) 0) raw-content)
                               ((> (string-length reasoning) 0)
                                (string-append "<think>" reasoning "</think>"))
                               (else "")))
-           (indices (list-sort < (hash-keys tc-table)))
+           (indices (if (unbox forced-tool-mismatch-box)
+                      '()
+                      (list-sort < (hash-keys tc-table))))
            (struct-tool-calls
              (map (lambda (idx)
                     (let ((acc (hash-ref tc-table idx)))
diff --git a/test/run.ss b/test/run.ss
index e3f9f91..d6addf5 100644
--- a/test/run.ss
+++ b/test/run.ss
@@ -13086,6 +13086,50 @@
       (tcp-close srv))))
 
 (let* ([sse-body
+         (string-append
+           "data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning_content\":\"I should check balance, but writing instead.\"},\"finish_reason\":null}]}\n\n"
+           "data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_bad\",\"type\":\"function\",\"function\":{\"name\":\"write\",\"arguments\":\"\"}}]},\"finish_reason\":null}]}\n\n"
+           "data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"path\\\":\\\"tetris.ss\\\",\\\"content\\\":\\\"large ignored write\\\"}\"}}]},\"finish_reason\":null}]}\n\n")]
+       [srv (tcp-listen "127.0.0.1" 0)]
+       [base-url (format "http://127.0.0.1:~a/v1" (tcp-server-port srv))]
+       [captured (vector #f)]
+       [balance-tool (uuid-recovery-tool "balance" (args "path" (args "type" "string")) '("path"))]
+       [old-max-tokens (getenv "JCODE_MAX_TOKENS")])
+  (dynamic-wind
+    (lambda () (putenv "JCODE_MAX_TOKENS" ""))
+    (lambda ()
+      (serve-one-captured-sse! srv captured sse-body)
+      (let* ([p (make-provider "mlx" "" "deepseek-v4-pro" base-url)]
+             [result (call-with-values
+                       (lambda ()
+                         (parameterize ([current-tool-choice-override
+                                         (openai-function-tool-choice "balance")])
+                           (provider-stream-chat
+                             p
+                             (list (make-user-message "hi"))
+                             (list balance-tool)
+                             (lambda (_token) #f))))
+                       (lambda (content tcs usage) (list content tcs usage)))]
+             [req (vector-ref captured 0)])
+        (check-pred! "mlx forced streaming hidden write returns retry marker"
+                     (car result)
+                     (lambda (s)
+                       (and (string? s)
+                            (str-contains? s
+                              "JCODE_INTERNAL_FORCE_FIRST_EDIT_RETRY")
+                            (str-contains? s "hidden tool write")
+                            (str-contains? s "required balance"))))
+        (check! "mlx forced streaming hidden write has no tool call"
+                (list-ref result 1) '())
+        (check! "mlx forced streaming request required balance"
+                (and req
+                     (str-contains? req "\"tool_choice\"")
+                     (str-contains? req "\"name\":\"balance\"")) #t)))
+    (lambda ()
+      (putenv "JCODE_MAX_TOKENS" (or old-max-tokens ""))
+      (tcp-close srv))))
+
+(let* ([sse-body
          "data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning_content\":\"partial\"},\"finish_reason\":null}]}\n\n"]
        [srv (tcp-listen "127.0.0.1" 0)]
        [base-url (format "http://127.0.0.1:~a/v1" (tcp-server-port srv))]