From 747b58d24cab351613c57a153bca9af0fc3a72fa Mon Sep 17 00:00:00 2001 From: Duc Nguyen <87276974+itsddvn@users.noreply.github.com> Date: Fri, 3 Jul 2026 06:24:21 +0700 Subject: [PATCH] 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. --- cmd/gateway.go | 101 +++++- cmd/gateway_methods.go | 4 +- docs/02-providers.md | 7 +- docs/10-tracing-observability.md | 14 +- docs/15-core-skills-system.md | 6 + docs/project-changelog.md | 116 +++++++ internal/agent/loop_pricing.go | 45 +++ internal/agent/loop_pricing_test.go | 185 +++++++++++ internal/agent/loop_run.go | 11 +- internal/agent/loop_tracing.go | 39 ++- internal/agent/usage_events.go | 8 +- internal/agent/usage_events_test.go | 45 +++ internal/gateway/methods/usage.go | 45 ++- internal/gateway/methods/usage_test.go | 106 ++++++ internal/http/openapi_spec.json | 4 +- internal/http/usage.go | 97 +----- internal/http/usage_live.go | 145 +++++++++ internal/http/usage_live_breakdown.go | 146 +++++++++ internal/http/usage_live_merge.go | 150 +++++++++ internal/http/usage_live_test.go | 205 ++++++++++++ internal/store/pg/snapshot.go | 4 + internal/store/pg/snapshot_cost_backfill.go | 145 +++++++++ .../store/pg/snapshot_cost_backfill_test.go | 90 +++++ internal/store/pg/tracing.go | 61 +++- internal/store/pg/tracing_cost_backfill.go | 307 ++++++++++++++++++ .../store/pg/tracing_cost_backfill_test.go | 175 ++++++++++ internal/store/pg/usage_caps_test.go | 38 +++ .../store/pg/usage_event_cost_backfill.go | 61 ++++ .../pg/usage_event_cost_backfill_test.go | 128 ++++++++ internal/store/pg/usage_pricing.go | 31 +- internal/store/snapshot_store.go | 8 +- internal/store/sqlitestore/snapshots.go | 4 + .../snapshots_provider_model_test.go | 53 +++ internal/store/sqlitestore/tracing.go | 58 ++++ .../sqlitestore/tracing_session_costs_test.go | 150 +++++++++ internal/store/sqlitestore/tracing_spans.go | 7 +- internal/store/tracing_store.go | 10 + internal/store/usage_event_store.go | 5 + internal/tracing/cost.go | 17 + internal/tracing/cost_test.go | 43 ++- internal/tracing/snapshot_worker.go | 53 ++- internal/tracing/snapshot_worker_test.go | 37 +++ internal/usage/caps/service.go | 12 + internal/usage/pricing/openrouter_sync.go | 61 ++++ internal/usage/pricing/openrouter_test.go | 100 ++++++ internal/webui/handler.go | 11 +- ui/web/src/lib/format.test.ts | 16 + ui/web/src/lib/format.ts | 10 +- ui/web/src/pages/overview/overview-page.tsx | 4 +- .../pages/usage/components/summary-cards.tsx | 4 +- .../usage/components/top-models-table.tsx | 31 +- .../pages/usage/components/usage-cap-row.tsx | 4 +- .../usage/components/usage-caps-panel.tsx | 4 +- .../usage-event-analytics-panel.tsx | 6 +- .../usage/context/usage-filter-context.tsx | 2 +- .../pages/usage/hooks/use-usage-analytics.ts | 56 +++- .../usage/hooks/use-usage-event-analytics.ts | 15 +- ui/web/src/pages/usage/hooks/use-usage.ts | 1 + ui/web/src/pages/usage/usage-page.tsx | 32 +- 59 files changed, 3136 insertions(+), 197 deletions(-) create mode 100644 internal/agent/loop_pricing.go create mode 100644 internal/agent/loop_pricing_test.go create mode 100644 internal/gateway/methods/usage_test.go create mode 100644 internal/http/usage_live.go create mode 100644 internal/http/usage_live_breakdown.go create mode 100644 internal/http/usage_live_merge.go create mode 100644 internal/http/usage_live_test.go create mode 100644 internal/store/pg/snapshot_cost_backfill.go create mode 100644 internal/store/pg/snapshot_cost_backfill_test.go create mode 100644 internal/store/pg/tracing_cost_backfill.go create mode 100644 internal/store/pg/tracing_cost_backfill_test.go create mode 100644 internal/store/pg/usage_event_cost_backfill.go create mode 100644 internal/store/pg/usage_event_cost_backfill_test.go create mode 100644 internal/store/sqlitestore/snapshots_provider_model_test.go create mode 100644 internal/store/sqlitestore/tracing_session_costs_test.go create mode 100644 internal/tracing/snapshot_worker_test.go create mode 100644 internal/usage/pricing/openrouter_sync.go create mode 100644 internal/usage/pricing/openrouter_test.go create mode 100644 ui/web/src/lib/format.test.ts diff --git a/cmd/gateway.go b/cmd/gateway.go index 3f65591c..9a6ae41f 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -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) diff --git a/cmd/gateway_methods.go b/cmd/gateway_methods.go index 28c5c976..15085d6c 100644 --- a/cmd/gateway_methods.go +++ b/cmd/gateway_methods.go @@ -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) diff --git a/docs/02-providers.md b/docs/02-providers.md index 5bcdf536..ed998e7f 100644 --- a/docs/02-providers.md +++ b/docs/02-providers.md @@ -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. diff --git a/docs/10-tracing-observability.md b/docs/10-tracing-observability.md index 6ed99ed6..a1a55de4 100644 --- a/docs/10-tracing-observability.md +++ b/docs/10-tracing-observability.md @@ -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. --- diff --git a/docs/15-core-skills-system.md b/docs/15-core-skills-system.md index a3e3fabb..04d363db 100644 --- a/docs/15-core-skills-system.md +++ b/docs/15-core-skills-system.md @@ -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 ``` diff --git a/docs/project-changelog.md b/docs/project-changelog.md index 8d9a68ef..c7b27765 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.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 diff --git a/internal/agent/loop_pricing.go b/internal/agent/loop_pricing.go new file mode 100644 index 00000000..06395aef --- /dev/null +++ b/internal/agent/loop_pricing.go @@ -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 +} diff --git a/internal/agent/loop_pricing_test.go b/internal/agent/loop_pricing_test.go new file mode 100644 index 00000000..bfe8b394 --- /dev/null +++ b/internal/agent/loop_pricing_test.go @@ -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 +} diff --git a/internal/agent/loop_run.go b/internal/agent/loop_run.go index 6231a6dd..fd573534 100644 --- a/internal/agent/loop_run.go +++ b/internal/agent/loop_run.go @@ -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 { diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index 604150bb..c6cc88bd 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -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) diff --git a/internal/agent/usage_events.go b/internal/agent/usage_events.go index 3d299701..fa543265 100644 --- a/internal/agent/usage_events.go +++ b/internal/agent/usage_events.go @@ -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) } diff --git a/internal/agent/usage_events_test.go b/internal/agent/usage_events_test.go index 3c53e497..165878c2 100644 --- a/internal/agent/usage_events_test.go +++ b/internal/agent/usage_events_test.go @@ -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()} diff --git a/internal/gateway/methods/usage.go b/internal/gateway/methods/usage.go index 2e817683..b2fea661 100644 --- a/internal/gateway/methods/usage.go +++ b/internal/gateway/methods/usage.go @@ -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, diff --git a/internal/gateway/methods/usage_test.go b/internal/gateway/methods/usage_test.go new file mode 100644 index 00000000..673ebf3c --- /dev/null +++ b/internal/gateway/methods/usage_test.go @@ -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 +} diff --git a/internal/http/openapi_spec.json b/internal/http/openapi_spec.json index db18fa97..ccb19a53 100644 --- a/internal/http/openapi_spec.json +++ b/internal/http/openapi_spec.json @@ -10533,7 +10533,9 @@ "enum": [ "agent", "provider", - "model" + "model", + "provider_model", + "channel" ] } } diff --git a/internal/http/usage.go b/internal/http/usage.go index 5b29e275..04e71727 100644 --- a/internal/http/usage.go +++ b/internal/http/usage.go @@ -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 { diff --git a/internal/http/usage_live.go b/internal/http/usage_live.go new file mode 100644 index 00000000..d92217ef --- /dev/null +++ b/internal/http/usage_live.go @@ -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 + } +} diff --git a/internal/http/usage_live_breakdown.go b/internal/http/usage_live_breakdown.go new file mode 100644 index 00000000..f6f1361d --- /dev/null +++ b/internal/http/usage_live_breakdown.go @@ -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() +} diff --git a/internal/http/usage_live_merge.go b/internal/http/usage_live_merge.go new file mode 100644 index 00000000..3dc79ec8 --- /dev/null +++ b/internal/http/usage_live_merge.go @@ -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 +} diff --git a/internal/http/usage_live_test.go b/internal/http/usage_live_test.go new file mode 100644 index 00000000..6808332b --- /dev/null +++ b/internal/http/usage_live_test.go @@ -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) + } +} diff --git a/internal/store/pg/snapshot.go b/internal/store/pg/snapshot.go index 40125058..71dcc726 100644 --- a/internal/store/pg/snapshot.go +++ b/internal/store/pg/snapshot.go @@ -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" diff --git a/internal/store/pg/snapshot_cost_backfill.go b/internal/store/pg/snapshot_cost_backfill.go new file mode 100644 index 00000000..de8aa71c --- /dev/null +++ b/internal/store/pg/snapshot_cost_backfill.go @@ -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 +} diff --git a/internal/store/pg/snapshot_cost_backfill_test.go b/internal/store/pg/snapshot_cost_backfill_test.go new file mode 100644 index 00000000..3c0b38a3 --- /dev/null +++ b/internal/store/pg/snapshot_cost_backfill_test.go @@ -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) + } +} diff --git a/internal/store/pg/tracing.go b/internal/store/pg/tracing.go index a3836672..cd12d009 100644 --- a/internal/store/pg/tracing.go +++ b/internal/store/pg/tracing.go @@ -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 diff --git a/internal/store/pg/tracing_cost_backfill.go b/internal/store/pg/tracing_cost_backfill.go new file mode 100644 index 00000000..7af68c87 --- /dev/null +++ b/internal/store/pg/tracing_cost_backfill.go @@ -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 +} diff --git a/internal/store/pg/tracing_cost_backfill_test.go b/internal/store/pg/tracing_cost_backfill_test.go new file mode 100644 index 00000000..1a91311b --- /dev/null +++ b/internal/store/pg/tracing_cost_backfill_test.go @@ -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) + } +} diff --git a/internal/store/pg/usage_caps_test.go b/internal/store/pg/usage_caps_test.go index 216ad6de..6db24490 100644 --- a/internal/store/pg/usage_caps_test.go +++ b/internal/store/pg/usage_caps_test.go @@ -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) { diff --git a/internal/store/pg/usage_event_cost_backfill.go b/internal/store/pg/usage_event_cost_backfill.go new file mode 100644 index 00000000..e0d8aa70 --- /dev/null +++ b/internal/store/pg/usage_event_cost_backfill.go @@ -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 +} diff --git a/internal/store/pg/usage_event_cost_backfill_test.go b/internal/store/pg/usage_event_cost_backfill_test.go new file mode 100644 index 00000000..75375961 --- /dev/null +++ b/internal/store/pg/usage_event_cost_backfill_test.go @@ -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) + } +} diff --git a/internal/store/pg/usage_pricing.go b/internal/store/pg/usage_pricing.go index 64ef138e..50e71086 100644 --- a/internal/store/pg/usage_pricing.go +++ b/internal/store/pg/usage_pricing.go @@ -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, " ", "-") diff --git a/internal/store/snapshot_store.go b/internal/store/snapshot_store.go index d4e1bc0b..ce9576f9 100644 --- a/internal/store/snapshot_store.go +++ b/internal/store/snapshot_store.go @@ -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. diff --git a/internal/store/sqlitestore/snapshots.go b/internal/store/sqlitestore/snapshots.go index 9f633f5f..28f81f9a 100644 --- a/internal/store/sqlitestore/snapshots.go +++ b/internal/store/sqlitestore/snapshots.go @@ -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" diff --git a/internal/store/sqlitestore/snapshots_provider_model_test.go b/internal/store/sqlitestore/snapshots_provider_model_test.go new file mode 100644 index 00000000..a00b0eea --- /dev/null +++ b/internal/store/sqlitestore/snapshots_provider_model_test.go @@ -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) + } +} diff --git a/internal/store/sqlitestore/tracing.go b/internal/store/sqlitestore/tracing.go index fcb912e0..410803e4 100644 --- a/internal/store/sqlitestore/tracing.go +++ b/internal/store/sqlitestore/tracing.go @@ -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) diff --git a/internal/store/sqlitestore/tracing_session_costs_test.go b/internal/store/sqlitestore/tracing_session_costs_test.go new file mode 100644 index 00000000..c3fe1ded --- /dev/null +++ b/internal/store/sqlitestore/tracing_session_costs_test.go @@ -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 +} diff --git a/internal/store/sqlitestore/tracing_spans.go b/internal/store/sqlitestore/tracing_spans.go index f5d0ac29..8042b372 100644 --- a/internal/store/sqlitestore/tracing_spans.go +++ b/internal/store/sqlitestore/tracing_spans.go @@ -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 } - diff --git a/internal/store/tracing_store.go b/internal/store/tracing_store.go index 2f652754..97a51313 100644 --- a/internal/store/tracing_store.go +++ b/internal/store/tracing_store.go @@ -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 diff --git a/internal/store/usage_event_store.go b/internal/store/usage_event_store.go index bd4a9be5..75d5f7a5 100644 --- a/internal/store/usage_event_store.go +++ b/internal/store/usage_event_store.go @@ -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 diff --git a/internal/tracing/cost.go b/internal/tracing/cost.go index c30aeaa3..7a3d9512 100644 --- a/internal/tracing/cost.go +++ b/internal/tracing/cost.go @@ -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 { diff --git a/internal/tracing/cost_test.go b/internal/tracing/cost_test.go index f36c566c..a2eb4b67 100644 --- a/internal/tracing/cost_test.go +++ b/internal/tracing/cost_test.go @@ -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) + } +} diff --git a/internal/tracing/snapshot_worker.go b/internal/tracing/snapshot_worker.go index 672c364e..e74b98af 100644 --- a/internal/tracing/snapshot_worker.go +++ b/internal/tracing/snapshot_worker.go @@ -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) diff --git a/internal/tracing/snapshot_worker_test.go b/internal/tracing/snapshot_worker_test.go new file mode 100644 index 00000000..71a47f2d --- /dev/null +++ b/internal/tracing/snapshot_worker_test.go @@ -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) + } +} diff --git a/internal/usage/caps/service.go b/internal/usage/caps/service.go index 30061bec..6af4c68a 100644 --- a/internal/usage/caps/service.go +++ b/internal/usage/caps/service.go @@ -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) } diff --git a/internal/usage/pricing/openrouter_sync.go b/internal/usage/pricing/openrouter_sync.go new file mode 100644 index 00000000..edcfa44b --- /dev/null +++ b/internal/usage/pricing/openrouter_sync.go @@ -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) + } +} diff --git a/internal/usage/pricing/openrouter_test.go b/internal/usage/pricing/openrouter_test.go new file mode 100644 index 00000000..891ae59b --- /dev/null +++ b/internal/usage/pricing/openrouter_test.go @@ -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) +} diff --git a/internal/webui/handler.go b/internal/webui/handler.go index 4516ce9a..4e70f8a3 100644 --- a/internal/webui/handler.go +++ b/internal/webui/handler.go @@ -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") +} diff --git a/ui/web/src/lib/format.test.ts b/ui/web/src/lib/format.test.ts new file mode 100644 index 00000000..01179a66 --- /dev/null +++ b/ui/web/src/lib/format.test.ts @@ -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"); + }); +}); diff --git a/ui/web/src/lib/format.ts b/ui/web/src/lib/format.ts index 753ba9de..2f6a60a6 100644 --- a/ui/web/src/lib/format.ts +++ b/ui/web/src/lib/format.ts @@ -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; -} \ No newline at end of file +} diff --git a/ui/web/src/pages/overview/overview-page.tsx b/ui/web/src/pages/overview/overview-page.tsx index dfe2afc9..089bede7 100644 --- a/ui/web/src/pages/overview/overview-page.tsx +++ b/ui/web/src/pages/overview/overview-page.tsx @@ -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() { diff --git a/ui/web/src/pages/usage/components/summary-cards.tsx b/ui/web/src/pages/usage/components/summary-cards.tsx index 7394d507..f6817603 100644 --- a/ui/web/src/pages/usage/components/summary-cards.tsx +++ b/ui/web/src/pages/usage/components/summary-cards.tsx @@ -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, }, diff --git a/ui/web/src/pages/usage/components/top-models-table.tsx b/ui/web/src/pages/usage/components/top-models-table.tsx index 89dce8d5..8b97c6af 100644 --- a/ui/web/src/pages/usage/components/top-models-table.tsx +++ b/ui/web/src/pages/usage/components/top-models-table.tsx @@ -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("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 (
@@ -90,10 +96,17 @@ export function TopModelsTable({ data, loading }: TopModelsTableProps) { {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 ( toggleFilter("model", row.key)} + onClick={toggleProviderModel} > {model} - {provider} + {provider || "—"} {row.llm_call_count.toLocaleString()} {formatTokens(row.input_tokens)} {formatTokens(row.output_tokens)} {formatDuration(row.avg_duration_ms)} - {formatCost(row.total_cost)} + {formatApiCost(row.total_cost)} ); })} diff --git a/ui/web/src/pages/usage/components/usage-cap-row.tsx b/ui/web/src/pages/usage/components/usage-cap-row.tsx index 99786c2d..ec3b5ad0 100644 --- a/ui/web/src/pages/usage/components/usage-cap-row.tsx +++ b/ui/web/src/pages/usage/components/usage-cap-row.tsx @@ -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) { {t(`caps.windows.${p.window}`)} {p.max_tokens ? `${formatTokens(tokenUsed)} / ${formatTokens(p.max_tokens)} (${tokenPct}%)` : "-"} - {p.max_cost_micros ? `${formatCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "-"} + {p.max_cost_micros ? `${formatApiCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "-"} {p.enabled ? t("caps.enabled") : t("caps.disabled")}
diff --git a/ui/web/src/pages/usage/components/usage-caps-panel.tsx b/ui/web/src/pages/usage/components/usage-caps-panel.tsx index 972f1c34..ce09b471 100644 --- a/ui/web/src/pages/usage/components/usage-caps-panel.tsx +++ b/ui/web/src/pages/usage/components/usage-caps-panel.tsx @@ -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) => (
{event.reason || t("caps.blocked")}
-
{formatDate(event.created_at)} · {formatTokens(event.estimated_tokens)} · {formatCost(event.estimated_cost_micros / 1_000_000)}
+
{formatDate(event.created_at)} · {formatTokens(event.estimated_tokens)} · {formatApiCost(event.estimated_cost_micros / 1_000_000)}
))}
diff --git a/ui/web/src/pages/usage/components/usage-event-analytics-panel.tsx b/ui/web/src/pages/usage/components/usage-event-analytics-panel.tsx index 186325c0..6089588a 100644 --- a/ui/web/src/pages/usage/components/usage-event-analytics-panel.tsx +++ b/ui/web/src/pages/usage/components/usage-event-analytics-panel.tsx @@ -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() { - +
@@ -116,7 +116,7 @@ export function UsageEventAnalyticsPanel() { {rowErrorRate.toFixed(1)}% {formatDuration(row.avg_duration_ms)} {formatTokens(row.total_tokens)} - {formatCost(row.cost_usd)} + {formatApiCost(row.cost_usd)} ); }) diff --git a/ui/web/src/pages/usage/context/usage-filter-context.tsx b/ui/web/src/pages/usage/context/usage-filter-context.tsx index 76bc9d0b..cea59803 100644 --- a/ui/web/src/pages/usage/context/usage-filter-context.tsx +++ b/ui/web/src/pages/usage/context/usage-filter-context.tsx @@ -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"; diff --git a/ui/web/src/pages/usage/hooks/use-usage-analytics.ts b/ui/web/src/pages/usage/hooks/use-usage-analytics.ts index 37e44e6b..72b40f54 100644 --- a/ui/web/src/pages/usage/hooks/use-usage-analytics.ts +++ b/ui/web/src/pages/usage/hooks/use-usage-analytics.ts @@ -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): Record { + const range = filters.period === "custom" ? filters : buildTimeRange(filters.period); const p: Record = { - 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, }; } diff --git a/ui/web/src/pages/usage/hooks/use-usage-event-analytics.ts b/ui/web/src/pages/usage/hooks/use-usage-event-analytics.ts index ba5fbf2c..835af8fa 100644 --- a/ui/web/src/pages/usage/hooks/use-usage-event-analytics.ts +++ b/ui/web/src/pages/usage/hooks/use-usage-event-analytics.ts @@ -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): Record { + const range = filters.period === "custom" ? filters : buildTimeRange(filters.period); const p: Record = { - 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(); diff --git a/ui/web/src/pages/usage/hooks/use-usage.ts b/ui/web/src/pages/usage/hooks/use-usage.ts index 950b2b06..454a131f 100644 --- a/ui/web/src/pages/usage/hooks/use-usage.ts +++ b/ui/web/src/pages/usage/hooks/use-usage.ts @@ -11,6 +11,7 @@ export interface UsageRecord { inputTokens: number; outputTokens: number; totalTokens: number; + cost: number; timestamp: number; } diff --git a/ui/web/src/pages/usage/usage-page.tsx b/ui/web/src/pages/usage/usage-page.tsx index b23b2b3e..9cf30359 100644 --- a/ui/web/src/pages/usage/usage-page.tsx +++ b/ui/web/src/pages/usage/usage-page.tsx @@ -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 (
@@ -86,8 +102,8 @@ function AnalyticsDashboard() { title={t("analytics.title")} description={t("description")} actions={ - } @@ -142,7 +158,7 @@ function AnalyticsDashboard() { - + {/* Legacy records table */} @@ -178,7 +194,7 @@ function AnalyticsDashboard() { — {formatTokens(r.inputTokens)} {formatTokens(r.outputTokens)} - {formatCost(0)} + {formatApiCost(r.cost)} ))}