Force balance after rejected create drafts

Jaime Fournier <jaimef@linbsd.org>

eae5c6b1c06b6fe1137c08afa247b31891d6642b

diff --git a/src/jcode/core/verified-run.ss b/src/jcode/core/verified-run.ss
index b07c476..9593bda 100644
--- a/src/jcode/core/verified-run.ss
+++ b/src/jcode/core/verified-run.ss
@@ -447,10 +447,16 @@
        (openai-function-tool-choice "edit")))
 
 (def (force-staged-create-repair-tool-name specs)
-  ;; Syntax-broken missing-file creates have no on-disk artifact yet. Let the
-  ;; model choose between a bounded repair and a corrected complete file body;
-  ;; forcing line_edit made local models fight the natural full-retry path.
-  #f)
+  ;; Syntax-broken missing-file creates have no on-disk artifact yet. The
+  ;; retained draft is the only inspectable artifact, and weaker reasoning
+  ;; models often burn the retry turn manually counting parens from memory.
+  ;; Force one compact balance read first, then return the broader staged
+  ;; repair tools on the following turn.
+  (and (syntax-broken-missing-create-repair-mode?)
+       (= (current-rejected-draft-inspections) 0)
+       (= (current-rejected-draft-path-reject-count) 1)
+       (tool-specs-include? specs "balance")
+       "balance"))
 
 (def (forced-first-edit-specs specs forced-choice)
   (if forced-choice
@@ -459,6 +465,13 @@
             specs)
     specs))
 
+(def (forced-staged-repair-specs specs staged-tool)
+  (if staged-tool
+    (filter (lambda (spec)
+              (string=? (tool-spec-name spec) staged-tool))
+            specs)
+    specs))
+
 (def (forced-tool-choice-response response forced-choice allowed-tools)
   (if (and forced-choice
 	           (pair? response)
@@ -476,6 +489,9 @@
 (def (forced-first-edit-response response forced-choice)
   (forced-tool-choice-response response forced-choice '("edit" "write")))
 
+(def (forced-staged-repair-response response staged-tool)
+  (forced-tool-choice-response response staged-tool (list staged-tool)))
+
 (def (forced-first-edit-messages messages specs forced-specs)
   (with-visible-tool-system-note messages specs forced-specs))
 
@@ -520,14 +536,15 @@
 	                            (force-first-edit-tool-choice specs))))
 		                 (forced-specs
 		                   (if staged-tool
-		                     (filter verified-current-staged-repair-tool-spec?
-		                             specs)
+		                     (forced-staged-repair-specs specs staged-tool)
 		                     (forced-first-edit-specs specs forced-choice))))
 		            (if staged-tool
-		              (backend (forced-first-edit-messages
-		                         messages specs forced-specs)
-	                       forced-specs
-	                       step)
+		              (forced-staged-repair-response
+		                (backend (forced-first-edit-messages
+		                           messages specs forced-specs)
+	                         forced-specs
+	                         step)
+		                staged-tool)
 		              (forced-first-edit-response
 	              (backend (forced-first-edit-messages
 	                         messages specs forced-specs)
@@ -563,18 +580,19 @@
 	                     (forced-choice (or staged-choice first-edit-choice))
 				                     (forced-specs
 				                       (if staged-tool
-				                         (filter verified-current-staged-repair-tool-spec?
-				                                 specs)
+				                         (forced-staged-repair-specs specs staged-tool)
 			                         (forced-first-edit-specs specs forced-choice))))
 		                (parameterize
 		                  ((current-tool-choice-override
 		                     (and (specified-tool-choice-supported? provider)
 	                          forced-choice)))
 			                  (if staged-tool
-			                    (backend (forced-first-edit-messages
-			                               messages specs forced-specs)
-			                             forced-specs
-			                             #f)
+			                    (forced-staged-repair-response
+			                      (backend (forced-first-edit-messages
+			                                 messages specs forced-specs)
+			                               forced-specs
+			                               #f)
+			                      staged-tool)
 			                    (forced-first-edit-response
 		                    (backend (forced-first-edit-messages
 		                               messages specs forced-specs)
@@ -3023,6 +3041,10 @@
     "The file was not written. Repair the retained draft for "
     path
     " with line_edit, replace_range, replace_def, exact old_str replacement, or edit/write with complete corrected contents. "
+    "Do not reason through the whole file in prose; code in reasoning or markdown is discarded. "
+    "Next structured call should be balance(path=\""
+    path
+    "\") for the retained draft unless the reported line fix is already obvious. "
     "Use the reported line/span directly when inspection tools are not listed in the current tool schema. "
     "MCP file-edit tools cannot repair this state because the file does not exist on disk. "
     "Do not call run/list/verify until that file exists."))
diff --git a/test/run.ss b/test/run.ss
index 1c0580a..acb00c9 100644
--- a/test/run.ss
+++ b/test/run.ss
@@ -3086,16 +3086,12 @@
 	                           (list (cons "path" target)
 	                                 (cons "content" bad1))
 	                           #f))
-	                   (list (make-wtool-call "edit"
-	                           (list (cons "path" target)
-	                                 (cons "content" bad2))
+	                   (list (make-wtool-call "balance"
+	                           (list (cons "path" target))
 	                           #f))
 	                   (list (make-wtool-call "edit"
 	                           (list (cons "path" target)
-	                                 (cons "content" good))
-	                           #f))
-	                   (list (make-wtool-call "balance"
-	                           (list (cons "path" target))
+	                                 (cons "content" bad2))
 	                           #f))
 	                   (list (make-wtool-call "edit"
 	                           (list (cons "path" target)
@@ -3851,28 +3847,33 @@
 	       [bad2 "(import (jerboa prelude))\n(def (main) (displayln \"two\")))\n"]
 	       [good-line "(def (main) (displayln \"two\"))"]
 	       [staged-tool-names '()]
+	       [calls 0]
 	       [slurp (lambda (p)
 	                (call-with-input-file p (lambda (i) (get-string-all i))))])
 	  (safe-delete-test-file! target-path)
 	  (let* ([resp
-	           (lambda (_messages specs step)
-	             (case step
-	               [(0)
+	           (lambda (_messages specs _step)
+	             (set! calls (+ calls 1))
+	             (case calls
+	               [(1)
 	                (list (make-wtool-call "write"
 	                        (list (cons "path" target)
 	                              (cons "content" bad1)) #f))]
-	               [(1)
+	               [(2)
+	                (list (make-wtool-call "balance"
+	                        (list (cons "path" target)) #f))]
+	               [(3)
 	                (list (make-wtool-call "write"
 	                        (list (cons "path" target)
 	                              (cons "content" bad2)) #f))]
-	               [(2)
+	               [(4)
 	                (let ([names (map tool-spec-name specs)])
 	                  (set! staged-tool-names names)
 	                  (list (make-wtool-call "line_edit"
 	                          (list (cons "path" target)
 	                                (cons "line" 2)
 	                                (cons "content" good-line)) #f)))]
-	               [(3) (list (make-wtool-call "verify" '() #f))]
+	               [(5) (list (make-wtool-call "verify" '() #f))]
 	               [else (error 'test "unexpected staged repair step")]))]
 	         [result
 	           (verified-run resp "repair a repeatedly rejected local draft"
@@ -4619,16 +4620,18 @@
 	       [responder
 	         (lambda (_messages tools _step)
 	           (set! calls (+ calls 1))
-	           (when (= calls 3)
+	           (when (= calls 4)
 	             (set! tool-names-after-repeat (map tool-spec-name tools)))
 	           (case calls
 	             [(1) (list (make-wtool-call "write"
 	                         (list (cons "path" target)
 	                               (cons "content" bad-a)) #f))]
-	             [(2) (list (make-wtool-call "write"
+	             [(2) (list (make-wtool-call "balance"
+	                         (list (cons "path" target)) #f))]
+	             [(3) (list (make-wtool-call "write"
 	                         (list (cons "path" target)
 	                               (cons "content" bad-b)) #f))]
-	             [(3) (list (make-wtool-call "line_edit"
+	             [(4) (list (make-wtool-call "line_edit"
 	                         (list (cons "path" target)
 	                               (cons "line" 2)
 	                               (cons "content" "(def (main) (displayln \"fixed\"))")) #f))]
@@ -4727,16 +4730,18 @@
 	       [responder
 	         (lambda (_messages tools _step)
 	           (set! calls (+ calls 1))
-	           (when (= calls 3)
+	           (when (= calls 4)
 	             (set! tool-names-after-repeat (map tool-spec-name tools)))
 	           (case calls
 	             [(1) (list (make-wtool-call "write"
 	                         (list (cons "path" target)
 	                               (cons "content" bad-a)) #f))]
-	             [(2) (list (make-wtool-call "write"
+	             [(2) (list (make-wtool-call "balance"
+	                         (list (cons "path" target)) #f))]
+	             [(3) (list (make-wtool-call "write"
 	                         (list (cons "path" target)
 	                               (cons "content" bad-b)) #f))]
-	             [(3) (list (make-wtool-call "replace_range"
+	             [(4) (list (make-wtool-call "replace_range"
 	                         (list (cons "path" target)
 	                               (cons "start" 1)
 	                               (cons "end" 3)
@@ -4773,18 +4778,20 @@
 	       [responder
 	         (lambda (_messages tools _step)
 	           (set! calls (+ calls 1))
-	           (when (= calls 3)
+	           (when (= calls 4)
 	             (set! tool-names-after-repeat (map tool-spec-name tools)))
 	           (case calls
 	             [(1) (list (make-wtool-call "write"
 	                         (list (cons "path" target)
 	                               (cons "content" bad-a)) #f))]
-	             [(2) (list (make-wtool-call "write"
+	             [(2) (list (make-wtool-call "balance"
+	                         (list (cons "path" target)) #f))]
+	             [(3) (list (make-wtool-call "write"
 	                         (list (cons "path" target)
 	                               (cons "content" bad-b)) #f))]
-	             [(3) (list (make-wtool-call "read"
+	             [(4) (list (make-wtool-call "read"
 	                         (list (cons "path" target)) #f))]
-	             [(4) (list (make-wtool-call "line_edit"
+	             [(5) (list (make-wtool-call "line_edit"
 	                         (list (cons "path" target)
 	                               (cons "line" 3)
 	                               (cons "content" "  (displayln \"fixed\"))")) #f))]
@@ -8204,10 +8211,13 @@
                       (list (cons "path" target)
                             (cons "content" bad)) #f))]
              [(= i 2)
+              (list (make-wtool-call "balance"
+                      (list (cons "path" target)) #f))]
+             [(= i 3)
               (list (make-wtool-call "edit"
                       (list (cons "path" target)
                             (cons "content" good)) #f))]
-             [(= i 3)
+             [(= i 4)
               (list (make-wtool-call "verify" '() #f))]
              [else
               (list (make-wtool-call "done"
@@ -8225,11 +8235,16 @@
                            (cons 'max-tool-errors 3))))])
     (check! "verified-run: missing create schema repair reaches done"
             result "missing-create-schema-ok")
-    (check-pred! "verified-run: syntax create repair keeps draft inspection"
+    (check-pred! "verified-run: syntax create repair forces draft balance first"
       (reverse seen-specs)
       (lambda (xs)
         (and (>= (length xs) 2)
-             (let ([names (list-ref xs 1)])
+             (equal? (list-ref xs 1) '("balance")))))
+    (check-pred! "verified-run: syntax create repair restores draft inspection"
+      (reverse seen-specs)
+      (lambda (xs)
+        (and (>= (length xs) 3)
+             (let ([names (list-ref xs 2)])
                (and (member "line_edit" names)
                     (member "replace_def" names)
                     (member "replace_range" names)
@@ -8391,12 +8406,15 @@
 	                        (list (cons "path" target)
 	                              (cons "content" broken)) #f))]
 	               [(= i 2)
+	                (list (make-wtool-call "balance"
+	                        (list (cons "path" target)) #f))]
+	               [(= i 3)
 	                (list (make-wtool-call "replace_range"
 	                        (list (cons "path" target)
 	                              (cons "start" 1)
 	                              (cons "end" 4)
 	                              (cons "content" fixed)) #f))]
-	               [(= i 3)
+	               [(= i 4)
 	                (list (make-wtool-call "verify" '() #f))]
 	               [else
 	                (list (make-wtool-call "done"
@@ -8412,11 +8430,16 @@
 	            result "VERIFIED: exit 0\n")
 	    (check! "verified-run: syntax create staged repair writes file"
 	            (slurp target-path) (string-append fixed "\n"))
-			    (check-pred! "verified-run: syntax create staged schema keeps full retry"
+			    (check-pred! "verified-run: syntax create staged schema forces balance first"
 			      (reverse seen-specs)
 			      (lambda (xs)
 			        (and (>= (length xs) 2)
-			             (let ([names (list-ref xs 1)])
+			             (equal? (list-ref xs 1) '("balance")))))
+			    (check-pred! "verified-run: syntax create staged schema restores full retry"
+			      (reverse seen-specs)
+			      (lambda (xs)
+			        (and (>= (length xs) 3)
+			             (let ([names (list-ref xs 2)])
 			               (and (member "line_edit" names)
 			                    (member "replace_range" names)
 			                    (member "write" names)
@@ -8443,10 +8466,13 @@
 		                        (list (cons "path" target)
 		                              (cons "content" broken)) #f))]
 			               [(= i 2)
+			                (list (make-wtool-call "balance"
+			                        (list (cons "path" target)) #f))]
+			               [(= i 3)
 			                (list (make-wtool-call "write"
 			                        (list (cons "path" target)
 			                              (cons "content" fixed)) #f))]
-			               [(= i 3)
+			               [(= i 4)
 			                (list (make-wtool-call "verify" '() #f))]
 			               [else
 			                (list (make-wtool-call "done"
@@ -9082,6 +9108,75 @@
 		  (safe-delete-test-file! target-path))
 
 		(let* ([vr-dir  "/tmp"]
+		       [target "jcode-verified-force-balance-after-create-reject.ss"]
+		       [target-path (string-append vr-dir "/" target)]
+		       [bad "(import (jerboa prelude))\n(define (main) 1))\n"]
+		       [good "(import (jerboa prelude))\n(define (main) 1)\n"]
+		       [seen-specs '()]
+		       [step 0]
+		       [slurp (lambda (p)
+		                (call-with-input-file p
+		                  (lambda (i) (get-string-all i))))])
+		  (safe-delete-test-file! target-path)
+		  (let* ([scope (parse-write-scope target)]
+		         [wf (coding-workflow (string-append "test -s " target)
+		                              vr-dir
+		                              (list (cons 'write-scope scope)))]
+		         [backend
+		           (lambda (messages specs current-step)
+		             (set! seen-specs
+		               (cons (map tool-spec-name specs) seen-specs))
+		             (set! step (+ step 1))
+		             (cond
+		               [(= step 1)
+		                (list
+		                  (make-wtool-call
+		                    "edit"
+		                    (list (cons "path" target)
+		                          (cons "content" bad))
+		                    #f))]
+		               [(= step 2)
+		                (list
+		                  (make-wtool-call
+		                    "balance"
+		                    (list (cons "path" target))
+		                    #f))]
+		               [(= step 3)
+		                (list
+		                  (make-wtool-call
+		                    "edit"
+		                    (list (cons "path" target)
+		                          (cons "content" good))
+		                    #f))]
+		               [(= step 4)
+		                (list (make-wtool-call "verify" '() #f))]
+		               [else
+		                (list
+		                  (make-wtool-call
+		                    "done"
+		                    '(("summary" . "forced-balance-after-create-reject-ok"))
+		                    #f))]))]
+		         [result
+		           (parameterize ((current-write-scope scope))
+		             (run-workflow wf
+		                           "force balance after rejected create"
+		                           (provider-responder backend)
+		                           (list (cons 'max-iterations 8)
+		                                 (cons 'max-tool-errors 3))))])
+		    (check! "verified-run: forced balance after rejected create reaches done"
+		            result
+		            "forced-balance-after-create-reject-ok")
+		    (check! "verified-run: forced balance repair writes final source"
+		            (slurp target-path)
+		            good)
+		    (check-pred! "verified-run: rejected create exposes only balance next"
+		      (reverse seen-specs)
+		      (lambda (xs)
+		        (and (>= (length xs) 2)
+		             (equal? (cadr xs) '("balance"))))))
+		  (safe-delete-test-file! target-path))
+
+		(let* ([vr-dir  "/tmp"]
 		       [target "jcode-verified-balance-autoclose.ss"]
 		       [target-path (string-append vr-dir "/" target)]
 		       [bad "(import (jerboa prelude))\n(define (main)\n  (displayln \"draft\")\n"]
@@ -9961,6 +10056,9 @@
 	                      (list (cons "path" target)
 	                            (cons "content" bad)) #f))]
 	             [(2)
+	              (list (make-wtool-call "balance"
+	                      (list (cons "path" target)) #f))]
+	             [(3)
 	              (list (make-wtool-call "write"
 	                      (list (cons "content" good)) #f))]
 	             [else