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
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue