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:
Jiang Bohan
parent
abcc7bf3cd
commit
8a8d3ea20e
14 changed files with 477 additions and 19 deletions
201
server/pkg/db/generated/task_usage.sql.go
Normal file
201
server/pkg/db/generated/task_usage.sql.go
Normal 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
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue