mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 04:14:56 +00:00
Duc Nguyen
747b58d24c
fix(usage): repair cost analytics and display precision - Automatic OpenRouter pricing sync with cost backfill for traces/snapshots/events - Live usage data merging for current-hour dashboard accuracy - 2-decimal API cost formatting across usage/overview pages - Comprehensive test coverage (PG, SQLite, HTTP, UI) Merged by github-maintain automation.
571 lines
21 KiB
Go
571 lines
21 KiB
Go
package pg
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"log/slog"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/lib/pq"
|
|
|
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
|
)
|
|
|
|
// PGTracingStore implements store.TracingStore backed by Postgres.
|
|
type PGTracingStore struct {
|
|
db *sql.DB
|
|
}
|
|
|
|
func NewPGTracingStore(db *sql.DB) *PGTracingStore {
|
|
return &PGTracingStore{db: db}
|
|
}
|
|
|
|
func (s *PGTracingStore) CreateTrace(ctx context.Context, trace *store.TraceData) error {
|
|
if trace.ID == uuid.Nil {
|
|
trace.ID = store.GenNewID()
|
|
}
|
|
tenantID := store.TenantIDFromContext(ctx)
|
|
if tenantID == uuid.Nil {
|
|
tenantID = store.MasterTenantID
|
|
}
|
|
_, err := s.db.ExecContext(ctx,
|
|
`INSERT INTO traces (id, parent_trace_id, agent_id, user_id, session_key, run_id, start_time, end_time,
|
|
duration_ms, name, channel, input_preview, output_preview,
|
|
total_input_tokens, total_output_tokens, total_cost, span_count, llm_call_count, tool_call_count,
|
|
status, error, metadata, tags, team_id, created_at, tenant_id)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26)`,
|
|
trace.ID, nilUUID(trace.ParentTraceID), nilUUID(trace.AgentID), nilStr(trace.UserID), nilStr(trace.SessionKey),
|
|
nilStr(trace.RunID), trace.StartTime, nilTime(trace.EndTime),
|
|
nilInt(trace.DurationMS), nilStr(trace.Name), nilStr(trace.Channel),
|
|
nilStr(trace.InputPreview), nilStr(trace.OutputPreview),
|
|
trace.TotalInputTokens, trace.TotalOutputTokens, trace.TotalCost, trace.SpanCount, trace.LLMCallCount, trace.ToolCallCount,
|
|
trace.Status, nilStr(trace.Error), jsonOrEmpty(trace.Metadata), pqStringArray(trace.Tags), nilUUID(trace.TeamID), trace.CreatedAt, tenantID,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (s *PGTracingStore) UpdateTrace(ctx context.Context, traceID uuid.UUID, updates map[string]any) error {
|
|
if store.IsCrossTenant(ctx) {
|
|
return execMapUpdate(ctx, s.db, "traces", traceID, updates)
|
|
}
|
|
tid := store.TenantIDFromContext(ctx)
|
|
if tid == uuid.Nil {
|
|
return fmt.Errorf("tenant_id required for update")
|
|
}
|
|
return execMapUpdateWhereTenant(ctx, s.db, "traces", updates, traceID, tid)
|
|
}
|
|
|
|
func (s *PGTracingStore) GetTrace(ctx context.Context, traceID uuid.UUID) (*store.TraceData, error) {
|
|
query := `SELECT id, parent_trace_id, agent_id, user_id, session_key, run_id, start_time, end_time,
|
|
duration_ms, name, channel, input_preview, output_preview,
|
|
total_input_tokens, total_output_tokens, COALESCE(total_cost, 0) AS total_cost, span_count, llm_call_count, tool_call_count,
|
|
status, error, COALESCE(metadata, '{}'::jsonb) AS metadata, COALESCE(tags, '{}') AS tags, team_id, created_at
|
|
FROM traces WHERE id = $1`
|
|
qArgs := []any{traceID}
|
|
if !store.IsCrossTenant(ctx) {
|
|
tenantID := store.TenantIDFromContext(ctx)
|
|
if tenantID == uuid.Nil {
|
|
return nil, sql.ErrNoRows
|
|
}
|
|
query += ` AND tenant_id = $2`
|
|
qArgs = append(qArgs, tenantID)
|
|
}
|
|
|
|
var row traceRow
|
|
if err := pkgSqlxDB.GetContext(ctx, &row, query, qArgs...); err != nil {
|
|
return nil, err
|
|
}
|
|
d := row.toTraceData()
|
|
return &d, nil
|
|
}
|
|
|
|
func buildTraceWhere(ctx context.Context, opts store.TraceListOpts) (string, []any) {
|
|
var conditions []string
|
|
var args []any
|
|
argIdx := 1
|
|
|
|
if !store.IsCrossTenant(ctx) {
|
|
tenantID := store.TenantIDFromContext(ctx)
|
|
if tenantID == uuid.Nil {
|
|
return " WHERE 1=0", nil // fail-closed: no tenant = no results
|
|
}
|
|
conditions = append(conditions, fmt.Sprintf("tenant_id = $%d", argIdx))
|
|
args = append(args, tenantID)
|
|
argIdx++
|
|
}
|
|
|
|
if opts.AgentID != nil {
|
|
conditions = append(conditions, fmt.Sprintf("agent_id = $%d", argIdx))
|
|
args = append(args, *opts.AgentID)
|
|
argIdx++
|
|
}
|
|
if opts.UserID != "" {
|
|
conditions = append(conditions, fmt.Sprintf("user_id = $%d", argIdx))
|
|
args = append(args, opts.UserID)
|
|
argIdx++
|
|
}
|
|
if opts.SessionKey != "" {
|
|
conditions = append(conditions, fmt.Sprintf("session_key = $%d", argIdx))
|
|
args = append(args, opts.SessionKey)
|
|
argIdx++
|
|
}
|
|
if opts.Status != "" {
|
|
conditions = append(conditions, fmt.Sprintf("status = $%d", argIdx))
|
|
args = append(args, opts.Status)
|
|
argIdx++
|
|
}
|
|
if opts.Channel != "" {
|
|
conditions = append(conditions, fmt.Sprintf("channel = $%d", argIdx))
|
|
args = append(args, opts.Channel)
|
|
argIdx++
|
|
}
|
|
if opts.From != nil {
|
|
conditions = append(conditions, fmt.Sprintf("start_time >= $%d", argIdx))
|
|
args = append(args, *opts.From)
|
|
argIdx++
|
|
}
|
|
if opts.To != nil {
|
|
conditions = append(conditions, fmt.Sprintf("start_time < $%d", argIdx))
|
|
args = append(args, *opts.To)
|
|
argIdx++
|
|
}
|
|
if opts.MinInputTokens != nil {
|
|
conditions = append(conditions, fmt.Sprintf("total_input_tokens >= $%d", argIdx))
|
|
args = append(args, *opts.MinInputTokens)
|
|
argIdx++
|
|
}
|
|
if opts.MaxInputTokens != nil {
|
|
conditions = append(conditions, fmt.Sprintf("total_input_tokens <= $%d", argIdx))
|
|
args = append(args, *opts.MaxInputTokens)
|
|
argIdx++
|
|
}
|
|
if opts.MinOutputTokens != nil {
|
|
conditions = append(conditions, fmt.Sprintf("total_output_tokens >= $%d", argIdx))
|
|
args = append(args, *opts.MinOutputTokens)
|
|
argIdx++
|
|
}
|
|
if opts.MaxOutputTokens != nil {
|
|
conditions = append(conditions, fmt.Sprintf("total_output_tokens <= $%d", argIdx))
|
|
args = append(args, *opts.MaxOutputTokens)
|
|
argIdx++
|
|
}
|
|
if opts.MinToolCalls != nil {
|
|
conditions = append(conditions, fmt.Sprintf("tool_call_count >= $%d", argIdx))
|
|
args = append(args, *opts.MinToolCalls)
|
|
argIdx++
|
|
}
|
|
if opts.MaxToolCalls != nil {
|
|
conditions = append(conditions, fmt.Sprintf("tool_call_count <= $%d", argIdx))
|
|
args = append(args, *opts.MaxToolCalls)
|
|
argIdx++
|
|
}
|
|
if opts.HasToolCalls != nil {
|
|
if *opts.HasToolCalls {
|
|
conditions = append(conditions, "tool_call_count > 0")
|
|
} else {
|
|
conditions = append(conditions, "tool_call_count = 0")
|
|
}
|
|
}
|
|
if opts.Query != "" {
|
|
placeholder := fmt.Sprintf("$%d", argIdx)
|
|
conditions = append(conditions, fmt.Sprintf(`(
|
|
CAST(id AS text) ILIKE %[1]s ESCAPE '\' OR
|
|
COALESCE(name, '') ILIKE %[1]s ESCAPE '\' OR
|
|
COALESCE(input_preview, '') ILIKE %[1]s ESCAPE '\' OR
|
|
COALESCE(output_preview, '') ILIKE %[1]s ESCAPE '\' OR
|
|
COALESCE(session_key, '') ILIKE %[1]s ESCAPE '\' OR
|
|
COALESCE(channel, '') ILIKE %[1]s ESCAPE '\' OR
|
|
EXISTS (SELECT 1 FROM agents a WHERE a.id = traces.agent_id AND a.tenant_id = traces.tenant_id AND (COALESCE(a.display_name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(a.agent_key, '') ILIKE %[1]s ESCAPE '\')) OR
|
|
EXISTS (SELECT 1 FROM channel_instances ci WHERE ci.name = traces.channel AND ci.tenant_id = traces.tenant_id AND (COALESCE(ci.display_name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(ci.name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(ci.channel_type, '') ILIKE %[1]s ESCAPE '\')) OR
|
|
EXISTS (SELECT 1 FROM spans s WHERE s.trace_id = traces.id AND s.tenant_id = traces.tenant_id AND (COALESCE(s.tool_name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(s.input_preview, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(s.output_preview, '') ILIKE %[1]s ESCAPE '\'))
|
|
)`, placeholder))
|
|
args = append(args, containsPattern(opts.Query))
|
|
argIdx++
|
|
}
|
|
if opts.AgentQuery != "" {
|
|
placeholder := fmt.Sprintf("$%d", argIdx)
|
|
conditions = append(conditions, fmt.Sprintf(`EXISTS (SELECT 1 FROM agents a WHERE a.id = traces.agent_id AND a.tenant_id = traces.tenant_id AND (COALESCE(a.display_name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(a.agent_key, '') ILIKE %[1]s ESCAPE '\'))`, placeholder))
|
|
args = append(args, containsPattern(opts.AgentQuery))
|
|
argIdx++
|
|
}
|
|
if opts.ChannelQuery != "" {
|
|
placeholder := fmt.Sprintf("$%d", argIdx)
|
|
conditions = append(conditions, fmt.Sprintf(`EXISTS (SELECT 1 FROM channel_instances ci WHERE ci.name = traces.channel AND ci.tenant_id = traces.tenant_id AND (COALESCE(ci.display_name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(ci.name, '') ILIKE %[1]s ESCAPE '\' OR COALESCE(ci.channel_type, '') ILIKE %[1]s ESCAPE '\'))`, placeholder))
|
|
args = append(args, containsPattern(opts.ChannelQuery))
|
|
argIdx++
|
|
}
|
|
if opts.ToolName != "" {
|
|
placeholder := fmt.Sprintf("$%d", argIdx)
|
|
conditions = append(conditions, fmt.Sprintf(`EXISTS (SELECT 1 FROM spans s WHERE s.trace_id = traces.id AND s.tenant_id = traces.tenant_id AND s.tool_name ILIKE %[1]s ESCAPE '\')`, placeholder))
|
|
args = append(args, containsPattern(opts.ToolName))
|
|
argIdx++
|
|
}
|
|
if opts.ChangedAfter != nil {
|
|
conditions = append(conditions, fmt.Sprintf("(created_at > $%d OR end_time > $%d OR status = $%d)", argIdx, argIdx, argIdx+1))
|
|
args = append(args, *opts.ChangedAfter, store.TraceStatusRunning)
|
|
argIdx += 2
|
|
}
|
|
|
|
where := ""
|
|
if len(conditions) > 0 {
|
|
where = " WHERE " + strings.Join(conditions, " AND ")
|
|
}
|
|
return where, args
|
|
}
|
|
|
|
func containsPattern(value string) string {
|
|
return "%" + escapeLike(value) + "%"
|
|
}
|
|
|
|
func escapeLike(value string) string {
|
|
var b strings.Builder
|
|
b.Grow(len(value))
|
|
for _, r := range value {
|
|
switch r {
|
|
case '\\', '%', '_':
|
|
b.WriteRune('\\')
|
|
}
|
|
b.WriteRune(r)
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
func (s *PGTracingStore) CountTraces(ctx context.Context, opts store.TraceListOpts) (int, error) {
|
|
where, args := buildTraceWhere(ctx, opts)
|
|
var count int
|
|
err := s.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM traces"+where, args...).Scan(&count)
|
|
return count, err
|
|
}
|
|
|
|
func (s *PGTracingStore) ListTraces(ctx context.Context, opts store.TraceListOpts) ([]store.TraceData, error) {
|
|
where, args := buildTraceWhere(ctx, opts)
|
|
|
|
q := `SELECT id, parent_trace_id, agent_id, user_id, session_key, run_id, start_time, end_time,
|
|
duration_ms, name, channel, input_preview, output_preview,
|
|
total_input_tokens, total_output_tokens, COALESCE(total_cost, 0) AS total_cost, span_count, llm_call_count, tool_call_count,
|
|
status, error, metadata, tags, team_id, created_at
|
|
FROM traces` + where
|
|
|
|
limit := opts.Limit
|
|
if limit <= 0 {
|
|
limit = 50
|
|
}
|
|
q += fmt.Sprintf(" ORDER BY created_at DESC OFFSET %d LIMIT %d", opts.Offset, limit)
|
|
|
|
var rows []traceRow
|
|
if err := pkgSqlxDB.SelectContext(ctx, &rows, q, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
return traceRowsToData(rows), nil
|
|
}
|
|
|
|
func (s *PGTracingStore) ListChildTraces(ctx context.Context, parentTraceID uuid.UUID) ([]store.TraceData, error) {
|
|
q := `SELECT id, parent_trace_id, agent_id, user_id, session_key, run_id, start_time, end_time,
|
|
duration_ms, name, channel, input_preview, output_preview,
|
|
total_input_tokens, total_output_tokens, COALESCE(total_cost, 0) AS total_cost, span_count, llm_call_count, tool_call_count,
|
|
status, error, metadata, tags, team_id, created_at
|
|
FROM traces WHERE parent_trace_id = $1`
|
|
qArgs := []any{parentTraceID}
|
|
|
|
if !store.IsCrossTenant(ctx) {
|
|
tid := store.TenantIDFromContext(ctx)
|
|
if tid == uuid.Nil {
|
|
return nil, fmt.Errorf("tenant_id required")
|
|
}
|
|
q += " AND tenant_id = $2"
|
|
qArgs = append(qArgs, tid)
|
|
}
|
|
q += " ORDER BY created_at"
|
|
|
|
var rows []traceRow
|
|
if err := pkgSqlxDB.SelectContext(ctx, &rows, q, qArgs...); err != nil {
|
|
return nil, err
|
|
}
|
|
return traceRowsToData(rows), nil
|
|
}
|
|
|
|
func (s *PGTracingStore) GetSessionCosts(ctx context.Context, sessionKeys []string) (map[string]float64, error) {
|
|
result := make(map[string]float64, len(sessionKeys))
|
|
keys := compactSessionKeys(sessionKeys)
|
|
if len(keys) == 0 {
|
|
return result, nil
|
|
}
|
|
q := `SELECT session_key, COALESCE(SUM(total_cost), 0)
|
|
FROM traces
|
|
WHERE session_key = ANY($1)
|
|
AND parent_trace_id IS NULL`
|
|
args := []any{pq.Array(keys)}
|
|
if !store.IsCrossTenant(ctx) {
|
|
tid := store.TenantIDFromContext(ctx)
|
|
if tid == uuid.Nil {
|
|
return result, nil
|
|
}
|
|
q += ` AND tenant_id = $2`
|
|
args = append(args, tid)
|
|
}
|
|
q += ` GROUP BY session_key`
|
|
|
|
rows, err := s.db.QueryContext(ctx, q, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var sessionKey string
|
|
var cost float64
|
|
if err := rows.Scan(&sessionKey, &cost); err != nil {
|
|
return nil, err
|
|
}
|
|
result[sessionKey] = cost
|
|
}
|
|
return result, rows.Err()
|
|
}
|
|
|
|
func compactSessionKeys(sessionKeys []string) []string {
|
|
seen := make(map[string]struct{}, len(sessionKeys))
|
|
keys := make([]string, 0, len(sessionKeys))
|
|
for _, key := range sessionKeys {
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
continue
|
|
}
|
|
if _, ok := seen[key]; ok {
|
|
continue
|
|
}
|
|
seen[key] = struct{}{}
|
|
keys = append(keys, key)
|
|
}
|
|
return keys
|
|
}
|
|
|
|
func (s *PGTracingStore) CreateSpan(ctx context.Context, span *store.SpanData) error {
|
|
if span.ID == uuid.Nil {
|
|
span.ID = store.GenNewID()
|
|
}
|
|
tenantID := span.TenantID
|
|
if tenantID == uuid.Nil {
|
|
tenantID = store.MasterTenantID
|
|
}
|
|
_, err := s.db.ExecContext(ctx,
|
|
`INSERT INTO spans (id, trace_id, parent_span_id, agent_id, span_type, name,
|
|
start_time, end_time, duration_ms, status, error, level,
|
|
model, provider, input_tokens, output_tokens, finish_reason,
|
|
model_params, tool_name, tool_call_id, input_preview, output_preview,
|
|
metadata, team_id, created_at, tenant_id)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26)`,
|
|
span.ID, span.TraceID, span.ParentSpanID, span.AgentID, span.SpanType, nilStr(span.Name),
|
|
span.StartTime, nilTime(span.EndTime), nilInt(span.DurationMS), span.Status, nilStr(span.Error), span.Level,
|
|
nilStr(span.Model), nilStr(span.Provider), nilInt(span.InputTokens), nilInt(span.OutputTokens), nilStr(span.FinishReason),
|
|
jsonOrNull(span.ModelParams), nilStr(span.ToolName), nilStr(span.ToolCallID), nilStr(span.InputPreview), nilStr(span.OutputPreview),
|
|
jsonOrNull(span.Metadata), nilUUID(span.TeamID), span.CreatedAt, tenantID,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (s *PGTracingStore) UpdateSpan(ctx context.Context, spanID uuid.UUID, updates map[string]any) error {
|
|
return execMapUpdate(ctx, s.db, "spans", spanID, updates)
|
|
}
|
|
|
|
func (s *PGTracingStore) GetTraceSpans(ctx context.Context, traceID uuid.UUID) ([]store.SpanData, error) {
|
|
var rows []spanRow
|
|
err := pkgSqlxDB.SelectContext(ctx, &rows,
|
|
`SELECT id, trace_id, parent_span_id, agent_id, span_type, name,
|
|
start_time, end_time, duration_ms, status, error, level,
|
|
model, provider, input_tokens, output_tokens, finish_reason,
|
|
COALESCE(model_params, '{}'::jsonb) AS model_params,
|
|
tool_name, tool_call_id, input_preview, output_preview,
|
|
COALESCE(metadata, '{}'::jsonb) AS metadata, team_id, created_at
|
|
FROM spans WHERE trace_id = $1 ORDER BY start_time`, traceID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return spanRowsToData(rows), nil
|
|
}
|
|
|
|
func (s *PGTracingStore) BatchCreateSpans(ctx context.Context, spans []store.SpanData) error {
|
|
if len(spans) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// Build multi-row INSERT
|
|
const cols = 26
|
|
valueGroups := make([]string, len(spans))
|
|
args := make([]any, 0, len(spans)*cols)
|
|
|
|
for i, span := range spans {
|
|
if span.ID == uuid.Nil {
|
|
span.ID = store.GenNewID()
|
|
spans[i].ID = span.ID
|
|
}
|
|
tenantID := span.TenantID
|
|
if tenantID == uuid.Nil {
|
|
tenantID = store.MasterTenantID
|
|
}
|
|
base := i * cols
|
|
placeholders := make([]string, cols)
|
|
for j := range cols {
|
|
placeholders[j] = fmt.Sprintf("$%d", base+j+1)
|
|
}
|
|
valueGroups[i] = "(" + strings.Join(placeholders, ", ") + ")"
|
|
|
|
args = append(args,
|
|
span.ID, span.TraceID, span.ParentSpanID, span.AgentID, span.SpanType, nilStr(span.Name),
|
|
span.StartTime, nilTime(span.EndTime), nilInt(span.DurationMS), span.Status, nilStr(span.Error), span.Level,
|
|
nilStr(span.Model), nilStr(span.Provider), nilInt(span.InputTokens), nilInt(span.OutputTokens), nilStr(span.FinishReason),
|
|
jsonOrNull(span.ModelParams), nilStr(span.ToolName), nilStr(span.ToolCallID), nilStr(span.InputPreview), nilStr(span.OutputPreview),
|
|
jsonOrNull(span.Metadata), nilUUID(span.TeamID), span.CreatedAt, tenantID,
|
|
)
|
|
}
|
|
|
|
q := `INSERT INTO spans (id, trace_id, parent_span_id, agent_id, span_type, name,
|
|
start_time, end_time, duration_ms, status, error, level,
|
|
model, provider, input_tokens, output_tokens, finish_reason,
|
|
model_params, tool_name, tool_call_id, input_preview, output_preview,
|
|
metadata, team_id, created_at, tenant_id)
|
|
VALUES ` + strings.Join(valueGroups, ", ")
|
|
|
|
_, err := s.db.ExecContext(ctx, q, args...)
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
|
|
// Batch failed — fallback to individual inserts
|
|
slog.Warn("tracing: batch insert failed, falling back to individual inserts", "count", len(spans), "error", err)
|
|
var firstErr error
|
|
for i := range spans {
|
|
if e := s.CreateSpan(ctx, &spans[i]); e != nil {
|
|
slog.Warn("tracing: individual span insert failed", "span_id", spans[i].ID, "error", e)
|
|
if firstErr == nil {
|
|
firstErr = e
|
|
}
|
|
}
|
|
}
|
|
return firstErr
|
|
}
|
|
|
|
func (s *PGTracingStore) BatchUpdateTraceAggregates(ctx context.Context, traceID uuid.UUID) error {
|
|
_, err := s.db.ExecContext(ctx, `
|
|
UPDATE traces SET
|
|
span_count = (SELECT COUNT(*) FROM spans WHERE trace_id = $1),
|
|
llm_call_count = (SELECT COUNT(*) FROM spans WHERE trace_id = $1 AND span_type = 'llm_call'),
|
|
tool_call_count = (SELECT COUNT(*) FROM spans WHERE trace_id = $1 AND span_type = 'tool_call'),
|
|
total_input_tokens = COALESCE((SELECT SUM(input_tokens) FROM spans WHERE trace_id = $1 AND span_type IN ('llm_call', 'tool_call') AND input_tokens IS NOT NULL), 0),
|
|
total_output_tokens = COALESCE((SELECT SUM(output_tokens) FROM spans WHERE trace_id = $1 AND span_type IN ('llm_call', 'tool_call') AND output_tokens IS NOT NULL), 0),
|
|
total_cost = COALESCE((SELECT SUM(total_cost) FROM spans WHERE trace_id = $1 AND total_cost IS NOT NULL), 0),
|
|
metadata = (
|
|
SELECT jsonb_build_object(
|
|
'total_cache_read_tokens', COALESCE(SUM((metadata->>'cache_read_tokens')::int), 0),
|
|
'total_cache_creation_tokens', COALESCE(SUM((metadata->>'cache_creation_tokens')::int), 0)
|
|
)
|
|
FROM spans WHERE trace_id = $1 AND span_type IN ('llm_call', 'tool_call') AND metadata IS NOT NULL
|
|
)
|
|
WHERE id = $1`, traceID)
|
|
return err
|
|
}
|
|
|
|
func (s *PGTracingStore) GetMonthlyAgentCost(ctx context.Context, agentID uuid.UUID, year int, month time.Month) (float64, error) {
|
|
start := time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
|
|
end := start.AddDate(0, 1, 0)
|
|
|
|
q := `SELECT COALESCE(SUM(total_cost), 0) FROM traces
|
|
WHERE agent_id = $1 AND created_at >= $2 AND created_at < $3 AND parent_trace_id IS NULL`
|
|
qArgs := []any{agentID, start, end}
|
|
|
|
if !store.IsCrossTenant(ctx) {
|
|
tid := store.TenantIDFromContext(ctx)
|
|
if tid != uuid.Nil {
|
|
q += " AND tenant_id = $4"
|
|
qArgs = append(qArgs, tid)
|
|
}
|
|
}
|
|
|
|
var cost float64
|
|
err := s.db.QueryRowContext(ctx, q, qArgs...).Scan(&cost)
|
|
return cost, err
|
|
}
|
|
|
|
func (s *PGTracingStore) GetCostSummary(ctx context.Context, opts store.CostSummaryOpts) ([]store.CostSummaryRow, error) {
|
|
var conditions []string
|
|
var args []any
|
|
argIdx := 1
|
|
|
|
// Only root traces (not delegations)
|
|
conditions = append(conditions, "parent_trace_id IS NULL")
|
|
|
|
if !store.IsCrossTenant(ctx) {
|
|
tenantID := store.TenantIDFromContext(ctx)
|
|
if tenantID != uuid.Nil {
|
|
conditions = append(conditions, fmt.Sprintf("tenant_id = $%d", argIdx))
|
|
args = append(args, tenantID)
|
|
argIdx++
|
|
}
|
|
}
|
|
|
|
if opts.AgentID != nil {
|
|
conditions = append(conditions, fmt.Sprintf("agent_id = $%d", argIdx))
|
|
args = append(args, *opts.AgentID)
|
|
argIdx++
|
|
}
|
|
if opts.From != nil {
|
|
conditions = append(conditions, fmt.Sprintf("created_at >= $%d", argIdx))
|
|
args = append(args, *opts.From)
|
|
argIdx++
|
|
}
|
|
if opts.To != nil {
|
|
conditions = append(conditions, fmt.Sprintf("created_at < $%d", argIdx))
|
|
args = append(args, *opts.To)
|
|
argIdx++
|
|
}
|
|
|
|
where := " WHERE " + strings.Join(conditions, " AND ")
|
|
|
|
q := `SELECT agent_id, COALESCE(SUM(total_cost), 0) AS total_cost,
|
|
COALESCE(SUM(total_input_tokens), 0) AS total_input_tokens,
|
|
COALESCE(SUM(total_output_tokens), 0) AS total_output_tokens,
|
|
COUNT(*) AS trace_count
|
|
FROM traces` + where + ` GROUP BY agent_id ORDER BY SUM(total_cost) DESC`
|
|
|
|
var result []store.CostSummaryRow
|
|
if err := pkgSqlxDB.SelectContext(ctx, &result, q, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// DeleteTracesOlderThan deletes traces and their spans older than cutoff.
|
|
// Spans are removed automatically via ON DELETE CASCADE on spans.trace_id → traces(id)
|
|
// added in migration 081. Returns total traces deleted.
|
|
func (s *PGTracingStore) DeleteTracesOlderThan(ctx context.Context, cutoff time.Time) (int64, error) {
|
|
res, err := s.db.ExecContext(ctx, `DELETE FROM traces WHERE created_at < $1`, cutoff)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("delete old traces: %w", err)
|
|
}
|
|
return res.RowsAffected()
|
|
}
|
|
|
|
// RecoverStaleRunningTraces marks traces stuck in "running" since before cutoff as "error".
|
|
// Also recovers their stuck spans. Called on startup to fix orphans from crashes.
|
|
func (s *PGTracingStore) RecoverStaleRunningTraces(ctx context.Context, cutoff time.Time) (int64, error) {
|
|
// Recover stuck spans first.
|
|
_, err := s.db.ExecContext(ctx,
|
|
`UPDATE spans SET status = 'error', error = 'recovered: server restart',
|
|
end_time = NOW(), duration_ms = EXTRACT(EPOCH FROM (NOW() - start_time))::int * 1000
|
|
WHERE status = 'running' AND start_time < $1`, cutoff)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("recover stale spans: %w", err)
|
|
}
|
|
|
|
res, err := s.db.ExecContext(ctx,
|
|
`UPDATE traces SET status = 'error',
|
|
error = 'recovered: stuck in running state (server restart)',
|
|
end_time = NOW(), duration_ms = EXTRACT(EPOCH FROM (NOW() - start_time))::int * 1000
|
|
WHERE status = 'running' AND start_time < $1`, cutoff)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("recover stale running traces: %w", err)
|
|
}
|
|
return res.RowsAffected()
|
|
}
|