From 8a47a6fb6759806495294599b2f798ea81127e05 Mon Sep 17 00:00:00 2001 From: Leandro Barbosa Date: Thu, 12 Feb 2026 14:16:50 -0300 Subject: [PATCH] fix: add tool call extraction to Codex CLI provider - Extract shared tool call parsing into tool_call_extract.go (extractToolCallsFromText, stripToolCallsFromText, findMatchingBrace) - Both ClaudeCliProvider and CodexCliProvider now share the same tool call extraction logic for PicoClaw-specific tools - Fix cache token accounting: include cached_input_tokens in total - Add 2 new tests for tool call extraction from JSONL events - Update existing tests for corrected token calculations --- pkg/providers/claude_cli_provider.go | 62 ++----------------- pkg/providers/codex_cli_provider.go | 18 ++++-- pkg/providers/codex_cli_provider_test.go | 79 ++++++++++++++++++++++-- pkg/providers/tool_call_extract.go | 72 +++++++++++++++++++++ 4 files changed, 163 insertions(+), 68 deletions(-) create mode 100644 pkg/providers/tool_call_extract.go diff --git a/pkg/providers/claude_cli_provider.go b/pkg/providers/claude_cli_provider.go index 242126aa2..91df4ccb7 100644 --- a/pkg/providers/claude_cli_provider.go +++ b/pkg/providers/claude_cli_provider.go @@ -171,68 +171,14 @@ func (p *ClaudeCliProvider) parseClaudeCliResponse(output string) (*LLMResponse, }, nil } -// extractToolCalls parses tool call JSON from the response text. +// extractToolCalls delegates to the shared extractToolCallsFromText function. func (p *ClaudeCliProvider) extractToolCalls(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 { - 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 extractToolCallsFromText(text) } -// stripToolCallsJSON removes tool call JSON from response text. +// stripToolCallsJSON delegates to the shared stripToolCallsFromText function. func (p *ClaudeCliProvider) stripToolCallsJSON(text string) string { - start := strings.Index(text, `{"tool_calls"`) - if start == -1 { - return text - } - - end := findMatchingBrace(text, start) - if end == start { - return text - } - - return strings.TrimSpace(text[:start] + text[end:]) + return stripToolCallsFromText(text) } // findMatchingBrace finds the index after the closing brace matching the opening brace at pos. diff --git a/pkg/providers/codex_cli_provider.go b/pkg/providers/codex_cli_provider.go index 46aa53b46..37cfd9314 100644 --- a/pkg/providers/codex_cli_provider.go +++ b/pkg/providers/codex_cli_provider.go @@ -199,10 +199,11 @@ func (p *CodexCliProvider) parseJSONLEvents(output string) (*LLMResponse, error) } case "turn.completed": if event.Usage != nil { + promptTokens := event.Usage.InputTokens + event.Usage.CachedInputTokens usage = &UsageInfo{ - PromptTokens: event.Usage.InputTokens, + PromptTokens: promptTokens, CompletionTokens: event.Usage.OutputTokens, - TotalTokens: event.Usage.InputTokens + event.Usage.OutputTokens, + TotalTokens: promptTokens + event.Usage.OutputTokens, } } case "error": @@ -220,10 +221,19 @@ func (p *CodexCliProvider) parseJSONLEvents(output string) (*LLMResponse, error) content := strings.Join(contentParts, "\n") + // Extract tool calls from response text (same pattern as ClaudeCliProvider) + toolCalls := extractToolCallsFromText(content) + + finishReason := "stop" + if len(toolCalls) > 0 { + finishReason = "tool_calls" + content = stripToolCallsFromText(content) + } + return &LLMResponse{ Content: strings.TrimSpace(content), - ToolCalls: nil, - FinishReason: "stop", + ToolCalls: toolCalls, + FinishReason: finishReason, Usage: usage, }, nil } diff --git a/pkg/providers/codex_cli_provider_test.go b/pkg/providers/codex_cli_provider_test.go index 84e1303ec..7e4e1bc15 100644 --- a/pkg/providers/codex_cli_provider_test.go +++ b/pkg/providers/codex_cli_provider_test.go @@ -2,6 +2,7 @@ package providers import ( "context" + "encoding/json" "fmt" "os" "os/exec" @@ -32,20 +33,86 @@ func TestParseJSONLEvents_AgentMessage(t *testing.T) { if resp.Usage == nil { t.Fatal("Usage should not be nil") } - if resp.Usage.PromptTokens != 100 { - t.Errorf("PromptTokens = %d, want 100", resp.Usage.PromptTokens) + if resp.Usage.PromptTokens != 150 { + t.Errorf("PromptTokens = %d, want 150", resp.Usage.PromptTokens) } if resp.Usage.CompletionTokens != 20 { t.Errorf("CompletionTokens = %d, want 20", resp.Usage.CompletionTokens) } - if resp.Usage.TotalTokens != 120 { - t.Errorf("TotalTokens = %d, want 120", resp.Usage.TotalTokens) + if resp.Usage.TotalTokens != 170 { + t.Errorf("TotalTokens = %d, want 170", resp.Usage.TotalTokens) } if len(resp.ToolCalls) != 0 { t.Errorf("ToolCalls should be empty, got %d", len(resp.ToolCalls)) } } +func TestParseJSONLEvents_ToolCallExtraction(t *testing.T) { + p := &CodexCliProvider{} + toolCallText := `Let me read that file. +{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"/tmp/test.txt\"}"}}]}` + // Build valid JSONL by marshaling the event + item := codexEvent{ + Type: "item.completed", + Item: &codexEventItem{ID: "item_1", Type: "agent_message", Text: toolCallText}, + } + itemJSON, _ := json.Marshal(item) + usageEvt := `{"type":"turn.completed","usage":{"input_tokens":50,"cached_input_tokens":0,"output_tokens":20}}` + events := `{"type":"turn.started"}` + "\n" + string(itemJSON) + "\n" + usageEvt + + resp, err := p.parseJSONLEvents(events) + if err != nil { + t.Fatalf("parseJSONLEvents() error: %v", err) + } + if resp.FinishReason != "tool_calls" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "tool_calls") + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("ToolCalls count = %d, want 1", len(resp.ToolCalls)) + } + if resp.ToolCalls[0].Name != "read_file" { + t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "read_file") + } + if resp.ToolCalls[0].ID != "call_1" { + t.Errorf("ToolCalls[0].ID = %q, want %q", resp.ToolCalls[0].ID, "call_1") + } + if resp.ToolCalls[0].Function.Arguments != `{"path":"/tmp/test.txt"}` { + t.Errorf("ToolCalls[0].Function.Arguments = %q", resp.ToolCalls[0].Function.Arguments) + } + // Content should have the tool call JSON stripped + if strings.Contains(resp.Content, "tool_calls") { + t.Errorf("Content should not contain tool_calls JSON, got: %q", resp.Content) + } +} + +func TestParseJSONLEvents_MultipleToolCalls(t *testing.T) { + p := &CodexCliProvider{} + toolCallText := `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"a.txt\"}"}},{"id":"call_2","type":"function","function":{"name":"write_file","arguments":"{\"path\":\"b.txt\",\"content\":\"hello\"}"}}]}` + item := codexEvent{ + Type: "item.completed", + Item: &codexEventItem{ID: "item_1", Type: "agent_message", Text: toolCallText}, + } + itemJSON, _ := json.Marshal(item) + events := `{"type":"turn.started"}` + "\n" + string(itemJSON) + "\n" + `{"type":"turn.completed"}` + + resp, err := p.parseJSONLEvents(events) + if err != nil { + t.Fatalf("parseJSONLEvents() error: %v", err) + } + if len(resp.ToolCalls) != 2 { + t.Fatalf("ToolCalls count = %d, want 2", len(resp.ToolCalls)) + } + if resp.ToolCalls[0].Name != "read_file" { + t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "read_file") + } + if resp.ToolCalls[1].Name != "write_file" { + t.Errorf("ToolCalls[1].Name = %q, want %q", resp.ToolCalls[1].Name, "write_file") + } + if resp.FinishReason != "tool_calls" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "tool_calls") + } +} + func TestParseJSONLEvents_MultipleMessages(t *testing.T) { p := &CodexCliProvider{} events := `{"type":"turn.started"} @@ -372,8 +439,8 @@ func TestCodexCliProvider_MockCLI_Success(t *testing.T) { if resp.Usage == nil { t.Fatal("Usage should not be nil") } - if resp.Usage.PromptTokens != 50 { - t.Errorf("PromptTokens = %d, want 50", resp.Usage.PromptTokens) + if resp.Usage.PromptTokens != 60 { + t.Errorf("PromptTokens = %d, want 60", resp.Usage.PromptTokens) } if resp.Usage.CompletionTokens != 15 { t.Errorf("CompletionTokens = %d, want 15", resp.Usage.CompletionTokens) diff --git a/pkg/providers/tool_call_extract.go b/pkg/providers/tool_call_extract.go new file mode 100644 index 000000000..97a219283 --- /dev/null +++ b/pkg/providers/tool_call_extract.go @@ -0,0 +1,72 @@ +package providers + +import ( + "encoding/json" + "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 + } + + 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 +} + +// 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 { + return text + } + + return strings.TrimSpace(text[:start] + text[end:]) +}