From 2c980c33a87b8cabbdb59b13db67edfd46bf7f14 Mon Sep 17 00:00:00 2001 From: Rahul Chand Date: Fri, 20 Feb 2026 02:14:40 +0530 Subject: [PATCH] feat: Enhance tool call extraction to robustly parse JSON blocks from text, including multiple occurrences. --- pkg/providers/claude_cli_provider_test.go | 21 ---- pkg/providers/tool_call_extract.go | 117 +++++++++++----------- pkg/providers/tool_call_extract_test.go | 16 +++ 3 files changed, 77 insertions(+), 77 deletions(-) diff --git a/pkg/providers/claude_cli_provider_test.go b/pkg/providers/claude_cli_provider_test.go index 063530deb..ba50ebde3 100644 --- a/pkg/providers/claude_cli_provider_test.go +++ b/pkg/providers/claude_cli_provider_test.go @@ -958,24 +958,3 @@ func TestStripToolCallsJSON_OnlyToolCalls(t *testing.T) { // --- findMatchingBrace tests --- -func TestFindMatchingBrace(t *testing.T) { - tests := []struct { - text string - pos int - want int - }{ - {`{"a":1}`, 0, 7}, - {`{"a":{"b":2}}`, 0, 13}, - {`text {"a":1} more`, 5, 12}, - {`{unclosed`, 0, 0}, // no match returns pos - {`{}`, 0, 2}, // empty object - {`{{{}}}`, 0, 6}, // deeply nested - {`{"a":"b{c}d"}`, 0, 13}, // braces in strings (simplified matcher) - } - for _, tt := range tests { - got := findMatchingBrace(tt.text, tt.pos) - if got != tt.want { - t.Errorf("findMatchingBrace(%q, %d) = %d, want %d", tt.text, tt.pos, got, tt.want) - } - } -} diff --git a/pkg/providers/tool_call_extract.go b/pkg/providers/tool_call_extract.go index 298876224..86e7aebd9 100644 --- a/pkg/providers/tool_call_extract.go +++ b/pkg/providers/tool_call_extract.go @@ -5,70 +5,75 @@ import ( "strings" ) -// 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 + for i := 0; i < len(text); i++ { + if text[i] == '{' { + end := findMatchingBrace(text, i) + if end > i { + jsonStr := text[i:end] + // Quick check to avoid expensive parsing if it doesn't mention tool_calls + if !strings.Contains(jsonStr, "tool_calls") { + continue + } + + 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 && len(wrapper.ToolCalls) > 0 { + var result []ToolCall + for _, tc := range wrapper.ToolCalls { + var args map[string]interface{} + json.Unmarshal([]byte(tc.Function.Arguments), &args) + + result = append(result, ToolCall{ + ID: tc.ID, + Type: tc.Type, + Name: tc.Function.Name, + Arguments: args, + Function: &FunctionCall{ + Name: tc.Function.Name, + Arguments: tc.Function.Arguments, + }, + }) + } + return result + } + } + } } - - 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 { - return nil - } - - var result []ToolCall - for _, tc := range wrapper.ToolCalls { - var args map[string]interface{} - json.Unmarshal([]byte(tc.Function.Arguments), &args) - - result = append(result, ToolCall{ - ID: tc.ID, - Type: tc.Type, - Name: tc.Function.Name, - Arguments: args, - Function: &FunctionCall{ - Name: tc.Function.Name, - Arguments: tc.Function.Arguments, - }, - }) - } - - return result + return nil } // stripToolCallsFromText removes tool call JSON from response text. func stripToolCallsFromText(text string) string { - start := strings.Index(text, `{"tool_calls"`) - if start == -1 { - return text - } + for i := 0; i < len(text); i++ { + if text[i] == '{' { + end := findMatchingBrace(text, i) + if end > i { + jsonStr := text[i:end] + if !strings.Contains(jsonStr, "tool_calls") { + continue + } - end := findMatchingBrace(text, start) - if end == start { - return text - } + var wrapper struct { + ToolCalls interface{} `json:"tool_calls"` + } - return strings.TrimSpace(text[:start] + text[end:]) + if err := json.Unmarshal([]byte(jsonStr), &wrapper); err == nil && wrapper.ToolCalls != nil { + return strings.TrimSpace(text[:i] + text[end:]) + } + } + } + } + return text } // findMatchingBrace finds the index after the closing brace matching the opening brace at pos. diff --git a/pkg/providers/tool_call_extract_test.go b/pkg/providers/tool_call_extract_test.go index 92f973226..fdd5b957f 100644 --- a/pkg/providers/tool_call_extract_test.go +++ b/pkg/providers/tool_call_extract_test.go @@ -65,6 +65,22 @@ func TestExtractToolCallsFromText(t *testing.T) { }, }, }, + { + name: "Multiple JSON blocks", + text: `Some config: {"debug": true}. Then the tool call: {"tool_calls":[{"id":"c1","type":"function","function":{"name":"search","arguments":"{}"}}]}.`, + want: []ToolCall{ + { + ID: "c1", + Type: "function", + Name: "search", + Arguments: map[string]interface{}{}, + Function: &FunctionCall{ + Name: "search", + Arguments: "{}", + }, + }, + }, + }, { name: "No tool calls", text: "Just some normal text here.",