marketing-shibata50/multica
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(usage): add per-task token usage tracking

Browse source 
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:
Jiang Bohan 2026-04-08 13:08:15 +08:00
parent abcc7bf3cd
commit 8a8d3ea20e
14 changed files with 477 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

9
server/pkg/agent/agent.go
View file

@ -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.

27
server/pkg/agent/claude.go
View file

@ -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 {

6
server/pkg/agent/claude_test.go
View file

@ -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())

12
server/pkg/db/generated/models.go
View file

@ -281,6 +281,18 @@ type TaskMessage struct {
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type TaskUsage struct {
ID pgtype.UUID `json:"id"`
TaskID pgtype.UUID `json:"task_id"`
Provider string `json:"provider"`
Model string `json:"model"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadTokens int64 `json:"cache_read_tokens"`
CacheWriteTokens int64 `json:"cache_write_tokens"`
CreatedAt pgtype.Timestamptz `json:"created_at"`
}
type User struct {
ID pgtype.UUID `json:"id"`
Name string `json:"name"`

201
server/pkg/db/generated/task_usage.sql.go Normal file
View file

@ -0,0 +1,201 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.30.0
// source: task_usage.sql
package db
import (
"context"
"github.com/jackc/pgx/v5/pgtype"
)
const getTaskUsage = `-- name: GetTaskUsage :many
SELECT id, task_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, created_at FROM task_usage
WHERE task_id = $1
ORDER BY model
`
func (q *Queries) GetTaskUsage(ctx context.Context, taskID pgtype.UUID) ([]TaskUsage, error) {
rows, err := q.db.Query(ctx, getTaskUsage, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
items := []TaskUsage{}
for rows.Next() {
var i TaskUsage
if err := rows.Scan(
&i.ID,
&i.TaskID,
&i.Provider,
&i.Model,
&i.InputTokens,
&i.OutputTokens,
&i.CacheReadTokens,
&i.CacheWriteTokens,
&i.CreatedAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getWorkspaceUsageByDay = `-- name: GetWorkspaceUsageByDay :many
SELECT
DATE(atq.created_at) AS date,
tu.model,
SUM(tu.input_tokens)::bigint AS total_input_tokens,
SUM(tu.output_tokens)::bigint AS total_output_tokens,
SUM(tu.cache_read_tokens)::bigint AS total_cache_read_tokens,
SUM(tu.cache_write_tokens)::bigint AS total_cache_write_tokens,
COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
JOIN agent a ON a.id = atq.agent_id
WHERE a.workspace_id = $1
AND atq.created_at >= $2::timestamptz
GROUP BY DATE(atq.created_at), tu.model
ORDER BY DATE(atq.created_at) DESC, tu.model
`
type GetWorkspaceUsageByDayParams struct {
WorkspaceID pgtype.UUID `json:"workspace_id"`
Since pgtype.Timestamptz `json:"since"`
}
type GetWorkspaceUsageByDayRow struct {
Date pgtype.Date `json:"date"`
Model string `json:"model"`
TotalInputTokens int64 `json:"total_input_tokens"`
TotalOutputTokens int64 `json:"total_output_tokens"`
TotalCacheReadTokens int64 `json:"total_cache_read_tokens"`
TotalCacheWriteTokens int64 `json:"total_cache_write_tokens"`
TaskCount int32 `json:"task_count"`
}
func (q *Queries) GetWorkspaceUsageByDay(ctx context.Context, arg GetWorkspaceUsageByDayParams) ([]GetWorkspaceUsageByDayRow, error) {
rows, err := q.db.Query(ctx, getWorkspaceUsageByDay, arg.WorkspaceID, arg.Since)
if err != nil {
return nil, err
}
defer rows.Close()
items := []GetWorkspaceUsageByDayRow{}
for rows.Next() {
var i GetWorkspaceUsageByDayRow
if err := rows.Scan(
&i.Date,
&i.Model,
&i.TotalInputTokens,
&i.TotalOutputTokens,
&i.TotalCacheReadTokens,
&i.TotalCacheWriteTokens,
&i.TaskCount,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getWorkspaceUsageSummary = `-- name: GetWorkspaceUsageSummary :many
SELECT
tu.model,
SUM(tu.input_tokens)::bigint AS total_input_tokens,
SUM(tu.output_tokens)::bigint AS total_output_tokens,
SUM(tu.cache_read_tokens)::bigint AS total_cache_read_tokens,
SUM(tu.cache_write_tokens)::bigint AS total_cache_write_tokens,
COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
JOIN agent a ON a.id = atq.agent_id
WHERE a.workspace_id = $1
AND atq.created_at >= $2::timestamptz
GROUP BY tu.model
ORDER BY (SUM(tu.input_tokens) + SUM(tu.output_tokens)) DESC
`
type GetWorkspaceUsageSummaryParams struct {
WorkspaceID pgtype.UUID `json:"workspace_id"`
Since pgtype.Timestamptz `json:"since"`
}
type GetWorkspaceUsageSummaryRow struct {
Model string `json:"model"`
TotalInputTokens int64 `json:"total_input_tokens"`
TotalOutputTokens int64 `json:"total_output_tokens"`
TotalCacheReadTokens int64 `json:"total_cache_read_tokens"`
TotalCacheWriteTokens int64 `json:"total_cache_write_tokens"`
TaskCount int32 `json:"task_count"`
}
func (q *Queries) GetWorkspaceUsageSummary(ctx context.Context, arg GetWorkspaceUsageSummaryParams) ([]GetWorkspaceUsageSummaryRow, error) {
rows, err := q.db.Query(ctx, getWorkspaceUsageSummary, arg.WorkspaceID, arg.Since)
if err != nil {
return nil, err
}
defer rows.Close()
items := []GetWorkspaceUsageSummaryRow{}
for rows.Next() {
var i GetWorkspaceUsageSummaryRow
if err := rows.Scan(
&i.Model,
&i.TotalInputTokens,
&i.TotalOutputTokens,
&i.TotalCacheReadTokens,
&i.TotalCacheWriteTokens,
&i.TaskCount,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const upsertTaskUsage = `-- name: UpsertTaskUsage :exec
INSERT INTO task_usage (task_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1, $2, $3, $4, $5, $6, $7)
ON CONFLICT (task_id, provider, model)
DO UPDATE SET
input_tokens = EXCLUDED.input_tokens,
output_tokens = EXCLUDED.output_tokens,
cache_read_tokens = EXCLUDED.cache_read_tokens,
cache_write_tokens = EXCLUDED.cache_write_tokens
`
type UpsertTaskUsageParams struct {
TaskID pgtype.UUID `json:"task_id"`
Provider string `json:"provider"`
Model string `json:"model"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CacheReadTokens int64 `json:"cache_read_tokens"`
CacheWriteTokens int64 `json:"cache_write_tokens"`
}
func (q *Queries) UpsertTaskUsage(ctx context.Context, arg UpsertTaskUsageParams) error {
_, err := q.db.Exec(ctx, upsertTaskUsage,
arg.TaskID,
arg.Provider,
arg.Model,
arg.InputTokens,
arg.OutputTokens,
arg.CacheReadTokens,
arg.CacheWriteTokens,
)
return err
}

47
server/pkg/db/queries/task_usage.sql Normal file
View file

@ -0,0 +1,47 @@
-- name: UpsertTaskUsage :exec
INSERT INTO task_usage (task_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
VALUES ($1, $2, $3, $4, $5, $6, $7)
ON CONFLICT (task_id, provider, model)
DO UPDATE SET
input_tokens = EXCLUDED.input_tokens,
output_tokens = EXCLUDED.output_tokens,
cache_read_tokens = EXCLUDED.cache_read_tokens,
cache_write_tokens = EXCLUDED.cache_write_tokens;
-- name: GetTaskUsage :many
SELECT * FROM task_usage
WHERE task_id = $1
ORDER BY model;
-- name: GetWorkspaceUsageByDay :many
SELECT
DATE(atq.created_at) AS date,
tu.model,
SUM(tu.input_tokens)::bigint AS total_input_tokens,
SUM(tu.output_tokens)::bigint AS total_output_tokens,
SUM(tu.cache_read_tokens)::bigint AS total_cache_read_tokens,
SUM(tu.cache_write_tokens)::bigint AS total_cache_write_tokens,
COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
JOIN agent a ON a.id = atq.agent_id
WHERE a.workspace_id = $1
AND atq.created_at >= @since::timestamptz
GROUP BY DATE(atq.created_at), tu.model
ORDER BY DATE(atq.created_at) DESC, tu.model;
-- name: GetWorkspaceUsageSummary :many
SELECT
tu.model,
SUM(tu.input_tokens)::bigint AS total_input_tokens,
SUM(tu.output_tokens)::bigint AS total_output_tokens,
SUM(tu.cache_read_tokens)::bigint AS total_cache_read_tokens,
SUM(tu.cache_write_tokens)::bigint AS total_cache_write_tokens,
COUNT(DISTINCT tu.task_id)::int AS task_count
FROM task_usage tu
JOIN agent_task_queue atq ON atq.id = tu.task_id
JOIN agent a ON a.id = atq.agent_id
WHERE a.workspace_id = $1
AND atq.created_at >= @since::timestamptz
GROUP BY tu.model
ORDER BY (SUM(tu.input_tokens) + SUM(tu.output_tokens)) DESC;