fix(usage): address review feedback — independent usage reporting + all providers

1. Separate ReportTaskUsage endpoint (POST /api/daemon/tasks/{id}/usage)
   so usage is captured independently of complete/fail — fixes usage loss
   for failed/blocked tasks.

2. Add usage tracking for all four providers:
   - Claude: already done (stream-json message.usage)
   - OpenCode: extract from step_finish.part.tokens
   - OpenClaw: extract from step_end.data token fields
   - Codex: extract from turn/completed and task_complete usage fields

3. Remove usage from CompleteTask payload — all usage goes through the
   dedicated endpoint now.
This commit is contained in:
Jiang Bohan 2026-04-08 13:23:54 +08:00
parent 8a8d3ea20e
commit fa0c0fe747
7 changed files with 196 additions and 25 deletions

View file

@ -220,11 +220,25 @@ func (b *codexBackend) Execute(ctx context.Context, prompt string, opts ExecOpti
finalOutput := output.String()
outputMu.Unlock()
// Build usage map from accumulated codex usage.
var usageMap map[string]TokenUsage
c.usageMu.Lock()
u := c.usage
c.usageMu.Unlock()
if u.InputTokens > 0 || u.OutputTokens > 0 || u.CacheReadTokens > 0 || u.CacheWriteTokens > 0 {
model := opts.Model
if model == "" {
model = "unknown"
}
usageMap = map[string]TokenUsage{model: u}
}
resCh <- Result{
Status: finalStatus,
Output: finalOutput,
Error: finalError,
DurationMs: duration.Milliseconds(),
Usage: usageMap,
}
}()
@ -247,6 +261,9 @@ type codexClient struct {
notificationProtocol string // "unknown", "legacy", "raw"
turnStarted bool
completedTurnIDs map[string]bool
usageMu sync.Mutex
usage TokenUsage // accumulated from turn events
}
type pendingRPC struct {
@ -498,6 +515,8 @@ func (c *codexClient) handleEvent(msg map[string]any) {
})
}
case "task_complete":
// Extract usage from legacy task_complete if present.
c.extractUsageFromMap(msg)
if c.onTurnDone != nil {
c.onTurnDone(false)
}
@ -535,6 +554,11 @@ func (c *codexClient) handleRawNotification(method string, params map[string]any
c.completedTurnIDs[turnID] = true
}
// Extract usage from turn/completed if present (e.g. params.turn.usage).
if turn, ok := params["turn"].(map[string]any); ok {
c.extractUsageFromMap(turn)
}
if c.onTurnDone != nil {
c.onTurnDone(aborted)
}
@ -618,6 +642,48 @@ func (c *codexClient) handleItemNotification(method string, params map[string]an
}
}
// extractUsageFromMap extracts token usage from a map that may contain
// "usage", "token_usage", or "tokens" fields. Handles various Codex formats.
func (c *codexClient) extractUsageFromMap(data map[string]any) {
// Try common field names for usage data.
var usageMap map[string]any
for _, key := range []string{"usage", "token_usage", "tokens"} {
if v, ok := data[key].(map[string]any); ok {
usageMap = v
break
}
}
if usageMap == nil {
return
}
c.usageMu.Lock()
defer c.usageMu.Unlock()
// Try various key conventions.
c.usage.InputTokens += codexInt64(usageMap, "input_tokens", "input", "prompt_tokens")
c.usage.OutputTokens += codexInt64(usageMap, "output_tokens", "output", "completion_tokens")
c.usage.CacheReadTokens += codexInt64(usageMap, "cache_read_tokens", "cache_read_input_tokens")
c.usage.CacheWriteTokens += codexInt64(usageMap, "cache_write_tokens", "cache_creation_input_tokens")
}
// codexInt64 returns the first non-zero int64 value from the map for the given keys.
func codexInt64(m map[string]any, keys ...string) int64 {
for _, key := range keys {
switch v := m[key].(type) {
case float64:
if v != 0 {
return int64(v)
}
case int64:
if v != 0 {
return v
}
}
}
return 0
}
// ── Helpers ──
func extractThreadID(result json.RawMessage) string {