From a9821d120e329d0d8b1709f1dc13052ee8cdb85a Mon Sep 17 00:00:00 2001
From: nyoungo <nyoungo@users.noreply.github.com>
Date: Thu, 8 Oct 2026 01:07:04 +0800
Subject: [PATCH] fix: accept prose-prefixed tool JSON and heal invalid escape
 sequences

---
 internal/openai/continue.go      | 54 +++++++++++++++++++++++++++++---
 internal/openai/continue_test.go | 39 +++++++++++++++++++++++
 internal/openai/toolcalls.go     |  8 +++++
 3 files changed, 97 insertions(+), 4 deletions(-)

diff --git a/internal/openai/continue.go b/internal/openai/continue.go
index 1d5649e..ceaeae9 100644
--- a/internal/openai/continue.go
+++ b/internal/openai/continue.go
@@ -64,7 +64,13 @@ func ToolCallTruncated(text string) bool {
 		}
 	}
 	if !strings.HasPrefix(cand, `{"tool_calls"`) {
-		return false
+		// The model sometimes opens with a sentence before the tool JSON:
+		// evaluate the payload itself, mirroring TryParseToolCalls.
+		i := strings.Index(cand, `{"tool_calls"`)
+		if i <= 0 {
+			return false
+		}
+		cand = cand[i:]
 	}
 	return !json.Valid([]byte(cand))
 }
@@ -118,8 +124,22 @@ func (st *escapeState) feed(s string, out *strings.Builder) {
 		}
 		switch c {
 		case '\\':
-			st.escaped = true
-			out.WriteByte(c)
+			if i+1 >= len(s) {
+				// dangling escape: carried into the next feed so a
+				// continuation piece can complete it
+				st.escaped = true
+				out.WriteByte('\\')
+				continue
+			}
+			if j := escapeEndAt(s, i); j > i {
+				out.WriteString(s[i : j+1])
+				i = j
+				continue
+			}
+			// invalid escape (…\d): double the backslash so the text
+			// parses — matches escapeStringBody used by forceCloseJSON,
+			// keeping streamed and repaired arguments identical
+			out.WriteString(`\\`)
 		case '"':
 			st.inString = false
 			out.WriteByte(c)
@@ -140,10 +160,36 @@ func (st *escapeState) feed(s string, out *strings.Builder) {
 }
 
 func needsEscape(s string) bool {
+	inString, escaped := false, false
 	for i := 0; i < len(s); i++ {
-		if s[i] < 0x20 {
+		c := s[i]
+		if c < 0x20 {
 			return true
 		}
+		if !inString {
+			if c == '"' {
+				inString = true
+			}
+			continue
+		}
+		if escaped {
+			escaped = false
+			continue
+		}
+		switch c {
+		case '\\':
+			if i+1 >= len(s) {
+				escaped = true
+				continue
+			}
+			j := escapeEndAt(s, i)
+			if j == i {
+				return true // invalid escape — feed doubles it
+			}
+			i = j
+		case '"':
+			inString = false
+		}
 	}
 	return false
 }
diff --git a/internal/openai/continue_test.go b/internal/openai/continue_test.go
index 111952c..0acadcd 100644
--- a/internal/openai/continue_test.go
+++ b/internal/openai/continue_test.go
@@ -24,6 +24,10 @@ func TestToolCallTruncated(t *testing.T) {
 		{fenced, true},
 		{"", false},
 		{`{"other": true}`, false},
+		// prose before the payload: mirrored from TryParseToolCalls
+		{"我先说明一下：\n\n" + valid, false},
+		{"我先说明一下：\n\n" + truncated, true},
+		{"我先说明一下：\n\n" + fenced, true},
 	}
 	for _, c := range cases {
 		if got := ToolCallTruncated(c.in); got != c.want {
@@ -64,6 +68,41 @@ func TestEscapeControlChars(t *testing.T) {
 	}
 }
 
+func TestEscapeControlCharsDoublesInvalidEscapes(t *testing.T) {
+	cases := []struct{ in, want string }{
+		// \q is not a JSON escape — the backslash must be doubled
+		{`{"a": "x\qy"}`, `{"a": "x\\qy"}`},
+		// valid pairs survive, invalid ones heal
+		{`{"a": "ok\nstill\q"}`, `{"a": "ok\nstill\\q"}`},
+		// dangling escape at the end is carried, not doubled
+		{`{"a": "dangling\`, `{"a": "dangling\`},
+		// no backslash at all: untouched
+		{`{"a": "plain"}`, `{"a": "plain"}`},
+	}
+	for _, c := range cases {
+		if got := EscapeControlChars(c.in); got != c.want {
+			t.Errorf("EscapeControlChars(%q) = %q, want %q", c.in, got, c.want)
+		}
+	}
+	if !json.Valid([]byte(EscapeControlChars(`{"a": "x\qy"}`))) {
+		t.Fatal("healed text is not valid JSON")
+	}
+}
+
+// TestInnerEscapeMatchesParseHeal pins the contract that keeps streamed tool
+// arguments identical to the final repaired parse.
+func TestInnerEscapeMatchesParseHeal(t *testing.T) {
+	inner := `{"content": "a\qb` + "\n" + `c\qd"}`
+	streamed := innerEscape(inner)
+	healed := EscapeControlChars(inner)
+	if streamed != healed {
+		t.Fatalf("stream/parse heal mismatch:\n stream=%q\n parse =%q", streamed, healed)
+	}
+	if !json.Valid([]byte(healed)) {
+		t.Fatalf("healed inner invalid: %q", healed)
+	}
+}
+
 func TestTrimOverlap(t *testing.T) {
 	cases := []struct{ acc, head, want string }{
 		{"...fillStyle='#F4", "F4E2B8';draw()", "E2B8';draw()"},
diff --git a/internal/openai/toolcalls.go b/internal/openai/toolcalls.go
index e183ef1..3b589ca 100644
--- a/internal/openai/toolcalls.go
+++ b/internal/openai/toolcalls.go
@@ -17,6 +17,9 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) {
 	var candidates []string
 	if strings.HasPrefix(trimmed, "{") {
 		candidates = append(candidates, trimmed)
+	} else if i := strings.Index(trimmed, `{"tool_calls"`); i > 0 {
+		// prose before the tool JSON: start the candidate at the payload
+		candidates = append(candidates, strings.TrimSpace(trimmed[i:]))
 	}
 	if strings.Contains(trimmed, "```json") {
 		for _, part := range strings.Split(trimmed, "```json")[1:] {
@@ -45,6 +48,11 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) {
 		if !decodeToolPayload(cand, &payload) {
 			continue
 		}
+		// A prose-prefix candidate runs to the end of the text; keep only
+		// the payload itself so prose after the JSON survives in remaining.
+		if full, ok := jsonPrefix(cand); ok && len(full) < len(cand) {
+			cand = full
+		}
 		if payload.ToolCalls == nil {
 			continue
 		}