feat(webhooks): return per-call usage breakdown with provider/model/cost (#1421)

feat(webhooks): per-call usage breakdown with provider/model/cost (#1421)
This commit is contained in:
thotam authored and GitHub committed 2026-07-10 22:56:05 +07:00
1 parent 6b738924d5
commit 0c1ededc92
18 files changed
+551 -59

No files matched your search

+53 -6
View File
@@ -264,13 +264,29 @@ Triggers an agent with an input prompt. Available in all editions.
"agent_id": "<uuid>",
"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
+21
View File
@@ -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)
}
}
+45
View File
@@ -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)
}
}
+1
View File
@@ -281,6 +281,7 @@ func convertRunResult(pr *pipeline.RunResult) *RunResult {
BlockReplies: pr.BlockReplies,
LastBlockReply: pr.LastBlockReply,
LoopKilled: pr.LoopKilled,
Calls: pr.Calls,
}
}
+11
View File
@@ -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
}
}
@@ -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
}
+4 -1
View File
@@ -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)
}
+11 -10
View File
@@ -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.
+43
View File
@@ -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)
}
}
+34 -7
View File
@@ -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,
+34 -7
View File
@@ -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 {
+14
View File
@@ -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,
}
}
+29
View File
@@ -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))
}
}
+1
View File
@@ -117,4 +117,5 @@ type RunResult struct {
Deliverables []string
BlockReplies int
LastBlockReply string
Calls []providers.CallUsage
}
+42
View File
@@ -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
}
+48
View File
@@ -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)
}
}
+49 -26
View File
@@ -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.
+93 -2
View File
@@ -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) {