feat: force-close truncated tool-call JSON into partial tool calls

This commit is contained in:
nyoungo committed 2026-10-07 22:41:18 +08:00
1 parent 413be4ba9e
commit 975882292c
4 files changed
+231 -1

No files matched your search

+3 -1
View File
@@ -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]
+40
View File
@@ -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)
}
}
+174
View File
@@ -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
}
+14
View File
@@ -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
}