mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 06:17:41 +00:00
fix(usage): repair cost analytics and display precision (#1330)
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.
This commit is contained in:
1 parent
fb4dadce8b
commit
747b58d24c
59 files changed
+3136
-197
No files matched your search
+100
-1
@@ -50,6 +50,7 @@ import (
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
|
||||
usagepricing "github.com/nextlevelbuilder/goclaw/internal/usage/pricing"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/vault"
|
||||
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
|
||||
|
||||
@@ -76,6 +77,100 @@ func gatewayLogOutput() io.Writer {
|
||||
return io.MultiWriter(os.Stdout, f)
|
||||
}
|
||||
|
||||
type traceCostBackfiller interface {
|
||||
BackfillLLMCosts(context.Context) (store.TraceCostBackfillStats, error)
|
||||
}
|
||||
|
||||
type traceUsageAggregateReconciler interface {
|
||||
ReconcileTraceUsageAggregates(context.Context) (store.TraceUsageAggregateStats, error)
|
||||
}
|
||||
|
||||
type usageEventCostBackfiller interface {
|
||||
BackfillUsageEventCosts(context.Context) (store.UsageEventCostBackfillStats, error)
|
||||
}
|
||||
|
||||
type snapshotCostBackfiller interface {
|
||||
BackfillSnapshotCosts(context.Context) (store.SnapshotCostBackfillStats, error)
|
||||
}
|
||||
|
||||
type snapshotBucketRefresher interface {
|
||||
RefreshBuckets(context.Context, []time.Time) (int, error)
|
||||
}
|
||||
|
||||
func backfillTraceCostsAfterPricingSync(ctx context.Context, stores *store.Stores, snapshots snapshotBucketRefresher) {
|
||||
if stores == nil {
|
||||
return
|
||||
}
|
||||
backfillCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
if stores.Tracing != nil {
|
||||
backfiller, ok := stores.Tracing.(traceCostBackfiller)
|
||||
if ok {
|
||||
stats, err := backfiller.BackfillLLMCosts(backfillCtx)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.trace_cost_backfill_failed", "error", err)
|
||||
} else {
|
||||
refreshedBuckets := 0
|
||||
if snapshots != nil && len(stats.SnapshotBuckets) > 0 {
|
||||
refreshedBuckets, err = snapshots.RefreshBuckets(backfillCtx, stats.SnapshotBuckets)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.trace_cost_snapshot_refresh_failed", "error", err, "buckets", len(stats.SnapshotBuckets))
|
||||
}
|
||||
}
|
||||
if stats.SpanRowsUpdated > 0 || stats.TraceRowsUpdated > 0 || refreshedBuckets > 0 {
|
||||
slog.Info("usage_pricing.trace_cost_backfill_complete",
|
||||
"spans", stats.SpanRowsUpdated,
|
||||
"traces", stats.TraceRowsUpdated,
|
||||
"snapshot_buckets", refreshedBuckets,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if stores.Tracing != nil {
|
||||
reconciler, ok := stores.Tracing.(traceUsageAggregateReconciler)
|
||||
if ok {
|
||||
stats, err := reconciler.ReconcileTraceUsageAggregates(backfillCtx)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.trace_usage_aggregate_reconcile_failed", "error", err)
|
||||
} else if stats.TraceRowsUpdated > 0 {
|
||||
slog.Info("usage_pricing.trace_usage_aggregate_reconcile_complete", "traces", stats.TraceRowsUpdated)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if stores.Snapshots != nil {
|
||||
backfiller, ok := stores.Snapshots.(snapshotCostBackfiller)
|
||||
if ok {
|
||||
stats, err := backfiller.BackfillSnapshotCosts(backfillCtx)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.snapshot_cost_backfill_failed", "error", err)
|
||||
} else if stats.SnapshotRowsUpdated > 0 {
|
||||
slog.Info("usage_pricing.snapshot_cost_backfill_complete", "snapshots", stats.SnapshotRowsUpdated)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if stores.UsageEvents != nil {
|
||||
backfiller, ok := stores.UsageEvents.(usageEventCostBackfiller)
|
||||
if ok {
|
||||
stats, err := backfiller.BackfillUsageEventCosts(backfillCtx)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.usage_event_cost_backfill_failed", "error", err)
|
||||
return
|
||||
}
|
||||
if stats.EventRowsUpdated > 0 || len(stats.RollupBuckets) > 0 {
|
||||
slog.Info("usage_pricing.usage_event_cost_backfill_complete",
|
||||
"events", stats.EventRowsUpdated,
|
||||
"rollup_buckets", len(stats.RollupBuckets),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runGateway() {
|
||||
// Setup structured logging
|
||||
logLevel := slog.LevelInfo
|
||||
@@ -525,7 +620,7 @@ func runGateway() {
|
||||
// Register all RPC methods
|
||||
server.SetLogTee(logTee)
|
||||
server.SetRuntimeLogsHandler(httpapi.NewRuntimeLogsHandler(logTee))
|
||||
pairingMethods, heartbeatMethods, chatMethods, cfgPermsMethods := registerAllMethods(server, agentRouter, pgStores.Sessions, pgStores.RunTimeline, pgStores.Cron, pgStores.Pairing, cfg, cfgPath, workspace, dataDir, msgBus, execApprovalMgr, pgStores.Agents, pgStores.Skills, pgStores.ConfigSecrets, pgStores.Teams, contextFileInterceptor, logTee, pgStores.Heartbeats, pgStores.ConfigPermissions, pgStores.SystemConfigs, pgStores.Tenants, pgStores.SkillTenantCfgs, audioMgr, usageCapSvc, providerRegistry)
|
||||
pairingMethods, heartbeatMethods, chatMethods, cfgPermsMethods := registerAllMethods(server, agentRouter, pgStores.Sessions, pgStores.Tracing, pgStores.RunTimeline, pgStores.Cron, pgStores.Pairing, cfg, cfgPath, workspace, dataDir, msgBus, execApprovalMgr, pgStores.Agents, pgStores.Skills, pgStores.ConfigSecrets, pgStores.Teams, contextFileInterceptor, logTee, pgStores.Heartbeats, pgStores.ConfigPermissions, pgStores.SystemConfigs, pgStores.Tenants, pgStores.SkillTenantCfgs, audioMgr, usageCapSvc, providerRegistry)
|
||||
|
||||
// Phase 3: Agent hooks RPC methods (hooks.list/create/update/delete/toggle/test/history).
|
||||
if hs, ok := pgStores.Hooks.(hooks.HookStore); ok && hs != nil {
|
||||
@@ -706,6 +801,10 @@ func runGateway() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
go backfillTraceCostsAfterPricingSync(ctx, pgStores, snapshotWorker)
|
||||
usagepricing.StartOpenRouterCatalogAutoSync(ctx, pgStores.UsageCaps, usagepricing.DefaultOpenRouterCatalogSyncInterval, func(syncCtx context.Context, _ int) {
|
||||
backfillTraceCostsAfterPricingSync(syncCtx, pgStores, snapshotWorker)
|
||||
})
|
||||
server.StartUpdateChecker(ctx)
|
||||
|
||||
sigCh := make(chan os.Signal, 1)
|
||||
|
||||
@@ -16,7 +16,7 @@ import (
|
||||
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
|
||||
)
|
||||
|
||||
func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, runTimeline store.RunTimelineStore, cronStore store.CronStore, pairingStore store.PairingStore, cfg *config.Config, cfgPath, workspace, dataDir string, msgBus *bus.MessageBus, execApprovalMgr *tools.ExecApprovalManager, agentStore store.AgentStore, skillStore store.SkillStore, configSecretsStore store.ConfigSecretsStore, teamStore store.TeamStore, contextFileInterceptor *tools.ContextFileInterceptor, logTee *gateway.LogTee, heartbeatStore store.HeartbeatStore, configPermStore store.ConfigPermissionStore, sysConfigStore store.SystemConfigStore, tenantStore store.TenantStore, skillTenantCfgStore store.SkillTenantConfigStore, audioMgr *audio.Manager, usageCapSvc *usagecaps.Service, providerReg *providers.Registry) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) {
|
||||
func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, tracingStore store.TracingStore, runTimeline store.RunTimelineStore, cronStore store.CronStore, pairingStore store.PairingStore, cfg *config.Config, cfgPath, workspace, dataDir string, msgBus *bus.MessageBus, execApprovalMgr *tools.ExecApprovalManager, agentStore store.AgentStore, skillStore store.SkillStore, configSecretsStore store.ConfigSecretsStore, teamStore store.TeamStore, contextFileInterceptor *tools.ContextFileInterceptor, logTee *gateway.LogTee, heartbeatStore store.HeartbeatStore, configPermStore store.ConfigPermissionStore, sysConfigStore store.SystemConfigStore, tenantStore store.TenantStore, skillTenantCfgStore store.SkillTenantConfigStore, audioMgr *audio.Manager, usageCapSvc *usagecaps.Service, providerReg *providers.Registry) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) {
|
||||
router := server.Router()
|
||||
|
||||
// Phase 1: Core methods
|
||||
@@ -65,7 +65,7 @@ func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore
|
||||
pairingMethods.Register(router)
|
||||
|
||||
// Phase 2: Usage (queries SessionStore for real token data)
|
||||
methods.NewUsageMethods(sessStore).Register(router)
|
||||
methods.NewUsageMethods(sessStore, tracingStore).Register(router)
|
||||
methods.NewLLMMethods(providerReg, cfg.Gateway.BackgroundProvider, cfg.Gateway.BackgroundModel).Register(router)
|
||||
|
||||
// Phase 2: Exec approval (always registered — returns empty when manager is nil)
|
||||
|
||||
@@ -56,16 +56,17 @@ Streaming fallback is conservative: backup models are tried only if the stream f
|
||||
|
||||
## Usage Cap Pricing Enforcement
|
||||
|
||||
Standard edition can enforce AI budget caps before billable provider dispatch. API-key providers use OpenRouter `/models` pricing as the catalog source, with optional tenant/provider/model overrides in the dashboard.
|
||||
Standard edition can enforce AI budget caps before billable provider dispatch. API-key providers use OpenRouter `/models` pricing as the catalog source, with optional tenant/provider/model overrides in the dashboard. The gateway syncs the OpenRouter catalog automatically at startup and then once per day; the dashboard sync action remains available for manual refresh.
|
||||
|
||||
Excluded provider classes:
|
||||
- `chatgpt_oauth`, `claude_cli`, and `bailian` are treated as subscription/non-API pricing in round one.
|
||||
- `chatgpt_oauth`, `claude_cli`, and `bailian` are skipped for budget-cap enforcement in round one.
|
||||
- Tracing observability can still map Bailian/DashScope Qwen model IDs to the OpenRouter catalog for market-price dashboard reporting.
|
||||
- local/no-key subprocess providers such as `acp` and `ollama` are skipped unless a future feature explicitly enables pricing for them.
|
||||
|
||||
Runtime flow:
|
||||
1. Resolve the stored provider by name and skip non-billable provider classes.
|
||||
2. Load matching policies for tenant, agent, provider, provider type, and model.
|
||||
3. Resolve custom pricing override first, then OpenRouter catalog pricing when a matching policy has a cost ceiling. Native provider model IDs are mapped to OpenRouter prefixes for common providers such as OpenAI, Anthropic, and Gemini.
|
||||
3. Resolve custom pricing override first, then OpenRouter catalog pricing when a matching policy has a cost ceiling. Native provider model IDs are mapped to OpenRouter prefixes for common providers such as OpenAI, Anthropic, and Gemini. Tracing cost calculation uses the same override/catalog resolver, with legacy `telemetry.model_pricing` kept only as a fallback.
|
||||
4. Reserve estimated tokens and cost atomically before each dispatch attempt.
|
||||
5. Reconcile reserved counters after the provider returns usage or after a failed call.
|
||||
|
||||
|
||||
@@ -121,7 +121,13 @@ The exporter lives in a separate sub-package (`internal/tracing/otelexport/`) so
|
||||
|
||||
## 5. Cost Calculation
|
||||
|
||||
Per-span cost is calculated using the `CalculateCost()` function in `internal/tracing/cost.go`. For each LLM call span:
|
||||
Per-span cost is calculated in `internal/tracing/cost.go`. For each LLM call span, pricing resolves in this order:
|
||||
|
||||
1. Tenant/provider/model override
|
||||
2. OpenRouter-backed pricing catalog
|
||||
3. Legacy `telemetry.model_pricing` config fallback
|
||||
|
||||
Catalog and override prices are stored as USD per token/unit. Legacy config prices remain USD per million tokens:
|
||||
|
||||
```
|
||||
Cost = (PromptTokens × InputCostPerMillion) / 1,000,000
|
||||
@@ -130,9 +136,11 @@ Cost = (PromptTokens × InputCostPerMillion) / 1,000,000
|
||||
+ (CacheCreationTokens × CacheCreateCostPerMillion) / 1,000,000
|
||||
```
|
||||
|
||||
Model pricing is loaded from `config.ModelPricing` and keyed by `provider/model` (with fallback to `model` only). Cost is stored in the `total_cost` field of each LLM call span. The trace aggregation sums costs from all child `llm_call` spans to compute the trace-level `total_cost`.
|
||||
Cost is stored in the `total_cost` field of each LLM call span. The trace aggregation sums costs from all child `llm_call` spans to compute the trace-level `total_cost`.
|
||||
|
||||
Cache token costs (read + create) are optional and only applied if the pricing config specifies non-zero values.
|
||||
Cache token costs (read + create) are optional and only applied when the resolved catalog, override, or fallback config provides those price dimensions.
|
||||
|
||||
At startup, after the OpenRouter catalog sync succeeds, PostgreSQL tracing runs an idempotent historical backfill for LLM spans whose `total_cost` is still zero. The backfill uses the same override/catalog precedence, refreshes trace totals, and refreshes the affected `usage_snapshots` buckets so the Usage dashboard updates without waiting for new traffic. Bailian/DashScope Qwen model IDs and common unprefixed OpenAI-compatible model IDs are mapped to OpenRouter catalog IDs for observability pricing; usage cap enforcement rules remain separate.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -25,6 +25,8 @@ Current bundled core skills:
|
||||
| `skill-creator` | Meta-skill for creating new skills |
|
||||
| `workspace-organizing` | Workspace layout and file organization guidance |
|
||||
| `goclaw` | Gateway CLI/runtime administration and troubleshooting |
|
||||
| `lark-pm` | Lark/Feishu project-management workflows through GoClaw MCP tools |
|
||||
| `lark-playbook` | Troubleshooting playbook for failed Lark/Feishu MCP operations |
|
||||
|
||||
Shared helper modules live in `skills/_shared/` and are copied alongside each skill but not registered as standalone skills.
|
||||
|
||||
@@ -92,6 +94,10 @@ skills/
|
||||
│ └── SKILL.md
|
||||
├── workspace-organizing/
|
||||
│ └── SKILL.md
|
||||
├── lark-pm/
|
||||
│ └── SKILL.md
|
||||
├── lark-playbook/
|
||||
│ └── SKILL.md
|
||||
└── goclaw/
|
||||
└── SKILL.md
|
||||
```
|
||||
|
||||
@@ -4,6 +4,122 @@ Significant changes, features, and fixes in reverse chronological order.
|
||||
|
||||
---
|
||||
|
||||
## 2026-07-03
|
||||
|
||||
### Usage cost display precision
|
||||
|
||||
**Fixes**
|
||||
|
||||
- Usage and overview cost values now render with two decimal places for API cost
|
||||
displays, matching dashboard expectations such as `$1.02` and `$5.93`.
|
||||
|
||||
**Tests**
|
||||
|
||||
- Added web formatter regression coverage for two-decimal API cost rendering.
|
||||
|
||||
### Usage analytics historical cost audit
|
||||
|
||||
**Fixes**
|
||||
|
||||
- Backfilled zero-cost historical `usage_snapshots` rows from the synced pricing
|
||||
catalog, fixing 7-day and 30-day totals that previously showed the same cost
|
||||
despite different token counts.
|
||||
- Usage analytics span aggregation now includes model-using `tool_call` spans in
|
||||
token and cost totals while keeping direct LLM call counts scoped to
|
||||
`llm_call` spans.
|
||||
- Added trace aggregate reconciliation so overview daily token/cost cards match
|
||||
the same `llm_call` + `tool_call` usage scope used by usage analytics.
|
||||
|
||||
**Tests**
|
||||
|
||||
- Added PostgreSQL and SQLite regressions for cost backfill and trace aggregate
|
||||
reconciliation with tool-call model usage.
|
||||
|
||||
### Live usage cost refresh
|
||||
|
||||
**Fixes**
|
||||
|
||||
- Usage summary and breakdown endpoints now include current-hour live trace/span
|
||||
data, so dashboard cost cards and top model costs update immediately after a
|
||||
completed LLM call instead of waiting for the hourly snapshot worker.
|
||||
|
||||
**Tests**
|
||||
|
||||
- Added HTTP regression coverage for current-hour summary cost and provider/model
|
||||
breakdown cost when no usage snapshot exists yet.
|
||||
|
||||
### Usage cost backfill after pricing sync
|
||||
|
||||
**Fixes**
|
||||
|
||||
- Added startup backfill for historical `llm_call` spans with zero cost after
|
||||
OpenRouter catalog sync succeeds.
|
||||
- Backfill now refreshes trace totals and affected usage snapshot buckets so
|
||||
dashboard cost numbers update for existing data.
|
||||
- Added catalog aliases for Bailian/DashScope Qwen models and common unprefixed
|
||||
OpenAI-compatible GPT/o-series model IDs.
|
||||
|
||||
**Tests**
|
||||
|
||||
- Added PG regression coverage for catalog-backed trace cost backfill.
|
||||
- Extended pricing resolver coverage for custom OpenAI-compatible and Bailian
|
||||
model aliases.
|
||||
|
||||
---
|
||||
|
||||
## 2026-07-02
|
||||
|
||||
### Automatic model pricing for usage costs
|
||||
|
||||
**Fixes**
|
||||
|
||||
- Tracing cost now resolves tenant/provider/model pricing overrides first, then
|
||||
the OpenRouter catalog, before falling back to legacy telemetry config.
|
||||
- OpenRouter model pricing catalog now syncs automatically at gateway startup
|
||||
and once per day, while preserving the manual dashboard sync action.
|
||||
|
||||
**Tests**
|
||||
|
||||
- Added regressions for catalog-backed trace cost precedence, config fallback,
|
||||
and cached-token price calculation.
|
||||
|
||||
### Zalo Personal nested reply context
|
||||
|
||||
**Fixes**
|
||||
|
||||
- Preserved Zalo Personal `quote` payloads on inbound messages so agents receive
|
||||
quoted-message context when users reply to a reply.
|
||||
- Rendered quoted text and quoted attachment placeholders into DM and group
|
||||
prompts before mention gating/history flush.
|
||||
- Added a channel-agnostic bounded reply-context cache so Zalo Personal can
|
||||
render multi-level reply chains without DB persistence or Zalo API backfill.
|
||||
- Moved Zalo Personal policy checks before attachment download/media extraction
|
||||
so rejected messages cannot trigger media fetches or reply-cache writes.
|
||||
|
||||
**Tests**
|
||||
|
||||
- Added Zalo Personal protocol and handler regressions for quoted text, quoted
|
||||
attachments, and nested-style replies.
|
||||
- Added reply-context cache coverage for scoped lookup, multi-level chains,
|
||||
fallback, depth/size caps, cycle guard, TTL expiry, and clear-on-stop behavior.
|
||||
- Added a Zalo Personal regression proving denied attachment messages do not
|
||||
download media before policy rejection.
|
||||
|
||||
### Lark MCP skill guidance
|
||||
|
||||
**Added**
|
||||
|
||||
- Added bundled `lark-pm` and `lark-playbook` skills so GoClaw agents handle
|
||||
Lark/Feishu PM operations through MCP only.
|
||||
- Documented schema-first LarkBase writes, writable field checks, person-field
|
||||
payload shape, per-record updates, and canonical user ID verification.
|
||||
|
||||
**Docs**
|
||||
|
||||
- Updated the core skills system list to include the new Lark MCP skills.
|
||||
|
||||
---
|
||||
|
||||
## 2026-07-01
|
||||
|
||||
### Usage event analytics schema gate
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tracing"
|
||||
)
|
||||
|
||||
func (l *Loop) calculateLLMCost(ctx context.Context, providerName, model string, usage *providers.Usage) float64 {
|
||||
if usage == nil {
|
||||
return 0
|
||||
}
|
||||
tenantID := store.TenantIDFromContext(ctx)
|
||||
if tenantID == uuid.Nil {
|
||||
tenantID = l.tenantID
|
||||
}
|
||||
if tenantID == uuid.Nil {
|
||||
tenantID = store.MasterTenantID
|
||||
}
|
||||
if l.usageCaps != nil {
|
||||
resolveCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer cancel()
|
||||
resolved, err := l.usageCaps.ResolvePricing(resolveCtx, tenantID, providerName, model)
|
||||
if err == nil && resolved != nil {
|
||||
cost, calcErr := tracing.CalculateCostFromUsagePricing(resolved.Pricing, usage)
|
||||
if calcErr == nil {
|
||||
return cost
|
||||
}
|
||||
slog.Warn("usage_pricing.trace_cost_calculation_failed", "provider", providerName, "model", model, "source", resolved.Source, "error", calcErr)
|
||||
} else if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
slog.Warn("usage_pricing.trace_pricing_resolve_failed", "provider", providerName, "model", model, "error", err)
|
||||
}
|
||||
}
|
||||
if pricing := tracing.LookupPricing(l.modelPricing, providerName, model); pricing != nil {
|
||||
return tracing.CalculateCost(pricing, usage)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/config"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
|
||||
)
|
||||
|
||||
func TestCalculateLLMCostUsesResolvedUsagePricingBeforeConfigFallback(t *testing.T) {
|
||||
tenantID := uuid.New()
|
||||
providerID := uuid.New()
|
||||
input := "0.000001"
|
||||
output := "0.000002"
|
||||
usageStore := &tracePricingUsageStore{resolved: &store.ResolvedUsagePricing{
|
||||
ModelID: "openai/gpt-4o-mini",
|
||||
Source: "catalog",
|
||||
Pricing: store.UsagePricingFields{Input: &input, Output: &output},
|
||||
}}
|
||||
providerStore := &tracePricingProviderStore{provider: &store.LLMProviderData{
|
||||
BaseModel: store.BaseModel{ID: providerID},
|
||||
TenantID: tenantID,
|
||||
Name: "openai",
|
||||
ProviderType: store.ProviderOpenAICompat,
|
||||
APIKey: "sk-test",
|
||||
Enabled: true,
|
||||
}}
|
||||
loop := &Loop{
|
||||
tenantID: tenantID,
|
||||
usageCaps: usagecaps.NewService(usageStore, providerStore),
|
||||
modelPricing: map[string]*config.ModelPricing{
|
||||
"openai/gpt-4o-mini": {InputPerMillion: 1000, OutputPerMillion: 1000},
|
||||
},
|
||||
}
|
||||
|
||||
got := loop.calculateLLMCost(context.Background(), "openai", "gpt-4o-mini", &providers.Usage{
|
||||
PromptTokens: 1000,
|
||||
CompletionTokens: 500,
|
||||
})
|
||||
if !floatClose(got, 0.002) {
|
||||
t.Fatalf("cost = %.9f, want 0.002", got)
|
||||
}
|
||||
if usageStore.resolveCalls != 1 {
|
||||
t.Fatalf("ResolvePricing calls = %d, want 1", usageStore.resolveCalls)
|
||||
}
|
||||
if usageStore.providerID != providerID {
|
||||
t.Fatalf("providerID = %s, want %s", usageStore.providerID, providerID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculateLLMCostFallsBackToConfigWhenCatalogMissing(t *testing.T) {
|
||||
tenantID := uuid.New()
|
||||
usageStore := &tracePricingUsageStore{resolveErr: sql.ErrNoRows}
|
||||
providerStore := &tracePricingProviderStore{provider: &store.LLMProviderData{
|
||||
BaseModel: store.BaseModel{ID: uuid.New()},
|
||||
TenantID: tenantID,
|
||||
Name: "openai",
|
||||
ProviderType: store.ProviderOpenAICompat,
|
||||
APIKey: "sk-test",
|
||||
Enabled: true,
|
||||
}}
|
||||
loop := &Loop{
|
||||
tenantID: tenantID,
|
||||
usageCaps: usagecaps.NewService(usageStore, providerStore),
|
||||
modelPricing: map[string]*config.ModelPricing{
|
||||
"openai/gpt-4o-mini": {InputPerMillion: 10, OutputPerMillion: 20},
|
||||
},
|
||||
}
|
||||
|
||||
got := loop.calculateLLMCost(context.Background(), "openai", "gpt-4o-mini", &providers.Usage{
|
||||
PromptTokens: 1000,
|
||||
CompletionTokens: 500,
|
||||
})
|
||||
if !floatClose(got, 0.02) {
|
||||
t.Fatalf("cost = %.9f, want 0.02", got)
|
||||
}
|
||||
}
|
||||
|
||||
func floatClose(a, b float64) bool {
|
||||
return math.Abs(a-b) < 1e-9
|
||||
}
|
||||
|
||||
type tracePricingUsageStore struct {
|
||||
resolved *store.ResolvedUsagePricing
|
||||
resolveErr error
|
||||
resolveCalls int
|
||||
tenantID uuid.UUID
|
||||
providerID uuid.UUID
|
||||
providerName string
|
||||
providerType string
|
||||
modelID string
|
||||
}
|
||||
|
||||
func (s *tracePricingUsageStore) UpsertPricingCatalog(context.Context, []store.UsagePricingCatalogEntry) (int, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ListPricingCatalog(context.Context, store.UsagePricingQuery) ([]store.UsagePricingCatalogEntry, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) PutPricingOverride(context.Context, *store.UsagePricingOverride) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ListPricingOverrides(context.Context, store.UsagePricingQuery) ([]store.UsagePricingOverride, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) DeletePricingOverride(context.Context, uuid.UUID, uuid.UUID) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ResolvePricing(_ context.Context, tenantID, providerID uuid.UUID, providerName, providerType, modelID string) (*store.ResolvedUsagePricing, error) {
|
||||
s.resolveCalls++
|
||||
s.tenantID = tenantID
|
||||
s.providerID = providerID
|
||||
s.providerName = providerName
|
||||
s.providerType = providerType
|
||||
s.modelID = modelID
|
||||
if s.resolveErr != nil {
|
||||
return nil, s.resolveErr
|
||||
}
|
||||
return s.resolved, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) CreateUsageCapPolicy(context.Context, *store.UsageCapPolicy) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ListUsageCapPolicies(context.Context, store.UsageCapScope, bool) ([]store.UsageCapPolicy, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) UpdateUsageCapPolicy(context.Context, uuid.UUID, uuid.UUID, store.UsageCapPolicyPatch) (*store.UsageCapPolicy, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) DeleteUsageCapPolicy(context.Context, uuid.UUID, uuid.UUID) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ReserveUsage(context.Context, store.UsageReserveRequest, []store.UsageCapPolicy) (*store.UsageReservationResult, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ReconcileUsage(context.Context, store.UsageReconcileRequest) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ListUsageCapUtilization(context.Context, uuid.UUID) ([]store.UsageCapUtilization, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) ListUsageCapEvents(context.Context, uuid.UUID, int) ([]store.UsageCapEvent, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingUsageStore) InsertUsageCapEvent(context.Context, *store.UsageCapEvent) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type tracePricingProviderStore struct {
|
||||
provider *store.LLMProviderData
|
||||
}
|
||||
|
||||
func (s *tracePricingProviderStore) CreateProvider(context.Context, *store.LLMProviderData) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingProviderStore) GetProvider(context.Context, uuid.UUID) (*store.LLMProviderData, error) {
|
||||
if s.provider == nil {
|
||||
return nil, sql.ErrNoRows
|
||||
}
|
||||
return s.provider, nil
|
||||
}
|
||||
func (s *tracePricingProviderStore) GetProviderByName(context.Context, string) (*store.LLMProviderData, error) {
|
||||
if s.provider == nil {
|
||||
return nil, sql.ErrNoRows
|
||||
}
|
||||
return s.provider, nil
|
||||
}
|
||||
func (s *tracePricingProviderStore) ListProviders(context.Context) ([]store.LLMProviderData, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingProviderStore) ListAllProviders(context.Context) ([]store.LLMProviderData, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (s *tracePricingProviderStore) UpdateProvider(context.Context, uuid.UUID, map[string]any) error {
|
||||
return nil
|
||||
}
|
||||
func (s *tracePricingProviderStore) DeleteProvider(context.Context, uuid.UUID) error {
|
||||
return nil
|
||||
}
|
||||
@@ -222,11 +222,12 @@ func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) {
|
||||
}
|
||||
if result != nil && result.Usage != nil {
|
||||
completedPayload["usage"] = map[string]any{
|
||||
"prompt_tokens": result.Usage.PromptTokens,
|
||||
"completion_tokens": result.Usage.CompletionTokens,
|
||||
"total_tokens": result.Usage.TotalTokens,
|
||||
"cache_creation_tokens": result.Usage.CacheCreationTokens,
|
||||
"cache_read_tokens": result.Usage.CacheReadTokens,
|
||||
"prompt_tokens": result.Usage.PromptTokens,
|
||||
"completion_tokens": result.Usage.CompletionTokens,
|
||||
"total_tokens": result.Usage.TotalTokens,
|
||||
"cache_creation_tokens": result.Usage.CacheCreationTokens,
|
||||
"cache_read_tokens": result.Usage.CacheReadTokens,
|
||||
"prompt_tokens_include_cached_segments": result.Usage.PromptTokensIncludeCachedSegments,
|
||||
}
|
||||
}
|
||||
if result != nil && len(result.Media) > 0 {
|
||||
|
||||
@@ -188,9 +188,9 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.
|
||||
if resp.Usage != nil {
|
||||
updates["input_tokens"] = resp.Usage.PromptTokens
|
||||
updates["output_tokens"] = resp.Usage.CompletionTokens
|
||||
hasMeta := resp.Usage.CacheCreationTokens > 0 || resp.Usage.CacheReadTokens > 0 || resp.Usage.ThinkingTokens > 0
|
||||
hasMeta := resp.Usage.CacheCreationTokens > 0 || resp.Usage.CacheReadTokens > 0 || resp.Usage.ThinkingTokens > 0 || resp.Usage.PromptTokensIncludeCachedSegments
|
||||
if hasMeta {
|
||||
meta := map[string]int{}
|
||||
meta := map[string]any{}
|
||||
if resp.Usage.CacheCreationTokens > 0 {
|
||||
meta["cache_creation_tokens"] = resp.Usage.CacheCreationTokens
|
||||
}
|
||||
@@ -200,17 +200,16 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.
|
||||
if resp.Usage.ThinkingTokens > 0 {
|
||||
meta["thinking_tokens"] = resp.Usage.ThinkingTokens
|
||||
}
|
||||
if resp.Usage.PromptTokensIncludeCachedSegments {
|
||||
meta["prompt_tokens_include_cached_segments"] = true
|
||||
}
|
||||
if b, err := json.Marshal(meta); err == nil {
|
||||
spanMetadata = b
|
||||
}
|
||||
}
|
||||
}
|
||||
// Calculate cost if pricing config is available.
|
||||
if pricing := tracing.LookupPricing(l.modelPricing, spanOpts.provider, spanOpts.model); pricing != nil {
|
||||
cost := tracing.CalculateCost(pricing, resp.Usage)
|
||||
if cost > 0 {
|
||||
updates["total_cost"] = cost
|
||||
}
|
||||
if cost := l.calculateLLMCost(ctx, spanOpts.provider, spanOpts.model, resp.Usage); cost > 0 {
|
||||
updates["total_cost"] = cost
|
||||
}
|
||||
updates["finish_reason"] = resp.FinishReason
|
||||
limit := previewLimitForVerbose(collector.Verbose())
|
||||
@@ -325,23 +324,29 @@ func (l *Loop) emitToolSpanEnd(ctx context.Context, spanID uuid.UUID, start time
|
||||
updates["output_tokens"] = result.Usage.CompletionTokens
|
||||
updates["provider"] = result.Provider
|
||||
updates["model"] = result.Model
|
||||
if result.Usage.CacheCreationTokens > 0 || result.Usage.CacheReadTokens > 0 {
|
||||
meta := map[string]int{
|
||||
hasMeta := result.Usage.CacheCreationTokens > 0 ||
|
||||
result.Usage.CacheReadTokens > 0 ||
|
||||
result.Usage.ThinkingTokens > 0 ||
|
||||
result.Usage.PromptTokensIncludeCachedSegments
|
||||
if hasMeta {
|
||||
meta := map[string]any{
|
||||
"cache_creation_tokens": result.Usage.CacheCreationTokens,
|
||||
"cache_read_tokens": result.Usage.CacheReadTokens,
|
||||
}
|
||||
if result.Usage.ThinkingTokens > 0 {
|
||||
meta["thinking_tokens"] = result.Usage.ThinkingTokens
|
||||
}
|
||||
if result.Usage.PromptTokensIncludeCachedSegments {
|
||||
meta["prompt_tokens_include_cached_segments"] = true
|
||||
}
|
||||
if b, err := json.Marshal(meta); err == nil {
|
||||
updates["metadata"] = b
|
||||
}
|
||||
}
|
||||
// Calculate cost for tool's internal LLM calls.
|
||||
provider := result.Provider
|
||||
model := result.Model
|
||||
if pricing := tracing.LookupPricing(l.modelPricing, provider, model); pricing != nil {
|
||||
cost := tracing.CalculateCost(pricing, result.Usage)
|
||||
if cost > 0 {
|
||||
updates["total_cost"] = cost
|
||||
}
|
||||
if cost := l.calculateLLMCost(ctx, provider, model, result.Usage); cost > 0 {
|
||||
updates["total_cost"] = cost
|
||||
}
|
||||
}
|
||||
|
||||
@@ -422,7 +427,7 @@ func (l *Loop) emitAgentSpanEnd(ctx context.Context, agentSpanID uuid.UUID, star
|
||||
limit := previewLimitForVerbose(collector.Verbose())
|
||||
updates["output_preview"] = tracing.TruncateMid(result.Content, limit)
|
||||
// Note: token counts are NOT set on agent spans to avoid double-counting
|
||||
// with child llm_call spans. Trace aggregation sums only llm_call spans.
|
||||
// with child spans that directly report model usage.
|
||||
}
|
||||
|
||||
collector.EmitSpanUpdate(agentSpanID, traceID, updates)
|
||||
|
||||
@@ -77,6 +77,8 @@ func (l *Loop) recordToolUsageEvent(ctx context.Context, req *RunRequest, canoni
|
||||
event := l.baseUsageEvent(ctx, req, start, eventType, resourceType, resourceName, resourceID, source)
|
||||
event.SpanID = uuidPtr(spanID)
|
||||
event.Status = "completed"
|
||||
event.Provider = result.Provider
|
||||
event.Model = result.Model
|
||||
if result.IsError {
|
||||
event.Status = "error"
|
||||
event.ErrorCount = 1
|
||||
@@ -92,9 +94,11 @@ func (l *Loop) recordToolUsageEvent(ctx context.Context, req *RunRequest, canoni
|
||||
if event.TotalTokens == 0 {
|
||||
event.TotalTokens = event.InputTokens + event.OutputTokens
|
||||
}
|
||||
if result.Usage.PromptTokensIncludeCachedSegments {
|
||||
metadata["prompt_tokens_include_cached_segments"] = true
|
||||
}
|
||||
event.CostUSD = l.calculateLLMCost(ctx, event.Provider, event.Model, result.Usage)
|
||||
}
|
||||
event.Provider = result.Provider
|
||||
event.Model = result.Model
|
||||
event.Metadata = usageMetadata(metadata)
|
||||
l.insertUsageEventBestEffort(ctx, event)
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tracing"
|
||||
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
|
||||
)
|
||||
|
||||
type fakeUsageEventStore struct {
|
||||
@@ -146,6 +147,50 @@ func TestRecordToolUsageEvent_PreservesProviderCacheUsage(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordToolUsageEvent_ComputesCost(t *testing.T) {
|
||||
tenantID := uuid.New()
|
||||
providerID := uuid.New()
|
||||
storeSpy := newFakeUsageEventStore()
|
||||
input := "0.000001"
|
||||
output := "0.000002"
|
||||
usageStore := &tracePricingUsageStore{resolved: &store.ResolvedUsagePricing{
|
||||
ModelID: "openai/gpt-4o-mini",
|
||||
Source: "catalog",
|
||||
Pricing: store.UsagePricingFields{Input: &input, Output: &output},
|
||||
}}
|
||||
providerStore := &tracePricingProviderStore{provider: &store.LLMProviderData{
|
||||
BaseModel: store.BaseModel{ID: providerID},
|
||||
TenantID: tenantID,
|
||||
Name: "openai",
|
||||
ProviderType: store.ProviderOpenAICompat,
|
||||
APIKey: "sk-test",
|
||||
Enabled: true,
|
||||
}}
|
||||
loop := &Loop{
|
||||
usageEvents: storeSpy,
|
||||
agentUUID: uuid.New(),
|
||||
tenantID: tenantID,
|
||||
usageCaps: usagecaps.NewService(usageStore, providerStore),
|
||||
}
|
||||
ctx := tracing.WithTraceID(store.WithTenantID(t.Context(), tenantID), uuid.New())
|
||||
result := tools.NewResult("ok")
|
||||
result.Provider = "openai"
|
||||
result.Model = "gpt-4o-mini"
|
||||
result.Usage = &providers.Usage{
|
||||
PromptTokens: 1000,
|
||||
CompletionTokens: 500,
|
||||
TotalTokens: 1500,
|
||||
}
|
||||
|
||||
loop.recordToolUsageEvent(ctx, &RunRequest{RunID: "run-1", SessionKey: "session-1"}, "read_image", "read_image", "call-1",
|
||||
nil, time.Now().Add(-time.Millisecond), result, uuid.New())
|
||||
|
||||
event := waitUsageEvent(t, storeSpy)
|
||||
if !floatClose(event.CostUSD, 0.002) {
|
||||
t.Fatalf("cost = %.9f, want 0.002", event.CostUSD)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordSkillSlashUsageEvent_RequiresTraceContext(t *testing.T) {
|
||||
storeSpy := newFakeUsageEventStore()
|
||||
loop := &Loop{usageEvents: storeSpy, agentUUID: uuid.New(), tenantID: uuid.New()}
|
||||
|
||||
@@ -13,23 +13,35 @@ import (
|
||||
// UsageMethods handles usage.get, usage.summary.
|
||||
// Queries SessionStore for real token data (accumulated via AccumulateTokens in agent loop).
|
||||
type UsageMethods struct {
|
||||
sessions store.SessionStore
|
||||
sessions store.SessionStore
|
||||
sessionCosts sessionCostReader
|
||||
}
|
||||
|
||||
type sessionCostReader interface {
|
||||
GetSessionCosts(ctx context.Context, sessionKeys []string) (map[string]float64, error)
|
||||
}
|
||||
|
||||
// UsageRecord is a single usage entry derived from session data.
|
||||
type UsageRecord struct {
|
||||
AgentID string `json:"agentId"`
|
||||
SessionKey string `json:"sessionKey"`
|
||||
Model string `json:"model"`
|
||||
Provider string `json:"provider"`
|
||||
InputTokens int64 `json:"inputTokens"`
|
||||
OutputTokens int64 `json:"outputTokens"`
|
||||
TotalTokens int64 `json:"totalTokens"`
|
||||
Timestamp int64 `json:"timestamp"`
|
||||
AgentID string `json:"agentId"`
|
||||
SessionKey string `json:"sessionKey"`
|
||||
Model string `json:"model"`
|
||||
Provider string `json:"provider"`
|
||||
InputTokens int64 `json:"inputTokens"`
|
||||
OutputTokens int64 `json:"outputTokens"`
|
||||
TotalTokens int64 `json:"totalTokens"`
|
||||
Cost float64 `json:"cost"`
|
||||
Timestamp int64 `json:"timestamp"`
|
||||
}
|
||||
|
||||
func NewUsageMethods(sessStore store.SessionStore) *UsageMethods {
|
||||
return &UsageMethods{sessions: sessStore}
|
||||
func NewUsageMethods(sessStore store.SessionStore, tracingStores ...store.TracingStore) *UsageMethods {
|
||||
m := &UsageMethods{sessions: sessStore}
|
||||
if len(tracingStores) > 0 && tracingStores[0] != nil {
|
||||
if costs, ok := tracingStores[0].(sessionCostReader); ok {
|
||||
m.sessionCosts = costs
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *UsageMethods) Register(router *gateway.MethodRouter) {
|
||||
@@ -86,6 +98,17 @@ func (m *UsageMethods) handleGet(ctx context.Context, client *gateway.Client, re
|
||||
offset := min(params.Offset, total)
|
||||
end := min(offset+params.Limit, total)
|
||||
records = records[offset:end]
|
||||
if m.sessionCosts != nil && len(records) > 0 {
|
||||
sessionKeys := make([]string, 0, len(records))
|
||||
for _, record := range records {
|
||||
sessionKeys = append(sessionKeys, record.SessionKey)
|
||||
}
|
||||
if costs, err := m.sessionCosts.GetSessionCosts(ctx, sessionKeys); err == nil {
|
||||
for i := range records {
|
||||
records[i].Cost = costs[records[i].SessionKey]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{
|
||||
"records": records,
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
package methods
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/gateway"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/permissions"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
|
||||
)
|
||||
|
||||
type usageSessionStore struct {
|
||||
store.SessionStore
|
||||
sessions []store.SessionInfoRich
|
||||
opts []store.SessionListOpts
|
||||
}
|
||||
|
||||
func (s *usageSessionStore) ListPagedRich(_ context.Context, opts store.SessionListOpts) store.SessionListRichResult {
|
||||
s.opts = append(s.opts, opts)
|
||||
return store.SessionListRichResult{Sessions: s.sessions, Total: len(s.sessions)}
|
||||
}
|
||||
|
||||
type usageTracingStore struct {
|
||||
store.TracingStore
|
||||
costs map[string]float64
|
||||
keys []string
|
||||
}
|
||||
|
||||
func (s *usageTracingStore) GetSessionCosts(_ context.Context, sessionKeys []string) (map[string]float64, error) {
|
||||
s.keys = append(s.keys[:0], sessionKeys...)
|
||||
return s.costs, nil
|
||||
}
|
||||
|
||||
func TestUsageGetIncludesTraceCost(t *testing.T) {
|
||||
now := time.Date(2026, 7, 2, 12, 0, 0, 0, time.UTC)
|
||||
sessions := &usageSessionStore{
|
||||
sessions: []store.SessionInfoRich{
|
||||
{
|
||||
SessionInfo: store.SessionInfo{Key: "agent:cppai-pm:direct:user-a", Updated: now},
|
||||
Model: "qwen3.7-plus",
|
||||
Provider: "bailian",
|
||||
InputTokens: 1_200,
|
||||
OutputTokens: 300,
|
||||
},
|
||||
{
|
||||
SessionInfo: store.SessionInfo{Key: "agent:itsddvn:direct:user-b", Updated: now.Add(-time.Minute)},
|
||||
Model: "gpt-5.5",
|
||||
Provider: "openai",
|
||||
InputTokens: 500,
|
||||
OutputTokens: 100,
|
||||
},
|
||||
},
|
||||
}
|
||||
tracing := &usageTracingStore{costs: map[string]float64{
|
||||
"agent:cppai-pm:direct:user-a": 0.0123,
|
||||
"agent:itsddvn:direct:user-b": 0.0045,
|
||||
}}
|
||||
methods := NewUsageMethods(sessions, tracing)
|
||||
tenantID := uuid.Must(uuid.NewV7())
|
||||
client, responses := gateway.NewCapturingTestClient(permissions.RoleViewer, tenantID, "caller", 1)
|
||||
ctx := store.WithTenantID(context.Background(), tenantID)
|
||||
|
||||
methods.handleGet(ctx, client, sessionReqFrame(t, protocol.MethodUsageGet, map[string]any{"limit": 1}))
|
||||
|
||||
resp := readUsageResponse(t, responses)
|
||||
if resp.Error != nil {
|
||||
t.Fatalf("unexpected error: %+v", resp.Error)
|
||||
}
|
||||
payload, ok := resp.Payload.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("payload type = %T", resp.Payload)
|
||||
}
|
||||
records, ok := payload["records"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("records type = %T", payload["records"])
|
||||
}
|
||||
if len(records) != 1 {
|
||||
t.Fatalf("len(records) = %d, want 1", len(records))
|
||||
}
|
||||
record, ok := records[0].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("record type = %T", records[0])
|
||||
}
|
||||
if got, want := record["cost"], 0.0123; got != want {
|
||||
t.Fatalf("cost = %#v, want %#v", got, want)
|
||||
}
|
||||
if want := []string{"agent:cppai-pm:direct:user-a"}; !reflect.DeepEqual(tracing.keys, want) {
|
||||
t.Fatalf("cost keys = %#v, want %#v", tracing.keys, want)
|
||||
}
|
||||
}
|
||||
|
||||
func readUsageResponse(t *testing.T, ch <-chan []byte) protocol.ResponseFrame {
|
||||
t.Helper()
|
||||
raw := <-ch
|
||||
var resp protocol.ResponseFrame
|
||||
if err := json.Unmarshal(raw, &resp); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
return resp
|
||||
}
|
||||
@@ -10533,7 +10533,9 @@
|
||||
"enum": [
|
||||
"agent",
|
||||
"provider",
|
||||
"model"
|
||||
"model",
|
||||
"provider_model",
|
||||
"channel"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
+16
-81
@@ -54,13 +54,12 @@ func (h *UsageHandler) handleTimeSeries(w http.ResponseWriter, r *http.Request)
|
||||
return
|
||||
}
|
||||
|
||||
// Gap-fill: if "to" extends into current incomplete hour, query live traces
|
||||
now := time.Now().UTC()
|
||||
currentHourStart := now.Truncate(time.Hour)
|
||||
if q.To.After(currentHourStart) && currentHourStart.After(q.From) {
|
||||
livePoint := h.queryLiveHour(r, currentHourStart, now, q)
|
||||
if livePoint != nil {
|
||||
points = append(points, *livePoint)
|
||||
if from, to, ok := liveUsageWindow(q, time.Now().UTC()); ok {
|
||||
livePoint, err := h.queryLiveTimeSeries(r, from, to, q)
|
||||
if err != nil {
|
||||
slog.Warn("usage.timeseries live query failed", "error", err)
|
||||
} else if livePoint != nil {
|
||||
points = mergeSnapshotTimeSeries(points, *livePoint, q.GroupBy)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -83,6 +82,14 @@ func (h *UsageHandler) handleBreakdown(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "internal server error"})
|
||||
return
|
||||
}
|
||||
if from, to, ok := liveUsageWindow(q, time.Now().UTC()); ok {
|
||||
liveRows, err := h.queryLiveBreakdown(r, from, to, q)
|
||||
if err != nil {
|
||||
slog.Warn("usage.breakdown live query failed", "error", err)
|
||||
} else {
|
||||
rows = mergeSnapshotBreakdowns(rows, liveRows, q.GroupBy)
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"rows": rows})
|
||||
}
|
||||
|
||||
@@ -124,7 +131,7 @@ func (h *UsageHandler) handleSummary(w http.ResponseWriter, r *http.Request) {
|
||||
previousQ.To = currentFrom
|
||||
previousQ.GroupBy = "hour"
|
||||
|
||||
currentSummary := h.aggregateTimeSeries(r, currentQ)
|
||||
currentSummary := h.aggregateTimeSeriesWithLive(r, currentQ, now)
|
||||
previousSummary := h.aggregateTimeSeries(r, previousQ)
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
@@ -227,79 +234,7 @@ func (h *UsageHandler) aggregateTimeSeries(r *http.Request, q store.SnapshotQuer
|
||||
return usageSummary{}
|
||||
}
|
||||
|
||||
var s usageSummary
|
||||
var totalWeightedDuration int64
|
||||
for _, p := range points {
|
||||
s.Requests += p.RequestCount
|
||||
s.InputTokens += p.InputTokens
|
||||
s.OutputTokens += p.OutputTokens
|
||||
s.Cost += p.TotalCost
|
||||
s.UniqueUsers += p.UniqueUsers
|
||||
s.Errors += p.ErrorCount
|
||||
s.LLMCalls += p.LLMCallCount
|
||||
s.ToolCalls += p.ToolCallCount
|
||||
totalWeightedDuration += int64(p.AvgDurationMS) * int64(p.RequestCount)
|
||||
}
|
||||
if s.Requests > 0 {
|
||||
s.AvgDurationMS = int(totalWeightedDuration / int64(s.Requests))
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// queryLiveHour runs a lightweight query on traces for the current incomplete hour.
|
||||
func (h *UsageHandler) queryLiveHour(r *http.Request, from, to time.Time, q store.SnapshotQuery) *store.SnapshotTimeSeries {
|
||||
query := `SELECT
|
||||
COUNT(*),
|
||||
COUNT(*) FILTER (WHERE status = 'error'),
|
||||
COUNT(DISTINCT user_id),
|
||||
COALESCE(SUM(total_input_tokens), 0),
|
||||
COALESCE(SUM(total_output_tokens), 0),
|
||||
COALESCE(SUM(total_cost), 0),
|
||||
COALESCE(SUM(llm_call_count), 0),
|
||||
COALESCE(SUM(tool_call_count), 0),
|
||||
COALESCE(AVG(duration_ms), 0)::INTEGER
|
||||
FROM traces
|
||||
WHERE start_time >= $1 AND start_time < $2
|
||||
AND parent_trace_id IS NULL`
|
||||
|
||||
args := []any{from, to}
|
||||
idx := 3
|
||||
|
||||
// Tenant isolation: scope to caller's tenant
|
||||
if !store.IsOwnerRole(r.Context()) {
|
||||
tid := store.TenantIDFromContext(r.Context())
|
||||
if tid != uuid.Nil {
|
||||
query += fmt.Sprintf(" AND tenant_id = $%d", idx)
|
||||
args = append(args, tid)
|
||||
idx++
|
||||
}
|
||||
}
|
||||
|
||||
if q.AgentID != nil {
|
||||
query += fmt.Sprintf(" AND agent_id = $%d", idx)
|
||||
args = append(args, *q.AgentID)
|
||||
idx++
|
||||
}
|
||||
if q.Channel != "" {
|
||||
query += fmt.Sprintf(" AND channel = $%d", idx)
|
||||
args = append(args, q.Channel)
|
||||
idx++
|
||||
}
|
||||
// Note: provider/model filters are not applied here because the traces table
|
||||
// does not have provider/model columns (those live on spans). The live gap-fill
|
||||
// is a rough approximation for the current incomplete hour.
|
||||
|
||||
var p store.SnapshotTimeSeries
|
||||
p.BucketTime = from
|
||||
err := h.db.QueryRowContext(r.Context(), query, args...).Scan(
|
||||
&p.RequestCount, &p.ErrorCount, &p.UniqueUsers,
|
||||
&p.InputTokens, &p.OutputTokens, &p.TotalCost,
|
||||
&p.LLMCallCount, &p.ToolCallCount, &p.AvgDurationMS,
|
||||
)
|
||||
if err != nil || p.RequestCount == 0 {
|
||||
return nil
|
||||
}
|
||||
return &p
|
||||
return summarizeTimeSeries(points)
|
||||
}
|
||||
|
||||
func parseSnapshotFilters(r *http.Request) store.SnapshotQuery {
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func liveUsageWindow(q store.SnapshotQuery, now time.Time) (time.Time, time.Time, bool) {
|
||||
now = now.UTC()
|
||||
from := now.Truncate(time.Hour)
|
||||
if q.From.After(from) {
|
||||
from = q.From.UTC()
|
||||
}
|
||||
to := now
|
||||
if !q.To.IsZero() && q.To.Before(to) {
|
||||
to = q.To.UTC()
|
||||
}
|
||||
if !to.After(from) {
|
||||
return time.Time{}, time.Time{}, false
|
||||
}
|
||||
return from, to, true
|
||||
}
|
||||
|
||||
func liveUsageBucket(from time.Time, groupBy string) time.Time {
|
||||
from = from.UTC()
|
||||
if groupBy == "day" {
|
||||
return time.Date(from.Year(), from.Month(), from.Day(), 0, 0, 0, 0, time.UTC)
|
||||
}
|
||||
return from.Truncate(time.Hour)
|
||||
}
|
||||
|
||||
func (h *UsageHandler) queryLiveTimeSeries(r *http.Request, from, to time.Time, q store.SnapshotQuery) (*store.SnapshotTimeSeries, error) {
|
||||
if h.db == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
point := store.SnapshotTimeSeries{BucketTime: liveUsageBucket(from, q.GroupBy)}
|
||||
if q.Provider == "" && q.Model == "" {
|
||||
if err := h.queryLiveTraceTotals(r, from, to, q, &point); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if err := h.queryLiveSpanTotals(r, from, to, q, &point); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !hasTimeSeriesUsage(point) {
|
||||
return nil, nil
|
||||
}
|
||||
return &point, nil
|
||||
}
|
||||
|
||||
func (h *UsageHandler) queryLiveTraceTotals(r *http.Request, from, to time.Time, q store.SnapshotQuery, point *store.SnapshotTimeSeries) error {
|
||||
var query strings.Builder
|
||||
query.WriteString(`SELECT
|
||||
COUNT(*),
|
||||
COALESCE(SUM(CASE WHEN t.status = 'error' THEN 1 ELSE 0 END), 0),
|
||||
COUNT(DISTINCT t.user_id),
|
||||
COALESCE(SUM(t.tool_call_count), 0),
|
||||
CAST(COALESCE(AVG(t.duration_ms), 0) AS INTEGER)
|
||||
FROM traces t
|
||||
WHERE t.start_time >= $1 AND t.start_time < $2
|
||||
AND t.parent_trace_id IS NULL`)
|
||||
|
||||
args := []any{from, to}
|
||||
idx := 3
|
||||
appendLiveTraceFilters(&query, &args, &idx, r, q)
|
||||
|
||||
return h.db.QueryRowContext(r.Context(), query.String(), args...).Scan(
|
||||
&point.RequestCount,
|
||||
&point.ErrorCount,
|
||||
&point.UniqueUsers,
|
||||
&point.ToolCallCount,
|
||||
&point.AvgDurationMS,
|
||||
)
|
||||
}
|
||||
|
||||
func (h *UsageHandler) queryLiveSpanTotals(r *http.Request, from, to time.Time, q store.SnapshotQuery, point *store.SnapshotTimeSeries) error {
|
||||
var query strings.Builder
|
||||
query.WriteString(`SELECT
|
||||
COALESCE(SUM(COALESCE(s.input_tokens, 0)), 0),
|
||||
COALESCE(SUM(COALESCE(s.output_tokens, 0)), 0),
|
||||
COALESCE(SUM(COALESCE(s.total_cost, 0)), 0),
|
||||
COUNT(*) FILTER (WHERE s.span_type = 'llm_call'),
|
||||
COALESCE(SUM(CAST(COALESCE(NULLIF(s.metadata->>'cache_read_tokens', ''), '0') AS INTEGER)), 0),
|
||||
COALESCE(SUM(CAST(COALESCE(NULLIF(s.metadata->>'cache_creation_tokens', ''), '0') AS INTEGER)), 0),
|
||||
COALESCE(SUM(CAST(COALESCE(NULLIF(s.metadata->>'thinking_tokens', ''), '0') AS INTEGER)), 0)
|
||||
FROM traces t
|
||||
JOIN spans s ON s.trace_id = t.id AND s.span_type IN ('llm_call', 'tool_call')
|
||||
WHERE t.start_time >= $1 AND t.start_time < $2
|
||||
AND t.parent_trace_id IS NULL`)
|
||||
|
||||
args := []any{from, to}
|
||||
idx := 3
|
||||
appendLiveTraceFilters(&query, &args, &idx, r, q)
|
||||
appendLiveSpanFilters(&query, &args, &idx, q)
|
||||
|
||||
return h.db.QueryRowContext(r.Context(), query.String(), args...).Scan(
|
||||
&point.InputTokens,
|
||||
&point.OutputTokens,
|
||||
&point.TotalCost,
|
||||
&point.LLMCallCount,
|
||||
&point.CacheReadTokens,
|
||||
&point.CacheCreateTokens,
|
||||
&point.ThinkingTokens,
|
||||
)
|
||||
}
|
||||
|
||||
func appendLiveTraceFilters(query *strings.Builder, args *[]any, idx *int, r *http.Request, q store.SnapshotQuery) {
|
||||
if !store.IsCrossTenant(r.Context()) {
|
||||
if tenantID := store.TenantIDFromContext(r.Context()); tenantID != uuid.Nil {
|
||||
fmt.Fprintf(query, " AND t.tenant_id = $%d", *idx)
|
||||
*args = append(*args, tenantID)
|
||||
*idx += 1
|
||||
}
|
||||
}
|
||||
if q.AgentID != nil {
|
||||
fmt.Fprintf(query, " AND t.agent_id = $%d", *idx)
|
||||
*args = append(*args, *q.AgentID)
|
||||
*idx += 1
|
||||
}
|
||||
if q.Channel != "" {
|
||||
fmt.Fprintf(query, " AND t.channel = $%d", *idx)
|
||||
*args = append(*args, q.Channel)
|
||||
*idx += 1
|
||||
}
|
||||
}
|
||||
|
||||
func appendLiveSpanFilters(query *strings.Builder, args *[]any, idx *int, q store.SnapshotQuery) {
|
||||
if q.Provider != "" {
|
||||
fmt.Fprintf(query, " AND s.provider = $%d", *idx)
|
||||
*args = append(*args, q.Provider)
|
||||
*idx += 1
|
||||
}
|
||||
if q.Model != "" {
|
||||
fmt.Fprintf(query, " AND s.model = $%d", *idx)
|
||||
*args = append(*args, q.Model)
|
||||
*idx += 1
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func (h *UsageHandler) queryLiveBreakdown(r *http.Request, from, to time.Time, q store.SnapshotQuery) ([]store.SnapshotBreakdown, error) {
|
||||
if h.db == nil {
|
||||
return nil, nil
|
||||
}
|
||||
groupBy := q.GroupBy
|
||||
if groupBy == "" {
|
||||
groupBy = "provider"
|
||||
}
|
||||
|
||||
switch groupBy {
|
||||
case "channel", "agent":
|
||||
spanRows, err := h.queryLiveSpanBreakdown(r, from, to, q, groupBy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if q.Provider != "" || q.Model != "" {
|
||||
return spanRows, nil
|
||||
}
|
||||
traceRows, err := h.queryLiveTraceBreakdown(r, from, to, q, groupBy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return mergeSnapshotBreakdowns(traceRows, spanRows, groupBy), nil
|
||||
default:
|
||||
return h.queryLiveSpanBreakdown(r, from, to, q, groupBy)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *UsageHandler) queryLiveSpanBreakdown(r *http.Request, from, to time.Time, q store.SnapshotQuery, groupBy string) ([]store.SnapshotBreakdown, error) {
|
||||
groupExpr := "COALESCE(s.provider, '')"
|
||||
extraFilter := " AND COALESCE(s.provider, '') != '' AND COALESCE(s.model, '') != ''"
|
||||
switch groupBy {
|
||||
case "model":
|
||||
groupExpr = "COALESCE(s.model, '')"
|
||||
case "provider_model":
|
||||
groupExpr = "COALESCE(s.provider, '') || '/' || COALESCE(s.model, '')"
|
||||
case "channel":
|
||||
groupExpr = "COALESCE(t.channel, '')"
|
||||
extraFilter = " AND COALESCE(t.channel, '') != ''"
|
||||
case "agent":
|
||||
groupExpr = "COALESCE(CAST(t.agent_id AS TEXT), '')"
|
||||
extraFilter = ""
|
||||
}
|
||||
|
||||
var query strings.Builder
|
||||
fmt.Fprintf(&query, `SELECT
|
||||
%s AS key,
|
||||
COALESCE(SUM(COALESCE(s.input_tokens, 0)), 0),
|
||||
COALESCE(SUM(COALESCE(s.output_tokens, 0)), 0),
|
||||
COALESCE(SUM(CAST(COALESCE(NULLIF(s.metadata->>'cache_read_tokens', ''), '0') AS INTEGER)), 0),
|
||||
COALESCE(SUM(CAST(COALESCE(NULLIF(s.metadata->>'cache_creation_tokens', ''), '0') AS INTEGER)), 0),
|
||||
COALESCE(SUM(COALESCE(s.total_cost, 0)), 0),
|
||||
0,
|
||||
COUNT(*) FILTER (WHERE s.span_type = 'llm_call'),
|
||||
0,
|
||||
0,
|
||||
0
|
||||
FROM traces t
|
||||
JOIN spans s ON s.trace_id = t.id AND s.span_type IN ('llm_call', 'tool_call')
|
||||
WHERE t.start_time >= $1 AND t.start_time < $2
|
||||
AND t.parent_trace_id IS NULL`, groupExpr)
|
||||
query.WriteString(extraFilter)
|
||||
|
||||
args := []any{from, to}
|
||||
idx := 3
|
||||
appendLiveTraceFilters(&query, &args, &idx, r, q)
|
||||
appendLiveSpanFilters(&query, &args, &idx, q)
|
||||
fmt.Fprintf(&query, " GROUP BY %s", groupExpr)
|
||||
|
||||
return scanLiveBreakdownRows(r, h, query.String(), args...)
|
||||
}
|
||||
|
||||
func (h *UsageHandler) queryLiveTraceBreakdown(r *http.Request, from, to time.Time, q store.SnapshotQuery, groupBy string) ([]store.SnapshotBreakdown, error) {
|
||||
groupExpr := "COALESCE(t.channel, '')"
|
||||
extraFilter := " AND COALESCE(t.channel, '') != ''"
|
||||
if groupBy == "agent" {
|
||||
groupExpr = "COALESCE(CAST(t.agent_id AS TEXT), '')"
|
||||
extraFilter = ""
|
||||
}
|
||||
|
||||
var query strings.Builder
|
||||
fmt.Fprintf(&query, `SELECT
|
||||
%s AS key,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
COUNT(*),
|
||||
0,
|
||||
COALESCE(SUM(t.tool_call_count), 0),
|
||||
COALESCE(SUM(CASE WHEN t.status = 'error' THEN 1 ELSE 0 END), 0),
|
||||
CAST(COALESCE(AVG(t.duration_ms), 0) AS INTEGER)
|
||||
FROM traces t
|
||||
WHERE t.start_time >= $1 AND t.start_time < $2
|
||||
AND t.parent_trace_id IS NULL`, groupExpr)
|
||||
query.WriteString(extraFilter)
|
||||
|
||||
args := []any{from, to}
|
||||
idx := 3
|
||||
appendLiveTraceFilters(&query, &args, &idx, r, q)
|
||||
fmt.Fprintf(&query, " GROUP BY %s", groupExpr)
|
||||
|
||||
return scanLiveBreakdownRows(r, h, query.String(), args...)
|
||||
}
|
||||
|
||||
func scanLiveBreakdownRows(r *http.Request, h *UsageHandler, query string, args ...any) ([]store.SnapshotBreakdown, error) {
|
||||
rows, err := h.db.QueryContext(r.Context(), query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var result []store.SnapshotBreakdown
|
||||
for rows.Next() {
|
||||
var b store.SnapshotBreakdown
|
||||
if err := rows.Scan(
|
||||
&b.Key,
|
||||
&b.InputTokens,
|
||||
&b.OutputTokens,
|
||||
&b.CacheReadTokens,
|
||||
&b.CacheCreateTokens,
|
||||
&b.TotalCost,
|
||||
&b.RequestCount,
|
||||
&b.LLMCallCount,
|
||||
&b.ToolCallCount,
|
||||
&b.ErrorCount,
|
||||
&b.AvgDurationMS,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result = append(result, b)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func (h *UsageHandler) aggregateTimeSeriesWithLive(r *http.Request, q store.SnapshotQuery, now time.Time) usageSummary {
|
||||
points, err := h.snapshots.GetTimeSeries(r.Context(), q)
|
||||
if err != nil {
|
||||
return usageSummary{}
|
||||
}
|
||||
if from, to, ok := liveUsageWindow(q, now); ok {
|
||||
livePoint, err := h.queryLiveTimeSeries(r, from, to, q)
|
||||
if err != nil {
|
||||
return summarizeTimeSeries(points)
|
||||
}
|
||||
if livePoint != nil {
|
||||
points = mergeSnapshotTimeSeries(points, *livePoint, q.GroupBy)
|
||||
}
|
||||
}
|
||||
return summarizeTimeSeries(points)
|
||||
}
|
||||
|
||||
func summarizeTimeSeries(points []store.SnapshotTimeSeries) usageSummary {
|
||||
var s usageSummary
|
||||
var totalWeightedDuration int64
|
||||
for _, p := range points {
|
||||
s.Requests += p.RequestCount
|
||||
s.InputTokens += p.InputTokens
|
||||
s.OutputTokens += p.OutputTokens
|
||||
s.Cost += p.TotalCost
|
||||
s.UniqueUsers += p.UniqueUsers
|
||||
s.Errors += p.ErrorCount
|
||||
s.LLMCalls += p.LLMCallCount
|
||||
s.ToolCalls += p.ToolCallCount
|
||||
totalWeightedDuration += int64(p.AvgDurationMS) * int64(p.RequestCount)
|
||||
}
|
||||
if s.Requests > 0 {
|
||||
s.AvgDurationMS = int(totalWeightedDuration / int64(s.Requests))
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func mergeSnapshotTimeSeries(points []store.SnapshotTimeSeries, live store.SnapshotTimeSeries, groupBy string) []store.SnapshotTimeSeries {
|
||||
for i := range points {
|
||||
if points[i].BucketTime.Equal(live.BucketTime) {
|
||||
if groupBy == "day" {
|
||||
addPointToTimeSeries(&points[i], live)
|
||||
} else {
|
||||
points[i] = live
|
||||
}
|
||||
return points
|
||||
}
|
||||
}
|
||||
points = append(points, live)
|
||||
sort.Slice(points, func(i, j int) bool {
|
||||
return points[i].BucketTime.Before(points[j].BucketTime)
|
||||
})
|
||||
return points
|
||||
}
|
||||
|
||||
func addPointToTimeSeries(dst *store.SnapshotTimeSeries, src store.SnapshotTimeSeries) {
|
||||
oldRequests := dst.RequestCount
|
||||
weightedDuration := int64(dst.AvgDurationMS)*int64(oldRequests) + int64(src.AvgDurationMS)*int64(src.RequestCount)
|
||||
dst.InputTokens += src.InputTokens
|
||||
dst.OutputTokens += src.OutputTokens
|
||||
dst.CacheReadTokens += src.CacheReadTokens
|
||||
dst.CacheCreateTokens += src.CacheCreateTokens
|
||||
dst.ThinkingTokens += src.ThinkingTokens
|
||||
dst.TotalCost += src.TotalCost
|
||||
dst.RequestCount += src.RequestCount
|
||||
dst.LLMCallCount += src.LLMCallCount
|
||||
dst.ToolCallCount += src.ToolCallCount
|
||||
dst.ErrorCount += src.ErrorCount
|
||||
dst.UniqueUsers += src.UniqueUsers
|
||||
dst.MemoryDocs += src.MemoryDocs
|
||||
dst.MemoryChunks += src.MemoryChunks
|
||||
dst.KGEntities += src.KGEntities
|
||||
dst.KGRelations += src.KGRelations
|
||||
if dst.RequestCount > 0 {
|
||||
dst.AvgDurationMS = int(weightedDuration / int64(dst.RequestCount))
|
||||
}
|
||||
}
|
||||
|
||||
func mergeSnapshotBreakdowns(base, extra []store.SnapshotBreakdown, groupBy string) []store.SnapshotBreakdown {
|
||||
if len(extra) == 0 {
|
||||
return base
|
||||
}
|
||||
byKey := make(map[string]int, len(base)+len(extra))
|
||||
var result []store.SnapshotBreakdown
|
||||
add := func(row store.SnapshotBreakdown) {
|
||||
if row.Key == "" && groupBy != "agent" {
|
||||
return
|
||||
}
|
||||
if idx, ok := byKey[row.Key]; ok {
|
||||
addBreakdown(&result[idx], row)
|
||||
return
|
||||
}
|
||||
result = append(result, row)
|
||||
byKey[row.Key] = len(result) - 1
|
||||
}
|
||||
for _, row := range base {
|
||||
add(row)
|
||||
}
|
||||
for _, row := range extra {
|
||||
add(row)
|
||||
}
|
||||
|
||||
sort.Slice(result, func(i, j int) bool {
|
||||
if groupBy == "channel" {
|
||||
return result[i].RequestCount > result[j].RequestCount
|
||||
}
|
||||
return result[i].InputTokens > result[j].InputTokens
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
func addBreakdown(dst *store.SnapshotBreakdown, src store.SnapshotBreakdown) {
|
||||
oldRequests := dst.RequestCount
|
||||
weightedDuration := int64(dst.AvgDurationMS)*int64(oldRequests) + int64(src.AvgDurationMS)*int64(src.RequestCount)
|
||||
dst.InputTokens += src.InputTokens
|
||||
dst.OutputTokens += src.OutputTokens
|
||||
dst.CacheReadTokens += src.CacheReadTokens
|
||||
dst.CacheCreateTokens += src.CacheCreateTokens
|
||||
dst.TotalCost += src.TotalCost
|
||||
dst.RequestCount += src.RequestCount
|
||||
dst.LLMCallCount += src.LLMCallCount
|
||||
dst.ToolCallCount += src.ToolCallCount
|
||||
dst.ErrorCount += src.ErrorCount
|
||||
if dst.RequestCount > 0 {
|
||||
dst.AvgDurationMS = int(weightedDuration / int64(dst.RequestCount))
|
||||
}
|
||||
}
|
||||
|
||||
func hasTimeSeriesUsage(point store.SnapshotTimeSeries) bool {
|
||||
return point.RequestCount > 0 ||
|
||||
point.LLMCallCount > 0 ||
|
||||
point.ToolCallCount > 0 ||
|
||||
point.InputTokens > 0 ||
|
||||
point.OutputTokens > 0 ||
|
||||
point.CacheReadTokens > 0 ||
|
||||
point.CacheCreateTokens > 0 ||
|
||||
point.ThinkingTokens > 0 ||
|
||||
point.TotalCost > 0 ||
|
||||
point.ErrorCount > 0
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
_ "modernc.org/sqlite"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
type emptyUsageSnapshotStore struct{}
|
||||
|
||||
func (emptyUsageSnapshotStore) UpsertSnapshots(context.Context, []store.UsageSnapshot) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (emptyUsageSnapshotStore) GetTimeSeries(context.Context, store.SnapshotQuery) ([]store.SnapshotTimeSeries, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (emptyUsageSnapshotStore) GetBreakdown(context.Context, store.SnapshotQuery) ([]store.SnapshotBreakdown, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (emptyUsageSnapshotStore) GetLatestBucket(context.Context) (*time.Time, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func TestUsageSummaryIncludesLiveCurrentHourCost(t *testing.T) {
|
||||
db, tenantID := usageLiveTestDB(t)
|
||||
insertUsageLiveTrace(t, db, tenantID, time.Now().UTC())
|
||||
|
||||
h := NewUsageHandler(emptyUsageSnapshotStore{}, nil, db)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/usage/summary?period=24h", nil)
|
||||
req = req.WithContext(store.WithTenantID(req.Context(), tenantID))
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
h.handleSummary(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Current usageSummary `json:"current"`
|
||||
}
|
||||
if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if body.Current.Cost != 0.25 {
|
||||
t.Fatalf("current cost = %v, want 0.25", body.Current.Cost)
|
||||
}
|
||||
if body.Current.InputTokens != 1100 || body.Current.OutputTokens != 250 {
|
||||
t.Fatalf("tokens = %d/%d, want 1100/250", body.Current.InputTokens, body.Current.OutputTokens)
|
||||
}
|
||||
if body.Current.Requests != 1 || body.Current.LLMCalls != 1 || body.Current.ToolCalls != 2 {
|
||||
t.Fatalf("counts = requests %d llm %d tools %d, want 1/1/2", body.Current.Requests, body.Current.LLMCalls, body.Current.ToolCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageSummaryHandlesEmptyLiveCurrentHour(t *testing.T) {
|
||||
db, tenantID := usageLiveTestDB(t)
|
||||
|
||||
h := NewUsageHandler(emptyUsageSnapshotStore{}, nil, db)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/usage/summary?period=24h", nil)
|
||||
req = req.WithContext(store.WithTenantID(req.Context(), tenantID))
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
h.handleSummary(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Current usageSummary `json:"current"`
|
||||
}
|
||||
if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if body.Current.Cost != 0 || body.Current.Requests != 0 || body.Current.LLMCalls != 0 {
|
||||
t.Fatalf("current = %+v, want empty live usage", body.Current)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageBreakdownIncludesLiveCurrentHourProviderModelCost(t *testing.T) {
|
||||
db, tenantID := usageLiveTestDB(t)
|
||||
now := time.Now().UTC()
|
||||
insertUsageLiveTrace(t, db, tenantID, now)
|
||||
|
||||
h := NewUsageHandler(emptyUsageSnapshotStore{}, nil, db)
|
||||
target := fmt.Sprintf(
|
||||
"/v1/usage/breakdown?from=%s&to=%s&group_by=provider_model",
|
||||
now.Add(-time.Hour).Format(time.RFC3339),
|
||||
now.Add(time.Hour).Format(time.RFC3339),
|
||||
)
|
||||
req := httptest.NewRequest(http.MethodGet, target, nil)
|
||||
req = req.WithContext(store.WithTenantID(req.Context(), tenantID))
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
h.handleBreakdown(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Rows []store.SnapshotBreakdown `json:"rows"`
|
||||
}
|
||||
if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if len(body.Rows) != 1 {
|
||||
t.Fatalf("rows = %+v, want one provider/model row", body.Rows)
|
||||
}
|
||||
row := body.Rows[0]
|
||||
if row.Key != "bailian/qwen3.7-plus" {
|
||||
t.Fatalf("key = %q, want bailian/qwen3.7-plus", row.Key)
|
||||
}
|
||||
if row.TotalCost != 0.25 || row.InputTokens != 1100 || row.OutputTokens != 250 || row.LLMCallCount != 1 {
|
||||
t.Fatalf("row = %+v, want live cost/tokens/llm count", row)
|
||||
}
|
||||
}
|
||||
|
||||
func usageLiveTestDB(t *testing.T) (*sql.DB, uuid.UUID) {
|
||||
t.Helper()
|
||||
db, err := sql.Open("sqlite", ":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { db.Close() })
|
||||
|
||||
schema := []string{
|
||||
`CREATE TABLE traces (
|
||||
id TEXT PRIMARY KEY,
|
||||
tenant_id TEXT NOT NULL,
|
||||
agent_id TEXT,
|
||||
user_id TEXT,
|
||||
start_time TEXT NOT NULL,
|
||||
duration_ms INTEGER,
|
||||
status TEXT,
|
||||
channel TEXT,
|
||||
parent_trace_id TEXT,
|
||||
tool_call_count INTEGER
|
||||
)`,
|
||||
`CREATE TABLE spans (
|
||||
id TEXT PRIMARY KEY,
|
||||
trace_id TEXT NOT NULL,
|
||||
tenant_id TEXT NOT NULL,
|
||||
span_type TEXT NOT NULL,
|
||||
provider TEXT,
|
||||
model TEXT,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
total_cost REAL,
|
||||
metadata TEXT,
|
||||
start_time TEXT
|
||||
)`,
|
||||
}
|
||||
for _, stmt := range schema {
|
||||
if _, err := db.Exec(stmt); err != nil {
|
||||
t.Fatalf("create schema: %v", err)
|
||||
}
|
||||
}
|
||||
return db, uuid.Must(uuid.NewV7())
|
||||
}
|
||||
|
||||
func insertUsageLiveTrace(t *testing.T, db *sql.DB, tenantID uuid.UUID, start time.Time) {
|
||||
t.Helper()
|
||||
traceID := uuid.Must(uuid.NewV7())
|
||||
spanID := uuid.Must(uuid.NewV7())
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO traces (
|
||||
id, tenant_id, agent_id, user_id, start_time, duration_ms, status, channel, tool_call_count
|
||||
) VALUES (?, ?, ?, 'user-a', ?, 1200, 'completed', 'web', 2)`,
|
||||
traceID, tenantID, uuid.Must(uuid.NewV7()), start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert trace: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO spans (
|
||||
id, trace_id, tenant_id, span_type, provider, model, input_tokens, output_tokens, total_cost, metadata, start_time
|
||||
) VALUES (?, ?, ?, 'llm_call', 'bailian', 'qwen3.7-plus', 1000, 200, 0.25, ?, ?)`,
|
||||
spanID, traceID, tenantID,
|
||||
`{"cache_read_tokens":7,"cache_creation_tokens":11,"thinking_tokens":13}`,
|
||||
start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert span: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO spans (
|
||||
id, trace_id, tenant_id, span_type, provider, model, input_tokens, output_tokens, total_cost, metadata, start_time
|
||||
) VALUES (?, ?, ?, 'tool_call', 'bailian', 'qwen3.7-plus', 100, 50, 0, ?, ?)`,
|
||||
uuid.Must(uuid.NewV7()), traceID, tenantID,
|
||||
`{"cache_read_tokens":0,"cache_creation_tokens":0,"thinking_tokens":0}`,
|
||||
start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert tool span: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -182,6 +182,10 @@ func (s *PGSnapshotStore) GetBreakdown(ctx context.Context, q store.SnapshotQuer
|
||||
groupCol = "model"
|
||||
orderExpr = "SUM(CASE WHEN provider != '' THEN input_tokens ELSE 0 END) DESC"
|
||||
extraFilter = " AND provider != '' AND model != ''"
|
||||
case "provider_model":
|
||||
groupCol = "provider || '/' || model"
|
||||
orderExpr = "SUM(CASE WHEN provider != '' THEN input_tokens ELSE 0 END) DESC"
|
||||
extraFilter = " AND provider != '' AND model != ''"
|
||||
case "channel":
|
||||
groupCol = "channel"
|
||||
orderExpr = "SUM(CASE WHEN provider = '' AND model = '' THEN request_count ELSE 0 END) DESC"
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
const pricedUsageSnapshotsCTE = `
|
||||
WITH priced_snapshots AS (
|
||||
SELECT
|
||||
us.id AS snapshot_id,
|
||||
calc.total_cost
|
||||
FROM usage_snapshots us
|
||||
LEFT JOIN llm_providers lp ON lp.tenant_id = us.tenant_id AND lp.name = us.provider
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
btrim(COALESCE(us.model, '')) AS model_id,
|
||||
replace(lower(btrim(COALESCE(lp.name, us.provider, ''))), ' ', '-') AS provider_name,
|
||||
COALESCE(lp.provider_type, '') AS provider_type
|
||||
) base
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ARRAY_REMOVE(ARRAY[
|
||||
NULLIF(base.model_id, '')::text,
|
||||
CASE
|
||||
WHEN position('/' IN base.model_id) > 0 THEN NULL
|
||||
WHEN base.provider_type = 'anthropic_native' THEN 'anthropic/' || base.model_id
|
||||
WHEN base.provider_type IN ('gemini_native', 'vertex') THEN 'google/' || base.model_id
|
||||
WHEN base.provider_type = 'openai_compat' AND base.provider_name IN ('openai', 'azure', 'azure-openai', 'azure_openai') THEN 'openai/' || base.model_id
|
||||
WHEN base.provider_type = 'openai_compat' AND base.provider_name = 'anthropic' THEN 'anthropic/' || base.model_id
|
||||
WHEN base.provider_type = 'openai_compat' AND base.provider_name IN ('gemini', 'google', 'vertex') THEN 'google/' || base.model_id
|
||||
WHEN base.provider_type = 'groq' THEN 'groq/' || base.model_id
|
||||
WHEN base.provider_type = 'deepseek' THEN 'deepseek/' || base.model_id
|
||||
WHEN base.provider_type = 'mistral' THEN 'mistralai/' || base.model_id
|
||||
WHEN base.provider_type = 'xai' THEN 'x-ai/' || base.model_id
|
||||
WHEN base.provider_type = 'minimax_native' THEN 'minimax/' || base.model_id
|
||||
WHEN base.provider_type = 'cohere' THEN 'cohere/' || base.model_id
|
||||
WHEN base.provider_type = 'perplexity' THEN 'perplexity/' || base.model_id
|
||||
WHEN base.provider_type IN ('dashscope', 'bailian') THEN 'qwen/' || base.model_id
|
||||
ELSE NULL
|
||||
END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'gpt-%' THEN 'openai/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND (lower(base.model_id) LIKE 'o1%' OR lower(base.model_id) LIKE 'o3%' OR lower(base.model_id) LIKE 'o4%' OR lower(base.model_id) LIKE 'o5%') THEN 'openai/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'qwen%' THEN 'qwen/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'claude-%' THEN 'anthropic/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'gemini-%' THEN 'google/' || base.model_id END
|
||||
], NULL)::text[] AS model_ids
|
||||
) candidates
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
p.input_price,
|
||||
p.output_price,
|
||||
p.cache_read_price,
|
||||
p.cache_write_price,
|
||||
p.reasoning_price
|
||||
FROM (
|
||||
SELECT
|
||||
0 AS source_priority,
|
||||
COALESCE(array_position(candidates.model_ids, o.model_id), 2147483647) AS candidate_order,
|
||||
o.input_price,
|
||||
o.output_price,
|
||||
o.cache_read_price,
|
||||
o.cache_write_price,
|
||||
o.reasoning_price
|
||||
FROM usage_pricing_overrides o
|
||||
WHERE lp.id IS NOT NULL
|
||||
AND o.tenant_id = us.tenant_id
|
||||
AND o.provider_id = lp.id
|
||||
AND o.enabled = true
|
||||
AND o.model_id = ANY(candidates.model_ids)
|
||||
UNION ALL
|
||||
SELECT
|
||||
1 AS source_priority,
|
||||
LEAST(
|
||||
COALESCE(array_position(candidates.model_ids, c.model_id), 2147483647),
|
||||
COALESCE(array_position(candidates.model_ids, c.canonical_model_id), 2147483647)
|
||||
) AS candidate_order,
|
||||
c.input_price,
|
||||
c.output_price,
|
||||
c.cache_read_price,
|
||||
c.cache_write_price,
|
||||
c.reasoning_price
|
||||
FROM usage_pricing_catalog c
|
||||
WHERE c.model_id = ANY(candidates.model_ids)
|
||||
OR c.canonical_model_id = ANY(candidates.model_ids)
|
||||
) p
|
||||
ORDER BY p.source_priority, p.candidate_order
|
||||
LIMIT 1
|
||||
) price
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
GREATEST(
|
||||
COALESCE(us.input_tokens, 0)::numeric -
|
||||
CASE
|
||||
WHEN base.provider_type IN ('openai_compat', 'openrouter', 'dashscope', 'bailian', 'chatgpt_oauth')
|
||||
THEN COALESCE(us.cache_read_tokens, 0)::numeric + COALESCE(us.cache_create_tokens, 0)::numeric
|
||||
ELSE 0
|
||||
END,
|
||||
0
|
||||
) AS input_tokens,
|
||||
COALESCE(us.output_tokens, 0)::numeric AS output_tokens,
|
||||
COALESCE(us.cache_read_tokens, 0)::numeric AS cache_read_tokens,
|
||||
COALESCE(us.cache_create_tokens, 0)::numeric AS cache_write_tokens,
|
||||
COALESCE(us.thinking_tokens, 0)::numeric AS reasoning_tokens
|
||||
) usage
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT (
|
||||
CASE WHEN usage.input_tokens > 0 THEN CEIL(usage.input_tokens * price.input_price * 1000000) / 1000000 ELSE 0 END +
|
||||
CASE
|
||||
WHEN usage.output_tokens > 0 AND usage.reasoning_tokens > 0 AND price.reasoning_price IS NOT NULL THEN
|
||||
CEIL(GREATEST(usage.output_tokens - usage.reasoning_tokens, 0) * price.output_price * 1000000) / 1000000 +
|
||||
CEIL(usage.reasoning_tokens * price.reasoning_price * 1000000) / 1000000
|
||||
WHEN usage.output_tokens > 0 THEN CEIL(usage.output_tokens * price.output_price * 1000000) / 1000000
|
||||
ELSE 0
|
||||
END +
|
||||
CASE WHEN usage.cache_read_tokens > 0 THEN CEIL(usage.cache_read_tokens * price.cache_read_price * 1000000) / 1000000 ELSE 0 END +
|
||||
CASE WHEN usage.cache_write_tokens > 0 THEN CEIL(usage.cache_write_tokens * price.cache_write_price * 1000000) / 1000000 ELSE 0 END
|
||||
)::numeric(12,8) AS total_cost
|
||||
) calc
|
||||
WHERE COALESCE(us.provider, '') <> ''
|
||||
AND base.model_id <> ''
|
||||
AND COALESCE(us.total_cost, 0) = 0
|
||||
AND COALESCE(us.input_tokens, 0) + COALESCE(us.output_tokens, 0) + COALESCE(us.cache_read_tokens, 0) + COALESCE(us.cache_create_tokens, 0) + COALESCE(us.thinking_tokens, 0) > 0
|
||||
AND (usage.input_tokens = 0 OR price.input_price IS NOT NULL)
|
||||
AND (usage.output_tokens = 0 OR price.output_price IS NOT NULL)
|
||||
AND (usage.cache_read_tokens = 0 OR price.cache_read_price IS NOT NULL)
|
||||
AND (usage.cache_write_tokens = 0 OR price.cache_write_price IS NOT NULL)
|
||||
AND calc.total_cost > 0
|
||||
)`
|
||||
|
||||
func (s *PGSnapshotStore) BackfillSnapshotCosts(ctx context.Context) (store.SnapshotCostBackfillStats, error) {
|
||||
var stats store.SnapshotCostBackfillStats
|
||||
res, err := s.db.ExecContext(ctx, pricedUsageSnapshotsCTE+`
|
||||
UPDATE usage_snapshots us
|
||||
SET total_cost = priced_snapshots.total_cost
|
||||
FROM priced_snapshots
|
||||
WHERE us.id = priced_snapshots.snapshot_id
|
||||
AND us.total_cost IS DISTINCT FROM priced_snapshots.total_cost`)
|
||||
if err != nil {
|
||||
return stats, fmt.Errorf("backfill snapshot costs: %w", err)
|
||||
}
|
||||
stats.SnapshotRowsUpdated, _ = res.RowsAffected()
|
||||
return stats, nil
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func TestPGSnapshotStoreBackfillSnapshotCosts(t *testing.T) {
|
||||
db := hooksTestDB(t)
|
||||
tenantID, agentID := seedTenantAndAgent(t, db)
|
||||
providerID := uuid.Must(uuid.NewV7())
|
||||
hour := time.Date(2026, 6, 20, 8, 0, 0, 0, time.UTC)
|
||||
|
||||
t.Cleanup(func() {
|
||||
db.Exec(`DELETE FROM usage_snapshots WHERE tenant_id=$1 AND bucket_hour=$2`, tenantID, hour)
|
||||
db.Exec(`DELETE FROM llm_providers WHERE id=$1`, providerID)
|
||||
db.Exec(`DELETE FROM usage_pricing_catalog WHERE model_id=$1`, "qwen/qwen3.7-plus")
|
||||
})
|
||||
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled)
|
||||
VALUES ($1,$2,'bailian','bailian','test',true)`,
|
||||
providerID, tenantID,
|
||||
); err != nil {
|
||||
t.Fatalf("insert provider: %v", err)
|
||||
}
|
||||
inputPrice := "0.000001"
|
||||
outputPrice := "0.000002"
|
||||
cacheReadPrice := "0.0000001"
|
||||
cacheWritePrice := "0.0000002"
|
||||
if _, err := NewPGUsageCapStore(db).UpsertPricingCatalog(context.Background(), []store.UsagePricingCatalogEntry{{
|
||||
ModelID: "qwen/qwen3.7-plus",
|
||||
CanonicalModelID: "qwen/qwen3.7-plus",
|
||||
Pricing: store.UsagePricingFields{
|
||||
Input: &inputPrice,
|
||||
Output: &outputPrice,
|
||||
CacheRead: &cacheReadPrice,
|
||||
CacheWrite: &cacheWritePrice,
|
||||
},
|
||||
SyncedAt: time.Now().UTC(),
|
||||
}}); err != nil {
|
||||
t.Fatalf("upsert pricing catalog: %v", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO usage_snapshots (
|
||||
bucket_hour, agent_id, provider, model, channel,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_create_tokens, thinking_tokens,
|
||||
total_cost, request_count, llm_call_count, tool_call_count,
|
||||
error_count, unique_users, avg_duration_ms, tenant_id
|
||||
) VALUES ($1,$2,'bailian','qwen3.7-plus','web',1000,500,10,5,100,0,0,1,0,0,0,0,$3)`,
|
||||
hour, agentID, tenantID,
|
||||
); err != nil {
|
||||
t.Fatalf("insert snapshot: %v", err)
|
||||
}
|
||||
|
||||
stats, err := NewPGSnapshotStore(db).BackfillSnapshotCosts(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("BackfillSnapshotCosts: %v", err)
|
||||
}
|
||||
if stats.SnapshotRowsUpdated == 0 {
|
||||
t.Fatalf("stats = %+v, want updated snapshot", stats)
|
||||
}
|
||||
|
||||
var cost float64
|
||||
if err := db.QueryRow(`SELECT total_cost FROM usage_snapshots WHERE tenant_id=$1 AND bucket_hour=$2 AND provider='bailian'`, tenantID, hour).Scan(&cost); err != nil {
|
||||
t.Fatalf("query snapshot cost: %v", err)
|
||||
}
|
||||
if math.Abs(cost-0.001987) > 0.0000001 {
|
||||
t.Fatalf("snapshot cost = %.8f, want 0.001987", cost)
|
||||
}
|
||||
|
||||
points, err := NewPGSnapshotStore(db).GetTimeSeries(store.WithTenantID(context.Background(), tenantID), store.SnapshotQuery{
|
||||
From: hour.Add(-time.Hour),
|
||||
To: hour.Add(time.Hour),
|
||||
GroupBy: "hour",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("GetTimeSeries: %v", err)
|
||||
}
|
||||
if len(points) != 1 || math.Abs(points[0].TotalCost-0.001987) > 0.0000001 {
|
||||
t.Fatalf("points = %+v, want backfilled cost", points)
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/lib/pq"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
@@ -286,6 +287,60 @@ func (s *PGTracingStore) ListChildTraces(ctx context.Context, parentTraceID uuid
|
||||
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()
|
||||
@@ -397,15 +452,15 @@ func (s *PGTracingStore) BatchUpdateTraceAggregates(ctx context.Context, traceID
|
||||
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 = 'llm_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 = 'llm_call' AND output_tokens IS NOT NULL), 0),
|
||||
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 = 'llm_call' AND metadata IS NOT NULL
|
||||
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
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/lib/pq"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
const pricedLLMSpansCTE = `
|
||||
WITH priced_spans AS (
|
||||
SELECT
|
||||
s.id AS span_id,
|
||||
s.trace_id,
|
||||
t.start_time AS trace_start_time,
|
||||
calc.total_cost
|
||||
FROM spans s
|
||||
JOIN traces t ON t.id = s.trace_id AND t.tenant_id = s.tenant_id
|
||||
LEFT JOIN llm_providers lp ON lp.tenant_id = s.tenant_id AND lp.name = s.provider
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
btrim(COALESCE(s.model, '')) AS model_id,
|
||||
replace(lower(btrim(COALESCE(lp.name, s.provider, ''))), ' ', '-') AS provider_name,
|
||||
COALESCE(lp.provider_type, '') AS provider_type
|
||||
) base
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ARRAY_REMOVE(ARRAY[
|
||||
NULLIF(base.model_id, '')::text,
|
||||
CASE
|
||||
WHEN position('/' IN base.model_id) > 0 THEN NULL
|
||||
WHEN base.provider_type = 'anthropic_native' THEN 'anthropic/' || base.model_id
|
||||
WHEN base.provider_type IN ('gemini_native', 'vertex') THEN 'google/' || base.model_id
|
||||
WHEN base.provider_type = 'openai_compat' AND base.provider_name IN ('openai', 'azure', 'azure-openai', 'azure_openai') THEN 'openai/' || base.model_id
|
||||
WHEN base.provider_type = 'openai_compat' AND base.provider_name = 'anthropic' THEN 'anthropic/' || base.model_id
|
||||
WHEN base.provider_type = 'openai_compat' AND base.provider_name IN ('gemini', 'google', 'vertex') THEN 'google/' || base.model_id
|
||||
WHEN base.provider_type = 'groq' THEN 'groq/' || base.model_id
|
||||
WHEN base.provider_type = 'deepseek' THEN 'deepseek/' || base.model_id
|
||||
WHEN base.provider_type = 'mistral' THEN 'mistralai/' || base.model_id
|
||||
WHEN base.provider_type = 'xai' THEN 'x-ai/' || base.model_id
|
||||
WHEN base.provider_type = 'minimax_native' THEN 'minimax/' || base.model_id
|
||||
WHEN base.provider_type = 'cohere' THEN 'cohere/' || base.model_id
|
||||
WHEN base.provider_type = 'perplexity' THEN 'perplexity/' || base.model_id
|
||||
WHEN base.provider_type IN ('dashscope', 'bailian') THEN 'qwen/' || base.model_id
|
||||
ELSE NULL
|
||||
END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'gpt-%' THEN 'openai/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND (lower(base.model_id) LIKE 'o1%' OR lower(base.model_id) LIKE 'o3%' OR lower(base.model_id) LIKE 'o4%' OR lower(base.model_id) LIKE 'o5%') THEN 'openai/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'qwen%' THEN 'qwen/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'claude-%' THEN 'anthropic/' || base.model_id END,
|
||||
CASE WHEN position('/' IN base.model_id) = 0 AND lower(base.model_id) LIKE 'gemini-%' THEN 'google/' || base.model_id END
|
||||
], NULL)::text[] AS model_ids
|
||||
) candidates
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
p.input_price,
|
||||
p.output_price,
|
||||
p.cache_read_price,
|
||||
p.cache_write_price,
|
||||
p.reasoning_price
|
||||
FROM (
|
||||
SELECT
|
||||
0 AS source_priority,
|
||||
COALESCE(array_position(candidates.model_ids, o.model_id), 2147483647) AS candidate_order,
|
||||
o.input_price,
|
||||
o.output_price,
|
||||
o.cache_read_price,
|
||||
o.cache_write_price,
|
||||
o.reasoning_price
|
||||
FROM usage_pricing_overrides o
|
||||
WHERE lp.id IS NOT NULL
|
||||
AND o.tenant_id = s.tenant_id
|
||||
AND o.provider_id = lp.id
|
||||
AND o.enabled = true
|
||||
AND o.model_id = ANY(candidates.model_ids)
|
||||
UNION ALL
|
||||
SELECT
|
||||
1 AS source_priority,
|
||||
LEAST(
|
||||
COALESCE(array_position(candidates.model_ids, c.model_id), 2147483647),
|
||||
COALESCE(array_position(candidates.model_ids, c.canonical_model_id), 2147483647)
|
||||
) AS candidate_order,
|
||||
c.input_price,
|
||||
c.output_price,
|
||||
c.cache_read_price,
|
||||
c.cache_write_price,
|
||||
c.reasoning_price
|
||||
FROM usage_pricing_catalog c
|
||||
WHERE c.model_id = ANY(candidates.model_ids)
|
||||
OR c.canonical_model_id = ANY(candidates.model_ids)
|
||||
) p
|
||||
ORDER BY p.source_priority, p.candidate_order
|
||||
LIMIT 1
|
||||
) price
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
GREATEST(
|
||||
COALESCE(s.input_tokens, 0)::numeric -
|
||||
CASE
|
||||
WHEN lower(COALESCE(s.metadata->>'prompt_tokens_include_cached_segments', '')) = 'true'
|
||||
OR (
|
||||
NOT (s.metadata ? 'prompt_tokens_include_cached_segments')
|
||||
AND base.provider_type IN ('openai_compat', 'openrouter', 'dashscope', 'bailian', 'chatgpt_oauth')
|
||||
)
|
||||
THEN
|
||||
(CASE WHEN s.metadata ? 'cache_read_tokens' AND (s.metadata->>'cache_read_tokens') ~ '^[0-9]+$'
|
||||
THEN (s.metadata->>'cache_read_tokens')::numeric ELSE 0 END) +
|
||||
(CASE WHEN s.metadata ? 'cache_creation_tokens' AND (s.metadata->>'cache_creation_tokens') ~ '^[0-9]+$'
|
||||
THEN (s.metadata->>'cache_creation_tokens')::numeric ELSE 0 END)
|
||||
ELSE 0
|
||||
END,
|
||||
0
|
||||
) AS input_tokens,
|
||||
COALESCE(s.output_tokens, 0)::numeric AS output_tokens,
|
||||
CASE WHEN s.metadata ? 'cache_read_tokens' AND (s.metadata->>'cache_read_tokens') ~ '^[0-9]+$'
|
||||
THEN (s.metadata->>'cache_read_tokens')::numeric ELSE 0 END AS cache_read_tokens,
|
||||
CASE WHEN s.metadata ? 'cache_creation_tokens' AND (s.metadata->>'cache_creation_tokens') ~ '^[0-9]+$'
|
||||
THEN (s.metadata->>'cache_creation_tokens')::numeric ELSE 0 END AS cache_write_tokens,
|
||||
CASE WHEN s.metadata ? 'thinking_tokens' AND (s.metadata->>'thinking_tokens') ~ '^[0-9]+$'
|
||||
THEN (s.metadata->>'thinking_tokens')::numeric ELSE 0 END AS reasoning_tokens
|
||||
) usage
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT (
|
||||
CASE WHEN usage.input_tokens > 0 THEN CEIL(usage.input_tokens * price.input_price * 1000000) / 1000000 ELSE 0 END +
|
||||
CASE
|
||||
WHEN usage.output_tokens > 0 AND usage.reasoning_tokens > 0 AND price.reasoning_price IS NOT NULL THEN
|
||||
CEIL(GREATEST(usage.output_tokens - usage.reasoning_tokens, 0) * price.output_price * 1000000) / 1000000 +
|
||||
CEIL(usage.reasoning_tokens * price.reasoning_price * 1000000) / 1000000
|
||||
WHEN usage.output_tokens > 0 THEN CEIL(usage.output_tokens * price.output_price * 1000000) / 1000000
|
||||
ELSE 0
|
||||
END +
|
||||
CASE WHEN usage.cache_read_tokens > 0 THEN CEIL(usage.cache_read_tokens * price.cache_read_price * 1000000) / 1000000 ELSE 0 END +
|
||||
CASE WHEN usage.cache_write_tokens > 0 THEN CEIL(usage.cache_write_tokens * price.cache_write_price * 1000000) / 1000000 ELSE 0 END
|
||||
)::numeric(12,8) AS total_cost
|
||||
) calc
|
||||
WHERE s.span_type IN ('llm_call', 'tool_call')
|
||||
AND base.model_id <> ''
|
||||
AND COALESCE(s.total_cost, 0) = 0
|
||||
AND (usage.input_tokens = 0 OR price.input_price IS NOT NULL)
|
||||
AND (usage.output_tokens = 0 OR price.output_price IS NOT NULL)
|
||||
AND (usage.cache_read_tokens = 0 OR price.cache_read_price IS NOT NULL)
|
||||
AND (usage.cache_write_tokens = 0 OR price.cache_write_price IS NOT NULL)
|
||||
AND calc.total_cost > 0
|
||||
)`
|
||||
|
||||
func (s *PGTracingStore) BackfillLLMCosts(ctx context.Context) (store.TraceCostBackfillStats, error) {
|
||||
var stats store.TraceCostBackfillStats
|
||||
traceIDs, buckets, err := s.listPricedLLMBackfillTargets(ctx)
|
||||
if err != nil {
|
||||
return stats, err
|
||||
}
|
||||
if len(traceIDs) == 0 {
|
||||
return stats, nil
|
||||
}
|
||||
stats.SnapshotBuckets = buckets
|
||||
|
||||
res, err := s.db.ExecContext(ctx, pricedLLMSpansCTE+`
|
||||
UPDATE spans s
|
||||
SET total_cost = priced_spans.total_cost
|
||||
FROM priced_spans
|
||||
WHERE s.id = priced_spans.span_id
|
||||
AND s.total_cost IS DISTINCT FROM priced_spans.total_cost`)
|
||||
if err != nil {
|
||||
return stats, fmt.Errorf("backfill span costs: %w", err)
|
||||
}
|
||||
stats.SpanRowsUpdated, _ = res.RowsAffected()
|
||||
|
||||
stats.TraceRowsUpdated, err = s.refreshTraceCostAggregates(ctx, traceIDs)
|
||||
if err != nil {
|
||||
return stats, err
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
func (s *PGTracingStore) listPricedLLMBackfillTargets(ctx context.Context) ([]uuid.UUID, []time.Time, error) {
|
||||
rows, err := s.db.QueryContext(ctx, pricedLLMSpansCTE+`
|
||||
SELECT DISTINCT trace_id, date_trunc('hour', trace_start_time) AS bucket_hour
|
||||
FROM priced_spans`)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("list priced span targets: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
traceSeen := make(map[uuid.UUID]struct{})
|
||||
bucketSeen := make(map[time.Time]struct{})
|
||||
var traceIDs []uuid.UUID
|
||||
var buckets []time.Time
|
||||
for rows.Next() {
|
||||
var traceID uuid.UUID
|
||||
var bucket time.Time
|
||||
if err := rows.Scan(&traceID, &bucket); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if _, ok := traceSeen[traceID]; !ok {
|
||||
traceSeen[traceID] = struct{}{}
|
||||
traceIDs = append(traceIDs, traceID)
|
||||
}
|
||||
bucket = bucket.UTC().Truncate(time.Hour)
|
||||
if _, ok := bucketSeen[bucket]; !ok {
|
||||
bucketSeen[bucket] = struct{}{}
|
||||
buckets = append(buckets, bucket)
|
||||
}
|
||||
}
|
||||
return traceIDs, buckets, rows.Err()
|
||||
}
|
||||
|
||||
func (s *PGTracingStore) refreshTraceCostAggregates(ctx context.Context, traceIDs []uuid.UUID) (int64, error) {
|
||||
if len(traceIDs) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
ids := make([]string, len(traceIDs))
|
||||
for i, id := range traceIDs {
|
||||
ids[i] = id.String()
|
||||
}
|
||||
res, err := s.db.ExecContext(ctx, `
|
||||
WITH agg AS (
|
||||
SELECT
|
||||
trace_id,
|
||||
COUNT(*) AS span_count,
|
||||
COUNT(*) FILTER (WHERE span_type = 'llm_call') AS llm_call_count,
|
||||
COUNT(*) FILTER (WHERE span_type = 'tool_call') AS tool_call_count,
|
||||
COALESCE(SUM(input_tokens) FILTER (WHERE span_type IN ('llm_call', 'tool_call') AND input_tokens IS NOT NULL), 0) AS total_input_tokens,
|
||||
COALESCE(SUM(output_tokens) FILTER (WHERE span_type IN ('llm_call', 'tool_call') AND output_tokens IS NOT NULL), 0) AS total_output_tokens,
|
||||
COALESCE(SUM(total_cost) FILTER (WHERE total_cost IS NOT NULL), 0) AS total_cost,
|
||||
COALESCE(SUM(CASE WHEN metadata ? 'cache_read_tokens' AND (metadata->>'cache_read_tokens') ~ '^[0-9]+$'
|
||||
THEN (metadata->>'cache_read_tokens')::int ELSE 0 END), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(CASE WHEN metadata ? 'cache_creation_tokens' AND (metadata->>'cache_creation_tokens') ~ '^[0-9]+$'
|
||||
THEN (metadata->>'cache_creation_tokens')::int ELSE 0 END), 0) AS cache_creation_tokens
|
||||
FROM spans
|
||||
WHERE trace_id = ANY($1::uuid[])
|
||||
GROUP BY trace_id
|
||||
)
|
||||
UPDATE traces t
|
||||
SET
|
||||
span_count = agg.span_count,
|
||||
llm_call_count = agg.llm_call_count,
|
||||
tool_call_count = agg.tool_call_count,
|
||||
total_input_tokens = agg.total_input_tokens,
|
||||
total_output_tokens = agg.total_output_tokens,
|
||||
total_cost = agg.total_cost,
|
||||
metadata = jsonb_build_object(
|
||||
'total_cache_read_tokens', agg.cache_read_tokens,
|
||||
'total_cache_creation_tokens', agg.cache_creation_tokens
|
||||
)
|
||||
FROM agg
|
||||
WHERE t.id = agg.trace_id`, pq.Array(ids))
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return 0, nil
|
||||
}
|
||||
return 0, fmt.Errorf("refresh trace cost aggregates: %w", err)
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (s *PGTracingStore) ReconcileTraceUsageAggregates(ctx context.Context) (store.TraceUsageAggregateStats, error) {
|
||||
var stats store.TraceUsageAggregateStats
|
||||
res, err := s.db.ExecContext(ctx, `
|
||||
WITH agg AS (
|
||||
SELECT
|
||||
trace_id,
|
||||
COUNT(*) AS span_count,
|
||||
COUNT(*) FILTER (WHERE span_type = 'llm_call') AS llm_call_count,
|
||||
COUNT(*) FILTER (WHERE span_type = 'tool_call') AS tool_call_count,
|
||||
COALESCE(SUM(input_tokens) FILTER (WHERE span_type IN ('llm_call', 'tool_call') AND input_tokens IS NOT NULL), 0) AS total_input_tokens,
|
||||
COALESCE(SUM(output_tokens) FILTER (WHERE span_type IN ('llm_call', 'tool_call') AND output_tokens IS NOT NULL), 0) AS total_output_tokens,
|
||||
COALESCE(SUM(total_cost) FILTER (WHERE total_cost IS NOT NULL), 0) AS total_cost,
|
||||
COALESCE(SUM(CASE WHEN metadata ? 'cache_read_tokens' AND (metadata->>'cache_read_tokens') ~ '^[0-9]+$'
|
||||
THEN (metadata->>'cache_read_tokens')::int ELSE 0 END), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(CASE WHEN metadata ? 'cache_creation_tokens' AND (metadata->>'cache_creation_tokens') ~ '^[0-9]+$'
|
||||
THEN (metadata->>'cache_creation_tokens')::int ELSE 0 END), 0) AS cache_creation_tokens
|
||||
FROM spans
|
||||
GROUP BY trace_id
|
||||
)
|
||||
UPDATE traces t
|
||||
SET
|
||||
span_count = agg.span_count,
|
||||
llm_call_count = agg.llm_call_count,
|
||||
tool_call_count = agg.tool_call_count,
|
||||
total_input_tokens = agg.total_input_tokens,
|
||||
total_output_tokens = agg.total_output_tokens,
|
||||
total_cost = agg.total_cost,
|
||||
metadata = jsonb_build_object(
|
||||
'total_cache_read_tokens', agg.cache_read_tokens,
|
||||
'total_cache_creation_tokens', agg.cache_creation_tokens
|
||||
)
|
||||
FROM agg
|
||||
WHERE t.id = agg.trace_id
|
||||
AND (
|
||||
t.span_count IS DISTINCT FROM agg.span_count OR
|
||||
t.llm_call_count IS DISTINCT FROM agg.llm_call_count OR
|
||||
t.tool_call_count IS DISTINCT FROM agg.tool_call_count OR
|
||||
t.total_input_tokens IS DISTINCT FROM agg.total_input_tokens OR
|
||||
t.total_output_tokens IS DISTINCT FROM agg.total_output_tokens OR
|
||||
t.total_cost IS DISTINCT FROM agg.total_cost
|
||||
)`)
|
||||
if err != nil {
|
||||
return stats, fmt.Errorf("reconcile trace usage aggregates: %w", err)
|
||||
}
|
||||
stats.TraceRowsUpdated, _ = res.RowsAffected()
|
||||
return stats, nil
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func TestPGTracingStoreBackfillLLMCosts(t *testing.T) {
|
||||
db := hooksTestDB(t)
|
||||
tenantID, agentID := seedTenantAndAgent(t, db)
|
||||
providerID := uuid.Must(uuid.NewV7())
|
||||
traceID := uuid.Must(uuid.NewV7())
|
||||
spanID := uuid.Must(uuid.NewV7())
|
||||
toolSpanID := uuid.Must(uuid.NewV7())
|
||||
start := time.Date(2026, 7, 3, 1, 23, 0, 0, time.UTC)
|
||||
|
||||
t.Cleanup(func() {
|
||||
db.Exec(`DELETE FROM spans WHERE id IN ($1,$2)`, spanID, toolSpanID)
|
||||
db.Exec(`DELETE FROM traces WHERE id=$1`, traceID)
|
||||
db.Exec(`DELETE FROM llm_providers WHERE id=$1`, providerID)
|
||||
db.Exec(`DELETE FROM usage_pricing_catalog WHERE model_id=$1`, "qwen/qwen3.7-plus")
|
||||
})
|
||||
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled)
|
||||
VALUES ($1,$2,'bailian','bailian','test',true)`,
|
||||
providerID, tenantID,
|
||||
); err != nil {
|
||||
t.Fatalf("insert provider: %v", err)
|
||||
}
|
||||
inputPrice := "0.000001"
|
||||
outputPrice := "0.000002"
|
||||
cacheReadPrice := "0.0000001"
|
||||
cacheWritePrice := "0.0000002"
|
||||
if _, err := NewPGUsageCapStore(db).UpsertPricingCatalog(context.Background(), []store.UsagePricingCatalogEntry{{
|
||||
ModelID: "qwen/qwen3.7-plus",
|
||||
CanonicalModelID: "qwen/qwen3.7-plus",
|
||||
Pricing: store.UsagePricingFields{
|
||||
Input: &inputPrice,
|
||||
Output: &outputPrice,
|
||||
CacheRead: &cacheReadPrice,
|
||||
CacheWrite: &cacheWritePrice,
|
||||
},
|
||||
SyncedAt: time.Now().UTC(),
|
||||
}}); err != nil {
|
||||
t.Fatalf("upsert pricing catalog: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO traces (id, tenant_id, agent_id, start_time, created_at, status)
|
||||
VALUES ($1,$2,$3,$4,$4,'completed')`,
|
||||
traceID, tenantID, agentID, start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert trace: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO spans (
|
||||
id, trace_id, tenant_id, agent_id, span_type, start_time, created_at, status,
|
||||
provider, model, input_tokens, output_tokens, total_cost, metadata
|
||||
) VALUES ($1,$2,$3,$4,'llm_call',$5,$5,'completed','bailian','qwen3.7-plus',1000,500,0,$6::jsonb)`,
|
||||
spanID, traceID, tenantID, agentID, start,
|
||||
`{"cache_read_tokens":10,"cache_creation_tokens":5,"thinking_tokens":100}`,
|
||||
); err != nil {
|
||||
t.Fatalf("insert span: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO spans (
|
||||
id, trace_id, tenant_id, agent_id, span_type, start_time, created_at, status,
|
||||
name, provider, model, input_tokens, output_tokens, total_cost
|
||||
) VALUES ($1,$2,$3,$4,'tool_call',$5,$5,'completed','read_image','bailian','qwen3.7-plus',100,50,0)`,
|
||||
toolSpanID, traceID, tenantID, agentID, start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert tool span: %v", err)
|
||||
}
|
||||
|
||||
stats, err := NewPGTracingStore(db).BackfillLLMCosts(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("BackfillLLMCosts: %v", err)
|
||||
}
|
||||
if stats.SpanRowsUpdated == 0 || stats.TraceRowsUpdated == 0 || len(stats.SnapshotBuckets) == 0 {
|
||||
t.Fatalf("stats = %+v, want updated span, trace, and snapshot bucket", stats)
|
||||
}
|
||||
|
||||
var spanCost, toolSpanCost, traceCost float64
|
||||
if err := db.QueryRow(`SELECT total_cost FROM spans WHERE id=$1`, spanID).Scan(&spanCost); err != nil {
|
||||
t.Fatalf("query span cost: %v", err)
|
||||
}
|
||||
if err := db.QueryRow(`SELECT total_cost FROM spans WHERE id=$1`, toolSpanID).Scan(&toolSpanCost); err != nil {
|
||||
t.Fatalf("query tool span cost: %v", err)
|
||||
}
|
||||
if err := db.QueryRow(`SELECT total_cost FROM traces WHERE id=$1`, traceID).Scan(&traceCost); err != nil {
|
||||
t.Fatalf("query trace cost: %v", err)
|
||||
}
|
||||
if spanCost <= 0 || traceCost <= 0 {
|
||||
t.Fatalf("costs = span %.8f trace %.8f, want positive", spanCost, traceCost)
|
||||
}
|
||||
if math.Abs(spanCost-0.001987) > 0.0000001 {
|
||||
t.Fatalf("span cost = %.8f, want 0.001987", spanCost)
|
||||
}
|
||||
if math.Abs(toolSpanCost-0.0002) > 0.0000001 {
|
||||
t.Fatalf("tool span cost = %.8f, want 0.0002", toolSpanCost)
|
||||
}
|
||||
if math.Abs(traceCost-0.002187) > 0.0000001 {
|
||||
t.Fatalf("trace cost = %.8f, want 0.002187", traceCost)
|
||||
}
|
||||
|
||||
var inputTokens, outputTokens, llmCalls, toolCalls int
|
||||
if err := db.QueryRow(`SELECT total_input_tokens, total_output_tokens, llm_call_count, tool_call_count FROM traces WHERE id=$1`, traceID).Scan(&inputTokens, &outputTokens, &llmCalls, &toolCalls); err != nil {
|
||||
t.Fatalf("query trace aggregates: %v", err)
|
||||
}
|
||||
if inputTokens != 1100 || outputTokens != 550 || llmCalls != 1 || toolCalls != 1 {
|
||||
t.Fatalf("trace aggregates = input %d output %d llm %d tool %d, want 1100/550/1/1", inputTokens, outputTokens, llmCalls, toolCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPGTracingStoreReconcileTraceUsageAggregatesIncludesToolUsageTokens(t *testing.T) {
|
||||
db := hooksTestDB(t)
|
||||
tenantID, agentID := seedTenantAndAgent(t, db)
|
||||
traceID := uuid.Must(uuid.NewV7())
|
||||
llmSpanID := uuid.Must(uuid.NewV7())
|
||||
toolSpanID := uuid.Must(uuid.NewV7())
|
||||
start := time.Date(2026, 7, 3, 2, 0, 0, 0, time.UTC)
|
||||
|
||||
t.Cleanup(func() {
|
||||
db.Exec(`DELETE FROM spans WHERE id IN ($1,$2)`, llmSpanID, toolSpanID)
|
||||
db.Exec(`DELETE FROM traces WHERE id=$1`, traceID)
|
||||
})
|
||||
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO traces (
|
||||
id, tenant_id, agent_id, start_time, created_at, status,
|
||||
total_input_tokens, total_output_tokens, total_cost,
|
||||
span_count, llm_call_count, tool_call_count
|
||||
) VALUES ($1,$2,$3,$4,$4,'completed',1000,200,0.25,1,1,0)`,
|
||||
traceID, tenantID, agentID, start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert trace: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO spans (
|
||||
id, trace_id, tenant_id, agent_id, span_type, start_time, created_at, status,
|
||||
provider, model, input_tokens, output_tokens, total_cost
|
||||
) VALUES
|
||||
($1,$3,$4,$5,'llm_call',$6,$6,'completed','bailian','qwen3.7-plus',1000,200,0.25),
|
||||
($2,$3,$4,$5,'tool_call',$6,$6,'completed','bailian','qwen3.7-plus',100,50,0.05)`,
|
||||
llmSpanID, toolSpanID, traceID, tenantID, agentID, start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert spans: %v", err)
|
||||
}
|
||||
|
||||
stats, err := NewPGTracingStore(db).ReconcileTraceUsageAggregates(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("ReconcileTraceUsageAggregates: %v", err)
|
||||
}
|
||||
if stats.TraceRowsUpdated == 0 {
|
||||
t.Fatalf("stats = %+v, want updated trace", stats)
|
||||
}
|
||||
|
||||
var inputTokens, outputTokens, llmCalls, toolCalls, spanCount int
|
||||
var traceCost float64
|
||||
if err := db.QueryRow(`
|
||||
SELECT total_input_tokens, total_output_tokens, total_cost, span_count, llm_call_count, tool_call_count
|
||||
FROM traces WHERE id=$1`, traceID).Scan(&inputTokens, &outputTokens, &traceCost, &spanCount, &llmCalls, &toolCalls); err != nil {
|
||||
t.Fatalf("query trace aggregates: %v", err)
|
||||
}
|
||||
if inputTokens != 1100 || outputTokens != 250 || math.Abs(traceCost-0.30) > 0.0000001 || spanCount != 2 || llmCalls != 1 || toolCalls != 1 {
|
||||
t.Fatalf("trace aggregates = input %d output %d cost %.8f spans %d llm %d tool %d, want 1100/250/0.30/2/1/1",
|
||||
inputTokens, outputTokens, traceCost, spanCount, llmCalls, toolCalls)
|
||||
}
|
||||
}
|
||||
@@ -139,6 +139,40 @@ func TestValidateUsagePricingFieldsRejectsNegativeValues(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPGUsageCapStoreUpsertPricingCatalogSkipsInvalidEntries(t *testing.T) {
|
||||
db := hooksTestDB(t)
|
||||
usageStore := NewPGUsageCapStore(db)
|
||||
inputPrice := "0.000001"
|
||||
outputPrice := "0.000002"
|
||||
negativePrice := "-0.000001"
|
||||
entries := []store.UsagePricingCatalogEntry{
|
||||
{ModelID: "invalid/negative-output", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &negativePrice}, SyncedAt: time.Now().UTC()},
|
||||
{ModelID: "openai/gpt-valid", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()},
|
||||
}
|
||||
|
||||
count, err := usageStore.UpsertPricingCatalog(context.Background(), entries)
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertPricingCatalog: %v", err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Fatalf("upserted count = %d, want 1", count)
|
||||
}
|
||||
listed, err := usageStore.ListPricingCatalog(context.Background(), store.UsagePricingQuery{ModelID: "gpt-valid"})
|
||||
if err != nil {
|
||||
t.Fatalf("ListPricingCatalog valid: %v", err)
|
||||
}
|
||||
if len(listed) != 1 || listed[0].ModelID != "openai/gpt-valid" {
|
||||
t.Fatalf("valid entries = %+v, want openai/gpt-valid", listed)
|
||||
}
|
||||
listed, err = usageStore.ListPricingCatalog(context.Background(), store.UsagePricingQuery{ModelID: "negative-output"})
|
||||
if err != nil {
|
||||
t.Fatalf("ListPricingCatalog invalid: %v", err)
|
||||
}
|
||||
if len(listed) != 0 {
|
||||
t.Fatalf("invalid entries = %+v, want none", listed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPGUsageCapStoreResolvePricingUsesOpenRouterAliases(t *testing.T) {
|
||||
db := hooksTestDB(t)
|
||||
usageStore := NewPGUsageCapStore(db)
|
||||
@@ -146,8 +180,10 @@ func TestPGUsageCapStoreResolvePricingUsesOpenRouterAliases(t *testing.T) {
|
||||
outputPrice := "0.000002"
|
||||
entries := []store.UsagePricingCatalogEntry{
|
||||
{ModelID: "openai/gpt-4o-mini", CanonicalModelID: "openai/gpt-4o-mini", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()},
|
||||
{ModelID: "openai/gpt-5.5", CanonicalModelID: "openai/gpt-5.5", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()},
|
||||
{ModelID: "anthropic/claude-3-5-haiku-latest", CanonicalModelID: "anthropic/claude-3-5-haiku-latest", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()},
|
||||
{ModelID: "google/gemini-2.5-flash", CanonicalModelID: "google/gemini-2.5-flash", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()},
|
||||
{ModelID: "qwen/qwen3.7-plus", CanonicalModelID: "qwen/qwen3.7-plus", Pricing: store.UsagePricingFields{Input: &inputPrice, Output: &outputPrice}, SyncedAt: time.Now().UTC()},
|
||||
}
|
||||
if _, err := usageStore.UpsertPricingCatalog(context.Background(), entries); err != nil {
|
||||
t.Fatalf("UpsertPricingCatalog: %v", err)
|
||||
@@ -161,8 +197,10 @@ func TestPGUsageCapStoreResolvePricingUsesOpenRouterAliases(t *testing.T) {
|
||||
wantModelID string
|
||||
}{
|
||||
{name: "openai compat", providerName: "openai", providerType: store.ProviderOpenAICompat, modelID: "gpt-4o-mini", wantModelID: "openai/gpt-4o-mini"},
|
||||
{name: "openai compat custom gpt model", providerName: "cppai", providerType: store.ProviderOpenAICompat, modelID: "gpt-5.5", wantModelID: "openai/gpt-5.5"},
|
||||
{name: "anthropic native", providerName: "anthropic", providerType: store.ProviderAnthropicNative, modelID: "claude-3-5-haiku-latest", wantModelID: "anthropic/claude-3-5-haiku-latest"},
|
||||
{name: "gemini native", providerName: "gemini", providerType: store.ProviderGeminiNative, modelID: "gemini-2.5-flash", wantModelID: "google/gemini-2.5-flash"},
|
||||
{name: "bailian qwen model", providerName: "bailian", providerType: store.ProviderBailian, modelID: "qwen3.7-plus", wantModelID: "qwen/qwen3.7-plus"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func (s *PGUsageEventStore) BackfillUsageEventCosts(ctx context.Context) (store.UsageEventCostBackfillStats, error) {
|
||||
var stats store.UsageEventCostBackfillStats
|
||||
rows, err := s.db.QueryContext(ctx, `
|
||||
WITH linked_costs AS (
|
||||
SELECT
|
||||
e.id,
|
||||
e.bucket_hour,
|
||||
s.total_cost
|
||||
FROM usage_events e
|
||||
JOIN spans s
|
||||
ON s.id = e.span_id
|
||||
AND s.trace_id = e.trace_id
|
||||
AND s.tenant_id = e.tenant_id
|
||||
WHERE COALESCE(e.cost_usd, 0) = 0
|
||||
AND COALESCE(s.total_cost, 0) > 0
|
||||
)
|
||||
UPDATE usage_events e
|
||||
SET cost_usd = linked_costs.total_cost
|
||||
FROM linked_costs
|
||||
WHERE e.id = linked_costs.id
|
||||
AND e.cost_usd IS DISTINCT FROM linked_costs.total_cost
|
||||
RETURNING e.bucket_hour`)
|
||||
if err != nil {
|
||||
return stats, fmt.Errorf("backfill usage event costs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
bucketSeen := make(map[time.Time]struct{})
|
||||
for rows.Next() {
|
||||
var bucket time.Time
|
||||
if err := rows.Scan(&bucket); err != nil {
|
||||
return stats, fmt.Errorf("scan usage event cost backfill bucket: %w", err)
|
||||
}
|
||||
stats.EventRowsUpdated++
|
||||
bucket = bucket.UTC().Truncate(time.Hour)
|
||||
if _, ok := bucketSeen[bucket]; !ok {
|
||||
bucketSeen[bucket] = struct{}{}
|
||||
stats.RollupBuckets = append(stats.RollupBuckets, bucket)
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return stats, err
|
||||
}
|
||||
|
||||
for _, bucket := range stats.RollupBuckets {
|
||||
if err := s.RefreshEventRollupHour(ctx, bucket); err != nil {
|
||||
return stats, fmt.Errorf("refresh usage event rollup %s: %w", bucket.Format(time.RFC3339), err)
|
||||
}
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func TestPGUsageEventStoreBackfillUsageEventCosts(t *testing.T) {
|
||||
db := hooksTestDB(t)
|
||||
tenantID, agentID := seedTenantAndAgent(t, db)
|
||||
providerID := uuid.Must(uuid.NewV7())
|
||||
traceID := uuid.Must(uuid.NewV7())
|
||||
spanID := uuid.Must(uuid.NewV7())
|
||||
start := time.Date(2026, 7, 3, 2, 15, 0, 0, time.UTC)
|
||||
modelID := "gpt-usage-event-backfill-" + spanID.String()[:8]
|
||||
catalogModelID := "openai/" + modelID
|
||||
|
||||
t.Cleanup(func() {
|
||||
db.Exec(`DELETE FROM usage_event_rollups WHERE tenant_id=$1`, tenantID)
|
||||
db.Exec(`DELETE FROM usage_events WHERE tenant_id=$1`, tenantID)
|
||||
db.Exec(`DELETE FROM spans WHERE id=$1`, spanID)
|
||||
db.Exec(`DELETE FROM traces WHERE id=$1`, traceID)
|
||||
db.Exec(`DELETE FROM llm_providers WHERE id=$1`, providerID)
|
||||
db.Exec(`DELETE FROM usage_pricing_catalog WHERE model_id=$1`, catalogModelID)
|
||||
})
|
||||
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled)
|
||||
VALUES ($1,$2,'cppai','openai_compat','test',true)`,
|
||||
providerID, tenantID,
|
||||
); err != nil {
|
||||
t.Fatalf("insert provider: %v", err)
|
||||
}
|
||||
inputPrice := "0.000005"
|
||||
outputPrice := "0.000030"
|
||||
if _, err := NewPGUsageCapStore(db).UpsertPricingCatalog(context.Background(), []store.UsagePricingCatalogEntry{{
|
||||
ModelID: catalogModelID,
|
||||
CanonicalModelID: catalogModelID,
|
||||
Pricing: store.UsagePricingFields{
|
||||
Input: &inputPrice,
|
||||
Output: &outputPrice,
|
||||
},
|
||||
SyncedAt: time.Now().UTC(),
|
||||
}}); err != nil {
|
||||
t.Fatalf("upsert pricing catalog: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO traces (id, tenant_id, agent_id, start_time, created_at, status)
|
||||
VALUES ($1,$2,$3,$4,$4,'completed')`,
|
||||
traceID, tenantID, agentID, start,
|
||||
); err != nil {
|
||||
t.Fatalf("insert trace: %v", err)
|
||||
}
|
||||
if _, err := db.Exec(
|
||||
`INSERT INTO spans (
|
||||
id, trace_id, tenant_id, agent_id, span_type, start_time, created_at, status,
|
||||
name, provider, model, input_tokens, output_tokens, total_cost
|
||||
) VALUES ($1,$2,$3,$4,'tool_call',$5,$5,'completed','read_image','cppai',$6,1000,500,0)`,
|
||||
spanID, traceID, tenantID, agentID, start, modelID,
|
||||
); err != nil {
|
||||
t.Fatalf("insert span: %v", err)
|
||||
}
|
||||
eventStore := NewPGUsageEventStore(db)
|
||||
ctx := store.WithTenantID(context.Background(), tenantID)
|
||||
if err := eventStore.InsertEvent(ctx, &store.UsageEvent{
|
||||
TenantID: tenantID,
|
||||
EventTime: start,
|
||||
BucketHour: start.Truncate(time.Hour),
|
||||
EventType: store.UsageEventTypeToolCall,
|
||||
ResourceType: store.UsageResourceTypeTool,
|
||||
ResourceName: "read_image",
|
||||
ResourceID: "read_image",
|
||||
Source: store.UsageSourceToolCall,
|
||||
AgentID: &agentID,
|
||||
TraceID: &traceID,
|
||||
SpanID: &spanID,
|
||||
Provider: "cppai",
|
||||
Model: modelID,
|
||||
Status: "completed",
|
||||
InputTokens: 1000,
|
||||
OutputTokens: 500,
|
||||
TotalTokens: 1500,
|
||||
CallCount: 1,
|
||||
}); err != nil {
|
||||
t.Fatalf("insert usage event: %v", err)
|
||||
}
|
||||
|
||||
traceStats, err := NewPGTracingStore(db).BackfillLLMCosts(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("BackfillLLMCosts: %v", err)
|
||||
}
|
||||
if traceStats.SpanRowsUpdated == 0 {
|
||||
t.Fatalf("trace stats = %+v, want tool span cost updated", traceStats)
|
||||
}
|
||||
|
||||
eventStats, err := eventStore.BackfillUsageEventCosts(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("BackfillUsageEventCosts: %v", err)
|
||||
}
|
||||
if eventStats.EventRowsUpdated != 1 || len(eventStats.RollupBuckets) != 1 {
|
||||
t.Fatalf("event stats = %+v, want 1 event and 1 rollup bucket", eventStats)
|
||||
}
|
||||
|
||||
var eventCost float64
|
||||
if err := db.QueryRow(`SELECT cost_usd FROM usage_events WHERE span_id=$1`, spanID).Scan(&eventCost); err != nil {
|
||||
t.Fatalf("query usage event cost: %v", err)
|
||||
}
|
||||
if math.Abs(eventCost-0.02) > 0.0000001 {
|
||||
t.Fatalf("event cost = %.8f, want 0.02000000", eventCost)
|
||||
}
|
||||
summary, err := eventStore.GetEventSummary(ctx, store.UsageEventQuery{
|
||||
From: start.Add(-time.Hour),
|
||||
To: start.Add(time.Hour),
|
||||
ResourceType: store.UsageResourceTypeTool,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("GetEventSummary: %v", err)
|
||||
}
|
||||
if math.Abs(summary.CostUSD-0.02) > 0.0000001 {
|
||||
t.Fatalf("summary cost = %.8f, want 0.02000000", summary.CostUSD)
|
||||
}
|
||||
}
|
||||
@@ -43,12 +43,13 @@ ON CONFLICT (model_id) DO UPDATE SET
|
||||
web_search_price = EXCLUDED.web_search_price,
|
||||
synced_at = EXCLUDED.synced_at,
|
||||
updated_at = now()`
|
||||
upserted := 0
|
||||
for _, e := range entries {
|
||||
if strings.TrimSpace(e.ModelID) == "" {
|
||||
continue
|
||||
}
|
||||
if err := validateUsagePricingFields(e.Pricing); err != nil {
|
||||
return 0, err
|
||||
continue
|
||||
}
|
||||
if len(e.RawPricing) == 0 {
|
||||
e.RawPricing = json.RawMessage(`{}`)
|
||||
@@ -65,8 +66,9 @@ ON CONFLICT (model_id) DO UPDATE SET
|
||||
); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
upserted++
|
||||
}
|
||||
return len(entries), nil
|
||||
return upserted, nil
|
||||
}
|
||||
|
||||
func (s *PGUsageCapStore) ListPricingCatalog(ctx context.Context, q store.UsagePricingQuery) ([]store.UsagePricingCatalogEntry, error) {
|
||||
@@ -216,6 +218,9 @@ func usagePricingModelCandidates(providerName, providerType, modelID string) []s
|
||||
for _, prefix := range openRouterProviderPrefixes(providerName, providerType) {
|
||||
out = appendUniqueString(out, prefix+"/"+modelID)
|
||||
}
|
||||
for _, prefix := range openRouterModelPrefixes(modelID) {
|
||||
out = appendUniqueString(out, prefix+"/"+modelID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -251,13 +256,33 @@ func openRouterProviderPrefixes(providerName, providerType string) []string {
|
||||
return []string{"cohere"}
|
||||
case store.ProviderPerplexity:
|
||||
return []string{"perplexity"}
|
||||
case store.ProviderDashScope:
|
||||
case store.ProviderDashScope, store.ProviderBailian:
|
||||
return []string{"qwen"}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func openRouterModelPrefixes(modelID string) []string {
|
||||
modelID = strings.ToLower(strings.TrimSpace(modelID))
|
||||
switch {
|
||||
case strings.HasPrefix(modelID, "gpt-"),
|
||||
strings.HasPrefix(modelID, "o1"),
|
||||
strings.HasPrefix(modelID, "o3"),
|
||||
strings.HasPrefix(modelID, "o4"),
|
||||
strings.HasPrefix(modelID, "o5"):
|
||||
return []string{"openai"}
|
||||
case strings.HasPrefix(modelID, "qwen"):
|
||||
return []string{"qwen"}
|
||||
case strings.HasPrefix(modelID, "claude-"):
|
||||
return []string{"anthropic"}
|
||||
case strings.HasPrefix(modelID, "gemini-"):
|
||||
return []string{"google"}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeProviderAlias(s string) string {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
s = strings.ReplaceAll(s, " ", "-")
|
||||
|
||||
@@ -42,7 +42,7 @@ type SnapshotQuery struct {
|
||||
Provider string // optional: cross-filter by provider
|
||||
Model string // optional: cross-filter by model
|
||||
Channel string // optional: cross-filter by channel
|
||||
GroupBy string // "hour" (default), "day", "provider", "model", "channel", "agent"
|
||||
GroupBy string // "hour" (default), "day", "provider", "model", "provider_model", "channel", "agent"
|
||||
}
|
||||
|
||||
// SnapshotTimeSeries is a single point in a time series response.
|
||||
@@ -81,6 +81,10 @@ type SnapshotBreakdown struct {
|
||||
AvgDurationMS int `json:"avg_duration_ms" db:"avg_duration_ms"`
|
||||
}
|
||||
|
||||
type SnapshotCostBackfillStats struct {
|
||||
SnapshotRowsUpdated int64
|
||||
}
|
||||
|
||||
// SnapshotStore manages pre-computed usage snapshots.
|
||||
type SnapshotStore interface {
|
||||
// UpsertSnapshots inserts or updates (on conflict, replace) a batch of snapshots.
|
||||
@@ -89,7 +93,7 @@ type SnapshotStore interface {
|
||||
// GetTimeSeries returns hourly (or daily) aggregated time series.
|
||||
GetTimeSeries(ctx context.Context, q SnapshotQuery) ([]SnapshotTimeSeries, error)
|
||||
|
||||
// GetBreakdown returns aggregated data grouped by a dimension (provider, model, channel, agent).
|
||||
// GetBreakdown returns aggregated data grouped by a dimension (provider, model, provider_model, channel, agent).
|
||||
GetBreakdown(ctx context.Context, q SnapshotQuery) ([]SnapshotBreakdown, error)
|
||||
|
||||
// GetLatestBucket returns the most recent bucket_hour, used by worker to know where to resume.
|
||||
|
||||
@@ -184,6 +184,10 @@ func (s *SQLiteSnapshotStore) GetBreakdown(ctx context.Context, q store.Snapshot
|
||||
groupCol = "model"
|
||||
orderExpr = "SUM(CASE WHEN provider != '' THEN input_tokens ELSE 0 END) DESC"
|
||||
extraFilter = " AND provider != '' AND model != ''"
|
||||
case "provider_model":
|
||||
groupCol = "provider || '/' || model"
|
||||
orderExpr = "SUM(CASE WHEN provider != '' THEN input_tokens ELSE 0 END) DESC"
|
||||
extraFilter = " AND provider != '' AND model != ''"
|
||||
case "channel":
|
||||
groupCol = "channel"
|
||||
orderExpr = "SUM(CASE WHEN provider = '' AND model = '' THEN request_count ELSE 0 END) DESC"
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
//go:build sqlite || sqliteonly
|
||||
|
||||
package sqlitestore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func TestSQLiteSnapshotStoreGetBreakdownProviderModelKeepsProviderAndModel(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
if err := EnsureSchema(db); err != nil {
|
||||
t.Fatalf("EnsureSchema: %v", err)
|
||||
}
|
||||
snapshots := NewSQLiteSnapshotStore(db)
|
||||
ctx := store.WithTenantID(context.Background(), store.MasterTenantID)
|
||||
bucket := time.Date(2026, 7, 2, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
if err := snapshots.UpsertSnapshots(ctx, []store.UsageSnapshot{
|
||||
{
|
||||
BucketHour: bucket,
|
||||
Provider: "openrouter",
|
||||
Model: "openai/gpt-5.5",
|
||||
InputTokens: 1200,
|
||||
OutputTokens: 300,
|
||||
LLMCallCount: 2,
|
||||
TotalCost: 0.42,
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("UpsertSnapshots: %v", err)
|
||||
}
|
||||
|
||||
rows, err := snapshots.GetBreakdown(ctx, store.SnapshotQuery{
|
||||
From: bucket.Add(-time.Hour),
|
||||
To: bucket.Add(time.Hour),
|
||||
GroupBy: "provider_model",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("GetBreakdown: %v", err)
|
||||
}
|
||||
if len(rows) != 1 {
|
||||
t.Fatalf("len(rows) = %d, want 1: %#v", len(rows), rows)
|
||||
}
|
||||
if got, want := rows[0].Key, "openrouter/openai/gpt-5.5"; got != want {
|
||||
t.Fatalf("key = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := rows[0].TotalCost, 0.42; got != want {
|
||||
t.Fatalf("total_cost = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
@@ -245,6 +245,64 @@ func (s *SQLiteTracingStore) ListChildTraces(ctx context.Context, parentTraceID
|
||||
return scanTraceRows(rows)
|
||||
}
|
||||
|
||||
func (s *SQLiteTracingStore) 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
|
||||
}
|
||||
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(keys)), ",")
|
||||
q := `SELECT session_key, COALESCE(SUM(total_cost), 0)
|
||||
FROM traces
|
||||
WHERE session_key IN (` + placeholders + `)
|
||||
AND parent_trace_id IS NULL`
|
||||
args := make([]any, 0, len(keys)+1)
|
||||
for _, key := range keys {
|
||||
args = append(args, key)
|
||||
}
|
||||
if !store.IsCrossTenant(ctx) {
|
||||
tid := store.TenantIDFromContext(ctx)
|
||||
if tid == uuid.Nil {
|
||||
return result, nil
|
||||
}
|
||||
q += ` AND tenant_id = ?`
|
||||
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 *SQLiteTracingStore) 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)
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
//go:build sqlite || sqliteonly
|
||||
|
||||
package sqlitestore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func TestSQLiteTracingStoreGetSessionCostsTenantScopedRootTraces(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
if err := EnsureSchema(db); err != nil {
|
||||
t.Fatalf("EnsureSchema: %v", err)
|
||||
}
|
||||
traces := NewSQLiteTracingStore(db)
|
||||
tenantA := uuid.Must(uuid.NewV7())
|
||||
tenantB := uuid.Must(uuid.NewV7())
|
||||
insertSessionCostTenant(t, db, tenantA, "tenant-a")
|
||||
insertSessionCostTenant(t, db, tenantB, "tenant-b")
|
||||
ctxA := store.WithTenantID(context.Background(), tenantA)
|
||||
ctxB := store.WithTenantID(context.Background(), tenantB)
|
||||
now := time.Date(2026, 7, 2, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
rootA1 := createSessionCostTrace(t, traces, ctxA, "session-a", nil, 1.25, now)
|
||||
createSessionCostTrace(t, traces, ctxA, "session-a", nil, 0.75, now.Add(time.Second))
|
||||
createSessionCostTrace(t, traces, ctxA, "session-a", &rootA1, 99, now.Add(2*time.Second))
|
||||
createSessionCostTrace(t, traces, ctxB, "session-a", nil, 8, now.Add(3*time.Second))
|
||||
|
||||
costs, err := traces.GetSessionCosts(ctxA, []string{"session-a", "session-a", " ", "missing"})
|
||||
if err != nil {
|
||||
t.Fatalf("GetSessionCosts: %v", err)
|
||||
}
|
||||
if got, want := costs["session-a"], 2.0; got != want {
|
||||
t.Fatalf("session-a cost = %v, want %v", got, want)
|
||||
}
|
||||
if _, ok := costs["missing"]; ok {
|
||||
t.Fatalf("missing session unexpectedly present: %#v", costs)
|
||||
}
|
||||
|
||||
empty, err := traces.GetSessionCosts(context.Background(), []string{"session-a"})
|
||||
if err != nil {
|
||||
t.Fatalf("GetSessionCosts without tenant: %v", err)
|
||||
}
|
||||
if len(empty) != 0 {
|
||||
t.Fatalf("nil-tenant result = %#v, want empty", empty)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSQLiteTracingStoreBatchUpdateTraceAggregatesIncludesToolUsageTokens(t *testing.T) {
|
||||
db := openTestDB(t)
|
||||
if err := EnsureSchema(db); err != nil {
|
||||
t.Fatalf("EnsureSchema: %v", err)
|
||||
}
|
||||
traces := NewSQLiteTracingStore(db)
|
||||
tenantID := uuid.Must(uuid.NewV7())
|
||||
insertSessionCostTenant(t, db, tenantID, "tenant-trace-aggregate")
|
||||
ctx := store.WithTenantID(context.Background(), tenantID)
|
||||
now := time.Date(2026, 7, 2, 13, 0, 0, 0, time.UTC)
|
||||
traceID := createSessionCostTrace(t, traces, ctx, "session-a", nil, 0, now)
|
||||
|
||||
llmSpanID := uuid.Must(uuid.NewV7())
|
||||
toolSpanID := uuid.Must(uuid.NewV7())
|
||||
llmCost := 0.25
|
||||
toolCost := 0.05
|
||||
if err := traces.CreateSpan(ctx, &store.SpanData{
|
||||
ID: llmSpanID,
|
||||
TraceID: traceID,
|
||||
SpanType: store.SpanTypeLLMCall,
|
||||
StartTime: now,
|
||||
Status: store.SpanStatusCompleted,
|
||||
Provider: "bailian",
|
||||
Model: "qwen3.7-plus",
|
||||
InputTokens: 1000,
|
||||
OutputTokens: 200,
|
||||
TotalCost: &llmCost,
|
||||
CreatedAt: now,
|
||||
}); err != nil {
|
||||
t.Fatalf("CreateSpan llm: %v", err)
|
||||
}
|
||||
if err := traces.CreateSpan(ctx, &store.SpanData{
|
||||
ID: toolSpanID,
|
||||
TraceID: traceID,
|
||||
SpanType: store.SpanTypeToolCall,
|
||||
Name: "read_image",
|
||||
StartTime: now,
|
||||
Status: store.SpanStatusCompleted,
|
||||
Provider: "bailian",
|
||||
Model: "qwen3.7-plus",
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
TotalCost: &toolCost,
|
||||
CreatedAt: now,
|
||||
}); err != nil {
|
||||
t.Fatalf("CreateSpan tool: %v", err)
|
||||
}
|
||||
if err := traces.UpdateSpan(ctx, llmSpanID, map[string]any{"total_cost": llmCost}); err != nil {
|
||||
t.Fatalf("UpdateSpan llm: %v", err)
|
||||
}
|
||||
if err := traces.UpdateSpan(ctx, toolSpanID, map[string]any{"total_cost": toolCost}); err != nil {
|
||||
t.Fatalf("UpdateSpan tool: %v", err)
|
||||
}
|
||||
|
||||
if err := traces.BatchUpdateTraceAggregates(ctx, traceID); err != nil {
|
||||
t.Fatalf("BatchUpdateTraceAggregates: %v", err)
|
||||
}
|
||||
got, err := traces.GetTrace(ctx, traceID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetTrace: %v", err)
|
||||
}
|
||||
if got.TotalInputTokens != 1100 || got.TotalOutputTokens != 250 || math.Abs(got.TotalCost-0.30) > 0.0000001 || got.LLMCallCount != 1 || got.ToolCallCount != 1 {
|
||||
t.Fatalf("trace = input %d output %d cost %.2f llm %d tool %d, want 1100/250/0.30/1/1",
|
||||
got.TotalInputTokens, got.TotalOutputTokens, got.TotalCost, got.LLMCallCount, got.ToolCallCount)
|
||||
}
|
||||
}
|
||||
|
||||
func insertSessionCostTenant(t *testing.T, db *sql.DB, id uuid.UUID, slug string) {
|
||||
t.Helper()
|
||||
_, err := db.Exec(
|
||||
`INSERT INTO tenants (id, name, slug, status) VALUES (?, ?, ?, 'active')`,
|
||||
id.String(), slug, slug,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("insert tenant: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func createSessionCostTrace(t *testing.T, traces *SQLiteTracingStore, ctx context.Context, sessionKey string, parentID *uuid.UUID, cost float64, ts time.Time) uuid.UUID {
|
||||
t.Helper()
|
||||
id := uuid.Must(uuid.NewV7())
|
||||
err := traces.CreateTrace(ctx, &store.TraceData{
|
||||
ID: id,
|
||||
ParentTraceID: parentID,
|
||||
SessionKey: sessionKey,
|
||||
StartTime: ts,
|
||||
Status: store.TraceStatusCompleted,
|
||||
TotalCost: cost,
|
||||
CreatedAt: ts,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateTrace: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
@@ -191,18 +191,17 @@ func (s *SQLiteTracingStore) BatchUpdateTraceAggregates(ctx context.Context, tra
|
||||
span_count = (SELECT COUNT(*) FROM spans WHERE trace_id = ?),
|
||||
llm_call_count = (SELECT COUNT(*) FROM spans WHERE trace_id = ? AND span_type = 'llm_call'),
|
||||
tool_call_count = (SELECT COUNT(*) FROM spans WHERE trace_id = ? AND span_type = 'tool_call'),
|
||||
total_input_tokens = COALESCE((SELECT SUM(input_tokens) FROM spans WHERE trace_id = ? AND span_type = 'llm_call' AND input_tokens IS NOT NULL), 0),
|
||||
total_output_tokens = COALESCE((SELECT SUM(output_tokens) FROM spans WHERE trace_id = ? AND span_type = 'llm_call' AND output_tokens IS NOT NULL), 0),
|
||||
total_input_tokens = COALESCE((SELECT SUM(input_tokens) FROM spans WHERE trace_id = ? 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 = ? 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 = ? AND total_cost IS NOT NULL), 0),
|
||||
metadata = (
|
||||
SELECT json_object(
|
||||
'total_cache_read_tokens', COALESCE(SUM(json_extract(metadata, '$.cache_read_tokens')), 0),
|
||||
'total_cache_creation_tokens', COALESCE(SUM(json_extract(metadata, '$.cache_creation_tokens')), 0)
|
||||
)
|
||||
FROM spans WHERE trace_id = ? AND span_type = 'llm_call' AND metadata IS NOT NULL
|
||||
FROM spans WHERE trace_id = ? AND span_type IN ('llm_call', 'tool_call') AND metadata IS NOT NULL
|
||||
)
|
||||
WHERE id = ?`,
|
||||
traceID, traceID, traceID, traceID, traceID, traceID, traceID, traceID)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -138,6 +138,16 @@ type CostSummaryRow struct {
|
||||
TraceCount int `json:"trace_count" db:"trace_count"`
|
||||
}
|
||||
|
||||
type TraceCostBackfillStats struct {
|
||||
SpanRowsUpdated int64
|
||||
TraceRowsUpdated int64
|
||||
SnapshotBuckets []time.Time
|
||||
}
|
||||
|
||||
type TraceUsageAggregateStats struct {
|
||||
TraceRowsUpdated int64
|
||||
}
|
||||
|
||||
// CodexPoolSpan holds the fields from a single LLM span for Codex pool activity analysis.
|
||||
type CodexPoolSpan struct {
|
||||
SpanID uuid.UUID
|
||||
|
||||
@@ -149,6 +149,11 @@ type UsageEventRollup struct {
|
||||
UpdatedAt time.Time `json:"updated_at" db:"updated_at"`
|
||||
}
|
||||
|
||||
type UsageEventCostBackfillStats struct {
|
||||
EventRowsUpdated int64
|
||||
RollupBuckets []time.Time
|
||||
}
|
||||
|
||||
type UsageEventStore interface {
|
||||
InsertEvent(ctx context.Context, event *UsageEvent) error
|
||||
InsertEvents(ctx context.Context, events []UsageEvent) error
|
||||
|
||||
@@ -3,6 +3,8 @@ package tracing
|
||||
import (
|
||||
"github.com/nextlevelbuilder/goclaw/internal/config"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
usagepricing "github.com/nextlevelbuilder/goclaw/internal/usage/pricing"
|
||||
)
|
||||
|
||||
// CalculateCost computes the USD cost for a single LLM call based on token usage and pricing.
|
||||
@@ -48,6 +50,21 @@ func CalculateCost(pricing *config.ModelPricing, usage *providers.Usage) float64
|
||||
return cost
|
||||
}
|
||||
|
||||
func CalculateCostFromUsagePricing(fields store.UsagePricingFields, usage *providers.Usage) (float64, error) {
|
||||
if usage == nil {
|
||||
return 0, nil
|
||||
}
|
||||
billable := usagepricing.FromProviderUsage(usage)
|
||||
if fields.Request == nil {
|
||||
billable.RequestCount = 0
|
||||
}
|
||||
micros, err := usagepricing.CostMicros(fields, billable)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return float64(micros) / 1_000_000, nil
|
||||
}
|
||||
|
||||
// LookupPricing finds the model pricing from config.
|
||||
// Tries "provider/model" first, then just "model".
|
||||
func LookupPricing(pricingMap map[string]*config.ModelPricing, provider, model string) *config.ModelPricing {
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/config"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
func floatEquals(a, b float64) bool {
|
||||
@@ -44,8 +45,8 @@ func TestCalculateCost_CacheTokens(t *testing.T) {
|
||||
pricing := &config.ModelPricing{
|
||||
InputPerMillion: 3.0,
|
||||
OutputPerMillion: 15.0,
|
||||
CacheReadPerMillion: 0.3, // 10% of input
|
||||
CacheCreatePerMillion: 3.75, // 25% premium
|
||||
CacheReadPerMillion: 0.3, // 10% of input
|
||||
CacheCreatePerMillion: 3.75, // 25% premium
|
||||
}
|
||||
usage := &providers.Usage{
|
||||
PromptTokens: 1_000_000,
|
||||
@@ -194,3 +195,41 @@ func TestLookupPricing_ProviderQualified(t *testing.T) {
|
||||
t.Errorf("expected nil for nil map, got %+v", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculateCostFromUsagePricingSeparatesIncludedCacheTokens(t *testing.T) {
|
||||
input := "0.000001"
|
||||
cacheRead := "0.0000001"
|
||||
got, err := CalculateCostFromUsagePricing(store.UsagePricingFields{
|
||||
Input: &input,
|
||||
CacheRead: &cacheRead,
|
||||
}, &providers.Usage{
|
||||
PromptTokens: 2006,
|
||||
CacheReadTokens: 1920,
|
||||
PromptTokensIncludeCachedSegments: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CalculateCostFromUsagePricing: %v", err)
|
||||
}
|
||||
if !floatEquals(got, 0.000278) {
|
||||
t.Fatalf("cost = %.9f, want 0.000278", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculateCostFromUsagePricingIgnoresUnpricedRequestTelemetry(t *testing.T) {
|
||||
input := "0.000001"
|
||||
output := "0.000002"
|
||||
got, err := CalculateCostFromUsagePricing(store.UsagePricingFields{
|
||||
Input: &input,
|
||||
Output: &output,
|
||||
}, &providers.Usage{
|
||||
PromptTokens: 1000,
|
||||
CompletionTokens: 500,
|
||||
RequestCount: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CalculateCostFromUsagePricing: %v", err)
|
||||
}
|
||||
if !floatEquals(got, 0.002) {
|
||||
t.Fatalf("cost = %.9f, want 0.002", got)
|
||||
}
|
||||
}
|
||||
@@ -21,6 +21,8 @@ type SnapshotWorker struct {
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
const recentUsageRefreshWindow = 2 * time.Hour
|
||||
|
||||
func NewSnapshotWorker(db *sql.DB, snapshots store.SnapshotStore, usageEvents store.UsageEventStore) *SnapshotWorker {
|
||||
return &SnapshotWorker{
|
||||
db: db,
|
||||
@@ -72,7 +74,7 @@ func (w *SnapshotWorker) loop() {
|
||||
}
|
||||
}
|
||||
|
||||
// catchUp computes snapshots for all missed hours between latest bucket and current hour.
|
||||
// catchUp computes missing buckets and refreshes recent closed buckets that may have late writes.
|
||||
func (w *SnapshotWorker) catchUp() {
|
||||
ctx := context.Background()
|
||||
w.catchUpSnapshots(ctx)
|
||||
@@ -89,13 +91,7 @@ func (w *SnapshotWorker) catchUpSnapshots(ctx context.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
var startHour time.Time
|
||||
if latest == nil {
|
||||
// No snapshots yet — only compute the previous hour (backfill handles history)
|
||||
startHour = targetHour
|
||||
} else {
|
||||
startHour = latest.Add(time.Hour)
|
||||
}
|
||||
startHour := usageCatchUpStartHour(latest, targetHour)
|
||||
|
||||
for h := startHour; !h.After(targetHour); h = h.Add(time.Hour) {
|
||||
start := time.Now()
|
||||
@@ -120,10 +116,7 @@ func (w *SnapshotWorker) catchUpUsageEvents(ctx context.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
startHour := targetHour
|
||||
if latest != nil {
|
||||
startHour = latest.Add(time.Hour)
|
||||
}
|
||||
startHour := usageCatchUpStartHour(latest, targetHour)
|
||||
|
||||
for h := startHour; !h.After(targetHour); h = h.Add(time.Hour) {
|
||||
start := time.Now()
|
||||
@@ -135,6 +128,21 @@ func (w *SnapshotWorker) catchUpUsageEvents(ctx context.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
func usageCatchUpStartHour(latest *time.Time, targetHour time.Time) time.Time {
|
||||
targetHour = targetHour.UTC().Truncate(time.Hour)
|
||||
if latest == nil {
|
||||
// Historical backfill handles older data; steady-state startup should refresh the last closed bucket.
|
||||
return targetHour
|
||||
}
|
||||
|
||||
nextMissing := latest.UTC().Truncate(time.Hour).Add(time.Hour)
|
||||
refreshStart := targetHour.Add(-recentUsageRefreshWindow)
|
||||
if nextMissing.Before(refreshStart) {
|
||||
return nextMissing
|
||||
}
|
||||
return refreshStart
|
||||
}
|
||||
|
||||
// Backfill populates usage_snapshots from historical trace/span data.
|
||||
// Returns the number of hours processed.
|
||||
func (w *SnapshotWorker) Backfill(ctx context.Context) (int, error) {
|
||||
@@ -173,6 +181,23 @@ func (w *SnapshotWorker) Backfill(ctx context.Context) (int, error) {
|
||||
return count + eventCount, nil
|
||||
}
|
||||
|
||||
func (w *SnapshotWorker) RefreshBuckets(ctx context.Context, buckets []time.Time) (int, error) {
|
||||
seen := make(map[time.Time]struct{}, len(buckets))
|
||||
count := 0
|
||||
for _, bucket := range buckets {
|
||||
hour := bucket.UTC().Truncate(time.Hour)
|
||||
if _, ok := seen[hour]; ok {
|
||||
continue
|
||||
}
|
||||
seen[hour] = struct{}{}
|
||||
if err := w.aggregateHour(ctx, hour); err != nil {
|
||||
return count, fmt.Errorf("refresh snapshot bucket %s: %w", hour.Format(time.RFC3339), err)
|
||||
}
|
||||
count++
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (w *SnapshotWorker) backfillUsageEvents(ctx context.Context) (int, error) {
|
||||
if w.usageEvents == nil {
|
||||
return 0, nil
|
||||
@@ -315,7 +340,7 @@ func querySpanAggregates(ctx context.Context, db *sql.DB, from, to time.Time) ([
|
||||
COALESCE(t.channel, '') as channel,
|
||||
COALESCE(s.provider, '') as provider,
|
||||
COALESCE(s.model, '') as model,
|
||||
COUNT(*) as llm_call_count,
|
||||
COUNT(*) FILTER (WHERE s.span_type = 'llm_call') as llm_call_count,
|
||||
COALESCE(SUM(s.input_tokens), 0) as span_input_tokens,
|
||||
COALESCE(SUM(s.output_tokens), 0) as span_output_tokens,
|
||||
COALESCE(SUM(s.total_cost), 0) as span_cost,
|
||||
@@ -323,7 +348,7 @@ func querySpanAggregates(ctx context.Context, db *sql.DB, from, to time.Time) ([
|
||||
COALESCE(SUM(CAST(s.metadata->>'cache_creation_tokens' AS INTEGER)), 0) as cache_create_tokens,
|
||||
COALESCE(SUM(CAST(s.metadata->>'thinking_tokens' AS INTEGER)), 0) as thinking_tokens
|
||||
FROM traces t
|
||||
JOIN spans s ON s.trace_id = t.id AND s.span_type = 'llm_call'
|
||||
JOIN spans s ON s.trace_id = t.id AND s.span_type IN ('llm_call', 'tool_call')
|
||||
WHERE t.start_time >= $1 AND t.start_time < $2
|
||||
AND t.parent_trace_id IS NULL
|
||||
GROUP BY t.agent_id, t.channel, s.provider, s.model`, from, to)
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
package tracing
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestUsageCatchUpStartHourRefreshesLatestClosedBucket(t *testing.T) {
|
||||
target := time.Date(2026, 7, 3, 1, 0, 0, 0, time.UTC)
|
||||
latest := target
|
||||
|
||||
got := usageCatchUpStartHour(&latest, target)
|
||||
want := target.Add(-recentUsageRefreshWindow)
|
||||
if !got.Equal(want) {
|
||||
t.Fatalf("start hour = %s, want %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageCatchUpStartHourContinuesLargeBacklog(t *testing.T) {
|
||||
target := time.Date(2026, 7, 3, 6, 0, 0, 0, time.UTC)
|
||||
latest := target.Add(-12 * time.Hour)
|
||||
|
||||
got := usageCatchUpStartHour(&latest, target)
|
||||
want := latest.Add(time.Hour)
|
||||
if !got.Equal(want) {
|
||||
t.Fatalf("start hour = %s, want %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageCatchUpStartHourWithoutSnapshotsComputesPreviousHourOnly(t *testing.T) {
|
||||
target := time.Date(2026, 7, 3, 1, 0, 0, 0, time.UTC)
|
||||
|
||||
got := usageCatchUpStartHour(nil, target)
|
||||
if !got.Equal(target) {
|
||||
t.Fatalf("start hour = %s, want %s", got, target)
|
||||
}
|
||||
}
|
||||
@@ -163,6 +163,18 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Service) ResolvePricing(ctx context.Context, tenantID uuid.UUID, providerName, modelID string) (*store.ResolvedUsagePricing, error) {
|
||||
if s == nil || s.store == nil {
|
||||
return nil, sql.ErrNoRows
|
||||
}
|
||||
ctx = scopedRequestContext(ctx, Request{TenantID: tenantID, ProviderName: providerName, ModelID: modelID})
|
||||
providerData, err := s.resolveProvider(ctx, tenantID, providerName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.store.ResolvePricing(ctx, tenantID, providerData.ID, providerData.Name, providerData.ProviderType, modelID)
|
||||
}
|
||||
|
||||
func (r *Reservation) Reconcile(ctx context.Context, resp *providers.ChatResponse, callErr error) {
|
||||
r.reconcile(ctx, resp, callErr, false)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
const DefaultOpenRouterCatalogSyncInterval = 24 * time.Hour
|
||||
|
||||
func StartOpenRouterCatalogAutoSync(ctx context.Context, s store.UsageCapStore, interval time.Duration, onSync ...func(context.Context, int)) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if interval <= 0 {
|
||||
interval = DefaultOpenRouterCatalogSyncInterval
|
||||
}
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
go func() {
|
||||
syncOpenRouterCatalog(ctx, s, client, firstOpenRouterSyncHook(onSync))
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
syncOpenRouterCatalog(ctx, s, client, firstOpenRouterSyncHook(onSync))
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func firstOpenRouterSyncHook(hooks []func(context.Context, int)) func(context.Context, int) {
|
||||
if len(hooks) == 0 {
|
||||
return nil
|
||||
}
|
||||
return hooks[0]
|
||||
}
|
||||
|
||||
func syncOpenRouterCatalog(ctx context.Context, s store.UsageCapStore, client *http.Client, onSync func(context.Context, int)) {
|
||||
syncCtx, cancel := context.WithTimeout(ctx, 45*time.Second)
|
||||
defer cancel()
|
||||
entries, err := FetchOpenRouterCatalog(syncCtx, client)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.openrouter_auto_sync_failed", "error", err)
|
||||
return
|
||||
}
|
||||
count, err := s.UpsertPricingCatalog(syncCtx, entries)
|
||||
if err != nil {
|
||||
slog.Warn("usage_pricing.openrouter_auto_store_failed", "error", err)
|
||||
return
|
||||
}
|
||||
slog.Info("usage_pricing.openrouter_auto_synced", "count", count)
|
||||
if onSync != nil {
|
||||
onSync(ctx, count)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFetchOpenRouterCatalogMapsPricingFields(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.Method != http.MethodGet {
|
||||
t.Fatalf("method = %s, want GET", req.Method)
|
||||
}
|
||||
if req.URL.String() != OpenRouterModelsURL {
|
||||
t.Fatalf("url = %s, want %s", req.URL.String(), OpenRouterModelsURL)
|
||||
}
|
||||
body := `{
|
||||
"data": [
|
||||
{
|
||||
"id": "openai/gpt-4o-mini",
|
||||
"name": "GPT-4o mini",
|
||||
"pricing": {
|
||||
"prompt": "0.00000015",
|
||||
"completion": "0.0000006",
|
||||
"input_cache_read": "0.000000075",
|
||||
"input_cache_write": "0.0000003",
|
||||
"internal_reasoning": "0.0000002",
|
||||
"request": "0",
|
||||
"image": "0.001",
|
||||
"web_search": "0.01"
|
||||
}
|
||||
},
|
||||
{"name": "missing id"}
|
||||
]
|
||||
}`
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})}
|
||||
|
||||
entries, err := FetchOpenRouterCatalog(context.Background(), client)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchOpenRouterCatalog: %v", err)
|
||||
}
|
||||
if len(entries) != 1 {
|
||||
t.Fatalf("entries = %d, want 1", len(entries))
|
||||
}
|
||||
got := entries[0]
|
||||
if got.ModelID != "openai/gpt-4o-mini" || got.CanonicalModelID != got.ModelID {
|
||||
t.Fatalf("model ids = %q/%q, want openai/gpt-4o-mini", got.ModelID, got.CanonicalModelID)
|
||||
}
|
||||
assertStringPtr(t, "input", got.Pricing.Input, "0.00000015")
|
||||
assertStringPtr(t, "output", got.Pricing.Output, "0.0000006")
|
||||
assertStringPtr(t, "cache_read", got.Pricing.CacheRead, "0.000000075")
|
||||
assertStringPtr(t, "cache_write", got.Pricing.CacheWrite, "0.0000003")
|
||||
assertStringPtr(t, "reasoning", got.Pricing.Reasoning, "0.0000002")
|
||||
assertStringPtr(t, "request", got.Pricing.Request, "0")
|
||||
assertStringPtr(t, "image", got.Pricing.Image, "0.001")
|
||||
assertStringPtr(t, "web_search", got.Pricing.WebSearch, "0.01")
|
||||
if len(got.RawPricing) == 0 || len(got.RawModel) == 0 {
|
||||
t.Fatal("expected raw pricing and raw model payloads")
|
||||
}
|
||||
if got.SyncedAt.IsZero() {
|
||||
t.Fatal("SyncedAt is zero")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchOpenRouterCatalogReturnsStatusError(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
Body: io.NopCloser(strings.NewReader(`bad gateway`)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})}
|
||||
|
||||
if _, err := FetchOpenRouterCatalog(context.Background(), client); err == nil {
|
||||
t.Fatal("expected non-2xx status error")
|
||||
}
|
||||
}
|
||||
|
||||
func assertStringPtr(t *testing.T, name string, got *string, want string) {
|
||||
t.Helper()
|
||||
if got == nil {
|
||||
t.Fatalf("%s = nil, want %q", name, want)
|
||||
}
|
||||
if *got != want {
|
||||
t.Fatalf("%s = %q, want %q", name, *got, want)
|
||||
}
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
@@ -47,9 +47,7 @@ func (h *spaHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.HasPrefix(r.URL.Path, "/assets/") {
|
||||
w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
|
||||
} else {
|
||||
// index.html and other root files must never be cached so the
|
||||
// browser always gets the latest asset manifest after a deploy.
|
||||
w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
|
||||
setIndexCacheHeaders(w)
|
||||
}
|
||||
h.fileServer.ServeHTTP(w, r)
|
||||
return
|
||||
@@ -57,6 +55,13 @@ func (h *spaHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// SPA fallback: serve index.html for any unmatched route.
|
||||
// This handles client-side routing (React Router).
|
||||
setIndexCacheHeaders(w)
|
||||
r.URL.Path = "/"
|
||||
h.fileServer.ServeHTTP(w, r)
|
||||
}
|
||||
|
||||
func setIndexCacheHeaders(w http.ResponseWriter) {
|
||||
// index.html and SPA fallback routes must never be cached so the browser
|
||||
// always fetches the latest Vite asset manifest after an embedded UI rebuild.
|
||||
w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { formatApiCost } from "./format";
|
||||
|
||||
describe("formatApiCost", () => {
|
||||
it("formats API costs with two decimal places", () => {
|
||||
expect(formatApiCost(1.024)).toBe("$1.02");
|
||||
expect(formatApiCost(5.934)).toBe("$5.93");
|
||||
expect(formatApiCost(304.451178)).toBe("$304.45");
|
||||
});
|
||||
|
||||
it("formats empty or zero API costs as zero dollars", () => {
|
||||
expect(formatApiCost(undefined)).toBe("$0.00");
|
||||
expect(formatApiCost(null)).toBe("$0.00");
|
||||
expect(formatApiCost(0)).toBe("$0.00");
|
||||
});
|
||||
});
|
||||
@@ -40,6 +40,14 @@ export function formatCost(cost: number | null | undefined): string {
|
||||
return `$${cost.toFixed(2)}`;
|
||||
}
|
||||
|
||||
export function formatApiCost(cost: number | null | undefined): string {
|
||||
if (cost == null || cost === 0) return "$0.00";
|
||||
return `$${cost.toLocaleString("en-US", {
|
||||
minimumFractionDigits: 2,
|
||||
maximumFractionDigits: 2,
|
||||
})}`;
|
||||
}
|
||||
|
||||
export function formatDuration(ms: number | undefined | null): string {
|
||||
if (ms == null || isNaN(ms)) return "—";
|
||||
if (ms < 1000) return `${ms}ms`;
|
||||
@@ -104,4 +112,4 @@ export function computeDurationMs(startTime?: string, endTime?: string): number
|
||||
const end = new Date(endTime).getTime();
|
||||
if (isNaN(start) || isNaN(end)) return null;
|
||||
return end - start;
|
||||
}
|
||||
}
|
||||
@@ -13,7 +13,7 @@ import { useProviders } from "@/pages/providers/hooks/use-providers";
|
||||
import { useTraces } from "@/pages/traces/hooks/use-traces";
|
||||
import { Methods, Events } from "@/api/protocol";
|
||||
import { ROUTES } from "@/lib/constants";
|
||||
import { formatTokens, formatCost } from "@/lib/format";
|
||||
import { formatTokens, formatApiCost } from "@/lib/format";
|
||||
|
||||
import type {
|
||||
HealthPayload,
|
||||
@@ -220,7 +220,7 @@ export function OverviewPage() {
|
||||
<StatCard
|
||||
icon={DollarSign}
|
||||
label={t("statCards.costToday", "Cost Today")}
|
||||
value={formatCost(quota?.costToday)}
|
||||
value={formatApiCost(quota?.costToday)}
|
||||
sparkline={sparklines?.costSparkline}
|
||||
trend={sparklines?.trends.cost}
|
||||
/>
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { TrendingUp, TrendingDown, Minus } from "lucide-react";
|
||||
import { formatTokens, formatCost } from "@/lib/format";
|
||||
import { formatTokens, formatApiCost } from "@/lib/format";
|
||||
import type { SummaryData } from "../hooks/use-usage-analytics";
|
||||
|
||||
interface SummaryCardsProps {
|
||||
@@ -82,7 +82,7 @@ export function SummaryCards({ current, previous, loading }: SummaryCardsProps)
|
||||
},
|
||||
{
|
||||
label: t("analytics.cost"),
|
||||
value: formatCost(current.cost),
|
||||
value: formatApiCost(current.cost),
|
||||
trend: trendPercent(current.cost, previous.cost),
|
||||
hint: allCostZero ? t("analytics.configurePricing") : undefined,
|
||||
},
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { useState } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { ArrowUpDown, ArrowUp, ArrowDown } from "lucide-react";
|
||||
import { formatTokens, formatCost, formatDuration } from "@/lib/format";
|
||||
import { formatTokens, formatApiCost, formatDuration } from "@/lib/format";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { useUsageFilterContext } from "../context/usage-filter-context";
|
||||
import type { SnapshotBreakdown } from "../hooks/use-usage-analytics";
|
||||
@@ -15,7 +15,7 @@ interface TopModelsTableProps {
|
||||
|
||||
export function TopModelsTable({ data, loading }: TopModelsTableProps) {
|
||||
const { t } = useTranslation("usage");
|
||||
const { filters, toggleFilter } = useUsageFilterContext();
|
||||
const { filters, setFilter } = useUsageFilterContext();
|
||||
const [sortKey, setSortKey] = useState<SortKey>("llm_call_count");
|
||||
const [sortDir, setSortDir] = useState<"asc" | "desc">("desc");
|
||||
|
||||
@@ -55,6 +55,12 @@ export function TopModelsTable({ data, loading }: TopModelsTableProps) {
|
||||
);
|
||||
}
|
||||
|
||||
function parseProviderModelKey(key: string) {
|
||||
const idx = key.indexOf("/");
|
||||
if (idx < 0) return { provider: "", model: key };
|
||||
return { provider: key.slice(0, idx), model: key.slice(idx + 1) };
|
||||
}
|
||||
|
||||
if (loading) {
|
||||
return (
|
||||
<div className="rounded-lg border bg-card p-4">
|
||||
@@ -90,10 +96,17 @@ export function TopModelsTable({ data, loading }: TopModelsTableProps) {
|
||||
</thead>
|
||||
<tbody>
|
||||
{sorted.map((row) => {
|
||||
const isActive = filters.model === row.key;
|
||||
const [model, provider] = row.key.includes("/")
|
||||
? row.key.split("/", 2)
|
||||
: [row.key, "—"];
|
||||
const { provider, model } = parseProviderModelKey(row.key);
|
||||
const isActive = filters.model === model && (!provider || filters.provider === provider);
|
||||
const toggleProviderModel = () => {
|
||||
if (isActive) {
|
||||
setFilter("provider", undefined);
|
||||
setFilter("model", undefined);
|
||||
return;
|
||||
}
|
||||
setFilter("provider", provider || undefined);
|
||||
setFilter("model", model);
|
||||
};
|
||||
return (
|
||||
<tr
|
||||
key={row.key}
|
||||
@@ -101,15 +114,15 @@ export function TopModelsTable({ data, loading }: TopModelsTableProps) {
|
||||
"border-b last:border-0 hover:bg-muted/30 cursor-pointer transition-colors",
|
||||
isActive && "bg-primary/5",
|
||||
)}
|
||||
onClick={() => toggleFilter("model", row.key)}
|
||||
onClick={toggleProviderModel}
|
||||
>
|
||||
<td className="px-3 py-2 font-medium">{model}</td>
|
||||
<td className="px-3 py-2 text-muted-foreground">{provider}</td>
|
||||
<td className="px-3 py-2 text-muted-foreground">{provider || "—"}</td>
|
||||
<td className="px-3 py-2 text-right">{row.llm_call_count.toLocaleString()}</td>
|
||||
<td className="px-3 py-2 text-right text-muted-foreground">{formatTokens(row.input_tokens)}</td>
|
||||
<td className="px-3 py-2 text-right text-muted-foreground">{formatTokens(row.output_tokens)}</td>
|
||||
<td className="px-3 py-2 text-right text-muted-foreground">{formatDuration(row.avg_duration_ms)}</td>
|
||||
<td className="px-3 py-2 text-right">{formatCost(row.total_cost)}</td>
|
||||
<td className="px-3 py-2 text-right">{formatApiCost(row.total_cost)}</td>
|
||||
</tr>
|
||||
);
|
||||
})}
|
||||
|
||||
@@ -2,7 +2,7 @@ import { useTranslation } from "react-i18next";
|
||||
import { Pencil, Trash2 } from "lucide-react";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { formatCost, formatTokens } from "@/lib/format";
|
||||
import { formatApiCost, formatCost, formatTokens } from "@/lib/format";
|
||||
import type { UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps";
|
||||
|
||||
interface UsageCapRowProps {
|
||||
@@ -31,7 +31,7 @@ export function UsageCapRow({ row, onEdit, onDelete }: UsageCapRowProps) {
|
||||
</td>
|
||||
<td className="px-3 py-2"><Badge variant="outline">{t(`caps.windows.${p.window}`)}</Badge></td>
|
||||
<td className="px-3 py-2 text-right">{p.max_tokens ? `${formatTokens(tokenUsed)} / ${formatTokens(p.max_tokens)} (${tokenPct}%)` : "-"}</td>
|
||||
<td className="px-3 py-2 text-right">{p.max_cost_micros ? `${formatCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "-"}</td>
|
||||
<td className="px-3 py-2 text-right">{p.max_cost_micros ? `${formatApiCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "-"}</td>
|
||||
<td className="px-3 py-2 text-right"><Badge variant={p.enabled ? "default" : "secondary"}>{p.enabled ? t("caps.enabled") : t("caps.disabled")}</Badge></td>
|
||||
<td className="px-3 py-2 text-right">
|
||||
<div className="flex justify-end gap-1">
|
||||
|
||||
@@ -7,7 +7,7 @@ import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { formatCost, formatDate, formatTokens } from "@/lib/format";
|
||||
import { formatApiCost, formatDate, formatTokens } from "@/lib/format";
|
||||
import { useAgents } from "@/pages/agents/hooks/use-agents";
|
||||
import { useProviders } from "@/pages/providers/hooks/use-providers";
|
||||
import { useUsageCaps } from "../hooks/use-usage-caps";
|
||||
@@ -183,7 +183,7 @@ export function UsageCapsPanel() {
|
||||
{blockedEvents.slice(0, 4).map((event) => (
|
||||
<div key={event.id} className="rounded-md border border-destructive/30 bg-destructive/5 px-3 py-2 text-xs">
|
||||
<div className="font-medium">{event.reason || t("caps.blocked")}</div>
|
||||
<div className="text-muted-foreground">{formatDate(event.created_at)} · {formatTokens(event.estimated_tokens)} · {formatCost(event.estimated_cost_micros / 1_000_000)}</div>
|
||||
<div className="text-muted-foreground">{formatDate(event.created_at)} · {formatTokens(event.estimated_tokens)} · {formatApiCost(event.estimated_cost_micros / 1_000_000)}</div>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
|
||||
@@ -3,7 +3,7 @@ import { useTranslation } from "react-i18next";
|
||||
import { Cable, Sparkles, Terminal, Wrench } from "lucide-react";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { formatCost, formatDuration, formatTokens } from "@/lib/format";
|
||||
import { formatApiCost, formatDuration, formatTokens } from "@/lib/format";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { useUsageFilterContext } from "../context/usage-filter-context";
|
||||
import { useUsageEventAnalytics, type UsageEventResourceType } from "../hooks/use-usage-event-analytics";
|
||||
@@ -66,7 +66,7 @@ export function UsageEventAnalyticsPanel() {
|
||||
<Metric label={t("analytics.events.metrics.errorRate")} value={`${errorRate.toFixed(1)}%`} loading={loading} />
|
||||
<Metric label={t("analytics.events.metrics.tokens")} value={formatTokens(current.total_tokens)} loading={loading} />
|
||||
<Metric label={t("analytics.events.metrics.avgDuration")} value={formatDuration(current.avg_duration_ms)} loading={loading} />
|
||||
<Metric label={t("analytics.events.metrics.cost")} value={formatCost(current.cost_usd)} loading={loading} />
|
||||
<Metric label={t("analytics.events.metrics.cost")} value={formatApiCost(current.cost_usd)} loading={loading} />
|
||||
</div>
|
||||
|
||||
<div className="flex flex-wrap items-center gap-2 text-xs text-muted-foreground">
|
||||
@@ -116,7 +116,7 @@ export function UsageEventAnalyticsPanel() {
|
||||
<td className="px-3 py-2 text-right text-muted-foreground">{rowErrorRate.toFixed(1)}%</td>
|
||||
<td className="px-3 py-2 text-right text-muted-foreground">{formatDuration(row.avg_duration_ms)}</td>
|
||||
<td className="px-3 py-2 text-right text-muted-foreground">{formatTokens(row.total_tokens)}</td>
|
||||
<td className="px-3 py-2 text-right">{formatCost(row.cost_usd)}</td>
|
||||
<td className="px-3 py-2 text-right">{formatApiCost(row.cost_usd)}</td>
|
||||
</tr>
|
||||
);
|
||||
})
|
||||
|
||||
@@ -23,7 +23,7 @@ interface UsageFilterContextValue {
|
||||
activeFilterCount: number;
|
||||
}
|
||||
|
||||
function buildTimeRange(period: Period): { from: string; to: string; granularity: "hour" | "day" } {
|
||||
export function buildTimeRange(period: Period): { from: string; to: string; granularity: "hour" | "day" } {
|
||||
const now = new Date();
|
||||
let from: Date;
|
||||
let granularity: "hour" | "day";
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import { useCallback } from "react";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { useHttp } from "@/hooks/use-ws";
|
||||
import type { UsageFilters } from "../context/usage-filter-context";
|
||||
import { buildTimeRange, type UsageFilters } from "../context/usage-filter-context";
|
||||
|
||||
export interface SnapshotTimeSeries {
|
||||
bucket_time: string;
|
||||
@@ -50,9 +51,10 @@ interface SummaryResponse {
|
||||
}
|
||||
|
||||
function buildParams(filters: UsageFilters, extra?: Record<string, string>): Record<string, string> {
|
||||
const range = filters.period === "custom" ? filters : buildTimeRange(filters.period);
|
||||
const p: Record<string, string> = {
|
||||
from: filters.from,
|
||||
to: filters.to,
|
||||
from: range.from,
|
||||
to: range.to,
|
||||
};
|
||||
if (filters.agentId) p.agent_id = filters.agentId;
|
||||
if (filters.provider) p.provider = filters.provider;
|
||||
@@ -66,9 +68,14 @@ function filterKey(f: UsageFilters) {
|
||||
return [f.from, f.to, f.agentId, f.provider, f.model, f.channel] as const;
|
||||
}
|
||||
|
||||
// Snapshots update hourly — no need to re-fetch on window focus or within 60s.
|
||||
const STALE_TIME = 60_000;
|
||||
const QUERY_OPTS = { staleTime: STALE_TIME, refetchOnWindowFocus: false } as const;
|
||||
// Current-hour usage is filled from live traces, so keep the open dashboard fresh.
|
||||
const REFRESH_INTERVAL = 30_000;
|
||||
const QUERY_OPTS = {
|
||||
staleTime: REFRESH_INTERVAL,
|
||||
refetchInterval: REFRESH_INTERVAL,
|
||||
refetchIntervalInBackground: false,
|
||||
refetchOnWindowFocus: true,
|
||||
} as const;
|
||||
|
||||
export function useUsageAnalytics(filters: UsageFilters) {
|
||||
const http = useHttp();
|
||||
@@ -98,6 +105,14 @@ export function useUsageAnalytics(filters: UsageFilters) {
|
||||
...QUERY_OPTS,
|
||||
});
|
||||
|
||||
const providerModelQuery = useQuery({
|
||||
queryKey: ["usage", "breakdown", "provider_model", ...fk],
|
||||
queryFn: () =>
|
||||
http.get<{ rows: SnapshotBreakdown[] }>("/v1/usage/breakdown", buildParams(filters, { group_by: "provider_model" })),
|
||||
placeholderData: (prev) => prev,
|
||||
...QUERY_OPTS,
|
||||
});
|
||||
|
||||
const channelQuery = useQuery({
|
||||
queryKey: ["usage", "breakdown", "channel", ...fk],
|
||||
queryFn: () =>
|
||||
@@ -114,22 +129,51 @@ export function useUsageAnalytics(filters: UsageFilters) {
|
||||
...QUERY_OPTS,
|
||||
});
|
||||
|
||||
const refreshAnalytics = useCallback(async () => {
|
||||
await Promise.all([
|
||||
timeseriesQuery.refetch(),
|
||||
providerQuery.refetch(),
|
||||
modelQuery.refetch(),
|
||||
providerModelQuery.refetch(),
|
||||
channelQuery.refetch(),
|
||||
summaryQuery.refetch(),
|
||||
]);
|
||||
}, [
|
||||
timeseriesQuery.refetch,
|
||||
providerQuery.refetch,
|
||||
modelQuery.refetch,
|
||||
providerModelQuery.refetch,
|
||||
channelQuery.refetch,
|
||||
summaryQuery.refetch,
|
||||
]);
|
||||
|
||||
// isLoading = first mount only (no cached data) → shows skeleton.
|
||||
// placeholderData keeps previous results visible during refetch → no flicker.
|
||||
const loading =
|
||||
timeseriesQuery.isLoading ||
|
||||
providerQuery.isLoading ||
|
||||
modelQuery.isLoading ||
|
||||
providerModelQuery.isLoading ||
|
||||
channelQuery.isLoading ||
|
||||
summaryQuery.isLoading;
|
||||
const refreshing =
|
||||
timeseriesQuery.isFetching ||
|
||||
providerQuery.isFetching ||
|
||||
modelQuery.isFetching ||
|
||||
providerModelQuery.isFetching ||
|
||||
channelQuery.isFetching ||
|
||||
summaryQuery.isFetching;
|
||||
|
||||
return {
|
||||
timeseries: timeseriesQuery.data?.points ?? [],
|
||||
providerBreakdown: providerQuery.data?.rows ?? [],
|
||||
modelBreakdown: modelQuery.data?.rows ?? [],
|
||||
providerModelBreakdown: providerModelQuery.data?.rows ?? [],
|
||||
channelBreakdown: channelQuery.data?.rows ?? [],
|
||||
summary: summaryQuery.data ?? null,
|
||||
loading,
|
||||
refreshing,
|
||||
refreshAnalytics,
|
||||
error: timeseriesQuery.error,
|
||||
};
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { useHttp } from "@/hooks/use-ws";
|
||||
import type { UsageFilters } from "../context/usage-filter-context";
|
||||
import { buildTimeRange, type UsageFilters } from "../context/usage-filter-context";
|
||||
|
||||
export type UsageEventResourceType = "tool" | "skill" | "mcp_tool" | "runtime_tool";
|
||||
|
||||
@@ -50,9 +50,10 @@ export interface UsageEventTimeSeries {
|
||||
}
|
||||
|
||||
function buildParams(filters: UsageFilters, resourceType: UsageEventResourceType, extra?: Record<string, string>): Record<string, string> {
|
||||
const range = filters.period === "custom" ? filters : buildTimeRange(filters.period);
|
||||
const p: Record<string, string> = {
|
||||
from: filters.from,
|
||||
to: filters.to,
|
||||
from: range.from,
|
||||
to: range.to,
|
||||
resource_type: resourceType,
|
||||
};
|
||||
if (filters.agentId) p.agent_id = filters.agentId;
|
||||
@@ -66,7 +67,13 @@ function filterKey(f: UsageFilters, resourceType: UsageEventResourceType) {
|
||||
return [resourceType, f.from, f.to, f.agentId, f.provider, f.model, f.channel, f.granularity] as const;
|
||||
}
|
||||
|
||||
const QUERY_OPTS = { staleTime: 60_000, refetchOnWindowFocus: false } as const;
|
||||
const REFRESH_INTERVAL = 30_000;
|
||||
const QUERY_OPTS = {
|
||||
staleTime: REFRESH_INTERVAL,
|
||||
refetchInterval: REFRESH_INTERVAL,
|
||||
refetchIntervalInBackground: false,
|
||||
refetchOnWindowFocus: true,
|
||||
} as const;
|
||||
|
||||
export function useUsageEventAnalytics(filters: UsageFilters, resourceType: UsageEventResourceType) {
|
||||
const http = useHttp();
|
||||
|
||||
@@ -11,6 +11,7 @@ export interface UsageRecord {
|
||||
inputTokens: number;
|
||||
outputTokens: number;
|
||||
totalTokens: number;
|
||||
cost: number;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ import { PageHeader } from "@/components/shared/page-header";
|
||||
import { EmptyState } from "@/components/shared/empty-state";
|
||||
import { Pagination } from "@/components/shared/pagination";
|
||||
import { ErrorBoundary } from "@/components/shared/error-boundary";
|
||||
import { formatTokens, formatCost } from "@/lib/format";
|
||||
import { formatTokens, formatApiCost } from "@/lib/format";
|
||||
import { useAgents } from "@/pages/agents/hooks/use-agents";
|
||||
import { useUsage } from "./hooks/use-usage";
|
||||
import { useUsageAnalytics } from "./hooks/use-usage-analytics";
|
||||
@@ -30,7 +30,7 @@ function AnalyticsDashboard() {
|
||||
const { t } = useTranslation("usage");
|
||||
const { filters, setFilter } = useUsageFilterContext();
|
||||
const { agents } = useAgents();
|
||||
const { timeseries, providerBreakdown, modelBreakdown, channelBreakdown, summary, loading, error } =
|
||||
const { timeseries, providerBreakdown, modelBreakdown, providerModelBreakdown, channelBreakdown, summary, loading, refreshing, refreshAnalytics, error } =
|
||||
useUsageAnalytics(filters);
|
||||
|
||||
// Legacy records table state
|
||||
@@ -42,9 +42,20 @@ function AnalyticsDashboard() {
|
||||
const setPageSize = (size: number) => { setPageSizeRaw(size); setPage(1); setGlobalPageSize(size); };
|
||||
const totalPages = Math.max(1, Math.ceil(total / pageSize));
|
||||
|
||||
const loadCurrentRecords = useCallback(() => {
|
||||
return loadRecords({ limit: pageSize, offset: (page - 1) * pageSize, agentId: filters.agentId });
|
||||
}, [loadRecords, page, pageSize, filters.agentId]);
|
||||
|
||||
useEffect(() => {
|
||||
loadRecords({ limit: pageSize, offset: (page - 1) * pageSize, agentId: filters.agentId });
|
||||
}, [page, pageSize, filters.agentId]);
|
||||
loadCurrentRecords();
|
||||
}, [loadCurrentRecords]);
|
||||
|
||||
useEffect(() => {
|
||||
const id = window.setInterval(() => {
|
||||
void loadCurrentRecords();
|
||||
}, 30_000);
|
||||
return () => window.clearInterval(id);
|
||||
}, [loadCurrentRecords]);
|
||||
|
||||
const agentList = agents.map((a) => ({
|
||||
id: a.id,
|
||||
@@ -75,10 +86,15 @@ function AnalyticsDashboard() {
|
||||
URL.revokeObjectURL(url);
|
||||
}, [timeseries, filters.period]);
|
||||
|
||||
const handleRefresh = useCallback(() => {
|
||||
void Promise.all([refreshAnalytics(), loadCurrentRecords()]);
|
||||
}, [refreshAnalytics, loadCurrentRecords]);
|
||||
|
||||
const current = summary?.current ?? EMPTY_SUMMARY;
|
||||
const previous = summary?.previous ?? EMPTY_SUMMARY;
|
||||
|
||||
const apiError = error instanceof Error ? error.message : error ? String(error) : null;
|
||||
const isRefreshing = recLoading || refreshing;
|
||||
|
||||
return (
|
||||
<div className="p-4 sm:p-6 space-y-4">
|
||||
@@ -86,8 +102,8 @@ function AnalyticsDashboard() {
|
||||
title={t("analytics.title")}
|
||||
description={t("description")}
|
||||
actions={
|
||||
<Button variant="outline" size="sm" onClick={() => loadRecords()} disabled={recLoading} className="gap-1">
|
||||
<RefreshCw className={`h-3.5 w-3.5${recLoading ? " animate-spin" : ""}`} />
|
||||
<Button variant="outline" size="sm" onClick={handleRefresh} disabled={isRefreshing} className="gap-1">
|
||||
<RefreshCw className={`h-3.5 w-3.5${isRefreshing ? " animate-spin" : ""}`} />
|
||||
{t("common:refresh", "Refresh")}
|
||||
</Button>
|
||||
}
|
||||
@@ -142,7 +158,7 @@ function AnalyticsDashboard() {
|
||||
</ErrorBoundary>
|
||||
|
||||
<ErrorBoundary>
|
||||
<TopModelsTable data={modelBreakdown} loading={loading} />
|
||||
<TopModelsTable data={providerModelBreakdown} loading={loading} />
|
||||
</ErrorBoundary>
|
||||
|
||||
{/* Legacy records table */}
|
||||
@@ -178,7 +194,7 @@ function AnalyticsDashboard() {
|
||||
<td className="px-4 py-3 text-muted-foreground">—</td>
|
||||
<td className="px-4 py-3 text-right text-muted-foreground">{formatTokens(r.inputTokens)}</td>
|
||||
<td className="px-4 py-3 text-right text-muted-foreground">{formatTokens(r.outputTokens)}</td>
|
||||
<td className="px-4 py-3 text-right">{formatCost(0)}</td>
|
||||
<td className="px-4 py-3 text-right">{formatApiCost(r.cost)}</td>
|
||||
</tr>
|
||||
))}
|
||||
</tbody>
|
||||
|
||||
Reference in new issue
Block a user