From fd3a0fa7824d1f1967287101092db243ee375886 Mon Sep 17 00:00:00 2001 From: Alix-007 <267018309+Alix-007@users.noreply.github.com> Date: Sun, 29 Mar 2026 16:13:35 +0800 Subject: [PATCH] fix(provider): harden text tool-call extraction parsing --- pkg/providers/claude_cli_provider.go | 26 +++++- pkg/providers/tool_call_extract.go | 89 ++++++++++++-------- pkg/providers/tool_call_extract_test.go | 105 ++++++++++++++++++++++++ 3 files changed, 185 insertions(+), 35 deletions(-) create mode 100644 pkg/providers/tool_call_extract_test.go diff --git a/pkg/providers/claude_cli_provider.go b/pkg/providers/claude_cli_provider.go index 40b581490..e2271cde3 100644 --- a/pkg/providers/claude_cli_provider.go +++ b/pkg/providers/claude_cli_provider.go @@ -164,10 +164,32 @@ func (p *ClaudeCliProvider) stripToolCallsJSON(text string) string { // findMatchingBrace finds the index after the closing brace matching the opening brace at pos. func findMatchingBrace(text string, pos int) int { depth := 0 + inString := false + escaped := false + for i := pos; i < len(text); i++ { - if text[i] == '{' { + ch := text[i] + if inString { + if escaped { + escaped = false + continue + } + if ch == '\\' { + escaped = true + continue + } + if ch == '"' { + inString = false + } + continue + } + + switch ch { + case '"': + inString = true + case '{': depth++ - } else if text[i] == '}' { + case '}': depth-- if depth == 0 { return i + 1 diff --git a/pkg/providers/tool_call_extract.go b/pkg/providers/tool_call_extract.go index 7ddea0e99..5ae13f1c8 100644 --- a/pkg/providers/tool_call_extract.go +++ b/pkg/providers/tool_call_extract.go @@ -5,41 +5,28 @@ import ( "strings" ) +type textToolCall struct { + ID string `json:"id"` + Type string `json:"type"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` +} + // extractToolCallsFromText parses tool call JSON from response text. // Both ClaudeCliProvider and CodexCliProvider use this to extract // tool calls that the model outputs in its response text. func extractToolCallsFromText(text string) []ToolCall { - start := strings.Index(text, `{"tool_calls"`) - if start == -1 { - return nil - } - - end := findMatchingBrace(text, start) - if end == start { - return nil - } - - jsonStr := text[start:end] - - var wrapper struct { - ToolCalls []struct { - ID string `json:"id"` - Type string `json:"type"` - Function struct { - Name string `json:"name"` - Arguments string `json:"arguments"` - } `json:"function"` - } `json:"tool_calls"` - } - - if err := json.Unmarshal([]byte(jsonStr), &wrapper); err != nil { + start, end, rawToolCalls, ok := locateToolCallObject(text) + if !ok || end <= start { return nil } var result []ToolCall - for _, tc := range wrapper.ToolCalls { + for _, tc := range rawToolCalls { var args map[string]any - json.Unmarshal([]byte(tc.Function.Arguments), &args) + _ = json.Unmarshal([]byte(tc.Function.Arguments), &args) result = append(result, ToolCall{ ID: tc.ID, @@ -58,15 +45,51 @@ func extractToolCallsFromText(text string) []ToolCall { // stripToolCallsFromText removes tool call JSON from response text. func stripToolCallsFromText(text string) string { - start := strings.Index(text, `{"tool_calls"`) - if start == -1 { - return text - } - - end := findMatchingBrace(text, start) - if end == start { + start, end, _, ok := locateToolCallObject(text) + if !ok || end <= start { return text } return strings.TrimSpace(text[:start] + text[end:]) } + +func locateToolCallObject(text string) (start int, end int, toolCalls []textToolCall, ok bool) { + for i := 0; i < len(text); i++ { + if text[i] != '{' { + continue + } + + candidateEnd := findMatchingBrace(text, i) + if candidateEnd <= i { + continue + } + + candidateToolCalls, matched := parseToolCallsObject(text[i:candidateEnd]) + if !matched { + continue + } + + return i, candidateEnd, candidateToolCalls, true + } + + return 0, 0, nil, false +} + +func parseToolCallsObject(jsonStr string) ([]textToolCall, bool) { + var obj map[string]json.RawMessage + if err := json.Unmarshal([]byte(jsonStr), &obj); err != nil { + return nil, false + } + + rawToolCalls, exists := obj["tool_calls"] + if !exists { + return nil, false + } + + var toolCalls []textToolCall + if err := json.Unmarshal(rawToolCalls, &toolCalls); err != nil { + return nil, false + } + + return toolCalls, true +} diff --git a/pkg/providers/tool_call_extract_test.go b/pkg/providers/tool_call_extract_test.go new file mode 100644 index 000000000..0933ac5f7 --- /dev/null +++ b/pkg/providers/tool_call_extract_test.go @@ -0,0 +1,105 @@ +package providers + +import "testing" + +func TestExtractToolCallsFromText_IndentedObjectWithExtraText(t *testing.T) { + text := `Planning tool execution... +{ + "message": "calling tools", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "read_file", + "arguments": "{\"path\":\"/tmp/test.txt\"}" + } + } + ], + "status": "ok" +} +Done.` + + got := extractToolCallsFromText(text) + if len(got) != 1 { + t.Fatalf("extractToolCallsFromText() len = %d, want 1", len(got)) + } + if got[0].Name != "read_file" { + t.Fatalf("tool name = %q, want %q", got[0].Name, "read_file") + } + if got[0].Arguments["path"] != "/tmp/test.txt" { + t.Fatalf("tool args[path] = %v, want /tmp/test.txt", got[0].Arguments["path"]) + } +} + +func TestExtractToolCallsFromText_ToolCallsNotFirstField(t *testing.T) { + text := `{"kind":"assistant_result","metadata":{"provider":"x"},"tool_calls":[{"id":"call_7","type":"function","function":{"name":"write_file","arguments":"{\"path\":\"a.txt\",\"content\":\"hello\"}"}}]}` + + got := extractToolCallsFromText(text) + if len(got) != 1 { + t.Fatalf("extractToolCallsFromText() len = %d, want 1", len(got)) + } + if got[0].ID != "call_7" { + t.Fatalf("tool id = %q, want %q", got[0].ID, "call_7") + } + if got[0].Name != "write_file" { + t.Fatalf("tool name = %q, want %q", got[0].Name, "write_file") + } +} + +func TestExtractToolCallsFromText_SkipsInvalidCandidateAndFindsValidObject(t *testing.T) { + text := `prefix {"tool_calls":invalid} middle {"note":"valid-json-without-tools"} tail { + "note": "valid-json-with-tools", + "tool_calls": [ + { + "id": "call_2", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Tokyo\"}" + } + } + ] +}` + + got := extractToolCallsFromText(text) + if len(got) != 1 { + t.Fatalf("extractToolCallsFromText() len = %d, want 1", len(got)) + } + if got[0].Name != "get_weather" { + t.Fatalf("tool name = %q, want %q", got[0].Name, "get_weather") + } +} + +func TestStripToolCallsFromText_DoesNotStripInvalidToolCallsObject(t *testing.T) { + text := `before {"tool_calls":"not-an-array","other":1} after` + + got := stripToolCallsFromText(text) + if got != text { + t.Fatalf("stripToolCallsFromText() = %q, want unchanged", got) + } +} + +func TestStripToolCallsFromText_StripsValidIndentedObject(t *testing.T) { + text := `before +{ + "message": "tool call follows", + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "list_files", + "arguments": "{}" + } + } + ] +} +after` + + got := stripToolCallsFromText(text) + want := "before\n\nafter" + if got != want { + t.Fatalf("stripToolCallsFromText() = %q, want %q", got, want) + } +}