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 2026-04-08 13:08:15 +08:00
parent abcc7bf3cd
commit 8a8d3ea20e
14 changed files with 477 additions and 19 deletions

View file

@ -341,10 +341,21 @@ func (h *Handler) ReportTaskProgress(w http.ResponseWriter, r *http.Request) {
// CompleteTask marks a running task as completed.
type TaskCompleteRequest struct {
PRURL string `json:"pr_url"`
Output string `json:"output"`
SessionID string `json:"session_id"` // Claude session ID for future resumption
WorkDir string `json:"work_dir"` // working directory used during execution
PRURL string `json:"pr_url"`
Output string `json:"output"`
SessionID string `json:"session_id"` // Claude session ID for future resumption
WorkDir string `json:"work_dir"` // working directory used during execution
Usage []TaskUsagePayload `json:"usage,omitempty"`
}
// TaskUsagePayload is the per-model token usage reported by the daemon.
type TaskUsagePayload struct {
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 (h *Handler) CompleteTask(w http.ResponseWriter, r *http.Request) {
@ -364,6 +375,21 @@ func (h *Handler) CompleteTask(w http.ResponseWriter, r *http.Request) {
return
}
// Store per-task token usage (best-effort, don't fail the request).
for _, u := range req.Usage {
if err := h.Queries.UpsertTaskUsage(r.Context(), db.UpsertTaskUsageParams{
TaskID: parseUUID(taskID),
Provider: u.Provider,
Model: u.Model,
InputTokens: u.InputTokens,
OutputTokens: u.OutputTokens,
CacheReadTokens: u.CacheReadTokens,
CacheWriteTokens: u.CacheWriteTokens,
}); err != nil {
slog.Warn("upsert task usage failed", "task_id", taskID, "model", u.Model, "error", err)
}
}
slog.Info("task completed", "task_id", taskID, "agent_id", uuidToString(task.AgentID))
writeJSON(w, http.StatusOK, taskToResponse(*task))
}

View file

@ -192,6 +192,96 @@ func (h *Handler) GetRuntimeTaskActivity(w http.ResponseWriter, r *http.Request)
writeJSON(w, http.StatusOK, resp)
}
// GetWorkspaceUsageByDay returns daily token usage aggregated by model for the workspace.
func (h *Handler) GetWorkspaceUsageByDay(w http.ResponseWriter, r *http.Request) {
workspaceID := resolveWorkspaceID(r)
since := parseSinceParam(r, 30)
rows, err := h.Queries.GetWorkspaceUsageByDay(r.Context(), db.GetWorkspaceUsageByDayParams{
WorkspaceID: parseUUID(workspaceID),
Since: since,
})
if err != nil {
writeError(w, http.StatusInternalServerError, "failed to get usage")
return
}
type DailyUsageRow struct {
Date string `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"`
}
resp := make([]DailyUsageRow, len(rows))
for i, row := range rows {
resp[i] = DailyUsageRow{
Date: row.Date.Time.Format("2006-01-02"),
Model: row.Model,
TotalInputTokens: row.TotalInputTokens,
TotalOutputTokens: row.TotalOutputTokens,
TotalCacheReadTokens: row.TotalCacheReadTokens,
TotalCacheWriteTokens: row.TotalCacheWriteTokens,
TaskCount: row.TaskCount,
}
}
writeJSON(w, http.StatusOK, resp)
}
// GetWorkspaceUsageSummary returns total token usage aggregated by model for the workspace.
func (h *Handler) GetWorkspaceUsageSummary(w http.ResponseWriter, r *http.Request) {
workspaceID := resolveWorkspaceID(r)
since := parseSinceParam(r, 30)
rows, err := h.Queries.GetWorkspaceUsageSummary(r.Context(), db.GetWorkspaceUsageSummaryParams{
WorkspaceID: parseUUID(workspaceID),
Since: since,
})
if err != nil {
writeError(w, http.StatusInternalServerError, "failed to get usage summary")
return
}
type UsageSummaryRow 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"`
}
resp := make([]UsageSummaryRow, len(rows))
for i, row := range rows {
resp[i] = UsageSummaryRow{
Model: row.Model,
TotalInputTokens: row.TotalInputTokens,
TotalOutputTokens: row.TotalOutputTokens,
TotalCacheReadTokens: row.TotalCacheReadTokens,
TotalCacheWriteTokens: row.TotalCacheWriteTokens,
TaskCount: row.TaskCount,
}
}
writeJSON(w, http.StatusOK, resp)
}
// parseSinceParam parses the "days" query parameter and returns a timestamptz.
func parseSinceParam(r *http.Request, defaultDays int) pgtype.Timestamptz {
days := defaultDays
if d := r.URL.Query().Get("days"); d != "" {
if parsed, err := strconv.Atoi(d); err == nil && parsed > 0 && parsed <= 365 {
days = parsed
}
}
t := time.Now().AddDate(0, 0, -days)
return pgtype.Timestamptz{Time: t, Valid: true}
}
func (h *Handler) ListAgentRuntimes(w http.ResponseWriter, r *http.Request) {
workspaceID := resolveWorkspaceID(r)