From 0c1ededc929567dc1740c822daea466478ddff7d Mon Sep 17 00:00:00 2001 From: thotam <78630736+thotam@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:56:05 +0700 Subject: [PATCH] feat(webhooks): return per-call usage breakdown with provider/model/cost (#1421) feat(webhooks): per-call usage breakdown with provider/model/cost (#1421) --- docs/webhooks.md | 59 ++++++++++-- internal/agent/adapter_calls_test.go | 21 ++++ internal/agent/llm_call_usage_test.go | 45 +++++++++ internal/agent/loop_pipeline_adapter.go | 1 + internal/agent/loop_pipeline_callbacks.go | 11 +++ .../agent/loop_pipeline_tool_callbacks.go | 18 ++++ internal/agent/loop_tracing.go | 5 +- internal/agent/loop_types.go | 21 ++-- internal/agent/tool_call_usage_test.go | 43 +++++++++ internal/http/webhooks_llm.go | 41 ++++++-- internal/http/webhooks_llm_test.go | 41 ++++++-- internal/pipeline/run_state.go | 14 +++ internal/pipeline/run_state_calls_test.go | 29 ++++++ internal/pipeline/substates.go | 1 + internal/providers/call_usage.go | 42 ++++++++ internal/providers/call_usage_test.go | 48 ++++++++++ internal/webhooks/worker.go | 75 ++++++++++----- internal/webhooks/worker_test.go | 95 ++++++++++++++++++- 18 files changed, 551 insertions(+), 59 deletions(-) create mode 100644 internal/agent/adapter_calls_test.go create mode 100644 internal/agent/llm_call_usage_test.go create mode 100644 internal/agent/tool_call_usage_test.go create mode 100644 internal/pipeline/run_state_calls_test.go create mode 100644 internal/providers/call_usage.go create mode 100644 internal/providers/call_usage_test.go diff --git a/docs/webhooks.md b/docs/webhooks.md index 96bb4478..69440f7c 100644 --- a/docs/webhooks.md +++ b/docs/webhooks.md @@ -264,13 +264,29 @@ Triggers an agent with an input prompt. Available in all editions. "agent_id": "", "output": "Here are the metrics: ...", "usage": { - "prompt_tokens": 150, - "completion_tokens": 200, - "total_tokens": 350, + "prompt_tokens": 250, + "completion_tokens": 220, + "total_tokens": 470, "cache_read_input_tokens": 120, "cache_creation_input_tokens": 30, "prompt_tokens_include_cached_segments": true }, + "total_cost_usd": 0.0279, + "calls": [ + { + "type": "llm_call", "name": "9router/cx/gpt-5.6 #1", + "provider": "9router", "model": "cx/gpt-5.6", + "prompt_tokens": 150, "completion_tokens": 200, "total_tokens": 350, + "cache_read_input_tokens": 120, "cache_creation_input_tokens": 30, + "prompt_tokens_include_cached_segments": true, "cost_usd": 0.0187 + }, + { + "type": "tool_call", "name": "read_image", + "provider": "9router", "model": "cx/gpt-5.5", + "prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, + "cost_usd": 0.0092 + } + ], "finish_reason": "stop" } ``` @@ -279,6 +295,17 @@ Triggers an agent with an input prompt. Available in all editions. > prompt caching was active (omitted otherwise). When > `prompt_tokens_include_cached_segments` is `true`, `prompt_tokens` already counts > the cached segments, so non-cached input = `prompt_tokens - cache_read_input_tokens`. +> +> `calls[]` lists every LLM call and every tool that makes a **direct** internal LLM +> call (e.g. `read_image`, `read_video`), each attributed to its `provider`/`model` +> with its own tokens and `cost_usd`. `usage` is the **sum of all calls** — so it +> includes tool-internal LLM tokens — and `total_cost_usd` is the sum of +> `calls[].cost_usd` (best-effort; `0` when a model has no configured pricing). +> +> Note: token spend inside **nested agent runs** (`subagent`/`delegate` tools, which +> spawn a separate child agent loop) is **not** itemized in `calls[]` and not included +> in `usage`/`total_cost_usd` — consistent with how the child run's usage has always +> been excluded from the parent's totals. Sync mode times out after the configured deadline (default **600s**). On timeout: `504 Gateway Timeout` with `webhook.llm_timeout`. @@ -434,18 +461,38 @@ User-Agent: goclaw-webhook/1 "status": "done", "output": "Agent response text...", "usage": { - "prompt_tokens": 150, - "completion_tokens": 200, - "total_tokens": 350, + "prompt_tokens": 250, + "completion_tokens": 220, + "total_tokens": 470, "cache_read_input_tokens": 120, "cache_creation_input_tokens": 30, "prompt_tokens_include_cached_segments": true }, + "total_cost_usd": 0.0279, + "calls": [ + { + "type": "llm_call", "name": "9router/cx/gpt-5.6 #1", + "provider": "9router", "model": "cx/gpt-5.6", + "prompt_tokens": 150, "completion_tokens": 200, "total_tokens": 350, + "cache_read_input_tokens": 120, "cache_creation_input_tokens": 30, + "prompt_tokens_include_cached_segments": true, "cost_usd": 0.0187 + }, + { + "type": "tool_call", "name": "read_image", + "provider": "9router", "model": "cx/gpt-5.5", + "prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, + "cost_usd": 0.0092 + } + ], "metadata": {}, "error": "" } ``` +> `calls[]` and `total_cost_usd` on the async callback follow the same semantics as +> the sync response above: `usage` is the sum of all `calls[]` (including +> tool-internal LLM calls), each call carries its own `provider`/`model`/tokens/`cost_usd`. + `status` is `"done"` on success, `"failed"` on agent error. `error` is non-empty on failure. ### Retry Schedule diff --git a/internal/agent/adapter_calls_test.go b/internal/agent/adapter_calls_test.go new file mode 100644 index 00000000..59081034 --- /dev/null +++ b/internal/agent/adapter_calls_test.go @@ -0,0 +1,21 @@ +package agent + +import ( + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/pipeline" + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +func TestConvertRunResultMapsCalls(t *testing.T) { + pr := &pipeline.RunResult{ + RunID: "r1", + Calls: []providers.CallUsage{ + {Type: "llm_call", Provider: "p", Model: "m", Usage: providers.Usage{TotalTokens: 5}}, + }, + } + got := convertRunResult(pr) + if len(got.Calls) != 1 || got.Calls[0].Type != "llm_call" { + t.Fatalf("Calls not mapped: %+v", got.Calls) + } +} diff --git a/internal/agent/llm_call_usage_test.go b/internal/agent/llm_call_usage_test.go new file mode 100644 index 00000000..d94f13b8 --- /dev/null +++ b/internal/agent/llm_call_usage_test.go @@ -0,0 +1,45 @@ +package agent + +import ( + "context" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/pipeline" + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +type usageStubProvider struct{} + +func (usageStubProvider) Chat(_ context.Context, _ providers.ChatRequest) (*providers.ChatResponse, error) { + return &providers.ChatResponse{Content: "hi", FinishReason: "stop", + Usage: &providers.Usage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120, CacheReadTokens: 80}}, nil +} +func (usageStubProvider) ChatStream(_ context.Context, _ providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) { + return usageStubProvider{}.Chat(context.Background(), providers.ChatRequest{}) +} +func (usageStubProvider) DefaultModel() string { return "stub-model" } +func (usageStubProvider) Name() string { return "stubprov" } + +func TestLLMCallUsageRecorded(t *testing.T) { + prov := usageStubProvider{} + l := &Loop{id: "a"} + req := &RunRequest{RunID: "r1", SessionKey: "s1", Channel: "ws"} + state := &pipeline.RunState{RunID: "r1", Provider: prov, Model: "stub-model"} + + _, err := l.makeCallLLM(req, func(AgentEvent) {})(context.Background(), state, providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: "hi"}}, + }) + if err != nil { + t.Fatalf("makeCallLLM: %v", err) + } + if len(state.Calls) != 1 { + t.Fatalf("Calls len = %d, want 1", len(state.Calls)) + } + c := state.Calls[0] + if c.Type != "llm_call" || c.Provider != "stubprov" || c.Model != "stub-model" { + t.Errorf("wrong attribution: %+v", c) + } + if c.PromptTokens != 100 || c.CacheReadTokens != 80 { + t.Errorf("wrong tokens: %+v", c.Usage) + } +} diff --git a/internal/agent/loop_pipeline_adapter.go b/internal/agent/loop_pipeline_adapter.go index 2abe66b4..61e07a42 100644 --- a/internal/agent/loop_pipeline_adapter.go +++ b/internal/agent/loop_pipeline_adapter.go @@ -281,6 +281,7 @@ func convertRunResult(pr *pipeline.RunResult) *RunResult { BlockReplies: pr.BlockReplies, LastBlockReply: pr.LastBlockReply, LoopKilled: pr.LoopKilled, + Calls: pr.Calls, } } diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index e780fafd..0ebde72b 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -591,6 +591,17 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } } l.emitLLMSpanEnd(ctx, spanID, start, resp, err, opts...) + if err == nil && resp != nil && resp.Usage != nil { + effModel, effProvider := l.resolveSpan(opts) + state.AppendCall(providers.CallUsage{ + Type: "llm_call", + Name: fmt.Sprintf("%s/%s #%d", effProvider, effModel, state.Iteration+1), + Provider: effProvider, + Model: effModel, + Usage: *resp.Usage, + CostUSD: l.calculateLLMCost(ctx, effProvider, effModel, resp.Usage), + }) + } return resp, err } } diff --git a/internal/agent/loop_pipeline_tool_callbacks.go b/internal/agent/loop_pipeline_tool_callbacks.go index f550e853..d7829b9d 100644 --- a/internal/agent/loop_pipeline_tool_callbacks.go +++ b/internal/agent/loop_pipeline_tool_callbacks.go @@ -17,6 +17,22 @@ import ( "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) +// recordToolCallUsage appends a tool's internal LLM call to the run breakdown. +// No-op unless the tool actually made an LLM call (result.Usage != nil). +func (l *Loop) recordToolCallUsage(ctx context.Context, state *pipeline.RunState, toolName string, result *tools.Result) { + if result == nil || result.Usage == nil { + return + } + state.AppendCall(providers.CallUsage{ + Type: "tool_call", + Name: toolName, + Provider: result.Provider, + Model: result.Model, + Usage: *result.Usage, + CostUSD: l.calculateLLMCost(ctx, result.Provider, result.Model, result.Usage), + }) +} + // makeExecuteToolCall wraps tool execution: name resolution, execute, process result. // Uses bridgeRS to share loop detection state between the pipeline and agent's processToolResult. func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx context.Context, state *pipeline.RunState, tc providers.ToolCall) ([]providers.Message, error) { @@ -65,6 +81,7 @@ func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx toolDuration := time.Since(toolStart) l.emitToolSpanEnd(ctx, toolSpanID, toolStart, result) + l.recordToolCallUsage(ctx, state, registryName, result) l.recordToolUsageEvent(ctx, req, registryName, tc.Name, tc.ID, tc.Arguments, toolStart, result, toolSpanID) // v3 evolution metrics: record tool execution non-blocking (best-effort). @@ -190,6 +207,7 @@ func (l *Loop) makeProcessToolResult(req *RunRequest, bridgeRS *runState) func(c if result == nil { return []providers.Message{rawMsg} } + l.recordToolCallUsage(ctx, state, registryName, result) if rawName == "" { rawName = tc.Name } diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index c6cc88bd..f9ee428b 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -83,7 +83,10 @@ func (l *Loop) resolveSpan(opts []spanOption) (string, string) { } func (l *Loop) resolveSpanOverrides(opts []spanOption) spanOverrides { - o := spanOverrides{model: l.model, provider: l.provider.Name()} + o := spanOverrides{model: l.model} + if l.provider != nil { + o.provider = l.provider.Name() + } for _, fn := range opts { fn(&o) } diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index ab5d6f8f..16a75f60 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -673,16 +673,17 @@ type RunRequest struct { // RunResult is the output of a completed agent run. type RunResult struct { - Content string `json:"content"` - Thinking string `json:"thinking,omitempty"` // reasoning content from thinking models (Claude, o3, DeepSeek-R1, Kimi) - RunID string `json:"runId"` - Iterations int `json:"iterations"` - Usage *providers.Usage `json:"usage,omitempty"` - Media []MediaResult `json:"media,omitempty"` // media files from tool results (MEDIA: prefix) - Deliverables []string `json:"deliverables,omitempty"` // actual content from tool outputs (for team task results) - BlockReplies int `json:"blockReplies,omitempty"` // number of block.reply events emitted - LastBlockReply string `json:"lastBlockReply,omitempty"` // last block reply content (for dedup) - LoopKilled bool `json:"loopKilled,omitempty"` // true when run was terminated by loop detector + Content string `json:"content"` + Thinking string `json:"thinking,omitempty"` // reasoning content from thinking models (Claude, o3, DeepSeek-R1, Kimi) + RunID string `json:"runId"` + Iterations int `json:"iterations"` + Usage *providers.Usage `json:"usage,omitempty"` + Media []MediaResult `json:"media,omitempty"` // media files from tool results (MEDIA: prefix) + Deliverables []string `json:"deliverables,omitempty"` // actual content from tool outputs (for team task results) + BlockReplies int `json:"blockReplies,omitempty"` // number of block.reply events emitted + LastBlockReply string `json:"lastBlockReply,omitempty"` // last block reply content (for dedup) + LoopKilled bool `json:"loopKilled,omitempty"` // true when run was terminated by loop detector + Calls []providers.CallUsage `json:"calls,omitempty"` // per-call usage breakdown } // MediaResult represents a media file produced by a tool during the agent run. diff --git a/internal/agent/tool_call_usage_test.go b/internal/agent/tool_call_usage_test.go new file mode 100644 index 00000000..d1f19efb --- /dev/null +++ b/internal/agent/tool_call_usage_test.go @@ -0,0 +1,43 @@ +package agent + +import ( + "context" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/pipeline" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/tools" +) + +type usageToolExecutor struct{} + +func (usageToolExecutor) ExecuteWithContext(_ context.Context, _ string, _ map[string]any, _, _, _, _ string, _ tools.AsyncCallback) *tools.Result { + return &tools.Result{ForLLM: "ok", Provider: "9router", Model: "cx/gpt-5.5", + Usage: &providers.Usage{PromptTokens: 4677, CompletionTokens: 1799, TotalTokens: 6476}} +} +func (usageToolExecutor) TryActivateDeferred(string) bool { return false } +func (usageToolExecutor) ProviderDefs() []providers.ToolDefinition { return nil } +func (usageToolExecutor) Get(string) (tools.Tool, bool) { return nil, false } +func (usageToolExecutor) List() []string { return nil } +func (usageToolExecutor) Aliases() map[string]string { return nil } + +func TestToolCallUsageRecorded(t *testing.T) { + l := &Loop{id: "a", tools: usageToolExecutor{}} + req := &RunRequest{RunID: "r1", SessionKey: "s1", Channel: "ws"} + state := &pipeline.RunState{RunID: "r1"} + tc := providers.ToolCall{ID: "tc-1", Name: "read_image", Arguments: map[string]any{}} + + if _, err := l.makeExecuteToolCall(req, &runState{})(context.Background(), state, tc); err != nil { + t.Fatalf("makeExecuteToolCall: %v", err) + } + if len(state.Calls) != 1 { + t.Fatalf("Calls len = %d, want 1", len(state.Calls)) + } + c := state.Calls[0] + if c.Type != "tool_call" || c.Name != "read_image" || c.Provider != "9router" || c.Model != "cx/gpt-5.5" { + t.Errorf("wrong attribution: %+v", c) + } + if c.PromptTokens != 4677 || c.CompletionTokens != 1799 { + t.Errorf("wrong tokens: %+v", c.Usage) + } +} diff --git a/internal/http/webhooks_llm.go b/internal/http/webhooks_llm.go index 65868d74..d43c3033 100644 --- a/internal/http/webhooks_llm.go +++ b/internal/http/webhooks_llm.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/agent" "github.com/nextlevelbuilder/goclaw/internal/i18n" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/scheduler" "github.com/nextlevelbuilder/goclaw/internal/security" "github.com/nextlevelbuilder/goclaw/internal/store" @@ -70,11 +71,13 @@ type webhookInputMessage struct { // webhookLLMSyncResp is the 200 response for synchronous LLM calls. type webhookLLMSyncResp struct { - CallID string `json:"call_id"` - AgentID string `json:"agent_id"` - Output string `json:"output"` - Usage *webhookLLMUsage `json:"usage,omitempty"` - FinishReason string `json:"finish_reason"` + CallID string `json:"call_id"` + AgentID string `json:"agent_id"` + Output string `json:"output"` + Usage *webhookLLMUsage `json:"usage,omitempty"` + FinishReason string `json:"finish_reason"` + Calls []providers.CallUsage `json:"calls,omitempty"` + TotalCostUSD float64 `json:"total_cost_usd,omitempty"` } // webhookLLMUsage mirrors providers.Usage for the response envelope. @@ -418,7 +421,19 @@ func (h *WebhookLLMHandler) handleSync( Output: out.result.Content, FinishReason: "stop", } - if out.result.Usage != nil { + if len(out.result.Calls) > 0 { + resp.Calls = out.result.Calls + resp.TotalCostUSD = providers.SumCallCost(out.result.Calls) + sum := providers.SumCallUsage(out.result.Calls) + resp.Usage = &webhookLLMUsage{ + PromptTokens: sum.PromptTokens, + CompletionTokens: sum.CompletionTokens, + TotalTokens: sum.TotalTokens, + CacheReadTokens: sum.CacheReadTokens, + CacheCreationTokens: sum.CacheCreationTokens, + PromptTokensIncludeCachedSegments: sum.PromptTokensIncludeCachedSegments, + } + } else if out.result.Usage != nil { resp.Usage = &webhookLLMUsage{ PromptTokens: out.result.Usage.PromptTokens, CompletionTokens: out.result.Usage.CompletionTokens, @@ -636,7 +651,19 @@ func (h *WebhookLLMHandler) RunTest(ctx context.Context, wh *store.WebhookData, Output: out.result.Content, FinishReason: "stop", } - if out.result.Usage != nil { + if len(out.result.Calls) > 0 { + resp.Calls = out.result.Calls + resp.TotalCostUSD = providers.SumCallCost(out.result.Calls) + sum := providers.SumCallUsage(out.result.Calls) + resp.Usage = &webhookLLMUsage{ + PromptTokens: sum.PromptTokens, + CompletionTokens: sum.CompletionTokens, + TotalTokens: sum.TotalTokens, + CacheReadTokens: sum.CacheReadTokens, + CacheCreationTokens: sum.CacheCreationTokens, + PromptTokensIncludeCachedSegments: sum.PromptTokensIncludeCachedSegments, + } + } else if out.result.Usage != nil { resp.Usage = &webhookLLMUsage{ PromptTokens: out.result.Usage.PromptTokens, CompletionTokens: out.result.Usage.CompletionTokens, diff --git a/internal/http/webhooks_llm_test.go b/internal/http/webhooks_llm_test.go index 1890d5eb..4823bfbd 100644 --- a/internal/http/webhooks_llm_test.go +++ b/internal/http/webhooks_llm_test.go @@ -28,22 +28,22 @@ type stubLLMAgent struct { runFn func(ctx context.Context, req agent.RunRequest) (*agent.RunResult, error) } -func (a *stubLLMAgent) ID() string { return a.id } -func (a *stubLLMAgent) UUID() uuid.UUID { return a.agentID } +func (a *stubLLMAgent) ID() string { return a.id } +func (a *stubLLMAgent) UUID() uuid.UUID { return a.agentID } func (a *stubLLMAgent) OtherConfig() json.RawMessage { return nil } func (a *stubLLMAgent) Run(ctx context.Context, req agent.RunRequest) (*agent.RunResult, error) { return a.runFn(ctx, req) } -func (a *stubLLMAgent) IsRunning() bool { return false } -func (a *stubLLMAgent) Model() string { return "test-model" } -func (a *stubLLMAgent) ProviderName() string { return "test" } +func (a *stubLLMAgent) IsRunning() bool { return false } +func (a *stubLLMAgent) Model() string { return "test-model" } +func (a *stubLLMAgent) ProviderName() string { return "test" } func (a *stubLLMAgent) Provider() providers.Provider { return nil } // ---- stub: store.WebhookCallStore for LLM tests ---- // llmCallStore captures Create calls for assertion. type llmCallStore struct { - created []*store.WebhookCallData + created []*store.WebhookCallData createErr error } @@ -226,6 +226,18 @@ func TestWebhookLLMHandler_SyncHappyPath(t *testing.T) { CacheCreationTokens: 2, PromptTokensIncludeCachedSegments: true, }, + // Cache fields live on call #1 only — mirrors production, where each + // CallUsage copies the full providers.Usage from the LLM response + // (loop_pipeline_callbacks.go), so SumCallUsage must OR/sum them + // across calls rather than reading them off the flat RunResult.Usage. + Calls: []providers.CallUsage{ + {Type: "llm_call", Name: "stub/m #1", Provider: "stub", Model: "m", + Usage: providers.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15, + CacheReadTokens: 8, CacheCreationTokens: 2, PromptTokensIncludeCachedSegments: true}, + CostUSD: 0.01}, + {Type: "tool_call", Name: "read_image", Provider: "9router", Model: "cx/gpt-5.5", + Usage: providers.Usage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120}, CostUSD: 0.02}, + }, }, nil }, } @@ -260,9 +272,12 @@ func TestWebhookLLMHandler_SyncHappyPath(t *testing.T) { if resp.Output != "42" { t.Errorf("expected output '42', got %q", resp.Output) } - if resp.Usage == nil || resp.Usage.TotalTokens != 15 { + // usage now equals SumCallUsage(resp.Calls): prompt 10+100=110, completion 5+20=25, total 15+120=135. + if resp.Usage == nil || resp.Usage.TotalTokens != 135 { t.Errorf("unexpected usage: %+v", resp.Usage) } + // Cache fields (set only on call #1) must survive the sum unchanged — regression + // coverage for cache-token propagation through the Calls -> SumCallUsage path. if resp.Usage.CacheReadTokens != 8 || resp.Usage.CacheCreationTokens != 2 { t.Errorf("cache tokens not propagated: read=%d create=%d", resp.Usage.CacheReadTokens, resp.Usage.CacheCreationTokens) } @@ -272,6 +287,18 @@ func TestWebhookLLMHandler_SyncHappyPath(t *testing.T) { if resp.AgentID != agentUUID.String() { t.Errorf("expected agent_id %s, got %s", agentUUID, resp.AgentID) } + if len(resp.Calls) != 2 { + t.Fatalf("Calls len = %d, want 2", len(resp.Calls)) + } + if resp.Calls[1].Provider != "9router" || resp.Calls[1].Model != "cx/gpt-5.5" { + t.Errorf("tool call attribution wrong: %+v", resp.Calls[1]) + } + if resp.Usage == nil || resp.Usage.PromptTokens != 110 { // 10 + 100 = sum of calls + t.Errorf("usage should equal SumCallUsage(calls).PromptTokens=110, got %+v", resp.Usage) + } + if resp.TotalCostUSD < 0.0299 || resp.TotalCostUSD > 0.0301 { + t.Errorf("TotalCostUSD = %f, want ~0.03", resp.TotalCostUSD) + } // Audit row must be written with status=done. if len(callStore.created) != 1 { diff --git a/internal/pipeline/run_state.go b/internal/pipeline/run_state.go index 98241c9f..4399f20f 100644 --- a/internal/pipeline/run_state.go +++ b/internal/pipeline/run_state.go @@ -2,6 +2,7 @@ package pipeline import ( "context" + "sync" "github.com/google/uuid" @@ -44,6 +45,11 @@ type RunState struct { CurrentLLMSpanID *uuid.UUID // CurrentToolSpanID is the most recent tool-call span; post-tool-use hook spans parent to it. CurrentToolSpanID *uuid.UUID + + // Calls is the per-call usage breakdown (LLM calls + tool-internal LLM calls), + // appended during the run. Guarded by callsMu for the parallel tool path. + Calls []providers.CallUsage + callsMu sync.Mutex } // NewRunState creates a RunState with identity fields set. @@ -58,6 +64,13 @@ func NewRunState(input *RunInput, ws *workspace.WorkspaceContext, model string, } } +// AppendCall records one call's usage in the run breakdown (thread-safe). +func (rs *RunState) AppendCall(c providers.CallUsage) { + rs.callsMu.Lock() + rs.Calls = append(rs.Calls, c) + rs.callsMu.Unlock() +} + // BuildResult converts final RunState into a RunResult. func (rs *RunState) BuildResult() *RunResult { return &RunResult{ @@ -73,6 +86,7 @@ func (rs *RunState) BuildResult() *RunResult { Deliverables: rs.Tool.Deliverables, BlockReplies: rs.Observe.BlockReplies, LastBlockReply: rs.Observe.LastBlockReply, + Calls: rs.Calls, } } diff --git a/internal/pipeline/run_state_calls_test.go b/internal/pipeline/run_state_calls_test.go new file mode 100644 index 00000000..2fa39864 --- /dev/null +++ b/internal/pipeline/run_state_calls_test.go @@ -0,0 +1,29 @@ +package pipeline + +import ( + "sync" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +func TestRunStateAppendCallAndBuildResult(t *testing.T) { + rs := &RunState{RunID: "r1"} + var wg sync.WaitGroup + for i := 0; i < 10; i++ { // race-safety: concurrent appends (parallel tools) + wg.Add(1) + go func() { + defer wg.Done() + rs.AppendCall(providers.CallUsage{Type: "tool_call", Provider: "p", Model: "m", + Usage: providers.Usage{PromptTokens: 1, TotalTokens: 1}}) + }() + } + wg.Wait() + if len(rs.Calls) != 10 { + t.Fatalf("Calls len = %d, want 10", len(rs.Calls)) + } + res := rs.BuildResult() + if len(res.Calls) != 10 { + t.Fatalf("BuildResult Calls len = %d, want 10", len(res.Calls)) + } +} diff --git a/internal/pipeline/substates.go b/internal/pipeline/substates.go index 8f72bbe0..1a840b3e 100644 --- a/internal/pipeline/substates.go +++ b/internal/pipeline/substates.go @@ -117,4 +117,5 @@ type RunResult struct { Deliverables []string BlockReplies int LastBlockReply string + Calls []providers.CallUsage } diff --git a/internal/providers/call_usage.go b/internal/providers/call_usage.go new file mode 100644 index 00000000..5f7b3eed --- /dev/null +++ b/internal/providers/call_usage.go @@ -0,0 +1,42 @@ +package providers + +// CallUsage attributes one LLM call (or a tool's internal LLM call) to its +// provider/model with the tokens it consumed and its cost. Usage is embedded +// so its token fields serialize flat (prompt_tokens, cache_read_input_tokens, …) +// alongside type/name/provider/model — this same value is both the in-memory +// accumulation record and the webhook `calls[]` element. +type CallUsage struct { + Type string `json:"type"` // "llm_call" | "tool_call" + Name string `json:"name"` // e.g. "9router/cx/gpt-5.6 #3" or "read_image" + Provider string `json:"provider"` + Model string `json:"model"` + Usage // embedded, NO json tag → promotes token fields to flat CallUsage JSON + CostUSD float64 `json:"cost_usd,omitempty"` +} + +// SumCallUsage folds per-call usage into one aggregate: summed tokens/cache, +// OR-ed PromptTokensIncludeCachedSegments. Zero value for an empty slice. +func SumCallUsage(calls []CallUsage) Usage { + var total Usage + for _, c := range calls { + total.PromptTokens += c.PromptTokens + total.CompletionTokens += c.CompletionTokens + total.TotalTokens += c.TotalTokens + total.CacheReadTokens += c.CacheReadTokens + total.CacheCreationTokens += c.CacheCreationTokens + total.ThinkingTokens += c.ThinkingTokens + if c.PromptTokensIncludeCachedSegments { + total.PromptTokensIncludeCachedSegments = true + } + } + return total +} + +// SumCallCost totals the per-call USD cost (0 when pricing was unavailable). +func SumCallCost(calls []CallUsage) float64 { + var total float64 + for _, c := range calls { + total += c.CostUSD + } + return total +} diff --git a/internal/providers/call_usage_test.go b/internal/providers/call_usage_test.go new file mode 100644 index 00000000..abd5b65d --- /dev/null +++ b/internal/providers/call_usage_test.go @@ -0,0 +1,48 @@ +package providers + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestSumCallUsage(t *testing.T) { + calls := []CallUsage{ + {Type: "llm_call", Provider: "p", Model: "m", + Usage: Usage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120, + CacheReadTokens: 80, PromptTokensIncludeCachedSegments: true}, CostUSD: 0.01}, + {Type: "tool_call", Provider: "q", Model: "n", + Usage: Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15}, CostUSD: 0.002}, + } + got := SumCallUsage(calls) + if got.PromptTokens != 110 || got.CompletionTokens != 25 || got.TotalTokens != 135 { + t.Errorf("base tokens = %+v, want 110/25/135", got) + } + if got.CacheReadTokens != 80 { + t.Errorf("CacheReadTokens = %d, want 80", got.CacheReadTokens) + } + if !got.PromptTokensIncludeCachedSegments { + t.Error("flag not OR-ed") + } + if cost := SumCallCost(calls); cost < 0.0119 || cost > 0.0121 { + t.Errorf("SumCallCost = %f, want ~0.012", cost) + } + if SumCallUsage(nil) != (Usage{}) { + t.Error("nil should sum to zero Usage") + } +} + +func TestCallUsageFlatJSON(t *testing.T) { + b, _ := json.Marshal(CallUsage{Type: "llm_call", Name: "x", Provider: "p", Model: "m", + Usage: Usage{PromptTokens: 5, CompletionTokens: 1, TotalTokens: 6, CacheReadTokens: 3}, CostUSD: 0.5}) + s := string(b) + for _, want := range []string{`"type":"llm_call"`, `"provider":"p"`, `"model":"m"`, + `"prompt_tokens":5`, `"cache_read_input_tokens":3`, `"cost_usd":0.5`} { + if !strings.Contains(s, want) { + t.Errorf("missing %s in %s", want, s) + } + } + if strings.Contains(s, `"usage"`) { + t.Errorf("usage must be flattened, not nested: %s", s) + } +} diff --git a/internal/webhooks/worker.go b/internal/webhooks/worker.go index d4d0116d..a921fce7 100644 --- a/internal/webhooks/worker.go +++ b/internal/webhooks/worker.go @@ -93,14 +93,16 @@ func decodeAsyncPayload(payload []byte) (asyncPayload, error) { // callbackPayload is the JSON body POSTed to the receiver's callback_url. type callbackPayload struct { - CallID string `json:"call_id"` - DeliveryID string `json:"delivery_id"` - AgentID string `json:"agent_id,omitempty"` - Status string `json:"status"` // "done" | "failed" - Output string `json:"output,omitempty"` - Usage *callbackUsage `json:"usage,omitempty"` - Metadata json.RawMessage `json:"metadata,omitempty"` - Error string `json:"error,omitempty"` + CallID string `json:"call_id"` + DeliveryID string `json:"delivery_id"` + AgentID string `json:"agent_id,omitempty"` + Status string `json:"status"` // "done" | "failed" + Output string `json:"output,omitempty"` + Usage *callbackUsage `json:"usage,omitempty"` + Calls []providers.CallUsage `json:"calls,omitempty"` + TotalCostUSD float64 `json:"total_cost_usd,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` + Error string `json:"error,omitempty"` } // callbackUsage mirrors providers.Usage for the callback payload. @@ -129,6 +131,16 @@ func newCallbackUsage(u *providers.Usage) *callbackUsage { } } +// callbackBreakdown derives the aggregate usage (sum of all calls) and total +// cost from a run's per-call breakdown. Returns (nil, 0) for an empty slice. +func callbackBreakdown(calls []providers.CallUsage) (*callbackUsage, float64) { + if len(calls) == 0 { + return nil, 0 + } + sum := providers.SumCallUsage(calls) + return newCallbackUsage(&sum), providers.SumCallCost(calls) +} + // WorkerConfig holds tunable parameters for WebhookWorker. type WorkerConfig struct { // WorkerConcurrency is the number of parallel claim-and-deliver goroutines. @@ -362,9 +374,11 @@ func (w *WebhookWorker) execute(ctx context.Context, call *store.WebhookCallData var output string var usageVal *callbackUsage var agentErrMsg string + var payloadCalls []providers.CallUsage + var payloadCost float64 if len(call.Response) == 0 && call.AgentID != nil { - out, usage, invokeErr := w.invokeAgentWithHeartbeat(tctx, call, req, lease) + out, usage, calls, invokeErr := w.invokeAgentWithHeartbeat(tctx, call, req, lease) if invokeErr != nil { agentErrMsg = invokeErr.Error() slog.Warn("webhook.worker.agent_invoke_failed", @@ -375,6 +389,11 @@ func (w *WebhookWorker) execute(ctx context.Context, call *store.WebhookCallData } else { output = out usageVal = usage + if bd, cost := callbackBreakdown(calls); bd != nil { + usageVal = bd // usage = sum of all calls (includes tool-internal LLM) + payloadCalls = calls + payloadCost = cost + } } } else if len(call.Response) > 0 { // Prior attempt stored a partial response; extract output for re-delivery. @@ -382,6 +401,8 @@ func (w *WebhookWorker) execute(ctx context.Context, call *store.WebhookCallData if err := json.Unmarshal(call.Response, &prevResp); err == nil { output = prevResp.Output usageVal = prevResp.Usage + payloadCalls = prevResp.Calls + payloadCost = prevResp.TotalCostUSD } } @@ -416,14 +437,16 @@ func (w *WebhookWorker) execute(ctx context.Context, call *store.WebhookCallData } payload := callbackPayload{ - CallID: call.ID.String(), - DeliveryID: call.DeliveryID.String(), - AgentID: agentIDStr, - Status: statusStr, - Output: output, - Usage: usageVal, - Metadata: req.Metadata, - Error: agentErrMsg, + CallID: call.ID.String(), + DeliveryID: call.DeliveryID.String(), + AgentID: agentIDStr, + Status: statusStr, + Output: output, + Usage: usageVal, + Calls: payloadCalls, + TotalCostUSD: payloadCost, + Metadata: req.Metadata, + Error: agentErrMsg, } bodyBytes, err := json.Marshal(payload) if err != nil { @@ -758,7 +781,7 @@ func (w *WebhookWorker) invokeAgentWithHeartbeat( call *store.WebhookCallData, req asyncPayload, lease string, -) (string, *callbackUsage, error) { +) (string, *callbackUsage, []providers.CallUsage, error) { runCtx, cancelRun := context.WithCancel(ctx) defer cancelRun() @@ -806,29 +829,29 @@ func (w *WebhookWorker) heartbeatLoop( } } -// invokeAgent runs the agent for an async call and returns (output, usage, error). +// invokeAgent runs the agent for an async call and returns (output, usage, calls, error). func (w *WebhookWorker) invokeAgent( ctx context.Context, call *store.WebhookCallData, req asyncPayload, -) (string, *callbackUsage, error) { +) (string, *callbackUsage, []providers.CallUsage, error) { if call.AgentID == nil { - return "", nil, fmt.Errorf("call has no agent_id") + return "", nil, nil, fmt.Errorf("call has no agent_id") } agentIDStr := call.AgentID.String() ag, err := w.router.Get(ctx, agentIDStr) if err != nil { - return "", nil, fmt.Errorf("agent lookup %s: %w", agentIDStr, err) + return "", nil, nil, fmt.Errorf("agent lookup %s: %w", agentIDStr, err) } // Parse input. userMessage, extraSystem, err := parseAsyncInput(req.Input) if err != nil { - return "", nil, fmt.Errorf("parse input: %w", err) + return "", nil, nil, fmt.Errorf("parse input: %w", err) } if userMessage == "" { - return "", nil, fmt.Errorf("empty user message in stored payload") + return "", nil, nil, fmt.Errorf("empty user message in stored payload") } runID := uuid.NewString() @@ -860,11 +883,11 @@ func (w *WebhookWorker) invokeAgent( result, runErr := ag.Run(agentCtx, rr) if runErr != nil { - return "", nil, runErr + return "", nil, nil, runErr } usage := newCallbackUsage(result.Usage) - return result.Content, usage, nil + return result.Content, usage, result.Calls, nil } // reclaimStale resets stale running rows back to queued. diff --git a/internal/webhooks/worker_test.go b/internal/webhooks/worker_test.go index bc2e3a9c..7fb0dbe0 100644 --- a/internal/webhooks/worker_test.go +++ b/internal/webhooks/worker_test.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "sync/atomic" "testing" "time" @@ -29,8 +30,8 @@ type stubCallStore struct { claimErr error // if non-nil, returned by ClaimNext reclaimN int64 // count returned by ReclaimStale casLeaseErr error // if non-nil, returned by UpdateStatusCAS - hbCount int32 // số lần Heartbeat được gọi (atomic) - hbErr error // nếu non-nil, Heartbeat trả về lỗi này + hbCount int32 // số lần Heartbeat được gọi (atomic) + hbErr error // nếu non-nil, Heartbeat trả về lỗi này } func newStubCallStore(initial *store.WebhookCallData) *stubCallStore { @@ -327,6 +328,78 @@ func TestHMACHeaderPresent(t *testing.T) { } } +// TestExecuteReDeliveryCarriesCallsBreakdown verifies that when a call already has +// a stored callbackPayload (call.Response set, simulating a re-delivery/retry), the +// re-sent HTTP callback body carries forward the per-call breakdown (`calls`) and +// `total_cost_usd` fields from the previously-stored payload. +func TestExecuteReDeliveryCarriesCallsBreakdown(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + + var gotBody []byte + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + agentID := uuid.New() + call := newTestCall(srv.URL, &agentID) + // Pre-populate response so agent invocation is skipped and the re-delivery + // path (execute's "else if len(call.Response) > 0" branch) is exercised. + prevResp, _ := json.Marshal(callbackPayload{ + Output: "prior output", + Status: "done", + Calls: []providers.CallUsage{ + { + Type: "tool_call", + Name: "read_image", + Provider: "9router", + Model: "cx/gpt-5.5", + Usage: providers.Usage{PromptTokens: 100, TotalTokens: 100}, + CostUSD: 0.02, + }, + }, + TotalCostUSD: 0.02, + }) + call.Response = prevResp + + wh, _ := newTestWebhook(call.WebhookID, testEncKey) + callStore := newStubCallStore(call) + whStore := &stubWebhookStore{wh: wh} + + w := newTestWorker(callStore, whStore) + w.execute(context.Background(), call, call.TenantID, "test-lease") + + if len(gotBody) == 0 { + t.Fatal("no callback body captured — request never reached server") + } + + var payload callbackPayload + if err := json.Unmarshal(gotBody, &payload); err != nil { + t.Fatalf("unmarshal captured body: %v\nbody: %s", err, gotBody) + } + + if len(payload.Calls) != 1 { + t.Fatalf("payload.Calls length: got %d, want 1 (body: %s)", len(payload.Calls), gotBody) + } + if payload.Calls[0].Name != "read_image" { + t.Errorf("payload.Calls[0].Name: got %q, want %q", payload.Calls[0].Name, "read_image") + } + if payload.TotalCostUSD < 0.0199 || payload.TotalCostUSD > 0.0201 { + t.Errorf("payload.TotalCostUSD: got %f, want ~0.02", payload.TotalCostUSD) + } + + // Belt-and-suspenders: also confirm the raw JSON carries the expected keys. + bodyStr := string(gotBody) + for _, want := range []string{`"calls"`, `"read_image"`, `"total_cost_usd"`} { + if !strings.Contains(bodyStr, want) { + t.Errorf("callback body missing %s: %s", want, bodyStr) + } + } +} + // TestDeliveryIDStableAcrossRetries verifies same delivery_id sent on attempt 1 and 3. func TestDeliveryIDStableAcrossRetries(t *testing.T) { security.SetAllowLoopbackForTest(true) @@ -760,6 +833,24 @@ func TestHeartbeatLoopRenews(t *testing.T) { } } +func TestBuildCallbackBreakdown(t *testing.T) { + calls := []providers.CallUsage{ + {Type: "llm_call", Provider: "p", Model: "m", Usage: providers.Usage{PromptTokens: 10, TotalTokens: 10}, CostUSD: 0.01}, + {Type: "tool_call", Name: "read_image", Provider: "9router", Model: "cx/gpt-5.5", + Usage: providers.Usage{PromptTokens: 100, TotalTokens: 100}, CostUSD: 0.02}, + } + usage, cost := callbackBreakdown(calls) + if usage == nil || usage.PromptTokens != 110 { + t.Errorf("usage sum wrong: %+v", usage) + } + if cost < 0.0299 || cost > 0.0301 { + t.Errorf("cost = %f, want ~0.03", cost) + } + if usage2, _ := callbackBreakdown(nil); usage2 != nil { + t.Error("nil calls → nil usage") + } +} + // TestHeartbeatLoopCancelsOnLeaseLost verifies that when Heartbeat returns ErrLeaseExpired, // the loop cancels runCtx and returns. func TestHeartbeatLoopCancelsOnLeaseLost(t *testing.T) {