From 975882292c786b0f71974fe7ecd16138646bab30 Mon Sep 17 00:00:00 2001
From: nyoungo <nyoungo@users.noreply.github.com>
Date: Wed, 7 Oct 2026 22:41:18 +0800
Subject: [PATCH] feat: force-close truncated tool-call JSON into partial tool
 calls

---
 config.example.toml              |   4 +-
 internal/openai/continue_test.go |  40 +++++++
 internal/openai/forceclose.go    | 174 +++++++++++++++++++++++++++++++
 internal/openai/toolcalls.go     |  14 +++
 4 files changed, 231 insertions(+), 1 deletion(-)
 create mode 100644 internal/openai/forceclose.go

diff --git a/config.example.toml b/config.example.toml
index c9cbd3b..f9a12e2 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -44,7 +44,9 @@ user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTM
 # 本站点配额耗尽时自动切换到其他站点的同款模型。
 cross_site_failover = true
 # 站点对单次回复有约 1k token 的硬上限，大参数的工具调用 JSON 会被截断。
-# 自动向上游追加「续写」轮次把 JSON 拼完整；0 = 关闭。
+# 自动向上游追加「续写」轮次把 JSON 拼完整；若最终仍不闭合，会在截断处
+# 强行补全闭合，返回内容写到哪算哪的部分工具调用，由客户端下一轮继续写。
+# 0 = 关闭续写（每次截断都直接强行补全返回）。
 continue_rounds = 20
 
 [turnstile]
diff --git a/internal/openai/continue_test.go b/internal/openai/continue_test.go
index fc385db..111952c 100644
--- a/internal/openai/continue_test.go
+++ b/internal/openai/continue_test.go
@@ -335,3 +335,43 @@ func TestTryParseToolCallsShortCloseSequence(t *testing.T) {
 		t.Fatalf("content = %q, want html", obj.Content)
 	}
 }
+
+func TestTryParseToolCallsLeakedQuoteCutoff(t *testing.T) {
+	raw := `{"tool_calls": [{"id": "call_001", "type": "function", "function": {"name": "write", "arguments": "{\"filePath\": \"tihu.html\", \"content\": \"<html><body>canvas scene leak "and then cut`
+	_, calls, ok := TryParseToolCalls(raw)
+	if !ok {
+		t.Fatal("TryParseToolCalls ok = false, want force-closed tool call")
+	}
+	args := calls[0].Function.Arguments
+	if !json.Valid([]byte(args)) {
+		t.Fatalf("args invalid after force close: %s", args)
+	}
+	var inner struct {
+		FilePath string `json:"filePath"`
+		Content  string `json:"content"`
+	}
+	if err := json.Unmarshal([]byte(args), &inner); err != nil {
+		t.Fatal(err)
+	}
+	if inner.FilePath != "tihu.html" {
+		t.Fatalf("filePath = %q", inner.FilePath)
+	}
+	if !strings.Contains(inner.Content, "canvas scene leak") || !strings.Contains(inner.Content, `"and then cut`) {
+		t.Fatalf("content lost pre-cut text: %q", inner.Content)
+	}
+}
+
+func TestTryParseToolCallsPlainCutoffReturnsPartial(t *testing.T) {
+	raw := `{"tool_calls": [{"id": "c", "type": "function", "function": {"name": "write", "arguments": "{\"filePath\": \"a.html\", \"content\": \"<html><body><canvas id=\"scene\">stopped mid`
+	_, calls, ok := TryParseToolCalls(raw)
+	if !ok {
+		t.Fatal("TryParseToolCalls ok = false, want partial tool call")
+	}
+	args := calls[0].Function.Arguments
+	if !json.Valid([]byte(args)) {
+		t.Fatalf("args invalid: %s", args)
+	}
+	if !strings.Contains(args, "stopped mid") {
+		t.Fatalf("args lost tail: %s", args)
+	}
+}
diff --git a/internal/openai/forceclose.go b/internal/openai/forceclose.go
new file mode 100644
index 0000000..7050dd8
--- /dev/null
+++ b/internal/openai/forceclose.go
@@ -0,0 +1,174 @@
+package openai
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+// forceCloseJSON repairs a tool-call document the channel cap cut off so
+// badly that suffix balancing cannot fix it — e.g. a raw quote leaked
+// mid-string and collapsed the structure after it. The fast path reuses
+// balanceJSON (append-only damage). Otherwise the rebuild re-enters the last
+// opened string literal, escapes every raw quote, stray backslash and
+// control character in the dangling tail, then appends the closers the cut
+// chopped off. The result is always structurally valid JSON carrying
+// whatever content survived to the truncation point, so callers deliver a
+// complete tool call with a partial payload instead of raw text the client
+// cannot parse. ok=false means no repair validated.
+func forceCloseJSON(s string) (string, bool) {
+	if b := balanceJSON(s); b != s {
+		return b, true
+	}
+
+	type anchor struct {
+		pos   int
+		stack []byte
+	}
+	var anchors []anchor
+	var stack []byte
+	inString, escaped := false, false
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		if !inString {
+			switch c {
+			case '"':
+				inString = true
+				anchors = append(anchors, anchor{pos: i, stack: append([]byte(nil), stack...)})
+			case '{', '[':
+				stack = append(stack, c)
+			case '}', ']':
+				if len(stack) > 0 {
+					stack = stack[:len(stack)-1]
+				}
+			}
+			continue
+		}
+		if escaped {
+			escaped = false
+			continue
+		}
+		switch c {
+		case '\\':
+			escaped = true
+		case '"':
+			inString = false
+		}
+	}
+	if len(anchors) == 0 {
+		return "", false
+	}
+
+	// Re-anchor on the most recent string literal first: it keeps the most
+	// surrounding structure verbatim, and anything further out folds into
+	// the escaped tail. Each anchor gets both closing forms.
+	const maxAnchors = 512
+	start := len(anchors) - maxAnchors
+	if start < 0 {
+		start = 0
+	}
+	const quoteByte = byte(34)
+	for a := len(anchors) - 1; a >= start; a-- {
+		an := anchors[a]
+		body := escapeStringBody(s[an.pos+1:])
+		for _, isKey := range []bool{false, true} {
+			// The key form requires a bare key: a quote in the dangling
+			// text means this anchor is garbage (e.g. an invalid value
+			// such as 12.), and inventing a nonsense key would silently
+			// discard real data instead of rejecting the reply.
+			if isKey && strings.IndexByte(body, quoteByte) >= 0 {
+				continue
+			}
+			b := strings.Builder{}
+			b.Grow(len(s) + len(body) + len(an.stack) + 8)
+			b.WriteString(s[:an.pos+1])
+			b.WriteString(body)
+			if isKey {
+				b.WriteString("\": \"\"")
+			} else {
+				b.WriteString("\"")
+			}
+			writeClosers(&b, an.stack)
+			if json.Valid([]byte(b.String())) {
+				return b.String(), true
+			}
+		}
+	}
+	return "", false
+}
+
+func writeClosers(b *strings.Builder, stack []byte) {
+	for i := len(stack) - 1; i >= 0; i-- {
+		if stack[i] == '{' {
+			b.WriteByte('}')
+		} else {
+			b.WriteByte(']')
+		}
+	}
+}
+
+// escapeStringBody makes dangling text safe inside a JSON string literal:
+// raw quotes become escaped quotes, backslashes that are not valid JSON
+// escapes get doubled, and control characters become the usual two-character
+// escape sequences.
+func escapeStringBody(s string) string {
+	var b strings.Builder
+	b.Grow(len(s) + 8)
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		switch {
+		case c == '"':
+			b.WriteString(`\"`)
+		case c == '\\':
+			if j := escapeEndAt(s, i); j > i {
+				b.WriteString(s[i : j+1])
+				i = j
+			} else {
+				b.WriteString(`\\`)
+			}
+		case c < 0x20:
+			switch c {
+			case '\n':
+				b.WriteString(`\n`)
+			case '\r':
+				b.WriteString(`\r`)
+			case '\t':
+				b.WriteString(`\t`)
+			default:
+				fmt.Fprintf(&b, `\u%04x`, c)
+			}
+		default:
+			b.WriteByte(c)
+		}
+	}
+	return b.String()
+}
+
+// escapeEndAt reports the index of the last byte of the JSON escape sequence
+// starting at the backslash s[i], or i when the sequence is missing,
+// incomplete or invalid.
+func escapeEndAt(s string, i int) int {
+	if i+1 >= len(s) {
+		return i
+	}
+	switch s[i+1] {
+	case '"', '\\', '/', 'b', 'f', 'n', 'r', 't':
+		return i + 1
+	case 'u':
+		if i+5 < len(s) && isHex4(s[i+2:i+6]) {
+			return i + 5
+		}
+	}
+	return i
+}
+
+func isHex4(s string) bool {
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		if (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F') {
+			continue
+		}
+		return false
+	}
+	return true
+}
diff --git a/internal/openai/toolcalls.go b/internal/openai/toolcalls.go
index 5bd37f8..e183ef1 100644
--- a/internal/openai/toolcalls.go
+++ b/internal/openai/toolcalls.go
@@ -75,6 +75,15 @@ func TryParseToolCalls(text string) (string, []ToolCall, bool) {
 							args = prefix
 						}
 					}
+					if !json.Valid([]byte(args)) {
+						// Last resort: the reply leaked a raw quote or got
+						// cut mid-escape — re-anchor on the final string
+						// literal and rebuild so the client still receives
+						// a tool call with whatever content survived.
+						if closed, ok := forceCloseJSON(s); ok {
+							args = closed
+						}
+					}
 					if !json.Valid([]byte(args)) {
 						bad = true
 						break
@@ -134,6 +143,11 @@ func decodeToolPayload(cand string, payload any) bool {
 			return true
 		}
 	}
+	if closed, ok := forceCloseJSON(cand); ok {
+		if err := json.Unmarshal([]byte(closed), payload); err == nil {
+			return true
+		}
+	}
 	return false
 }