From 9ff5a629d2a469b02eefbccb76487c783aa3dc98 Mon Sep 17 00:00:00 2001 From: Goon Date: Sat, 23 May 2026 00:16:03 +0700 Subject: [PATCH 01/10] feat(usage): add AI budget usage caps --- cmd/gateway.go | 12 +- cmd/gateway_agents.go | 7 +- cmd/gateway_http_wiring.go | 3 + cmd/gateway_managed.go | 15 +- docs/02-providers.md | 21 + docs/06-store-data-model.md | 26 ++ docs/18-http-api.md | 51 +++ docs/project-changelog.md | 17 + internal/agent/loop_compact.go | 5 +- internal/agent/loop_history_sanitize.go | 5 +- internal/agent/loop_pipeline_callbacks.go | 143 ++++-- internal/agent/loop_types.go | 112 ++--- internal/agent/memoryflush.go | 5 +- internal/agent/resolver.go | 3 + internal/agent/usage_caps_runtime.go | 88 ++++ internal/gateway/server.go | 5 + internal/http/usage_caps.go | 354 +++++++++++++++ internal/http/usage_caps_test.go | 17 + internal/providers/anthropic_stream.go | 2 +- internal/providers/codex.go | 2 +- internal/providers/model_fallback.go | 48 +++ internal/providers/model_fallback_test.go | 27 ++ internal/providers/openai_chat.go | 8 +- internal/providers/openai_http.go | 6 + internal/providers/openai_types.go | 9 +- internal/providers/types.go | 16 +- internal/store/pg/factory.go | 73 ++-- internal/store/pg/usage_caps.go | 406 ++++++++++++++++++ internal/store/pg/usage_caps_test.go | 180 ++++++++ internal/store/pg/usage_pricing.go | 362 ++++++++++++++++ internal/store/stores.go | 73 ++-- internal/store/usage_caps.go | 195 +++++++++ internal/tools/read_audio.go | 6 + internal/tools/read_audio_resolve.go | 48 ++- internal/tools/read_document.go | 6 + internal/tools/read_document_resolve.go | 28 +- internal/tools/read_image.go | 20 +- internal/tools/read_video.go | 6 + internal/tools/read_video_resolve.go | 24 +- internal/tools/subagent.go | 65 +-- internal/tools/subagent_exec.go | 47 +- internal/tools/subagent_spawn.go | 22 +- internal/tools/usage_caps.go | 44 ++ internal/upgrade/version.go | 2 +- internal/usage/caps/service.go | 275 ++++++++++++ internal/usage/caps/service_test.go | 337 +++++++++++++++ internal/usage/pricing/decimal.go | 119 +++++ internal/usage/pricing/decimal_test.go | 66 +++ internal/usage/pricing/openrouter.go | 91 ++++ migrations/000069_usage_caps_pricing.down.sql | 2 + migrations/000069_usage_caps_pricing.up.sql | 45 ++ migrations/000070_usage_cap_policies.down.sql | 4 + migrations/000070_usage_cap_policies.up.sql | 68 +++ ui/web/src/i18n/locales/en/providers.json | 32 ++ ui/web/src/i18n/locales/en/usage.json | 40 +- ui/web/src/i18n/locales/vi/providers.json | 32 ++ ui/web/src/i18n/locales/vi/usage.json | 40 +- ui/web/src/i18n/locales/zh/providers.json | 32 ++ ui/web/src/i18n/locales/zh/usage.json | 40 +- ui/web/src/lib/query-keys.ts | 7 + .../providers/hooks/use-model-pricing.ts | 93 ++++ .../provider-detail/provider-overview.tsx | 3 + .../provider-pricing-section.tsx | 158 +++++++ .../usage/components/usage-caps-panel.tsx | 179 ++++++++ .../src/pages/usage/hooks/use-usage-caps.ts | 94 ++++ ui/web/src/pages/usage/usage-page.tsx | 5 + ui/web/src/types/usage-caps.ts | 68 +++ 67 files changed, 4195 insertions(+), 249 deletions(-) create mode 100644 internal/agent/usage_caps_runtime.go create mode 100644 internal/http/usage_caps.go create mode 100644 internal/http/usage_caps_test.go create mode 100644 internal/store/pg/usage_caps.go create mode 100644 internal/store/pg/usage_caps_test.go create mode 100644 internal/store/pg/usage_pricing.go create mode 100644 internal/store/usage_caps.go create mode 100644 internal/tools/usage_caps.go create mode 100644 internal/usage/caps/service.go create mode 100644 internal/usage/caps/service_test.go create mode 100644 internal/usage/pricing/decimal.go create mode 100644 internal/usage/pricing/decimal_test.go create mode 100644 internal/usage/pricing/openrouter.go create mode 100644 migrations/000069_usage_caps_pricing.down.sql create mode 100644 migrations/000069_usage_caps_pricing.up.sql create mode 100644 migrations/000070_usage_cap_policies.down.sql create mode 100644 migrations/000070_usage_cap_policies.up.sql create mode 100644 ui/web/src/pages/providers/hooks/use-model-pricing.ts create mode 100644 ui/web/src/pages/providers/provider-detail/provider-pricing-section.tsx create mode 100644 ui/web/src/pages/usage/components/usage-caps-panel.tsx create mode 100644 ui/web/src/pages/usage/hooks/use-usage-caps.ts create mode 100644 ui/web/src/types/usage-caps.ts diff --git a/cmd/gateway.go b/cmd/gateway.go index da35b29e..31428735 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -46,6 +46,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/internal/vault" "github.com/nextlevelbuilder/goclaw/pkg/protocol" @@ -282,8 +283,15 @@ func runGateway() { slog.Info("bootstrap: capabilities backfill complete", "agents", count) } + usageCapSvc := usagecaps.NewService(pgStores.UsageCaps, pgStores.Providers) + if readImage, ok := toolsReg.Get("read_image"); ok { + if t, ok := readImage.(*tools.ReadImageTool); ok { + t.SetUsageCapService(usageCapSvc) + } + } + // Subagent system (secureCLI store wired so subagent ExecTools enforce the gate) - subagentMgr := setupSubagents(providerRegistry, cfg, msgBus, toolsReg, workspace, sandboxMgr, pgStores.SecureCLI) + subagentMgr := setupSubagents(providerRegistry, cfg, msgBus, toolsReg, workspace, sandboxMgr, pgStores.SecureCLI, usageCapSvc) if subagentMgr != nil { // Wire announce queue for batched subagent result delivery (matching TS debounce pattern). announceQueue := tools.NewAnnounceQueue(1000, 20, makeDelegateAnnounceCallback(subagentMgr, msgBus)) @@ -336,7 +344,7 @@ func runGateway() { var mcpPool *mcpbridge.Pool var mediaStore *media.Store var postTurn tools.PostTurnProcessor - contextFileInterceptor, mcpPool, mediaStore, postTurn = wireExtras(pgStores, agentRouter, providerRegistry, modelReg, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, redisClient, domainBus) + contextFileInterceptor, mcpPool, mediaStore, postTurn = wireExtras(pgStores, agentRouter, providerRegistry, modelReg, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, redisClient, domainBus, usageCapSvc) if mcpPool != nil { defer mcpPool.Stop() } diff --git a/cmd/gateway_agents.go b/cmd/gateway_agents.go index 0dcbda6b..2efec76a 100644 --- a/cmd/gateway_agents.go +++ b/cmd/gateway_agents.go @@ -15,6 +15,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tts" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // resolveEmbeddingProvider selects an embedding provider from DB only. @@ -159,7 +160,7 @@ func buildEmbeddingProvider( return nil } -func setupSubagents(providerReg *providers.Registry, cfg *config.Config, msgBus *bus.MessageBus, toolsReg *tools.Registry, workspace string, sandboxMgr sandbox.Manager, secureCLIStore store.SecureCLIStore) *tools.SubagentManager { +func setupSubagents(providerReg *providers.Registry, cfg *config.Config, msgBus *bus.MessageBus, toolsReg *tools.Registry, workspace string, sandboxMgr sandbox.Manager, secureCLIStore store.SecureCLIStore, usageCapSvc *usagecaps.Service) *tools.SubagentManager { names := providerReg.List(context.Background()) if len(names) == 0 { return nil @@ -207,7 +208,9 @@ func setupSubagents(providerReg *providers.Registry, cfg *config.Config, msgBus return reg } - return tools.NewSubagentManager(provider, providerReg, agentCfg.Model, msgBus, toolsFactory, subCfg) + manager := tools.NewSubagentManager(provider, providerReg, agentCfg.Model, msgBus, toolsFactory, subCfg) + manager.SetUsageCapService(usageCapSvc) + return manager } // buildSubagentToolsRegistry produces a cloned tool registry for a subagent diff --git a/cmd/gateway_http_wiring.go b/cmd/gateway_http_wiring.go index 94d52245..76c782a3 100644 --- a/cmd/gateway_http_wiring.go +++ b/cmd/gateway_http_wiring.go @@ -134,6 +134,9 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer( if d.pgStores.Snapshots != nil { d.server.SetUsageHandler(httpapi.NewUsageHandler(d.pgStores.Snapshots, d.pgStores.DB)) } + if d.pgStores.UsageCaps != nil { + d.server.SetUsageCapsHandler(httpapi.NewUsageCapsHandler(d.pgStores.UsageCaps, d.pgStores.Tenants)) + } // Runtime package management (install/uninstall system/pip/npm/github packages) // Wire the update registry AFTER initGitHubInstaller so DefaultGitHubInstaller() is set. diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index 18411a0f..0db67050 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -31,6 +31,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store/pg" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -58,6 +59,7 @@ func wireExtras( sandboxMgr sandbox.Manager, redisClient any, // nil when built without -tags redis or when Redis is unconfigured domainBus eventbus.DomainEventBus, + usageCapSvc *usagecaps.Service, ) (*tools.ContextFileInterceptor, *mcpbridge.Pool, *media.Store, tools.PostTurnProcessor) { // 1. Build cache instances (in-memory or Redis depending on build tags) agentCtxCache, userCtxCache := makeCaches(redisClient) @@ -82,9 +84,15 @@ func wireExtras( } } // Register media analysis tools (need mediaStore for file access). - toolsReg.Register(tools.NewReadDocumentTool(providerReg, mediaStore)) - toolsReg.Register(tools.NewReadAudioTool(providerReg, mediaStore)) - toolsReg.Register(tools.NewReadVideoTool(providerReg, mediaStore)) + readDocumentTool := tools.NewReadDocumentTool(providerReg, mediaStore) + readDocumentTool.SetUsageCapService(usageCapSvc) + toolsReg.Register(readDocumentTool) + readAudioTool := tools.NewReadAudioTool(providerReg, mediaStore) + readAudioTool.SetUsageCapService(usageCapSvc) + toolsReg.Register(readAudioTool) + readVideoTool := tools.NewReadVideoTool(providerReg, mediaStore) + readVideoTool.SetUsageCapService(usageCapSvc) + toolsReg.Register(readVideoTool) toolsReg.Register(tools.NewCreateVideoTool(providerReg)) slog.Info("media tools registered", "tools", "read_document,read_audio,read_video,create_video") } @@ -228,6 +236,7 @@ func wireExtras( MediaStore: mediaStore, ModelPricing: appCfg.Telemetry.ModelPricing, TracingStore: stores.Tracing, + UsageCaps: usageCapSvc, MemoryStore: stores.Memory, ContactStore: stores.Contacts, TenantStore: stores.Tenants, diff --git a/docs/02-providers.md b/docs/02-providers.md index b9633d0f..003b2e5a 100644 --- a/docs/02-providers.md +++ b/docs/02-providers.md @@ -54,6 +54,27 @@ 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. + +Excluded provider classes: +- `chatgpt_oauth`, `claude_cli`, and `bailian` are treated as subscription/non-API pricing in round one. +- 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. +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. + +Token-only policies do not require catalog pricing. Model fallback routes reserve against the actual candidate provider/model before each attempt. Cached input is separated from uncached input for OpenAI-compatible usage accounting. Partial stream failures keep the estimate, or actual provider usage when available, instead of clearing billed output to zero. Internal LLM calls for memory flush, compaction, media reading tools (`read_image`, `read_document`, `read_audio`, `read_video`), and subagents use the same preflight/reconcile path. + +Supported price units: input, output, cache read, cache write, reasoning, request, image, and web search. + +--- + ## 2. Supported Providers ### Six Core Provider Types diff --git a/docs/06-store-data-model.md b/docs/06-store-data-model.md index eddc0118..eb826417 100644 --- a/docs/06-store-data-model.md +++ b/docs/06-store-data-model.md @@ -45,6 +45,7 @@ The `Stores` struct is the top-level container holding all PostgreSQL-backed sto | ContactStore | `PGContactStore` | Channel contacts (auto-collected), cross-channel deduplication, merge | | ActivityStore | `PGActivityStore` | Audit logs, action tracking, compliance | | SnapshotStore | `PGSnapshotStore` | Hourly usage snapshots, cost aggregation, time series queries | +| UsageCapStore | `PGUsageCapStore` | OpenRouter pricing catalog, pricing overrides, cap policies, reservations, counters, events | | SecureCLIStore | `PGSecureCLIStore` | CLI binary configs with encrypted credential injection | | APIKeyStore | `PGAPIKeyStore` | Gateway API keys, scopes, expiration, revocation | | HookStore | `PGHookStore` | Lifecycle hook definitions (event, handler type, matcher, config), execution audit log | @@ -85,6 +86,31 @@ Migration versions: --- +## Usage Cap Storage + +Usage cap enforcement is Standard/PostgreSQL-only in round one. The `UsageCapStore` is wired on the PostgreSQL store factory and left nil in SQLite/Lite builds. + +Tables: +- `usage_pricing_catalog`: OpenRouter model catalog prices, raw upstream model payload, sync time. +- `usage_pricing_overrides`: tenant/provider/model override prices for custom billing assumptions. +- `usage_cap_policies`: cap definitions scoped by tenant, agent, provider, provider type, model, and window. +- `usage_cap_counters`: current window used and reserved token/cost counters. +- `usage_cap_reservations`: preflight reservations keyed by LLM call attempt. +- `usage_cap_events`: allow/block/reconcile/skip audit events. + +Reservation updates are atomic: counters are updated only when `used + reserved + estimate` remains below configured token and cost ceilings. +Reservation keys are idempotent per policy, so a retry using the same key does not double-increment reserved counters. +Policy `agent_id` references must belong to the same tenant as the policy. Policy and pricing override `provider_id` references may belong to the same tenant or the master tenant for default provider fallback, but not another non-master tenant. +Catalog and override price fields are nullable decimal strings with non-negative validation in the store layer and database checks. +Pricing resolution checks exact override/catalog model IDs first, then provider-derived OpenRouter aliases for native unprefixed model IDs. + +Migration versions: + +- PostgreSQL: `000069_usage_caps_pricing`, `000070_usage_cap_policies`. +- SQLite: no schema change; feature is not active in Lite. + +--- + ## 3. Session Caching The session store uses an in-memory write-behind cache to minimize database I/O during the agent tool loop. All reads and writes happen in memory; data is flushed to the persistent backend only when `Save()` is called at the end of a run. diff --git a/docs/18-http-api.md b/docs/18-http-api.md index 333be42d..f24f4d41 100644 --- a/docs/18-http-api.md +++ b/docs/18-http-api.md @@ -357,6 +357,36 @@ LLM provider management. API keys are encrypted with AES-256-GCM in the database | `GET` | `/v1/embedding/status` | Check global embedding availability | | `GET` | `/v1/providers/claude-cli/auth-status` | Check Claude CLI login status | +### Model Pricing + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/v1/model-pricing/sync-openrouter` | Sync OpenRouter `/models` pricing catalog (master scope) | +| `GET` | `/v1/model-pricing` | Search catalog models by `model` query | +| `GET` | `/v1/model-pricing/overrides` | List tenant pricing overrides, optionally by `provider_id` | +| `PUT` | `/v1/model-pricing/overrides` | Upsert provider/model custom pricing | +| `DELETE` | `/v1/model-pricing/overrides/{id}` | Delete pricing override | + +Override body: + +```json +{ + "provider_id": "0193a5b0-7000-7000-8000-000000000123", + "provider_type": "openrouter", + "model_id": "anthropic/claude-sonnet-4-5", + "pricing": { + "input": "0.000003", + "output": "0.000015", + "cache_read": "0.0000003", + "cache_write": "0.00000375", + "reasoning": "0.000015", + "request": "0", + "image": "0", + "web_search": "0" + } +} +``` + **Supported types:** `anthropic_native`, `openai_compat`, `chatgpt_oauth`, `gemini_native`, `dashscope`, `bailian`, `minimax`, `claude_cli`, `acp` Reconnect response: @@ -1238,11 +1268,32 @@ Follow response: | `GET` | `/v1/usage/timeseries` | Time-series usage points | | `GET` | `/v1/usage/breakdown` | Breakdown by provider/model/channel | | `GET` | `/v1/usage/summary` | Summary with period comparison | +| `GET` | `/v1/usage-caps/policies` | List usage cap policies | +| `POST` | `/v1/usage-caps/policies` | Create token/cost cap policy | +| `PATCH` | `/v1/usage-caps/policies/{id}` | Update cap policy | +| `DELETE` | `/v1/usage-caps/policies/{id}` | Delete cap policy | +| `GET` | `/v1/usage-caps/utilization` | Current-window used and reserved counters | +| `GET` | `/v1/usage-caps/events` | Recent allow/block/reconcile/skip events | **Query params:** `from`, `to` (RFC 3339), `agent_id`, `provider`, `model`, `channel`, `group_by` **Periods:** `24h`, `today`, `7d`, `30d` +Usage cap policy body: + +```json +{ + "agent_id": "optional-agent-uuid", + "provider_id": "optional-provider-uuid", + "provider_type": "openrouter", + "model_id": "anthropic/claude-sonnet-4-5", + "window": "day", + "max_tokens": 500000, + "max_cost_usd": 25, + "enabled": true +} +``` + --- ## 24. Activity & Audit diff --git a/docs/project-changelog.md b/docs/project-changelog.md index bd9954e2..dd7bca47 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -6,6 +6,23 @@ Significant changes, features, and fixes in reverse chronological order. ## 2026-05-22 +### Usage Cap budget controls + +**Features** + +- Added Standard/PostgreSQL usage caps for AI budget control by hour, day, week, or month. +- Caps support token and USD cost ceilings at tenant, agent, provider, provider type, and model scopes. +- Added OpenRouter catalog sync plus tenant/provider/model pricing overrides for input, output, cache read/write, reasoning, request, image, and web search units. +- Enforced caps in agent, fallback model, subagent, memory flush, compaction, and media reading tools (`read_image`, `read_document`, `read_audio`, `read_video`) with preflight reservation and post-call reconciliation. +- Added non-negative validation for catalog and override pricing fields. +- Added OpenRouter alias resolution for native model IDs, cached-input accounting normalization, and partial-stream failure reconciliation. +- Added web dashboard controls on Usage and Provider detail pages. + +**Tests** + +- Added pricing and cap service coverage. +- Verified Go builds, SQLite build compatibility, full Go test suite, integration race suite, and Web UI production build. + ### CLI P6 backend API unblock **Features** diff --git a/internal/agent/loop_compact.go b/internal/agent/loop_compact.go index c6c31753..7adb6382 100644 --- a/internal/agent/loop_compact.go +++ b/internal/agent/loop_compact.go @@ -89,14 +89,15 @@ func (l *Loop) compactMessagesInPlace(ctx context.Context, messages []providers. inTokens := l.estimateSummaryInputTokens(toSummarize) slog.Info("compact_budget", "agent", l.id, "in_tokens", inTokens, "out_tokens", dynamicSummaryMax(inTokens)) - resp, err := l.provider.Chat(sctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{{ Role: "user", Content: compactionSummaryPrompt + sb.String(), }}, Model: l.model, Options: map[string]any{"max_tokens": dynamicSummaryMax(inTokens), "temperature": 0.3}, - }) + } + resp, err := l.callInternalLLMWithUsage(sctx, chatReq, "mid-loop-compaction") if err != nil { slog.Warn("mid_loop_compaction_failed", "agent", l.id, "error", err) return nil diff --git a/internal/agent/loop_history_sanitize.go b/internal/agent/loop_history_sanitize.go index 6e0834b8..706f0dc9 100644 --- a/internal/agent/loop_history_sanitize.go +++ b/internal/agent/loop_history_sanitize.go @@ -284,11 +284,12 @@ func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { inTokens := l.estimateSummaryInputTokens(toSummarize) slog.Info("compact_budget", "agent", l.id, "in_tokens", inTokens, "out_tokens", dynamicSummaryMax(inTokens)) - resp, err := l.provider.Chat(sctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{{Role: "user", Content: prompt.String()}}, Model: l.model, Options: map[string]any{"max_tokens": dynamicSummaryMax(inTokens), "temperature": 0.3}, - }) + } + resp, err := l.callInternalLLMWithUsage(sctx, chatReq, "session-summarization") if err != nil { slog.Warn("summarization failed", "session", sessionKey, "error", err) return diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index c8abda9a..6befa4d1 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -2,6 +2,7 @@ package agent import ( "context" + "fmt" "log/slog" "strings" "time" @@ -15,6 +16,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/internal/workspace" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -294,30 +296,75 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } spanID := l.emitLLMSpanStart(ctx, start, state.Iteration+1, chatReq.Messages, opts...) - var resp *providers.ChatResponse - var err error - if req.Stream { - resp, err = provider.ChatStream(ctx, chatReq, func(chunk providers.StreamChunk) { - if chunk.Thinking != "" { - emitRun(AgentEvent{ - Type: protocol.ChatEventThinking, - AgentID: l.id, - RunID: req.RunID, - Payload: map[string]string{"content": chunk.Thinking}, - }) - } - if chunk.Content != "" { - emitRun(AgentEvent{ - Type: protocol.ChatEventChunk, - AgentID: l.id, - RunID: req.RunID, - Payload: map[string]string{"content": chunk.Content}, - }) - } - }) - } else { - resp, err = provider.Chat(ctx, chatReq) + emitChunk := func(chunk providers.StreamChunk) { + if chunk.Thinking != "" { + emitRun(AgentEvent{ + Type: protocol.ChatEventThinking, + AgentID: l.id, + RunID: req.RunID, + Payload: map[string]string{"content": chunk.Thinking}, + }) + } + if chunk.Content != "" { + emitRun(AgentEvent{ + Type: protocol.ChatEventChunk, + AgentID: l.id, + RunID: req.RunID, + Payload: map[string]string{"content": chunk.Content}, + }) + } } + callProvider := func(attempt string, request providers.ChatRequest) (*providers.ChatResponse, error) { + if fallbackProvider, ok := provider.(*providers.ModelFallbackProvider); ok { + before := func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + candidateAttempt := fmt.Sprintf("%s:%s:%s", attempt, entry.ProviderName, actualReq.Model) + reservation, reserveErr := l.reserveLLMUsageFor(callCtx, req, state.Iteration, actualReq, candidateAttempt, entry.ProviderName, actualReq.Model) + if reserveErr != nil { + return nil, reserveErr + } + return func(callResp *providers.ChatResponse, callErr error, info providers.FallbackCallInfo) { + if reservation != nil { + if info.Streamed { + reservation.ReconcileStream(callCtx, callResp, callErr, true) + } else { + reservation.Reconcile(callCtx, callResp, callErr) + } + } + }, nil + } + if req.Stream { + return fallbackProvider.ChatStreamWithHook(ctx, request, emitChunk, before) + } + return fallbackProvider.ChatWithHook(ctx, request, before) + } + reservation, reserveErr := l.reserveLLMUsage(ctx, req, state, request, attempt) + if reserveErr != nil { + return nil, reserveErr + } + var callResp *providers.ChatResponse + var callErr error + if req.Stream { + streamed := false + callResp, callErr = provider.ChatStream(ctx, request, func(chunk providers.StreamChunk) { + if chunk.Content != "" || chunk.Thinking != "" || len(chunk.Images) > 0 { + streamed = true + } + emitChunk(chunk) + }) + if reservation != nil { + reservation.ReconcileStream(ctx, callResp, callErr, streamed) + } + return callResp, callErr + } else { + callResp, callErr = provider.Chat(ctx, request) + } + if reservation != nil { + reservation.Reconcile(ctx, callResp, callErr) + } + return callResp, callErr + } + + resp, err := callProvider("initial", chatReq) slog.Info("debug.llm.first_response", "has_error", err != nil, "tool_calls_count", func() int { @@ -342,28 +389,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c Role: "system", Content: "MCP task tools are available in this turn. Do not ask for CRM identifier/email first. Call the relevant MCP task tool immediately, then answer with the tool result.", }) - if req.Stream { - resp, err = provider.ChatStream(ctx, retryReq, func(chunk providers.StreamChunk) { - if chunk.Thinking != "" { - emitRun(AgentEvent{ - Type: protocol.ChatEventThinking, - AgentID: l.id, - RunID: req.RunID, - Payload: map[string]string{"content": chunk.Thinking}, - }) - } - if chunk.Content != "" { - emitRun(AgentEvent{ - Type: protocol.ChatEventChunk, - AgentID: l.id, - RunID: req.RunID, - Payload: map[string]string{"content": chunk.Content}, - }) - } - }) - } else { - resp, err = provider.Chat(ctx, retryReq) - } + resp, err = callProvider("retry-tool-choice", retryReq) slog.Info("debug.llm.retry_response", "has_error", err != nil, "tool_calls_count", func() int { @@ -536,3 +562,30 @@ func (l *Loop) makeBootstrapCleanup() func(ctx context.Context, state *pipeline. return l.bootstrapCleanup(ctx, l.agentUUID, state.Input.UserID) } } + +func (l *Loop) reserveLLMUsage(ctx context.Context, req *RunRequest, state *pipeline.RunState, chatReq providers.ChatRequest, attempt string) (*usagecaps.Reservation, error) { + if l.usageCaps == nil || state.Provider == nil { + return nil, nil + } + return l.reserveLLMUsageFor(ctx, req, state.Iteration, chatReq, attempt, state.Provider.Name(), state.Model) +} + +func (l *Loop) reserveLLMUsageFor(ctx context.Context, req *RunRequest, iteration int, chatReq providers.ChatRequest, attempt, providerName, model string) (*usagecaps.Reservation, error) { + if l.usageCaps == nil { + return nil, nil + } + tenantID := store.TenantIDFromContext(ctx) + if tenantID == uuid.Nil { + tenantID = l.tenantID + } + key := fmt.Sprintf("%s:%s:%d:%s", req.RunID, l.agentUUID.String(), iteration+1, attempt) + return l.usageCaps.Preflight(ctx, usagecaps.Request{ + TenantID: tenantID, + AgentID: l.agentUUID, + ProviderName: providerName, + ModelID: model, + ReservationKey: key, + Messages: chatReq.Messages, + MaxOutputTokens: l.maxOutputTokensFromRequest(chatReq), + }) +} diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index 50c50228..c62d8c3d 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -23,6 +23,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/tokencount" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // bootstrapAutoCleanupTurns is the number of user messages after which @@ -80,13 +81,13 @@ type Loop struct { // agentUUID is the canonical DB primary key. Use for SQL WHERE/JOIN, // DomainEvent.AgentID, OTel span attributes, and context propagation via // store.WithAgentID. See docs/agent-identity-conventions.md. - agentUUID uuid.UUID - tenantID uuid.UUID // agent's owning tenant + agentUUID uuid.UUID + tenantID uuid.UUID // agent's owning tenant // agentOtherConfig is a defensive byte copy of agents.other_config JSONB. // Copied once at Loop construction; used to build AgentAudioSnapshot at tool dispatch. agentOtherConfig json.RawMessage - agentType string // "open" or "predefined" - defaultTimezone string // system default timezone for bootstrap pre-fill + agentType string // "open" or "predefined" + defaultTimezone string // system default timezone for bootstrap pre-fill provider providers.Provider model string modelRegistry providers.ModelRegistry // resolves per-model context window at run time (nil = use static contextWindow) @@ -108,7 +109,7 @@ type Loop struct { // Memory flush runs if callback != nil; auto-inject runs if AutoInjector != nil. autoInjector memory.AutoInjector // v3 L0 memory auto-inject (nil = disabled) - eventPub bus.EventPublisher // currently unused by Loop; kept for future use + eventPub bus.EventPublisher // currently unused by Loop; kept for future use domainBus eventbus.DomainEventBus // V3 domain event bus for consolidation pipeline sessions store.SessionStore tools tools.ToolExecutor @@ -137,11 +138,11 @@ type Loop struct { userSetups sync.Map // userID → *userSetup (workspace + seeding state, per Loop instance) // Per-user MCP tools: servers requiring user credentials get connected per-request. - mcpStore store.MCPServerStore // for credential lookup - mcpPool *mcpbridge.Pool // user-keyed connection pool - mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds - mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools) - mcpGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip) + mcpStore store.MCPServerStore // for credential lookup + mcpPool *mcpbridge.Pool // user-keyed connection pool + mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds + mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools) + mcpGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip) // Compaction config (memory flush settings) compactionCfg *config.CompactionConfig @@ -239,13 +240,14 @@ type Loop struct { // Budget enforcement: monthly spending limit in cents (0 = unlimited) budgetMonthlyCents int tracingStore store.TracingStore + usageCaps *usagecaps.Service // Memory store for extractive memory fallback (writes directly when LLM flush fails) memStore store.MemoryStore // v3 orchestration mode (spawn/delegate/team) — controls tool visibility - orchMode OrchestrationMode - delegateTargets []DelegateTargetEntry // delegation targets for prompt injection + orchMode OrchestrationMode + delegateTargets []DelegateTargetEntry // delegation targets for prompt injection // v3 evolution metrics store (nil = disabled) evolutionMetricsStore store.EvolutionMetricsStore @@ -350,11 +352,11 @@ type LoopConfig struct { // Agent UUID + tenant for context propagation to tools AgentUUID uuid.UUID - TenantID uuid.UUID // agent's owning tenant — injected into execution context - AgentOtherConfig json.RawMessage // raw other_config JSONB — copied defensively in NewLoop - AgentType string // "open" or "predefined" - DisplayName string // human-readable agent display name (for runtime section) - IsTeamLead bool // agent leads a team (from resolver detection) + TenantID uuid.UUID // agent's owning tenant — injected into execution context + AgentOtherConfig json.RawMessage // raw other_config JSONB — copied defensively in NewLoop + AgentType string // "open" or "predefined" + DisplayName string // human-readable agent display name (for runtime section) + IsTeamLead bool // agent leads a team (from resolver detection) // Per-user profile + file seeding + dynamic context loading EnsureUserProfile EnsureUserProfileFunc // preferred: separate profile + workspace @@ -430,19 +432,20 @@ type LoopConfig struct { // Budget enforcement BudgetMonthlyCents int TracingStore store.TracingStore + UsageCaps *usagecaps.Service // Memory store for extractive memory fallback (writes directly when LLM flush fails) MemoryStore store.MemoryStore // Per-user MCP tools (servers requiring per-user credentials) - MCPStore store.MCPServerStore // for credential lookup - MCPPool *mcpbridge.Pool // user-keyed connection pool - MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds - MCPGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip) + MCPStore store.MCPServerStore // for credential lookup + MCPPool *mcpbridge.Pool // user-keyed connection pool + MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds + MCPGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip) // V3 orchestration mode (resolved by resolver, controls tool visibility) - OrchMode OrchestrationMode - DelegateTargets []DelegateTargetEntry // delegation targets for prompt injection + OrchMode OrchestrationMode + DelegateTargets []DelegateTargetEntry // delegation targets for prompt injection // V3 evolution metrics store for recording tool/retrieval/feedback metrics EvolutionMetricsStore store.EvolutionMetricsStore @@ -568,6 +571,7 @@ func NewLoop(cfg LoopConfig) *Loop { modelPricing: cfg.ModelPricing, budgetMonthlyCents: cfg.BudgetMonthlyCents, tracingStore: cfg.TracingStore, + usageCaps: cfg.UsageCaps, memStore: cfg.MemoryStore, mcpStore: cfg.MCPStore, mcpPool: cfg.MCPPool, @@ -582,36 +586,36 @@ func NewLoop(cfg LoopConfig) *Loop { // RunRequest is the input for processing a message through the agent. type RunRequest struct { - SessionKey string // composite key: agent:{agentId}:{channel}:{peerKind}:{chatId} - Message string // user message - Media []bus.MediaFile // local media files with MIME types - ForwardMedia []bus.MediaFile // media files to forward to output (from delegation results) - Channel string // source channel instance name (e.g. "my-telegram-bot") - ChannelType string // platform type (e.g. "zalo_personal", "telegram") — for system prompt context - BitrixPortalDomain string // bitrix24-only: portal domain (e.g. "tamgiac.bitrix24.com") for entity URL construction - ChatTitle string // group chat display name (e.g. Telegram group title) - ChatID string // source chat ID - PeerKind string // "direct" or "group" (for session key building and tool context) - RunID string // unique run identifier - UserID string // external user ID (TEXT, free-form) for multi-tenant scoping - SenderID string // original individual sender ID (preserved in group chats for permission checks) - SenderName string // display name from channel metadata (for bootstrap auto-contact) - Role string // caller's RBAC role (admin/operator/viewer/owner); bypasses per-user grants for authenticated admins (#915) - Stream bool // whether to stream response chunks - ExtraSystemPrompt string // optional: injected into system prompt (skills, subagent context, etc.) - SkillFilter []string // per-request skill override: nil=use agent default, []=no skills, ["x","y"]=whitelist - HistoryLimit int // max user turns to keep in context (0=unlimited, from channel config) - ToolAllow []string // per-group tool allow list (nil = no restriction, supports "group:xxx") - LocalKey string // composite key with topic/thread suffix for routing (e.g. "-100123:topic:42") - ParentTraceID uuid.UUID // if set, reuse parent trace instead of creating new (announce runs) - ParentRootSpanID uuid.UUID // if set, nest announce agent span under this parent span - LinkedTraceID uuid.UUID // if set, create new trace with parent_trace_id pointing to this (team task runs) - TraceName string // override trace name (default: "chat ") - TraceTags []string // additional tags for the trace (e.g. "cron") - MaxIterations int // per-request override (0 = use agent default, must be lower) - ModelOverride string // per-request model override (heartbeat uses cheaper model) - ProviderOverride providers.Provider // per-request provider override (heartbeat uses different provider) - LightContext bool // skip loading context files (only inject ExtraSystemPrompt) + SessionKey string // composite key: agent:{agentId}:{channel}:{peerKind}:{chatId} + Message string // user message + Media []bus.MediaFile // local media files with MIME types + ForwardMedia []bus.MediaFile // media files to forward to output (from delegation results) + Channel string // source channel instance name (e.g. "my-telegram-bot") + ChannelType string // platform type (e.g. "zalo_personal", "telegram") — for system prompt context + BitrixPortalDomain string // bitrix24-only: portal domain (e.g. "tamgiac.bitrix24.com") for entity URL construction + ChatTitle string // group chat display name (e.g. Telegram group title) + ChatID string // source chat ID + PeerKind string // "direct" or "group" (for session key building and tool context) + RunID string // unique run identifier + UserID string // external user ID (TEXT, free-form) for multi-tenant scoping + SenderID string // original individual sender ID (preserved in group chats for permission checks) + SenderName string // display name from channel metadata (for bootstrap auto-contact) + Role string // caller's RBAC role (admin/operator/viewer/owner); bypasses per-user grants for authenticated admins (#915) + Stream bool // whether to stream response chunks + ExtraSystemPrompt string // optional: injected into system prompt (skills, subagent context, etc.) + SkillFilter []string // per-request skill override: nil=use agent default, []=no skills, ["x","y"]=whitelist + HistoryLimit int // max user turns to keep in context (0=unlimited, from channel config) + ToolAllow []string // per-group tool allow list (nil = no restriction, supports "group:xxx") + LocalKey string // composite key with topic/thread suffix for routing (e.g. "-100123:topic:42") + ParentTraceID uuid.UUID // if set, reuse parent trace instead of creating new (announce runs) + ParentRootSpanID uuid.UUID // if set, nest announce agent span under this parent span + LinkedTraceID uuid.UUID // if set, create new trace with parent_trace_id pointing to this (team task runs) + TraceName string // override trace name (default: "chat ") + TraceTags []string // additional tags for the trace (e.g. "cron") + MaxIterations int // per-request override (0 = use agent default, must be lower) + ModelOverride string // per-request model override (heartbeat uses cheaper model) + ProviderOverride providers.Provider // per-request provider override (heartbeat uses different provider) + LightContext bool // skip loading context files (only inject ExtraSystemPrompt) // Run classification RunKind string // "delegation", "announce" — empty for user-initiated runs @@ -646,7 +650,7 @@ type RunRequest struct { // RunResult is the output of a completed agent run. type RunResult struct { Content string `json:"content"` - Thinking string `json:"thinking,omitempty"` // reasoning content from thinking models (Claude, o3, DeepSeek-R1, Kimi) + Thinking string `json:"thinking,omitempty"` // reasoning content from thinking models (Claude, o3, DeepSeek-R1, Kimi) RunID string `json:"runId"` Iterations int `json:"iterations"` Usage *providers.Usage `json:"usage,omitempty"` diff --git a/internal/agent/memoryflush.go b/internal/agent/memoryflush.go index d5106100..7db493a0 100644 --- a/internal/agent/memoryflush.go +++ b/internal/agent/memoryflush.go @@ -186,7 +186,7 @@ func (l *Loop) runMemoryFlush(ctx context.Context, sessionKey string, settings * // Run LLM iteration loop (max 5 iterations for flush) maxFlushIter := 5 for range maxFlushIter { - resp, err := l.provider.Chat(flushCtx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: messages, Tools: toolDefs, Model: l.model, @@ -194,7 +194,8 @@ func (l *Loop) runMemoryFlush(ctx context.Context, sessionKey string, settings * "max_tokens": 4096, "temperature": 0.3, }, - }) + } + resp, err := l.callInternalLLMWithUsage(flushCtx, chatReq, "memory-flush") if err != nil { slog.Warn("memory flush: LLM call failed", "error", err) l.extractiveMemoryFallback(flushCtx, sessionKey, history, "LLM error") diff --git a/internal/agent/resolver.go b/internal/agent/resolver.go index afd3a936..5a82fb8d 100644 --- a/internal/agent/resolver.go +++ b/internal/agent/resolver.go @@ -25,6 +25,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" ) // ResolverDeps holds shared dependencies for the agent resolver. @@ -97,6 +98,7 @@ type ResolverDeps struct { // Tracing store for budget enforcement queries TracingStore store.TracingStore + UsageCaps *usagecaps.Service // Memory store for extractive memory fallback MemoryStore store.MemoryStore @@ -528,6 +530,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { ModelPricing: deps.ModelPricing, BudgetMonthlyCents: derefInt(ag.BudgetMonthlyCents), TracingStore: deps.TracingStore, + UsageCaps: deps.UsageCaps, MemoryStore: deps.MemoryStore, MCPStore: deps.MCPStore, MCPPool: deps.MCPPool, diff --git a/internal/agent/usage_caps_runtime.go b/internal/agent/usage_caps_runtime.go new file mode 100644 index 00000000..e883a9c3 --- /dev/null +++ b/internal/agent/usage_caps_runtime.go @@ -0,0 +1,88 @@ +package agent + +import ( + "context" + "fmt" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" +) + +func (l *Loop) reserveInternalLLMUsage(ctx context.Context, chatReq providers.ChatRequest, purpose string) (*usagecaps.Reservation, error) { + if l.usageCaps == nil || l.provider == nil { + return nil, nil + } + return l.reserveInternalLLMUsageFor(ctx, chatReq, purpose, l.provider.Name(), chatReq.Model) +} + +func (l *Loop) reserveInternalLLMUsageFor(ctx context.Context, chatReq providers.ChatRequest, purpose, providerName, model string) (*usagecaps.Reservation, error) { + if l.usageCaps == nil { + return nil, nil + } + if model == "" { + model = chatReq.Model + } + if model == "" { + model = l.model + } + tenantID := store.TenantIDFromContext(ctx) + if tenantID == uuid.Nil { + tenantID = l.tenantID + } + return l.usageCaps.Preflight(ctx, usagecaps.Request{ + TenantID: tenantID, + AgentID: l.agentUUID, + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("%s:%s:%s", purpose, l.agentUUID.String(), uuid.NewString()), + Messages: chatReq.Messages, + MaxOutputTokens: l.maxOutputTokensFromRequest(chatReq), + }) +} + +func (l *Loop) callInternalLLMWithUsage(ctx context.Context, chatReq providers.ChatRequest, purpose string) (*providers.ChatResponse, error) { + if fallbackProvider, ok := l.provider.(*providers.ModelFallbackProvider); ok { + before := func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + candidatePurpose := fmt.Sprintf("%s:%s:%s", purpose, entry.ProviderName, actualReq.Model) + reservation, err := l.reserveInternalLLMUsageFor(callCtx, actualReq, candidatePurpose, entry.ProviderName, actualReq.Model) + if err != nil { + return nil, err + } + return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) { + if reservation != nil { + reservation.Reconcile(callCtx, resp, callErr) + } + }, nil + } + return fallbackProvider.ChatWithHook(ctx, chatReq, before) + } + reservation, reserveErr := l.reserveInternalLLMUsage(ctx, chatReq, purpose) + if reserveErr != nil { + return nil, reserveErr + } + resp, err := l.provider.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } + return resp, err +} + +func (l *Loop) maxOutputTokensFromRequest(chatReq providers.ChatRequest) int { + maxTokens := l.effectiveMaxTokens() + if chatReq.Options == nil { + return maxTokens + } + if v, ok := chatReq.Options[providers.OptMaxTokens]; ok { + switch n := v.(type) { + case int: + maxTokens = n + case int64: + maxTokens = int(n) + case float64: + maxTokens = int(n) + } + } + return maxTokens +} diff --git a/internal/gateway/server.go b/internal/gateway/server.go index 5301e58c..db0d617d 100644 --- a/internal/gateway/server.go +++ b/internal/gateway/server.go @@ -634,6 +634,11 @@ func (s *Server) SetSystemConfigsHandler(h *httpapi.SystemConfigsHandler) { // SetUsageHandler sets the usage analytics handler. func (s *Server) SetUsageHandler(h *httpapi.UsageHandler) { s.handlers = append(s.handlers, h) } +// SetUsageCapsHandler sets usage cap and model pricing handlers. +func (s *Server) SetUsageCapsHandler(h *httpapi.UsageCapsHandler) { + s.handlers = append(s.handlers, h) +} + // SetBackupHandler sets the system backup handler. func (s *Server) SetBackupHandler(h *httpapi.BackupHandler) { s.handlers = append(s.handlers, h) } diff --git a/internal/http/usage_caps.go b/internal/http/usage_caps.go new file mode 100644 index 00000000..7e027528 --- /dev/null +++ b/internal/http/usage_caps.go @@ -0,0 +1,354 @@ +package http + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "strconv" + "time" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/permissions" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/usage/pricing" +) + +type UsageCapsHandler struct { + store store.UsageCapStore + tenants store.TenantStore + client *http.Client +} + +func NewUsageCapsHandler(s store.UsageCapStore, tenants store.TenantStore) *UsageCapsHandler { + return &UsageCapsHandler{store: s, tenants: tenants, client: &http.Client{Timeout: 30 * time.Second}} +} + +func (h *UsageCapsHandler) RegisterRoutes(mux *http.ServeMux) { + mux.HandleFunc("GET /v1/usage-caps/policies", h.auth(h.handleListPolicies)) + mux.HandleFunc("POST /v1/usage-caps/policies", h.writeAuth(h.handleCreatePolicy)) + mux.HandleFunc("PATCH /v1/usage-caps/policies/{id}", h.writeAuth(h.handleUpdatePolicy)) + mux.HandleFunc("DELETE /v1/usage-caps/policies/{id}", h.writeAuth(h.handleDeletePolicy)) + mux.HandleFunc("GET /v1/usage-caps/utilization", h.auth(h.handleUtilization)) + mux.HandleFunc("GET /v1/usage-caps/events", h.auth(h.handleEvents)) + mux.HandleFunc("POST /v1/model-pricing/sync-openrouter", h.masterAuth(h.handleSyncOpenRouter)) + mux.HandleFunc("GET /v1/model-pricing", h.auth(h.handleListPricing)) + mux.HandleFunc("PUT /v1/model-pricing/overrides", h.writeAuth(h.handlePutOverride)) + mux.HandleFunc("GET /v1/model-pricing/overrides", h.auth(h.handleListOverrides)) + mux.HandleFunc("DELETE /v1/model-pricing/overrides/{id}", h.writeAuth(h.handleDeleteOverride)) +} + +func (h *UsageCapsHandler) auth(next http.HandlerFunc) http.HandlerFunc { + return requireAuth(permissions.RoleAdmin, next) +} + +func (h *UsageCapsHandler) writeAuth(next http.HandlerFunc) http.HandlerFunc { + return h.auth(func(w http.ResponseWriter, r *http.Request) { + if !requireTenantAdmin(w, r, h.tenants) { + return + } + next(w, r) + }) +} + +func (h *UsageCapsHandler) masterAuth(next http.HandlerFunc) http.HandlerFunc { + return h.auth(func(w http.ResponseWriter, r *http.Request) { + if !requireMasterScope(w, r) { + return + } + next(w, r) + }) +} + +func (h *UsageCapsHandler) handleListPolicies(w http.ResponseWriter, r *http.Request) { + scope := store.UsageCapScope{TenantID: tenantIDOrMaster(r)} + policies, err := h.store.ListUsageCapPolicies(r.Context(), scope, true) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "list policies failed"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"policies": policies}) +} + +func (h *UsageCapsHandler) handleCreatePolicy(w http.ResponseWriter, r *http.Request) { + var body policyBody + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) + return + } + p, err := body.toPolicy(tenantIDOrMaster(r)) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + return + } + if err := h.store.CreateUsageCapPolicy(r.Context(), &p); err != nil { + slog.Warn("usage_caps.create_policy_failed", "error", err) + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "usage cap policy validation failed"}) + return + } + writeJSON(w, http.StatusCreated, p) +} + +func (h *UsageCapsHandler) handleUpdatePolicy(w http.ResponseWriter, r *http.Request) { + id, err := uuid.Parse(r.PathValue("id")) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid policy id"}) + return + } + var body policyBody + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) + return + } + patch, err := body.toPatch() + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + return + } + p, err := h.store.UpdateUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id, patch) + if err != nil { + slog.Warn("usage_caps.update_policy_failed", "error", err) + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "usage cap policy validation failed"}) + return + } + writeJSON(w, http.StatusOK, p) +} + +func (h *UsageCapsHandler) handleDeletePolicy(w http.ResponseWriter, r *http.Request) { + id, err := uuid.Parse(r.PathValue("id")) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid policy id"}) + return + } + if err := h.store.DeleteUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id); err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "delete failed"}) + return + } + w.WriteHeader(http.StatusNoContent) +} + +func (h *UsageCapsHandler) handleUtilization(w http.ResponseWriter, r *http.Request) { + rows, err := h.store.ListUsageCapUtilization(r.Context(), tenantIDOrMaster(r)) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "utilization failed"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"rows": rows}) +} + +func (h *UsageCapsHandler) handleEvents(w http.ResponseWriter, r *http.Request) { + limit, _ := strconv.Atoi(r.URL.Query().Get("limit")) + events, err := h.store.ListUsageCapEvents(r.Context(), tenantIDOrMaster(r), limit) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "events failed"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"events": events}) +} + +func (h *UsageCapsHandler) handleSyncOpenRouter(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithTimeout(r.Context(), 45*time.Second) + defer cancel() + entries, err := pricing.FetchOpenRouterCatalog(ctx, h.client) + if err != nil { + slog.Warn("usage_pricing.openrouter_sync", "error", err) + writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()}) + return + } + count, err := h.store.UpsertPricingCatalog(r.Context(), entries) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "store catalog failed"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"count": count}) +} + +func (h *UsageCapsHandler) handleListPricing(w http.ResponseWriter, r *http.Request) { + rows, err := h.store.ListPricingCatalog(r.Context(), store.UsagePricingQuery{ + ModelID: r.URL.Query().Get("model"), + Limit: queryInt(r, "limit", 100), + }) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "list pricing failed"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"models": rows}) +} + +func (h *UsageCapsHandler) handlePutOverride(w http.ResponseWriter, r *http.Request) { + var body overrideBody + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) + return + } + providerID, err := uuid.Parse(body.ProviderID) + if err != nil || body.ModelID == "" { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "provider_id and model_id are required"}) + return + } + o := &store.UsagePricingOverride{ + TenantID: tenantIDOrMaster(r), ProviderID: providerID, + ProviderType: body.ProviderType, ModelID: body.ModelID, + Pricing: body.Pricing, Enabled: body.Enabled == nil || *body.Enabled, + } + if err := h.store.PutPricingOverride(r.Context(), o); err != nil { + slog.Warn("usage_pricing.put_override_failed", "error", err) + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "pricing override validation failed"}) + return + } + writeJSON(w, http.StatusOK, o) +} + +func (h *UsageCapsHandler) handleListOverrides(w http.ResponseWriter, r *http.Request) { + var providerID uuid.UUID + if raw := r.URL.Query().Get("provider_id"); raw != "" { + var err error + providerID, err = uuid.Parse(raw) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid provider_id"}) + return + } + } + rows, err := h.store.ListPricingOverrides(r.Context(), store.UsagePricingQuery{TenantID: tenantIDOrMaster(r), ProviderID: providerID}) + if err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "list overrides failed"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"overrides": rows}) +} + +func (h *UsageCapsHandler) handleDeleteOverride(w http.ResponseWriter, r *http.Request) { + id, err := uuid.Parse(r.PathValue("id")) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid override id"}) + return + } + if err := h.store.DeletePricingOverride(r.Context(), tenantIDOrMaster(r), id); err != nil { + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "delete failed"}) + return + } + w.WriteHeader(http.StatusNoContent) +} + +type policyBody struct { + AgentID *string `json:"agent_id"` + ProviderID *string `json:"provider_id"` + ProviderType *string `json:"provider_type"` + ModelID *string `json:"model_id"` + Window string `json:"window"` + MaxTokens *int64 `json:"max_tokens"` + MaxCostMicros *int64 `json:"max_cost_micros"` + MaxCostUSD *float64 `json:"max_cost_usd"` + Enabled *bool `json:"enabled"` + Priority *int `json:"priority"` +} + +type overrideBody struct { + ProviderID string `json:"provider_id"` + ProviderType string `json:"provider_type"` + ModelID string `json:"model_id"` + Pricing store.UsagePricingFields `json:"pricing"` + Enabled *bool `json:"enabled"` +} + +func (b policyBody) toPolicy(tenantID uuid.UUID) (store.UsageCapPolicy, error) { + p := store.UsageCapPolicy{TenantID: tenantID, Window: b.Window, Enabled: true, Priority: 100} + if b.Enabled != nil { + p.Enabled = *b.Enabled + } + if b.Priority != nil { + p.Priority = *b.Priority + } + if p.Window == "" { + p.Window = store.UsageCapWindowDay + } + var err error + p.AgentID, err = parseOptionalUUIDStrict("agent_id", b.AgentID) + if err != nil { + return p, err + } + p.ProviderID, err = parseOptionalUUIDStrict("provider_id", b.ProviderID) + if err != nil { + return p, err + } + if b.ProviderType != nil { + p.ProviderType = *b.ProviderType + } + if b.ModelID != nil { + p.ModelID = *b.ModelID + } + p.MaxTokens = b.MaxTokens + p.MaxCostMicros = maxCostMicros(b.MaxCostMicros, b.MaxCostUSD) + return p, nil +} + +func (b policyBody) toPatch() (store.UsageCapPolicyPatch, error) { + var patch store.UsageCapPolicyPatch + if b.AgentID != nil { + v, err := parseOptionalUUIDStrict("agent_id", b.AgentID) + if err != nil { + return patch, err + } + patch.AgentID = &v + } + if b.ProviderID != nil { + v, err := parseOptionalUUIDStrict("provider_id", b.ProviderID) + if err != nil { + return patch, err + } + patch.ProviderID = &v + } + patch.ProviderType = b.ProviderType + patch.ModelID = b.ModelID + if b.Window != "" { + patch.Window = &b.Window + } + if b.MaxTokens != nil { + patch.MaxTokens = &b.MaxTokens + } + if b.MaxCostMicros != nil || b.MaxCostUSD != nil { + v := maxCostMicros(b.MaxCostMicros, b.MaxCostUSD) + patch.MaxCostMicros = &v + } + patch.Enabled = b.Enabled + patch.Priority = b.Priority + return patch, nil +} + +func tenantIDOrMaster(r *http.Request) uuid.UUID { + if tid := store.TenantIDFromContext(r.Context()); tid != uuid.Nil { + return tid + } + return store.MasterTenantID +} + +func queryInt(r *http.Request, key string, fallback int) int { + n, err := strconv.Atoi(r.URL.Query().Get(key)) + if err != nil || n <= 0 { + return fallback + } + return n +} + +func parseOptionalUUIDStrict(field string, raw *string) (*uuid.UUID, error) { + if raw == nil || *raw == "" { + return nil, nil + } + id, err := uuid.Parse(*raw) + if err != nil { + return nil, fmt.Errorf("invalid %s", field) + } + return &id, nil +} + +func maxCostMicros(micros *int64, usd *float64) *int64 { + if micros != nil { + return micros + } + if usd == nil { + return nil + } + v := int64(*usd * 1_000_000) + return &v +} diff --git a/internal/http/usage_caps_test.go b/internal/http/usage_caps_test.go new file mode 100644 index 00000000..3923972a --- /dev/null +++ b/internal/http/usage_caps_test.go @@ -0,0 +1,17 @@ +package http + +import ( + "testing" + + "github.com/google/uuid" +) + +func TestUsageCapPolicyBodyRejectsInvalidUUID(t *testing.T) { + invalid := "not-a-uuid" + if _, err := (policyBody{AgentID: &invalid}).toPolicy(uuid.New()); err == nil { + t.Fatal("toPolicy accepted invalid agent_id") + } + if _, err := (policyBody{ProviderID: &invalid}).toPatch(); err == nil { + t.Fatal("toPatch accepted invalid provider_id") + } +} diff --git a/internal/providers/anthropic_stream.go b/internal/providers/anthropic_stream.go index f9936c32..0d2451c0 100644 --- a/internal/providers/anthropic_stream.go +++ b/internal/providers/anthropic_stream.go @@ -150,7 +150,7 @@ func (p *AnthropicProvider) ChatStream(ctx context.Context, req ChatRequest, onC } if err := sse.Err(); err != nil { - return nil, fmt.Errorf("anthropic stream read error: %w", err) + return result, fmt.Errorf("anthropic stream read error: %w", err) } // Parse accumulated tool call JSON arguments diff --git a/internal/providers/codex.go b/internal/providers/codex.go index e876ff20..8687cbb5 100644 --- a/internal/providers/codex.go +++ b/internal/providers/codex.go @@ -168,7 +168,7 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk } if err := sse.Err(); err != nil { - return nil, fmt.Errorf("%s: stream read error: %w", p.name, err) + return result, fmt.Errorf("%s: stream read error: %w", p.name, err) } // Assemble generated images from image accumulator into ChatResponse. diff --git a/internal/providers/model_fallback.go b/internal/providers/model_fallback.go index 9df97e03..2f9310c7 100644 --- a/internal/providers/model_fallback.go +++ b/internal/providers/model_fallback.go @@ -21,6 +21,13 @@ type ModelFallbackProvider struct { maxAttempts int } +type FallbackCallInfo struct { + Streamed bool +} + +type FallbackAfterCall func(*ChatResponse, error, FallbackCallInfo) +type FallbackBeforeCall func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (after FallbackAfterCall, err error) + func NewModelFallbackProvider(primary FallbackCandidate, fallbacks []FallbackCandidate, maxAttempts int, cooldownEnabled bool) *ModelFallbackProvider { var tracker *CooldownTracker if cooldownEnabled { @@ -64,6 +71,22 @@ func (p *ModelFallbackProvider) Chat(ctx context.Context, req ChatRequest) (*Cha }) } +func (p *ModelFallbackProvider) ChatWithHook(ctx context.Context, req ChatRequest, before FallbackBeforeCall) (*ChatResponse, error) { + return p.runOrdered(ctx, req, func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (*ChatResponse, error) { + nextReq := req + nextReq.Model = entry.Model + after, err := before(ctx, entry, nextReq) + if err != nil { + return nil, err + } + resp, err := entry.Provider.Chat(ctx, nextReq) + if after != nil { + after(resp, err, FallbackCallInfo{}) + } + return resp, err + }) +} + func (p *ModelFallbackProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk func(StreamChunk)) (*ChatResponse, error) { return p.runOrdered(ctx, req, func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (*ChatResponse, error) { nextReq := req @@ -82,6 +105,31 @@ func (p *ModelFallbackProvider) ChatStream(ctx context.Context, req ChatRequest, }) } +func (p *ModelFallbackProvider) ChatStreamWithHook(ctx context.Context, req ChatRequest, onChunk func(StreamChunk), before FallbackBeforeCall) (*ChatResponse, error) { + return p.runOrdered(ctx, req, func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (*ChatResponse, error) { + nextReq := req + nextReq.Model = entry.Model + after, err := before(ctx, entry, nextReq) + if err != nil { + return nil, err + } + streamed := false + resp, err := entry.Provider.ChatStream(ctx, nextReq, func(chunk StreamChunk) { + if chunk.Content != "" || chunk.Thinking != "" || len(chunk.Images) > 0 { + streamed = true + } + onChunk(chunk) + }) + if after != nil { + after(resp, err, FallbackCallInfo{Streamed: streamed}) + } + if streamed && err != nil { + return nil, noFallbackAfterStreamError{err: err} + } + return resp, err + }) +} + func (p *ModelFallbackProvider) runOrdered( ctx context.Context, req ChatRequest, diff --git a/internal/providers/model_fallback_test.go b/internal/providers/model_fallback_test.go index 5bea45bb..a891374d 100644 --- a/internal/providers/model_fallback_test.go +++ b/internal/providers/model_fallback_test.go @@ -94,6 +94,33 @@ func TestModelFallbackProviderDoesNotFallbackAfterStreamChunk(t *testing.T) { } } +func TestModelFallbackProviderChatStreamWithHookReportsStreamedChunks(t *testing.T) { + streamErr := &HTTPError{Status: 429, Body: "rate limited"} + primary := &testFallbackProvider{ + name: "primary", + model: "primary-model", + streamErr: streamErr, + } + provider := NewModelFallbackProvider(FallbackCandidate{ + ProviderName: "primary", + Provider: primary, + Model: "primary-model", + }, nil, 1, false) + + var streamed bool + _, err := provider.ChatStreamWithHook(context.Background(), ChatRequest{}, func(StreamChunk) {}, func(context.Context, FallbackCandidate, ChatRequest) (FallbackAfterCall, error) { + return func(_ *ChatResponse, _ error, info FallbackCallInfo) { + streamed = info.Streamed + }, nil + }) + if err == nil { + t.Fatal("ChatStreamWithHook() error = nil, want stream error") + } + if !streamed { + t.Fatal("FallbackCallInfo.Streamed = false, want true after partial stream") + } +} + func TestModelFallbackProviderFallsBackToSameModelOnDifferentProvider(t *testing.T) { primary := &testFallbackProvider{ name: "primary", diff --git a/internal/providers/openai_chat.go b/internal/providers/openai_chat.go index b3991990..c5b1f960 100644 --- a/internal/providers/openai_chat.go +++ b/internal/providers/openai_chat.go @@ -108,13 +108,19 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun PromptTokens: chunk.Usage.PromptTokens, CompletionTokens: chunk.Usage.CompletionTokens, TotalTokens: chunk.Usage.TotalTokens, + RequestCount: 1, } if chunk.Usage.PromptTokensDetails != nil { result.Usage.CacheReadTokens = chunk.Usage.PromptTokensDetails.CachedTokens + result.Usage.CacheCreationTokens = chunk.Usage.PromptTokensDetails.CacheWriteTokens + result.Usage.PromptTokensIncludeCachedSegments = true } if chunk.Usage.CompletionTokensDetails != nil && chunk.Usage.CompletionTokensDetails.ReasoningTokens > 0 { result.Usage.ThinkingTokens = chunk.Usage.CompletionTokensDetails.ReasoningTokens } + if chunk.Usage.ServerToolUse != nil { + result.Usage.WebSearchCount = chunk.Usage.ServerToolUse.WebSearchRequests + } } if len(chunk.Choices) == 0 { @@ -181,7 +187,7 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun // Check for scanner errors (timeout, connection reset, etc.) if err := sse.Err(); err != nil { - return nil, fmt.Errorf("%s: stream read error: %w", p.name, err) + return result, fmt.Errorf("%s: stream read error: %w", p.name, err) } // Parse accumulated tool call arguments diff --git a/internal/providers/openai_http.go b/internal/providers/openai_http.go index 896021e0..99b269a8 100644 --- a/internal/providers/openai_http.go +++ b/internal/providers/openai_http.go @@ -125,13 +125,19 @@ func (p *OpenAIProvider) parseResponse(resp *openAIResponse) *ChatResponse { PromptTokens: resp.Usage.PromptTokens, CompletionTokens: resp.Usage.CompletionTokens, TotalTokens: resp.Usage.TotalTokens, + RequestCount: 1, } if resp.Usage.PromptTokensDetails != nil { result.Usage.CacheReadTokens = resp.Usage.PromptTokensDetails.CachedTokens + result.Usage.CacheCreationTokens = resp.Usage.PromptTokensDetails.CacheWriteTokens + result.Usage.PromptTokensIncludeCachedSegments = true } if resp.Usage.CompletionTokensDetails != nil && resp.Usage.CompletionTokensDetails.ReasoningTokens > 0 { result.Usage.ThinkingTokens = resp.Usage.CompletionTokensDetails.ReasoningTokens } + if resp.Usage.ServerToolUse != nil { + result.Usage.WebSearchCount = resp.Usage.ServerToolUse.WebSearchRequests + } } return result diff --git a/internal/providers/openai_types.go b/internal/providers/openai_types.go index 57361e1c..48389cfd 100644 --- a/internal/providers/openai_types.go +++ b/internal/providers/openai_types.go @@ -49,18 +49,25 @@ type openAIUsage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` TotalTokens int `json:"total_tokens"` + Cost float64 `json:"cost,omitempty"` PromptTokensDetails *openAIPromptDetails `json:"prompt_tokens_details,omitempty"` CompletionTokensDetails *openAICompletionDetails `json:"completion_tokens_details,omitempty"` + ServerToolUse *openAIServerToolUse `json:"server_tool_use,omitempty"` } type openAIPromptDetails struct { - CachedTokens int `json:"cached_tokens"` + CachedTokens int `json:"cached_tokens"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` } type openAICompletionDetails struct { ReasoningTokens int `json:"reasoning_tokens"` } +type openAIServerToolUse struct { + WebSearchRequests int `json:"web_search_requests,omitempty"` +} + // Streaming types type openAIStreamChunk struct { diff --git a/internal/providers/types.go b/internal/providers/types.go index 555a2bf0..d7f03413 100644 --- a/internal/providers/types.go +++ b/internal/providers/types.go @@ -186,10 +186,14 @@ type ToolFunctionSchema struct { // Usage tracks token consumption. type Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - CacheCreationTokens int `json:"cache_creation_input_tokens,omitempty"` - CacheReadTokens int `json:"cache_read_input_tokens,omitempty"` - ThinkingTokens int `json:"thinking_tokens,omitempty"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + CacheCreationTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadTokens int `json:"cache_read_input_tokens,omitempty"` + PromptTokensIncludeCachedSegments bool `json:"prompt_tokens_include_cached_segments,omitempty"` + ThinkingTokens int `json:"thinking_tokens,omitempty"` + RequestCount int `json:"request_count,omitempty"` + ImageCount int `json:"image_count,omitempty"` + WebSearchCount int `json:"web_search_count,omitempty"` } diff --git a/internal/store/pg/factory.go b/internal/store/pg/factory.go index 22ef792c..b22c1a90 100644 --- a/internal/store/pg/factory.go +++ b/internal/store/pg/factory.go @@ -24,48 +24,49 @@ func NewPGStores(cfg store.StoreConfig) (*store.Stores, error) { } pgStores := &store.Stores{ - DB: db, - Sessions: NewPGSessionStore(db), - Memory: NewPGMemoryStore(db, memCfg), - Cron: NewPGCronStore(db), - Pairing: NewPGPairingStore(db), - Skills: NewPGSkillStore(db, skillsDir), - Agents: NewPGAgentStore(db), - Providers: NewPGProviderStore(db, cfg.EncryptionKey), - Tracing: NewPGTracingStore(db), - MCP: NewPGMCPServerStore(db, cfg.EncryptionKey), - ChannelInstances: NewPGChannelInstanceStore(db, cfg.EncryptionKey), - ConfigSecrets: NewPGConfigSecretsStore(db, cfg.EncryptionKey), - AgentLinks: NewPGAgentLinkStore(db), - Teams: NewPGTeamStore(db), - BuiltinTools: NewPGBuiltinToolStore(db), - PendingMessages: NewPGPendingMessageStore(db), - KnowledgeGraph: NewPGKnowledgeGraphStore(db), - Contacts: NewPGContactStore(db), - Activity: NewPGActivityStore(db), - Snapshots: NewPGSnapshotStore(db), - SecureCLI: NewPGSecureCLIStore(db, cfg.EncryptionKey), - SecureCLIGrants: NewPGSecureCLIAgentGrantStore(db, cfg.EncryptionKey), - APIKeys: NewPGAPIKeyStore(db), - Heartbeats: NewPGHeartbeatStore(db), - ConfigPermissions: NewPGConfigPermissionStore(db), - Tenants: NewPGTenantStore(db), - BuiltinToolTenantCfgs: NewPGBuiltinToolTenantConfigStore(db), - SkillTenantCfgs: NewPGSkillTenantConfigStore(db), - SystemConfigs: NewPGSystemConfigStore(db), - SubagentTasks: NewPGSubagentTaskStore(db), - Vault: NewPGVaultStore(db), - Episodic: NewPGEpisodicStore(db), - EvolutionMetrics: NewPGEvolutionMetricsStore(db), - EvolutionSuggestions: NewPGEvolutionSuggestionStore(db), - BitrixPortals: NewPGBitrixPortalStore(db, cfg.EncryptionKey), - Hooks: NewPGHookStore(db), + DB: db, + Sessions: NewPGSessionStore(db), + Memory: NewPGMemoryStore(db, memCfg), + Cron: NewPGCronStore(db), + Pairing: NewPGPairingStore(db), + Skills: NewPGSkillStore(db, skillsDir), + Agents: NewPGAgentStore(db), + Providers: NewPGProviderStore(db, cfg.EncryptionKey), + Tracing: NewPGTracingStore(db), + MCP: NewPGMCPServerStore(db, cfg.EncryptionKey), + ChannelInstances: NewPGChannelInstanceStore(db, cfg.EncryptionKey), + ConfigSecrets: NewPGConfigSecretsStore(db, cfg.EncryptionKey), + AgentLinks: NewPGAgentLinkStore(db), + Teams: NewPGTeamStore(db), + BuiltinTools: NewPGBuiltinToolStore(db), + PendingMessages: NewPGPendingMessageStore(db), + KnowledgeGraph: NewPGKnowledgeGraphStore(db), + Contacts: NewPGContactStore(db), + Activity: NewPGActivityStore(db), + Snapshots: NewPGSnapshotStore(db), + SecureCLI: NewPGSecureCLIStore(db, cfg.EncryptionKey), + SecureCLIGrants: NewPGSecureCLIAgentGrantStore(db, cfg.EncryptionKey), + APIKeys: NewPGAPIKeyStore(db), + Heartbeats: NewPGHeartbeatStore(db), + ConfigPermissions: NewPGConfigPermissionStore(db), + Tenants: NewPGTenantStore(db), + BuiltinToolTenantCfgs: NewPGBuiltinToolTenantConfigStore(db), + SkillTenantCfgs: NewPGSkillTenantConfigStore(db), + SystemConfigs: NewPGSystemConfigStore(db), + SubagentTasks: NewPGSubagentTaskStore(db), + Vault: NewPGVaultStore(db), + Episodic: NewPGEpisodicStore(db), + EvolutionMetrics: NewPGEvolutionMetricsStore(db), + EvolutionSuggestions: NewPGEvolutionSuggestionStore(db), + BitrixPortals: NewPGBitrixPortalStore(db, cfg.EncryptionKey), + Hooks: NewPGHookStore(db), Webhooks: NewPGWebhookStore(db), WebhookCalls: NewPGWebhookCallStore(db), Workstations: NewPGWorkstationStore(db, cfg.EncryptionKey), WorkstationLinks: NewPGAgentWorkstationLinkStore(db), WorkstationPermissions: NewPGWorkstationPermissionStore(db), WorkstationActivity: NewPGWorkstationActivityStore(db), + UsageCaps: NewPGUsageCapStore(db), } // Wire permStore into WorkstationStore so Create seeds allowlist atomically (H5 fix). // Must happen after both stores are constructed. diff --git a/internal/store/pg/usage_caps.go b/internal/store/pg/usage_caps.go new file mode 100644 index 00000000..833edb49 --- /dev/null +++ b/internal/store/pg/usage_caps.go @@ -0,0 +1,406 @@ +package pg + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func (s *PGUsageCapStore) CreateUsageCapPolicy(ctx context.Context, p *store.UsageCapPolicy) error { + if p.ID == uuid.Nil { + p.ID = uuid.New() + } + if err := s.validateUsageCapRefs(ctx, p.TenantID, p.AgentID, p.ProviderID); err != nil { + return err + } + const q = ` +INSERT INTO usage_cap_policies ( + id, tenant_id, agent_id, provider_id, provider_type, model_id, window, + max_tokens, max_cost_micros, enabled, priority +) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) +RETURNING created_at, updated_at` + return s.db.QueryRowContext(ctx, q, + p.ID, p.TenantID, uuidPtrVal(p.AgentID), uuidPtrVal(p.ProviderID), + nullEmpty(p.ProviderType), nullEmpty(p.ModelID), p.Window, + intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros), p.Enabled, p.Priority, + ).Scan(&p.CreatedAt, &p.UpdatedAt) +} + +func (s *PGUsageCapStore) ListUsageCapPolicies(ctx context.Context, scope store.UsageCapScope, includeDisabled bool) ([]store.UsageCapPolicy, error) { + args := []any{scope.TenantID} + conds := []string{"tenant_id = $1"} + if !includeDisabled { + conds = append(conds, "enabled = true") + } + if scope.AgentID != uuid.Nil { + args = append(args, scope.AgentID) + conds = append(conds, "(agent_id IS NULL OR agent_id = $"+fmt.Sprint(len(args))+")") + } else if !includeDisabled { + conds = append(conds, "agent_id IS NULL") + } + if scope.ProviderID != uuid.Nil { + args = append(args, scope.ProviderID) + conds = append(conds, "(provider_id IS NULL OR provider_id = $"+fmt.Sprint(len(args))+")") + } else if !includeDisabled { + conds = append(conds, "provider_id IS NULL") + } + if scope.ProviderType != "" { + args = append(args, scope.ProviderType) + conds = append(conds, "(provider_type IS NULL OR provider_type = $"+fmt.Sprint(len(args))+")") + } else if !includeDisabled { + conds = append(conds, "provider_type IS NULL") + } + if scope.ModelID != "" { + args = append(args, scope.ModelID) + conds = append(conds, "(model_id IS NULL OR model_id = $"+fmt.Sprint(len(args))+")") + } else if !includeDisabled { + conds = append(conds, "model_id IS NULL") + } + rows, err := s.db.QueryContext(ctx, policySelectSQL+" WHERE "+strings.Join(conds, " AND ")+" ORDER BY priority ASC, created_at ASC", args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []store.UsageCapPolicy + for rows.Next() { + p, err := scanPolicy(rows) + if err != nil { + return nil, err + } + out = append(out, p) + } + return out, rows.Err() +} + +func (s *PGUsageCapStore) UpdateUsageCapPolicy(ctx context.Context, tenantID, id uuid.UUID, patch store.UsageCapPolicyPatch) (*store.UsageCapPolicy, error) { + p, err := s.getPolicy(ctx, tenantID, id) + if err != nil { + return nil, err + } + if patch.AgentID != nil { + p.AgentID = *patch.AgentID + } + if patch.ProviderID != nil { + p.ProviderID = *patch.ProviderID + } + if patch.ProviderType != nil { + p.ProviderType = *patch.ProviderType + } + if patch.ModelID != nil { + p.ModelID = *patch.ModelID + } + if patch.Window != nil { + p.Window = *patch.Window + } + if patch.MaxTokens != nil { + p.MaxTokens = *patch.MaxTokens + } + if patch.MaxCostMicros != nil { + p.MaxCostMicros = *patch.MaxCostMicros + } + if patch.Enabled != nil { + p.Enabled = *patch.Enabled + } + if patch.Priority != nil { + p.Priority = *patch.Priority + } + if err := s.validateUsageCapRefs(ctx, tenantID, p.AgentID, p.ProviderID); err != nil { + return nil, err + } + const q = ` +UPDATE usage_cap_policies SET agent_id=$3, provider_id=$4, provider_type=$5, + model_id=$6, window=$7, max_tokens=$8, max_cost_micros=$9, + enabled=$10, priority=$11, updated_at=now() +WHERE tenant_id=$1 AND id=$2 +RETURNING updated_at` + if err := s.db.QueryRowContext(ctx, q, tenantID, id, uuidPtrVal(p.AgentID), uuidPtrVal(p.ProviderID), + nullEmpty(p.ProviderType), nullEmpty(p.ModelID), p.Window, intPtrVal(p.MaxTokens), + intPtrVal(p.MaxCostMicros), p.Enabled, p.Priority).Scan(&p.UpdatedAt); err != nil { + return nil, err + } + return p, nil +} + +func (s *PGUsageCapStore) DeleteUsageCapPolicy(ctx context.Context, tenantID, id uuid.UUID) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM usage_cap_policies WHERE tenant_id=$1 AND id=$2`, tenantID, id) + return err +} + +func (s *PGUsageCapStore) ReserveUsage(ctx context.Context, req store.UsageReserveRequest, policies []store.UsageCapPolicy) (*store.UsageReservationResult, error) { + if len(policies) == 0 { + return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Skipped: true, Reason: "no_policy"}, nil + } + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer tx.Rollback() + for _, p := range policies { + start, end := usageWindow(time.Now().UTC(), p.Window) + if _, err := tx.ExecContext(ctx, `INSERT INTO usage_cap_counters (policy_id, window_start, window_end) VALUES ($1,$2,$3) ON CONFLICT DO NOTHING`, p.ID, start, end); err != nil { + return nil, err + } + meta := req.Metadata + if len(meta) == 0 { + meta = json.RawMessage(`{}`) + } + var inserted bool + err := tx.QueryRowContext(ctx, ` +INSERT INTO usage_cap_reservations (reservation_key, policy_id, window_start, reserved_tokens, reserved_cost_micros, metadata) +VALUES ($1,$2,$3,$4,$5,$6) +ON CONFLICT (reservation_key, policy_id) DO NOTHING +RETURNING true`, + req.ReservationKey, p.ID, start, req.EstimatedTokens, req.EstimatedCostMicros, meta).Scan(&inserted) + if errors.Is(err, sql.ErrNoRows) { + continue + } + if err != nil { + return nil, err + } + var ok bool + err = tx.QueryRowContext(ctx, ` +UPDATE usage_cap_counters SET + reserved_tokens = reserved_tokens + $4, + reserved_cost_micros = reserved_cost_micros + $5, + updated_at = now() +WHERE policy_id=$1 AND window_start=$2 + AND ($6::bigint IS NULL OR used_tokens + reserved_tokens + $4 <= $6) + AND ($7::bigint IS NULL OR used_cost_micros + reserved_cost_micros + $5 <= $7) + RETURNING true`, p.ID, start, end, req.EstimatedTokens, req.EstimatedCostMicros, intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros)).Scan(&ok) + if err != nil { + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + return nil, &store.UsageCapExceededError{PolicyID: p.ID, Reason: "cap_exceeded"} + } + } + if err := tx.Commit(); err != nil { + return nil, err + } + return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Policies: policies}, nil +} + +func (s *PGUsageCapStore) ReconcileUsage(ctx context.Context, req store.UsageReconcileRequest) error { + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() + meta := req.Metadata + if len(meta) == 0 { + meta = json.RawMessage(`{}`) + } + rows, err := tx.QueryContext(ctx, ` +UPDATE usage_cap_reservations +SET status=$2, actual_tokens=$3, actual_cost_micros=$4, metadata=$5, updated_at=now() +WHERE reservation_key=$1 AND status='reserved' +RETURNING policy_id, window_start, reserved_tokens, reserved_cost_micros`, + req.ReservationKey, nullStatus(req.Status), req.ActualTokens, req.ActualCostMicros, meta) + if err != nil { + return err + } + type resv struct { + policyID uuid.UUID + start time.Time + tokens int64 + cost int64 + } + var reservations []resv + for rows.Next() { + var r resv + if err := rows.Scan(&r.policyID, &r.start, &r.tokens, &r.cost); err != nil { + rows.Close() + return err + } + reservations = append(reservations, r) + } + rows.Close() + if err := rows.Err(); err != nil { + return err + } + for _, r := range reservations { + if _, err := tx.ExecContext(ctx, ` +UPDATE usage_cap_counters SET + reserved_tokens = GREATEST(reserved_tokens - $3, 0), + reserved_cost_micros = GREATEST(reserved_cost_micros - $4, 0), + used_tokens = used_tokens + $5, + used_cost_micros = used_cost_micros + $6, + updated_at = now() +WHERE policy_id=$1 AND window_start=$2`, r.policyID, r.start, r.tokens, r.cost, req.ActualTokens, req.ActualCostMicros); err != nil { + return err + } + } + return tx.Commit() +} + +func (s *PGUsageCapStore) ListUsageCapUtilization(ctx context.Context, tenantID uuid.UUID) ([]store.UsageCapUtilization, error) { + policies, err := s.ListUsageCapPolicies(ctx, store.UsageCapScope{TenantID: tenantID}, true) + if err != nil { + return nil, err + } + out := make([]store.UsageCapUtilization, 0, len(policies)) + for _, p := range policies { + start, end := usageWindow(time.Now().UTC(), p.Window) + u := store.UsageCapUtilization{Policy: p, WindowStart: start, WindowEnd: end} + _ = s.db.QueryRowContext(ctx, ` +SELECT used_tokens, reserved_tokens, used_cost_micros, reserved_cost_micros +FROM usage_cap_counters WHERE policy_id=$1 AND window_start=$2`, p.ID, start).Scan(&u.UsedTokens, &u.ReservedTokens, &u.UsedCostMicros, &u.ReservedCostMicros) + out = append(out, u) + } + return out, nil +} + +func (s *PGUsageCapStore) ListUsageCapEvents(ctx context.Context, tenantID uuid.UUID, limit int) ([]store.UsageCapEvent, error) { + if limit <= 0 || limit > 200 { + limit = 50 + } + rows, err := s.db.QueryContext(ctx, ` +SELECT id, tenant_id, policy_id, COALESCE(reservation_key,''), decision, COALESCE(reason,''), + estimated_tokens, estimated_cost_micros, actual_tokens, actual_cost_micros, metadata, created_at +FROM usage_cap_events WHERE tenant_id=$1 ORDER BY created_at DESC LIMIT $2`, tenantID, limit) + if err != nil { + return nil, err + } + defer rows.Close() + var out []store.UsageCapEvent + for rows.Next() { + var e store.UsageCapEvent + var pid uuid.NullUUID + if err := rows.Scan(&e.ID, &e.TenantID, &pid, &e.ReservationKey, &e.Decision, &e.Reason, + &e.EstimatedTokens, &e.EstimatedCostMicros, &e.ActualTokens, &e.ActualCostMicros, &e.Metadata, &e.CreatedAt); err != nil { + return nil, err + } + if pid.Valid { + e.PolicyID = &pid.UUID + } + out = append(out, e) + } + return out, rows.Err() +} + +func (s *PGUsageCapStore) InsertUsageCapEvent(ctx context.Context, event *store.UsageCapEvent) error { + if event.ID == uuid.Nil { + event.ID = uuid.New() + } + if len(event.Metadata) == 0 { + event.Metadata = json.RawMessage(`{}`) + } + return s.db.QueryRowContext(ctx, ` +INSERT INTO usage_cap_events (id, tenant_id, policy_id, reservation_key, decision, reason, + estimated_tokens, estimated_cost_micros, actual_tokens, actual_cost_micros, metadata) +VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) +RETURNING created_at`, event.ID, event.TenantID, uuidPtrVal(event.PolicyID), nullEmpty(event.ReservationKey), + event.Decision, nullEmpty(event.Reason), event.EstimatedTokens, event.EstimatedCostMicros, + event.ActualTokens, event.ActualCostMicros, event.Metadata).Scan(&event.CreatedAt) +} + +func (s *PGUsageCapStore) getPolicy(ctx context.Context, tenantID, id uuid.UUID) (*store.UsageCapPolicy, error) { + row := s.db.QueryRowContext(ctx, policySelectSQL+" WHERE tenant_id=$1 AND id=$2", tenantID, id) + p, err := scanPolicy(row) + if err != nil { + return nil, err + } + return &p, nil +} + +func (s *PGUsageCapStore) validateUsageCapRefs(ctx context.Context, tenantID uuid.UUID, agentID, providerID *uuid.UUID) error { + if agentID != nil && *agentID != uuid.Nil { + var ok bool + if err := s.db.QueryRowContext(ctx, ` +SELECT EXISTS ( + SELECT 1 FROM agents WHERE id=$1 AND tenant_id=$2 AND deleted_at IS NULL +)`, *agentID, tenantID).Scan(&ok); err != nil { + return err + } + if !ok { + return errors.New("agent_id does not belong to tenant") + } + } + if providerID != nil && *providerID != uuid.Nil { + var ok bool + if err := s.db.QueryRowContext(ctx, ` +SELECT EXISTS ( + SELECT 1 FROM llm_providers WHERE id=$1 AND tenant_id IN ($2, $3) +)`, *providerID, tenantID, store.MasterTenantID).Scan(&ok); err != nil { + return err + } + if !ok { + return errors.New("provider_id does not belong to tenant") + } + } + return nil +} + +const policySelectSQL = `SELECT id, tenant_id, agent_id, provider_id, COALESCE(provider_type,''), COALESCE(model_id,''), + window, max_tokens, max_cost_micros, enabled, priority, created_at, updated_at FROM usage_cap_policies` + +func scanPolicy(row scanner) (store.UsageCapPolicy, error) { + var p store.UsageCapPolicy + var agentID, providerID uuid.NullUUID + var maxTokens, maxCost sql.NullInt64 + err := row.Scan(&p.ID, &p.TenantID, &agentID, &providerID, &p.ProviderType, &p.ModelID, + &p.Window, &maxTokens, &maxCost, &p.Enabled, &p.Priority, &p.CreatedAt, &p.UpdatedAt) + if agentID.Valid { + p.AgentID = &agentID.UUID + } + if providerID.Valid { + p.ProviderID = &providerID.UUID + } + if maxTokens.Valid { + p.MaxTokens = &maxTokens.Int64 + } + if maxCost.Valid { + p.MaxCostMicros = &maxCost.Int64 + } + return p, err +} + +func usageWindow(now time.Time, window string) (time.Time, time.Time) { + now = now.UTC() + switch window { + case store.UsageCapWindowDay: + start := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC) + return start, start.AddDate(0, 0, 1) + case store.UsageCapWindowWeek: + weekday := int(now.Weekday()) + if weekday == 0 { + weekday = 7 + } + start := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, -(weekday - 1)) + return start, start.AddDate(0, 0, 7) + case store.UsageCapWindowMonth: + start := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + return start, start.AddDate(0, 1, 0) + default: + start := now.Truncate(time.Hour) + return start, start.Add(time.Hour) + } +} + +func uuidPtrVal(v *uuid.UUID) any { + if v == nil || *v == uuid.Nil { + return nil + } + return *v +} + +func intPtrVal(v *int64) any { + if v == nil { + return nil + } + return *v +} + +func nullStatus(s string) string { + if strings.TrimSpace(s) == "" { + return "reconciled" + } + return s +} diff --git a/internal/store/pg/usage_caps_test.go b/internal/store/pg/usage_caps_test.go new file mode 100644 index 00000000..a6702d31 --- /dev/null +++ b/internal/store/pg/usage_caps_test.go @@ -0,0 +1,180 @@ +package pg + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func TestPGUsageCapStoreReserveUsageIdempotent(t *testing.T) { + db := hooksTestDB(t) + tenantID, _ := seedTenantAndAgent(t, db) + usageStore := NewPGUsageCapStore(db) + maxTokens := int64(100) + policy := &store.UsageCapPolicy{ + TenantID: tenantID, Window: store.UsageCapWindowHour, + MaxTokens: &maxTokens, Enabled: true, Priority: 100, + } + if err := usageStore.CreateUsageCapPolicy(context.Background(), policy); err != nil { + t.Fatalf("CreateUsageCapPolicy: %v", err) + } + req := store.UsageReserveRequest{ + UsageCapScope: store.UsageCapScope{TenantID: tenantID}, + ReservationKey: "duplicate-reservation", + EstimatedTokens: 10, + } + for i := range 2 { + if _, err := usageStore.ReserveUsage(context.Background(), req, []store.UsageCapPolicy{*policy}); err != nil { + t.Fatalf("ReserveUsage call %d: %v", i+1, err) + } + } + + var reservedTokens, reservationRows int64 + if err := db.QueryRow(`SELECT COALESCE(SUM(reserved_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&reservedTokens); err != nil { + t.Fatalf("query reserved tokens: %v", err) + } + if err := db.QueryRow(`SELECT COUNT(*) FROM usage_cap_reservations WHERE policy_id=$1 AND reservation_key=$2`, policy.ID, req.ReservationKey).Scan(&reservationRows); err != nil { + t.Fatalf("query reservation rows: %v", err) + } + if reservedTokens != 10 { + t.Fatalf("reserved_tokens = %d, want 10", reservedTokens) + } + if reservationRows != 1 { + t.Fatalf("reservation rows = %d, want 1", reservationRows) + } + + var wg sync.WaitGroup + errs := make(chan error, 2) + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + errs <- usageStore.ReconcileUsage(context.Background(), store.UsageReconcileRequest{ + ReservationKey: req.ReservationKey, + ActualTokens: 7, + Status: "reconciled", + }) + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("ReconcileUsage: %v", err) + } + } + var usedTokens int64 + if err := db.QueryRow(`SELECT COALESCE(SUM(used_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&usedTokens); err != nil { + t.Fatalf("query used tokens: %v", err) + } + if err := db.QueryRow(`SELECT COALESCE(SUM(reserved_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&reservedTokens); err != nil { + t.Fatalf("query reserved tokens after reconcile: %v", err) + } + if usedTokens != 7 { + t.Fatalf("used_tokens = %d, want 7", usedTokens) + } + if reservedTokens != 0 { + t.Fatalf("reserved_tokens after reconcile = %d, want 0", reservedTokens) + } +} + +func TestPGUsageCapStoreRejectsCrossTenantRefs(t *testing.T) { + db := hooksTestDB(t) + tenantA, agentA := seedTenantAndAgent(t, db) + tenantB, _ := seedTenantAndAgent(t, db) + usageStore := NewPGUsageCapStore(db) + + policy := &store.UsageCapPolicy{ + TenantID: tenantB, AgentID: &agentA, Window: store.UsageCapWindowDay, + MaxTokens: int64PtrPG(100), Enabled: true, + } + if err := usageStore.CreateUsageCapPolicy(context.Background(), policy); err == nil { + t.Fatal("CreateUsageCapPolicy accepted agent_id from another tenant") + } + + providerID := uuid.New() + if _, err := db.Exec( + `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) + VALUES ($1,$2,$3,'openrouter','sk-test',true)`, + providerID, tenantA, "usage-cap-provider-"+providerID.String(), + ); err != nil { + t.Fatalf("seed provider: %v", err) + } + override := &store.UsagePricingOverride{ + TenantID: tenantB, ProviderID: providerID, ProviderType: store.ProviderOpenRouter, + ModelID: "openai/gpt-test", Enabled: true, + } + if err := usageStore.PutPricingOverride(context.Background(), override); err == nil { + t.Fatal("PutPricingOverride accepted provider_id from another tenant") + } + + masterProviderID := uuid.New() + if _, err := db.Exec( + `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) + VALUES ($1,$2,$3,'openrouter','sk-test',true)`, + masterProviderID, store.MasterTenantID, "usage-cap-master-provider-"+masterProviderID.String(), + ); err != nil { + t.Fatalf("seed master provider: %v", err) + } + t.Cleanup(func() { + db.Exec("DELETE FROM llm_providers WHERE id=$1", masterProviderID) + }) + masterScopedPolicy := &store.UsageCapPolicy{ + TenantID: tenantB, ProviderID: &masterProviderID, Window: store.UsageCapWindowDay, + MaxTokens: int64PtrPG(100), Enabled: true, + } + if err := usageStore.CreateUsageCapPolicy(context.Background(), masterScopedPolicy); err != nil { + t.Fatalf("CreateUsageCapPolicy rejected master provider ref: %v", err) + } +} + +func TestValidateUsagePricingFieldsRejectsNegativeValues(t *testing.T) { + negative := "-0.01" + if err := validateUsagePricingFields(store.UsagePricingFields{Input: &negative}); err == nil { + t.Fatal("validateUsagePricingFields accepted negative price") + } +} + +func TestPGUsageCapStoreResolvePricingUsesOpenRouterAliases(t *testing.T) { + db := hooksTestDB(t) + usageStore := NewPGUsageCapStore(db) + inputPrice := "0.000001" + 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: "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()}, + } + if _, err := usageStore.UpsertPricingCatalog(context.Background(), entries); err != nil { + t.Fatalf("UpsertPricingCatalog: %v", err) + } + + cases := []struct { + name string + providerName string + providerType string + modelID string + wantModelID string + }{ + {name: "openai compat", providerName: "openai", providerType: store.ProviderOpenAICompat, modelID: "gpt-4o-mini", wantModelID: "openai/gpt-4o-mini"}, + {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"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resolved, err := usageStore.ResolvePricing(context.Background(), uuid.New(), uuid.New(), tc.providerName, tc.providerType, tc.modelID) + if err != nil { + t.Fatalf("ResolvePricing: %v", err) + } + if resolved.ModelID != tc.wantModelID { + t.Fatalf("resolved model = %q, want %q", resolved.ModelID, tc.wantModelID) + } + }) + } +} + +func int64PtrPG(v int64) *int64 { return &v } diff --git a/internal/store/pg/usage_pricing.go b/internal/store/pg/usage_pricing.go new file mode 100644 index 00000000..64ef138e --- /dev/null +++ b/internal/store/pg/usage_pricing.go @@ -0,0 +1,362 @@ +package pg + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "math/big" + "strconv" + "strings" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +type PGUsageCapStore struct { + db *sql.DB +} + +func NewPGUsageCapStore(db *sql.DB) *PGUsageCapStore { + return &PGUsageCapStore{db: db} +} + +func (s *PGUsageCapStore) UpsertPricingCatalog(ctx context.Context, entries []store.UsagePricingCatalogEntry) (int, error) { + const q = ` +INSERT INTO usage_pricing_catalog ( + model_id, canonical_model_id, raw_pricing, raw_model, + input_price, output_price, cache_read_price, cache_write_price, + reasoning_price, request_price, image_price, web_search_price, synced_at +) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13) +ON CONFLICT (model_id) DO UPDATE SET + canonical_model_id = EXCLUDED.canonical_model_id, + raw_pricing = EXCLUDED.raw_pricing, + raw_model = EXCLUDED.raw_model, + input_price = EXCLUDED.input_price, + output_price = EXCLUDED.output_price, + cache_read_price = EXCLUDED.cache_read_price, + cache_write_price = EXCLUDED.cache_write_price, + reasoning_price = EXCLUDED.reasoning_price, + request_price = EXCLUDED.request_price, + image_price = EXCLUDED.image_price, + web_search_price = EXCLUDED.web_search_price, + synced_at = EXCLUDED.synced_at, + updated_at = now()` + for _, e := range entries { + if strings.TrimSpace(e.ModelID) == "" { + continue + } + if err := validateUsagePricingFields(e.Pricing); err != nil { + return 0, err + } + if len(e.RawPricing) == 0 { + e.RawPricing = json.RawMessage(`{}`) + } + if len(e.RawModel) == 0 { + e.RawModel = json.RawMessage(`{}`) + } + if _, err := s.db.ExecContext(ctx, q, + e.ModelID, nullEmpty(e.CanonicalModelID), e.RawPricing, e.RawModel, + priceVal(e.Pricing.Input), priceVal(e.Pricing.Output), + priceVal(e.Pricing.CacheRead), priceVal(e.Pricing.CacheWrite), + priceVal(e.Pricing.Reasoning), priceVal(e.Pricing.Request), + priceVal(e.Pricing.Image), priceVal(e.Pricing.WebSearch), e.SyncedAt, + ); err != nil { + return 0, err + } + } + return len(entries), nil +} + +func (s *PGUsageCapStore) ListPricingCatalog(ctx context.Context, q store.UsagePricingQuery) ([]store.UsagePricingCatalogEntry, error) { + limit := q.Limit + if limit <= 0 || limit > 500 { + limit = 100 + } + args := []any{} + where := "TRUE" + if q.ModelID != "" { + args = append(args, "%"+q.ModelID+"%") + where = "model_id ILIKE $1" + } + args = append(args, limit) + rows, err := s.db.QueryContext(ctx, ` +SELECT id, model_id, COALESCE(canonical_model_id,''), raw_pricing, raw_model, + input_price::text, output_price::text, cache_read_price::text, cache_write_price::text, + reasoning_price::text, request_price::text, image_price::text, web_search_price::text, + synced_at, created_at, updated_at +FROM usage_pricing_catalog +WHERE `+where+` +ORDER BY model_id +LIMIT $`+strconv.Itoa(len(args))+``, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []store.UsagePricingCatalogEntry + for rows.Next() { + e, err := scanCatalog(rows) + if err != nil { + return nil, err + } + out = append(out, e) + } + return out, rows.Err() +} + +func (s *PGUsageCapStore) PutPricingOverride(ctx context.Context, o *store.UsagePricingOverride) error { + if o.ID == uuid.Nil { + o.ID = uuid.New() + } + if o.ProviderID == uuid.Nil || o.TenantID == uuid.Nil || o.ModelID == "" { + return errors.New("tenant_id, provider_id, and model_id are required") + } + if err := validateUsagePricingFields(o.Pricing); err != nil { + return err + } + if err := s.validateUsageCapRefs(ctx, o.TenantID, nil, &o.ProviderID); err != nil { + return err + } + const q = ` +INSERT INTO usage_pricing_overrides ( + id, tenant_id, provider_id, provider_type, model_id, + input_price, output_price, cache_read_price, cache_write_price, + reasoning_price, request_price, image_price, web_search_price, enabled +) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) +ON CONFLICT (tenant_id, provider_id, model_id) DO UPDATE SET + provider_type = EXCLUDED.provider_type, + input_price = EXCLUDED.input_price, + output_price = EXCLUDED.output_price, + cache_read_price = EXCLUDED.cache_read_price, + cache_write_price = EXCLUDED.cache_write_price, + reasoning_price = EXCLUDED.reasoning_price, + request_price = EXCLUDED.request_price, + image_price = EXCLUDED.image_price, + web_search_price = EXCLUDED.web_search_price, + enabled = EXCLUDED.enabled, + updated_at = now() +RETURNING id, created_at, updated_at` + return s.db.QueryRowContext(ctx, q, + o.ID, o.TenantID, o.ProviderID, o.ProviderType, o.ModelID, + priceVal(o.Pricing.Input), priceVal(o.Pricing.Output), + priceVal(o.Pricing.CacheRead), priceVal(o.Pricing.CacheWrite), + priceVal(o.Pricing.Reasoning), priceVal(o.Pricing.Request), + priceVal(o.Pricing.Image), priceVal(o.Pricing.WebSearch), o.Enabled, + ).Scan(&o.ID, &o.CreatedAt, &o.UpdatedAt) +} + +func (s *PGUsageCapStore) ListPricingOverrides(ctx context.Context, q store.UsagePricingQuery) ([]store.UsagePricingOverride, error) { + args := []any{q.TenantID} + where := "tenant_id = $1" + if q.ProviderID != uuid.Nil { + args = append(args, q.ProviderID) + where += " AND provider_id = $" + itoa(len(args)) + } + rows, err := s.db.QueryContext(ctx, overrideSelectSQL+" WHERE "+where+" ORDER BY updated_at DESC", args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []store.UsagePricingOverride + for rows.Next() { + o, err := scanOverride(rows) + if err != nil { + return nil, err + } + out = append(out, o) + } + return out, rows.Err() +} + +func (s *PGUsageCapStore) DeletePricingOverride(ctx context.Context, tenantID, id uuid.UUID) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM usage_pricing_overrides WHERE tenant_id=$1 AND id=$2`, tenantID, id) + return err +} + +func (s *PGUsageCapStore) ResolvePricing(ctx context.Context, tenantID, providerID uuid.UUID, providerName, providerType, modelID string) (*store.ResolvedUsagePricing, error) { + candidates := usagePricingModelCandidates(providerName, providerType, modelID) + if tenantID != uuid.Nil && providerID != uuid.Nil { + for _, candidate := range candidates { + row := s.db.QueryRowContext(ctx, overrideSelectSQL+` WHERE tenant_id=$1 AND provider_id=$2 AND model_id=$3 AND enabled=true`, tenantID, providerID, candidate) + if o, err := scanOverride(row); err == nil { + return &store.ResolvedUsagePricing{ModelID: o.ModelID, ProviderID: providerID, ProviderType: providerType, Source: "override", Pricing: o.Pricing, OverrideID: o.ID}, nil + } else if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + } + } + for _, candidate := range candidates { + row := s.db.QueryRowContext(ctx, ` +SELECT id, model_id, COALESCE(canonical_model_id,''), raw_pricing, raw_model, + input_price::text, output_price::text, cache_read_price::text, cache_write_price::text, + reasoning_price::text, request_price::text, image_price::text, web_search_price::text, + synced_at, created_at, updated_at +FROM usage_pricing_catalog WHERE model_id=$1 OR canonical_model_id=$1 LIMIT 1`, candidate) + e, err := scanCatalog(row) + if err == nil { + return &store.ResolvedUsagePricing{ModelID: e.ModelID, ProviderID: providerID, ProviderType: providerType, Source: "catalog", Pricing: e.Pricing, CatalogSynced: &e.SyncedAt}, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + } + return nil, sql.ErrNoRows +} + +func usagePricingModelCandidates(providerName, providerType, modelID string) []string { + modelID = strings.TrimSpace(modelID) + if modelID == "" { + return nil + } + out := []string{modelID} + if strings.Contains(modelID, "/") { + return out + } + for _, prefix := range openRouterProviderPrefixes(providerName, providerType) { + out = appendUniqueString(out, prefix+"/"+modelID) + } + return out +} + +func openRouterProviderPrefixes(providerName, providerType string) []string { + switch providerType { + case store.ProviderAnthropicNative: + return []string{"anthropic"} + case store.ProviderGeminiNative, store.ProviderVertex: + return []string{"google"} + case store.ProviderOpenAICompat: + switch normalizeProviderAlias(providerName) { + case "openai", "azure", "azure-openai", "azure_openai": + return []string{"openai"} + case "anthropic": + return []string{"anthropic"} + case "gemini", "google", "vertex": + return []string{"google"} + } + return nil + case store.ProviderOpenRouter: + return nil + case store.ProviderGroq: + return []string{"groq"} + case store.ProviderDeepSeek: + return []string{"deepseek"} + case store.ProviderMistral: + return []string{"mistralai"} + case store.ProviderXAI: + return []string{"x-ai"} + case store.ProviderMiniMax: + return []string{"minimax"} + case store.ProviderCohere: + return []string{"cohere"} + case store.ProviderPerplexity: + return []string{"perplexity"} + case store.ProviderDashScope: + return []string{"qwen"} + default: + return nil + } +} + +func normalizeProviderAlias(s string) string { + s = strings.ToLower(strings.TrimSpace(s)) + s = strings.ReplaceAll(s, " ", "-") + return s +} + +func appendUniqueString(values []string, next string) []string { + next = strings.TrimSpace(next) + if next == "" { + return values + } + for _, existing := range values { + if existing == next { + return values + } + } + return append(values, next) +} + +const overrideSelectSQL = `SELECT id, tenant_id, provider_id, provider_type, model_id, + input_price::text, output_price::text, cache_read_price::text, cache_write_price::text, + reasoning_price::text, request_price::text, image_price::text, web_search_price::text, + enabled, created_at, updated_at FROM usage_pricing_overrides` + +type scanner interface{ Scan(dest ...any) error } + +func scanOverride(row scanner) (store.UsagePricingOverride, error) { + var o store.UsagePricingOverride + var prices [8]sql.NullString + err := row.Scan(&o.ID, &o.TenantID, &o.ProviderID, &o.ProviderType, &o.ModelID, + &prices[0], &prices[1], &prices[2], &prices[3], &prices[4], &prices[5], &prices[6], &prices[7], + &o.Enabled, &o.CreatedAt, &o.UpdatedAt) + o.Pricing = pricingFromNulls(prices) + return o, err +} + +func scanCatalog(row scanner) (store.UsagePricingCatalogEntry, error) { + var e store.UsagePricingCatalogEntry + var prices [8]sql.NullString + err := row.Scan(&e.ID, &e.ModelID, &e.CanonicalModelID, &e.RawPricing, &e.RawModel, + &prices[0], &prices[1], &prices[2], &prices[3], &prices[4], &prices[5], &prices[6], &prices[7], + &e.SyncedAt, &e.CreatedAt, &e.UpdatedAt) + e.Pricing = pricingFromNulls(prices) + return e, err +} + +func pricingFromNulls(p [8]sql.NullString) store.UsagePricingFields { + return store.UsagePricingFields{ + Input: pricePtr(p[0]), Output: pricePtr(p[1]), + CacheRead: pricePtr(p[2]), CacheWrite: pricePtr(p[3]), + Reasoning: pricePtr(p[4]), Request: pricePtr(p[5]), + Image: pricePtr(p[6]), WebSearch: pricePtr(p[7]), + } +} + +func pricePtr(v sql.NullString) *string { + if !v.Valid { + return nil + } + s := v.String + return &s +} + +func priceVal(v *string) any { + if v == nil || strings.TrimSpace(*v) == "" { + return nil + } + return strings.TrimSpace(*v) +} + +func validateUsagePricingFields(fields store.UsagePricingFields) error { + values := map[string]*string{ + "input": fields.Input, + "output": fields.Output, + "cache_read": fields.CacheRead, + "cache_write": fields.CacheWrite, + "reasoning": fields.Reasoning, + "request": fields.Request, + "image": fields.Image, + "web_search": fields.WebSearch, + } + for name, raw := range values { + if raw == nil || strings.TrimSpace(*raw) == "" { + continue + } + rat, ok := new(big.Rat).SetString(strings.TrimSpace(*raw)) + if !ok { + return fmt.Errorf("invalid %s price", name) + } + if rat.Sign() < 0 { + return fmt.Errorf("%s price must be non-negative", name) + } + } + return nil +} + +func nullEmpty(s string) any { + if strings.TrimSpace(s) == "" { + return nil + } + return s +} diff --git a/internal/store/stores.go b/internal/store/stores.go index 3f2614a2..3c008a55 100644 --- a/internal/store/stores.go +++ b/internal/store/stores.go @@ -4,41 +4,41 @@ import "database/sql" // Stores is the top-level container for all storage backends. type Stores struct { - DB *sql.DB // underlying connection - Sessions SessionStore - Memory MemoryStore - Cron CronStore - Pairing PairingStore - Skills SkillStore - Agents AgentStore - Providers ProviderStore - Tracing TracingStore - MCP MCPServerStore - ChannelInstances ChannelInstanceStore - ConfigSecrets ConfigSecretsStore - AgentLinks AgentLinkStore - Teams TeamStore - BuiltinTools BuiltinToolStore - PendingMessages PendingMessageStore - KnowledgeGraph KnowledgeGraphStore - Contacts ContactStore - Activity ActivityStore - Snapshots SnapshotStore - SecureCLI SecureCLIStore - SecureCLIGrants SecureCLIAgentGrantStore - APIKeys APIKeyStore - Heartbeats HeartbeatStore - ConfigPermissions ConfigPermissionStore - Tenants TenantStore - BuiltinToolTenantCfgs BuiltinToolTenantConfigStore - SkillTenantCfgs SkillTenantConfigStore - SystemConfigs SystemConfigStore - SubagentTasks SubagentTaskStore - Vault VaultStore - Episodic EpisodicStore - EvolutionMetrics EvolutionMetricsStore - EvolutionSuggestions EvolutionSuggestionStore - BitrixPortals BitrixPortalStore + DB *sql.DB // underlying connection + Sessions SessionStore + Memory MemoryStore + Cron CronStore + Pairing PairingStore + Skills SkillStore + Agents AgentStore + Providers ProviderStore + Tracing TracingStore + MCP MCPServerStore + ChannelInstances ChannelInstanceStore + ConfigSecrets ConfigSecretsStore + AgentLinks AgentLinkStore + Teams TeamStore + BuiltinTools BuiltinToolStore + PendingMessages PendingMessageStore + KnowledgeGraph KnowledgeGraphStore + Contacts ContactStore + Activity ActivityStore + Snapshots SnapshotStore + SecureCLI SecureCLIStore + SecureCLIGrants SecureCLIAgentGrantStore + APIKeys APIKeyStore + Heartbeats HeartbeatStore + ConfigPermissions ConfigPermissionStore + Tenants TenantStore + BuiltinToolTenantCfgs BuiltinToolTenantConfigStore + SkillTenantCfgs SkillTenantConfigStore + SystemConfigs SystemConfigStore + SubagentTasks SubagentTaskStore + Vault VaultStore + Episodic EpisodicStore + EvolutionMetrics EvolutionMetricsStore + EvolutionSuggestions EvolutionSuggestionStore + BitrixPortals BitrixPortalStore // Hooks is hooks.HookStore — typed as any to avoid import cycle // (hooks package imports store for context helpers). // Callers: type-assert to hooks.HookStore before use. @@ -52,4 +52,7 @@ type Stores struct { WorkstationLinks AgentWorkstationLinkStore WorkstationPermissions WorkstationPermissionStore WorkstationActivity WorkstationActivityStore + + // UsageCaps is Standard/PostgreSQL only in the first budget-control rollout. + UsageCaps UsageCapStore } diff --git a/internal/store/usage_caps.go b/internal/store/usage_caps.go new file mode 100644 index 00000000..1c123690 --- /dev/null +++ b/internal/store/usage_caps.go @@ -0,0 +1,195 @@ +package store + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/google/uuid" +) + +const ( + UsageCapWindowHour = "hour" + UsageCapWindowDay = "day" + UsageCapWindowWeek = "week" + UsageCapWindowMonth = "month" + + UsageCapEventAllow = "allow" + UsageCapEventBlock = "block" + UsageCapEventReconcile = "reconcile" + UsageCapEventSkip = "skip" +) + +var ErrUsageCapExceeded = errors.New("usage cap exceeded") + +type UsageCapExceededError struct { + PolicyID uuid.UUID + Reason string +} + +func (e *UsageCapExceededError) Error() string { + return ErrUsageCapExceeded.Error() +} + +func (e *UsageCapExceededError) Unwrap() error { + return ErrUsageCapExceeded +} + +// UsagePricingFields stores OpenRouter-compatible USD prices as decimal strings. +// Nil means unknown; "0" means explicitly free. +type UsagePricingFields struct { + Input *string `json:"input,omitempty"` + Output *string `json:"output,omitempty"` + CacheRead *string `json:"cache_read,omitempty"` + CacheWrite *string `json:"cache_write,omitempty"` + Reasoning *string `json:"reasoning,omitempty"` + Request *string `json:"request,omitempty"` + Image *string `json:"image,omitempty"` + WebSearch *string `json:"web_search,omitempty"` +} + +type UsagePricingCatalogEntry struct { + ID uuid.UUID `json:"id" db:"id"` + ModelID string `json:"model_id" db:"model_id"` + CanonicalModelID string `json:"canonical_model_id,omitempty" db:"canonical_model_id"` + Pricing UsagePricingFields `json:"pricing"` + RawPricing json.RawMessage `json:"raw_pricing,omitempty" db:"raw_pricing"` + RawModel json.RawMessage `json:"raw_model,omitempty" db:"raw_model"` + SyncedAt time.Time `json:"synced_at" db:"synced_at"` + CreatedAt time.Time `json:"created_at" db:"created_at"` + UpdatedAt time.Time `json:"updated_at" db:"updated_at"` +} + +type UsagePricingOverride struct { + ID uuid.UUID `json:"id" db:"id"` + TenantID uuid.UUID `json:"tenant_id" db:"tenant_id"` + ProviderID uuid.UUID `json:"provider_id" db:"provider_id"` + ProviderType string `json:"provider_type" db:"provider_type"` + ModelID string `json:"model_id" db:"model_id"` + Pricing UsagePricingFields `json:"pricing"` + Enabled bool `json:"enabled" db:"enabled"` + CreatedAt time.Time `json:"created_at" db:"created_at"` + UpdatedAt time.Time `json:"updated_at" db:"updated_at"` +} + +type UsagePricingQuery struct { + TenantID uuid.UUID + ProviderID uuid.UUID + ProviderType string + ModelID string + Limit int +} + +type ResolvedUsagePricing struct { + ModelID string `json:"model_id"` + ProviderID uuid.UUID `json:"provider_id,omitempty"` + ProviderType string `json:"provider_type,omitempty"` + Source string `json:"source"` + Pricing UsagePricingFields `json:"pricing"` + CatalogSynced *time.Time `json:"catalog_synced_at,omitempty"` + OverrideID uuid.UUID `json:"override_id,omitempty"` +} + +type UsageCapPolicy struct { + ID uuid.UUID `json:"id" db:"id"` + TenantID uuid.UUID `json:"tenant_id" db:"tenant_id"` + AgentID *uuid.UUID `json:"agent_id,omitempty" db:"agent_id"` + ProviderID *uuid.UUID `json:"provider_id,omitempty" db:"provider_id"` + ProviderType string `json:"provider_type,omitempty" db:"provider_type"` + ModelID string `json:"model_id,omitempty" db:"model_id"` + Window string `json:"window" db:"window"` + MaxTokens *int64 `json:"max_tokens,omitempty" db:"max_tokens"` + MaxCostMicros *int64 `json:"max_cost_micros,omitempty" db:"max_cost_micros"` + Enabled bool `json:"enabled" db:"enabled"` + Priority int `json:"priority" db:"priority"` + CreatedAt time.Time `json:"created_at" db:"created_at"` + UpdatedAt time.Time `json:"updated_at" db:"updated_at"` +} + +type UsageCapPolicyPatch struct { + AgentID **uuid.UUID + ProviderID **uuid.UUID + ProviderType *string + ModelID *string + Window *string + MaxTokens **int64 + MaxCostMicros **int64 + Enabled *bool + Priority *int +} + +type UsageCapScope struct { + TenantID uuid.UUID + AgentID uuid.UUID + ProviderID uuid.UUID + ProviderType string + ModelID string +} + +type UsageReserveRequest struct { + UsageCapScope + ReservationKey string + EstimatedTokens int64 + EstimatedCostMicros int64 + Metadata json.RawMessage +} + +type UsageReservationResult struct { + ReservationKey string `json:"reservation_key"` + Policies []UsageCapPolicy `json:"policies"` + Skipped bool `json:"skipped,omitempty"` + Reason string `json:"reason,omitempty"` +} + +type UsageReconcileRequest struct { + ReservationKey string + ActualTokens int64 + ActualCostMicros int64 + Status string + Metadata json.RawMessage +} + +type UsageCapUtilization struct { + Policy UsageCapPolicy `json:"policy"` + WindowStart time.Time `json:"window_start"` + WindowEnd time.Time `json:"window_end"` + UsedTokens int64 `json:"used_tokens"` + ReservedTokens int64 `json:"reserved_tokens"` + UsedCostMicros int64 `json:"used_cost_micros"` + ReservedCostMicros int64 `json:"reserved_cost_micros"` +} + +type UsageCapEvent struct { + ID uuid.UUID `json:"id" db:"id"` + TenantID uuid.UUID `json:"tenant_id" db:"tenant_id"` + PolicyID *uuid.UUID `json:"policy_id,omitempty" db:"policy_id"` + ReservationKey string `json:"reservation_key,omitempty" db:"reservation_key"` + Decision string `json:"decision" db:"decision"` + Reason string `json:"reason,omitempty" db:"reason"` + EstimatedTokens int64 `json:"estimated_tokens" db:"estimated_tokens"` + EstimatedCostMicros int64 `json:"estimated_cost_micros" db:"estimated_cost_micros"` + ActualTokens int64 `json:"actual_tokens" db:"actual_tokens"` + ActualCostMicros int64 `json:"actual_cost_micros" db:"actual_cost_micros"` + Metadata json.RawMessage `json:"metadata,omitempty" db:"metadata"` + CreatedAt time.Time `json:"created_at" db:"created_at"` +} + +type UsageCapStore interface { + UpsertPricingCatalog(ctx context.Context, entries []UsagePricingCatalogEntry) (int, error) + ListPricingCatalog(ctx context.Context, q UsagePricingQuery) ([]UsagePricingCatalogEntry, error) + PutPricingOverride(ctx context.Context, o *UsagePricingOverride) error + ListPricingOverrides(ctx context.Context, q UsagePricingQuery) ([]UsagePricingOverride, error) + DeletePricingOverride(ctx context.Context, tenantID, id uuid.UUID) error + ResolvePricing(ctx context.Context, tenantID, providerID uuid.UUID, providerName, providerType, modelID string) (*ResolvedUsagePricing, error) + + CreateUsageCapPolicy(ctx context.Context, p *UsageCapPolicy) error + ListUsageCapPolicies(ctx context.Context, scope UsageCapScope, includeDisabled bool) ([]UsageCapPolicy, error) + UpdateUsageCapPolicy(ctx context.Context, tenantID, id uuid.UUID, patch UsageCapPolicyPatch) (*UsageCapPolicy, error) + DeleteUsageCapPolicy(ctx context.Context, tenantID, id uuid.UUID) error + ReserveUsage(ctx context.Context, req UsageReserveRequest, policies []UsageCapPolicy) (*UsageReservationResult, error) + ReconcileUsage(ctx context.Context, req UsageReconcileRequest) error + ListUsageCapUtilization(ctx context.Context, tenantID uuid.UUID) ([]UsageCapUtilization, error) + ListUsageCapEvents(ctx context.Context, tenantID uuid.UUID, limit int) ([]UsageCapEvent, error) + InsertUsageCapEvent(ctx context.Context, event *UsageCapEvent) error +} diff --git a/internal/tools/read_audio.go b/internal/tools/read_audio.go index e0a92283..a329dd26 100644 --- a/internal/tools/read_audio.go +++ b/internal/tools/read_audio.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // --- Context helpers for media audio --- @@ -45,12 +46,17 @@ var audioModelDefaults = map[string]string{ type ReadAudioTool struct { registry *providers.Registry mediaLoader MediaPathLoader + usageCaps *usagecaps.Service } func NewReadAudioTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadAudioTool { return &ReadAudioTool{registry: registry, mediaLoader: mediaLoader} } +func (t *ReadAudioTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadAudioTool) Name() string { return "read_audio" } func (t *ReadAudioTool) Description() string { diff --git a/internal/tools/read_audio_resolve.go b/internal/tools/read_audio_resolve.go index 65e27d74..f7f62a1e 100644 --- a/internal/tools/read_audio_resolve.go +++ b/internal/tools/read_audio_resolve.go @@ -93,7 +93,19 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, // providers exposing a /v1/audio/transcriptions endpoint. if isTranscriptionModel(model) { slog.Info("read_audio: using openai transcription API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := openaiTranscriptionCall(ctx, cp.APIKey(), cp.APIBase(), model, data, mime) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("openai transcription call: %w", err) } @@ -103,7 +115,19 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, // Gemini: use File API (inlineData doesn't work for audio). if ptype == "gemini" { slog.Info("read_audio: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 120*time.Second) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("gemini file API: %w", err) } @@ -113,7 +137,19 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, // Native OpenAI chat-audio (gpt-4o-audio-preview etc.): input_audio content part. if ptype == "openai" { slog.Info("read_audio: using openai input_audio API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := openaiAudioCall(ctx, cp.APIKey(), cp.APIBase(), model, prompt, data, mime) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("openai audio call: %w", err) } @@ -128,7 +164,7 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, } slog.Info("read_audio: using chat API fallback", "provider", providerName, "model", model, "size", len(data)) - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -141,7 +177,15 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, "max_tokens": 16384, "temperature": 0.2, }, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("chat API: %w", err) } diff --git a/internal/tools/read_document.go b/internal/tools/read_document.go index 7b3952df..4167d76a 100644 --- a/internal/tools/read_document.go +++ b/internal/tools/read_document.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // textReadableMIMEs are MIME types whose content can be returned directly without LLM analysis. @@ -68,12 +69,17 @@ var documentModelDefaults = map[string]string{ type ReadDocumentTool struct { registry *providers.Registry mediaLoader MediaPathLoader + usageCaps *usagecaps.Service } func NewReadDocumentTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadDocumentTool { return &ReadDocumentTool{registry: registry, mediaLoader: mediaLoader} } +func (t *ReadDocumentTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadDocumentTool) Name() string { return "read_document" } func (t *ReadDocumentTool) Description() string { diff --git a/internal/tools/read_document_resolve.go b/internal/tools/read_document_resolve.go index 47d72cf9..6f5bfd9f 100644 --- a/internal/tools/read_document_resolve.go +++ b/internal/tools/read_document_resolve.go @@ -152,7 +152,23 @@ func (t *ReadDocumentTool) callProvider(ctx context.Context, cp credentialProvid slog.Info("read_document: using gemini native API", "provider", providerName, "model", model, "doc_size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{ + Role: "user", + Content: prompt, + Images: []providers.ImageContent{{MimeType: mime}}, + }}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := geminiNativeDocumentCall(ctx, cp.APIKey(), model, prompt, data, mime) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("gemini native call: %w", err) } @@ -179,7 +195,7 @@ func (t *ReadDocumentTool) callProvider(ctx context.Context, cp credentialProvid opts["disable_tools"] = true } - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -189,7 +205,15 @@ func (t *ReadDocumentTool) callProvider(ctx context.Context, cp credentialProvid }, Model: model, Options: opts, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("chat call: %w", err) } diff --git a/internal/tools/read_image.go b/internal/tools/read_image.go index 9d754d5c..892a6dfb 100644 --- a/internal/tools/read_image.go +++ b/internal/tools/read_image.go @@ -10,6 +10,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // --- Context helpers for media images --- @@ -46,13 +47,18 @@ var visionModelDefaults = map[string]string{ // ReadImageTool uses a vision-capable provider to describe images attached to the current message. type ReadImageTool struct { - registry *providers.Registry + registry *providers.Registry + usageCaps *usagecaps.Service } func NewReadImageTool(registry *providers.Registry) *ReadImageTool { return &ReadImageTool{registry: registry} } +func (t *ReadImageTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadImageTool) Name() string { return "read_image" } func (t *ReadImageTool) Description() string { @@ -152,7 +158,7 @@ func (t *ReadImageTool) callProvider(ctx context.Context, cp credentialProvider, opts["disable_tools"] = true } - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -162,7 +168,15 @@ func (t *ReadImageTool) callProvider(ctx context.Context, cp credentialProvider, }, Model: model, Options: opts, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("vision provider error: %w", err) } diff --git a/internal/tools/read_video.go b/internal/tools/read_video.go index b6be69fb..46df93c4 100644 --- a/internal/tools/read_video.go +++ b/internal/tools/read_video.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // --- Context helpers for media video --- @@ -45,12 +46,17 @@ var videoModelDefaults = map[string]string{ type ReadVideoTool struct { registry *providers.Registry mediaLoader MediaPathLoader + usageCaps *usagecaps.Service } func NewReadVideoTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadVideoTool { return &ReadVideoTool{registry: registry, mediaLoader: mediaLoader} } +func (t *ReadVideoTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadVideoTool) Name() string { return "read_video" } func (t *ReadVideoTool) Description() string { diff --git a/internal/tools/read_video_resolve.go b/internal/tools/read_video_resolve.go index 8a67d925..f912fb96 100644 --- a/internal/tools/read_video_resolve.go +++ b/internal/tools/read_video_resolve.go @@ -70,7 +70,19 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, ptype := GetParamString(params, "_provider_type", providerTypeFromName(providerName)) if cp != nil && ptype == "gemini" { slog.Info("read_video: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 180*time.Second) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("gemini file API: %w", err) } @@ -84,7 +96,7 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, } slog.Info("read_video: using chat API fallback", "provider", providerName, "model", model, "size", len(data)) - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -97,7 +109,15 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, "max_tokens": 16384, "temperature": 0.2, }, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("chat API: %w", err) } diff --git a/internal/tools/subagent.go b/internal/tools/subagent.go index 865ee178..b642e1d9 100644 --- a/internal/tools/subagent.go +++ b/internal/tools/subagent.go @@ -18,6 +18,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // SubagentConfig configures the subagent system. @@ -40,33 +41,34 @@ const ( // SubagentTask tracks a running or completed subagent. type SubagentTask struct { - ID string `json:"id"` - ParentID string `json:"parentId"` - Task string `json:"task"` - Label string `json:"label"` - Status string `json:"status"` // "running", "completed", "failed", "cancelled" - Result string `json:"result,omitempty"` - Depth int `json:"depth"` - Model string `json:"model,omitempty"` // model override for this subagent - TotalInputTokens int64 `json:"totalInputTokens,omitempty"` - TotalOutputTokens int64 `json:"totalOutputTokens,omitempty"` - OriginChannel string `json:"originChannel,omitempty"` - OriginChatID string `json:"originChatId,omitempty"` - OriginPeerKind string `json:"originPeerKind,omitempty"` // "direct" or "group" (for session key building) - OriginLocalKey string `json:"originLocalKey,omitempty"` // composite key with topic/thread suffix for routing - OriginUserID string `json:"originUserId,omitempty"` // parent's userID for per-user scoping propagation - OriginSenderID string `json:"originSenderId,omitempty"` // real acting sender; preserves permission attribution in announce re-ingress (#915) - OriginRole string `json:"originRole,omitempty"` // parent's RBAC role; bypasses per-user grants for admin/operator/owner in re-ingress (#915) - OriginSessionKey string `json:"originSessionKey,omitempty"` // exact parent session key for announce routing (WS uses non-standard format) - CreatedAt int64 `json:"createdAt"` - CompletedAt int64 `json:"completedAt,omitempty"` - Media []bus.MediaFile `json:"-"` // media files from tool results - OriginTenantID uuid.UUID `json:"-"` // parent's tenant for announce routing - OriginTraceID uuid.UUID `json:"-"` // parent trace for announce linking - OriginRootSpanID uuid.UUID `json:"-"` // parent agent's root span ID - cancelFunc context.CancelFunc `json:"-"` // per-task context cancel - spawnConfig SubagentConfig `json:"-"` // resolved config at spawn time (per-agent override merged) - dbID uuid.UUID `json:"-"` // persistent DB UUID (zero if not persisted) + ID string `json:"id"` + ParentID string `json:"parentId"` + Task string `json:"task"` + Label string `json:"label"` + Status string `json:"status"` // "running", "completed", "failed", "cancelled" + Result string `json:"result,omitempty"` + Depth int `json:"depth"` + Model string `json:"model,omitempty"` // model override for this subagent + TotalInputTokens int64 `json:"totalInputTokens,omitempty"` + TotalOutputTokens int64 `json:"totalOutputTokens,omitempty"` + OriginChannel string `json:"originChannel,omitempty"` + OriginChatID string `json:"originChatId,omitempty"` + OriginPeerKind string `json:"originPeerKind,omitempty"` // "direct" or "group" (for session key building) + OriginLocalKey string `json:"originLocalKey,omitempty"` // composite key with topic/thread suffix for routing + OriginUserID string `json:"originUserId,omitempty"` // parent's userID for per-user scoping propagation + OriginSenderID string `json:"originSenderId,omitempty"` // real acting sender; preserves permission attribution in announce re-ingress (#915) + OriginRole string `json:"originRole,omitempty"` // parent's RBAC role; bypasses per-user grants for admin/operator/owner in re-ingress (#915) + OriginSessionKey string `json:"originSessionKey,omitempty"` // exact parent session key for announce routing (WS uses non-standard format) + CreatedAt int64 `json:"createdAt"` + CompletedAt int64 `json:"completedAt,omitempty"` + Media []bus.MediaFile `json:"-"` // media files from tool results + OriginAgentID uuid.UUID `json:"-"` // parent agent UUID for usage caps and scoped tools + OriginTenantID uuid.UUID `json:"-"` // parent's tenant for announce routing + OriginTraceID uuid.UUID `json:"-"` // parent trace for announce linking + OriginRootSpanID uuid.UUID `json:"-"` // parent agent's root span ID + cancelFunc context.CancelFunc `json:"-"` // per-task context cancel + spawnConfig SubagentConfig `json:"-"` // resolved config at spawn time (per-agent override merged) + dbID uuid.UUID `json:"-"` // persistent DB UUID (zero if not persisted) } // SubagentManager manages the lifecycle of spawned subagents. @@ -74,8 +76,8 @@ type SubagentManager struct { mu sync.RWMutex tasks map[string]*SubagentTask config SubagentConfig - provider providers.Provider // default provider (fallback) - providerReg *providers.Registry // registry for resolving parent's provider + provider providers.Provider // default provider (fallback) + providerReg *providers.Registry // registry for resolving parent's provider model string msgBus *bus.MessageBus @@ -83,6 +85,7 @@ type SubagentManager struct { createTools func() *Registry announceQueue *AnnounceQueue // optional: batches announces with debounce taskStore store.SubagentTaskStore // optional: persists tasks to DB (fire-and-forget) + usageCaps *usagecaps.Service } // NewSubagentManager creates a new subagent manager. @@ -116,6 +119,10 @@ func (sm *SubagentManager) SetTaskStore(s store.SubagentTaskStore) { sm.taskStore = s } +func (sm *SubagentManager) SetUsageCapService(s *usagecaps.Service) { + sm.usageCaps = s +} + // effectiveConfig returns the per-agent context override merged with defaults, // or falls back to sm.config when no override is present. func (sm *SubagentManager) effectiveConfig(ctx context.Context) SubagentConfig { diff --git a/internal/tools/subagent_exec.go b/internal/tools/subagent_exec.go index cfd38de1..15cc7c49 100644 --- a/internal/tools/subagent_exec.go +++ b/internal/tools/subagent_exec.go @@ -13,6 +13,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // runTask executes the subagent in a goroutine. @@ -253,10 +254,7 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) } slog.Info("subagent LLM retry", "id", task.ID, "iteration", iteration, "attempt", attempt+1) } - // ctx is the parent agent's run context — cancelling the parent (e.g. agent abort) - // cascades here and to all subsequent tool calls in this iteration. - // Do NOT replace ctx with context.Background() here; that would detach abort propagation. - resp, err = activeProvider.Chat(ctx, chatReq) + resp, err = sm.chatSubagentWithUsageCap(ctx, task, activeProvider, model, chatReq, iteration, attempt+1) if err == nil { break } @@ -344,3 +342,44 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) return iteration } + +func (sm *SubagentManager) chatSubagentWithUsageCap(ctx context.Context, task *SubagentTask, activeProvider providers.Provider, model string, chatReq providers.ChatRequest, iteration, attempt int) (*providers.ChatResponse, error) { + if fallbackProvider, ok := activeProvider.(*providers.ModelFallbackProvider); ok { + before := func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + reservation, err := sm.reserveSubagentUsage(callCtx, task, entry.ProviderName, actualReq.Model, actualReq, fmt.Sprintf("%d:%d:%s:%s", iteration, attempt, entry.ProviderName, actualReq.Model)) + if err != nil { + return nil, err + } + return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) { + if reservation != nil { + reservation.Reconcile(callCtx, resp, callErr) + } + }, nil + } + return fallbackProvider.ChatWithHook(ctx, chatReq, before) + } + reservation, err := sm.reserveSubagentUsage(ctx, task, activeProvider.Name(), model, chatReq, fmt.Sprintf("%d:%d", iteration, attempt)) + if err != nil { + return nil, err + } + resp, err := activeProvider.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } + return resp, err +} + +func (sm *SubagentManager) reserveSubagentUsage(ctx context.Context, task *SubagentTask, providerName, model string, chatReq providers.ChatRequest, suffix string) (*usagecaps.Reservation, error) { + if sm.usageCaps == nil { + return nil, nil + } + return sm.usageCaps.Preflight(ctx, usagecaps.Request{ + TenantID: task.OriginTenantID, + AgentID: task.OriginAgentID, + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("%s:%s", task.ID, suffix), + Messages: chatReq.Messages, + MaxOutputTokens: 4096, + }) +} diff --git a/internal/tools/subagent_spawn.go b/internal/tools/subagent_spawn.go index c9e43d05..0f44430b 100644 --- a/internal/tools/subagent_spawn.go +++ b/internal/tools/subagent_spawn.go @@ -82,16 +82,17 @@ func (sm *SubagentManager) Spawn( OriginChannel: channel, OriginChatID: chatID, OriginPeerKind: peerKind, - OriginLocalKey: ToolLocalKeyFromCtx(ctx), - OriginUserID: store.UserIDFromContext(ctx), - OriginSenderID: store.SenderIDFromContext(ctx), - OriginRole: store.RoleFromContext(ctx), - OriginSessionKey: ToolSessionKeyFromCtx(ctx), - OriginTenantID: store.TenantIDFromContext(ctx), - OriginTraceID: tracing.TraceIDFromContext(ctx), - OriginRootSpanID: tracing.ParentSpanIDFromContext(ctx), - CreatedAt: time.Now().UnixMilli(), - spawnConfig: cfg, + OriginLocalKey: ToolLocalKeyFromCtx(ctx), + OriginUserID: store.UserIDFromContext(ctx), + OriginSenderID: store.SenderIDFromContext(ctx), + OriginRole: store.RoleFromContext(ctx), + OriginSessionKey: ToolSessionKeyFromCtx(ctx), + OriginAgentID: store.AgentIDFromContext(ctx), + OriginTenantID: store.TenantIDFromContext(ctx), + OriginTraceID: tracing.TraceIDFromContext(ctx), + OriginRootSpanID: tracing.ParentSpanIDFromContext(ctx), + CreatedAt: time.Now().UnixMilli(), + spawnConfig: cfg, } // Detach from parent's cancellation chain so subagent survives after parent run completes. // WithoutCancel preserves all context values (agent ID, workspace, trace info, etc.) @@ -166,6 +167,7 @@ func (sm *SubagentManager) RunSync( OriginSenderID: store.SenderIDFromContext(ctx), OriginRole: store.RoleFromContext(ctx), OriginSessionKey: ToolSessionKeyFromCtx(ctx), + OriginAgentID: store.AgentIDFromContext(ctx), OriginTenantID: store.TenantIDFromContext(ctx), OriginTraceID: tracing.TraceIDFromContext(ctx), OriginRootSpanID: tracing.ParentSpanIDFromContext(ctx), diff --git a/internal/tools/usage_caps.go b/internal/tools/usage_caps.go new file mode 100644 index 00000000..0dc59e94 --- /dev/null +++ b/internal/tools/usage_caps.go @@ -0,0 +1,44 @@ +package tools + +import ( + "context" + "fmt" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" +) + +func reserveToolLLMUsage(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest) (*usagecaps.Reservation, error) { + if svc == nil { + return nil, nil + } + return svc.Preflight(ctx, usagecaps.Request{ + TenantID: store.TenantIDFromContext(ctx), + AgentID: store.AgentIDFromContext(ctx), + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("tool:%s:%s", toolName, uuid.NewString()), + Messages: req.Messages, + MaxOutputTokens: maxOutputTokensFromOptions(req.Options), + }) +} + +func maxOutputTokensFromOptions(options map[string]any) int { + maxTokens := 1024 + if options == nil { + return maxTokens + } + if v, ok := options[providers.OptMaxTokens]; ok { + switch n := v.(type) { + case int: + maxTokens = n + case int64: + maxTokens = int(n) + case float64: + maxTokens = int(n) + } + } + return maxTokens +} diff --git a/internal/upgrade/version.go b/internal/upgrade/version.go index db629515..ad21b8a5 100644 --- a/internal/upgrade/version.go +++ b/internal/upgrade/version.go @@ -2,4 +2,4 @@ package upgrade // RequiredSchemaVersion is the schema migration version this binary requires. // Bump this whenever adding a new SQL migration file. -const RequiredSchemaVersion uint = 68 +const RequiredSchemaVersion uint = 70 diff --git a/internal/usage/caps/service.go b/internal/usage/caps/service.go new file mode 100644 index 00000000..12db8dd9 --- /dev/null +++ b/internal/usage/caps/service.go @@ -0,0 +1,275 @@ +package caps + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "log/slog" + "strings" + "time" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/usage/pricing" +) + +var ( + ErrCapExceeded = errors.New("usage cap exceeded") + ErrPricingUnknown = pricing.ErrUnknownPricing +) + +type Service struct { + store store.UsageCapStore + providers store.ProviderStore +} + +func NewService(s store.UsageCapStore, providers store.ProviderStore) *Service { + if s == nil { + return nil + } + return &Service{store: s, providers: providers} +} + +type Request struct { + TenantID uuid.UUID + AgentID uuid.UUID + ProviderName string + ModelID string + ReservationKey string + Messages []providers.Message + MaxOutputTokens int +} + +type Reservation struct { + key string + result *store.UsageReservationResult + svc *Service + usage pricing.BillableUsage + prices store.UsagePricingFields + estimatedCostMicros int64 + skipped bool + reason string +} + +func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, error) { + if s == nil || s.store == nil { + return &Reservation{skipped: true, reason: "service_disabled"}, nil + } + providerData, err := s.resolveProvider(ctx, req.TenantID, req.ProviderName) + if err != nil { + return &Reservation{skipped: true, reason: "provider_metadata_missing"}, nil + } + scope := store.UsageCapScope{ + TenantID: req.TenantID, AgentID: req.AgentID, ProviderID: providerData.ID, + ProviderType: providerData.ProviderType, ModelID: req.ModelID, + } + if !ShouldEnforceProvider(providerData.ProviderType, providerData.APIKey != "") { + _ = s.store.InsertUsageCapEvent(ctx, &store.UsageCapEvent{ + TenantID: req.TenantID, Decision: store.UsageCapEventSkip, + Reason: "provider_not_billable_api", Metadata: mustJSON(scope), + }) + return &Reservation{skipped: true, reason: "provider_not_billable_api"}, nil + } + policies, err := s.store.ListUsageCapPolicies(ctx, scope, false) + if err != nil { + return nil, err + } + if len(policies) == 0 { + return &Reservation{skipped: true, reason: "no_policy"}, nil + } + usage := pricing.BillableUsage{ + InputTokens: int64(EstimateInputTokens(req.Messages)), + OutputTokens: int64(req.MaxOutputTokens), + ImageCount: int64(CountImages(req.Messages)), + } + if usage.OutputTokens <= 0 { + usage.OutputTokens = 1 + } + metadata := map[string]any{"model_id": req.ModelID, "pricing_source": "token_only"} + var prices store.UsagePricingFields + var costMicros int64 + if requiresCostCap(policies) { + resolved, err := s.store.ResolvePricing(ctx, req.TenantID, providerData.ID, providerData.Name, providerData.ProviderType, req.ModelID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID) + } + return nil, err + } + prices = resolved.Pricing + if prices.Request != nil { + usage.RequestCount = 1 + } + costMicros, err = pricing.CostMicros(prices, usage) + if err != nil { + return nil, err + } + metadata = map[string]any{"source": resolved.Source, "model_id": resolved.ModelID} + } + key := req.ReservationKey + if key == "" { + key = uuid.NewString() + } + result, err := s.store.ReserveUsage(ctx, store.UsageReserveRequest{ + UsageCapScope: scope, ReservationKey: key, + EstimatedTokens: usage.TotalTokens(), EstimatedCostMicros: costMicros, + Metadata: mustJSON(metadata), + }, policies) + if err != nil { + var capErr *store.UsageCapExceededError + if errors.As(err, &capErr) || errors.Is(err, store.ErrUsageCapExceeded) { + blockedPolicy := uuid.Nil + reason := "cap_exceeded" + if capErr != nil { + blockedPolicy = capErr.PolicyID + if capErr.Reason != "" { + reason = capErr.Reason + } + } + _ = s.store.InsertUsageCapEvent(ctx, &store.UsageCapEvent{ + TenantID: req.TenantID, PolicyID: optionalPolicyID(blockedPolicy), + ReservationKey: key, Decision: store.UsageCapEventBlock, Reason: reason, + EstimatedTokens: usage.TotalTokens(), EstimatedCostMicros: costMicros, + Metadata: mustJSON(map[string]any{"model_id": req.ModelID, "provider": req.ProviderName}), + }) + slog.Warn("usage_caps.blocked", "tenant_id", req.TenantID, "policy_id", blockedPolicy, "reason", reason) + return nil, ErrCapExceeded + } + return nil, err + } + return &Reservation{key: key, result: result, svc: s, usage: usage, prices: prices, estimatedCostMicros: costMicros}, nil +} + +func (r *Reservation) Reconcile(ctx context.Context, resp *providers.ChatResponse, callErr error) { + r.reconcile(ctx, resp, callErr, false) +} + +func (r *Reservation) ReconcileStream(ctx context.Context, resp *providers.ChatResponse, callErr error, streamed bool) { + r.reconcile(ctx, resp, callErr, streamed) +} + +func (r *Reservation) reconcile(ctx context.Context, resp *providers.ChatResponse, callErr error, keepEstimateOnError bool) { + if r == nil || r.svc == nil || r.key == "" || r.skipped || r.result == nil || len(r.result.Policies) == 0 { + return + } + actual := r.usage + if resp != nil && resp.Usage != nil { + actual = pricing.FromProviderUsage(resp.Usage) + if r.prices.Request == nil { + actual.RequestCount = 0 + } + if actual.RequestCount == 0 && r.usage.RequestCount > 0 { + actual.RequestCount = r.usage.RequestCount + } + if actual.ImageCount == 0 { + actual.ImageCount = r.usage.ImageCount + } + if actual.OutputTokens == 0 { + actual.OutputTokens = r.usage.OutputTokens + } + if actual.InputTokens == 0 { + actual.InputTokens = r.usage.InputTokens + } + } + status := "reconciled" + if callErr != nil { + status = "failed" + if resp == nil || resp.Usage == nil { + if keepEstimateOnError { + actual = r.usage + } else { + actual = pricing.BillableUsage{} + } + } + } + cost, err := pricing.CostMicros(r.prices, actual) + if err != nil { + cost = r.estimatedCostMicros + } + reconcileCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if err := r.svc.store.ReconcileUsage(reconcileCtx, store.UsageReconcileRequest{ + ReservationKey: r.key, ActualTokens: actual.TotalTokens(), + ActualCostMicros: cost, Status: status, + }); err != nil { + slog.Warn("usage_caps.reconcile_failed", "reservation_key", r.key, "error", err) + } +} + +func (s *Service) resolveProvider(ctx context.Context, tenantID uuid.UUID, name string) (*store.LLMProviderData, error) { + if s.providers == nil || strings.TrimSpace(name) == "" { + return nil, sql.ErrNoRows + } + p, err := s.providers.GetProviderByName(ctx, name) + if err == nil { + return p, nil + } + if tenantID != uuid.Nil && tenantID != store.MasterTenantID { + if fallback, fallbackErr := s.providers.GetProviderByName(store.WithTenantID(ctx, store.MasterTenantID), name); fallbackErr == nil { + return fallback, nil + } + } + return nil, err +} + +func ShouldEnforceProvider(providerType string, hasAPIKey bool) bool { + switch providerType { + case store.ProviderChatGPTOAuth, store.ProviderClaudeCLI, store.ProviderBailian, store.ProviderACP, store.ProviderOllama: + return false + default: + return hasAPIKey + } +} + +func CountImages(messages []providers.Message) int { + count := 0 + for _, msg := range messages { + for _, img := range msg.Images { + if strings.HasPrefix(strings.ToLower(img.MimeType), "image/") { + count++ + } + } + } + return count +} + +func EstimateInputTokens(messages []providers.Message) int { + total := 0 + for _, msg := range messages { + total += len(msg.Content) / 4 + if len(msg.Content)%4 != 0 { + total++ + } + } + if total <= 0 { + return 1 + } + return total +} + +func requiresCostCap(policies []store.UsageCapPolicy) bool { + for _, p := range policies { + if p.MaxCostMicros != nil { + return true + } + } + return false +} + +func mustJSON(v any) json.RawMessage { + b, _ := json.Marshal(v) + if len(b) == 0 { + return json.RawMessage(`{}`) + } + return b +} + +func optionalPolicyID(id uuid.UUID) *uuid.UUID { + if id == uuid.Nil { + return nil + } + return &id +} diff --git a/internal/usage/caps/service_test.go b/internal/usage/caps/service_test.go new file mode 100644 index 00000000..e1350d6c --- /dev/null +++ b/internal/usage/caps/service_test.go @@ -0,0 +1,337 @@ +package caps + +import ( + "context" + "database/sql" + "testing" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func TestShouldEnforceProvider(t *testing.T) { + cases := []struct { + providerType string + hasKey bool + want bool + }{ + {store.ProviderChatGPTOAuth, false, false}, + {store.ProviderClaudeCLI, false, false}, + {store.ProviderBailian, false, false}, + {store.ProviderOllama, false, false}, + {store.ProviderACP, false, false}, + {store.ProviderOpenAICompat, true, true}, + {store.ProviderOpenRouter, true, true}, + {store.ProviderOpenRouter, false, false}, + } + for _, tc := range cases { + if got := ShouldEnforceProvider(tc.providerType, tc.hasKey); got != tc.want { + t.Fatalf("ShouldEnforceProvider(%q,%v) = %v, want %v", tc.providerType, tc.hasKey, got, tc.want) + } + } +} + +func TestPreflightTokenOnlyCapDoesNotRequirePricing(t *testing.T) { + providerID := uuid.New() + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(1000), Enabled: true} + usageStore := &fakeUsageCapStore{policies: []store.UsageCapPolicy{policy}, resolveErr: sql.ErrNoRows} + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: providerID}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + + reservation, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "missing/model", + ReservationKey: "token-only", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 10, + }) + if err != nil { + t.Fatalf("Preflight returned error: %v", err) + } + if reservation == nil || reservation.skipped { + t.Fatalf("Preflight skipped token-only policy") + } + if usageStore.resolveCalls != 0 { + t.Fatalf("ResolvePricing called %d time(s), want 0", usageStore.resolveCalls) + } + if usageStore.reserved.EstimatedCostMicros != 0 { + t.Fatalf("EstimatedCostMicros = %d, want 0", usageStore.reserved.EstimatedCostMicros) + } +} + +func TestPreflightIncludesRequestPricingWhenConfigured(t *testing.T) { + zero := "0" + requestPrice := "0.01" + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxCostMicros: int64Ptr(20_000), Enabled: true} + usageStore := &fakeUsageCapStore{ + policies: []store.UsageCapPolicy{policy}, + resolved: &store.ResolvedUsagePricing{ + ModelID: "priced/model", + Source: "catalog", + Pricing: store.UsagePricingFields{Input: &zero, Output: &zero, Request: &requestPrice}, + }, + } + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + + _, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "priced/model", + ReservationKey: "request-fee", Messages: []providers.Message{{Role: "user", Content: "abcd"}}, + MaxOutputTokens: 1, + }) + if err != nil { + t.Fatalf("Preflight returned error: %v", err) + } + if got := usageStore.reserved.EstimatedCostMicros; got != 10_000 { + t.Fatalf("EstimatedCostMicros = %d, want 10000", got) + } +} + +func TestPreflightFallsBackToMasterProviderMetadata(t *testing.T) { + tenantID := uuid.New() + masterProviderID := uuid.New() + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: tenantID, MaxTokens: int64Ptr(1000), Enabled: true} + usageStore := &fakeUsageCapStore{policies: []store.UsageCapPolicy{policy}} + providerStore := &fakeProviderStore{ + masterProvider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: masterProviderID}, + TenantID: store.MasterTenantID, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }, + } + svc := NewService(usageStore, providerStore) + + reservation, err := svc.Preflight(store.WithTenantID(context.Background(), tenantID), Request{ + TenantID: tenantID, ProviderName: "openrouter", ModelID: "openai/gpt-test", + ReservationKey: "master-provider", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 10, + }) + if err != nil { + t.Fatalf("Preflight returned error: %v", err) + } + if reservation == nil || reservation.skipped { + t.Fatal("Preflight skipped master provider fallback") + } + if usageStore.reserved.ProviderID != masterProviderID { + t.Fatalf("ProviderID = %s, want %s", usageStore.reserved.ProviderID, masterProviderID) + } +} + +func TestReservationReconcileUsesDetachedContext(t *testing.T) { + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(1000), Enabled: true} + usageStore := &fakeUsageCapStore{policies: []store.UsageCapPolicy{policy}} + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + reservation, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "token/model", + ReservationKey: "reconcile", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 10, + }) + if err != nil { + t.Fatalf("Preflight returned error: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + reservation.Reconcile(ctx, &providers.ChatResponse{Usage: &providers.Usage{PromptTokens: 2, CompletionTokens: 3}}, nil) + + if usageStore.reconcileCalls != 1 { + t.Fatalf("ReconcileUsage calls = %d, want 1", usageStore.reconcileCalls) + } + if usageStore.reconcileCtxCanceled { + t.Fatal("ReconcileUsage received canceled context") + } +} + +func TestReservationReconcileStreamKeepsEstimateAfterPartialError(t *testing.T) { + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(1000), Enabled: true} + usageStore := &fakeUsageCapStore{policies: []store.UsageCapPolicy{policy}} + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + reservation, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "token/model", + ReservationKey: "partial-stream", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 10, + }) + if err != nil { + t.Fatalf("Preflight returned error: %v", err) + } + + reservation.ReconcileStream(context.Background(), nil, context.Canceled, true) + + if usageStore.reconciled.ActualTokens == 0 { + t.Fatal("ReconcileStream zeroed actual tokens after partial stream error") + } + if usageStore.reconciled.Status != "failed" { + t.Fatalf("Status = %q, want failed", usageStore.reconciled.Status) + } +} + +func TestReservationReconcileIgnoresUnpricedRequestCount(t *testing.T) { + tokenPrice := "0.000001" + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxCostMicros: int64Ptr(1_000_000), Enabled: true} + usageStore := &fakeUsageCapStore{ + policies: []store.UsageCapPolicy{policy}, + resolved: &store.ResolvedUsagePricing{ + ModelID: "priced/model", + Source: "catalog", + Pricing: store.UsagePricingFields{Input: &tokenPrice, Output: &tokenPrice}, + }, + } + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + reservation, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "priced/model", + ReservationKey: "unpriced-request", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 100, + }) + if err != nil { + t.Fatalf("Preflight returned error: %v", err) + } + + reservation.Reconcile(context.Background(), &providers.ChatResponse{Usage: &providers.Usage{ + PromptTokens: 2, + CompletionTokens: 3, + RequestCount: 1, + }}, nil) + + if usageStore.reconciled.ActualCostMicros != 5 { + t.Fatalf("ActualCostMicros = %d, want 5", usageStore.reconciled.ActualCostMicros) + } +} + +func TestCountImagesOnlyCountsImageMIMEs(t *testing.T) { + messages := []providers.Message{{ + Role: "user", + Images: []providers.ImageContent{ + {MimeType: "image/png"}, + {MimeType: "application/pdf"}, + {MimeType: "video/mp4"}, + }, + }} + if got := CountImages(messages); got != 1 { + t.Fatalf("CountImages = %d, want 1", got) + } +} + +type fakeUsageCapStore struct { + policies []store.UsageCapPolicy + resolved *store.ResolvedUsagePricing + resolveErr error + resolveCalls int + reserved store.UsageReserveRequest + reconciled store.UsageReconcileRequest + reconcileCalls int + reconcileCtxCanceled bool +} + +func (s *fakeUsageCapStore) UpsertPricingCatalog(context.Context, []store.UsagePricingCatalogEntry) (int, error) { + return 0, nil +} +func (s *fakeUsageCapStore) ListPricingCatalog(context.Context, store.UsagePricingQuery) ([]store.UsagePricingCatalogEntry, error) { + return nil, nil +} +func (s *fakeUsageCapStore) PutPricingOverride(context.Context, *store.UsagePricingOverride) error { + return nil +} +func (s *fakeUsageCapStore) ListPricingOverrides(context.Context, store.UsagePricingQuery) ([]store.UsagePricingOverride, error) { + return nil, nil +} +func (s *fakeUsageCapStore) DeletePricingOverride(context.Context, uuid.UUID, uuid.UUID) error { + return nil +} +func (s *fakeUsageCapStore) ResolvePricing(context.Context, uuid.UUID, uuid.UUID, string, string, string) (*store.ResolvedUsagePricing, error) { + s.resolveCalls++ + if s.resolveErr != nil { + return nil, s.resolveErr + } + return s.resolved, nil +} +func (s *fakeUsageCapStore) CreateUsageCapPolicy(context.Context, *store.UsageCapPolicy) error { + return nil +} +func (s *fakeUsageCapStore) ListUsageCapPolicies(context.Context, store.UsageCapScope, bool) ([]store.UsageCapPolicy, error) { + return s.policies, nil +} +func (s *fakeUsageCapStore) UpdateUsageCapPolicy(context.Context, uuid.UUID, uuid.UUID, store.UsageCapPolicyPatch) (*store.UsageCapPolicy, error) { + return nil, nil +} +func (s *fakeUsageCapStore) DeleteUsageCapPolicy(context.Context, uuid.UUID, uuid.UUID) error { + return nil +} +func (s *fakeUsageCapStore) ReserveUsage(_ context.Context, req store.UsageReserveRequest, policies []store.UsageCapPolicy) (*store.UsageReservationResult, error) { + s.reserved = req + return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Policies: policies}, nil +} +func (s *fakeUsageCapStore) ReconcileUsage(ctx context.Context, req store.UsageReconcileRequest) error { + s.reconciled = req + s.reconcileCalls++ + s.reconcileCtxCanceled = ctx.Err() != nil + return nil +} +func (s *fakeUsageCapStore) ListUsageCapUtilization(context.Context, uuid.UUID) ([]store.UsageCapUtilization, error) { + return nil, nil +} +func (s *fakeUsageCapStore) ListUsageCapEvents(context.Context, uuid.UUID, int) ([]store.UsageCapEvent, error) { + return nil, nil +} +func (s *fakeUsageCapStore) InsertUsageCapEvent(context.Context, *store.UsageCapEvent) error { + return nil +} + +type fakeProviderStore struct { + provider *store.LLMProviderData + masterProvider *store.LLMProviderData +} + +func (s *fakeProviderStore) CreateProvider(context.Context, *store.LLMProviderData) error { return nil } +func (s *fakeProviderStore) GetProvider(context.Context, uuid.UUID) (*store.LLMProviderData, error) { + return s.provider, nil +} +func (s *fakeProviderStore) GetProviderByName(ctx context.Context, _ string) (*store.LLMProviderData, error) { + if store.TenantIDFromContext(ctx) == store.MasterTenantID && s.masterProvider != nil { + return s.masterProvider, nil + } + if s.provider == nil { + return nil, sql.ErrNoRows + } + return s.provider, nil +} +func (s *fakeProviderStore) ListProviders(context.Context) ([]store.LLMProviderData, error) { + return nil, nil +} +func (s *fakeProviderStore) ListAllProviders(context.Context) ([]store.LLMProviderData, error) { + return nil, nil +} +func (s *fakeProviderStore) UpdateProvider(context.Context, uuid.UUID, map[string]any) error { + return nil +} +func (s *fakeProviderStore) DeleteProvider(context.Context, uuid.UUID) error { return nil } + +func int64Ptr(v int64) *int64 { return &v } diff --git a/internal/usage/pricing/decimal.go b/internal/usage/pricing/decimal.go new file mode 100644 index 00000000..21976542 --- /dev/null +++ b/internal/usage/pricing/decimal.go @@ -0,0 +1,119 @@ +package pricing + +import ( + "errors" + "fmt" + "math/big" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +var ErrUnknownPricing = errors.New("usage pricing unknown") + +type BillableUsage struct { + InputTokens int64 + OutputTokens int64 + CacheReadTokens int64 + CacheWriteTokens int64 + ReasoningTokens int64 + RequestCount int64 + ImageCount int64 + WebSearchCount int64 +} + +func FromProviderUsage(u *providers.Usage) BillableUsage { + if u == nil { + return BillableUsage{} + } + inputTokens := int64(u.PromptTokens) + cacheReadTokens := int64(u.CacheReadTokens) + cacheWriteTokens := int64(u.CacheCreationTokens) + if u.PromptTokensIncludeCachedSegments { + inputTokens -= cacheReadTokens + cacheWriteTokens + if inputTokens < 0 { + inputTokens = 0 + } + } + return BillableUsage{ + InputTokens: inputTokens, + OutputTokens: int64(u.CompletionTokens), + CacheReadTokens: cacheReadTokens, + CacheWriteTokens: cacheWriteTokens, + ReasoningTokens: int64(u.ThinkingTokens), + RequestCount: int64(u.RequestCount), + ImageCount: int64(u.ImageCount), + WebSearchCount: int64(u.WebSearchCount), + } +} + +func (u BillableUsage) TotalTokens() int64 { + return u.InputTokens + u.OutputTokens + u.CacheReadTokens + u.CacheWriteTokens +} + +func CostMicros(fields store.UsagePricingFields, usage BillableUsage) (int64, error) { + var total int64 + add := func(name string, price *string, units int64) error { + if units <= 0 { + return nil + } + if price == nil || strings.TrimSpace(*price) == "" { + return fmt.Errorf("%w: %s", ErrUnknownPricing, name) + } + micros, err := microsForUnits(*price, units) + if err != nil { + return fmt.Errorf("%s pricing: %w", name, err) + } + total += micros + return nil + } + if err := add("input", fields.Input, usage.InputTokens); err != nil { + return 0, err + } + if usage.ReasoningTokens > 0 && fields.Reasoning != nil { + visible := max(usage.OutputTokens-usage.ReasoningTokens, 0) + if err := add("output", fields.Output, visible); err != nil { + return 0, err + } + if err := add("reasoning", fields.Reasoning, usage.ReasoningTokens); err != nil { + return 0, err + } + } else if err := add("output", fields.Output, usage.OutputTokens); err != nil { + return 0, err + } + if err := add("cache_read", fields.CacheRead, usage.CacheReadTokens); err != nil { + return 0, err + } + if err := add("cache_write", fields.CacheWrite, usage.CacheWriteTokens); err != nil { + return 0, err + } + if err := add("request", fields.Request, usage.RequestCount); err != nil { + return 0, err + } + if err := add("image", fields.Image, usage.ImageCount); err != nil { + return 0, err + } + if err := add("web_search", fields.WebSearch, usage.WebSearchCount); err != nil { + return 0, err + } + return total, nil +} + +func microsForUnits(price string, units int64) (int64, error) { + r, ok := new(big.Rat).SetString(strings.TrimSpace(price)) + if !ok { + return 0, fmt.Errorf("invalid decimal %q", price) + } + r.Mul(r, big.NewRat(units, 1)) + r.Mul(r, big.NewRat(1_000_000, 1)) + q := new(big.Int).Quo(r.Num(), r.Denom()) + rem := new(big.Int).Rem(r.Num(), r.Denom()) + if rem.Sign() > 0 { + q.Add(q, big.NewInt(1)) + } + if !q.IsInt64() { + return 0, errors.New("cost overflows int64 micros") + } + return q.Int64(), nil +} diff --git a/internal/usage/pricing/decimal_test.go b/internal/usage/pricing/decimal_test.go new file mode 100644 index 00000000..515330fb --- /dev/null +++ b/internal/usage/pricing/decimal_test.go @@ -0,0 +1,66 @@ +package pricing + +import ( + "errors" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func strp(s string) *string { return &s } + +func TestCostMicrosTokenAndFlatDimensions(t *testing.T) { + got, err := CostMicros(store.UsagePricingFields{ + Input: strp("0.000001"), + Output: strp("0.000002"), + CacheRead: strp("0.0000001"), + CacheWrite: strp("0.0000002"), + Request: strp("0.01"), + Image: strp("0.02"), + WebSearch: strp("0.03"), + }, BillableUsage{ + InputTokens: 1000, + OutputTokens: 500, + CacheReadTokens: 100, + CacheWriteTokens: 50, + RequestCount: 1, + ImageCount: 2, + WebSearchCount: 1, + }) + if err != nil { + t.Fatal(err) + } + want := int64(1000 + 1000 + 10 + 10 + 10_000 + 40_000 + 30_000) + if got != want { + t.Fatalf("cost micros = %d, want %d", got, want) + } +} + +func TestCostMicrosMissingRequiredPrice(t *testing.T) { + _, err := CostMicros(store.UsagePricingFields{Input: strp("0.000001")}, BillableUsage{ + InputTokens: 1, + OutputTokens: 1, + }) + if !errors.Is(err, ErrUnknownPricing) { + t.Fatalf("err = %v, want ErrUnknownPricing", err) + } +} + +func TestFromProviderUsageSeparatesIncludedCacheTokens(t *testing.T) { + got := FromProviderUsage(&providers.Usage{ + PromptTokens: 2006, + CompletionTokens: 100, + CacheReadTokens: 1920, + PromptTokensIncludeCachedSegments: true, + }) + if got.InputTokens != 86 { + t.Fatalf("InputTokens = %d, want 86", got.InputTokens) + } + if got.CacheReadTokens != 1920 { + t.Fatalf("CacheReadTokens = %d, want 1920", got.CacheReadTokens) + } + if got.TotalTokens() != 2106 { + t.Fatalf("TotalTokens = %d, want 2106", got.TotalTokens()) + } +} diff --git a/internal/usage/pricing/openrouter.go b/internal/usage/pricing/openrouter.go new file mode 100644 index 00000000..98cb84e9 --- /dev/null +++ b/internal/usage/pricing/openrouter.go @@ -0,0 +1,91 @@ +package pricing + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +const OpenRouterModelsURL = "https://openrouter.ai/api/v1/models" + +type openRouterModelsResponse struct { + Data []json.RawMessage `json:"data"` +} + +type openRouterModel struct { + ID string `json:"id"` + Name string `json:"name"` + Pricing map[string]any `json:"pricing"` +} + +func FetchOpenRouterCatalog(ctx context.Context, client *http.Client) ([]store.UsagePricingCatalogEntry, error) { + if client == nil { + client = http.DefaultClient + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, OpenRouterModelsURL, nil) + if err != nil { + return nil, err + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("openrouter models status %d", resp.StatusCode) + } + var payload openRouterModelsResponse + if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil { + return nil, err + } + now := time.Now().UTC() + entries := make([]store.UsagePricingCatalogEntry, 0, len(payload.Data)) + for _, raw := range payload.Data { + var model openRouterModel + if err := json.Unmarshal(raw, &model); err != nil || model.ID == "" { + continue + } + rawPricing, _ := json.Marshal(model.Pricing) + entries = append(entries, store.UsagePricingCatalogEntry{ + ModelID: model.ID, + CanonicalModelID: model.ID, + Pricing: mapOpenRouterPricing(model.Pricing), + RawPricing: rawPricing, + RawModel: raw, + SyncedAt: now, + }) + } + return entries, nil +} + +func mapOpenRouterPricing(raw map[string]any) store.UsagePricingFields { + return store.UsagePricingFields{ + Input: decimalStringPtr(raw["prompt"]), + Output: decimalStringPtr(raw["completion"]), + CacheRead: decimalStringPtr(raw["input_cache_read"]), + CacheWrite: decimalStringPtr(raw["input_cache_write"]), + Reasoning: decimalStringPtr(raw["internal_reasoning"]), + Request: decimalStringPtr(raw["request"]), + Image: decimalStringPtr(raw["image"]), + WebSearch: decimalStringPtr(raw["web_search"]), + } +} + +func decimalStringPtr(v any) *string { + switch x := v.(type) { + case string: + if x == "" { + return nil + } + return &x + case float64: + s := fmt.Sprintf("%.18g", x) + return &s + default: + return nil + } +} diff --git a/migrations/000069_usage_caps_pricing.down.sql b/migrations/000069_usage_caps_pricing.down.sql new file mode 100644 index 00000000..18c41ec8 --- /dev/null +++ b/migrations/000069_usage_caps_pricing.down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS usage_pricing_overrides; +DROP TABLE IF EXISTS usage_pricing_catalog; diff --git a/migrations/000069_usage_caps_pricing.up.sql b/migrations/000069_usage_caps_pricing.up.sql new file mode 100644 index 00000000..d69e35d1 --- /dev/null +++ b/migrations/000069_usage_caps_pricing.up.sql @@ -0,0 +1,45 @@ +CREATE TABLE IF NOT EXISTS usage_pricing_catalog ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + model_id TEXT NOT NULL UNIQUE, + canonical_model_id TEXT, + raw_pricing JSONB NOT NULL DEFAULT '{}'::jsonb, + raw_model JSONB NOT NULL DEFAULT '{}'::jsonb, + input_price NUMERIC(30, 18) CHECK (input_price IS NULL OR input_price >= 0), + output_price NUMERIC(30, 18) CHECK (output_price IS NULL OR output_price >= 0), + cache_read_price NUMERIC(30, 18) CHECK (cache_read_price IS NULL OR cache_read_price >= 0), + cache_write_price NUMERIC(30, 18) CHECK (cache_write_price IS NULL OR cache_write_price >= 0), + reasoning_price NUMERIC(30, 18) CHECK (reasoning_price IS NULL OR reasoning_price >= 0), + request_price NUMERIC(30, 18) CHECK (request_price IS NULL OR request_price >= 0), + image_price NUMERIC(30, 18) CHECK (image_price IS NULL OR image_price >= 0), + web_search_price NUMERIC(30, 18) CHECK (web_search_price IS NULL OR web_search_price >= 0), + synced_at TIMESTAMPTZ NOT NULL DEFAULT now(), + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_usage_pricing_catalog_synced_at + ON usage_pricing_catalog (synced_at DESC); + +CREATE TABLE IF NOT EXISTS usage_pricing_overrides ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, + provider_id UUID NOT NULL REFERENCES llm_providers(id) ON DELETE CASCADE, + provider_type TEXT NOT NULL, + model_id TEXT NOT NULL, + input_price NUMERIC(30, 18) CHECK (input_price IS NULL OR input_price >= 0), + output_price NUMERIC(30, 18) CHECK (output_price IS NULL OR output_price >= 0), + cache_read_price NUMERIC(30, 18) CHECK (cache_read_price IS NULL OR cache_read_price >= 0), + cache_write_price NUMERIC(30, 18) CHECK (cache_write_price IS NULL OR cache_write_price >= 0), + reasoning_price NUMERIC(30, 18) CHECK (reasoning_price IS NULL OR reasoning_price >= 0), + request_price NUMERIC(30, 18) CHECK (request_price IS NULL OR request_price >= 0), + image_price NUMERIC(30, 18) CHECK (image_price IS NULL OR image_price >= 0), + web_search_price NUMERIC(30, 18) CHECK (web_search_price IS NULL OR web_search_price >= 0), + enabled BOOLEAN NOT NULL DEFAULT true, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (tenant_id, provider_id, model_id) +); + +CREATE INDEX IF NOT EXISTS idx_usage_pricing_overrides_tenant_provider + ON usage_pricing_overrides (tenant_id, provider_id, model_id) + WHERE enabled; diff --git a/migrations/000070_usage_cap_policies.down.sql b/migrations/000070_usage_cap_policies.down.sql new file mode 100644 index 00000000..49e817d5 --- /dev/null +++ b/migrations/000070_usage_cap_policies.down.sql @@ -0,0 +1,4 @@ +DROP TABLE IF EXISTS usage_cap_events; +DROP TABLE IF EXISTS usage_cap_reservations; +DROP TABLE IF EXISTS usage_cap_counters; +DROP TABLE IF EXISTS usage_cap_policies; diff --git a/migrations/000070_usage_cap_policies.up.sql b/migrations/000070_usage_cap_policies.up.sql new file mode 100644 index 00000000..5a101d89 --- /dev/null +++ b/migrations/000070_usage_cap_policies.up.sql @@ -0,0 +1,68 @@ +CREATE TABLE IF NOT EXISTS usage_cap_policies ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, + agent_id UUID REFERENCES agents(id) ON DELETE CASCADE, + provider_id UUID REFERENCES llm_providers(id) ON DELETE CASCADE, + provider_type TEXT, + model_id TEXT, + window TEXT NOT NULL CHECK (window IN ('hour', 'day', 'week', 'month')), + max_tokens BIGINT CHECK (max_tokens IS NULL OR max_tokens >= 0), + max_cost_micros BIGINT CHECK (max_cost_micros IS NULL OR max_cost_micros >= 0), + enabled BOOLEAN NOT NULL DEFAULT true, + priority INTEGER NOT NULL DEFAULT 100, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CHECK (max_tokens IS NOT NULL OR max_cost_micros IS NOT NULL) +); + +CREATE INDEX IF NOT EXISTS idx_usage_cap_policies_scope + ON usage_cap_policies (tenant_id, enabled, agent_id, provider_id, provider_type, model_id); + +CREATE TABLE IF NOT EXISTS usage_cap_counters ( + policy_id UUID NOT NULL REFERENCES usage_cap_policies(id) ON DELETE CASCADE, + window_start TIMESTAMPTZ NOT NULL, + window_end TIMESTAMPTZ NOT NULL, + used_tokens BIGINT NOT NULL DEFAULT 0, + reserved_tokens BIGINT NOT NULL DEFAULT 0, + used_cost_micros BIGINT NOT NULL DEFAULT 0, + reserved_cost_micros BIGINT NOT NULL DEFAULT 0, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (policy_id, window_start) +); + +CREATE TABLE IF NOT EXISTS usage_cap_reservations ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + reservation_key TEXT NOT NULL, + policy_id UUID NOT NULL REFERENCES usage_cap_policies(id) ON DELETE CASCADE, + window_start TIMESTAMPTZ NOT NULL, + reserved_tokens BIGINT NOT NULL DEFAULT 0, + reserved_cost_micros BIGINT NOT NULL DEFAULT 0, + actual_tokens BIGINT NOT NULL DEFAULT 0, + actual_cost_micros BIGINT NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'reserved', + metadata JSONB NOT NULL DEFAULT '{}'::jsonb, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (reservation_key, policy_id) +); + +CREATE INDEX IF NOT EXISTS idx_usage_cap_reservations_key + ON usage_cap_reservations (reservation_key); + +CREATE TABLE IF NOT EXISTS usage_cap_events ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, + policy_id UUID REFERENCES usage_cap_policies(id) ON DELETE SET NULL, + reservation_key TEXT, + decision TEXT NOT NULL, + reason TEXT, + estimated_tokens BIGINT NOT NULL DEFAULT 0, + estimated_cost_micros BIGINT NOT NULL DEFAULT 0, + actual_tokens BIGINT NOT NULL DEFAULT 0, + actual_cost_micros BIGINT NOT NULL DEFAULT 0, + metadata JSONB NOT NULL DEFAULT '{}'::jsonb, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_usage_cap_events_tenant_created + ON usage_cap_events (tenant_id, created_at DESC); diff --git a/ui/web/src/i18n/locales/en/providers.json b/ui/web/src/i18n/locales/en/providers.json index fb4fbbca..e9c38633 100644 --- a/ui/web/src/i18n/locales/en/providers.json +++ b/ui/web/src/i18n/locales/en/providers.json @@ -258,6 +258,38 @@ "verifyFailed": "Embedding verification failed", "dimensionsMismatch": "{{count}} dimensions — does not match required 1536" }, + "pricing": { + "title": "Model Pricing", + "description": "Sync OpenRouter prices and override this provider's model costs for usage caps.", + "skippedDescription": "Usage caps skip subscription and local providers, but prices can still be prepared for later API-based routing.", + "sync": "Sync OpenRouter", + "model": "Model", + "saveOverride": "Save Override", + "catalogMatches": "OpenRouter catalog matches", + "source": "Source", + "configuredFields": "Configured fields", + "override": "Override", + "emptyOverrides": "No pricing overrides for this provider.", + "deleteOverride": "Delete override", + "fields": { + "input": "Input", + "output": "Output", + "cache_read": "Cache read", + "cache_write": "Cache write", + "reasoning": "Reasoning", + "request": "Request", + "image": "Image", + "web_search": "Web search" + }, + "toast": { + "synced": "Synced {{count}} OpenRouter models", + "syncFailed": "Could not sync OpenRouter pricing", + "saved": "Pricing override saved", + "saveFailed": "Could not save pricing override", + "deleted": "Pricing override deleted", + "deleteFailed": "Could not delete pricing override" + } + }, "reasoning": { "off": "Off", "offDesc": "No default extended thinking", diff --git a/ui/web/src/i18n/locales/en/usage.json b/ui/web/src/i18n/locales/en/usage.json index 740e92e2..cc82afce 100644 --- a/ui/web/src/i18n/locales/en/usage.json +++ b/ui/web/src/i18n/locales/en/usage.json @@ -19,7 +19,45 @@ "total": "Total", "channel": "Channel", "cost": "Cost", - "status": "Status" + "status": "Status", + "actions": "Actions" + }, + "caps": { + "title": "Usage Caps", + "description": "Limit AI budget by tokens and cost across tenant, agent, provider, and model scopes.", + "window": "Window", + "windows": { + "hour": "Hourly", + "day": "Daily", + "week": "Weekly", + "month": "Monthly" + }, + "agent": "Agent", + "provider": "Provider", + "model": "Model", + "allAgents": "All Agents", + "allProviders": "All Providers", + "maxTokens": "Max tokens", + "maxCost": "Max cost USD", + "create": "Create cap", + "scope": "Scope", + "tokens": "Tokens", + "cost": "Cost", + "empty": "No cap policies yet.", + "tenantScope": "Tenant cap", + "tenantScoped": "Tenant scoped", + "agentScoped": "Agent scoped", + "enabled": "Enabled", + "disabled": "Disabled", + "delete": "Delete cap", + "recentBlocks": "Recent blocks", + "blocked": "Blocked", + "toast": { + "created": "Usage cap created", + "createFailed": "Could not create usage cap", + "deleted": "Usage cap deleted", + "deleteFailed": "Could not delete usage cap" + } }, "analytics": { "title": "Usage Analytics", diff --git a/ui/web/src/i18n/locales/vi/providers.json b/ui/web/src/i18n/locales/vi/providers.json index a50b7edf..dc898ade 100644 --- a/ui/web/src/i18n/locales/vi/providers.json +++ b/ui/web/src/i18n/locales/vi/providers.json @@ -293,6 +293,38 @@ "verifyFailed": "Kiểm tra embedding thất bại", "dimensionsMismatch": "{{count}} chiều — không khớp yêu cầu 1536" }, + "pricing": { + "title": "Giá Model", + "description": "Đồng bộ giá OpenRouter và ghi đè chi phí model của provider này cho usage cap.", + "skippedDescription": "Usage cap bỏ qua provider subscription và local, nhưng vẫn có thể chuẩn bị giá cho route API-based sau này.", + "sync": "Đồng bộ OpenRouter", + "model": "Model", + "saveOverride": "Lưu override", + "catalogMatches": "Kết quả trong catalog OpenRouter", + "source": "Nguồn", + "configuredFields": "Trường đã cấu hình", + "override": "Override", + "emptyOverrides": "Provider này chưa có pricing override.", + "deleteOverride": "Xóa override", + "fields": { + "input": "Input", + "output": "Output", + "cache_read": "Cache read", + "cache_write": "Cache write", + "reasoning": "Reasoning", + "request": "Request", + "image": "Image", + "web_search": "Web search" + }, + "toast": { + "synced": "Đã đồng bộ {{count}} model OpenRouter", + "syncFailed": "Không thể đồng bộ giá OpenRouter", + "saved": "Đã lưu pricing override", + "saveFailed": "Không thể lưu pricing override", + "deleted": "Đã xóa pricing override", + "deleteFailed": "Không thể xóa pricing override" + } + }, "reasoning": { "off": "Off", "offDesc": "No default extended thinking", diff --git a/ui/web/src/i18n/locales/vi/usage.json b/ui/web/src/i18n/locales/vi/usage.json index 22aefdb2..a7beab3f 100644 --- a/ui/web/src/i18n/locales/vi/usage.json +++ b/ui/web/src/i18n/locales/vi/usage.json @@ -19,7 +19,45 @@ "total": "Tổng", "channel": "Kênh", "cost": "Chi phí", - "status": "Trạng thái" + "status": "Trạng thái", + "actions": "Thao tác" + }, + "caps": { + "title": "Usage Cap", + "description": "Giới hạn AI budget theo token và chi phí ở cấp tenant, agent, provider và model.", + "window": "Chu kỳ", + "windows": { + "hour": "Theo giờ", + "day": "Theo ngày", + "week": "Theo tuần", + "month": "Theo tháng" + }, + "agent": "Agent", + "provider": "Provider", + "model": "Model", + "allAgents": "Tất cả Agent", + "allProviders": "Tất cả Provider", + "maxTokens": "Token tối đa", + "maxCost": "Chi phí tối đa USD", + "create": "Tạo cap", + "scope": "Phạm vi", + "tokens": "Token", + "cost": "Chi phí", + "empty": "Chưa có policy cap.", + "tenantScope": "Cap tenant", + "tenantScoped": "Phạm vi tenant", + "agentScoped": "Phạm vi agent", + "enabled": "Đã bật", + "disabled": "Đã tắt", + "delete": "Xóa cap", + "recentBlocks": "Lần chặn gần đây", + "blocked": "Đã chặn", + "toast": { + "created": "Đã tạo usage cap", + "createFailed": "Không thể tạo usage cap", + "deleted": "Đã xóa usage cap", + "deleteFailed": "Không thể xóa usage cap" + } }, "analytics": { "title": "Phân tích sử dụng", diff --git a/ui/web/src/i18n/locales/zh/providers.json b/ui/web/src/i18n/locales/zh/providers.json index 282dcb05..c4f99013 100644 --- a/ui/web/src/i18n/locales/zh/providers.json +++ b/ui/web/src/i18n/locales/zh/providers.json @@ -311,6 +311,38 @@ "verifyFailed": "Embedding验证失败", "dimensionsMismatch": "{{count}}维 — 不符合要求的1536" }, + "pricing": { + "title": "模型定价", + "description": "同步 OpenRouter 价格，并为此 Provider 的模型费用设置覆盖以供用量上限使用。", + "skippedDescription": "用量上限会跳过订阅和本地 Provider，但仍可为之后的 API 路由预先配置价格。", + "sync": "同步 OpenRouter", + "model": "模型", + "saveOverride": "保存覆盖", + "catalogMatches": "OpenRouter 目录匹配", + "source": "来源", + "configuredFields": "已配置字段", + "override": "覆盖", + "emptyOverrides": "此 Provider 暂无定价覆盖。", + "deleteOverride": "删除覆盖", + "fields": { + "input": "输入", + "output": "输出", + "cache_read": "缓存读取", + "cache_write": "缓存写入", + "reasoning": "推理", + "request": "请求", + "image": "图片", + "web_search": "网页搜索" + }, + "toast": { + "synced": "已同步 {{count}} 个 OpenRouter 模型", + "syncFailed": "无法同步 OpenRouter 定价", + "saved": "定价覆盖已保存", + "saveFailed": "无法保存定价覆盖", + "deleted": "定价覆盖已删除", + "deleteFailed": "无法删除定价覆盖" + } + }, "reasoning": { "off": "Off", "offDesc": "No default extended thinking", diff --git a/ui/web/src/i18n/locales/zh/usage.json b/ui/web/src/i18n/locales/zh/usage.json index 8dd29a7d..1810eb43 100644 --- a/ui/web/src/i18n/locales/zh/usage.json +++ b/ui/web/src/i18n/locales/zh/usage.json @@ -19,7 +19,45 @@ "total": "合计", "channel": "渠道", "cost": "费用", - "status": "状态" + "status": "状态", + "actions": "操作" + }, + "caps": { + "title": "用量上限", + "description": "按租户、Agent、Provider 和模型范围限制 AI 预算的令牌与费用。", + "window": "窗口", + "windows": { + "hour": "每小时", + "day": "每天", + "week": "每周", + "month": "每月" + }, + "agent": "Agent", + "provider": "Provider", + "model": "模型", + "allAgents": "全部Agent", + "allProviders": "全部Provider", + "maxTokens": "最大令牌", + "maxCost": "最大费用 USD", + "create": "创建上限", + "scope": "范围", + "tokens": "令牌", + "cost": "费用", + "empty": "暂无上限策略。", + "tenantScope": "租户上限", + "tenantScoped": "租户范围", + "agentScoped": "Agent范围", + "enabled": "已启用", + "disabled": "已禁用", + "delete": "删除上限", + "recentBlocks": "近期拦截", + "blocked": "已拦截", + "toast": { + "created": "用量上限已创建", + "createFailed": "无法创建用量上限", + "deleted": "用量上限已删除", + "deleteFailed": "无法删除用量上限" + } }, "analytics": { "title": "用量分析", diff --git a/ui/web/src/lib/query-keys.ts b/ui/web/src/lib/query-keys.ts index e58a2f88..3ce184e1 100644 --- a/ui/web/src/lib/query-keys.ts +++ b/ui/web/src/lib/query-keys.ts @@ -5,6 +5,8 @@ export const queryKeys = { providers: { all: ["providers"] as const, models: (providerId: string) => ["providers", providerId, "models"] as const, + pricing: (providerId: string) => ["providers", providerId, "pricing"] as const, + pricingCatalog: (model: string) => ["providers", "pricing-catalog", model] as const, chatgptOAuthStatuses: (providerKeys: string[]) => ["providers", "chatgpt-oauth-statuses", ...providerKeys] as const, chatgptOAuthQuotas: (providerNames: string[]) => ["providers", "chatgpt-oauth-quotas", ...providerNames] as const, codexPoolActivity: (providerId: string, limit: number) => ["providers", providerId, "codex-pool-activity", limit] as const, @@ -69,6 +71,11 @@ export const queryKeys = { usage: { all: ["usage"] as const, records: (params: Record) => ["usage", "records", params] as const, + caps: { + policies: ["usage", "caps", "policies"] as const, + utilization: ["usage", "caps", "utilization"] as const, + events: ["usage", "caps", "events"] as const, + }, }, teams: { all: ["teams"] as const, diff --git a/ui/web/src/pages/providers/hooks/use-model-pricing.ts b/ui/web/src/pages/providers/hooks/use-model-pricing.ts new file mode 100644 index 00000000..b8d72365 --- /dev/null +++ b/ui/web/src/pages/providers/hooks/use-model-pricing.ts @@ -0,0 +1,93 @@ +import { useCallback } from "react"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; +import i18next from "i18next"; +import { useHttp } from "@/hooks/use-ws"; +import { queryKeys } from "@/lib/query-keys"; +import { toast } from "@/stores/use-toast-store"; +import type { PricingCatalogEntry, PricingOverride, UsagePricingFields } from "@/types/usage-caps"; + +export interface PricingOverrideInput { + provider_id: string; + provider_type: string; + model_id: string; + pricing: UsagePricingFields; + enabled?: boolean; +} + +export function useModelPricing(providerId: string, modelSearch: string) { + const http = useHttp(); + const queryClient = useQueryClient(); + const pricingKey = queryKeys.providers.pricing(providerId); + + const catalogQuery = useQuery({ + queryKey: queryKeys.providers.pricingCatalog(modelSearch), + queryFn: async () => { + const res = await http.get<{ models: PricingCatalogEntry[] }>("/v1/model-pricing", { model: modelSearch, limit: "8" }); + return res.models ?? []; + }, + }); + + const overridesQuery = useQuery({ + queryKey: pricingKey, + queryFn: async () => { + const res = await http.get<{ overrides: PricingOverride[] }>("/v1/model-pricing/overrides", { provider_id: providerId }); + return res.overrides ?? []; + }, + }); + + const invalidate = useCallback(async () => { + await Promise.all([ + queryClient.invalidateQueries({ queryKey: pricingKey }), + queryClient.invalidateQueries({ queryKey: queryKeys.providers.pricingCatalog(modelSearch) }), + ]); + }, [modelSearch, pricingKey, queryClient]); + + const syncOpenRouter = useCallback(async () => { + try { + const res = await http.post<{ count: number }>("/v1/model-pricing/sync-openrouter"); + await invalidate(); + toast.success(i18next.t("providers:pricing.toast.synced", { count: res.count })); + } catch (err) { + toast.error(i18next.t("providers:pricing.toast.syncFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, [http, invalidate]); + + const saveOverride = useCallback( + async (input: PricingOverrideInput) => { + try { + await http.put("/v1/model-pricing/overrides", input); + await invalidate(); + toast.success(i18next.t("providers:pricing.toast.saved")); + } catch (err) { + toast.error(i18next.t("providers:pricing.toast.saveFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, invalidate], + ); + + const deleteOverride = useCallback( + async (id: string) => { + try { + await http.delete(`/v1/model-pricing/overrides/${id}`); + await invalidate(); + toast.success(i18next.t("providers:pricing.toast.deleted")); + } catch (err) { + toast.error(i18next.t("providers:pricing.toast.deleteFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, invalidate], + ); + + return { + catalog: catalogQuery.data ?? [], + overrides: overridesQuery.data ?? [], + loading: catalogQuery.isLoading || overridesQuery.isLoading, + refreshing: catalogQuery.isFetching || overridesQuery.isFetching, + syncOpenRouter, + saveOverride, + deleteOverride, + }; +} diff --git a/ui/web/src/pages/providers/provider-detail/provider-overview.tsx b/ui/web/src/pages/providers/provider-detail/provider-overview.tsx index ddf1ebac..ca482a08 100644 --- a/ui/web/src/pages/providers/provider-detail/provider-overview.tsx +++ b/ui/web/src/pages/providers/provider-detail/provider-overview.tsx @@ -21,6 +21,7 @@ import { ProviderOAuthAccountSection } from "./provider-oauth-account-section"; import { ProviderReasoningSection } from "./provider-reasoning-section"; import { ProviderEmbeddingSection } from "./provider-embedding-section"; import { ProviderPoolActivitySection } from "./provider-pool-activity-section"; +import { ProviderPricingSection } from "./provider-pricing-section"; import { buildProviderSettingsWithChatGPTOAuthRouting, buildProviderSettingsWithReasoningDefaults, @@ -285,6 +286,8 @@ export function ProviderOverview({ provider, onUpdate }: ProviderOverviewProps) ) : null} + + {showEmbedding ? ( verifyEmbedding(provider.id, embModel.trim() || undefined, undefined)} verifying={embVerifying} verifyResult={embResult} /> ) : null} diff --git a/ui/web/src/pages/providers/provider-detail/provider-pricing-section.tsx b/ui/web/src/pages/providers/provider-detail/provider-pricing-section.tsx new file mode 100644 index 00000000..6594821d --- /dev/null +++ b/ui/web/src/pages/providers/provider-detail/provider-pricing-section.tsx @@ -0,0 +1,158 @@ +import { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; +import { DatabaseZap, RefreshCw, Save, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Badge } from "@/components/ui/badge"; +import { useAuthStore } from "@/stores/use-auth-store"; +import type { ProviderData } from "@/types/provider"; +import type { UsagePricingFields } from "@/types/usage-caps"; +import { useModelPricing } from "../hooks/use-model-pricing"; + +const PRICE_FIELDS: Array = [ + "input", "output", "cache_read", "cache_write", "reasoning", "request", "image", "web_search", +]; +const SUBSCRIPTION_TYPES = new Set(["chatgpt_oauth", "claude_cli", "bailian", "acp", "ollama"]); + +export function ProviderPricingSection({ provider }: { provider: ProviderData }) { + const { t } = useTranslation("providers"); + const [modelId, setModelId] = useState(""); + const [prices, setPrices] = useState>({}); + const [saving, setSaving] = useState(false); + const [syncing, setSyncing] = useState(false); + const isMasterScope = useAuthStore((s) => s.isMasterScope); + const { catalog, overrides, refreshing, syncOpenRouter, saveOverride, deleteOverride } = useModelPricing(provider.id, modelId.trim()); + const isSkipped = SUBSCRIPTION_TYPES.has(provider.provider_type); + + const pricing = useMemo(() => { + const out: UsagePricingFields = {}; + for (const field of PRICE_FIELDS) { + const value = prices[field]?.trim(); + if (value) out[field] = value; + } + return out; + }, [prices]); + + const onSave = async () => { + if (!modelId.trim() || Object.keys(pricing).length === 0) return; + setSaving(true); + try { + await saveOverride({ + provider_id: provider.id, + provider_type: provider.provider_type, + model_id: modelId.trim(), + pricing, + enabled: true, + }); + } finally { + setSaving(false); + } + }; + + const onSync = async () => { + setSyncing(true); + try { + await syncOpenRouter(); + } finally { + setSyncing(false); + } + }; + + return ( +
+
+
+

{t("pricing.title")}

+

{isSkipped ? t("pricing.skippedDescription") : t("pricing.description")}

+
+ {isMasterScope ? ( + + ) : null} +
+ +
+
+ + setModelId(e.target.value)} placeholder="openai/gpt-4o-mini" className="text-base md:text-sm" /> +
+
+ {PRICE_FIELDS.map((field) => ( +
+ + setPrices((current) => ({ ...current, [field]: e.target.value }))} + inputMode="decimal" + placeholder="0.00000015" + className="text-base md:text-sm" + /> +
+ ))} +
+ +
+ + {catalog.length > 0 ? ( +
+
+ + {t("pricing.catalogMatches")} +
+
+ {catalog.map((entry) => ( + + ))} +
+
+ ) : null} + +
+ + + + + + + + + + + {overrides.length === 0 ? ( + + ) : overrides.map((override) => ( + + + + + + + ))} + +
{t("pricing.model")}{t("pricing.source")}{t("pricing.configuredFields")}{t("columns.actions")}
{t("pricing.emptyOverrides")}
{override.model_id}{t("pricing.override")}{Object.keys(override.pricing ?? {}).join(", ") || "—"} + +
+
+
+ ); +} + +function cleanPricing(fields: UsagePricingFields): Record { + const out: Record = {}; + for (const field of PRICE_FIELDS) { + const value = fields[field]; + if (value) out[field] = value; + } + return out; +} diff --git a/ui/web/src/pages/usage/components/usage-caps-panel.tsx b/ui/web/src/pages/usage/components/usage-caps-panel.tsx new file mode 100644 index 00000000..6f39de72 --- /dev/null +++ b/ui/web/src/pages/usage/components/usage-caps-panel.tsx @@ -0,0 +1,179 @@ +import { useMemo, useState } from "react"; +import type { ReactNode } from "react"; +import { useTranslation } from "react-i18next"; +import { Plus, RefreshCw, ShieldAlert, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Badge } from "@/components/ui/badge"; +import { formatCost, 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"; +import type { UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps"; + +const ALL = "__all__"; + +export function UsageCapsPanel() { + const { t } = useTranslation("usage"); + const { agents } = useAgents(); + const { providers } = useProviders(); + const { utilization, events, refreshing, refresh, createPolicy, deletePolicy } = useUsageCaps(); + const [windowValue, setWindowValue] = useState("day"); + const [agentId, setAgentId] = useState(ALL); + const [providerId, setProviderId] = useState(ALL); + const [modelId, setModelId] = useState(""); + const [maxTokens, setMaxTokens] = useState(""); + const [maxCost, setMaxCost] = useState(""); + const [saving, setSaving] = useState(false); + + const provider = useMemo(() => providers.find((p) => p.id === providerId), [providerId, providers]); + const blockedEvents = events.filter((event) => event.decision === "block"); + + const onSubmit = async () => { + const tokens = Number(maxTokens); + const cost = Number(maxCost); + if ((!Number.isFinite(tokens) || tokens <= 0) && (!Number.isFinite(cost) || cost <= 0)) return; + setSaving(true); + try { + await createPolicy({ + window: windowValue, + agent_id: agentId === ALL ? undefined : agentId, + provider_id: providerId === ALL ? undefined : providerId, + provider_type: provider?.provider_type, + model_id: modelId.trim() || undefined, + max_tokens: Number.isFinite(tokens) && tokens > 0 ? Math.floor(tokens) : undefined, + max_cost_usd: Number.isFinite(cost) && cost > 0 ? cost : undefined, + enabled: true, + }); + setMaxTokens(""); + setMaxCost(""); + setModelId(""); + } finally { + setSaving(false); + } + }; + + return ( +
+
+
+

{t("caps.title")}

+

{t("caps.description")}

+
+ +
+ +
+ + + + + + + + + + + setModelId(e.target.value)} placeholder="openai/gpt-4o-mini" className="text-base md:text-sm" /> + + + setMaxTokens(e.target.value)} inputMode="numeric" placeholder="500000" className="text-base md:text-sm" /> + + +
+ setMaxCost(e.target.value)} inputMode="decimal" placeholder="25" className="text-base md:text-sm" /> + +
+ +
+ +
+ + + + + + + + + + + + + {utilization.length === 0 ? ( + + ) : utilization.map((row) => ( + void deletePolicy(row.policy.id)} /> + ))} + +
{t("caps.scope")}{t("caps.window")}{t("caps.tokens")}{t("caps.cost")}{t("columns.status")}{t("columns.actions", "Actions")}
{t("caps.empty")}
+
+ + {blockedEvents.length > 0 ? ( +
+
{t("caps.recentBlocks")}
+
+ {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)}
+
+ ))} +
+
+ ) : null} +
+ ); +} + +function Field({ label, children }: { label: string; children: ReactNode }) { + return
{children}
; +} + +function UsageCapRow({ row, onDelete }: { row: UsageCapUtilization; onDelete: () => void }) { + const { t } = useTranslation("usage"); + const p = row.policy; + const tokenUsed = row.used_tokens + row.reserved_tokens; + const costUsed = row.used_cost_micros + row.reserved_cost_micros; + const tokenPct = p.max_tokens ? Math.min(100, Math.round((tokenUsed / p.max_tokens) * 100)) : 0; + const costPct = p.max_cost_micros ? Math.min(100, Math.round((costUsed / p.max_cost_micros) * 100)) : 0; + return ( + + +
{p.model_id || p.provider_type || t("caps.tenantScope")}
+
{p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")}
+ + {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.enabled ? t("caps.enabled") : t("caps.disabled")} + + + + + ); +} diff --git a/ui/web/src/pages/usage/hooks/use-usage-caps.ts b/ui/web/src/pages/usage/hooks/use-usage-caps.ts new file mode 100644 index 00000000..64c0e4c4 --- /dev/null +++ b/ui/web/src/pages/usage/hooks/use-usage-caps.ts @@ -0,0 +1,94 @@ +import { useCallback } from "react"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; +import i18next from "i18next"; +import { useHttp } from "@/hooks/use-ws"; +import { queryKeys } from "@/lib/query-keys"; +import { toast } from "@/stores/use-toast-store"; +import type { UsageCapEvent, UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps"; + +export interface UsageCapPolicyInput { + agent_id?: string; + provider_id?: string; + provider_type?: string; + model_id?: string; + window: UsageCapPolicy["window"]; + max_tokens?: number; + max_cost_usd?: number; + enabled?: boolean; +} + +export function useUsageCaps() { + const http = useHttp(); + const queryClient = useQueryClient(); + + const policiesQuery = useQuery({ + queryKey: queryKeys.usage.caps.policies, + queryFn: async () => { + const res = await http.get<{ policies: UsageCapPolicy[] }>("/v1/usage-caps/policies"); + return res.policies ?? []; + }, + }); + + const utilizationQuery = useQuery({ + queryKey: queryKeys.usage.caps.utilization, + queryFn: async () => { + const res = await http.get<{ rows: UsageCapUtilization[] }>("/v1/usage-caps/utilization"); + return res.rows ?? []; + }, + }); + + const eventsQuery = useQuery({ + queryKey: queryKeys.usage.caps.events, + queryFn: async () => { + const res = await http.get<{ events: UsageCapEvent[] }>("/v1/usage-caps/events", { limit: "10" }); + return res.events ?? []; + }, + }); + + const refresh = useCallback(async () => { + await Promise.all([ + queryClient.invalidateQueries({ queryKey: queryKeys.usage.caps.policies }), + queryClient.invalidateQueries({ queryKey: queryKeys.usage.caps.utilization }), + queryClient.invalidateQueries({ queryKey: queryKeys.usage.caps.events }), + ]); + }, [queryClient]); + + const createPolicy = useCallback( + async (input: UsageCapPolicyInput) => { + try { + await http.post("/v1/usage-caps/policies", input); + await refresh(); + toast.success(i18next.t("usage:caps.toast.created")); + } catch (err) { + toast.error(i18next.t("usage:caps.toast.createFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, refresh], + ); + + const deletePolicy = useCallback( + async (id: string) => { + try { + await http.delete(`/v1/usage-caps/policies/${id}`); + await refresh(); + toast.success(i18next.t("usage:caps.toast.deleted")); + } catch (err) { + toast.error(i18next.t("usage:caps.toast.deleteFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, refresh], + ); + + return { + policies: policiesQuery.data ?? [], + utilization: utilizationQuery.data ?? [], + events: eventsQuery.data ?? [], + loading: policiesQuery.isLoading || utilizationQuery.isLoading || eventsQuery.isLoading, + refreshing: policiesQuery.isFetching || utilizationQuery.isFetching || eventsQuery.isFetching, + refresh, + createPolicy, + deletePolicy, + }; +} diff --git a/ui/web/src/pages/usage/usage-page.tsx b/ui/web/src/pages/usage/usage-page.tsx index c052175b..71a36584 100644 --- a/ui/web/src/pages/usage/usage-page.tsx +++ b/ui/web/src/pages/usage/usage-page.tsx @@ -21,6 +21,7 @@ import { DistributionRow } from "./components/distribution-row"; import { DurationChart } from "./components/duration-chart"; import { KnowledgeChart } from "./components/knowledge-chart"; import { TopModelsTable } from "./components/top-models-table"; +import { UsageCapsPanel } from "./components/usage-caps-panel"; const EMPTY_SUMMARY = { requests: 0, input_tokens: 0, output_tokens: 0, cost: 0, errors: 0, unique_users: 0, llm_calls: 0, tool_calls: 0, avg_duration_ms: 0 }; @@ -108,6 +109,10 @@ function AnalyticsDashboard() { + + + + diff --git a/ui/web/src/types/usage-caps.ts b/ui/web/src/types/usage-caps.ts new file mode 100644 index 00000000..670687be --- /dev/null +++ b/ui/web/src/types/usage-caps.ts @@ -0,0 +1,68 @@ +export interface UsagePricingFields { + input?: string; + output?: string; + cache_read?: string; + cache_write?: string; + reasoning?: string; + request?: string; + image?: string; + web_search?: string; +} + +export interface UsageCapPolicy { + id: string; + tenant_id: string; + agent_id?: string; + provider_id?: string; + provider_type?: string; + model_id?: string; + window: "hour" | "day" | "week" | "month"; + max_tokens?: number; + max_cost_micros?: number; + enabled: boolean; + priority: number; + created_at: string; + updated_at: string; +} + +export interface UsageCapUtilization { + policy: UsageCapPolicy; + window_start: string; + window_end: string; + used_tokens: number; + reserved_tokens: number; + used_cost_micros: number; + reserved_cost_micros: number; +} + +export interface UsageCapEvent { + id: string; + policy_id?: string; + reservation_key?: string; + decision: "allow" | "block" | "reconcile" | "skip"; + reason?: string; + estimated_tokens: number; + estimated_cost_micros: number; + actual_tokens: number; + actual_cost_micros: number; + created_at: string; +} + +export interface PricingCatalogEntry { + id: string; + model_id: string; + canonical_model_id?: string; + pricing: UsagePricingFields; + synced_at: string; +} + +export interface PricingOverride { + id: string; + provider_id: string; + provider_type: string; + model_id: string; + pricing: UsagePricingFields; + enabled: boolean; + created_at: string; + updated_at: string; +} From ade1c1b8386bfbf1b02d695d72958e42617ebb1b Mon Sep 17 00:00:00 2001 From: Goon Date: Sat, 23 May 2026 00:39:14 +0700 Subject: [PATCH 02/10] fix(store): avoid usage cap migration keyword --- docs/06-store-data-model.md | 2 +- internal/store/pg/usage_caps.go | 18 +++++++++--------- internal/store/pg/usage_caps_test.go | 4 ++-- migrations/000070_usage_cap_policies.up.sql | 2 +- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/06-store-data-model.md b/docs/06-store-data-model.md index eb826417..fa8b696a 100644 --- a/docs/06-store-data-model.md +++ b/docs/06-store-data-model.md @@ -93,7 +93,7 @@ Usage cap enforcement is Standard/PostgreSQL-only in round one. The `UsageCapSto Tables: - `usage_pricing_catalog`: OpenRouter model catalog prices, raw upstream model payload, sync time. - `usage_pricing_overrides`: tenant/provider/model override prices for custom billing assumptions. -- `usage_cap_policies`: cap definitions scoped by tenant, agent, provider, provider type, model, and window. +- `usage_cap_policies`: cap definitions scoped by tenant, agent, provider, provider type, model, and `window_key`. - `usage_cap_counters`: current window used and reserved token/cost counters. - `usage_cap_reservations`: preflight reservations keyed by LLM call attempt. - `usage_cap_events`: allow/block/reconcile/skip audit events. diff --git a/internal/store/pg/usage_caps.go b/internal/store/pg/usage_caps.go index 833edb49..4e8f8e45 100644 --- a/internal/store/pg/usage_caps.go +++ b/internal/store/pg/usage_caps.go @@ -22,7 +22,7 @@ func (s *PGUsageCapStore) CreateUsageCapPolicy(ctx context.Context, p *store.Usa } const q = ` INSERT INTO usage_cap_policies ( - id, tenant_id, agent_id, provider_id, provider_type, model_id, window, + id, tenant_id, agent_id, provider_id, provider_type, model_id, window_key, max_tokens, max_cost_micros, enabled, priority ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) RETURNING created_at, updated_at` @@ -116,7 +116,7 @@ func (s *PGUsageCapStore) UpdateUsageCapPolicy(ctx context.Context, tenantID, id } const q = ` UPDATE usage_cap_policies SET agent_id=$3, provider_id=$4, provider_type=$5, - model_id=$6, window=$7, max_tokens=$8, max_cost_micros=$9, + model_id=$6, window_key=$7, max_tokens=$8, max_cost_micros=$9, enabled=$10, priority=$11, updated_at=now() WHERE tenant_id=$1 AND id=$2 RETURNING updated_at` @@ -167,13 +167,13 @@ RETURNING true`, var ok bool err = tx.QueryRowContext(ctx, ` UPDATE usage_cap_counters SET - reserved_tokens = reserved_tokens + $4, - reserved_cost_micros = reserved_cost_micros + $5, - updated_at = now() + reserved_tokens = reserved_tokens + $3, + reserved_cost_micros = reserved_cost_micros + $4, + updated_at = now() WHERE policy_id=$1 AND window_start=$2 - AND ($6::bigint IS NULL OR used_tokens + reserved_tokens + $4 <= $6) - AND ($7::bigint IS NULL OR used_cost_micros + reserved_cost_micros + $5 <= $7) - RETURNING true`, p.ID, start, end, req.EstimatedTokens, req.EstimatedCostMicros, intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros)).Scan(&ok) + AND ($5::bigint IS NULL OR used_tokens + reserved_tokens + $3 <= $5) + AND ($6::bigint IS NULL OR used_cost_micros + reserved_cost_micros + $4 <= $6) + RETURNING true`, p.ID, start, req.EstimatedTokens, req.EstimatedCostMicros, intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros)).Scan(&ok) if err != nil { if !errors.Is(err, sql.ErrNoRows) { return nil, err @@ -339,7 +339,7 @@ SELECT EXISTS ( } const policySelectSQL = `SELECT id, tenant_id, agent_id, provider_id, COALESCE(provider_type,''), COALESCE(model_id,''), - window, max_tokens, max_cost_micros, enabled, priority, created_at, updated_at FROM usage_cap_policies` + window_key, max_tokens, max_cost_micros, enabled, priority, created_at, updated_at FROM usage_cap_policies` func scanPolicy(row scanner) (store.UsageCapPolicy, error) { var p store.UsageCapPolicy diff --git a/internal/store/pg/usage_caps_test.go b/internal/store/pg/usage_caps_test.go index a6702d31..216ad6de 100644 --- a/internal/store/pg/usage_caps_test.go +++ b/internal/store/pg/usage_caps_test.go @@ -100,7 +100,7 @@ func TestPGUsageCapStoreRejectsCrossTenantRefs(t *testing.T) { if _, err := db.Exec( `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) VALUES ($1,$2,$3,'openrouter','sk-test',true)`, - providerID, tenantA, "usage-cap-provider-"+providerID.String(), + providerID, tenantA, "ucp-"+providerID.String(), ); err != nil { t.Fatalf("seed provider: %v", err) } @@ -116,7 +116,7 @@ func TestPGUsageCapStoreRejectsCrossTenantRefs(t *testing.T) { if _, err := db.Exec( `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) VALUES ($1,$2,$3,'openrouter','sk-test',true)`, - masterProviderID, store.MasterTenantID, "usage-cap-master-provider-"+masterProviderID.String(), + masterProviderID, store.MasterTenantID, "ucpm-"+masterProviderID.String(), ); err != nil { t.Fatalf("seed master provider: %v", err) } diff --git a/migrations/000070_usage_cap_policies.up.sql b/migrations/000070_usage_cap_policies.up.sql index 5a101d89..09a8a578 100644 --- a/migrations/000070_usage_cap_policies.up.sql +++ b/migrations/000070_usage_cap_policies.up.sql @@ -5,7 +5,7 @@ CREATE TABLE IF NOT EXISTS usage_cap_policies ( provider_id UUID REFERENCES llm_providers(id) ON DELETE CASCADE, provider_type TEXT, model_id TEXT, - window TEXT NOT NULL CHECK (window IN ('hour', 'day', 'week', 'month')), + window_key TEXT NOT NULL CHECK (window_key IN ('hour', 'day', 'week', 'month')), max_tokens BIGINT CHECK (max_tokens IS NULL OR max_tokens >= 0), max_cost_micros BIGINT CHECK (max_cost_micros IS NULL OR max_cost_micros >= 0), enabled BOOLEAN NOT NULL DEFAULT true, From a9a6463f71287c2d2477fe03919eff55222064b4 Mon Sep 17 00:00:00 2001 From: Goon Date: Sat, 23 May 2026 08:51:04 +0700 Subject: [PATCH 03/10] feat(usage): bridge agent monthly budget caps --- docs/02-providers.md | 2 + docs/06-store-data-model.md | 6 +- docs/18-http-api.md | 2 + docs/project-changelog.md | 1 + internal/gateway/methods/agents_update.go | 8 +- internal/http/usage_caps.go | 9 ++ internal/store/pg/agents.go | 95 ++++++++++++++++++ .../pg/agents_update_null_coerce_test.go | 96 +++++++++++++++++++ internal/store/pg/usage_caps.go | 38 ++++++-- internal/store/usage_caps.go | 9 +- internal/upgrade/version.go | 2 +- ...071_agent_budget_usage_cap_bridge.down.sql | 7 ++ ...00071_agent_budget_usage_cap_bridge.up.sql | 23 +++++ ui/web/src/i18n/locales/en/usage.json | 2 + ui/web/src/i18n/locales/vi/usage.json | 2 + ui/web/src/i18n/locales/zh/usage.json | 2 + .../usage/components/usage-caps-panel.tsx | 8 +- ui/web/src/types/usage-caps.ts | 1 + 18 files changed, 296 insertions(+), 17 deletions(-) create mode 100644 migrations/000071_agent_budget_usage_cap_bridge.down.sql create mode 100644 migrations/000071_agent_budget_usage_cap_bridge.up.sql diff --git a/docs/02-providers.md b/docs/02-providers.md index 003b2e5a..642464e7 100644 --- a/docs/02-providers.md +++ b/docs/02-providers.md @@ -71,6 +71,8 @@ Runtime flow: Token-only policies do not require catalog pricing. Model fallback routes reserve against the actual candidate provider/model before each attempt. Cached input is separated from uncached input for OpenAI-compatible usage accounting. Partial stream failures keep the estimate, or actual provider usage when available, instead of clearing billed output to zero. Internal LLM calls for memory flush, compaction, media reading tools (`read_image`, `read_document`, `read_audio`, `read_video`), and subagents use the same preflight/reconcile path. +The legacy agent-level `budget_monthly_cents` field is treated as a generated monthly agent USD cap. Existing values are backfilled during migration, and later agent budget edits update or remove the generated cap policy. + Supported price units: input, output, cache read, cache write, reasoning, request, image, and web search. --- diff --git a/docs/06-store-data-model.md b/docs/06-store-data-model.md index fa8b696a..efdac94a 100644 --- a/docs/06-store-data-model.md +++ b/docs/06-store-data-model.md @@ -93,7 +93,7 @@ Usage cap enforcement is Standard/PostgreSQL-only in round one. The `UsageCapSto Tables: - `usage_pricing_catalog`: OpenRouter model catalog prices, raw upstream model payload, sync time. - `usage_pricing_overrides`: tenant/provider/model override prices for custom billing assumptions. -- `usage_cap_policies`: cap definitions scoped by tenant, agent, provider, provider type, model, and `window_key`. +- `usage_cap_policies`: cap definitions scoped by tenant, agent, provider, provider type, model, `window_key`, and `source`. - `usage_cap_counters`: current window used and reserved token/cost counters. - `usage_cap_reservations`: preflight reservations keyed by LLM call attempt. - `usage_cap_events`: allow/block/reconcile/skip audit events. @@ -104,9 +104,11 @@ Policy `agent_id` references must belong to the same tenant as the policy. Polic Catalog and override price fields are nullable decimal strings with non-negative validation in the store layer and database checks. Pricing resolution checks exact override/catalog model IDs first, then provider-derived OpenRouter aliases for native unprefixed model IDs. +Agent `budget_monthly_cents` values are bridged into `usage_cap_policies` with `source = 'agent_budget_monthly_cents'`, an agent scope, `window_key = 'month'`, and `max_cost_micros = budget_monthly_cents * 10000`. Updating or clearing the agent budget keeps that generated policy in sync; manual cap policies continue to use `source = 'manual'`. + Migration versions: -- PostgreSQL: `000069_usage_caps_pricing`, `000070_usage_cap_policies`. +- PostgreSQL: `000069_usage_caps_pricing`, `000070_usage_cap_policies`, `000071_agent_budget_usage_cap_bridge`. - SQLite: no schema change; feature is not active in Lite. --- diff --git a/docs/18-http-api.md b/docs/18-http-api.md index f24f4d41..de1069f5 100644 --- a/docs/18-http-api.md +++ b/docs/18-http-api.md @@ -1294,6 +1294,8 @@ Usage cap policy body: } ``` +Policy responses include read-only `source`. `source="agent_budget_monthly_cents"` means the policy is generated from the agent monthly budget field and must be changed there; direct policy updates/deletes return `409 Conflict`. + --- ## 24. Activity & Audit diff --git a/docs/project-changelog.md b/docs/project-changelog.md index f6b68bdd..b965c81a 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -16,6 +16,7 @@ Significant changes, features, and fixes in reverse chronological order. - Enforced caps in agent, fallback model, subagent, memory flush, compaction, and media reading tools (`read_image`, `read_document`, `read_audio`, `read_video`) with preflight reservation and post-call reconciliation. - Added non-negative validation for catalog and override pricing fields. - Added OpenRouter alias resolution for native model IDs, cached-input accounting normalization, and partial-stream failure reconciliation. +- Bridged legacy `budget_monthly_cents` into generated monthly agent USD cap policies, including migration backfill and save-time sync. - Added web dashboard controls on Usage and Provider detail pages. **Tests** diff --git a/internal/gateway/methods/agents_update.go b/internal/gateway/methods/agents_update.go index e1f96651..d0d76d69 100644 --- a/internal/gateway/methods/agents_update.go +++ b/internal/gateway/methods/agents_update.go @@ -60,6 +60,10 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } + var rawParams map[string]json.RawMessage + if req.Params != nil { + _ = json.Unmarshal(req.Params, &rawParams) + } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) @@ -104,7 +108,9 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client if params.IsDefault != nil { updates["is_default"] = *params.IsDefault } - if params.BudgetCents != nil { + if rawBudget, ok := rawParams["budget_monthly_cents"]; ok && strings.TrimSpace(string(rawBudget)) == "null" { + updates["budget_monthly_cents"] = nil + } else if params.BudgetCents != nil { updates["budget_monthly_cents"] = *params.BudgetCents } // Per-agent JSONB config overrides diff --git a/internal/http/usage_caps.go b/internal/http/usage_caps.go index 7e027528..1fc64c8f 100644 --- a/internal/http/usage_caps.go +++ b/internal/http/usage_caps.go @@ -3,6 +3,7 @@ package http import ( "context" "encoding/json" + "errors" "fmt" "log/slog" "net/http" @@ -108,6 +109,10 @@ func (h *UsageCapsHandler) handleUpdatePolicy(w http.ResponseWriter, r *http.Req } p, err := h.store.UpdateUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id, patch) if err != nil { + if errors.Is(err, store.ErrUsageCapPolicyManaged) { + writeJSON(w, http.StatusConflict, map[string]string{"error": err.Error()}) + return + } slog.Warn("usage_caps.update_policy_failed", "error", err) writeJSON(w, http.StatusBadRequest, map[string]string{"error": "usage cap policy validation failed"}) return @@ -122,6 +127,10 @@ func (h *UsageCapsHandler) handleDeletePolicy(w http.ResponseWriter, r *http.Req return } if err := h.store.DeleteUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id); err != nil { + if errors.Is(err, store.ErrUsageCapPolicyManaged) { + writeJSON(w, http.StatusConflict, map[string]string{"error": err.Error()}) + return + } writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "delete failed"}) return } diff --git a/internal/store/pg/agents.go b/internal/store/pg/agents.go index 234cf1a6..e5a19437 100644 --- a/internal/store/pg/agents.go +++ b/internal/store/pg/agents.go @@ -133,6 +133,12 @@ func (s *PGAgentStore) Create(ctx context.Context, agent *store.AgentData) error if err != nil { return err } + if agent.BudgetMonthlyCents != nil { + if err := s.syncAgentBudgetUsageCap(ctx, tenantID, agent.ID, agent.BudgetMonthlyCents); err != nil { + _, _ = s.db.ExecContext(ctx, "DELETE FROM agents WHERE id = $1", agent.ID) + return err + } + } // Generate embedding for new agent with frontmatter if agent.Frontmatter != "" && s.embProvider != nil { @@ -189,6 +195,22 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str if len(updates) == 0 { return nil } + var err error + var budgetCents *int + rawBudget, syncBudget := updates["budget_monthly_cents"] + if syncBudget { + budgetCents, err = budgetCentsFromUpdate(rawBudget) + if err != nil { + return err + } + } + var budgetTenantID uuid.UUID + if syncBudget { + budgetTenantID, err = s.agentTenantID(ctx, id) + if err != nil { + return err + } + } // Coerce NOT NULL columns: null → default to prevent constraint violations. // Promoted TEXT columns (migration 000037): null → empty string. @@ -248,6 +270,11 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str return err } } + if syncBudget { + if err := s.syncAgentBudgetUsageCap(ctx, budgetTenantID, id, budgetCents); err != nil { + return err + } + } // Regenerate embedding when frontmatter changes if _, hasFrontmatter := updates["frontmatter"]; hasFrontmatter && s.embProvider != nil { @@ -262,6 +289,74 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str return nil } +func (s *PGAgentStore) agentTenantID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { + if !store.IsCrossTenant(ctx) { + tid := store.TenantIDFromContext(ctx) + if tid != uuid.Nil { + return tid, nil + } + } + var tenantID uuid.UUID + if err := s.db.QueryRowContext(ctx, "SELECT tenant_id FROM agents WHERE id = $1 AND deleted_at IS NULL", id).Scan(&tenantID); err != nil { + return uuid.Nil, fmt.Errorf("agent not found: %s", id) + } + return tenantID, nil +} + +func budgetCentsFromUpdate(v any) (*int, error) { + if v == nil { + return nil, nil + } + var cents int + switch n := v.(type) { + case int: + cents = n + case int32: + cents = int(n) + case int64: + cents = int(n) + case float64: + if n != float64(int(n)) { + return nil, fmt.Errorf("budget_monthly_cents must be an integer") + } + cents = int(n) + default: + return nil, fmt.Errorf("budget_monthly_cents must be an integer") + } + if cents < 0 { + return nil, fmt.Errorf("budget_monthly_cents must be non-negative") + } + return ¢s, nil +} + +func (s *PGAgentStore) syncAgentBudgetUsageCap(ctx context.Context, tenantID, agentID uuid.UUID, budgetCents *int) error { + if budgetCents == nil || *budgetCents <= 0 { + _, err := s.db.ExecContext(ctx, + `DELETE FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget) + return err + } + costMicros := int64(*budgetCents) * 10000 + _, err := s.db.ExecContext(ctx, ` +INSERT INTO usage_cap_policies ( + tenant_id, agent_id, window_key, max_cost_micros, enabled, priority, source +) VALUES ($1,$2,'month',$3,true,90,$4) +ON CONFLICT (tenant_id, agent_id) WHERE source = 'agent_budget_monthly_cents' +DO UPDATE SET + window_key='month', + provider_id=NULL, + provider_type=NULL, + model_id=NULL, + max_tokens=NULL, + max_cost_micros=EXCLUDED.max_cost_micros, + enabled=true, + priority=EXCLUDED.priority, + updated_at=now()`, + tenantID, agentID, costMicros, store.UsageCapSourceAgentBudget) + return err +} + func isEmptyOrNullJSONUpdate(v any) bool { if v == nil { return true diff --git a/internal/store/pg/agents_update_null_coerce_test.go b/internal/store/pg/agents_update_null_coerce_test.go index 3148b68e..5782d545 100644 --- a/internal/store/pg/agents_update_null_coerce_test.go +++ b/internal/store/pg/agents_update_null_coerce_test.go @@ -1,8 +1,14 @@ package pg import ( + "context" + "database/sql" "encoding/json" + "errors" "testing" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" ) func TestIsEmptyOrNullJSONUpdate(t *testing.T) { @@ -30,3 +36,93 @@ func TestIsEmptyOrNullJSONUpdate(t *testing.T) { }) } } + +func TestPGAgentStoreSyncsMonthlyBudgetUsageCap(t *testing.T) { + db := hooksTestDB(t) + tenantID, _ := seedTenantAndAgent(t, db) + ctx := tenantScopedCtx(tenantID) + agentStore := NewPGAgentStore(db) + budgetCents := 123 + agentID := uuid.New() + agent := &store.AgentData{ + BaseModel: store.BaseModel{ID: agentID}, + TenantID: tenantID, AgentKey: "budget-agent-" + agentID.String(), + OwnerID: "owner", Provider: "openai", Model: "gpt-test", + AgentType: store.AgentTypePredefined, Status: store.AgentStatusActive, + BudgetMonthlyCents: &budgetCents, + } + if err := agentStore.Create(ctx, agent); err != nil { + t.Fatalf("Create: %v", err) + } + assertAgentBudgetPolicy(t, db, tenantID, agentID, 1230000) + var generatedPolicyID uuid.UUID + if err := db.QueryRowContext(context.Background(), + `SELECT id FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&generatedPolicyID); err != nil { + t.Fatalf("query generated policy id: %v", err) + } + if err := NewPGUsageCapStore(db).DeleteUsageCapPolicy(ctx, tenantID, generatedPolicyID); err == nil { + t.Fatal("DeleteUsageCapPolicy deleted generated agent budget policy") + } else if !errors.Is(err, store.ErrUsageCapPolicyManaged) { + t.Fatalf("DeleteUsageCapPolicy error = %v, want managed policy error", err) + } + disablePolicy := false + if _, err := NewPGUsageCapStore(db).UpdateUsageCapPolicy(ctx, tenantID, generatedPolicyID, store.UsageCapPolicyPatch{ + Enabled: &disablePolicy, + }); err == nil { + t.Fatal("UpdateUsageCapPolicy updated generated agent budget policy") + } else if !errors.Is(err, store.ErrUsageCapPolicyManaged) { + t.Fatalf("UpdateUsageCapPolicy error = %v, want managed policy error", err) + } + assertAgentBudgetPolicy(t, db, tenantID, agentID, 1230000) + + if err := agentStore.Update(ctx, agentID, map[string]any{"budget_monthly_cents": float64(456)}); err != nil { + t.Fatalf("Update budget: %v", err) + } + assertAgentBudgetPolicy(t, db, tenantID, agentID, 4560000) + + if err := agentStore.Update(ctx, agentID, map[string]any{"budget_monthly_cents": nil}); err != nil { + t.Fatalf("Clear budget: %v", err) + } + var count int + if err := db.QueryRowContext(context.Background(), + `SELECT COUNT(*) FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&count); err != nil { + t.Fatalf("count cleared policy: %v", err) + } + if count != 0 { + t.Fatalf("agent budget policy count after clear = %d, want 0", count) + } +} + +func assertAgentBudgetPolicy(t *testing.T, db rowQueryer, tenantID, agentID uuid.UUID, wantCostMicros int64) { + t.Helper() + var windowKey, source string + var costMicros int64 + var count int + if err := db.QueryRowContext(context.Background(), + `SELECT COUNT(*), COALESCE(MAX(window_key), ''), COALESCE(MAX(source), ''), COALESCE(MAX(max_cost_micros), 0) + FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&count, &windowKey, &source, &costMicros); err != nil { + t.Fatalf("query agent budget policy: %v", err) + } + if count != 1 { + t.Fatalf("agent budget policy count = %d, want 1", count) + } + if windowKey != store.UsageCapWindowMonth { + t.Fatalf("window_key = %q, want month", windowKey) + } + if source != store.UsageCapSourceAgentBudget { + t.Fatalf("source = %q, want %q", source, store.UsageCapSourceAgentBudget) + } + if costMicros != wantCostMicros { + t.Fatalf("max_cost_micros = %d, want %d", costMicros, wantCostMicros) + } +} + +type rowQueryer interface { + QueryRowContext(context.Context, string, ...any) *sql.Row +} diff --git a/internal/store/pg/usage_caps.go b/internal/store/pg/usage_caps.go index 4e8f8e45..06bdeaa4 100644 --- a/internal/store/pg/usage_caps.go +++ b/internal/store/pg/usage_caps.go @@ -20,16 +20,19 @@ func (s *PGUsageCapStore) CreateUsageCapPolicy(ctx context.Context, p *store.Usa if err := s.validateUsageCapRefs(ctx, p.TenantID, p.AgentID, p.ProviderID); err != nil { return err } + if p.Source == "" { + p.Source = store.UsageCapSourceManual + } const q = ` -INSERT INTO usage_cap_policies ( - id, tenant_id, agent_id, provider_id, provider_type, model_id, window_key, - max_tokens, max_cost_micros, enabled, priority -) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) -RETURNING created_at, updated_at` + INSERT INTO usage_cap_policies ( + id, tenant_id, agent_id, provider_id, provider_type, model_id, window_key, + max_tokens, max_cost_micros, source, enabled, priority + ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) + RETURNING created_at, updated_at` return s.db.QueryRowContext(ctx, q, p.ID, p.TenantID, uuidPtrVal(p.AgentID), uuidPtrVal(p.ProviderID), nullEmpty(p.ProviderType), nullEmpty(p.ModelID), p.Window, - intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros), p.Enabled, p.Priority, + intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros), p.Source, p.Enabled, p.Priority, ).Scan(&p.CreatedAt, &p.UpdatedAt) } @@ -84,6 +87,9 @@ func (s *PGUsageCapStore) UpdateUsageCapPolicy(ctx context.Context, tenantID, id if err != nil { return nil, err } + if p.Source == store.UsageCapSourceAgentBudget { + return nil, fmt.Errorf("%w: agent monthly budget", store.ErrUsageCapPolicyManaged) + } if patch.AgentID != nil { p.AgentID = *patch.AgentID } @@ -129,8 +135,20 @@ RETURNING updated_at` } func (s *PGUsageCapStore) DeleteUsageCapPolicy(ctx context.Context, tenantID, id uuid.UUID) error { - _, err := s.db.ExecContext(ctx, `DELETE FROM usage_cap_policies WHERE tenant_id=$1 AND id=$2`, tenantID, id) - return err + res, err := s.db.ExecContext(ctx, + `DELETE FROM usage_cap_policies + WHERE tenant_id=$1 AND id=$2 AND COALESCE(source,'manual') <> $3`, + tenantID, id, store.UsageCapSourceAgentBudget) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + if _, err := s.getPolicy(ctx, tenantID, id); err == nil { + return fmt.Errorf("%w: agent monthly budget", store.ErrUsageCapPolicyManaged) + } + return sql.ErrNoRows + } + return nil } func (s *PGUsageCapStore) ReserveUsage(ctx context.Context, req store.UsageReserveRequest, policies []store.UsageCapPolicy) (*store.UsageReservationResult, error) { @@ -339,14 +357,14 @@ SELECT EXISTS ( } const policySelectSQL = `SELECT id, tenant_id, agent_id, provider_id, COALESCE(provider_type,''), COALESCE(model_id,''), - window_key, max_tokens, max_cost_micros, enabled, priority, created_at, updated_at FROM usage_cap_policies` + window_key, max_tokens, max_cost_micros, COALESCE(source,'manual'), enabled, priority, created_at, updated_at FROM usage_cap_policies` func scanPolicy(row scanner) (store.UsageCapPolicy, error) { var p store.UsageCapPolicy var agentID, providerID uuid.NullUUID var maxTokens, maxCost sql.NullInt64 err := row.Scan(&p.ID, &p.TenantID, &agentID, &providerID, &p.ProviderType, &p.ModelID, - &p.Window, &maxTokens, &maxCost, &p.Enabled, &p.Priority, &p.CreatedAt, &p.UpdatedAt) + &p.Window, &maxTokens, &maxCost, &p.Source, &p.Enabled, &p.Priority, &p.CreatedAt, &p.UpdatedAt) if agentID.Valid { p.AgentID = &agentID.UUID } diff --git a/internal/store/usage_caps.go b/internal/store/usage_caps.go index 1c123690..6848a4c6 100644 --- a/internal/store/usage_caps.go +++ b/internal/store/usage_caps.go @@ -19,9 +19,15 @@ const ( UsageCapEventBlock = "block" UsageCapEventReconcile = "reconcile" UsageCapEventSkip = "skip" + + UsageCapSourceManual = "manual" + UsageCapSourceAgentBudget = "agent_budget_monthly_cents" ) -var ErrUsageCapExceeded = errors.New("usage cap exceeded") +var ( + ErrUsageCapExceeded = errors.New("usage cap exceeded") + ErrUsageCapPolicyManaged = errors.New("usage cap policy is managed by another setting") +) type UsageCapExceededError struct { PolicyID uuid.UUID @@ -101,6 +107,7 @@ type UsageCapPolicy struct { Window string `json:"window" db:"window"` MaxTokens *int64 `json:"max_tokens,omitempty" db:"max_tokens"` MaxCostMicros *int64 `json:"max_cost_micros,omitempty" db:"max_cost_micros"` + Source string `json:"source,omitempty" db:"source"` Enabled bool `json:"enabled" db:"enabled"` Priority int `json:"priority" db:"priority"` CreatedAt time.Time `json:"created_at" db:"created_at"` diff --git a/internal/upgrade/version.go b/internal/upgrade/version.go index ad21b8a5..3ce25186 100644 --- a/internal/upgrade/version.go +++ b/internal/upgrade/version.go @@ -2,4 +2,4 @@ package upgrade // RequiredSchemaVersion is the schema migration version this binary requires. // Bump this whenever adding a new SQL migration file. -const RequiredSchemaVersion uint = 70 +const RequiredSchemaVersion uint = 71 diff --git a/migrations/000071_agent_budget_usage_cap_bridge.down.sql b/migrations/000071_agent_budget_usage_cap_bridge.down.sql new file mode 100644 index 00000000..e26c4b4a --- /dev/null +++ b/migrations/000071_agent_budget_usage_cap_bridge.down.sql @@ -0,0 +1,7 @@ +DELETE FROM usage_cap_policies +WHERE source = 'agent_budget_monthly_cents'; + +DROP INDEX IF EXISTS idx_usage_cap_policies_agent_budget_source; + +ALTER TABLE usage_cap_policies + DROP COLUMN IF EXISTS source; diff --git a/migrations/000071_agent_budget_usage_cap_bridge.up.sql b/migrations/000071_agent_budget_usage_cap_bridge.up.sql new file mode 100644 index 00000000..700d1b9d --- /dev/null +++ b/migrations/000071_agent_budget_usage_cap_bridge.up.sql @@ -0,0 +1,23 @@ +ALTER TABLE usage_cap_policies + ADD COLUMN IF NOT EXISTS source TEXT NOT NULL DEFAULT 'manual'; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_cap_policies_agent_budget_source + ON usage_cap_policies (tenant_id, agent_id) + WHERE source = 'agent_budget_monthly_cents'; + +INSERT INTO usage_cap_policies ( + tenant_id, agent_id, window_key, max_cost_micros, enabled, priority, source +) +SELECT + tenant_id, + id, + 'month', + budget_monthly_cents::BIGINT * 10000, + true, + 90, + 'agent_budget_monthly_cents' +FROM agents +WHERE deleted_at IS NULL + AND budget_monthly_cents IS NOT NULL + AND budget_monthly_cents > 0 +ON CONFLICT DO NOTHING; diff --git a/ui/web/src/i18n/locales/en/usage.json b/ui/web/src/i18n/locales/en/usage.json index cc82afce..70a45e09 100644 --- a/ui/web/src/i18n/locales/en/usage.json +++ b/ui/web/src/i18n/locales/en/usage.json @@ -47,6 +47,8 @@ "tenantScope": "Tenant cap", "tenantScoped": "Tenant scoped", "agentScoped": "Agent scoped", + "agentBudgetSource": "Agent budget", + "agentBudgetManaged": "Clear this from the agent monthly budget field.", "enabled": "Enabled", "disabled": "Disabled", "delete": "Delete cap", diff --git a/ui/web/src/i18n/locales/vi/usage.json b/ui/web/src/i18n/locales/vi/usage.json index a7beab3f..28b72586 100644 --- a/ui/web/src/i18n/locales/vi/usage.json +++ b/ui/web/src/i18n/locales/vi/usage.json @@ -47,6 +47,8 @@ "tenantScope": "Cap tenant", "tenantScoped": "Phạm vi tenant", "agentScoped": "Phạm vi agent", + "agentBudgetSource": "Ngân sách agent", + "agentBudgetManaged": "Xóa giới hạn này từ ô ngân sách hàng tháng của agent.", "enabled": "Đã bật", "disabled": "Đã tắt", "delete": "Xóa cap", diff --git a/ui/web/src/i18n/locales/zh/usage.json b/ui/web/src/i18n/locales/zh/usage.json index 1810eb43..58d6b2f4 100644 --- a/ui/web/src/i18n/locales/zh/usage.json +++ b/ui/web/src/i18n/locales/zh/usage.json @@ -47,6 +47,8 @@ "tenantScope": "租户上限", "tenantScoped": "租户范围", "agentScoped": "Agent范围", + "agentBudgetSource": "Agent预算", + "agentBudgetManaged": "请在 Agent 月度预算字段中清除此上限。", "enabled": "已启用", "disabled": "已禁用", "delete": "删除上限", 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 6f39de72..2371cf4b 100644 --- a/ui/web/src/pages/usage/components/usage-caps-panel.tsx +++ b/ui/web/src/pages/usage/components/usage-caps-panel.tsx @@ -161,18 +161,22 @@ function UsageCapRow({ row, onDelete }: { row: UsageCapUtilization; onDelete: () const costUsed = row.used_cost_micros + row.reserved_cost_micros; const tokenPct = p.max_tokens ? Math.min(100, Math.round((tokenUsed / p.max_tokens) * 100)) : 0; const costPct = p.max_cost_micros ? Math.min(100, Math.round((costUsed / p.max_cost_micros) * 100)) : 0; + const isAgentBudget = p.source === "agent_budget_monthly_cents"; return (
{p.model_id || p.provider_type || t("caps.tenantScope")}
-
{p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")}
+
+ {p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")} + {isAgentBudget ? {t("caps.agentBudgetSource")} : null} +
{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.enabled ? t("caps.enabled") : t("caps.disabled")} - + ); diff --git a/ui/web/src/types/usage-caps.ts b/ui/web/src/types/usage-caps.ts index 670687be..882e38ea 100644 --- a/ui/web/src/types/usage-caps.ts +++ b/ui/web/src/types/usage-caps.ts @@ -19,6 +19,7 @@ export interface UsageCapPolicy { window: "hour" | "day" | "week" | "month"; max_tokens?: number; max_cost_micros?: number; + source?: "manual" | "agent_budget_monthly_cents"; enabled: boolean; priority: number; created_at: string; From e4acb147d6be789955fed6f88fe70d80c8c3c1a1 Mon Sep 17 00:00:00 2001 From: Goon Date: Sat, 23 May 2026 09:34:29 +0700 Subject: [PATCH 04/10] feat(usage): edit cap policies and trace decisions --- docs/project-changelog.md | 2 + internal/agent/loop_pipeline_callbacks.go | 10 ++ internal/agent/loop_tracing.go | 27 +++- internal/http/usage_caps.go | 35 ++++- internal/http/usage_caps_policy_patch_test.go | 35 +++++ internal/usage/caps/reservation_metadata.go | 125 ++++++++++++++++++ internal/usage/caps/service.go | 38 ++++-- internal/usage/caps/service_test.go | 83 ++++++++++++ ui/web/src/i18n/locales/en/usage.json | 5 + ui/web/src/i18n/locales/vi/usage.json | 5 + ui/web/src/i18n/locales/zh/usage.json | 5 + .../pages/usage/components/usage-cap-row.tsx | 44 ++++++ .../usage/components/usage-caps-panel.tsx | 117 +++++++++------- .../src/pages/usage/hooks/use-usage-caps.ts | 19 ++- 14 files changed, 478 insertions(+), 72 deletions(-) create mode 100644 internal/http/usage_caps_policy_patch_test.go create mode 100644 internal/usage/caps/reservation_metadata.go create mode 100644 ui/web/src/pages/usage/components/usage-cap-row.tsx diff --git a/docs/project-changelog.md b/docs/project-changelog.md index b965c81a..84b9c87c 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -18,6 +18,8 @@ Significant changes, features, and fixes in reverse chronological order. - Added OpenRouter alias resolution for native model IDs, cached-input accounting normalization, and partial-stream failure reconciliation. - Bridged legacy `budget_monthly_cents` into generated monthly agent USD cap policies, including migration backfill and save-time sync. - Added web dashboard controls on Usage and Provider detail pages. +- Added Usage page editing for manual cap policies, including enable/disable, scope clearing, and token/USD limit clearing while keeping generated agent-budget caps read-only. +- Added usage-cap decision metadata to LLM spans, including allow/skip/block reason, policy IDs, estimates, actuals, and reconcile status. **Tests** diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index 6befa4d1..9d19e0b6 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -295,6 +295,11 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c opts = append(opts, withProvider(provider.Name())) } spanID := l.emitLLMSpanStart(ctx, start, state.Iteration+1, chatReq.Messages, opts...) + recordUsageCapAttempt := func(reservation *usagecaps.Reservation) { + if reservation != nil { + opts = append(opts, withUsageCapMetadata(reservation.TraceMetadata())) + } + } emitChunk := func(chunk providers.StreamChunk) { if chunk.Thinking != "" { @@ -320,6 +325,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c candidateAttempt := fmt.Sprintf("%s:%s:%s", attempt, entry.ProviderName, actualReq.Model) reservation, reserveErr := l.reserveLLMUsageFor(callCtx, req, state.Iteration, actualReq, candidateAttempt, entry.ProviderName, actualReq.Model) if reserveErr != nil { + recordUsageCapAttempt(reservation) return nil, reserveErr } return func(callResp *providers.ChatResponse, callErr error, info providers.FallbackCallInfo) { @@ -329,6 +335,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } else { reservation.Reconcile(callCtx, callResp, callErr) } + recordUsageCapAttempt(reservation) } }, nil } @@ -339,6 +346,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } reservation, reserveErr := l.reserveLLMUsage(ctx, req, state, request, attempt) if reserveErr != nil { + recordUsageCapAttempt(reservation) return nil, reserveErr } var callResp *providers.ChatResponse @@ -353,6 +361,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c }) if reservation != nil { reservation.ReconcileStream(ctx, callResp, callErr, streamed) + recordUsageCapAttempt(reservation) } return callResp, callErr } else { @@ -360,6 +369,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } if reservation != nil { reservation.Reconcile(ctx, callResp, callErr) + recordUsageCapAttempt(reservation) } return callResp, callErr } diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index b21c1817..469362a7 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -14,6 +14,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" ) func (l *Loop) emit(event AgentEvent) { @@ -51,20 +52,33 @@ func (l *Loop) IsRunning() bool { return l.activeRuns.Load() > 0 } type spanOption func(*spanOverrides) type spanOverrides struct { - model string - provider string + model string + provider string + usageCapAttempts []usagecaps.TraceMetadata } func withModel(m string) spanOption { return func(o *spanOverrides) { o.model = m } } func withProvider(p string) spanOption { return func(o *spanOverrides) { o.provider = p } } +func withUsageCapMetadata(metadata usagecaps.TraceMetadata) spanOption { + return func(o *spanOverrides) { + if !metadata.Empty() { + o.usageCapAttempts = append(o.usageCapAttempts, metadata) + } + } +} // resolveSpan returns (model, provider) applying any overrides on top of agent defaults. func (l *Loop) resolveSpan(opts []spanOption) (string, string) { + o := l.resolveSpanOverrides(opts) + return o.model, o.provider +} + +func (l *Loop) resolveSpanOverrides(opts []spanOption) spanOverrides { o := spanOverrides{model: l.model, provider: l.provider.Name()} for _, fn := range opts { fn(&o) } - return o.model, o.provider + return o } // --------------------------------------------------------------------------- @@ -150,6 +164,7 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time. "status": store.SpanStatusCompleted, } var spanMetadata json.RawMessage + spanOpts := l.resolveSpanOverrides(opts) if callErr != nil { updates["status"] = store.SpanStatusError @@ -176,8 +191,7 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time. } } // Calculate cost if pricing config is available. - model, providerName := l.resolveSpan(opts) - if pricing := tracing.LookupPricing(l.modelPricing, providerName, model); pricing != nil { + 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 @@ -203,6 +217,9 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time. if decision := providers.ReasoningDecisionFromContext(ctx); decision != nil { spanMetadata = providers.MergeReasoningMetadata(spanMetadata, *decision) } + if len(spanOpts.usageCapAttempts) > 0 { + spanMetadata = usagecaps.MergeTraceMetadata(spanMetadata, spanOpts.usageCapAttempts) + } if len(spanMetadata) > 0 { updates["metadata"] = spanMetadata } diff --git a/internal/http/usage_caps.go b/internal/http/usage_caps.go index 1fc64c8f..cf1940ca 100644 --- a/internal/http/usage_caps.go +++ b/internal/http/usage_caps.go @@ -5,9 +5,11 @@ import ( "encoding/json" "errors" "fmt" + "io" "log/slog" "net/http" "strconv" + "strings" "time" "github.com/google/uuid" @@ -97,12 +99,12 @@ func (h *UsageCapsHandler) handleUpdatePolicy(w http.ResponseWriter, r *http.Req writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid policy id"}) return } - var body policyBody - if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + bodyBytes, err := io.ReadAll(r.Body) + if err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) return } - patch, err := body.toPatch() + patch, err := policyPatchFromBody(bodyBytes) if err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) return @@ -325,6 +327,33 @@ func (b policyBody) toPatch() (store.UsageCapPolicyPatch, error) { return patch, nil } +func policyPatchFromBody(bodyBytes []byte) (store.UsageCapPolicyPatch, error) { + var body policyBody + if err := json.Unmarshal(bodyBytes, &body); err != nil { + return store.UsageCapPolicyPatch{}, err + } + patch, err := body.toPatch() + if err != nil { + return patch, err + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(bodyBytes, &raw); err == nil { + if isJSONNull(raw["max_tokens"]) { + var v *int64 + patch.MaxTokens = &v + } + if isJSONNull(raw["max_cost_micros"]) || isJSONNull(raw["max_cost_usd"]) { + var v *int64 + patch.MaxCostMicros = &v + } + } + return patch, nil +} + +func isJSONNull(raw json.RawMessage) bool { + return strings.TrimSpace(string(raw)) == "null" +} + func tenantIDOrMaster(r *http.Request) uuid.UUID { if tid := store.TenantIDFromContext(r.Context()); tid != uuid.Nil { return tid diff --git a/internal/http/usage_caps_policy_patch_test.go b/internal/http/usage_caps_policy_patch_test.go new file mode 100644 index 00000000..6038370e --- /dev/null +++ b/internal/http/usage_caps_policy_patch_test.go @@ -0,0 +1,35 @@ +package http + +import "testing" + +func TestPolicyPatchFromBodyClearsTokenAndCostLimits(t *testing.T) { + patch, err := policyPatchFromBody([]byte(`{"max_tokens":null,"max_cost_usd":null}`)) + if err != nil { + t.Fatalf("policyPatchFromBody: %v", err) + } + if patch.MaxTokens == nil { + t.Fatal("MaxTokens patch missing") + } + if *patch.MaxTokens != nil { + t.Fatalf("MaxTokens = %v, want nil clear", **patch.MaxTokens) + } + if patch.MaxCostMicros == nil { + t.Fatal("MaxCostMicros patch missing") + } + if *patch.MaxCostMicros != nil { + t.Fatalf("MaxCostMicros = %v, want nil clear", **patch.MaxCostMicros) + } +} + +func TestPolicyPatchFromBodyPreservesProvidedCostUSD(t *testing.T) { + patch, err := policyPatchFromBody([]byte(`{"max_cost_usd":12.5}`)) + if err != nil { + t.Fatalf("policyPatchFromBody: %v", err) + } + if patch.MaxCostMicros == nil || *patch.MaxCostMicros == nil { + t.Fatal("MaxCostMicros patch missing") + } + if got := **patch.MaxCostMicros; got != 12_500_000 { + t.Fatalf("MaxCostMicros = %d, want 12500000", got) + } +} diff --git a/internal/usage/caps/reservation_metadata.go b/internal/usage/caps/reservation_metadata.go new file mode 100644 index 00000000..ad60a88d --- /dev/null +++ b/internal/usage/caps/reservation_metadata.go @@ -0,0 +1,125 @@ +package caps + +import ( + "encoding/json" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/usage/pricing" +) + +const TraceMetadataKey = "usage_caps" + +type TraceMetadata struct { + Decision string `json:"decision"` + Reason string `json:"reason,omitempty"` + ReservationKey string `json:"reservation_key,omitempty"` + ProviderName string `json:"provider_name,omitempty"` + ProviderType string `json:"provider_type,omitempty"` + ModelID string `json:"model_id,omitempty"` + EstimatedTokens int64 `json:"estimated_tokens"` + EstimatedCostMicros int64 `json:"estimated_cost_micros"` + ActualTokens int64 `json:"actual_tokens"` + ActualCostMicros int64 `json:"actual_cost_micros"` + ReconcileStatus string `json:"reconcile_status,omitempty"` + PolicyCount int `json:"policy_count"` + PolicyIDs []string `json:"policy_ids,omitempty"` +} + +func (r *Reservation) TraceMetadata() TraceMetadata { + if r == nil { + return TraceMetadata{} + } + decision := r.decision + if decision == "" { + if r.skipped { + decision = store.UsageCapEventSkip + } else { + decision = store.UsageCapEventAllow + } + } + policyCount := 0 + policyIDs := make([]string, 0) + if r.result != nil { + policyCount = len(r.result.Policies) + for _, policy := range r.result.Policies { + if policy.ID != uuid.Nil { + policyIDs = append(policyIDs, policy.ID.String()) + } + } + } + if r.blockedPolicyID != uuid.Nil { + policyIDs = append(policyIDs, r.blockedPolicyID.String()) + if policyCount == 0 { + policyCount = 1 + } + } + return TraceMetadata{ + Decision: decision, + Reason: r.reason, + ReservationKey: r.key, + ProviderName: r.providerName, + ProviderType: r.providerType, + ModelID: r.modelID, + EstimatedTokens: r.usage.TotalTokens(), + EstimatedCostMicros: r.estimatedCostMicros, + ActualTokens: r.actualTokens, + ActualCostMicros: r.actualCostMicros, + ReconcileStatus: r.reconcileStatus, + PolicyCount: policyCount, + PolicyIDs: policyIDs, + } +} + +func (m TraceMetadata) Empty() bool { + return m.Decision == "" && m.Reason == "" && m.ReservationKey == "" && + m.ProviderName == "" && m.ProviderType == "" && m.ModelID == "" && + m.EstimatedTokens == 0 && m.EstimatedCostMicros == 0 && + m.ActualTokens == 0 && m.ActualCostMicros == 0 && + m.ReconcileStatus == "" && m.PolicyCount == 0 && len(m.PolicyIDs) == 0 +} + +func MergeTraceMetadata(existing json.RawMessage, entries []TraceMetadata) json.RawMessage { + clean := make([]TraceMetadata, 0, len(entries)) + for _, entry := range entries { + if !entry.Empty() { + clean = append(clean, entry) + } + } + if len(clean) == 0 { + return existing + } + payload := map[string]any{} + if len(existing) > 0 { + _ = json.Unmarshal(existing, &payload) + } + payload[TraceMetadataKey] = map[string]any{"attempts": clean} + data, err := json.Marshal(payload) + if err != nil { + return existing + } + return json.RawMessage(data) +} + +func skippedReservation(req Request, reason string) *Reservation { + return &Reservation{ + skipped: true, decision: store.UsageCapEventSkip, reason: reason, + providerName: req.ProviderName, modelID: req.ModelID, + } +} + +func skippedScopedReservation(req Request, scope store.UsageCapScope, reason string) *Reservation { + r := skippedReservation(req, reason) + r.providerType = scope.ProviderType + r.modelID = scope.ModelID + return r +} + +func blockedReservation(req Request, scope store.UsageCapScope, key string, usage pricing.BillableUsage, costMicros int64, policyID uuid.UUID, reason string) *Reservation { + return &Reservation{ + key: key, usage: usage, estimatedCostMicros: costMicros, + decision: store.UsageCapEventBlock, reason: reason, + blockedPolicyID: policyID, + providerName: req.ProviderName, providerType: scope.ProviderType, modelID: scope.ModelID, + } +} diff --git a/internal/usage/caps/service.go b/internal/usage/caps/service.go index 12db8dd9..d3fbbaf3 100644 --- a/internal/usage/caps/service.go +++ b/internal/usage/caps/service.go @@ -50,17 +50,25 @@ type Reservation struct { usage pricing.BillableUsage prices store.UsagePricingFields estimatedCostMicros int64 + actualTokens int64 + actualCostMicros int64 + reconcileStatus string skipped bool + decision string reason string + blockedPolicyID uuid.UUID + providerName string + providerType string + modelID string } func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, error) { if s == nil || s.store == nil { - return &Reservation{skipped: true, reason: "service_disabled"}, nil + return skippedReservation(req, "service_disabled"), nil } providerData, err := s.resolveProvider(ctx, req.TenantID, req.ProviderName) if err != nil { - return &Reservation{skipped: true, reason: "provider_metadata_missing"}, nil + return skippedReservation(req, "provider_metadata_missing"), nil } scope := store.UsageCapScope{ TenantID: req.TenantID, AgentID: req.AgentID, ProviderID: providerData.ID, @@ -71,14 +79,14 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err TenantID: req.TenantID, Decision: store.UsageCapEventSkip, Reason: "provider_not_billable_api", Metadata: mustJSON(scope), }) - return &Reservation{skipped: true, reason: "provider_not_billable_api"}, nil + return skippedScopedReservation(req, scope, "provider_not_billable_api"), nil } policies, err := s.store.ListUsageCapPolicies(ctx, scope, false) if err != nil { return nil, err } if len(policies) == 0 { - return &Reservation{skipped: true, reason: "no_policy"}, nil + return skippedScopedReservation(req, scope, "no_policy"), nil } usage := pricing.BillableUsage{ InputTokens: int64(EstimateInputTokens(req.Messages)), @@ -88,6 +96,10 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err if usage.OutputTokens <= 0 { usage.OutputTokens = 1 } + key := req.ReservationKey + if key == "" { + key = uuid.NewString() + } metadata := map[string]any{"model_id": req.ModelID, "pricing_source": "token_only"} var prices store.UsagePricingFields var costMicros int64 @@ -95,7 +107,7 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err resolved, err := s.store.ResolvePricing(ctx, req.TenantID, providerData.ID, providerData.Name, providerData.ProviderType, req.ModelID) if err != nil { if errors.Is(err, sql.ErrNoRows) { - return nil, fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID) + return blockedReservation(req, scope, key, usage, 0, uuid.Nil, "pricing_unknown"), fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID) } return nil, err } @@ -109,10 +121,6 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err } metadata = map[string]any{"source": resolved.Source, "model_id": resolved.ModelID} } - key := req.ReservationKey - if key == "" { - key = uuid.NewString() - } result, err := s.store.ReserveUsage(ctx, store.UsageReserveRequest{ UsageCapScope: scope, ReservationKey: key, EstimatedTokens: usage.TotalTokens(), EstimatedCostMicros: costMicros, @@ -136,11 +144,16 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err Metadata: mustJSON(map[string]any{"model_id": req.ModelID, "provider": req.ProviderName}), }) slog.Warn("usage_caps.blocked", "tenant_id", req.TenantID, "policy_id", blockedPolicy, "reason", reason) - return nil, ErrCapExceeded + return blockedReservation(req, scope, key, usage, costMicros, blockedPolicy, reason), ErrCapExceeded } return nil, err } - return &Reservation{key: key, result: result, svc: s, usage: usage, prices: prices, estimatedCostMicros: costMicros}, nil + return &Reservation{ + key: key, result: result, svc: s, usage: usage, prices: prices, + estimatedCostMicros: costMicros, decision: store.UsageCapEventAllow, + reason: "reserved", providerName: req.ProviderName, + providerType: scope.ProviderType, modelID: scope.ModelID, + }, nil } func (r *Reservation) Reconcile(ctx context.Context, resp *providers.ChatResponse, callErr error) { @@ -189,6 +202,9 @@ func (r *Reservation) reconcile(ctx context.Context, resp *providers.ChatRespons if err != nil { cost = r.estimatedCostMicros } + r.actualTokens = actual.TotalTokens() + r.actualCostMicros = cost + r.reconcileStatus = status reconcileCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() if err := r.svc.store.ReconcileUsage(reconcileCtx, store.UsageReconcileRequest{ diff --git a/internal/usage/caps/service_test.go b/internal/usage/caps/service_test.go index e1350d6c..7a14a084 100644 --- a/internal/usage/caps/service_test.go +++ b/internal/usage/caps/service_test.go @@ -3,6 +3,8 @@ package caps import ( "context" "database/sql" + "encoding/json" + "errors" "testing" "github.com/google/uuid" @@ -61,6 +63,13 @@ func TestPreflightTokenOnlyCapDoesNotRequirePricing(t *testing.T) { if usageStore.reserved.EstimatedCostMicros != 0 { t.Fatalf("EstimatedCostMicros = %d, want 0", usageStore.reserved.EstimatedCostMicros) } + metadata := reservation.TraceMetadata() + if metadata.Decision != store.UsageCapEventAllow { + t.Fatalf("Decision = %q, want allow", metadata.Decision) + } + if metadata.PolicyCount != 1 { + t.Fatalf("PolicyCount = %d, want 1", metadata.PolicyCount) + } } func TestPreflightIncludesRequestPricingWhenConfigured(t *testing.T) { @@ -224,6 +233,76 @@ func TestReservationReconcileIgnoresUnpricedRequestCount(t *testing.T) { if usageStore.reconciled.ActualCostMicros != 5 { t.Fatalf("ActualCostMicros = %d, want 5", usageStore.reconciled.ActualCostMicros) } + metadata := reservation.TraceMetadata() + if metadata.ActualTokens != 5 { + t.Fatalf("ActualTokens = %d, want 5", metadata.ActualTokens) + } + if metadata.ReconcileStatus != "reconciled" { + t.Fatalf("ReconcileStatus = %q, want reconciled", metadata.ReconcileStatus) + } +} + +func TestPreflightTraceMetadataForCapExceeded(t *testing.T) { + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(10), Enabled: true} + usageStore := &fakeUsageCapStore{ + policies: []store.UsageCapPolicy{policy}, + reserveErr: &store.UsageCapExceededError{PolicyID: policy.ID, Reason: "token_cap_exceeded"}, + } + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + + reservation, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "token/model", + ReservationKey: "blocked", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 10, + }) + if !errors.Is(err, ErrCapExceeded) { + t.Fatalf("Preflight error = %v, want ErrCapExceeded", err) + } + metadata := reservation.TraceMetadata() + if metadata.Decision != store.UsageCapEventBlock { + t.Fatalf("Decision = %q, want block", metadata.Decision) + } + if metadata.Reason != "token_cap_exceeded" { + t.Fatalf("Reason = %q, want token_cap_exceeded", metadata.Reason) + } + if metadata.ReservationKey != "blocked" { + t.Fatalf("ReservationKey = %q, want blocked", metadata.ReservationKey) + } + if len(metadata.PolicyIDs) != 1 || metadata.PolicyIDs[0] != policy.ID.String() { + t.Fatalf("PolicyIDs = %v, want [%s]", metadata.PolicyIDs, policy.ID) + } +} + +func TestMergeTraceMetadataPreservesExistingSections(t *testing.T) { + existing := json.RawMessage(`{"thinking":{"effort":"high"}}`) + merged := MergeTraceMetadata(existing, []TraceMetadata{{ + Decision: store.UsageCapEventAllow, + Reason: "reserved", + ModelID: "openai/gpt-test", + }}) + + var payload map[string]json.RawMessage + if err := json.Unmarshal(merged, &payload); err != nil { + t.Fatalf("Unmarshal merged metadata: %v", err) + } + if len(payload["thinking"]) == 0 { + t.Fatal("existing thinking metadata was removed") + } + var usagePayload struct { + Attempts []TraceMetadata `json:"attempts"` + } + if err := json.Unmarshal(payload[TraceMetadataKey], &usagePayload); err != nil { + t.Fatalf("Unmarshal usage caps metadata: %v", err) + } + if len(usagePayload.Attempts) != 1 || usagePayload.Attempts[0].Decision != store.UsageCapEventAllow { + t.Fatalf("usage cap attempts = %+v, want one allow attempt", usagePayload.Attempts) + } } func TestCountImagesOnlyCountsImageMIMEs(t *testing.T) { @@ -245,6 +324,7 @@ type fakeUsageCapStore struct { resolved *store.ResolvedUsagePricing resolveErr error resolveCalls int + reserveErr error reserved store.UsageReserveRequest reconciled store.UsageReconcileRequest reconcileCalls int @@ -287,6 +367,9 @@ func (s *fakeUsageCapStore) DeleteUsageCapPolicy(context.Context, uuid.UUID, uui } func (s *fakeUsageCapStore) ReserveUsage(_ context.Context, req store.UsageReserveRequest, policies []store.UsageCapPolicy) (*store.UsageReservationResult, error) { s.reserved = req + if s.reserveErr != nil { + return nil, s.reserveErr + } return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Policies: policies}, nil } func (s *fakeUsageCapStore) ReconcileUsage(ctx context.Context, req store.UsageReconcileRequest) error { diff --git a/ui/web/src/i18n/locales/en/usage.json b/ui/web/src/i18n/locales/en/usage.json index 70a45e09..477b8608 100644 --- a/ui/web/src/i18n/locales/en/usage.json +++ b/ui/web/src/i18n/locales/en/usage.json @@ -40,6 +40,9 @@ "maxTokens": "Max tokens", "maxCost": "Max cost USD", "create": "Create cap", + "edit": "Edit cap", + "save": "Save cap", + "cancel": "Cancel edit", "scope": "Scope", "tokens": "Tokens", "cost": "Cost", @@ -56,7 +59,9 @@ "blocked": "Blocked", "toast": { "created": "Usage cap created", + "updated": "Usage cap updated", "createFailed": "Could not create usage cap", + "updateFailed": "Could not update usage cap", "deleted": "Usage cap deleted", "deleteFailed": "Could not delete usage cap" } diff --git a/ui/web/src/i18n/locales/vi/usage.json b/ui/web/src/i18n/locales/vi/usage.json index 28b72586..dd2a9f45 100644 --- a/ui/web/src/i18n/locales/vi/usage.json +++ b/ui/web/src/i18n/locales/vi/usage.json @@ -40,6 +40,9 @@ "maxTokens": "Token tối đa", "maxCost": "Chi phí tối đa USD", "create": "Tạo cap", + "edit": "Sửa cap", + "save": "Lưu cap", + "cancel": "Hủy sửa", "scope": "Phạm vi", "tokens": "Token", "cost": "Chi phí", @@ -56,7 +59,9 @@ "blocked": "Đã chặn", "toast": { "created": "Đã tạo usage cap", + "updated": "Đã cập nhật usage cap", "createFailed": "Không thể tạo usage cap", + "updateFailed": "Không thể cập nhật usage cap", "deleted": "Đã xóa usage cap", "deleteFailed": "Không thể xóa usage cap" } diff --git a/ui/web/src/i18n/locales/zh/usage.json b/ui/web/src/i18n/locales/zh/usage.json index 58d6b2f4..b200c31f 100644 --- a/ui/web/src/i18n/locales/zh/usage.json +++ b/ui/web/src/i18n/locales/zh/usage.json @@ -40,6 +40,9 @@ "maxTokens": "最大令牌", "maxCost": "最大费用 USD", "create": "创建上限", + "edit": "编辑上限", + "save": "保存上限", + "cancel": "取消编辑", "scope": "范围", "tokens": "令牌", "cost": "费用", @@ -56,7 +59,9 @@ "blocked": "已拦截", "toast": { "created": "用量上限已创建", + "updated": "用量上限已更新", "createFailed": "无法创建用量上限", + "updateFailed": "无法更新用量上限", "deleted": "用量上限已删除", "deleteFailed": "无法删除用量上限" } diff --git a/ui/web/src/pages/usage/components/usage-cap-row.tsx b/ui/web/src/pages/usage/components/usage-cap-row.tsx new file mode 100644 index 00000000..99786c2d --- /dev/null +++ b/ui/web/src/pages/usage/components/usage-cap-row.tsx @@ -0,0 +1,44 @@ +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 type { UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps"; + +interface UsageCapRowProps { + row: UsageCapUtilization; + onEdit: (policy: UsageCapPolicy) => void; + onDelete: () => void; +} + +export function UsageCapRow({ row, onEdit, onDelete }: UsageCapRowProps) { + const { t } = useTranslation("usage"); + const p = row.policy; + const tokenUsed = row.used_tokens + row.reserved_tokens; + const costUsed = row.used_cost_micros + row.reserved_cost_micros; + const tokenPct = p.max_tokens ? Math.min(100, Math.round((tokenUsed / p.max_tokens) * 100)) : 0; + const costPct = p.max_cost_micros ? Math.min(100, Math.round((costUsed / p.max_cost_micros) * 100)) : 0; + const isAgentBudget = p.source === "agent_budget_monthly_cents"; + + return ( + + +
{p.model_id || p.provider_type || t("caps.tenantScope")}
+
+ {p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")} + {isAgentBudget ? {t("caps.agentBudgetSource")} : null} +
+ + {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.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 2371cf4b..972f1c34 100644 --- a/ui/web/src/pages/usage/components/usage-caps-panel.tsx +++ b/ui/web/src/pages/usage/components/usage-caps-panel.tsx @@ -1,17 +1,18 @@ import { useMemo, useState } from "react"; import type { ReactNode } from "react"; import { useTranslation } from "react-i18next"; -import { Plus, RefreshCw, ShieldAlert, Trash2 } from "lucide-react"; +import { Plus, RefreshCw, Save, ShieldAlert, X } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { Badge } from "@/components/ui/badge"; +import { Switch } from "@/components/ui/switch"; import { formatCost, 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"; -import type { UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps"; +import type { UsageCapPolicy } from "@/types/usage-caps"; +import { UsageCapRow } from "./usage-cap-row"; const ALL = "__all__"; @@ -19,37 +20,68 @@ export function UsageCapsPanel() { const { t } = useTranslation("usage"); const { agents } = useAgents(); const { providers } = useProviders(); - const { utilization, events, refreshing, refresh, createPolicy, deletePolicy } = useUsageCaps(); + const { utilization, events, refreshing, refresh, createPolicy, updatePolicy, deletePolicy } = useUsageCaps(); const [windowValue, setWindowValue] = useState("day"); const [agentId, setAgentId] = useState(ALL); const [providerId, setProviderId] = useState(ALL); const [modelId, setModelId] = useState(""); const [maxTokens, setMaxTokens] = useState(""); const [maxCost, setMaxCost] = useState(""); + const [enabled, setEnabled] = useState(true); + const [editingPolicy, setEditingPolicy] = useState(null); const [saving, setSaving] = useState(false); const provider = useMemo(() => providers.find((p) => p.id === providerId), [providerId, providers]); + const providerType = provider?.provider_type ?? (editingPolicy && providerId === (editingPolicy.provider_id ?? ALL) ? editingPolicy.provider_type : undefined); const blockedEvents = events.filter((event) => event.decision === "block"); + const isEditing = editingPolicy != null; + + const resetForm = () => { + setEditingPolicy(null); + setWindowValue("day"); + setAgentId(ALL); + setProviderId(ALL); + setModelId(""); + setMaxTokens(""); + setMaxCost(""); + setEnabled(true); + }; + + const startEdit = (policy: UsageCapPolicy) => { + setEditingPolicy(policy); + setWindowValue(policy.window); + setAgentId(policy.agent_id ?? ALL); + setProviderId(policy.provider_id ?? ALL); + setModelId(policy.model_id ?? ""); + setMaxTokens(policy.max_tokens ? String(policy.max_tokens) : ""); + setMaxCost(policy.max_cost_micros ? String(policy.max_cost_micros / 1_000_000) : ""); + setEnabled(policy.enabled); + }; const onSubmit = async () => { - const tokens = Number(maxTokens); - const cost = Number(maxCost); - if ((!Number.isFinite(tokens) || tokens <= 0) && (!Number.isFinite(cost) || cost <= 0)) return; + const tokenValue = maxTokens.trim(); + const costValue = maxCost.trim(); + const tokens = Number(tokenValue); + const cost = Number(costValue); + if ((!tokenValue || !Number.isFinite(tokens) || tokens <= 0) && (!costValue || !Number.isFinite(cost) || cost <= 0)) return; setSaving(true); try { - await createPolicy({ + const input = { window: windowValue, - agent_id: agentId === ALL ? undefined : agentId, - provider_id: providerId === ALL ? undefined : providerId, - provider_type: provider?.provider_type, - model_id: modelId.trim() || undefined, - max_tokens: Number.isFinite(tokens) && tokens > 0 ? Math.floor(tokens) : undefined, - max_cost_usd: Number.isFinite(cost) && cost > 0 ? cost : undefined, - enabled: true, - }); - setMaxTokens(""); - setMaxCost(""); - setModelId(""); + agent_id: agentId === ALL ? (isEditing ? "" : undefined) : agentId, + provider_id: providerId === ALL ? (isEditing ? "" : undefined) : providerId, + provider_type: providerType ?? (isEditing ? "" : undefined), + model_id: modelId.trim() || (isEditing ? "" : undefined), + max_tokens: tokenValue && Number.isFinite(tokens) && tokens > 0 ? Math.floor(tokens) : (isEditing ? null : undefined), + max_cost_usd: costValue && Number.isFinite(cost) && cost > 0 ? cost : (isEditing ? null : undefined), + enabled, + }; + if (editingPolicy) { + await updatePolicy(editingPolicy.id, input); + } else { + await createPolicy(input); + } + resetForm(); } finally { setSaving(false); } @@ -68,7 +100,7 @@ export function UsageCapsPanel() { -
+
setMaxCost(e.target.value)} inputMode="decimal" placeholder="25" className="text-base md:text-sm" /> - + {isEditing ? ( + + ) : null} +
+ + +
+ + {enabled ? t("caps.enabled") : t("caps.disabled")}
@@ -127,7 +170,7 @@ export function UsageCapsPanel() { {utilization.length === 0 ? ( {t("caps.empty")} ) : utilization.map((row) => ( - void deletePolicy(row.policy.id)} /> + void deletePolicy(row.policy.id)} /> ))} @@ -153,31 +196,3 @@ export function UsageCapsPanel() { function Field({ label, children }: { label: string; children: ReactNode }) { return
{children}
; } - -function UsageCapRow({ row, onDelete }: { row: UsageCapUtilization; onDelete: () => void }) { - const { t } = useTranslation("usage"); - const p = row.policy; - const tokenUsed = row.used_tokens + row.reserved_tokens; - const costUsed = row.used_cost_micros + row.reserved_cost_micros; - const tokenPct = p.max_tokens ? Math.min(100, Math.round((tokenUsed / p.max_tokens) * 100)) : 0; - const costPct = p.max_cost_micros ? Math.min(100, Math.round((costUsed / p.max_cost_micros) * 100)) : 0; - const isAgentBudget = p.source === "agent_budget_monthly_cents"; - return ( - - -
{p.model_id || p.provider_type || t("caps.tenantScope")}
-
- {p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")} - {isAgentBudget ? {t("caps.agentBudgetSource")} : null} -
- - {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.enabled ? t("caps.enabled") : t("caps.disabled")} - - - - - ); -} diff --git a/ui/web/src/pages/usage/hooks/use-usage-caps.ts b/ui/web/src/pages/usage/hooks/use-usage-caps.ts index 64c0e4c4..9d307e3b 100644 --- a/ui/web/src/pages/usage/hooks/use-usage-caps.ts +++ b/ui/web/src/pages/usage/hooks/use-usage-caps.ts @@ -12,8 +12,8 @@ export interface UsageCapPolicyInput { provider_type?: string; model_id?: string; window: UsageCapPolicy["window"]; - max_tokens?: number; - max_cost_usd?: number; + max_tokens?: number | null; + max_cost_usd?: number | null; enabled?: boolean; } @@ -81,6 +81,20 @@ export function useUsageCaps() { [http, refresh], ); + const updatePolicy = useCallback( + async (id: string, input: UsageCapPolicyInput) => { + try { + await http.patch(`/v1/usage-caps/policies/${id}`, input); + await refresh(); + toast.success(i18next.t("usage:caps.toast.updated")); + } catch (err) { + toast.error(i18next.t("usage:caps.toast.updateFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, refresh], + ); + return { policies: policiesQuery.data ?? [], utilization: utilizationQuery.data ?? [], @@ -89,6 +103,7 @@ export function useUsageCaps() { refreshing: policiesQuery.isFetching || utilizationQuery.isFetching || eventsQuery.isFetching, refresh, createPolicy, + updatePolicy, deletePolicy, }; } From d64a31ebdb598aab806e24e18a9a422f71ec5e0b Mon Sep 17 00:00:00 2001 From: Goon Date: Sun, 24 May 2026 10:42:26 +0700 Subject: [PATCH 05/10] fix(usage): enforce caps on auxiliary llm calls --- cmd/gateway.go | 11 +- cmd/gateway_consumer.go | 4 +- cmd/gateway_consumer_deps.go | 2 + cmd/gateway_consumer_normal.go | 6 +- cmd/gateway_deps.go | 8 +- cmd/gateway_hooks.go | 6 +- cmd/gateway_http_handlers.go | 7 +- cmd/gateway_http_wiring.go | 4 +- cmd/gateway_lifecycle.go | 2 +- cmd/gateway_managed.go | 7 +- cmd/gateway_methods.go | 4 +- internal/agent/intent_classify.go | 12 +- internal/agent/title_generate.go | 16 ++- internal/channels/history_compaction.go | 13 +- internal/channels/instance_loader.go | 11 +- internal/consolidation/dreaming_worker.go | 16 ++- internal/consolidation/episodic_worker.go | 18 ++- internal/consolidation/workers.go | 4 + internal/gateway/methods/chat.go | 11 +- internal/hooks/handlers/prompt.go | 34 ++++- internal/http/knowledge_graph.go | 10 +- internal/http/knowledge_graph_handlers.go | 17 ++- internal/http/pending_messages.go | 8 +- internal/http/provider_verify.go | 11 +- internal/http/providers.go | 6 + internal/http/summoner.go | 6 +- internal/http/summoner_regenerate.go | 10 +- internal/http/usage_caps.go | 53 ++++---- internal/i18n/catalog_en.go | 23 +++- internal/i18n/catalog_vi.go | 23 +++- internal/i18n/catalog_zh.go | 23 +++- internal/i18n/keys.go | 111 +++++++++------- internal/knowledgegraph/extractor.go | 20 ++- internal/store/pg/usage_caps.go | 9 +- internal/usage/caps/chat_call.go | 150 ++++++++++++++++++++++ internal/usage/caps/service.go | 7 + internal/usage/caps/service_test.go | 111 +++++++++++++++- internal/vault/enrich_classify.go | 9 ++ internal/vault/enrich_worker.go | 55 +++++--- 39 files changed, 701 insertions(+), 157 deletions(-) create mode 100644 internal/usage/caps/chat_call.go diff --git a/cmd/gateway.go b/cmd/gateway.go index 31428735..5a2e6006 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -223,6 +223,7 @@ func runGateway() { } } setupMemoryEmbeddings(pgStores, providerRegistry) + usageCapSvc := usagecaps.NewService(pgStores.UsageCaps, pgStores.Providers) // Resolve background provider for consolidation + vault enrichment. // Fallback: background.provider → agent.default_provider → first registered provider. @@ -234,6 +235,7 @@ func runGateway() { var kgExtractor *kg.Extractor if pgStores.KnowledgeGraph != nil { kgExtractor = kg.NewExtractor(bgProvider, bgModel, 0) + kgExtractor.SetUsageCapService(usageCapSvc) } cleanupConsolidation := consolidation.Register(consolidation.ConsolidationDeps{ EpisodicStore: pgStores.Episodic, @@ -245,6 +247,7 @@ func runGateway() { Registry: providerRegistry, Extractor: kgExtractor, AlertDeps: bgalert.AlertDeps{SystemConfigs: pgStores.SystemConfigs, MsgBus: msgBus}, + UsageCaps: usageCapSvc, AgentStore: pgStores.Agents, }) defer cleanupConsolidation() @@ -267,6 +270,7 @@ func runGateway() { MsgBus: msgBus, TeamStore: pgStores.Teams, AlertDeps: bgalert.AlertDeps{SystemConfigs: pgStores.SystemConfigs, MsgBus: msgBus}, + UsageCaps: usageCapSvc, }) enrichProgress = ep enrichWorker = ew @@ -283,7 +287,6 @@ func runGateway() { slog.Info("bootstrap: capabilities backfill complete", "agents", count) } - usageCapSvc := usagecaps.NewService(pgStores.UsageCaps, pgStores.Providers) if readImage, ok := toolsReg.Get("read_image"); ok { if t, ok := readImage.(*tools.ReadImageTool); ok { t.SetUsageCapService(usageCapSvc) @@ -364,6 +367,7 @@ func runGateway() { workspace: workspace, dataDir: dataDir, domainBus: domainBus, + usageCapSvc: usageCapSvc, audioMgr: audioMgr, } @@ -376,7 +380,7 @@ func runGateway() { httpapi.InitGatewayNoAuthFallbackAllowed(config.GatewayNoAuthFallbackAllowed(cfg.Gateway)) exportTokenStore := httpapi.InitExportTokenStore() defer exportTokenStore.Stop() - agentsH, skillsH, tracesH, mcpH, channelInstancesH, providersH, builtinToolsH, pendingMessagesH, teamEventsH, secureCLIH, secureCLIGrantH, mcpUserCredsH := wireHTTP(pgStores, cfg.Agents.Defaults.Workspace, dataDir, bundledSkillsDir, msgBus, toolsReg, providerRegistry, modelReg, permPE.IsOwner, gatewayAddr, mcpToolLister) + agentsH, skillsH, tracesH, mcpH, channelInstancesH, providersH, builtinToolsH, pendingMessagesH, teamEventsH, secureCLIH, secureCLIGrantH, mcpUserCredsH := wireHTTP(pgStores, cfg.Agents.Defaults.Workspace, dataDir, bundledSkillsDir, msgBus, toolsReg, providerRegistry, modelReg, permPE.IsOwner, gatewayAddr, mcpToolLister, usageCapSvc) // Wire dependencies for system prompt preview parity. if agentsH != nil { @@ -433,7 +437,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.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) + pairingMethods, heartbeatMethods, chatMethods, cfgPermsMethods := registerAllMethods(server, agentRouter, pgStores.Sessions, 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) // Phase 3: Agent hooks RPC methods (hooks.list/create/update/delete/toggle/test/history). if hs, ok := pgStores.Hooks.(hooks.HookStore); ok && hs != nil { @@ -519,6 +523,7 @@ func runGateway() { instanceLoader = channels.NewInstanceLoader(pgStores.ChannelInstances, pgStores.Agents, channelMgr, msgBus, pgStores.Pairing) instanceLoader.SetProviderRegistry(providerRegistry) instanceLoader.SetPendingCompactionConfig(cfg.Channels.PendingCompaction) + instanceLoader.SetUsageCapService(usageCapSvc) instanceLoader.RegisterFactory(channels.TypeTelegram, telegram.FactoryWithStoresAndAudio(pgStores.Agents, pgStores.ConfigPermissions, pgStores.Teams, pgStores.SubagentTasks, pgStores.PendingMessages, audioMgr)) instanceLoader.RegisterFactory(channels.TypeDiscord, discord.FactoryWithStoresAndAudio(pgStores.Agents, pgStores.ConfigPermissions, pgStores.PendingMessages, audioMgr)) instanceLoader.RegisterFactory(channels.TypeFeishu, feishu.FactoryWithPendingStoreAndAudio(pgStores.PendingMessages, audioMgr)) diff --git a/cmd/gateway_consumer.go b/cmd/gateway_consumer.go index cbca58e3..3bd7e78a 100644 --- a/cmd/gateway_consumer.go +++ b/cmd/gateway_consumer.go @@ -19,6 +19,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/scheduler" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -26,7 +27,7 @@ import ( // and routes them through the scheduler/agent loop, then publishes the response back. // Also handles subagent announcements: routes them through the parent agent's session // (matching TS subagent-announce.ts pattern) so the agent can reformulate for the user. -func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents *agent.Router, cfg *config.Config, sched *scheduler.Scheduler, channelMgr *channels.Manager, teamStore store.TeamStore, quotaChecker *channels.QuotaChecker, sessStore store.SessionStore, agentStore store.AgentStore, contactCollector *store.ContactCollector, postTurn tools.PostTurnProcessor, subagentMgr *tools.SubagentManager) { +func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents *agent.Router, cfg *config.Config, sched *scheduler.Scheduler, channelMgr *channels.Manager, teamStore store.TeamStore, quotaChecker *channels.QuotaChecker, sessStore store.SessionStore, agentStore store.AgentStore, contactCollector *store.ContactCollector, postTurn tools.PostTurnProcessor, subagentMgr *tools.SubagentManager, usageCapSvc *usagecaps.Service) { slog.Info("inbound message consumer started") // Inbound message deduplication (matching TS src/infra/dedupe.ts + inbound-dedupe.ts). @@ -58,6 +59,7 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents QuotaChecker: quotaChecker, ContactCollector: contactCollector, SubagentMgr: subagentMgr, + UsageCaps: usageCapSvc, GetAnnounceMu: getAnnounceMu, } diff --git a/cmd/gateway_consumer_deps.go b/cmd/gateway_consumer_deps.go index faf3755c..11174f2a 100644 --- a/cmd/gateway_consumer_deps.go +++ b/cmd/gateway_consumer_deps.go @@ -10,6 +10,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/scheduler" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ConsumerDeps bundles shared dependencies for consumer message handlers. @@ -28,6 +29,7 @@ type ConsumerDeps struct { ContactCollector *store.ContactCollector TaskRunSessions sync.Map SubagentMgr *tools.SubagentManager + UsageCaps *usagecaps.Service BgWg sync.WaitGroup GetAnnounceMu func(string) *sync.Mutex } diff --git a/cmd/gateway_consumer_normal.go b/cmd/gateway_consumer_normal.go index 0b243cf8..a719c446 100644 --- a/cmd/gateway_consumer_normal.go +++ b/cmd/gateway_consumer_normal.go @@ -333,7 +333,11 @@ func processNormalMessage( if locale == "" { locale = "en" } - intent := agent.ClassifyIntent(ctx, loop.Provider(), loop.Model(), msg.Content) + classifyCtx := ctx + if uid := loop.UUID(); uid != uuid.Nil { + classifyCtx = store.WithAgentID(classifyCtx, uid) + } + intent := agent.ClassifyIntentWithUsageCaps(classifyCtx, deps.UsageCaps, loop.Provider(), loop.Model(), msg.Content) switch intent { case agent.IntentStatusQuery: status := deps.Agents.GetActivity(sessionKey) diff --git a/cmd/gateway_deps.go b/cmd/gateway_deps.go index 0487f358..c30ca4f2 100644 --- a/cmd/gateway_deps.go +++ b/cmd/gateway_deps.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/internal/vault" ) @@ -28,13 +29,14 @@ type gatewayDeps struct { channelMgr *channels.Manager agentRouter *agent.Router toolsReg *tools.Registry - skillsLoader *skills.Loader // optional: enables skill creation in evolution approval + skillsLoader *skills.Loader // optional: enables skill creation in evolution approval permCache *cache.PermissionCache // nil if no tenant store; closed on shutdown to stop sweep goroutines - enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered - enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue + enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered + enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue workspace string dataDir string domainBus eventbus.DomainEventBus + usageCapSvc *usagecaps.Service audioMgr *audio.Manager // nil if TTS not configured; used by TTSHandler ttsHandler *httpapi.TTSHandler // nil if TTS not configured; for hot-reload } diff --git a/cmd/gateway_hooks.go b/cmd/gateway_hooks.go index 26b35887..29340bc5 100644 --- a/cmd/gateway_hooks.go +++ b/cmd/gateway_hooks.go @@ -7,12 +7,13 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/edition" "github.com/nextlevelbuilder/goclaw/internal/hooks" - hookhandlers "github.com/nextlevelbuilder/goclaw/internal/hooks/handlers" "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" + hookhandlers "github.com/nextlevelbuilder/goclaw/internal/hooks/handlers" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/security" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/store/pg" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // sharedHookHandlers is populated by wireExtras so the gateway.go router @@ -27,7 +28,7 @@ var sharedHookHandlers map[hooks.HandlerType]hooks.Handler // Budget wiring (C1 fix): the PromptHandler receives a budget.Store bound // to pg.NewPGHookBudget so token spend is atomically deducted per tenant. // When the DB handle is unavailable, budget falls back to nil (Lite desktop). -func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, hooksCfg config.HooksConfig) map[hooks.HandlerType]hooks.Handler { +func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, hooksCfg config.HooksConfig, usageCapSvc *usagecaps.Service) map[hooks.HandlerType]hooks.Handler { encryptKey := os.Getenv("GOCLAW_ENCRYPTION_KEY") var budgetStore *budget.Store @@ -38,6 +39,7 @@ func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, ho promptHandler := &hookhandlers.PromptHandler{ Resolver: hookhandlers.NewRegistryResolver(providerReg, stores.SystemConfigs), Budget: budgetStore, + UsageCaps: usageCapSvc, DefaultModel: "haiku", } diff --git a/cmd/gateway_http_handlers.go b/cmd/gateway_http_handlers.go index 5ad49409..ce8586ca 100644 --- a/cmd/gateway_http_handlers.go +++ b/cmd/gateway_http_handlers.go @@ -6,10 +6,11 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // wireHTTP creates HTTP handlers (agents + skills + traces + MCP + channel instances + providers + builtin tools + pending messages). -func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir string, msgBus *bus.MessageBus, toolsReg *tools.Registry, providerReg *providers.Registry, modelReg providers.ModelRegistry, isOwner func(string) bool, gatewayAddr string, mcpToolLister httpapi.MCPToolLister) (*httpapi.AgentsHandler, *httpapi.SkillsHandler, *httpapi.TracesHandler, *httpapi.MCPHandler, *httpapi.ChannelInstancesHandler, *httpapi.ProvidersHandler, *httpapi.BuiltinToolsHandler, *httpapi.PendingMessagesHandler, *httpapi.TeamEventsHandler, *httpapi.SecureCLIHandler, *httpapi.SecureCLIGrantHandler, *httpapi.MCPUserCredentialsHandler) { +func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir string, msgBus *bus.MessageBus, toolsReg *tools.Registry, providerReg *providers.Registry, modelReg providers.ModelRegistry, isOwner func(string) bool, gatewayAddr string, mcpToolLister httpapi.MCPToolLister, usageCapSvc *usagecaps.Service) (*httpapi.AgentsHandler, *httpapi.SkillsHandler, *httpapi.TracesHandler, *httpapi.MCPHandler, *httpapi.ChannelInstancesHandler, *httpapi.ProvidersHandler, *httpapi.BuiltinToolsHandler, *httpapi.PendingMessagesHandler, *httpapi.TeamEventsHandler, *httpapi.SecureCLIHandler, *httpapi.SecureCLIGrantHandler, *httpapi.MCPUserCredentialsHandler) { var agentsH *httpapi.AgentsHandler var skillsH *httpapi.SkillsHandler var tracesH *httpapi.TracesHandler @@ -24,7 +25,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir if stores != nil && stores.Agents != nil { var summoner *httpapi.AgentSummoner if providerReg != nil { - summoner = httpapi.NewAgentSummoner(stores.Agents, providerReg, msgBus) + summoner = httpapi.NewAgentSummoner(stores.Agents, providerReg, msgBus, usageCapSvc) } agentsH = httpapi.NewAgentsHandler(stores.Agents, stores.Providers, providerReg, stores.DB, stores.Tracing, defaultWorkspace, msgBus, summoner, isOwner) agentsH.SetImportStores(stores.Memory, stores.KnowledgeGraph) @@ -61,6 +62,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir if stores != nil && stores.Providers != nil { providersH = httpapi.NewProvidersHandler(stores.Providers, stores.ConfigSecrets, providerReg, gatewayAddr) providersH.SetMessageBus(msgBus) + providersH.SetUsageCapService(usageCapSvc) if modelReg != nil { providersH.SetModelRegistry(modelReg) } @@ -90,6 +92,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir if stores != nil && stores.PendingMessages != nil { pendingMessagesH = httpapi.NewPendingMessagesHandler(stores.PendingMessages, stores.Agents, providerReg) + pendingMessagesH.SetUsageCapService(usageCapSvc) } if stores != nil && stores.SecureCLI != nil { diff --git a/cmd/gateway_http_wiring.go b/cmd/gateway_http_wiring.go index 76c782a3..3eac9db3 100644 --- a/cmd/gateway_http_wiring.go +++ b/cmd/gateway_http_wiring.go @@ -234,7 +234,9 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer( // Knowledge graph API if d.pgStores != nil && d.pgStores.KnowledgeGraph != nil { - d.server.SetKnowledgeGraphHandler(httpapi.NewKnowledgeGraphHandler(d.pgStores.KnowledgeGraph, d.providerRegistry)) + kgHandler := httpapi.NewKnowledgeGraphHandler(d.pgStores.KnowledgeGraph, d.providerRegistry) + kgHandler.SetUsageCapService(d.usageCapSvc) + d.server.SetKnowledgeGraphHandler(kgHandler) } // V3: Evolution metrics + suggestions API diff --git a/cmd/gateway_lifecycle.go b/cmd/gateway_lifecycle.go index 3a8ef20a..631bdd4e 100644 --- a/cmd/gateway_lifecycle.go +++ b/cmd/gateway_lifecycle.go @@ -140,7 +140,7 @@ func (d *gatewayDeps) runLifecycle( d.channelMgr.SetContactCollector(contactCollector) } - go consumeInboundMessages(ctx, d.msgBus, d.agentRouter, d.cfg, deps.sched, d.channelMgr, deps.consumerTeamStore, deps.quotaChecker, d.pgStores.Sessions, d.pgStores.Agents, contactCollector, deps.postTurn, deps.subagentMgr) + go consumeInboundMessages(ctx, d.msgBus, d.agentRouter, d.cfg, deps.sched, d.channelMgr, deps.consumerTeamStore, deps.quotaChecker, d.pgStores.Sessions, d.pgStores.Agents, contactCollector, deps.postTurn, deps.subagentMgr, d.usageCapSvc) // Webhook callback worker — delivers async webhook_calls rows to receiver callback_url. // Runs in both editions: Standard (PG, concurrency=4) and Lite (SQLite, concurrency=1). diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index 0db67050..798bff9f 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -185,7 +185,7 @@ func wireExtras( "disabled_count", n, "edition", edition.Current().Name) } - handlers := buildHookHandlers(stores, providerReg, appCfg.Hooks) + handlers := buildHookHandlers(stores, providerReg, appCfg.Hooks, usageCapSvc) stdOpts := hooks.StdDispatcherOpts{ Store: hs, Audit: hooks.NewAuditWriter(hs, ""), @@ -304,7 +304,7 @@ func wireExtras( writeMemIntc = tools.NewMemoryInterceptor(stores.Memory, workspace) // Hook KG extraction on memory writes if KG store is available if stores.KnowledgeGraph != nil && stores.BuiltinTools != nil { - writeMemIntc.SetKGExtractFunc(buildKGExtractFunc(stores.KnowledgeGraph, stores.BuiltinTools, providerReg)) + writeMemIntc.SetKGExtractFunc(buildKGExtractFunc(stores.KnowledgeGraph, stores.BuiltinTools, providerReg, usageCapSvc)) } } if readTool, ok := toolsReg.Get("read_file"); ok { @@ -712,7 +712,7 @@ type kgSettings struct { // buildKGExtractFunc returns a callback that extracts entities from memory content. // Settings are read from the builtin_tools table on each invocation (not cached), // so changes take effect immediately without restart. -func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinToolStore, providerReg *providers.Registry) tools.KGExtractFunc { +func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinToolStore, providerReg *providers.Registry, usageCapSvc *usagecaps.Service) tools.KGExtractFunc { return func(ctx context.Context, agentID, userID, content string) { slog.Info("kg extract: triggered", "agent", agentID, "user", userID, "content_len", len(content)) // Read settings from DB on each call so admin changes take effect immediately @@ -736,6 +736,7 @@ func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinTool return } extractor := kg.NewExtractor(p, settings.ExtractionModel, settings.MinConfidence) + extractor.SetUsageCapService(usageCapSvc) result, err := extractor.Extract(ctx, content) if err != nil { slog.Warn("kg extract: extraction failed", "agent", agentID, "error", err) diff --git a/cmd/gateway_methods.go b/cmd/gateway_methods.go index d994217e..7d7f0664 100644 --- a/cmd/gateway_methods.go +++ b/cmd/gateway_methods.go @@ -12,14 +12,16 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/gateway/methods" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) -func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, 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) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) { +func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, 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) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) { router := server.Router() // Phase 1: Core methods chatMethods := methods.NewChatMethods(agents, sessStore, cfg, server.RateLimiter(), msgBus) chatMethods.SetAudioManager(audioMgr) // Wire TTS auto-apply for WS responses + chatMethods.SetUsageCapService(usageCapSvc) chatMethods.Register(router) methods.NewAgentsMethods(agents, cfg, cfgPath, workspace, agentStore, contextFileInterceptor, msgBus).Register(router) methods.NewSessionsMethods(sessStore, msgBus, cfg).Register(router) diff --git a/internal/agent/intent_classify.go b/internal/agent/intent_classify.go index 2fd4e660..46201708 100644 --- a/internal/agent/intent_classify.go +++ b/internal/agent/intent_classify.go @@ -8,6 +8,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // IntentType represents the classified intent of a user message. @@ -108,6 +109,10 @@ func containsWholeWord(s, kw string) bool { // Uses keyword fast-path first, then falls back to LLM classification. // Falls back to IntentNewTask on any error. func ClassifyIntent(ctx context.Context, provider providers.Provider, model, userMessage string) IntentType { + return ClassifyIntentWithUsageCaps(ctx, nil, provider, model, userMessage) +} + +func ClassifyIntentWithUsageCaps(ctx context.Context, usageCaps *usagecaps.Service, provider providers.Provider, model, userMessage string) IntentType { // Fast-path: keyword matching for obvious patterns (no LLM cost). if intent, ok := quickClassify(userMessage); ok { return intent @@ -116,7 +121,7 @@ func ClassifyIntent(ctx context.Context, provider providers.Provider, model, use ctx, cancel := context.WithTimeout(ctx, intentClassifyTimeout) defer cancel() - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: intentSystemPrompt}, {Role: "user", Content: userMessage}, @@ -126,6 +131,11 @@ func ClassifyIntent(ctx context.Context, provider providers.Provider, model, use providers.OptMaxTokens: 20, providers.OptTemperature: 0.0, }, + } + resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "intent-classify", + MaxOutputTokens: 20, }) if err != nil { return IntentNewTask diff --git a/internal/agent/title_generate.go b/internal/agent/title_generate.go index eab223cf..67fcd302 100644 --- a/internal/agent/title_generate.go +++ b/internal/agent/title_generate.go @@ -7,6 +7,7 @@ import ( "time" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) const titleGenerateTimeout = 120 * time.Second @@ -16,10 +17,14 @@ const titleSystemPrompt = `Generate a short title (max 15 words) for this conver // GenerateTitle uses a lightweight LLM call to create a short conversation title // from the user's first message. Returns empty string on error. func GenerateTitle(ctx context.Context, provider providers.Provider, model, userMessage string) string { + return GenerateTitleWithUsageCaps(ctx, nil, provider, model, userMessage) +} + +func GenerateTitleWithUsageCaps(ctx context.Context, usageCaps *usagecaps.Service, provider providers.Provider, model, userMessage string) string { ctx, cancel := context.WithTimeout(ctx, titleGenerateTimeout) defer cancel() - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: titleSystemPrompt}, {Role: "user", Content: userMessage}, @@ -29,13 +34,18 @@ func GenerateTitle(ctx context.Context, provider providers.Provider, model, user // Larger budget: thinking-capable models (Gemini 2.5/3, GPT-5 reasoning) // can consume output tokens on reasoning traces. 256 leaves room for a // 15-word title even when the provider allocates some budget to thinking. - providers.OptMaxTokens: 256, - providers.OptTemperature: 0.3, + providers.OptMaxTokens: 256, + providers.OptTemperature: 0.3, // Disable extended thinking for title generation — it's a trivial task // that doesn't benefit from reasoning and defaults (esp. Gemini's "high") // otherwise eat the entire max_tokens budget, truncating the title to 1 word. providers.OptThinkingLevel: "off", }, + } + resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "session-title", + MaxOutputTokens: 256, }) if err != nil { slog.Warn("title generation failed", "error", err) diff --git a/internal/channels/history_compaction.go b/internal/channels/history_compaction.go index 6480a6ad..79b20624 100644 --- a/internal/channels/history_compaction.go +++ b/internal/channels/history_compaction.go @@ -11,6 +11,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // CompactionConfig configures LLM-based history compaction. @@ -20,6 +21,7 @@ type CompactionConfig struct { MaxTokens int // max output tokens for summarization (default 4096) Provider providers.Provider // LLM provider for summarization Model string // model to use for summarization + UsageCaps *usagecaps.Service } // MaybeCompact checks if compaction is needed for a history key and triggers it in background. @@ -55,7 +57,7 @@ func (ph *PendingHistory) MaybeCompact(historyKey string, currentCount int, cfg // CompactGroup performs LLM-based compaction on a pending message group. // Reused by both auto-compact (channel) and HTTP compact endpoint. // Returns the number of entries remaining after compaction. -func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, historyKey string, provider providers.Provider, model string, keepRecent, maxTokens int) (int, error) { +func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, historyKey string, provider providers.Provider, model string, keepRecent, maxTokens int, usageCaps *usagecaps.Service) (int, error) { if keepRecent <= 0 { keepRecent = 40 } @@ -94,13 +96,18 @@ func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, fmt.Fprintf(&sb, "%s%s: %s\n", prefix, ts, e.Body) } - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{{ Role: "user", Content: "Summarize these group chat messages concisely, preserving key topics, decisions, names, and important context:\n\n" + sb.String(), }}, Model: model, Options: map[string]any{"max_tokens": maxTokens, "temperature": 0.3}, + } + resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "pending-history-compaction", + MaxOutputTokens: maxTokens, }) if err != nil { return 0, fmt.Errorf("llm summarize: %w", err) @@ -185,7 +192,7 @@ func (ph *PendingHistory) runCompaction(historyKey string, cfg *CompactionConfig return } - _, err = CompactGroup(ctx, ph.store, ph.channelName, historyKey, cfg.Provider, cfg.Model, cfg.KeepRecent, cfg.MaxTokens) + _, err = CompactGroup(ctx, ph.store, ph.channelName, historyKey, cfg.Provider, cfg.Model, cfg.KeepRecent, cfg.MaxTokens, cfg.UsageCaps) if err != nil { slog.Warn("compaction.failed", "channel", ph.channelName, "key", historyKey, "error", err) return diff --git a/internal/channels/instance_loader.go b/internal/channels/instance_loader.go index df6d677f..7f845d80 100644 --- a/internal/channels/instance_loader.go +++ b/internal/channels/instance_loader.go @@ -16,6 +16,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providerresolve" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // reloadStartTimeout bounds how long Reload() will wait for a single channel's @@ -39,6 +40,7 @@ type InstanceLoader struct { agentStore store.AgentStore providerReg *providers.Registry pendingCompactCfg *config.PendingCompactionConfig + usageCaps *usagecaps.Service factories map[string]ChannelFactory manager *Manager msgBus *bus.MessageBus @@ -78,6 +80,10 @@ func (l *InstanceLoader) SetPendingCompactionConfig(cfg *config.PendingCompactio l.pendingCompactCfg = cfg } +func (l *InstanceLoader) SetUsageCapService(s *usagecaps.Service) { + l.usageCaps = s +} + // RegisterFactory registers a factory for a channel type (e.g., "telegram", "discord"). func (l *InstanceLoader) RegisterFactory(channelType string, factory ChannelFactory) { l.factories[channelType] = factory @@ -306,8 +312,9 @@ func (l *InstanceLoader) loadInstance(ctx context.Context, inst store.ChannelIns if p != nil && model != "" { cc := &CompactionConfig{ - Provider: p, - Model: model, + Provider: p, + Model: model, + UsageCaps: l.usageCaps, } if l.pendingCompactCfg != nil { cc.Threshold = l.pendingCompactCfg.Threshold diff --git a/internal/consolidation/dreaming_worker.go b/internal/consolidation/dreaming_worker.go index 4a97bf76..4a981793 100644 --- a/internal/consolidation/dreaming_worker.go +++ b/internal/consolidation/dreaming_worker.go @@ -11,9 +11,10 @@ import ( "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/bgalert" "github.com/nextlevelbuilder/goclaw/internal/eventbus" - "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/providerresolve" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) const ( @@ -31,6 +32,7 @@ type dreamingWorker struct { systemConfigs store.SystemConfigStore // per-tenant provider config registry *providers.Registry // provider resolution alertDeps bgalert.AlertDeps + usageCaps *usagecaps.Service // threshold/debounce are the global defaults. Per-agent overrides come // from resolveConfig which reads the agent's MemoryConfig.Dreaming JSONB. @@ -95,6 +97,11 @@ func (w *dreamingWorker) Handle(ctx context.Context, event eventbus.DomainEvent) ctx = store.WithTenantID(ctx, tid) } } + if event.AgentID != "" { + if aid, err := uuid.Parse(event.AgentID); err == nil { + ctx = store.WithAgentID(ctx, aid) + } + } agentID := event.AgentID userID := event.UserID @@ -205,7 +212,7 @@ func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Prov } body := strings.Join(summaries, "\n---\n") - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: dreamingSystemPrompt}, {Role: "user", Content: "Session summaries:\n---\n" + body + "\n---"}, @@ -214,6 +221,11 @@ func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Prov Options: map[string]any{ providers.OptMaxTokens: dreamingMaxTokens, }, + } + resp, err := w.usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "dreaming-synthesis", + MaxOutputTokens: dreamingMaxTokens, }) if err != nil { return "", fmt.Errorf("dreaming chat: %w", err) diff --git a/internal/consolidation/episodic_worker.go b/internal/consolidation/episodic_worker.go index 8a5f8cf5..b93a5383 100644 --- a/internal/consolidation/episodic_worker.go +++ b/internal/consolidation/episodic_worker.go @@ -10,19 +10,21 @@ import ( "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/bgalert" "github.com/nextlevelbuilder/goclaw/internal/eventbus" - "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/providerresolve" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // episodicWorker handles session.completed events → creates episodic summaries. type episodicWorker struct { store store.EpisodicStore - sessions store.SessionCoreStore // for reading session messages during summarization - systemConfigs store.SystemConfigStore // per-tenant provider config - registry *providers.Registry // provider resolution + sessions store.SessionCoreStore // for reading session messages during summarization + systemConfigs store.SystemConfigStore // per-tenant provider config + registry *providers.Registry // provider resolution eventBus eventbus.DomainEventBus alertDeps bgalert.AlertDeps + usageCaps *usagecaps.Service } // resolveProvider delegates to shared background provider resolution. @@ -53,6 +55,7 @@ func (w *episodicWorker) Handle(ctx context.Context, event eventbus.DomainEvent) if err != nil { return fmt.Errorf("episodic: invalid agent_id %q: %w", event.AgentID, err) } + ctx = store.WithAgentID(ctx, agentUUID) // Build source_id for idempotency sourceID := fmt.Sprintf("%s:%d", payload.SessionKey, payload.CompactionCount) @@ -168,13 +171,18 @@ func (w *episodicWorker) summarizeFromMessages(ctx context.Context, provider pro sctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - resp, err := provider.Chat(sctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: summarizationPrompt}, {Role: "user", Content: sb.String()}, }, Model: model, Options: map[string]any{"max_tokens": 1024, "temperature": 0.3}, + } + resp, err := w.usageCaps.Chat(sctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "episodic-summary", + MaxOutputTokens: 1024, }) if err != nil { return "", err diff --git a/internal/consolidation/workers.go b/internal/consolidation/workers.go index 1c5fb82a..2c3be24e 100644 --- a/internal/consolidation/workers.go +++ b/internal/consolidation/workers.go @@ -13,6 +13,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/eventbus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ConsolidationDeps bundles all dependencies for the consolidation pipeline. @@ -26,6 +27,7 @@ type ConsolidationDeps struct { Registry *providers.Registry // provider resolution Extractor EntityExtractor AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + UsageCaps *usagecaps.Service // AgentStore is optional: when present, the dreaming worker reads // per-agent overrides from MemoryConfig.Dreaming. If nil, the worker // uses its built-in defaults for every agent. @@ -42,6 +44,7 @@ func Register(deps ConsolidationDeps) func() { registry: deps.Registry, eventBus: deps.EventBus, alertDeps: deps.AlertDeps, + usageCaps: deps.UsageCaps, } semantic := &semanticWorker{ kgStore: deps.KGStore, @@ -59,6 +62,7 @@ func Register(deps ConsolidationDeps) func() { systemConfigs: deps.SystemConfigs, registry: deps.Registry, alertDeps: deps.AlertDeps, + usageCaps: deps.UsageCaps, threshold: dreamingDefaultThreshold, debounce: dreamingDefaultDebounce, resolveConfig: newAgentStoreResolver(deps.AgentStore), diff --git a/internal/gateway/methods/chat.go b/internal/gateway/methods/chat.go index 1281a9c5..aa0e98cc 100644 --- a/internal/gateway/methods/chat.go +++ b/internal/gateway/methods/chat.go @@ -20,6 +20,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/sessions" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -32,6 +33,7 @@ type ChatMethods struct { eventBus bus.EventPublisher postTurn tools.PostTurnProcessor audioMgr *audio.Manager // for TTS auto-apply on WS responses (nil = disabled) + usageCaps *usagecaps.Service debouncer *chatDebouncer } @@ -46,6 +48,10 @@ func (m *ChatMethods) SetAudioManager(mgr *audio.Manager) { m.audioMgr = mgr } +func (m *ChatMethods) SetUsageCapService(s *usagecaps.Service) { + m.usageCaps = s +} + // SetPostTurnProcessor sets the post-turn processor for team task dispatch. func (m *ChatMethods) SetPostTurnProcessor(pt tools.PostTurnProcessor) { m.postTurn = pt @@ -352,7 +358,10 @@ func (m *ChatMethods) dispatchChatSends(requests []chatSendRequest) { // Use runCtxBase (WithoutCancel + tenant-aware) so title save uses correct tenant. titleCtx := runCtxBase go func() { - title := agent.GenerateTitle(titleCtx, agentProvider, agentModel, userMsg) + if uid := loop.UUID(); uid != uuid.Nil { + titleCtx = store.WithAgentID(titleCtx, uid) + } + title := agent.GenerateTitleWithUsageCaps(titleCtx, m.usageCaps, agentProvider, agentModel, userMsg) if title == "" { return } diff --git a/internal/hooks/handlers/prompt.go b/internal/hooks/handlers/prompt.go index 2ec9fb1f..062c3c1f 100644 --- a/internal/hooks/handlers/prompt.go +++ b/internal/hooks/handlers/prompt.go @@ -17,6 +17,8 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/hooks" "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ── Public surface ────────────────────────────────────────────────────────── @@ -46,6 +48,9 @@ type PromptHandler struct { // budget checks are skipped (Lite edition behavior). Budget *budget.Store + // UsageCaps enforces provider/model/tenant cost caps before the hook LLM call. + UsageCaps *usagecaps.Service + // DefaultModel is used when a hook config does not specify one. // Recommended: "haiku" for cheap evaluation. DefaultModel string @@ -155,7 +160,17 @@ func (h *PromptHandler) Execute(ctx context.Context, cfg hooks.HookConfig, ev ho req := h.buildChatRequest(cfg, ev, resolvedModel) // 6. Call provider. - resp, err := provider.Chat(ctx, req) + callCtx := ctx + if ev.TenantID != uuid.Nil { + callCtx = store.WithTenantID(callCtx, ev.TenantID) + } + resp, err := h.UsageCaps.Chat(callCtx, provider, req, usagecaps.ChatOptions{ + TenantID: ev.TenantID, + ProviderName: provider.Name(), + ModelID: resolvedModel, + Purpose: "hook-prompt", + MaxOutputTokens: promptMaxTokens(req), + }) if err != nil { // Fail-closed on transport/provider error for blocking events. if ev.HookEvent.IsBlocking() { @@ -300,6 +315,23 @@ func (h *PromptHandler) buildChatRequest(cfg hooks.HookConfig, ev hooks.Event, m } } +func promptMaxTokens(req providers.ChatRequest) int { + if req.Options == nil { + return 512 + } + if v, ok := req.Options[providers.OptMaxTokens]; ok { + switch n := v.(type) { + case int: + return n + case int64: + return int(n) + case float64: + return int(n) + } + } + return 512 +} + // sanitizeToolInput canonicalizes a tool_input map into stable JSON with // sorted keys, stripping only structural noise. Actual injection-attack // detection is delegated to the evaluator LLM (which has the anti-injection diff --git a/internal/http/knowledge_graph.go b/internal/http/knowledge_graph.go index 7333f851..fcc61b90 100644 --- a/internal/http/knowledge_graph.go +++ b/internal/http/knowledge_graph.go @@ -7,12 +7,14 @@ import ( kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // KnowledgeGraphHandler handles KG entity/relation management endpoints. type KnowledgeGraphHandler struct { store store.KnowledgeGraphStore providerReg *providers.Registry + usageCaps *usagecaps.Service } // NewKnowledgeGraphHandler creates a handler for KG management endpoints. @@ -20,6 +22,10 @@ func NewKnowledgeGraphHandler(s store.KnowledgeGraphStore, providerReg *provider return &KnowledgeGraphHandler{store: s, providerReg: providerReg} } +func (h *KnowledgeGraphHandler) SetUsageCapService(s *usagecaps.Service) { + h.usageCaps = s +} + // NewExtractor creates an Extractor from the given provider name and model. func (h *KnowledgeGraphHandler) NewExtractor(ctx context.Context, providerName, model string, minConfidence float64) *kg.Extractor { if h.providerReg == nil || providerName == "" || model == "" { @@ -29,7 +35,9 @@ func (h *KnowledgeGraphHandler) NewExtractor(ctx context.Context, providerName, if err != nil { return nil } - return kg.NewExtractor(p, model, minConfidence) + extractor := kg.NewExtractor(p, model, minConfidence) + extractor.SetUsageCapService(h.usageCaps) + return extractor } // RegisterRoutes registers all KG routes on the given mux. diff --git a/internal/http/knowledge_graph_handlers.go b/internal/http/knowledge_graph_handlers.go index 80cb420f..8e31b319 100644 --- a/internal/http/knowledge_graph_handlers.go +++ b/internal/http/knowledge_graph_handlers.go @@ -5,6 +5,7 @@ import ( "net/http" "strconv" + "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -156,6 +157,10 @@ func (h *KnowledgeGraphHandler) handleTraverse(w http.ResponseWriter, r *http.Re func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Request) { locale := extractLocale(r) agentID := r.PathValue("agentID") + callCtx := r.Context() + if parsedAgentID, err := uuid.Parse(agentID); err == nil { + callCtx = store.WithAgentID(callCtx, parsedAgentID) + } var body struct { Text string `json:"text"` @@ -176,13 +181,13 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req return } - extractor := h.NewExtractor(r.Context(), body.Provider, body.Model, body.MinConf) + extractor := h.NewExtractor(callCtx, body.Provider, body.Model, body.MinConf) if extractor == nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidProviderOrModel)}) return } - result, err := extractor.Extract(r.Context(), body.Text) + result, err := extractor.Extract(callCtx, body.Text) if err != nil { slog.Warn("kg.extract failed", "error", err) writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) @@ -213,10 +218,10 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req } writeJSON(w, http.StatusOK, map[string]any{ - "entities": len(result.Entities), - "relations": len(result.Relations), - "dedup_merged": dedupMerged, - "dedup_flagged": dedupFlagged, + "entities": len(result.Entities), + "relations": len(result.Relations), + "dedup_merged": dedupMerged, + "dedup_flagged": dedupFlagged, }) } diff --git a/internal/http/pending_messages.go b/internal/http/pending_messages.go index 6bc7d0a6..3f144fa7 100644 --- a/internal/http/pending_messages.go +++ b/internal/http/pending_messages.go @@ -11,6 +11,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providerresolve" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // PendingMessagesHandler handles pending message HTTP endpoints. @@ -22,6 +23,7 @@ type PendingMessagesHandler struct { maxTokens int // max output tokens for LLM summarization (0 = use default) cfgProvider string // config-level provider override (empty = resolve from agent) cfgModel string // config-level model override (empty = resolve from agent) + usageCaps *usagecaps.Service } func NewPendingMessagesHandler(s store.PendingMessageStore, agentStore store.AgentStore, providerReg *providers.Registry) *PendingMessagesHandler { @@ -40,6 +42,10 @@ func (h *PendingMessagesHandler) SetProviderModel(provider, model string) { h.cfgModel = model } +func (h *PendingMessagesHandler) SetUsageCapService(s *usagecaps.Service) { + h.usageCaps = s +} + func (h *PendingMessagesHandler) RegisterRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /v1/pending-messages", h.authMiddleware(h.handleListGroups)) mux.HandleFunc("GET /v1/pending-messages/messages", h.authMiddleware(h.handleListMessages)) @@ -148,7 +154,7 @@ func (h *PendingMessagesHandler) handleCompact(w http.ResponseWriter, r *http.Re go func() { ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 180*time.Second) defer cancel() - remaining, err := channels.CompactGroup(ctx, h.store, req.ChannelName, req.HistoryKey, provider, model, keepRecent, h.maxTokens) + remaining, err := channels.CompactGroup(ctx, h.store, req.ChannelName, req.HistoryKey, provider, model, keepRecent, h.maxTokens, h.usageCaps) if err != nil { slog.Warn("compact.failed", "channel", req.ChannelName, "key", req.HistoryKey, "error", err) } else { diff --git a/internal/http/provider_verify.go b/internal/http/provider_verify.go index 48e4e591..8b3e402e 100644 --- a/internal/http/provider_verify.go +++ b/internal/http/provider_verify.go @@ -16,6 +16,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // HandleVerifyProviderForTest invokes the verify handler directly without auth @@ -126,8 +127,9 @@ func (h *ProvidersHandler) handleVerifyProvider(w http.ResponseWriter, r *http.R ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) defer cancel() + ctx = store.WithTenantID(ctx, p.TenantID) - _, err = provider.Chat(ctx, providers.ChatRequest{ + reqChat := providers.ChatRequest{ Messages: []providers.Message{ {Role: "user", Content: "hi"}, }, @@ -136,6 +138,13 @@ func (h *ProvidersHandler) handleVerifyProvider(w http.ResponseWriter, r *http.R // Use a small but safe value — reasoning models need headroom beyond 1 token. "max_tokens": 50, }, + } + _, err = h.usageCaps.Chat(ctx, provider, reqChat, usagecaps.ChatOptions{ + TenantID: p.TenantID, + ProviderName: p.Name, + ModelID: req.Model, + Purpose: "provider-verify", + MaxOutputTokens: 50, }) if err != nil { writeJSON(w, http.StatusOK, map[string]any{"valid": false, "error": friendlyVerifyError(err)}) diff --git a/internal/http/providers.go b/internal/http/providers.go index bee754f0..9127982f 100644 --- a/internal/http/providers.go +++ b/internal/http/providers.go @@ -22,6 +22,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/permissions" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -39,6 +40,7 @@ type ProvidersHandler struct { tracingStore store.TracingStore // optional: for provider-scoped pool activity agents store.AgentCRUDStore // optional: for provider pool activity agent lookup modelReg providers.ModelRegistry // optional: forward-compat model resolver for Anthropic + usageCaps *usagecaps.Service } // NewProvidersHandler creates a handler for provider management endpoints. @@ -85,6 +87,10 @@ func (h *ProvidersHandler) SetModelRegistry(r providers.ModelRegistry) { h.modelReg = r } +func (h *ProvidersHandler) SetUsageCapService(s *usagecaps.Service) { + h.usageCaps = s +} + // resolveAPIBase returns the provider's api_base, falling back to config/env if empty. // For Ollama/OllamaCloud providers, applies a safety-net normalization: if the stored // value is missing the /v1 suffix (pre-existing record before write-time normalization), diff --git a/internal/http/summoner.go b/internal/http/summoner.go index 540257c1..64f2effb 100644 --- a/internal/http/summoner.go +++ b/internal/http/summoner.go @@ -12,6 +12,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // Summoning event type constants. @@ -48,14 +49,16 @@ type AgentSummoner struct { agents store.AgentStore providerReg *providers.Registry msgBus *bus.MessageBus + usageCaps *usagecaps.Service } // NewAgentSummoner creates a summoner backed by the given stores and provider registry. -func NewAgentSummoner(agents store.AgentStore, providerReg *providers.Registry, msgBus *bus.MessageBus) *AgentSummoner { +func NewAgentSummoner(agents store.AgentStore, providerReg *providers.Registry, msgBus *bus.MessageBus, usageCaps *usagecaps.Service) *AgentSummoner { return &AgentSummoner{ agents: agents, providerReg: providerReg, msgBus: msgBus, + usageCaps: usageCaps, } } @@ -73,6 +76,7 @@ const singleCallTimeout = 300 * time.Second func (s *AgentSummoner) SummonAgent(agentID uuid.UUID, tenantID uuid.UUID, providerName, model, description string) { ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 600*time.Second) defer cancel() + ctx = store.WithAgentID(ctx, agentID) s.ensureBackfillFiles(ctx, agentID) s.emitEvent(agentID, tenantID, SummonEventStarted, "", "") diff --git a/internal/http/summoner_regenerate.go b/internal/http/summoner_regenerate.go index e3237719..126f87db 100644 --- a/internal/http/summoner_regenerate.go +++ b/internal/http/summoner_regenerate.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -23,6 +24,7 @@ import ( func (s *AgentSummoner) RegenerateAgent(agentID uuid.UUID, tenantID uuid.UUID, providerName, model, editPrompt string) { ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 300*time.Second) defer cancel() + ctx = store.WithAgentID(ctx, agentID) s.ensureBackfillFiles(ctx, agentID) @@ -121,7 +123,7 @@ func (s *AgentSummoner) generateFiles(ctx context.Context, providerName, model, slog.Info("summoning: calling LLM", "provider", providerName, "model", model, "prompt_len", len(prompt)) - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: "You are a file generator. Output ONLY the requested XML-tagged files. No extra commentary."}, {Role: "user", Content: prompt}, @@ -133,6 +135,12 @@ func (s *AgentSummoner) generateFiles(ctx context.Context, providerName, model, providers.OptSessionKey: summonSessionKey, providers.OptDisableTools: true, }, + } + resp, err := s.usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ProviderName: providerName, + ModelID: model, + Purpose: "agent-summoner", + MaxOutputTokens: 8192, }) if err != nil { return nil, fmt.Errorf("%s: %w", providerName, err) diff --git a/internal/http/usage_caps.go b/internal/http/usage_caps.go index cf1940ca..bf2c3718 100644 --- a/internal/http/usage_caps.go +++ b/internal/http/usage_caps.go @@ -13,6 +13,7 @@ import ( "time" "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/permissions" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/usage/pricing" @@ -64,11 +65,15 @@ func (h *UsageCapsHandler) masterAuth(next http.HandlerFunc) http.HandlerFunc { }) } +func writeUsageCapError(w http.ResponseWriter, r *http.Request, status int, key string, args ...any) { + writeJSON(w, status, map[string]string{"error": i18n.T(store.LocaleFromContext(r.Context()), key, args...)}) +} + func (h *UsageCapsHandler) handleListPolicies(w http.ResponseWriter, r *http.Request) { scope := store.UsageCapScope{TenantID: tenantIDOrMaster(r)} policies, err := h.store.ListUsageCapPolicies(r.Context(), scope, true) if err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "list policies failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsListPoliciesFailed) return } writeJSON(w, http.StatusOK, map[string]any{"policies": policies}) @@ -77,17 +82,17 @@ func (h *UsageCapsHandler) handleListPolicies(w http.ResponseWriter, r *http.Req func (h *UsageCapsHandler) handleCreatePolicy(w http.ResponseWriter, r *http.Request) { var body policyBody if err := json.NewDecoder(r.Body).Decode(&body); err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidJSON) return } p, err := body.toPolicy(tenantIDOrMaster(r)) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidRequest, err.Error()) return } if err := h.store.CreateUsageCapPolicy(r.Context(), &p); err != nil { slog.Warn("usage_caps.create_policy_failed", "error", err) - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "usage cap policy validation failed"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsageCapPolicyValidationFailed) return } writeJSON(w, http.StatusCreated, p) @@ -96,27 +101,27 @@ func (h *UsageCapsHandler) handleCreatePolicy(w http.ResponseWriter, r *http.Req func (h *UsageCapsHandler) handleUpdatePolicy(w http.ResponseWriter, r *http.Request) { id, err := uuid.Parse(r.PathValue("id")) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid policy id"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "policy") return } bodyBytes, err := io.ReadAll(r.Body) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidJSON) return } patch, err := policyPatchFromBody(bodyBytes) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidRequest, err.Error()) return } p, err := h.store.UpdateUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id, patch) if err != nil { if errors.Is(err, store.ErrUsageCapPolicyManaged) { - writeJSON(w, http.StatusConflict, map[string]string{"error": err.Error()}) + writeUsageCapError(w, r, http.StatusConflict, i18n.MsgUsageCapPolicyManaged) return } slog.Warn("usage_caps.update_policy_failed", "error", err) - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "usage cap policy validation failed"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsageCapPolicyValidationFailed) return } writeJSON(w, http.StatusOK, p) @@ -125,15 +130,15 @@ func (h *UsageCapsHandler) handleUpdatePolicy(w http.ResponseWriter, r *http.Req func (h *UsageCapsHandler) handleDeletePolicy(w http.ResponseWriter, r *http.Request) { id, err := uuid.Parse(r.PathValue("id")) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid policy id"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "policy") return } if err := h.store.DeleteUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id); err != nil { if errors.Is(err, store.ErrUsageCapPolicyManaged) { - writeJSON(w, http.StatusConflict, map[string]string{"error": err.Error()}) + writeUsageCapError(w, r, http.StatusConflict, i18n.MsgUsageCapPolicyManaged) return } - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "delete failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsDeletePolicyFailed) return } w.WriteHeader(http.StatusNoContent) @@ -142,7 +147,7 @@ func (h *UsageCapsHandler) handleDeletePolicy(w http.ResponseWriter, r *http.Req func (h *UsageCapsHandler) handleUtilization(w http.ResponseWriter, r *http.Request) { rows, err := h.store.ListUsageCapUtilization(r.Context(), tenantIDOrMaster(r)) if err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "utilization failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsUtilizationFailed) return } writeJSON(w, http.StatusOK, map[string]any{"rows": rows}) @@ -152,7 +157,7 @@ func (h *UsageCapsHandler) handleEvents(w http.ResponseWriter, r *http.Request) limit, _ := strconv.Atoi(r.URL.Query().Get("limit")) events, err := h.store.ListUsageCapEvents(r.Context(), tenantIDOrMaster(r), limit) if err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "events failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsEventsFailed) return } writeJSON(w, http.StatusOK, map[string]any{"events": events}) @@ -164,12 +169,12 @@ func (h *UsageCapsHandler) handleSyncOpenRouter(w http.ResponseWriter, r *http.R entries, err := pricing.FetchOpenRouterCatalog(ctx, h.client) if err != nil { slog.Warn("usage_pricing.openrouter_sync", "error", err) - writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()}) + writeUsageCapError(w, r, http.StatusBadGateway, i18n.MsgUsagePricingSyncOpenRouterFailed, err.Error()) return } count, err := h.store.UpsertPricingCatalog(r.Context(), entries) if err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "store catalog failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingStoreCatalogFailed) return } writeJSON(w, http.StatusOK, map[string]any{"count": count}) @@ -181,7 +186,7 @@ func (h *UsageCapsHandler) handleListPricing(w http.ResponseWriter, r *http.Requ Limit: queryInt(r, "limit", 100), }) if err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "list pricing failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingListFailed) return } writeJSON(w, http.StatusOK, map[string]any{"models": rows}) @@ -190,12 +195,12 @@ func (h *UsageCapsHandler) handleListPricing(w http.ResponseWriter, r *http.Requ func (h *UsageCapsHandler) handlePutOverride(w http.ResponseWriter, r *http.Request) { var body overrideBody if err := json.NewDecoder(r.Body).Decode(&body); err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidJSON) return } providerID, err := uuid.Parse(body.ProviderID) if err != nil || body.ModelID == "" { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "provider_id and model_id are required"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsagePricingProviderModelRequired) return } o := &store.UsagePricingOverride{ @@ -205,7 +210,7 @@ func (h *UsageCapsHandler) handlePutOverride(w http.ResponseWriter, r *http.Requ } if err := h.store.PutPricingOverride(r.Context(), o); err != nil { slog.Warn("usage_pricing.put_override_failed", "error", err) - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "pricing override validation failed"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsagePricingOverrideValidationFailed) return } writeJSON(w, http.StatusOK, o) @@ -217,13 +222,13 @@ func (h *UsageCapsHandler) handleListOverrides(w http.ResponseWriter, r *http.Re var err error providerID, err = uuid.Parse(raw) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid provider_id"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "provider") return } } rows, err := h.store.ListPricingOverrides(r.Context(), store.UsagePricingQuery{TenantID: tenantIDOrMaster(r), ProviderID: providerID}) if err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "list overrides failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingListOverridesFailed) return } writeJSON(w, http.StatusOK, map[string]any{"overrides": rows}) @@ -232,11 +237,11 @@ func (h *UsageCapsHandler) handleListOverrides(w http.ResponseWriter, r *http.Re func (h *UsageCapsHandler) handleDeleteOverride(w http.ResponseWriter, r *http.Request) { id, err := uuid.Parse(r.PathValue("id")) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid override id"}) + writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "override") return } if err := h.store.DeletePricingOverride(r.Context(), tenantIDOrMaster(r), id); err != nil { - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "delete failed"}) + writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingDeleteOverrideFailed) return } w.WriteHeader(http.StatusNoContent) diff --git a/internal/i18n/catalog_en.go b/internal/i18n/catalog_en.go index 4013913d..c87a7a03 100644 --- a/internal/i18n/catalog_en.go +++ b/internal/i18n/catalog_en.go @@ -87,6 +87,21 @@ func init() { // Provider MsgProviderReqFailed: "%s: request failed: %s", + // Usage caps / pricing + MsgUsageCapsListPoliciesFailed: "failed to list usage cap policies", + MsgUsageCapPolicyValidationFailed: "usage cap policy validation failed", + MsgUsageCapPolicyManaged: "managed usage cap policies cannot be modified", + MsgUsageCapsDeletePolicyFailed: "failed to delete usage cap policy", + MsgUsageCapsUtilizationFailed: "failed to load usage cap utilization", + MsgUsageCapsEventsFailed: "failed to load usage cap events", + MsgUsagePricingSyncOpenRouterFailed: "failed to sync OpenRouter pricing: %s", + MsgUsagePricingStoreCatalogFailed: "failed to store pricing catalog", + MsgUsagePricingListFailed: "failed to list model pricing", + MsgUsagePricingProviderModelRequired: "provider_id and model_id are required", + MsgUsagePricingOverrideValidationFailed: "pricing override validation failed", + MsgUsagePricingListOverridesFailed: "failed to list pricing overrides", + MsgUsagePricingDeleteOverrideFailed: "failed to delete pricing override", + // Unknown method MsgUnknownMethod: "unknown method: %s", @@ -200,10 +215,10 @@ func init() { MsgTenantScopeRequired: "tenant scope is required for this operation", // TTS / Voices - MsgTtsUnknownModel: "unknown tts model: %s", - MsgVoicesListFailed: "failed to list voices: %s", - MsgTtsGeminiInvalidVoice: "invalid Gemini voice: %s", - MsgTtsGeminiSpeakerLimit: "Gemini TTS supports at most 2 speakers", + MsgTtsUnknownModel: "unknown tts model: %s", + MsgVoicesListFailed: "failed to list voices: %s", + MsgTtsGeminiInvalidVoice: "invalid Gemini voice: %s", + MsgTtsGeminiSpeakerLimit: "Gemini TTS supports at most 2 speakers", MsgTtsGeminiInvalidModel: "invalid Gemini TTS model: %s", MsgTtsGeminiTextOnly: "Gemini refused to generate audio. Try simpler text without translation or commentary.", MsgTtsParamOutOfRange: "TTS param %q value %v is out of range [%v, %v]", diff --git a/internal/i18n/catalog_vi.go b/internal/i18n/catalog_vi.go index fe5c1073..db80a909 100644 --- a/internal/i18n/catalog_vi.go +++ b/internal/i18n/catalog_vi.go @@ -87,6 +87,21 @@ func init() { // Provider MsgProviderReqFailed: "%s: yêu cầu thất bại: %s", + // Usage caps / pricing + MsgUsageCapsListPoliciesFailed: "không thể liệt kê chính sách usage cap", + MsgUsageCapPolicyValidationFailed: "xác thực chính sách usage cap thất bại", + MsgUsageCapPolicyManaged: "không thể chỉnh sửa chính sách usage cap do hệ thống quản lý", + MsgUsageCapsDeletePolicyFailed: "không thể xóa chính sách usage cap", + MsgUsageCapsUtilizationFailed: "không thể tải mức sử dụng usage cap", + MsgUsageCapsEventsFailed: "không thể tải sự kiện usage cap", + MsgUsagePricingSyncOpenRouterFailed: "không thể đồng bộ giá OpenRouter: %s", + MsgUsagePricingStoreCatalogFailed: "không thể lưu catalog giá", + MsgUsagePricingListFailed: "không thể liệt kê giá model", + MsgUsagePricingProviderModelRequired: "provider_id và model_id là bắt buộc", + MsgUsagePricingOverrideValidationFailed: "xác thực override giá thất bại", + MsgUsagePricingListOverridesFailed: "không thể liệt kê override giá", + MsgUsagePricingDeleteOverrideFailed: "không thể xóa override giá", + // Unknown method MsgUnknownMethod: "phương thức không xác định: %s", @@ -200,10 +215,10 @@ func init() { MsgTenantScopeRequired: "cần xác định tenant để thực hiện thao tác này", // TTS / Giọng đọc - MsgTtsUnknownModel: "model tts không hỗ trợ: %s", - MsgVoicesListFailed: "không tải được danh sách giọng đọc: %s", - MsgTtsGeminiInvalidVoice: "giọng đọc Gemini không hợp lệ: %s", - MsgTtsGeminiSpeakerLimit: "Gemini TTS hỗ trợ tối đa 2 người nói", + MsgTtsUnknownModel: "model tts không hỗ trợ: %s", + MsgVoicesListFailed: "không tải được danh sách giọng đọc: %s", + MsgTtsGeminiInvalidVoice: "giọng đọc Gemini không hợp lệ: %s", + MsgTtsGeminiSpeakerLimit: "Gemini TTS hỗ trợ tối đa 2 người nói", MsgTtsGeminiInvalidModel: "mô hình Gemini TTS không hợp lệ: %s", MsgTtsGeminiTextOnly: "Gemini từ chối tạo âm thanh. Vui lòng thử văn bản đơn giản hơn, không dịch hay bình luận.", MsgTtsParamOutOfRange: "tham số TTS %q có giá trị %v nằm ngoài phạm vi [%v, %v]", diff --git a/internal/i18n/catalog_zh.go b/internal/i18n/catalog_zh.go index 0fac3cbb..698ff34c 100644 --- a/internal/i18n/catalog_zh.go +++ b/internal/i18n/catalog_zh.go @@ -87,6 +87,21 @@ func init() { // Provider MsgProviderReqFailed: "%s：请求失败：%s", + // Usage caps / pricing + MsgUsageCapsListPoliciesFailed: "无法列出 usage cap 策略", + MsgUsageCapPolicyValidationFailed: "usage cap 策略验证失败", + MsgUsageCapPolicyManaged: "无法修改系统托管的 usage cap 策略", + MsgUsageCapsDeletePolicyFailed: "无法删除 usage cap 策略", + MsgUsageCapsUtilizationFailed: "无法加载 usage cap 使用量", + MsgUsageCapsEventsFailed: "无法加载 usage cap 事件", + MsgUsagePricingSyncOpenRouterFailed: "无法同步 OpenRouter 价格：%s", + MsgUsagePricingStoreCatalogFailed: "无法保存价格目录", + MsgUsagePricingListFailed: "无法列出模型价格", + MsgUsagePricingProviderModelRequired: "provider_id 和 model_id 是必填项", + MsgUsagePricingOverrideValidationFailed: "价格覆盖验证失败", + MsgUsagePricingListOverridesFailed: "无法列出价格覆盖", + MsgUsagePricingDeleteOverrideFailed: "无法删除价格覆盖", + // Unknown method MsgUnknownMethod: "未知方法：%s", @@ -200,10 +215,10 @@ func init() { MsgTenantScopeRequired: "此操作需要指定租户范围", // TTS / 声音 - MsgTtsUnknownModel: "未知的 tts 模型：%s", - MsgVoicesListFailed: "获取声音列表失败：%s", - MsgTtsGeminiInvalidVoice: "无效的 Gemini 声音：%s", - MsgTtsGeminiSpeakerLimit: "Gemini TTS 最多支持 2 位发言人", + MsgTtsUnknownModel: "未知的 tts 模型：%s", + MsgVoicesListFailed: "获取声音列表失败：%s", + MsgTtsGeminiInvalidVoice: "无效的 Gemini 声音：%s", + MsgTtsGeminiSpeakerLimit: "Gemini TTS 最多支持 2 位发言人", MsgTtsGeminiInvalidModel: "无效的 Gemini TTS 模型：%s", MsgTtsGeminiTextOnly: "Gemini 拒绝生成音频。请尝试更简单的文本，不要翻译或添加评论。", MsgTtsParamOutOfRange: "TTS 参数 %q 的值 %v 超出范围 [%v, %v]", diff --git a/internal/i18n/keys.go b/internal/i18n/keys.go index f6644b51..b53a7648 100644 --- a/internal/i18n/keys.go +++ b/internal/i18n/keys.go @@ -88,6 +88,21 @@ const ( // --- Provider --- MsgProviderReqFailed = "error.provider_request_failed" // "%s: request failed: %s" + // --- Usage caps / pricing --- + MsgUsageCapsListPoliciesFailed = "usage_caps.list_policies_failed" + MsgUsageCapPolicyValidationFailed = "usage_caps.policy_validation_failed" + MsgUsageCapPolicyManaged = "usage_caps.policy_managed" + MsgUsageCapsDeletePolicyFailed = "usage_caps.delete_policy_failed" + MsgUsageCapsUtilizationFailed = "usage_caps.utilization_failed" + MsgUsageCapsEventsFailed = "usage_caps.events_failed" + MsgUsagePricingSyncOpenRouterFailed = "usage_pricing.sync_openrouter_failed" + MsgUsagePricingStoreCatalogFailed = "usage_pricing.store_catalog_failed" + MsgUsagePricingListFailed = "usage_pricing.list_failed" + MsgUsagePricingProviderModelRequired = "usage_pricing.provider_model_required" + MsgUsagePricingOverrideValidationFailed = "usage_pricing.override_validation_failed" + MsgUsagePricingListOverridesFailed = "usage_pricing.list_overrides_failed" + MsgUsagePricingDeleteOverrideFailed = "usage_pricing.delete_override_failed" + // --- Unknown method --- MsgUnknownMethod = "error.unknown_method" // "unknown method: %s" @@ -117,14 +132,14 @@ const ( MsgInvalidVisibility = "error.invalid_visibility" // "invalid visibility %q: must be one of private, public" // --- Package updates (Phase 4+5) --- - MsgPackageNotInstalled = "packages.update.not_installed" // "Package {name} is not installed" - MsgPackageUpdateLocked = "packages.update.locked" // "Package {name} is being updated by another request" + MsgPackageNotInstalled = "packages.update.not_installed" // "Package {name} is not installed" + MsgPackageUpdateLocked = "packages.update.locked" // "Package {name} is being updated by another request" MsgReleaseNotFound = "packages.update.release_not_found" // "Release {tag} not found for {repo}" - MsgAssetNotFound = "packages.update.asset_not_found" // "No compatible asset for {os}/{arch}" + MsgAssetNotFound = "packages.update.asset_not_found" // "No compatible asset for {os}/{arch}" MsgChecksumMismatch = "packages.update.checksum_mismatch" // "Checksum mismatch for {name}" - MsgUpdateSwapFailed = "packages.update.swap_failed" // "Failed to install {name}; previous version restored" - MsgUpdateManifestDesync = "packages.update.manifest_desync" // "Binary updated but manifest save failed — manual recovery required for {name}" - MsgUpdateCacheStale = "packages.update.cache_stale" // "Updates cache stale; run refresh before applying an update" + MsgUpdateSwapFailed = "packages.update.swap_failed" // "Failed to install {name}; previous version restored" + MsgUpdateManifestDesync = "packages.update.manifest_desync" // "Binary updated but manifest save failed — manual recovery required for {name}" + MsgUpdateCacheStale = "packages.update.cache_stale" // "Updates cache stale; run refresh before applying an update" // Package update source labels MsgPackagesUpdatesSourceGithub = "packages.updates.source.github" // "GitHub" @@ -234,15 +249,15 @@ const ( MsgInvalidRole = "error.invalid_role" // "invalid role: allowed values are owner, admin, operator, member, viewer" // --- TTS / Voices --- - MsgTtsUnknownModel = "error.tts_unknown_model" // "unknown tts model: %s" - MsgVoicesListFailed = "error.voices_list_failed" // "failed to list voices: %s" - MsgTtsGeminiInvalidVoice = "error.tts_gemini_invalid_voice" // "invalid Gemini voice: %s" - MsgTtsGeminiSpeakerLimit = "error.tts_gemini_speaker_limit" // "Gemini TTS supports at most 2 speakers" - MsgTtsGeminiInvalidModel = "error.tts_gemini_invalid_model" // "invalid Gemini TTS model: %s" - MsgTtsGeminiTextOnly = "error.tts_gemini_text_only" // "Gemini refused to generate audio; try simpler text without translation or commentary" - MsgTtsParamOutOfRange = "error.tts_param_out_of_range" // "TTS param %q value %v is out of range [%v, %v]" - MsgTtsParamUnknownKey = "error.tts_param_unknown_key" // "TTS param %q is not supported by this provider" - MsgTtsMiniMaxVoicesFailed = "error.tts_minimax_voices_failed" // "failed to fetch MiniMax voices: %s" + MsgTtsUnknownModel = "error.tts_unknown_model" // "unknown tts model: %s" + MsgVoicesListFailed = "error.voices_list_failed" // "failed to list voices: %s" + MsgTtsGeminiInvalidVoice = "error.tts_gemini_invalid_voice" // "invalid Gemini voice: %s" + MsgTtsGeminiSpeakerLimit = "error.tts_gemini_speaker_limit" // "Gemini TTS supports at most 2 speakers" + MsgTtsGeminiInvalidModel = "error.tts_gemini_invalid_model" // "invalid Gemini TTS model: %s" + MsgTtsGeminiTextOnly = "error.tts_gemini_text_only" // "Gemini refused to generate audio; try simpler text without translation or commentary" + MsgTtsParamOutOfRange = "error.tts_param_out_of_range" // "TTS param %q value %v is out of range [%v, %v]" + MsgTtsParamUnknownKey = "error.tts_param_unknown_key" // "TTS param %q is not supported by this provider" + MsgTtsMiniMaxVoicesFailed = "error.tts_minimax_voices_failed" // "failed to fetch MiniMax voices: %s" // --- STT --- MsgSTTAllProvidersFailed = "error.stt_all_providers_failed" // "All STT providers failed" @@ -258,50 +273,50 @@ const ( MsgTenantScopeRequired = "error.tenant_scope_required" // "tenant scope is required for this operation" // --- Webhooks --- - MsgWebhookAuthFailed = "webhook.auth_failed" // "webhook authentication failed" - MsgWebhookHMACInvalid = "webhook.hmac_invalid" // "HMAC signature is invalid" - MsgWebhookHMACTimestampSkew = "webhook.hmac_timestamp_skew" // "request timestamp outside acceptable window" - MsgWebhookBearerRequiredHMAC = "webhook.bearer_required_hmac" // "this webhook requires HMAC authentication" - MsgWebhookRevoked = "webhook.revoked" // "webhook has been revoked" - MsgWebhookKindMismatch = "webhook.kind_mismatch" // "request kind does not match webhook configuration" - MsgWebhookRateLimited = "webhook.rate_limited" // "webhook rate limit exceeded" - MsgWebhookBodyTooLarge = "webhook.body_too_large" // "request body exceeds size limit" - MsgWebhookIdempotencyConflict = "webhook.idempotency_conflict" // "idempotency key conflict: request body mismatch" - MsgWebhookTenantMismatch = "webhook.tenant_mismatch" // "webhook tenant mismatch" - MsgWebhookAgentNotFound = "webhook.agent_not_found" // "webhook agent not found" - MsgWebhookChannelNotFound = "webhook.channel_not_found" // "webhook channel not found" - MsgWebhookMediaSSRFBlocked = "webhook.media_ssrf_blocked" // "media URL blocked by SSRF policy" - MsgWebhookMediaTooLarge = "webhook.media_too_large" // "media file exceeds size limit" - MsgWebhookMediaMIMEDenied = "webhook.media_mime_denied" // "media MIME type is not allowed" - MsgWebhookCallbackURLInvalid = "webhook.callback_url_invalid" // "callback URL is invalid or blocked" - MsgWebhookLLMTimeout = "webhook.llm_timeout" // "LLM processing timed out" - MsgWebhookLaneSaturated = "webhook.lane_saturated" // "webhook processing lane is at capacity" + MsgWebhookAuthFailed = "webhook.auth_failed" // "webhook authentication failed" + MsgWebhookHMACInvalid = "webhook.hmac_invalid" // "HMAC signature is invalid" + MsgWebhookHMACTimestampSkew = "webhook.hmac_timestamp_skew" // "request timestamp outside acceptable window" + MsgWebhookBearerRequiredHMAC = "webhook.bearer_required_hmac" // "this webhook requires HMAC authentication" + MsgWebhookRevoked = "webhook.revoked" // "webhook has been revoked" + MsgWebhookKindMismatch = "webhook.kind_mismatch" // "request kind does not match webhook configuration" + MsgWebhookRateLimited = "webhook.rate_limited" // "webhook rate limit exceeded" + MsgWebhookBodyTooLarge = "webhook.body_too_large" // "request body exceeds size limit" + MsgWebhookIdempotencyConflict = "webhook.idempotency_conflict" // "idempotency key conflict: request body mismatch" + MsgWebhookTenantMismatch = "webhook.tenant_mismatch" // "webhook tenant mismatch" + MsgWebhookAgentNotFound = "webhook.agent_not_found" // "webhook agent not found" + MsgWebhookChannelNotFound = "webhook.channel_not_found" // "webhook channel not found" + MsgWebhookMediaSSRFBlocked = "webhook.media_ssrf_blocked" // "media URL blocked by SSRF policy" + MsgWebhookMediaTooLarge = "webhook.media_too_large" // "media file exceeds size limit" + MsgWebhookMediaMIMEDenied = "webhook.media_mime_denied" // "media MIME type is not allowed" + MsgWebhookCallbackURLInvalid = "webhook.callback_url_invalid" // "callback URL is invalid or blocked" + MsgWebhookLLMTimeout = "webhook.llm_timeout" // "LLM processing timed out" + MsgWebhookLaneSaturated = "webhook.lane_saturated" // "webhook processing lane is at capacity" MsgWebhookLocalhostOnlyViolation = "webhook.localhost_only_violation" // "this webhook is restricted to localhost callers" MsgWebhookMediaChannelUnsupported = "webhook.media_channel_unsupported" // "channel does not support media attachments" MsgWebhookIPDenied = "webhook.ip_denied" // "request origin is not in the IP allowlist" MsgWebhookEncryptionUnavailable = "webhook.encryption_unavailable" // "webhook encryption key not configured; set GOCLAW_ENCRYPTION_KEY to enable webhooks" // --- Workstation permissions --- - MsgWorkstationCmdDenied = "error.workstation_cmd_denied" // "command denied by workstation policy: %s" - MsgWorkstationEnvDenied = "error.workstation_env_denied" // "env var denied by policy: %s" - MsgWorkstationInputInvalid = "error.workstation_input_invalid" // "command contains invalid characters: %s" - MsgWorkstationRateLimit = "error.workstation_rate_limit" // "workstation rate limit exceeded" - MsgWorkstationPermNotFound = "error.workstation_perm_not_found" // "permission entry not found: %s" + MsgWorkstationCmdDenied = "error.workstation_cmd_denied" // "command denied by workstation policy: %s" + MsgWorkstationEnvDenied = "error.workstation_env_denied" // "env var denied by policy: %s" + MsgWorkstationInputInvalid = "error.workstation_input_invalid" // "command contains invalid characters: %s" + MsgWorkstationRateLimit = "error.workstation_rate_limit" // "workstation rate limit exceeded" + MsgWorkstationPermNotFound = "error.workstation_perm_not_found" // "permission entry not found: %s" // --- Workstation activity (Phase 7) --- - MsgWorkstationActivityTitle = "ui.workstations.activity.title" // "Recent Activity" - MsgWorkstationActionExec = "ui.workstations.activity.action_exec" // "Exec" - MsgWorkstationActionDeny = "ui.workstations.activity.action_deny" // "Denied" + MsgWorkstationActivityTitle = "ui.workstations.activity.title" // "Recent Activity" + MsgWorkstationActionExec = "ui.workstations.activity.action_exec" // "Exec" + MsgWorkstationActionDeny = "ui.workstations.activity.action_deny" // "Denied" // --- Workstation --- - MsgWorkstationNotFound = "error.workstation_not_found" // "workstation not found: %s" - MsgWorkstationKeyExists = "error.workstation_key_exists" // "workstation key already in use: %s" - MsgInvalidBackend = "error.invalid_backend" // "invalid backend type: %s (must be ssh|docker)" - MsgWorkstationInactive = "error.workstation_inactive" // "workstation is inactive: %s" - MsgInvalidMetadataShape = "error.invalid_metadata_shape" // "invalid metadata for %s backend: %s" - MsgWorkstationRequired = "error.workstation_required" // "no workstation bound to agent; pass workstation_id" + MsgWorkstationNotFound = "error.workstation_not_found" // "workstation not found: %s" + MsgWorkstationKeyExists = "error.workstation_key_exists" // "workstation key already in use: %s" + MsgInvalidBackend = "error.invalid_backend" // "invalid backend type: %s (must be ssh|docker)" + MsgWorkstationInactive = "error.workstation_inactive" // "workstation is inactive: %s" + MsgInvalidMetadataShape = "error.invalid_metadata_shape" // "invalid metadata for %s backend: %s" + MsgWorkstationRequired = "error.workstation_required" // "no workstation bound to agent; pass workstation_id" MsgWorkstationAccessDenied = "error.workstation_access_denied" // "agent %s not authorized for workstation %s" - MsgBackendNotReady = "error.backend_not_ready" // "workstation backend not ready: %s" + MsgBackendNotReady = "error.backend_not_ready" // "workstation backend not ready: %s" // --- Hooks --- MsgHookInvalidMatcher = "hook.invalid_matcher" // "invalid matcher regex: %s" diff --git a/internal/knowledgegraph/extractor.go b/internal/knowledgegraph/extractor.go index 491c6c2c..addb80df 100644 --- a/internal/knowledgegraph/extractor.go +++ b/internal/knowledgegraph/extractor.go @@ -9,6 +9,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ExtractionResult holds entities and relations extracted from text. @@ -22,6 +23,7 @@ type Extractor struct { provider providers.Provider model string minConfidence float64 + usageCaps *usagecaps.Service } // NewExtractor creates a new Extractor with the given provider, model, and confidence threshold. @@ -32,6 +34,11 @@ func NewExtractor(provider providers.Provider, model string, minConfidence float return &Extractor{provider: provider, model: model, minConfidence: minConfidence} } +// SetUsageCapService enables cost enforcement for LLM extraction calls. +func (e *Extractor) SetUsageCapService(s *usagecaps.Service) { + e.usageCaps = s +} + const maxChunkChars = 12000 // Extract calls the LLM to extract entities and relations from text. @@ -72,7 +79,11 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR }, } - resp, err := e.provider.Chat(ctx, req) + resp, err := e.usageCaps.Chat(ctx, e.provider, req, usagecaps.ChatOptions{ + ModelID: e.model, + Purpose: "knowledge-graph-extract", + MaxOutputTokens: 8192, + }) if err != nil { return nil, fmt.Errorf("kg extraction LLM call: %w", err) } @@ -85,7 +96,11 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR text = text[:retryMaxChars] + "\n\n[...truncated]" } req.Messages[1].Content = text - resp, err = e.provider.Chat(ctx, req) + resp, err = e.usageCaps.Chat(ctx, e.provider, req, usagecaps.ChatOptions{ + ModelID: e.model, + Purpose: "knowledge-graph-extract-retry", + MaxOutputTokens: 8192, + }) if err != nil { return nil, fmt.Errorf("kg extraction LLM retry: %w", err) } @@ -286,4 +301,3 @@ func stripCodeBlock(s string) string { } return strings.TrimSpace(s) } - diff --git a/internal/store/pg/usage_caps.go b/internal/store/pg/usage_caps.go index 06bdeaa4..f693cb63 100644 --- a/internal/store/pg/usage_caps.go +++ b/internal/store/pg/usage_caps.go @@ -267,9 +267,16 @@ func (s *PGUsageCapStore) ListUsageCapUtilization(ctx context.Context, tenantID for _, p := range policies { start, end := usageWindow(time.Now().UTC(), p.Window) u := store.UsageCapUtilization{Policy: p, WindowStart: start, WindowEnd: end} - _ = s.db.QueryRowContext(ctx, ` + err := s.db.QueryRowContext(ctx, ` SELECT used_tokens, reserved_tokens, used_cost_micros, reserved_cost_micros FROM usage_cap_counters WHERE policy_id=$1 AND window_start=$2`, p.ID, start).Scan(&u.UsedTokens, &u.ReservedTokens, &u.UsedCostMicros, &u.ReservedCostMicros) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + out = append(out, u) + continue + } + return nil, err + } out = append(out, u) } return out, nil diff --git a/internal/usage/caps/chat_call.go b/internal/usage/caps/chat_call.go new file mode 100644 index 00000000..5883d263 --- /dev/null +++ b/internal/usage/caps/chat_call.go @@ -0,0 +1,150 @@ +package caps + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +// ChatOptions identifies a billable non-agent LLM call for usage-cap enforcement. +type ChatOptions struct { + TenantID uuid.UUID + AgentID uuid.UUID + ProviderName string + ModelID string + ReservationKey string + Purpose string + MaxOutputTokens int +} + +// Chat wraps Provider.Chat with the same usage-cap preflight and reconciliation +// used by agent loops. A nil service intentionally falls back to direct calls +// for Lite/subscription-only runtimes. +func (s *Service) Chat(ctx context.Context, provider providers.Provider, req providers.ChatRequest, opts ChatOptions) (*providers.ChatResponse, error) { + if provider == nil { + return nil, errors.New("usage cap chat: provider is nil") + } + if s == nil || s.store == nil { + return provider.Chat(ctx, req) + } + if fallback, ok := provider.(*providers.ModelFallbackProvider); ok { + return fallback.ChatWithHook(ctx, req, func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + callOpts := opts + callOpts.ProviderName = entry.ProviderName + if callOpts.ProviderName == "" && entry.Provider != nil { + callOpts.ProviderName = entry.Provider.Name() + } + callOpts.ModelID = actualReq.Model + callOpts.ReservationKey = "" + usageReq := s.chatRequest(callCtx, entry.Provider, actualReq, callOpts) + scopedCtx := scopedRequestContext(callCtx, usageReq) + reservation, err := s.Preflight(scopedCtx, usageReq) + if err != nil { + return nil, err + } + return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) { + if reservation != nil { + reservation.Reconcile(scopedCtx, resp, callErr) + } + }, nil + }) + } + + usageReq := s.chatRequest(ctx, provider, req, opts) + scopedCtx := scopedRequestContext(ctx, usageReq) + reservation, err := s.Preflight(scopedCtx, usageReq) + if err != nil { + return nil, err + } + resp, err := provider.Chat(scopedCtx, req) + if reservation != nil { + reservation.Reconcile(scopedCtx, resp, err) + } + return resp, err +} + +func (s *Service) chatRequest(ctx context.Context, provider providers.Provider, req providers.ChatRequest, opts ChatOptions) Request { + tenantID := opts.TenantID + if tenantID == uuid.Nil { + tenantID = store.TenantIDFromContext(ctx) + } + if tenantID == uuid.Nil { + tenantID = store.MasterTenantID + } + agentID := opts.AgentID + if agentID == uuid.Nil { + agentID = store.AgentIDFromContext(ctx) + } + providerName := opts.ProviderName + if providerName == "" && provider != nil { + providerName = provider.Name() + } + modelID := opts.ModelID + if modelID == "" { + modelID = req.Model + } + if modelID == "" && provider != nil { + modelID = provider.DefaultModel() + } + return Request{ + TenantID: tenantID, + AgentID: agentID, + ProviderName: providerName, + ModelID: modelID, + ReservationKey: reservationKey(opts), + Messages: req.Messages, + MaxOutputTokens: maxOutputTokens(req, opts.MaxOutputTokens), + } +} + +func scopedRequestContext(ctx context.Context, req Request) context.Context { + if req.TenantID != uuid.Nil && store.TenantIDFromContext(ctx) != req.TenantID { + ctx = store.WithTenantID(ctx, req.TenantID) + } + if req.AgentID != uuid.Nil && store.AgentIDFromContext(ctx) != req.AgentID { + ctx = store.WithAgentID(ctx, req.AgentID) + } + return ctx +} + +func reservationKey(opts ChatOptions) string { + if opts.ReservationKey != "" { + return opts.ReservationKey + } + purpose := opts.Purpose + if purpose == "" { + purpose = "llm" + } + return fmt.Sprintf("%s:%s", purpose, uuid.NewString()) +} + +func maxOutputTokens(req providers.ChatRequest, fallback int) int { + if fallback <= 0 { + fallback = 1024 + } + if req.Options == nil { + return fallback + } + v, ok := req.Options[providers.OptMaxTokens] + if !ok { + return fallback + } + switch n := v.(type) { + case int: + return n + case int64: + return int(n) + case int32: + return int(n) + case float64: + return int(n) + case float32: + return int(n) + default: + return fallback + } +} diff --git a/internal/usage/caps/service.go b/internal/usage/caps/service.go index d3fbbaf3..30061bec 100644 --- a/internal/usage/caps/service.go +++ b/internal/usage/caps/service.go @@ -66,6 +66,7 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err if s == nil || s.store == nil { return skippedReservation(req, "service_disabled"), nil } + ctx = scopedRequestContext(ctx, req) providerData, err := s.resolveProvider(ctx, req.TenantID, req.ProviderName) if err != nil { return skippedReservation(req, "provider_metadata_missing"), nil @@ -107,6 +108,12 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err resolved, err := s.store.ResolvePricing(ctx, req.TenantID, providerData.ID, providerData.Name, providerData.ProviderType, req.ModelID) if err != nil { if errors.Is(err, sql.ErrNoRows) { + _ = s.store.InsertUsageCapEvent(ctx, &store.UsageCapEvent{ + TenantID: req.TenantID, + ReservationKey: key, Decision: store.UsageCapEventBlock, Reason: "pricing_unknown", + EstimatedTokens: usage.TotalTokens(), EstimatedCostMicros: 0, + Metadata: mustJSON(map[string]any{"model_id": req.ModelID, "provider": req.ProviderName}), + }) return blockedReservation(req, scope, key, usage, 0, uuid.Nil, "pricing_unknown"), fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID) } return nil, err diff --git a/internal/usage/caps/service_test.go b/internal/usage/caps/service_test.go index 7a14a084..6a597404 100644 --- a/internal/usage/caps/service_test.go +++ b/internal/usage/caps/service_test.go @@ -89,7 +89,7 @@ func TestPreflightIncludesRequestPricingWhenConfigured(t *testing.T) { Name: "openrouter", ProviderType: store.ProviderOpenRouter, APIKey: "sk-test", - }} + }, requireTenant: policy.TenantID} svc := NewService(usageStore, providerStore) _, err := svc.Preflight(context.Background(), Request{ @@ -279,6 +279,79 @@ func TestPreflightTraceMetadataForCapExceeded(t *testing.T) { } } +func TestPreflightRecordsPricingUnknownBlockEvent(t *testing.T) { + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxCostMicros: int64Ptr(1000), Enabled: true} + usageStore := &fakeUsageCapStore{ + policies: []store.UsageCapPolicy{policy}, + resolveErr: sql.ErrNoRows, + } + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := NewService(usageStore, providerStore) + + reservation, err := svc.Preflight(context.Background(), Request{ + TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "missing/model", + ReservationKey: "pricing-missing", Messages: []providers.Message{{Role: "user", Content: "hello"}}, + MaxOutputTokens: 10, + }) + if !errors.Is(err, ErrPricingUnknown) { + t.Fatalf("Preflight error = %v, want ErrPricingUnknown", err) + } + if reservation == nil { + t.Fatal("Preflight returned nil reservation") + } + metadata := reservation.TraceMetadata() + if metadata.Reason != "pricing_unknown" { + t.Fatalf("reservation metadata = %+v, want pricing_unknown", metadata) + } + if len(usageStore.events) != 1 { + t.Fatalf("events = %d, want 1", len(usageStore.events)) + } + event := usageStore.events[0] + if event.Decision != store.UsageCapEventBlock || event.Reason != "pricing_unknown" { + t.Fatalf("event decision/reason = %q/%q, want block/pricing_unknown", event.Decision, event.Reason) + } + if event.ReservationKey != "pricing-missing" { + t.Fatalf("event reservation_key = %q, want pricing-missing", event.ReservationKey) + } +} + +func TestServiceChatBlocksBeforeProviderCall(t *testing.T) { + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(10), Enabled: true} + usageStore := &fakeUsageCapStore{ + policies: []store.UsageCapPolicy{policy}, + reserveErr: &store.UsageCapExceededError{PolicyID: policy.ID, Reason: "token_cap_exceeded"}, + } + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }, requireTenant: policy.TenantID} + svc := NewService(usageStore, providerStore) + provider := &fakeChatProvider{name: "openrouter", model: "token/model"} + + _, err := svc.Chat(context.Background(), provider, providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: "hello"}}, + Model: "token/model", + Options: map[string]any{providers.OptMaxTokens: 20}, + }, ChatOptions{ + TenantID: policy.TenantID, + ProviderName: "openrouter", + Purpose: "test-block", + }) + if !errors.Is(err, ErrCapExceeded) { + t.Fatalf("Chat error = %v, want ErrCapExceeded", err) + } + if provider.calls != 0 { + t.Fatalf("provider calls = %d, want 0", provider.calls) + } +} + func TestMergeTraceMetadataPreservesExistingSections(t *testing.T) { existing := json.RawMessage(`{"thinking":{"effort":"high"}}`) merged := MergeTraceMetadata(existing, []TraceMetadata{{ @@ -319,6 +392,32 @@ func TestCountImagesOnlyCountsImageMIMEs(t *testing.T) { } } +type fakeChatProvider struct { + name string + model string + calls int + resp *providers.ChatResponse + err error +} + +func (p *fakeChatProvider) Chat(context.Context, providers.ChatRequest) (*providers.ChatResponse, error) { + p.calls++ + if p.resp != nil || p.err != nil { + return p.resp, p.err + } + return &providers.ChatResponse{ + Content: "ok", + Usage: &providers.Usage{PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2}, + }, nil +} + +func (p *fakeChatProvider) ChatStream(ctx context.Context, req providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) { + return p.Chat(ctx, req) +} + +func (p *fakeChatProvider) DefaultModel() string { return p.model } +func (p *fakeChatProvider) Name() string { return p.name } + type fakeUsageCapStore struct { policies []store.UsageCapPolicy resolved *store.ResolvedUsagePricing @@ -329,6 +428,7 @@ type fakeUsageCapStore struct { reconciled store.UsageReconcileRequest reconcileCalls int reconcileCtxCanceled bool + events []store.UsageCapEvent } func (s *fakeUsageCapStore) UpsertPricingCatalog(context.Context, []store.UsagePricingCatalogEntry) (int, error) { @@ -384,13 +484,17 @@ func (s *fakeUsageCapStore) ListUsageCapUtilization(context.Context, uuid.UUID) func (s *fakeUsageCapStore) ListUsageCapEvents(context.Context, uuid.UUID, int) ([]store.UsageCapEvent, error) { return nil, nil } -func (s *fakeUsageCapStore) InsertUsageCapEvent(context.Context, *store.UsageCapEvent) error { +func (s *fakeUsageCapStore) InsertUsageCapEvent(_ context.Context, event *store.UsageCapEvent) error { + if event != nil { + s.events = append(s.events, *event) + } return nil } type fakeProviderStore struct { provider *store.LLMProviderData masterProvider *store.LLMProviderData + requireTenant uuid.UUID } func (s *fakeProviderStore) CreateProvider(context.Context, *store.LLMProviderData) error { return nil } @@ -398,6 +502,9 @@ func (s *fakeProviderStore) GetProvider(context.Context, uuid.UUID) (*store.LLMP return s.provider, nil } func (s *fakeProviderStore) GetProviderByName(ctx context.Context, _ string) (*store.LLMProviderData, error) { + if s.requireTenant != uuid.Nil && store.TenantIDFromContext(ctx) != s.requireTenant { + return nil, sql.ErrNoRows + } if store.TenantIDFromContext(ctx) == store.MasterTenantID && s.masterProvider != nil { return s.masterProvider, nil } diff --git a/internal/vault/enrich_classify.go b/internal/vault/enrich_classify.go index b8aaf857..227a27cc 100644 --- a/internal/vault/enrich_classify.go +++ b/internal/vault/enrich_classify.go @@ -7,6 +7,7 @@ import ( "maps" "slices" + "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -40,6 +41,14 @@ func (w *EnrichWorker) classifyLinks(ctx context.Context, provider providers.Pro if provider == nil { return } + if tid, err := uuid.Parse(tenantID); err == nil { + ctx = store.WithTenantID(ctx, tid) + } + if agentID != "" { + if aid, err := uuid.Parse(agentID); err == nil { + ctx = store.WithAgentID(ctx, aid) + } + } capped := results if len(capped) > classifyMaxSourceDocs { diff --git a/internal/vault/enrich_worker.go b/internal/vault/enrich_worker.go index 18d9f49e..2fd4cb1c 100644 --- a/internal/vault/enrich_worker.go +++ b/internal/vault/enrich_worker.go @@ -16,20 +16,21 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bgalert" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/eventbus" - "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/providerresolve" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "golang.org/x/sync/semaphore" ) const ( - enrichMaxDedupEntries = 10000 - enrichSimilarityLimit = 10 - enrichSimilarityMin = 0.7 - enrichMaxConcurrent = 3 // max concurrent batch summarize calls across chunks - enrichBatchSize = 5 // docs per enrichment chunk (1 LLM call per chunk) - enrichBatchItemMaxRunes = 3000 // per-file content limit in batch summarize - enrichMaxRetries = 3 // shared retry count for LLM calls (summarize + classify) + enrichMaxDedupEntries = 10000 + enrichSimilarityLimit = 10 + enrichSimilarityMin = 0.7 + enrichMaxConcurrent = 3 // max concurrent batch summarize calls across chunks + enrichBatchSize = 5 // docs per enrichment chunk (1 LLM call per chunk) + enrichBatchItemMaxRunes = 3000 // per-file content limit in batch summarize + enrichMaxRetries = 3 // shared retry count for LLM calls (summarize + classify) ) // Shared retry config for all enrichment LLM calls. @@ -41,12 +42,13 @@ var ( // EnrichWorkerDeps bundles dependencies for the vault enrichment worker. type EnrichWorkerDeps struct { VaultStore store.VaultStore - SystemConfigs store.SystemConfigStore // per-tenant provider config - Registry *providers.Registry // provider resolution + SystemConfigs store.SystemConfigStore // per-tenant provider config + Registry *providers.Registry // provider resolution EventBus eventbus.DomainEventBus - MsgBus bus.EventPublisher // for WS event broadcast - TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe) - AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + MsgBus bus.EventPublisher // for WS event broadcast + TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe) + AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + UsageCaps *usagecaps.Service } // RegisterEnrichWorker subscribes the enrichment worker to vault doc events. @@ -60,6 +62,7 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *Enri registry: deps.Registry, msgBus: deps.MsgBus, alertDeps: deps.AlertDeps, + usageCaps: deps.UsageCaps, dedup: make(map[string]string), sem: semaphore.NewWeighted(enrichMaxConcurrent), progress: progress, @@ -74,11 +77,12 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *Enri // Exported so HTTP handlers can call Stop/EnqueueUnenriched. type EnrichWorker struct { vault store.VaultStore - teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil - systemConfigs store.SystemConfigStore // per-tenant provider config - registry *providers.Registry // provider resolution - msgBus bus.EventPublisher // for error event broadcast - alertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil + systemConfigs store.SystemConfigStore // per-tenant provider config + registry *providers.Registry // provider resolution + msgBus bus.EventPublisher // for error event broadcast + alertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + usageCaps *usagecaps.Service queue enrichBatchQueue progress *EnrichProgress @@ -296,6 +300,14 @@ func (w *EnrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD // Batch-fetch all existing docs in a single query. tenantID := pending[0].TenantID + if tid, err := uuid.Parse(tenantID); err == nil { + ctx = store.WithTenantID(ctx, tid) + } + if pending[0].AgentID != "" { + if aid, err := uuid.Parse(pending[0].AgentID); err == nil { + ctx = store.WithAgentID(ctx, aid) + } + } // Resolve provider once per chunk (all items share tenantID) provider, model := w.resolveProviderForTenant(ctx, tenantID) @@ -515,7 +527,11 @@ func (w *EnrichWorker) chatWithRetry(ctx context.Context, provider providers.Pro } } cctx, cancel := context.WithTimeout(ctx, enrichRetryTimeouts[attempt]) - resp, err := provider.Chat(cctx, req) + resp, err := w.usageCaps.Chat(cctx, provider, req, usagecaps.ChatOptions{ + ModelID: req.Model, + Purpose: logPrefix, + MaxOutputTokens: 4096, + }) cancel() if err != nil { lastErr = err @@ -563,7 +579,6 @@ func (w *EnrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUps } } - // recordDedup stores a processed hash and evicts ~25% entries if over capacity. func (w *EnrichWorker) recordDedup(docID, hash string) { w.dedupMu.Lock() From 5ae58157383010e839733a56b20ec5ad3ca39910 Mon Sep 17 00:00:00 2001 From: Duy /zuey/ Date: Sun, 24 May 2026 12:48:51 +0700 Subject: [PATCH 06/10] fix(telegram): preserve archive upload paths (#58) --- docs/14-skills-runtime.md | 10 +-- internal/agent/loop_input_media.go | 5 +- internal/agent/media_persist_test.go | 8 +++ internal/agent/media_test.go | 28 ++++++++ internal/channels/history.go | 8 ++- internal/channels/telegram/handlers.go | 18 ++---- internal/channels/telegram/media.go | 67 +++++++++++++++++--- internal/channels/telegram/media_test.go | 54 ++++++++++++++++ internal/media/store.go | 6 ++ internal/media/store_test.go | 23 +++++++ internal/tools/read_document_resolve.go | 3 +- internal/tools/read_document_resolve_test.go | 19 +++++- 12 files changed, 217 insertions(+), 32 deletions(-) create mode 100644 internal/media/store_test.go diff --git a/docs/14-skills-runtime.md b/docs/14-skills-runtime.md index ac8fe740..7b91ff02 100644 --- a/docs/14-skills-runtime.md +++ b/docs/14-skills-runtime.md @@ -135,7 +135,7 @@ To install additional packages: pip3 install or npm install -g - Run Python/Node scripts via exec tool - Install packages via `pip3 install` / `npm install -g` -- Access files in `/app/workspace/` including `.media/` subdirectory +- Access files in `/app/workspace/`, including `.uploads/` for current user uploads and `.media/` for legacy media refs - Read skill files from `.goclaw/skills-store/` ### What Agents CANNOT Do @@ -152,16 +152,18 @@ To install additional packages: pip3 install or npm install -g Uploaded files (from web chat, Telegram, Discord, etc.) are persisted to: ``` -/app/workspace/.media/{sessionHash}/{uuid}.{ext} +/app/workspace/.uploads/{safe-original-name}-{8hex}.{ext} ``` +Uploads without a usable original filename fall back to `{uuid}.{ext}`. Legacy media refs may still resolve from `.media/{sessionHash}/{uuid}.{ext}`. + The `enrichDocumentPaths()` function injects the full path into `` tags: ``` - + ``` -Agents can read these files directly via exec — no copy to `/tmp` needed. +Agents can read these files directly via exec — no copy to `/tmp` needed. For archive uploads such as `.zip`, inspect or extract with commands like `unzip -l ""` or `unzip -q "" -d `. --- diff --git a/internal/agent/loop_input_media.go b/internal/agent/loop_input_media.go index fa7989eb..49cb3394 100644 --- a/internal/agent/loop_input_media.go +++ b/internal/agent/loop_input_media.go @@ -27,10 +27,11 @@ func isTextMime(mime string) bool { } // collectRefsByKind gathers MediaRefs of a given kind from message history -// (reverse order) and current-turn refs. Historical first, current last. +// in chronological order, then appends current-turn refs. The last ref is the +// newest document for read_document's omitted media_id fallback. func collectRefsByKind(messages []providers.Message, currentRefs []providers.MediaRef, kind string) []providers.MediaRef { var refs []providers.MediaRef - for i := len(messages) - 1; i >= 0; i-- { + for i := range messages { for _, ref := range messages[i].MediaRefs { if ref.Kind == kind { refs = append(refs, ref) diff --git a/internal/agent/media_persist_test.go b/internal/agent/media_persist_test.go index 297f4a94..b819d208 100644 --- a/internal/agent/media_persist_test.go +++ b/internal/agent/media_persist_test.go @@ -72,6 +72,14 @@ func TestPersistMedia_NamingScheme(t *testing.T) { wantSlug: "", wantExtPat: `\.pdf`, }, + { + name: "zip_archive_preserves_archive_extension", + filename: "codex.zip", + content: "zip bytes", + mime: "application/zip", + wantSlug: "codex", + wantExtPat: `\.zip`, + }, } var loop Loop diff --git a/internal/agent/media_test.go b/internal/agent/media_test.go index cd66173c..f5fd4144 100644 --- a/internal/agent/media_test.go +++ b/internal/agent/media_test.go @@ -73,6 +73,34 @@ func TestEnrichImageIDs_SkipsAlreadyEnriched(t *testing.T) { } } +func TestCollectRefsByKindOrdersOldestToNewestThenCurrent(t *testing.T) { + messages := []providers.Message{ + { + Role: "user", + Content: "old", + MediaRefs: []providers.MediaRef{{ID: "old-doc", Kind: "document"}}, + }, + { + Role: "user", + Content: "latest-history", + MediaRefs: []providers.MediaRef{{ID: "latest-history-doc", Kind: "document"}}, + }, + } + current := []providers.MediaRef{{ID: "current-doc", Kind: "document"}} + + got := collectRefsByKind(messages, current, "document") + want := []string{"old-doc", "latest-history-doc", "current-doc"} + + if len(got) != len(want) { + t.Fatalf("got %d refs, want %d", len(got), len(want)) + } + for i := range want { + if got[i].ID != want[i] { + t.Fatalf("ref[%d] = %q, want %q", i, got[i].ID, want[i]) + } + } +} + // testMediaStore creates a temporary media.Store for tests. func testMediaStore(t *testing.T) *media.Store { t.Helper() diff --git a/internal/channels/history.go b/internal/channels/history.go index 36692857..054d61a4 100644 --- a/internal/channels/history.go +++ b/internal/channels/history.go @@ -36,9 +36,11 @@ const ( // Stored in RAM only (not persisted to DB) — used by channels that defer media // download until the bot is actually mentioned (e.g. Telegram). type MediaRef struct { - Type string // "image", "video", "audio", "voice", "document", "animation" - FileID string // platform-specific file ID for lazy download - FileSize int64 // file size in bytes (0 if unknown) — used to skip large files + Type string // "image", "video", "audio", "voice", "document", "animation" + FileID string // platform-specific file ID for lazy download + FileSize int64 // file size in bytes (0 if unknown) — used to skip large files + FileName string // original filename, if the platform exposes one + ContentType string // MIME type, if the platform exposes one } // HistoryEntry represents a single tracked group message. diff --git a/internal/channels/telegram/handlers.go b/internal/channels/telegram/handlers.go index 7999103c..685d621f 100644 --- a/internal/channels/telegram/handlers.go +++ b/internal/channels/telegram/handlers.go @@ -513,14 +513,8 @@ func (c *Channel) handleMessage(ctx context.Context, update telego.Update) { // Resolve deferred media from history entries (lazy download). if histRefs := c.GroupHistory().CollectMediaRefs(localKey); len(histRefs) > 0 { histMedia, histErrors := c.resolveMediaRefs(ctx, histRefs) - for _, m := range histMedia { - mediaFiles = append(mediaFiles, bus.MediaFile{ - Path: m.FilePath, - MimeType: m.ContentType, - Filename: m.FileName, - }) - } if len(histMedia) > 0 { + mediaFiles = prependMediaInfoFiles(mediaFiles, histMedia) histTags := buildMediaTags(histMedia) annotated = "[Media from recent group messages — only analyze if user asks about them]\n" + histTags + "\n[/Media]\n\n" + annotated } @@ -577,12 +571,12 @@ func (c *Channel) handleMessage(ctx context.Context, update telego.Update) { // user sees typing indicator → first content appears directly. metadata := map[string]string{ - "message_id": fmt.Sprintf("%d", message.MessageID), - "user_id": fmt.Sprintf("%d", user.ID), + "message_id": fmt.Sprintf("%d", message.MessageID), + "user_id": fmt.Sprintf("%d", user.ID), tools.MetaUsername: user.Username, - "first_name": user.FirstName, - "is_group": fmt.Sprintf("%t", isGroup), - "local_key": localKey, + "first_name": user.FirstName, + "is_group": fmt.Sprintf("%t", isGroup), + "local_key": localKey, } if message.Chat.Title != "" { metadata[tools.MetaChatTitle] = message.Chat.Title diff --git a/internal/channels/telegram/media.go b/internal/channels/telegram/media.go index a7d92d02..1a1faa8c 100644 --- a/internal/channels/telegram/media.go +++ b/internal/channels/telegram/media.go @@ -14,6 +14,7 @@ import ( "github.com/mymmrac/telego" + "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/channels" "github.com/nextlevelbuilder/goclaw/internal/channels/media" "github.com/nextlevelbuilder/goclaw/internal/tools" @@ -359,6 +360,24 @@ func buildMediaTags(mediaList []MediaInfo) string { return media.BuildMediaTags(mediaList) } +func prependMediaInfoFiles(current []bus.MediaFile, history []MediaInfo) []bus.MediaFile { + if len(history) == 0 { + return current + } + ordered := make([]bus.MediaFile, 0, len(history)+len(current)) + for _, m := range history { + if m.FilePath == "" { + continue + } + ordered = append(ordered, bus.MediaFile{ + Path: m.FilePath, + MimeType: m.ContentType, + Filename: m.FileName, + }) + } + return append(ordered, current...) +} + // extractDocumentContent delegates to the shared media package. func extractDocumentContent(filePath, fileName string) (string, error) { return media.ExtractDocumentContent(filePath, fileName) @@ -412,22 +431,51 @@ func extractMediaRefs(msg *telego.Message) []channels.MediaRef { refs = append(refs, channels.MediaRef{Type: "image", FileID: photo.FileID, FileSize: int64(photo.FileSize)}) } if msg.Video != nil { - refs = append(refs, channels.MediaRef{Type: "video", FileID: msg.Video.FileID, FileSize: int64(msg.Video.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "video", + FileID: msg.Video.FileID, + FileSize: int64(msg.Video.FileSize), + FileName: msg.Video.FileName, + ContentType: msg.Video.MimeType, + }) } if msg.VideoNote != nil { refs = append(refs, channels.MediaRef{Type: "video", FileID: msg.VideoNote.FileID, FileSize: int64(msg.VideoNote.FileSize)}) } if msg.Animation != nil { - refs = append(refs, channels.MediaRef{Type: "animation", FileID: msg.Animation.FileID, FileSize: int64(msg.Animation.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "animation", + FileID: msg.Animation.FileID, + FileSize: int64(msg.Animation.FileSize), + FileName: msg.Animation.FileName, + ContentType: msg.Animation.MimeType, + }) } if msg.Audio != nil { - refs = append(refs, channels.MediaRef{Type: "audio", FileID: msg.Audio.FileID, FileSize: int64(msg.Audio.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "audio", + FileID: msg.Audio.FileID, + FileSize: int64(msg.Audio.FileSize), + FileName: msg.Audio.FileName, + ContentType: msg.Audio.MimeType, + }) } if msg.Voice != nil { - refs = append(refs, channels.MediaRef{Type: "voice", FileID: msg.Voice.FileID, FileSize: int64(msg.Voice.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "voice", + FileID: msg.Voice.FileID, + FileSize: int64(msg.Voice.FileSize), + ContentType: msg.Voice.MimeType, + }) } if msg.Document != nil { - refs = append(refs, channels.MediaRef{Type: "document", FileID: msg.Document.FileID, FileSize: int64(msg.Document.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "document", + FileID: msg.Document.FileID, + FileSize: int64(msg.Document.FileSize), + FileName: msg.Document.FileName, + ContentType: msg.Document.MimeType, + }) } return refs } @@ -492,9 +540,12 @@ func (c *Channel) resolveMediaRefs(ctx context.Context, refs []channels.MediaRef continue } results = append(results, MediaInfo{ - Type: ref.Type, - FilePath: filePath, - FileID: ref.FileID, + Type: ref.Type, + FilePath: filePath, + FileID: ref.FileID, + ContentType: ref.ContentType, + FileName: ref.FileName, + FileSize: ref.FileSize, }) } return results, errs diff --git a/internal/channels/telegram/media_test.go b/internal/channels/telegram/media_test.go index f568db34..a9822c70 100644 --- a/internal/channels/telegram/media_test.go +++ b/internal/channels/telegram/media_test.go @@ -3,6 +3,11 @@ package telegram import ( "strings" "testing" + + "github.com/mymmrac/telego" + + "github.com/nextlevelbuilder/goclaw/internal/bus" + "github.com/nextlevelbuilder/goclaw/internal/channels" ) // --- buildMediaTags tests --- @@ -148,3 +153,52 @@ func TestBuildMediaTags_UnknownType(t *testing.T) { t.Errorf("expected empty string for unknown type, got: %q", got) } } + +func TestExtractMediaRefsPreservesDocumentMetadata(t *testing.T) { + msg := &telego.Message{ + Document: &telego.Document{ + FileID: "telegram-file-id", + FileName: "codex.zip", + MimeType: "application/zip", + FileSize: 1234, + }, + } + + got := extractMediaRefs(msg) + want := []channels.MediaRef{{ + Type: "document", + FileID: "telegram-file-id", + FileSize: 1234, + FileName: "codex.zip", + ContentType: "application/zip", + }} + + if len(got) != len(want) { + t.Fatalf("got %d refs, want %d", len(got), len(want)) + } + if got[0] != want[0] { + t.Fatalf("ref = %+v, want %+v", got[0], want[0]) + } +} + +func TestPrependMediaInfoFilesKeepsHistoryBeforeCurrent(t *testing.T) { + current := []bus.MediaFile{{ + Path: "/workspace/.uploads/current.pdf", + MimeType: "application/pdf", + Filename: "current.pdf", + }} + history := []MediaInfo{{ + Type: "document", + FilePath: "/workspace/.uploads/codex.zip", + ContentType: "application/zip", + FileName: "codex.zip", + }} + + got := prependMediaInfoFiles(current, history) + if len(got) != 2 { + t.Fatalf("got %d files, want 2", len(got)) + } + if got[0].Path != history[0].FilePath || got[1].Path != current[0].Path { + t.Fatalf("order = [%q, %q], want history before current", got[0].Path, got[1].Path) + } +} diff --git a/internal/media/store.go b/internal/media/store.go index 162871b6..f9cada68 100644 --- a/internal/media/store.go +++ b/internal/media/store.go @@ -109,6 +109,12 @@ func ExtFromMime(mime string) string { return ".wav" case strings.HasPrefix(mime, "application/pdf"): return ".pdf" + case strings.HasPrefix(mime, "application/zip"), strings.HasPrefix(mime, "application/x-zip-compressed"): + return ".zip" + case strings.HasPrefix(mime, "application/x-tar"): + return ".tar" + case strings.HasPrefix(mime, "application/gzip"), strings.HasPrefix(mime, "application/x-gzip"): + return ".gz" case mime == "application/vnd.openxmlformats-officedocument.wordprocessingml.document": return ".docx" case mime == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": diff --git a/internal/media/store_test.go b/internal/media/store_test.go new file mode 100644 index 00000000..e8d1a12d --- /dev/null +++ b/internal/media/store_test.go @@ -0,0 +1,23 @@ +package media + +import "testing" + +func TestExtFromMimeArchiveTypes(t *testing.T) { + tests := []struct { + mime string + want string + }{ + {mime: "application/zip", want: ".zip"}, + {mime: "application/x-zip-compressed", want: ".zip"}, + {mime: "application/x-tar", want: ".tar"}, + {mime: "application/gzip", want: ".gz"}, + } + + for _, tt := range tests { + t.Run(tt.mime, func(t *testing.T) { + if got := ExtFromMime(tt.mime); got != tt.want { + t.Fatalf("ExtFromMime(%q) = %q, want %q", tt.mime, got, tt.want) + } + }) + } +} diff --git a/internal/tools/read_document_resolve.go b/internal/tools/read_document_resolve.go index 47d72cf9..a84b2870 100644 --- a/internal/tools/read_document_resolve.go +++ b/internal/tools/read_document_resolve.go @@ -43,8 +43,7 @@ func (t *ReadDocumentTool) resolveDocumentFile(ctx context.Context, mediaID, doc } } if ref == nil { - slog.Warn("read_document: media_id not found, falling back to most recent", "media_id", mediaID) - ref = &refs[len(refs)-1] + return "", "", fmt.Errorf("document media_id %q not found in this conversation", mediaID) } } else { // Use the last (most recent) document ref. diff --git a/internal/tools/read_document_resolve_test.go b/internal/tools/read_document_resolve_test.go index bc2ac937..172ac87f 100644 --- a/internal/tools/read_document_resolve_test.go +++ b/internal/tools/read_document_resolve_test.go @@ -66,7 +66,7 @@ func TestResolveDocumentFileMatchesUploadedFilenameAlias(t *testing.T) { } } -func TestResolveDocumentFileInvalidMediaIDFallsBackToMostRecent(t *testing.T) { +func TestResolveDocumentFileInvalidMediaIDReturnsError(t *testing.T) { refs := []providers.MediaRef{ {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/old.pdf", MimeType: "application/pdf"}, {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/latest.pdf", MimeType: "application/pdf"}, @@ -75,6 +75,23 @@ func TestResolveDocumentFileInvalidMediaIDFallsBackToMostRecent(t *testing.T) { tool := NewReadDocumentTool(nil, nil) ctx := WithMediaDocRefs(context.Background(), refs) gotPath, _, err := tool.resolveDocumentFile(ctx, "not-a-real-media-id", "") + if err == nil { + t.Fatalf("resolveDocumentFile returned path %q, want explicit media_id error", gotPath) + } + if !strings.Contains(err.Error(), "not-a-real-media-id") { + t.Fatalf("error = %q, want requested media_id", err.Error()) + } +} + +func TestResolveDocumentFileOmittedMediaIDUsesLastRef(t *testing.T) { + refs := []providers.MediaRef{ + {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/old.pdf", MimeType: "application/pdf"}, + {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/latest.pdf", MimeType: "application/pdf"}, + } + + tool := NewReadDocumentTool(nil, nil) + ctx := WithMediaDocRefs(context.Background(), refs) + gotPath, _, err := tool.resolveDocumentFile(ctx, "", "") if err != nil { t.Fatalf("resolveDocumentFile returned error: %v", err) } From 61fdcdcb8122a5a1b894d0da8533a8b7f8b7ed66 Mon Sep 17 00:00:00 2001 From: Duy /zuey/ Date: Sun, 24 May 2026 13:52:23 +0700 Subject: [PATCH 07/10] feat(tools): add command keyword allowlist Closes #42 --- cmd/gateway_lifecycle_shell_deny_groups.go | 6 +- ...ateway_lifecycle_shell_deny_groups_test.go | 31 +++ cmd/gateway_setup.go | 4 +- docs/03-tools-system.md | 28 ++- docs/project-changelog.md | 15 ++ internal/config/config_channels.go | 164 ++++++++------- internal/tools/command_keyword_allowlist.go | 187 ++++++++++++++++++ internal/tools/credentialed_exec.go | 17 +- internal/tools/credentialed_exec_test.go | 109 ++++++++++ internal/tools/shell.go | 65 ++++-- ui/web/src/i18n/locales/en/config.json | 4 + ui/web/src/i18n/locales/vi/config.json | 4 + ui/web/src/i18n/locales/zh/config.json | 4 + .../config/sections/tools-exec-section.tsx | 59 +++++- 14 files changed, 600 insertions(+), 97 deletions(-) create mode 100644 internal/tools/command_keyword_allowlist.go diff --git a/cmd/gateway_lifecycle_shell_deny_groups.go b/cmd/gateway_lifecycle_shell_deny_groups.go index 28c8be7f..46e0e0df 100644 --- a/cmd/gateway_lifecycle_shell_deny_groups.go +++ b/cmd/gateway_lifecycle_shell_deny_groups.go @@ -30,6 +30,10 @@ func subscribeShellDenyGroupsReload(msgBus *bus.MessageBus, toolsReg *tools.Regi return } et.SetGlobalShellDenyGroups(updatedCfg.Tools.ShellDenyGroups) - slog.Info("shell deny groups reloaded via pub/sub", "groups", len(updatedCfg.Tools.ShellDenyGroups)) + et.SetCommandKeywordAllowlist(updatedCfg.Tools.CommandKeywordAllowlist) + slog.Info("shell deny groups reloaded via pub/sub", + "groups", len(updatedCfg.Tools.ShellDenyGroups), + "command_keyword_allowlist_rules", len(updatedCfg.Tools.CommandKeywordAllowlist), + ) }) } diff --git a/cmd/gateway_lifecycle_shell_deny_groups_test.go b/cmd/gateway_lifecycle_shell_deny_groups_test.go index 13368dd6..e66be5c2 100644 --- a/cmd/gateway_lifecycle_shell_deny_groups_test.go +++ b/cmd/gateway_lifecycle_shell_deny_groups_test.go @@ -35,6 +35,37 @@ func TestShellDenyGroupsConfigReload_UpdatesGlobal(t *testing.T) { if v, ok := got["package_install"]; !ok || v != true { t.Fatalf("expected pub/sub to set global package_install=true, got %v", got) } + rules := execTool.CommandKeywordAllowlistForTest() + if len(rules) != 0 { + t.Fatalf("expected empty command keyword allowlist, got %v", rules) + } +} + +func TestShellDenyGroupsConfigReload_UpdatesCommandKeywordAllowlist(t *testing.T) { + msgBus := bus.New() + defer msgBus.Unsubscribe("shell-deny-groups-config-reload") + + toolsReg := tools.NewRegistry() + execTool := tools.NewExecTool("/tmp", false) + toolsReg.Register(execTool) + + subscribeShellDenyGroupsReload(msgBus, toolsReg) + + msgBus.Broadcast(bus.Event{ + Name: bus.TopicConfigChanged, + Payload: &config.Config{ + Tools: config.ToolsConfig{ + CommandKeywordAllowlist: []config.CommandKeywordAllowlistRule{ + {ID: "github-content", Command: "gh", Args: []string{"--body"}, Keywords: []string{"secret"}}, + }, + }, + }, + }) + + rules := execTool.CommandKeywordAllowlistForTest() + if len(rules) != 1 || rules[0].ID != "github-content" { + t.Fatalf("expected command keyword allowlist to reload, got %v", rules) + } } // TestShellDenyGroupsConfigReload_IgnoresOtherEvents: subscriber must guard diff --git a/cmd/gateway_setup.go b/cmd/gateway_setup.go index 491bd680..aba394db 100644 --- a/cmd/gateway_setup.go +++ b/cmd/gateway_setup.go @@ -14,11 +14,11 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bootstrap" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/edition" mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp" "github.com/nextlevelbuilder/goclaw/internal/permissions" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/sandbox" - "github.com/nextlevelbuilder/goclaw/internal/edition" "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/store/pg" @@ -214,6 +214,7 @@ func setupToolRegistry( // Apply global shell deny-group toggles before any request can arrive. // Per-agent overrides via store.WithShellDenyGroups still win per-key. et.SetGlobalShellDenyGroups(cfg.Tools.ShellDenyGroups) + et.SetCommandKeywordAllowlist(cfg.Tools.CommandKeywordAllowlist) et.DenyPaths(dataDir, ".goclaw/") // Allow skills execution: master-tenant skills-store + all tenant-scoped skills-store dirs. et.AllowPathExemptions( @@ -603,4 +604,3 @@ func setupSkillsSystem( return skillsLoader, skillSearchTool, globalSkillsDir, bundledSkillsDir, builtinSkillsDir } - diff --git a/docs/03-tools-system.md b/docs/03-tools-system.md index ea7a5d2d..c7940b6d 100644 --- a/docs/03-tools-system.md +++ b/docs/03-tools-system.md @@ -246,6 +246,32 @@ User-facing parameter schemas for the most commonly configured tools. ``` Available presets: `gh`, `gcloud`, `aws`, `kubectl`, `terraform`. +### Credentialed CLI keyword allowlist + +`config.tools.commandKeywordAllowlist` lets operators allow specific product or security vocabulary inside selected credentialed CLI content arguments without disabling `deny_args`. + +Example: + +```json +{ + "tools": { + "commandKeywordAllowlist": [ + { + "id": "github-content", + "command": "gh", + "subcommands": ["issue create", "issue edit", "pr create", "pr comment"], + "args": ["--body", "--title"], + "argPositions": [], + "keywords": ["secret", "secrets", "token", "credential"], + "reason": "Allow security vocabulary in GitHub issue and PR prose" + } + ] + } +} +``` + +The rule above allows `gh issue create --body "secret rotation notes"` but still blocks command paths like `gh secret set TOKEN`. `argPositions` are 0-based after the matched subcommand. The scanner evaluates command arguments only; it does not read the contents of files passed through arguments such as `--body-file`. + --- ## 6. Interception Layer @@ -406,7 +432,7 @@ Current adopters: `web_search`, `web_fetch`, `tts`, `create_image`, `read_image` - Per-key: agent value takes precedence over global value - Multi-tenant invariant: each tenant's config is isolated -**Live reload:** Changes to `config.tools.shellDenyGroups` propagate via `bus.TopicConfigChanged` pub/sub. Next agent turn automatically applies new toggles. +**Live reload:** Changes to `config.tools.shellDenyGroups` and `config.tools.commandKeywordAllowlist` propagate via `bus.TopicConfigChanged` pub/sub. Next agent turn automatically applies new toggles. **Deny-group classes** (from `internal/tools/shell_deny_groups.go` — all denied by default): diff --git a/docs/project-changelog.md b/docs/project-changelog.md index 68fcb9fd..1e234716 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -6,6 +6,21 @@ Significant changes, features, and fixes in reverse chronological order. ## 2026-05-24 +### Command keyword allowlist + +**Features** + +- Added scoped credentialed CLI keyword allowlist config for content arguments and positional arguments, with runtime reload and web config editing. + +**Fixes** + +- Kept credentialed CLI `deny_args` active for command paths such as `gh secret set` while allowing approved GitHub issue/PR prose to mention security vocabulary. +- Added security audit logging for real allowlisted pass-throughs without logging full argument values. + +**Tests** + +- Added regression coverage for scoped keyword masking, disabled rules, unsafe positional rules, config reload, and race-safe policy snapshots. + ### Browser cookie sync and config UI **Features** diff --git a/internal/config/config_channels.go b/internal/config/config_channels.go index a125d50e..6a9c5c70 100644 --- a/internal/config/config_channels.go +++ b/internal/config/config_channels.go @@ -24,25 +24,25 @@ type ChannelsConfig struct { } type TelegramConfig struct { - Enabled bool `json:"enabled"` - Token string `json:"token"` - Proxy string `json:"proxy,omitempty"` - APIServer string `json:"api_server,omitempty"` // custom Telegram Bot API server URL (e.g. "http://localhost:8081") - AllowFrom FlexibleStringSlice `json:"allow_from"` - DMPolicy string `json:"dm_policy,omitempty"` // "pairing" (default), "allowlist", "open", "disabled" - GroupPolicy string `json:"group_policy,omitempty"` // "open" (default), "allowlist", "disabled" - RequireMention *bool `json:"require_mention,omitempty"` // require @bot mention in groups (default true) - MentionMode string `json:"mention_mode,omitempty"` // "strict" (default) = only respond when mentioned; "yield" = respond unless another bot is mentioned - HistoryLimit int `json:"history_limit,omitempty"` // max pending group messages for context (default 50, 0=disabled) - DMStream *bool `json:"dm_stream,omitempty"` // enable streaming for DMs (default false) — edits placeholder progressively - GroupStream *bool `json:"group_stream,omitempty"` // enable streaming for groups (default false) — sends new message, edits progressively - DraftTransport *bool `json:"draft_transport,omitempty"` // use sendMessageDraft for DM streaming (default true) — stealth preview, no notifications per edit - ReasoningStream *bool `json:"reasoning_stream,omitempty"` // show reasoning as separate message when provider emits thinking events (default true) - ReactionLevel string `json:"reaction_level,omitempty"` // "off" (default), "minimal", "full" — status emoji reactions - MediaMaxBytes int64 `json:"media_max_bytes,omitempty"` // max media download size in bytes (default 20MB) - LinkPreview *bool `json:"link_preview,omitempty"` // enable URL previews in messages (default true) - BlockReply *bool `json:"block_reply,omitempty"` // override gateway block_reply (nil = inherit) - ForceIPv4 bool `json:"force_ipv4,omitempty"` // force IPv4 for all Telegram API requests (use when IPv6 routing is broken) + Enabled bool `json:"enabled"` + Token string `json:"token"` + Proxy string `json:"proxy,omitempty"` + APIServer string `json:"api_server,omitempty"` // custom Telegram Bot API server URL (e.g. "http://localhost:8081") + AllowFrom FlexibleStringSlice `json:"allow_from"` + DMPolicy string `json:"dm_policy,omitempty"` // "pairing" (default), "allowlist", "open", "disabled" + GroupPolicy string `json:"group_policy,omitempty"` // "open" (default), "allowlist", "disabled" + RequireMention *bool `json:"require_mention,omitempty"` // require @bot mention in groups (default true) + MentionMode string `json:"mention_mode,omitempty"` // "strict" (default) = only respond when mentioned; "yield" = respond unless another bot is mentioned + HistoryLimit int `json:"history_limit,omitempty"` // max pending group messages for context (default 50, 0=disabled) + DMStream *bool `json:"dm_stream,omitempty"` // enable streaming for DMs (default false) — edits placeholder progressively + GroupStream *bool `json:"group_stream,omitempty"` // enable streaming for groups (default false) — sends new message, edits progressively + DraftTransport *bool `json:"draft_transport,omitempty"` // use sendMessageDraft for DM streaming (default true) — stealth preview, no notifications per edit + ReasoningStream *bool `json:"reasoning_stream,omitempty"` // show reasoning as separate message when provider emits thinking events (default true) + ReactionLevel string `json:"reaction_level,omitempty"` // "off" (default), "minimal", "full" — status emoji reactions + MediaMaxBytes int64 `json:"media_max_bytes,omitempty"` // max media download size in bytes (default 20MB) + LinkPreview *bool `json:"link_preview,omitempty"` // enable URL previews in messages (default true) + BlockReply *bool `json:"block_reply,omitempty"` // override gateway block_reply (nil = inherit) + ForceIPv4 bool `json:"force_ipv4,omitempty"` // force IPv4 for all Telegram API requests (use when IPv6 routing is broken) // Optional STT (Speech-to-Text) pipeline for voice/audio inbound messages. // When stt_proxy_url is set, audio/voice messages are transcribed before being forwarded to the agent. @@ -133,7 +133,7 @@ type SlackConfig struct { type WhatsAppConfig struct { Enabled bool `json:"enabled"` - AuthDir string `json:"auth_dir,omitempty"` // optional: SQLite auth dir override (desktop) + AuthDir string `json:"auth_dir,omitempty"` // optional: SQLite auth dir override (desktop) AllowFrom FlexibleStringSlice `json:"allow_from"` DMPolicy string `json:"dm_policy,omitempty"` // "pairing" (default for DB instances), "open", "allowlist", "disabled" GroupPolicy string `json:"group_policy,omitempty"` // "pairing" (default for DB instances), "open" (default for config), "allowlist", "disabled" @@ -196,25 +196,25 @@ type FeishuConfig struct { // ProvidersConfig maps provider name to its config. type ProvidersConfig struct { - Anthropic ProviderConfig `json:"anthropic"` - OpenAI ProviderConfig `json:"openai"` - OpenRouter ProviderConfig `json:"openrouter"` - Groq ProviderConfig `json:"groq"` - Gemini ProviderConfig `json:"gemini"` - DeepSeek ProviderConfig `json:"deepseek"` - Mistral ProviderConfig `json:"mistral"` - XAI ProviderConfig `json:"xai"` - MiniMax ProviderConfig `json:"minimax"` - Cohere ProviderConfig `json:"cohere"` - Perplexity ProviderConfig `json:"perplexity"` - DashScope ProviderConfig `json:"dashscope"` - Bailian ProviderConfig `json:"bailian"` - Zai ProviderConfig `json:"zai"` - ZaiCoding ProviderConfig `json:"zai_coding"` - Ollama OllamaConfig `json:"ollama"` // local Ollama instance (no API key needed) - OllamaCloud ProviderConfig `json:"ollama_cloud"` // Ollama Cloud (API key required) - ClaudeCLI ClaudeCLIConfig `json:"claude_cli"` - ACP ACPConfig `json:"acp"` + Anthropic ProviderConfig `json:"anthropic"` + OpenAI ProviderConfig `json:"openai"` + OpenRouter ProviderConfig `json:"openrouter"` + Groq ProviderConfig `json:"groq"` + Gemini ProviderConfig `json:"gemini"` + DeepSeek ProviderConfig `json:"deepseek"` + Mistral ProviderConfig `json:"mistral"` + XAI ProviderConfig `json:"xai"` + MiniMax ProviderConfig `json:"minimax"` + Cohere ProviderConfig `json:"cohere"` + Perplexity ProviderConfig `json:"perplexity"` + DashScope ProviderConfig `json:"dashscope"` + Bailian ProviderConfig `json:"bailian"` + Zai ProviderConfig `json:"zai"` + ZaiCoding ProviderConfig `json:"zai_coding"` + Ollama OllamaConfig `json:"ollama"` // local Ollama instance (no API key needed) + OllamaCloud ProviderConfig `json:"ollama_cloud"` // Ollama Cloud (API key required) + ClaudeCLI ClaudeCLIConfig `json:"claude_cli"` + ACP ACPConfig `json:"acp"` Novita ProviderConfig `json:"novita"` // Novita AI (OpenAI-compatible endpoint) BytePlus ProviderConfig `json:"byteplus"` // BytePlus ModelArk (Seed 2.0) BytePlusCoding ProviderConfig `json:"byteplus_coding"` // BytePlus ModelArk Coding Plan @@ -362,16 +362,16 @@ type QuotaConfig struct { // GatewayConfig controls the gateway server. type GatewayConfig struct { - Host string `json:"host"` - Port int `json:"port"` - Token string `json:"token,omitempty"` // bearer token for WS/HTTP auth - OwnerIDs []string `json:"owner_ids,omitempty"` // sender IDs considered "owner" - AllowedOrigins []string `json:"allowed_origins,omitempty"` // WebSocket CORS whitelist (empty = allow all) - MaxMessageChars int `json:"max_message_chars,omitempty"` // max user message characters (default 32000) - RateLimitRPM int `json:"rate_limit_rpm,omitempty"` // rate limit: requests per minute per user (default 20, 0 = disabled) - InjectionAction string `json:"injection_action,omitempty"` // prompt injection action: "log", "warn" (default), "block", "off" - InboundDebounceMs int `json:"inbound_debounce_ms,omitempty"` // merge rapid channel/Web Chat messages from same sender/session (0 = no wait) - Quota *QuotaConfig `json:"quota,omitempty"` // per-user/group request quotas + Host string `json:"host"` + Port int `json:"port"` + Token string `json:"token,omitempty"` // bearer token for WS/HTTP auth + OwnerIDs []string `json:"owner_ids,omitempty"` // sender IDs considered "owner" + AllowedOrigins []string `json:"allowed_origins,omitempty"` // WebSocket CORS whitelist (empty = allow all) + MaxMessageChars int `json:"max_message_chars,omitempty"` // max user message characters (default 32000) + RateLimitRPM int `json:"rate_limit_rpm,omitempty"` // rate limit: requests per minute per user (default 20, 0 = disabled) + InjectionAction string `json:"injection_action,omitempty"` // prompt injection action: "log", "warn" (default), "block", "off" + InboundDebounceMs int `json:"inbound_debounce_ms,omitempty"` // merge rapid channel/Web Chat messages from same sender/session (0 = no wait) + Quota *QuotaConfig `json:"quota,omitempty"` // per-user/group request quotas BlockReply *bool `json:"block_reply,omitempty"` // deliver intermediate text during tool iterations (default false) ToolStatus *bool `json:"tool_status,omitempty"` // show tool name in streaming preview during tool execution (default true) TaskRecoveryIntervalSec int `json:"task_recovery_interval_sec,omitempty"` // team task recovery ticker interval in seconds (default 300 = 5min) @@ -381,18 +381,32 @@ type GatewayConfig struct { // ToolsConfig controls tool availability, policy, and web search. type ToolsConfig struct { - Profile string `json:"profile,omitempty"` // global profile: "minimal", "coding", "messaging", "full" - Allow []string `json:"allow,omitempty"` // global allow list (tool names or "group:xxx") - Deny []string `json:"deny,omitempty"` // global deny list - AlsoAllow []string `json:"alsoAllow,omitempty"` // additive: adds without removing existing - ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` // per-provider overrides - ShellDenyGroups map[string]bool `json:"shellDenyGroups,omitempty"` // global shell deny-group toggles (group name -> denied); per-agent overrides win per-key - ExecApproval ExecApprovalCfg `json:"execApproval"` // exec command approval settings - WebFetch WebFetchPolicyConfig `json:"web_fetch"` // domain policy for URL fetching - Browser BrowserToolConfig `json:"browser"` - RateLimitPerHour int `json:"rate_limit_per_hour,omitempty"` // max tool executions per hour per session (0 = disabled) - ScrubCredentials *bool `json:"scrub_credentials,omitempty"` // auto-redact API keys/tokens in tool output (default true) - McpServers map[string]*MCPServerConfig `json:"mcp_servers,omitempty"` // external MCP server connections + Profile string `json:"profile,omitempty"` // global profile: "minimal", "coding", "messaging", "full" + Allow []string `json:"allow,omitempty"` // global allow list (tool names or "group:xxx") + Deny []string `json:"deny,omitempty"` // global deny list + AlsoAllow []string `json:"alsoAllow,omitempty"` // additive: adds without removing existing + ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` // per-provider overrides + ShellDenyGroups map[string]bool `json:"shellDenyGroups,omitempty"` // global shell deny-group toggles (group name -> denied); per-agent overrides win per-key + CommandKeywordAllowlist []CommandKeywordAllowlistRule `json:"commandKeywordAllowlist,omitempty"` // scoped bypass for credentialed CLI content keywords + ExecApproval ExecApprovalCfg `json:"execApproval"` // exec command approval settings + WebFetch WebFetchPolicyConfig `json:"web_fetch"` // domain policy for URL fetching + Browser BrowserToolConfig `json:"browser"` + RateLimitPerHour int `json:"rate_limit_per_hour,omitempty"` // max tool executions per hour per session (0 = disabled) + ScrubCredentials *bool `json:"scrub_credentials,omitempty"` // auto-redact API keys/tokens in tool output (default true) + McpServers map[string]*MCPServerConfig `json:"mcp_servers,omitempty"` // external MCP server connections +} + +// CommandKeywordAllowlistRule scopes product/security vocabulary that may appear +// in credentialed CLI content args without disabling command-path deny patterns. +type CommandKeywordAllowlistRule struct { + ID string `json:"id,omitempty"` + Command string `json:"command"` + Subcommands []string `json:"subcommands,omitempty"` + Args []string `json:"args,omitempty"` + ArgPositions []int `json:"argPositions,omitempty"` // 0-based after matched subcommand + Keywords []string `json:"keywords,omitempty"` + Reason string `json:"reason,omitempty"` + Enabled *bool `json:"enabled,omitempty"` // nil means enabled } // MCPServerConfig configures a single external MCP server connection. @@ -429,24 +443,24 @@ type WebFetchPolicyConfig struct { // BrowserToolConfig controls the browser automation tool. type BrowserToolConfig struct { - Enabled bool `json:"enabled"` // enable the browser tool (default false) - Headless bool `json:"headless,omitempty"` // run Chrome in headless mode (ignored when RemoteURL is set) - RemoteURL string `json:"remote_url,omitempty"` // CDP endpoint for remote Chrome sidecar, e.g. "ws://chrome:9222" - ActionTimeoutMs int `json:"action_timeout_ms,omitempty"` // per-action timeout in ms (default 30000) - IdleTimeoutMs int `json:"idle_timeout_ms,omitempty"` // idle page auto-close in ms (default 600000, 0=disabled) - MaxPages int `json:"max_pages,omitempty"` // max open pages per tenant (default 5) - CookieSyncEnabled bool `json:"cookie_sync_enabled"` // apply selected synced cookies to scoped browser sessions + Enabled bool `json:"enabled"` // enable the browser tool (default false) + Headless bool `json:"headless,omitempty"` // run Chrome in headless mode (ignored when RemoteURL is set) + RemoteURL string `json:"remote_url,omitempty"` // CDP endpoint for remote Chrome sidecar, e.g. "ws://chrome:9222" + ActionTimeoutMs int `json:"action_timeout_ms,omitempty"` // per-action timeout in ms (default 30000) + IdleTimeoutMs int `json:"idle_timeout_ms,omitempty"` // idle page auto-close in ms (default 600000, 0=disabled) + MaxPages int `json:"max_pages,omitempty"` // max open pages per tenant (default 5) + CookieSyncEnabled bool `json:"cookie_sync_enabled"` // apply selected synced cookies to scoped browser sessions } // ToolPolicySpec defines a tool policy at any level (global, per-agent, per-provider). type ToolPolicySpec struct { - Profile string `json:"profile,omitempty"` - Allow []string `json:"allow,omitempty"` - Deny []string `json:"deny,omitempty"` - AlsoAllow []string `json:"alsoAllow,omitempty"` - ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` - Wait *WaitToolPolicy `json:"wait,omitempty"` - ToolCallPrefix string `json:"toolCallPrefix,omitempty"` // prefix to strip from model's tool call names before registry lookup + Profile string `json:"profile,omitempty"` + Allow []string `json:"allow,omitempty"` + Deny []string `json:"deny,omitempty"` + AlsoAllow []string `json:"alsoAllow,omitempty"` + ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` + Wait *WaitToolPolicy `json:"wait,omitempty"` + ToolCallPrefix string `json:"toolCallPrefix,omitempty"` // prefix to strip from model's tool call names before registry lookup } // WaitToolPolicy configures per-agent safety bounds for the wait tool. diff --git a/internal/tools/command_keyword_allowlist.go b/internal/tools/command_keyword_allowlist.go new file mode 100644 index 00000000..b03490eb --- /dev/null +++ b/internal/tools/command_keyword_allowlist.go @@ -0,0 +1,187 @@ +package tools + +import ( + "regexp" + "strconv" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/config" +) + +type commandKeywordAllowAudit struct { + RuleID string + Command string + Subcommand string + Arg string + Keyword string + Reason string +} + +func applyCommandKeywordAllowlist(command string, args []string, rules []config.CommandKeywordAllowlistRule) ([]string, []commandKeywordAllowAudit) { + if len(args) == 0 || len(rules) == 0 { + return args, nil + } + + var out []string + var audits []commandKeywordAllowAudit + for _, rule := range rules { + if !commandKeywordAllowlistRuleEnabled(rule) { + continue + } + if normalizeBinaryName(rule.Command) != normalizeBinaryName(command) { + continue + } + subcommand, argStart, ok := matchCommandKeywordSubcommand(args, rule.Subcommands) + if !ok { + continue + } + argNames := commandKeywordSet(rule.Args) + argPositions := commandKeywordPositionSet(rule.ArgPositions) + keywords := commandKeywordSet(rule.Keywords) + if (len(argNames) == 0 && len(argPositions) == 0) || len(keywords) == 0 { + continue + } + if len(argPositions) > 0 && subcommand == "" { + continue + } + if out == nil { + out = slicesClone(args) + } + for i := argStart; i < len(out); i++ { + argName, valueIndex, inlineValue, ok := commandKeywordArgValue(out, i, argNames) + if !ok { + if _, positional := argPositions[i-argStart]; positional { + argName = "argPositions[" + strconv.Itoa(i-argStart) + "]" + valueIndex = i + inlineValue = out[i] + ok = true + } + } + if !ok { + continue + } + value := inlineValue + if valueIndex != i { + value = out[valueIndex] + } + masked, hits := maskCommandKeywords(value, keywords) + if len(hits) == 0 { + continue + } + if valueIndex == i { + prefix := out[i][:strings.Index(out[i], "=")+1] + out[i] = prefix + masked + } else { + out[valueIndex] = masked + i = valueIndex + } + for _, keyword := range hits { + audits = append(audits, commandKeywordAllowAudit{ + RuleID: rule.ID, + Command: normalizeBinaryName(command), + Subcommand: subcommand, + Arg: argName, + Keyword: keyword, + Reason: rule.Reason, + }) + } + } + } + if out == nil { + return args, audits + } + return out, audits +} + +func commandKeywordAllowlistRuleEnabled(rule config.CommandKeywordAllowlistRule) bool { + return rule.Enabled == nil || *rule.Enabled +} + +func matchCommandKeywordSubcommand(args []string, subcommands []string) (string, int, bool) { + if len(subcommands) == 0 { + return "", 0, true + } + for _, subcommand := range subcommands { + parts := strings.Fields(strings.ToLower(strings.TrimSpace(subcommand))) + if len(parts) == 0 || len(parts) > len(args) { + continue + } + matched := true + for i, part := range parts { + if strings.ToLower(args[i]) != part { + matched = false + break + } + } + if matched { + return strings.Join(parts, " "), len(parts), true + } + } + return "", 0, false +} + +func commandKeywordSet(values []string) map[string]struct{} { + if len(values) == 0 { + return nil + } + set := make(map[string]struct{}, len(values)) + for _, value := range values { + value = strings.ToLower(strings.TrimSpace(value)) + if value != "" { + set[value] = struct{}{} + } + } + return set +} + +func commandKeywordPositionSet(values []int) map[int]struct{} { + if len(values) == 0 { + return nil + } + set := make(map[int]struct{}, len(values)) + for _, value := range values { + if value >= 0 { + set[value] = struct{}{} + } + } + return set +} + +func commandKeywordArgValue(args []string, i int, argNames map[string]struct{}) (string, int, string, bool) { + arg := args[i] + if eq := strings.Index(arg, "="); eq > 0 { + name := strings.ToLower(arg[:eq]) + if _, ok := argNames[name]; ok { + return name, i, arg[eq+1:], true + } + return "", 0, "", false + } + name := strings.ToLower(arg) + if _, ok := argNames[name]; !ok || i+1 >= len(args) { + return "", 0, "", false + } + return name, i + 1, "", true +} + +func maskCommandKeywords(value string, keywords map[string]struct{}) (string, []string) { + if value == "" || len(keywords) == 0 { + return value, nil + } + var hits []string + masked := value + for keyword := range keywords { + re := regexp.MustCompile(`(?i)\b` + regexp.QuoteMeta(keyword) + `\b`) + if !re.MatchString(masked) { + continue + } + hits = append(hits, keyword) + masked = re.ReplaceAllString(masked, "__allowlisted_keyword__") + } + return masked, hits +} + +func slicesClone[T any](src []T) []T { + dst := make([]T, len(src)) + copy(dst, src) + return dst +} diff --git a/internal/tools/credentialed_exec.go b/internal/tools/credentialed_exec.go index b811ca22..a15b5dbc 100644 --- a/internal/tools/credentialed_exec.go +++ b/internal/tools/credentialed_exec.go @@ -381,7 +381,8 @@ func (t *ExecTool) executeCredentialed(ctx context.Context, cred *store.SecureCL } // Step 3: Per-binary deny check (deny_args) - if p := matchesBinaryDeny(args, cred.DenyArgs); p != "" { + denyArgs, allowAudits := applyCommandKeywordAllowlist(binary, args, t.commandKeywordAllowlistSnapshot()) + if p := matchesBinaryDeny(denyArgs, cred.DenyArgs); p != "" { return credentialedDenyError(binary, args, p) } // Per-binary verbose deny check (deny_verbose) — per-arg start-anchored match @@ -389,6 +390,20 @@ func (t *ExecTool) executeCredentialed(ctx context.Context, cred *store.SecureCL if p := matchesBinaryVerbose(args, cred.DenyVerbose); p != "" { return credentialedDenyError(binary, args, p) } + for _, audit := range allowAudits { + slog.Info("security.command_keyword_allowlist", + "binary", audit.Command, + "subcommand", audit.Subcommand, + "arg", audit.Arg, + "keyword", audit.Keyword, + "rule_id", audit.RuleID, + "reason", audit.Reason, + "agent_id", store.AgentIDFromContext(ctx), + "user_id", store.UserIDFromContext(ctx), + "credential_user_id", store.CredentialUserIDFromContext(ctx), + "tenant_id", store.TenantIDFromContext(ctx), + ) + } // Step 4: Decrypt env vars from store (already decrypted by store layer). // Per-user env overrides take priority over binary/grant env. diff --git a/internal/tools/credentialed_exec_test.go b/internal/tools/credentialed_exec_test.go index 484d3dbe..6390dc18 100644 --- a/internal/tools/credentialed_exec_test.go +++ b/internal/tools/credentialed_exec_test.go @@ -5,6 +5,8 @@ import ( "os" "path/filepath" "testing" + + "github.com/nextlevelbuilder/goclaw/internal/config" ) func TestDetectShellOperators(t *testing.T) { @@ -208,6 +210,113 @@ func TestMatchesBinaryDenyJoinedArgs(t *testing.T) { } } +func TestApplyCommandKeywordAllowlistScopesContentArgs(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`auth\s+`, `repo\s+delete`, `secret\s+`, `token\s+`}) + rules := []config.CommandKeywordAllowlistRule{ + { + ID: "github-content", + Command: "gh", + Subcommands: []string{"issue create", "pr create"}, + Args: []string{"--body", "--title"}, + ArgPositions: []int{0}, + Keywords: []string{"secret", "token"}, + Reason: "GitHub issue and PR prose may discuss security terms.", + }, + } + + tests := []struct { + name string + args []string + wantHit bool + wantAudit int + }{ + { + name: "issue body content allowed", + args: []string{"issue", "create", "--body", "secret rotation details"}, + wantHit: false, + wantAudit: 1, + }, + { + name: "pr title content allowed", + args: []string{"pr", "create", "--title=token handling notes"}, + wantHit: false, + wantAudit: 1, + }, + { + name: "positional content allowed", + args: []string{"issue", "create", "token handling notes"}, + wantHit: false, + wantAudit: 1, + }, + { + name: "command path stays blocked", + args: []string{"secret", "set", "TOKEN"}, + wantHit: true, + }, + { + name: "non-allowlisted arg stays blocked", + args: []string{"issue", "create", "--label", "secret incident"}, + wantHit: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sanitized, audits := applyCommandKeywordAllowlist("gh", tt.args, rules) + got := matchesBinaryDeny(sanitized, ghPatterns) + if (got != "") != tt.wantHit { + t.Fatalf("matchesBinaryDeny(%v) after allowlist = %q, wantHit=%v", sanitized, got, tt.wantHit) + } + if len(audits) != tt.wantAudit { + t.Fatalf("audit count = %d, want %d", len(audits), tt.wantAudit) + } + }) + } +} + +func TestApplyCommandKeywordAllowlistIgnoresDisabledRules(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`secret\s+`}) + enabled := false + rules := []config.CommandKeywordAllowlistRule{ + { + ID: "disabled-github-content", + Command: "gh", + Subcommands: []string{"issue create"}, + Args: []string{"--body"}, + Keywords: []string{"secret"}, + Enabled: &enabled, + }, + } + + sanitized, audits := applyCommandKeywordAllowlist("gh", []string{"issue", "create", "--body", "secret notes"}, rules) + if got := matchesBinaryDeny(sanitized, ghPatterns); got == "" { + t.Fatalf("disabled rule bypassed deny_args; sanitized args = %v", sanitized) + } + if len(audits) != 0 { + t.Fatalf("disabled rule emitted audit records: %v", audits) + } +} + +func TestApplyCommandKeywordAllowlistRequiresSubcommandForPositions(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`secret\s+`}) + rules := []config.CommandKeywordAllowlistRule{ + { + ID: "unsafe-position", + Command: "gh", + ArgPositions: []int{0}, + Keywords: []string{"secret"}, + }, + } + + sanitized, audits := applyCommandKeywordAllowlist("gh", []string{"secret", "set", "TOKEN"}, rules) + if got := matchesBinaryDeny(sanitized, ghPatterns); got == "" { + t.Fatalf("position rule without subcommand bypassed command-path deny; sanitized args = %v", sanitized) + } + if len(audits) != 0 { + t.Fatalf("position rule without subcommand emitted audit records: %v", audits) + } +} + func TestResolveAndMatchBinaryUsesConfiguredExecutablePath(t *testing.T) { t.Setenv("PATH", "/usr/bin") binDir := t.TempDir() diff --git a/internal/tools/shell.go b/internal/tools/shell.go index 6be167cd..4e07182f 100644 --- a/internal/tools/shell.go +++ b/internal/tools/shell.go @@ -13,8 +13,10 @@ import ( "regexp" "runtime" "strings" + "sync" "time" + "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/sandbox" "github.com/nextlevelbuilder/goclaw/internal/store" "golang.org/x/text/unicode/norm" @@ -35,27 +37,56 @@ func DefaultDenyPatterns() []*regexp.Regexp { // ExecTool executes shell commands, optionally inside a sandbox container. type ExecTool struct { - workspace string - timeout time.Duration - pathDenyPatterns []*regexp.Regexp // always-on path-based denials (DenyPaths) - pathDenyRoots []string // raw deny roots for nested workspace exemptions - denyExemptions []string // substrings that exempt a command from deny - restrict bool - sandboxMgr sandbox.Manager // nil = no sandbox, execute on host - approvalMgr *ExecApprovalManager // nil = no approval needed - agentID string // for approval request context - secureCLIStore store.SecureCLIStore // nil = no credentialed exec + workspace string + timeout time.Duration + pathDenyPatterns []*regexp.Regexp // always-on path-based denials (DenyPaths) + pathDenyRoots []string // raw deny roots for nested workspace exemptions + denyExemptions []string // substrings that exempt a command from deny + restrict bool + sandboxMgr sandbox.Manager // nil = no sandbox, execute on host + approvalMgr *ExecApprovalManager // nil = no approval needed + agentID string // for approval request context + secureCLIStore store.SecureCLIStore // nil = no credentialed exec + policyMu sync.RWMutex + commandKeywordAllowlist []config.CommandKeywordAllowlistRule // globalDenyGroups holds global shell deny-group toggles from config.tools. // Per-agent overrides from context (store.WithShellDenyGroups) win per-key. // Updated at startup and via TopicConfigChanged pub/sub for runtime reload. globalDenyGroups map[string]bool } +// SetCommandKeywordAllowlist replaces the scoped credentialed CLI keyword +// allowlist. The slice is defensively copied so config reload callers cannot +// mutate the running tool policy after assignment. +func (t *ExecTool) SetCommandKeywordAllowlist(rules []config.CommandKeywordAllowlistRule) { + t.policyMu.Lock() + defer t.policyMu.Unlock() + if len(rules) == 0 { + t.commandKeywordAllowlist = nil + return + } + t.commandKeywordAllowlist = slicesClone(rules) +} + +// CommandKeywordAllowlistForTest exposes the effective global allowlist for +// cross-package config reload tests. Not for production callers. +func (t *ExecTool) CommandKeywordAllowlistForTest() []config.CommandKeywordAllowlistRule { + return t.commandKeywordAllowlistSnapshot() +} + +func (t *ExecTool) commandKeywordAllowlistSnapshot() []config.CommandKeywordAllowlistRule { + t.policyMu.RLock() + defer t.policyMu.RUnlock() + return slicesClone(t.commandKeywordAllowlist) +} + // SetGlobalShellDenyGroups replaces the global shell deny-group toggles. The // caller's map is defensively copied so later mutations cannot leak into the // tool's internal state. Passing nil or an empty map clears the global config // (per-agent context overrides, if any, still apply on their own). func (t *ExecTool) SetGlobalShellDenyGroups(groups map[string]bool) { + t.policyMu.Lock() + defer t.policyMu.Unlock() if len(groups) == 0 { t.globalDenyGroups = nil return @@ -70,14 +101,20 @@ func (t *ExecTool) SetGlobalShellDenyGroups(groups map[string]bool) { // empty, the other is returned directly (no allocation). func (t *ExecTool) effectiveDenyGroups(ctx context.Context) map[string]bool { agent := store.ShellDenyGroupsFromContext(ctx) - if len(t.globalDenyGroups) == 0 { + t.policyMu.RLock() + global := t.globalDenyGroups + if len(global) > 0 { + global = maps.Clone(global) + } + t.policyMu.RUnlock() + if len(global) == 0 { return agent } if len(agent) == 0 { - return t.globalDenyGroups + return global } - merged := make(map[string]bool, len(t.globalDenyGroups)+len(agent)) - maps.Copy(merged, t.globalDenyGroups) + merged := make(map[string]bool, len(global)+len(agent)) + maps.Copy(merged, global) // agent wins per-key maps.Copy(merged, agent) return merged diff --git a/ui/web/src/i18n/locales/en/config.json b/ui/web/src/i18n/locales/en/config.json index a1059e77..f22f94bc 100644 --- a/ui/web/src/i18n/locales/en/config.json +++ b/ui/web/src/i18n/locales/en/config.json @@ -208,6 +208,10 @@ "tools.execAskMode": "Ask Mode", "tools.execAskModeTip": "When to prompt the user for approval before executing a command.", "tools.execAllowlistLabel": "Allowed Commands (one pattern per line)", + "tools.commandKeywordAllowlist": "Command Keyword Allowlist", + "tools.commandKeywordAllowlistTip": "JSON rules that allow configured keywords only inside selected credentialed CLI content arguments. Command-path deny rules still apply.", + "tools.commandKeywordAllowlistJsonError": "Invalid JSON.", + "tools.commandKeywordAllowlistArrayError": "Use a JSON array of rules.", "tools.scrubCredentials": "Scrub Credentials", "browser.title": "Browser Automation", diff --git a/ui/web/src/i18n/locales/vi/config.json b/ui/web/src/i18n/locales/vi/config.json index b0c6a77d..aaace75b 100644 --- a/ui/web/src/i18n/locales/vi/config.json +++ b/ui/web/src/i18n/locales/vi/config.json @@ -208,6 +208,10 @@ "tools.execAskMode": "Chế độ hỏi", "tools.execAskModeTip": "Khi nào nhắc người dùng phê duyệt trước khi thực thi lệnh.", "tools.execAllowlistLabel": "Lệnh được phép (mỗi dòng một mẫu)", + "tools.commandKeywordAllowlist": "Allowlist từ khóa theo lệnh", + "tools.commandKeywordAllowlistTip": "Rule JSON cho phép từ khóa đã cấu hình chỉ trong argument nội dung của CLI có credential. Deny rule theo command path vẫn áp dụng.", + "tools.commandKeywordAllowlistJsonError": "JSON không hợp lệ.", + "tools.commandKeywordAllowlistArrayError": "Dùng mảng JSON các rule.", "tools.scrubCredentials": "Ẩn thông tin xác thực", "browser.title": "Tự động hóa trình duyệt", diff --git a/ui/web/src/i18n/locales/zh/config.json b/ui/web/src/i18n/locales/zh/config.json index 1994f1df..21fe8e1d 100644 --- a/ui/web/src/i18n/locales/zh/config.json +++ b/ui/web/src/i18n/locales/zh/config.json @@ -208,6 +208,10 @@ "tools.execAskMode": "询问模式", "tools.execAskModeTip": "在执行命令之前何时提示用户审批。", "tools.execAllowlistLabel": "允许的命令（每行一个模式）", + "tools.commandKeywordAllowlist": "命令关键词允许列表", + "tools.commandKeywordAllowlistTip": "JSON 规则只允许配置的关键词出现在指定的凭证 CLI 内容参数中。命令路径拒绝规则仍然生效。", + "tools.commandKeywordAllowlistJsonError": "JSON 无效。", + "tools.commandKeywordAllowlistArrayError": "请使用规则 JSON 数组。", "tools.scrubCredentials": "擦除凭证", "browser.title": "浏览器自动化", diff --git a/ui/web/src/pages/config/sections/tools-exec-section.tsx b/ui/web/src/pages/config/sections/tools-exec-section.tsx index a0caf55a..7486f094 100644 --- a/ui/web/src/pages/config/sections/tools-exec-section.tsx +++ b/ui/web/src/pages/config/sections/tools-exec-section.tsx @@ -19,7 +19,6 @@ import { } from "@/components/ui/select"; import { InfoLabel } from "@/components/shared/info-label"; - type ToolsData = Record; interface Props { @@ -31,10 +30,14 @@ interface Props { export function ToolsExecSection({ data, onSave, saving }: Props) { const { t } = useTranslation("config"); const [draft, setDraft] = useState(data ?? {}); + const [allowlistText, setAllowlistText] = useState(""); + const [allowlistError, setAllowlistError] = useState(null); const [dirty, setDirty] = useState(false); useEffect(() => { setDraft(data ?? {}); + setAllowlistText(JSON.stringify(data?.commandKeywordAllowlist ?? [], null, 2)); + setAllowlistError(null); setDirty(false); }, [data]); @@ -46,6 +49,30 @@ export function ToolsExecSection({ data, onSave, saving }: Props) { setDirty(true); }; + const updateCommandKeywordAllowlist = (value: string) => { + setAllowlistText(value); + setDirty(true); + + const trimmed = value.trim(); + if (!trimmed) { + setAllowlistError(null); + setDraft((prev) => ({ ...prev, commandKeywordAllowlist: [] })); + return; + } + + try { + const parsed = JSON.parse(trimmed); + if (!Array.isArray(parsed)) { + setAllowlistError(t("tools.commandKeywordAllowlistArrayError")); + return; + } + setAllowlistError(null); + setDraft((prev) => ({ ...prev, commandKeywordAllowlist: parsed })); + } catch { + setAllowlistError(t("tools.commandKeywordAllowlistJsonError")); + } + }; + if (!data) return null; const exec = draft.execApproval ?? {}; @@ -91,15 +118,41 @@ export function ToolsExecSection({ data, onSave, saving }: Props) { allowlist: e.target.value.split("\n").filter(Boolean), }) } - className="min-h-[80px] font-mono text-xs" + className="min-h-[80px] font-mono text-base md:text-sm" placeholder="git * npm * ls *" /> )} +
+ + {t("tools.commandKeywordAllowlist")} + +