feat(usage): add AI budget usage caps

This commit is contained in:
Goon committed 2026-05-23 00:16:03 +07:00
1 parent 43049d3b3f
commit 9ff5a629d2
67 files changed
+4195 -249

No files matched your search

+10 -2
View File
@@ -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()
}
+5 -2
View File
@@ -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
+3
View File
@@ -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.
+12 -3
View File
@@ -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,
+21
View File
@@ -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
+26
View File
@@ -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.
+51
View File
@@ -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
+17
View File
@@ -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**
+3 -2
View File
@@ -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
+3 -2
View File
@@ -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
+98 -45
View File
@@ -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),
})
}
+58 -54
View File
@@ -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 <agentID>")
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 <agentID>")
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"`
+3 -2
View File
@@ -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")
+3
View File
@@ -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,
+88
View File
@@ -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
}
+5
View File
@@ -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) }
+354
View File
@@ -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
}
+17
View File
@@ -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")
}
}
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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.
+48
View File
@@ -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,
+27
View File
@@ -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",
+7 -1
View File
@@ -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
+6
View File
@@ -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
+8 -1
View File
@@ -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 {
+10 -6
View File
@@ -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"`
}
+37 -36
View File
@@ -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.
+406
View File
@@ -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
}
+180
View File
@@ -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 }
+362
View File
@@ -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
}
+38 -35
View File
@@ -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
}
+195
View File
@@ -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
}
+6
View File
@@ -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 {
+46 -2
View File
@@ -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)
}
+6
View File
@@ -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 {
+26 -2
View File
@@ -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)
}
+17 -3
View File
@@ -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)
}
+6
View File
@@ -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 {
+22 -2
View File
@@ -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)
}
+36 -29
View File
@@ -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 {
+43 -4
View File
@@ -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,
})
}
+12 -10
View File
@@ -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),
+44
View File
@@ -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
}
+1 -1
View File
@@ -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
+275
View File
@@ -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
}
+337
View File
@@ -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 }
+119
View File
@@ -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
}
+66
View File
@@ -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())
}
}
+91
View File
@@ -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
}
}
@@ -0,0 +1,2 @@
DROP TABLE IF EXISTS usage_pricing_overrides;
DROP TABLE IF EXISTS usage_pricing_catalog;
@@ -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;
@@ -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;
@@ -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);
+32
View File
@@ -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",
+39 -1
View File
@@ -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",
+32
View File
@@ -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",
+39 -1
View File
@@ -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",
+32
View File
@@ -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",
+39 -1
View File
@@ -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": "用量分析",
+7
View File
@@ -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<string, unknown>) => ["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,
@@ -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<PricingOverride>("/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,
};
}
@@ -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)
</section>
) : null}
<ProviderPricingSection provider={provider} />
{showEmbedding ? (
<ProviderEmbeddingSection embEnabled={embEnabled} setEmbEnabled={setEmbEnabled} embModel={embModel} setEmbModel={setEmbModel} embApiBase={embApiBase} setEmbApiBase={setEmbApiBase} onVerify={() => verifyEmbedding(provider.id, embModel.trim() || undefined, undefined)} verifying={embVerifying} verifyResult={embResult} />
) : null}
@@ -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<keyof UsagePricingFields> = [
"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<Record<string, string>>({});
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 (
<section className="space-y-4 rounded-lg border p-3 sm:p-4 overflow-hidden">
<div className="flex flex-col gap-3 sm:flex-row sm:items-start sm:justify-between">
<div className="space-y-1">
<h3 className="text-sm font-medium">{t("pricing.title")}</h3>
<p className="text-xs text-muted-foreground">{isSkipped ? t("pricing.skippedDescription") : t("pricing.description")}</p>
</div>
{isMasterScope ? (
<Button type="button" variant="outline" size="sm" onClick={() => void onSync()} disabled={syncing} className="gap-1 self-start">
<DatabaseZap className={`h-3.5 w-3.5${syncing ? " animate-pulse" : ""}`} />
{t("pricing.sync")}
</Button>
) : null}
</div>
<div className="space-y-3">
<div className="space-y-1.5">
<Label htmlFor="pricingModel">{t("pricing.model")}</Label>
<Input id="pricingModel" value={modelId} onChange={(e) => setModelId(e.target.value)} placeholder="openai/gpt-4o-mini" className="text-base md:text-sm" />
</div>
<div className="grid gap-3 sm:grid-cols-2 lg:grid-cols-4">
{PRICE_FIELDS.map((field) => (
<div key={field} className="space-y-1.5">
<Label htmlFor={`pricing-${field}`} className="text-xs">{t(`pricing.fields.${field}`)}</Label>
<Input
id={`pricing-${field}`}
value={prices[field] ?? ""}
onChange={(e) => setPrices((current) => ({ ...current, [field]: e.target.value }))}
inputMode="decimal"
placeholder="0.00000015"
className="text-base md:text-sm"
/>
</div>
))}
</div>
<Button type="button" onClick={() => void onSave()} disabled={saving || !modelId.trim() || Object.keys(pricing).length === 0} className="gap-1">
<Save className="h-4 w-4" />
{t("pricing.saveOverride")}
</Button>
</div>
{catalog.length > 0 ? (
<div className="space-y-2">
<div className="flex items-center gap-2 text-xs font-medium text-muted-foreground">
<RefreshCw className={`h-3.5 w-3.5${refreshing ? " animate-spin" : ""}`} />
{t("pricing.catalogMatches")}
</div>
<div className="flex flex-wrap gap-2">
{catalog.map((entry) => (
<button key={entry.id} type="button" onClick={() => { setModelId(entry.model_id); setPrices(cleanPricing(entry.pricing)); }} className="rounded-md border px-2 py-1 text-left text-xs hover:bg-muted">
{entry.model_id}
</button>
))}
</div>
</div>
) : null}
<div className="overflow-x-auto">
<table className="w-full min-w-[680px] text-sm">
<thead>
<tr className="border-b bg-muted/40">
<th className="px-3 py-2 text-left font-medium">{t("pricing.model")}</th>
<th className="px-3 py-2 text-left font-medium">{t("pricing.source")}</th>
<th className="px-3 py-2 text-left font-medium">{t("pricing.configuredFields")}</th>
<th className="px-3 py-2 text-right font-medium">{t("columns.actions")}</th>
</tr>
</thead>
<tbody>
{overrides.length === 0 ? (
<tr><td colSpan={4} className="px-3 py-5 text-center text-muted-foreground">{t("pricing.emptyOverrides")}</td></tr>
) : overrides.map((override) => (
<tr key={override.id} className="border-b last:border-0">
<td className="px-3 py-2 font-medium">{override.model_id}</td>
<td className="px-3 py-2"><Badge variant="outline">{t("pricing.override")}</Badge></td>
<td className="px-3 py-2 text-muted-foreground">{Object.keys(override.pricing ?? {}).join(", ") || "—"}</td>
<td className="px-3 py-2 text-right">
<Button type="button" variant="ghost" size="icon" onClick={() => void deleteOverride(override.id)} aria-label={t("pricing.deleteOverride")}>
<Trash2 className="h-4 w-4" />
</Button>
</td>
</tr>
))}
</tbody>
</table>
</div>
</section>
);
}
function cleanPricing(fields: UsagePricingFields): Record<string, string> {
const out: Record<string, string> = {};
for (const field of PRICE_FIELDS) {
const value = fields[field];
if (value) out[field] = value;
}
return out;
}
@@ -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<UsageCapPolicy["window"]>("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 (
<section className="space-y-4 rounded-md border p-3 sm:p-4">
<div className="flex flex-col gap-3 sm:flex-row sm:items-center sm:justify-between">
<div>
<h3 className="text-sm font-semibold">{t("caps.title")}</h3>
<p className="text-xs text-muted-foreground">{t("caps.description")}</p>
</div>
<Button type="button" variant="outline" size="sm" onClick={() => void refresh()} disabled={refreshing} className="gap-1 self-start sm:self-auto">
<RefreshCw className={`h-3.5 w-3.5${refreshing ? " animate-spin" : ""}`} />
{t("common:refresh", "Refresh")}
</Button>
</div>
<div className="grid gap-3 md:grid-cols-2 lg:grid-cols-6">
<Field label={t("caps.window")}>
<Select value={windowValue} onValueChange={(value) => setWindowValue(value as UsageCapPolicy["window"])}>
<SelectTrigger className="w-full text-base md:text-sm"><SelectValue /></SelectTrigger>
<SelectContent>
{["hour", "day", "week", "month"].map((value) => <SelectItem key={value} value={value}>{t(`caps.windows.${value}`)}</SelectItem>)}
</SelectContent>
</Select>
</Field>
<Field label={t("caps.agent")}>
<Select value={agentId} onValueChange={setAgentId}>
<SelectTrigger className="w-full text-base md:text-sm"><SelectValue /></SelectTrigger>
<SelectContent>
<SelectItem value={ALL}>{t("caps.allAgents")}</SelectItem>
{agents.map((a) => <SelectItem key={a.id} value={a.id}>{a.display_name || a.agent_key || a.id}</SelectItem>)}
</SelectContent>
</Select>
</Field>
<Field label={t("caps.provider")}>
<Select value={providerId} onValueChange={setProviderId}>
<SelectTrigger className="w-full text-base md:text-sm"><SelectValue /></SelectTrigger>
<SelectContent>
<SelectItem value={ALL}>{t("caps.allProviders")}</SelectItem>
{providers.map((p) => <SelectItem key={p.id} value={p.id}>{p.display_name || p.name}</SelectItem>)}
</SelectContent>
</Select>
</Field>
<Field label={t("caps.model")}>
<Input value={modelId} onChange={(e) => setModelId(e.target.value)} placeholder="openai/gpt-4o-mini" className="text-base md:text-sm" />
</Field>
<Field label={t("caps.maxTokens")}>
<Input value={maxTokens} onChange={(e) => setMaxTokens(e.target.value)} inputMode="numeric" placeholder="500000" className="text-base md:text-sm" />
</Field>
<Field label={t("caps.maxCost")}>
<div className="flex gap-2">
<Input value={maxCost} onChange={(e) => setMaxCost(e.target.value)} inputMode="decimal" placeholder="25" className="text-base md:text-sm" />
<Button type="button" size="icon" onClick={() => void onSubmit()} disabled={saving} aria-label={t("caps.create")}>
<Plus className="h-4 w-4" />
</Button>
</div>
</Field>
</div>
<div className="overflow-x-auto">
<table className="w-full min-w-[760px] text-sm">
<thead>
<tr className="border-b bg-muted/40">
<th className="px-3 py-2 text-left font-medium">{t("caps.scope")}</th>
<th className="px-3 py-2 text-left font-medium">{t("caps.window")}</th>
<th className="px-3 py-2 text-right font-medium">{t("caps.tokens")}</th>
<th className="px-3 py-2 text-right font-medium">{t("caps.cost")}</th>
<th className="px-3 py-2 text-right font-medium">{t("columns.status")}</th>
<th className="px-3 py-2 text-right font-medium">{t("columns.actions", "Actions")}</th>
</tr>
</thead>
<tbody>
{utilization.length === 0 ? (
<tr><td colSpan={6} className="px-3 py-6 text-center text-muted-foreground">{t("caps.empty")}</td></tr>
) : utilization.map((row) => (
<UsageCapRow key={row.policy.id} row={row} onDelete={() => void deletePolicy(row.policy.id)} />
))}
</tbody>
</table>
</div>
{blockedEvents.length > 0 ? (
<div className="space-y-2">
<div className="flex items-center gap-2 text-sm font-medium"><ShieldAlert className="h-4 w-4 text-destructive" />{t("caps.recentBlocks")}</div>
<div className="grid gap-2 sm:grid-cols-2">
{blockedEvents.slice(0, 4).map((event) => (
<div key={event.id} className="rounded-md border border-destructive/30 bg-destructive/5 px-3 py-2 text-xs">
<div className="font-medium">{event.reason || t("caps.blocked")}</div>
<div className="text-muted-foreground">{formatDate(event.created_at)} · {formatTokens(event.estimated_tokens)} · {formatCost(event.estimated_cost_micros / 1_000_000)}</div>
</div>
))}
</div>
</div>
) : null}
</section>
);
}
function Field({ label, children }: { label: string; children: ReactNode }) {
return <div className="space-y-1.5"><Label className="text-xs">{label}</Label>{children}</div>;
}
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 (
<tr className="border-b last:border-0">
<td className="px-3 py-2">
<div className="font-medium">{p.model_id || p.provider_type || t("caps.tenantScope")}</div>
<div className="text-xs text-muted-foreground">{p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")}</div>
</td>
<td className="px-3 py-2"><Badge variant="outline">{t(`caps.windows.${p.window}`)}</Badge></td>
<td className="px-3 py-2 text-right">{p.max_tokens ? `${formatTokens(tokenUsed)} / ${formatTokens(p.max_tokens)} (${tokenPct}%)` : "—"}</td>
<td className="px-3 py-2 text-right">{p.max_cost_micros ? `${formatCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "—"}</td>
<td className="px-3 py-2 text-right"><Badge variant={p.enabled ? "default" : "secondary"}>{p.enabled ? t("caps.enabled") : t("caps.disabled")}</Badge></td>
<td className="px-3 py-2 text-right">
<Button type="button" variant="ghost" size="icon" onClick={onDelete} aria-label={t("caps.delete")}><Trash2 className="h-4 w-4" /></Button>
</td>
</tr>
);
}
@@ -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<UsageCapPolicy>("/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,
};
}
+5
View File
@@ -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() {
<SummaryCards current={current} previous={previous} loading={loading} />
</ErrorBoundary>
<ErrorBoundary>
<UsageCapsPanel />
</ErrorBoundary>
<ErrorBoundary>
<TokenAreaChart data={timeseries} loading={loading} granularity={filters.granularity} />
</ErrorBoundary>
+68
View File
@@ -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;
}