feat(usage): add per-task token usage tracking
Extract token usage from Claude Code's stream-json output in real-time during task execution, replacing the inaccurate global JSONL log scanner. - New `task_usage` table: tracks (task_id, provider, model) level usage - Agent SDK: parse `message.usage` from assistant messages, accumulate per-model and return in Result - Daemon: convert agent usage to entries, send with CompleteTask - Server: store usage on task completion, expose workspace-level aggregation APIs (GET /api/usage/daily, GET /api/usage/summary)
This commit is contained in:
parent
abcc7bf3cd
commit
8a8d3ea20e
14 changed files with 477 additions and 19 deletions
|
|
@ -62,6 +62,14 @@ type Message struct {
|
|||
Level string // log level (Log)
|
||||
}
|
||||
|
||||
// TokenUsage tracks token consumption for a single model.
|
||||
type TokenUsage struct {
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
CacheReadTokens int64
|
||||
CacheWriteTokens int64
|
||||
}
|
||||
|
||||
// Result is the final outcome after an agent session completes.
|
||||
type Result struct {
|
||||
Status string // "completed", "failed", "aborted", "timeout"
|
||||
|
|
@ -69,6 +77,7 @@ type Result struct {
|
|||
Error string // error message if failed
|
||||
DurationMs int64
|
||||
SessionID string
|
||||
Usage map[string]TokenUsage // keyed by model name
|
||||
}
|
||||
|
||||
// Config configures a Backend instance.
|
||||
|
|
|
|||
|
|
@ -91,6 +91,7 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt
|
|||
var sessionID string
|
||||
finalStatus := "completed"
|
||||
var finalError string
|
||||
usage := make(map[string]TokenUsage)
|
||||
|
||||
scanner := bufio.NewScanner(stdout)
|
||||
scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024)
|
||||
|
|
@ -108,7 +109,7 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt
|
|||
|
||||
switch msg.Type {
|
||||
case "assistant":
|
||||
b.handleAssistant(msg, msgCh, &output)
|
||||
b.handleAssistant(msg, msgCh, &output, usage)
|
||||
case "user":
|
||||
b.handleUser(msg, msgCh)
|
||||
case "system":
|
||||
|
|
@ -162,18 +163,29 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt
|
|||
Error: finalError,
|
||||
DurationMs: duration.Milliseconds(),
|
||||
SessionID: sessionID,
|
||||
Usage: usage,
|
||||
}
|
||||
}()
|
||||
|
||||
return &Session{Messages: msgCh, Result: resCh}, nil
|
||||
}
|
||||
|
||||
func (b *claudeBackend) handleAssistant(msg claudeSDKMessage, ch chan<- Message, output *strings.Builder) {
|
||||
func (b *claudeBackend) handleAssistant(msg claudeSDKMessage, ch chan<- Message, output *strings.Builder, usage map[string]TokenUsage) {
|
||||
var content claudeMessageContent
|
||||
if err := json.Unmarshal(msg.Message, &content); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Accumulate token usage per model.
|
||||
if content.Usage != nil && content.Model != "" {
|
||||
u := usage[content.Model]
|
||||
u.InputTokens += content.Usage.InputTokens
|
||||
u.OutputTokens += content.Usage.OutputTokens
|
||||
u.CacheReadTokens += content.Usage.CacheReadInputTokens
|
||||
u.CacheWriteTokens += content.Usage.CacheCreationInputTokens
|
||||
usage[content.Model] = u
|
||||
}
|
||||
|
||||
for _, block := range content.Content {
|
||||
switch block.Type {
|
||||
case "text":
|
||||
|
|
@ -287,8 +299,17 @@ type claudeLogEntry struct {
|
|||
}
|
||||
|
||||
type claudeMessageContent struct {
|
||||
Role string `json:"role"`
|
||||
Role string `json:"role"`
|
||||
Model string `json:"model"`
|
||||
Content []claudeContentBlock `json:"content"`
|
||||
Usage *claudeUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type claudeUsage struct {
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
|
||||
}
|
||||
|
||||
type claudeContentBlock struct {
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ func TestClaudeHandleAssistantText(t *testing.T) {
|
|||
}),
|
||||
}
|
||||
|
||||
b.handleAssistant(msg, ch, &output)
|
||||
b.handleAssistant(msg, ch, &output, make(map[string]TokenUsage))
|
||||
|
||||
if output.String() != "Hello world" {
|
||||
t.Fatalf("expected output 'Hello world', got %q", output.String())
|
||||
|
|
@ -62,7 +62,7 @@ func TestClaudeHandleAssistantToolUse(t *testing.T) {
|
|||
}),
|
||||
}
|
||||
|
||||
b.handleAssistant(msg, ch, &output)
|
||||
b.handleAssistant(msg, ch, &output, make(map[string]TokenUsage))
|
||||
|
||||
if output.String() != "" {
|
||||
t.Fatalf("tool_use should not add to output, got %q", output.String())
|
||||
|
|
@ -162,7 +162,7 @@ func TestClaudeHandleAssistantInvalidJSON(t *testing.T) {
|
|||
}
|
||||
|
||||
// Should not panic
|
||||
b.handleAssistant(msg, ch, &output)
|
||||
b.handleAssistant(msg, ch, &output, make(map[string]TokenUsage))
|
||||
|
||||
if output.String() != "" {
|
||||
t.Fatalf("expected empty output for invalid JSON, got %q", output.String())
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue