tiennm99/goclaw
1
0 Fork 0
mirror of https://github.com/tiennm99/goclaw.git synced 2026-10-12 18:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): add AI budget usage caps

Browse Source 
Merge usage-cap budget controls with conflict resolution against dev.
This commit is contained in:
Duy /zuey/ authored and GitHub committed 2026-05-24 13:57:31 +07:00
2 parents 61fdcdcb81 7fee3073df
commit 2bb736cf18
108 files changed
+5562 -387
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cmd/gateway.go
+17 -4
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"
@@ -225,6 +226,7 @@ func runGateway() {
}
}
setupMemoryEmbeddings(pgStores, providerRegistry)
usageCapSvc := usagecaps.NewService(pgStores.UsageCaps, pgStores.Providers)
// Resolve background provider for consolidation + vault enrichment.
// Fallback: background.provider → agent.default_provider → first registered provider.
@@ -236,6 +238,7 @@ func runGateway() {
var kgExtractor *kg.Extractor
if pgStores.KnowledgeGraph != nil {
kgExtractor = kg.NewExtractor(bgProvider, bgModel, 0)
kgExtractor.SetUsageCapService(usageCapSvc)
}
cleanupConsolidation := consolidation.Register(consolidation.ConsolidationDeps{
EpisodicStore: pgStores.Episodic,
@@ -247,6 +250,7 @@ func runGateway() {
Registry: providerRegistry,
Extractor: kgExtractor,
AlertDeps: bgalert.AlertDeps{SystemConfigs: pgStores.SystemConfigs, MsgBus: msgBus},
UsageCaps: usageCapSvc,
AgentStore: pgStores.Agents,
})
defer cleanupConsolidation()
@@ -269,6 +273,7 @@ func runGateway() {
MsgBus: msgBus,
TeamStore: pgStores.Teams,
AlertDeps: bgalert.AlertDeps{SystemConfigs: pgStores.SystemConfigs, MsgBus: msgBus},
UsageCaps: usageCapSvc,
})
enrichProgress = ep
enrichWorker = ew
@@ -285,8 +290,14 @@ func runGateway() {
slog.Info("bootstrap: capabilities backfill complete", "agents", count)
}
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))
@@ -339,7 +350,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()
}
@@ -359,6 +370,7 @@ func runGateway() {
workspace: workspace,
dataDir: dataDir,
domainBus: domainBus,
usageCapSvc: usageCapSvc,
audioMgr: audioMgr,
}
@@ -371,7 +383,7 @@ func runGateway() {
httpapi.InitGatewayNoAuthFallbackAllowed(config.GatewayNoAuthFallbackAllowed(cfg.Gateway))
exportTokenStore := httpapi.InitExportTokenStore()
defer exportTokenStore.Stop()
agentsH, skillsH, tracesH, mcpH, channelInstancesH, providersH, builtinToolsH, pendingMessagesH, teamEventsH, secureCLIH, secureCLIGrantH, mcpUserCredsH := wireHTTP(pgStores, cfg.Agents.Defaults.Workspace, dataDir, bundledSkillsDir, msgBus, toolsReg, providerRegistry, modelReg, permPE.IsOwner, gatewayAddr, mcpToolLister)
agentsH, skillsH, tracesH, mcpH, channelInstancesH, providersH, builtinToolsH, pendingMessagesH, teamEventsH, secureCLIH, secureCLIGrantH, mcpUserCredsH := wireHTTP(pgStores, cfg.Agents.Defaults.Workspace, dataDir, bundledSkillsDir, msgBus, toolsReg, providerRegistry, modelReg, permPE.IsOwner, gatewayAddr, mcpToolLister, usageCapSvc)
// Wire dependencies for system prompt preview parity.
if agentsH != nil {
@@ -428,7 +440,7 @@ func runGateway() {
// Register all RPC methods
server.SetLogTee(logTee)
server.SetRuntimeLogsHandler(httpapi.NewRuntimeLogsHandler(logTee))
pairingMethods, heartbeatMethods, chatMethods, cfgPermsMethods := registerAllMethods(server, agentRouter, pgStores.Sessions, pgStores.Cron, pgStores.Pairing, cfg, cfgPath, workspace, dataDir, msgBus, execApprovalMgr, pgStores.Agents, pgStores.Skills, pgStores.ConfigSecrets, pgStores.Teams, contextFileInterceptor, logTee, pgStores.Heartbeats, pgStores.ConfigPermissions, pgStores.SystemConfigs, pgStores.Tenants, pgStores.SkillTenantCfgs, audioMgr)
pairingMethods, heartbeatMethods, chatMethods, cfgPermsMethods := registerAllMethods(server, agentRouter, pgStores.Sessions, pgStores.Cron, pgStores.Pairing, cfg, cfgPath, workspace, dataDir, msgBus, execApprovalMgr, pgStores.Agents, pgStores.Skills, pgStores.ConfigSecrets, pgStores.Teams, contextFileInterceptor, logTee, pgStores.Heartbeats, pgStores.ConfigPermissions, pgStores.SystemConfigs, pgStores.Tenants, pgStores.SkillTenantCfgs, audioMgr, usageCapSvc)
// Phase 3: Agent hooks RPC methods (hooks.list/create/update/delete/toggle/test/history).
if hs, ok := pgStores.Hooks.(hooks.HookStore); ok && hs != nil {
@@ -514,6 +526,7 @@ func runGateway() {
instanceLoader = channels.NewInstanceLoader(pgStores.ChannelInstances, pgStores.Agents, channelMgr, msgBus, pgStores.Pairing)
instanceLoader.SetProviderRegistry(providerRegistry)
instanceLoader.SetPendingCompactionConfig(cfg.Channels.PendingCompaction)
instanceLoader.SetUsageCapService(usageCapSvc)
instanceLoader.RegisterFactory(channels.TypeTelegram, telegram.FactoryWithStoresAndAudio(pgStores.Agents, pgStores.ConfigPermissions, pgStores.Teams, pgStores.SubagentTasks, pgStores.PendingMessages, audioMgr))
instanceLoader.RegisterFactory(channels.TypeDiscord, discord.FactoryWithStoresAndAudio(pgStores.Agents, pgStores.ConfigPermissions, pgStores.PendingMessages, audioMgr))
instanceLoader.RegisterFactory(channels.TypeFeishu, feishu.FactoryWithPendingStoreAndAudio(pgStores.PendingMessages, audioMgr))
cmd/gateway_agents.go
+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
cmd/gateway_consumer.go
+3 -1
View File
@@ -19,6 +19,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/scheduler"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
)
@@ -26,7 +27,7 @@ import (
// and routes them through the scheduler/agent loop, then publishes the response back.
// Also handles subagent announcements: routes them through the parent agent's session
// (matching TS subagent-announce.ts pattern) so the agent can reformulate for the user.
func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents *agent.Router, cfg *config.Config, sched *scheduler.Scheduler, channelMgr *channels.Manager, teamStore store.TeamStore, quotaChecker *channels.QuotaChecker, sessStore store.SessionStore, agentStore store.AgentStore, contactCollector *store.ContactCollector, postTurn tools.PostTurnProcessor, subagentMgr *tools.SubagentManager) {
func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents *agent.Router, cfg *config.Config, sched *scheduler.Scheduler, channelMgr *channels.Manager, teamStore store.TeamStore, quotaChecker *channels.QuotaChecker, sessStore store.SessionStore, agentStore store.AgentStore, contactCollector *store.ContactCollector, postTurn tools.PostTurnProcessor, subagentMgr *tools.SubagentManager, usageCapSvc *usagecaps.Service) {
slog.Info("inbound message consumer started")
// Inbound message deduplication (matching TS src/infra/dedupe.ts + inbound-dedupe.ts).
@@ -58,6 +59,7 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents
QuotaChecker: quotaChecker,
ContactCollector: contactCollector,
SubagentMgr: subagentMgr,
UsageCaps: usageCapSvc,
GetAnnounceMu: getAnnounceMu,
}
cmd/gateway_consumer_deps.go
+2
View File
@@ -10,6 +10,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/scheduler"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// ConsumerDeps bundles shared dependencies for consumer message handlers.
@@ -28,6 +29,7 @@ type ConsumerDeps struct {
ContactCollector *store.ContactCollector
TaskRunSessions sync.Map
SubagentMgr *tools.SubagentManager
UsageCaps *usagecaps.Service
BgWg sync.WaitGroup
GetAnnounceMu func(string) *sync.Mutex
}
cmd/gateway_consumer_normal.go
+5 -1
View File
@@ -333,7 +333,11 @@ func processNormalMessage(
if locale == "" {
locale = "en"
}
intent := agent.ClassifyIntent(ctx, loop.Provider(), loop.Model(), msg.Content)
classifyCtx := ctx
if uid := loop.UUID(); uid != uuid.Nil {
classifyCtx = store.WithAgentID(classifyCtx, uid)
}
intent := agent.ClassifyIntentWithUsageCaps(classifyCtx, deps.UsageCaps, loop.Provider(), loop.Model(), msg.Content)
switch intent {
case agent.IntentStatusQuery:
status := deps.Agents.GetActivity(sessionKey)
cmd/gateway_deps.go
+5 -3
View File
@@ -14,6 +14,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/skills"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/internal/vault"
)
@@ -28,13 +29,14 @@ type gatewayDeps struct {
channelMgr *channels.Manager
agentRouter *agent.Router
toolsReg *tools.Registry
skillsLoader *skills.Loader // optional: enables skill creation in evolution approval
skillsLoader *skills.Loader // optional: enables skill creation in evolution approval
permCache *cache.PermissionCache // nil if no tenant store; closed on shutdown to stop sweep goroutines
enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered
enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue
enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered
enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue
workspace string
dataDir string
domainBus eventbus.DomainEventBus
usageCapSvc *usagecaps.Service
audioMgr *audio.Manager // nil if TTS not configured; used by TTSHandler
ttsHandler *httpapi.TTSHandler // nil if TTS not configured; for hot-reload
}
cmd/gateway_hooks.go
+4 -2
View File
@@ -7,12 +7,13 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/config"
"github.com/nextlevelbuilder/goclaw/internal/edition"
"github.com/nextlevelbuilder/goclaw/internal/hooks"
hookhandlers "github.com/nextlevelbuilder/goclaw/internal/hooks/handlers"
"github.com/nextlevelbuilder/goclaw/internal/hooks/budget"
hookhandlers "github.com/nextlevelbuilder/goclaw/internal/hooks/handlers"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/security"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/store/pg"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// sharedHookHandlers is populated by wireExtras so the gateway.go router
@@ -27,7 +28,7 @@ var sharedHookHandlers map[hooks.HandlerType]hooks.Handler
// Budget wiring (C1 fix): the PromptHandler receives a budget.Store bound
// to pg.NewPGHookBudget so token spend is atomically deducted per tenant.
// When the DB handle is unavailable, budget falls back to nil (Lite desktop).
func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, hooksCfg config.HooksConfig) map[hooks.HandlerType]hooks.Handler {
func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, hooksCfg config.HooksConfig, usageCapSvc *usagecaps.Service) map[hooks.HandlerType]hooks.Handler {
encryptKey := os.Getenv("GOCLAW_ENCRYPTION_KEY")
var budgetStore *budget.Store
@@ -38,6 +39,7 @@ func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, ho
promptHandler := &hookhandlers.PromptHandler{
Resolver: hookhandlers.NewRegistryResolver(providerReg, stores.SystemConfigs),
Budget: budgetStore,
UsageCaps: usageCapSvc,
DefaultModel: "haiku",
}
cmd/gateway_http_handlers.go
+5 -2
View File
@@ -6,10 +6,11 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// wireHTTP creates HTTP handlers (agents + skills + traces + MCP + channel instances + providers + builtin tools + pending messages).
func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir string, msgBus *bus.MessageBus, toolsReg *tools.Registry, providerReg *providers.Registry, modelReg providers.ModelRegistry, isOwner func(string) bool, gatewayAddr string, mcpToolLister httpapi.MCPToolLister) (*httpapi.AgentsHandler, *httpapi.SkillsHandler, *httpapi.TracesHandler, *httpapi.MCPHandler, *httpapi.ChannelInstancesHandler, *httpapi.ProvidersHandler, *httpapi.BuiltinToolsHandler, *httpapi.PendingMessagesHandler, *httpapi.TeamEventsHandler, *httpapi.SecureCLIHandler, *httpapi.SecureCLIGrantHandler, *httpapi.MCPUserCredentialsHandler) {
func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir string, msgBus *bus.MessageBus, toolsReg *tools.Registry, providerReg *providers.Registry, modelReg providers.ModelRegistry, isOwner func(string) bool, gatewayAddr string, mcpToolLister httpapi.MCPToolLister, usageCapSvc *usagecaps.Service) (*httpapi.AgentsHandler, *httpapi.SkillsHandler, *httpapi.TracesHandler, *httpapi.MCPHandler, *httpapi.ChannelInstancesHandler, *httpapi.ProvidersHandler, *httpapi.BuiltinToolsHandler, *httpapi.PendingMessagesHandler, *httpapi.TeamEventsHandler, *httpapi.SecureCLIHandler, *httpapi.SecureCLIGrantHandler, *httpapi.MCPUserCredentialsHandler) {
var agentsH *httpapi.AgentsHandler
var skillsH *httpapi.SkillsHandler
var tracesH *httpapi.TracesHandler
@@ -24,7 +25,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir
if stores != nil && stores.Agents != nil {
var summoner *httpapi.AgentSummoner
if providerReg != nil {
summoner = httpapi.NewAgentSummoner(stores.Agents, providerReg, msgBus)
summoner = httpapi.NewAgentSummoner(stores.Agents, providerReg, msgBus, usageCapSvc)
}
agentsH = httpapi.NewAgentsHandler(stores.Agents, stores.Providers, providerReg, stores.DB, stores.Tracing, defaultWorkspace, msgBus, summoner, isOwner)
agentsH.SetImportStores(stores.Memory, stores.KnowledgeGraph)
@@ -61,6 +62,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir
if stores != nil && stores.Providers != nil {
providersH = httpapi.NewProvidersHandler(stores.Providers, stores.ConfigSecrets, providerReg, gatewayAddr)
providersH.SetMessageBus(msgBus)
providersH.SetUsageCapService(usageCapSvc)
if modelReg != nil {
providersH.SetModelRegistry(modelReg)
}
@@ -90,6 +92,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir
if stores != nil && stores.PendingMessages != nil {
pendingMessagesH = httpapi.NewPendingMessagesHandler(stores.PendingMessages, stores.Agents, providerReg)
pendingMessagesH.SetUsageCapService(usageCapSvc)
}
if stores != nil && stores.SecureCLI != nil {
cmd/gateway_http_wiring.go
+6 -1
View File
@@ -137,6 +137,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.
@@ -234,7 +237,9 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer(
// Knowledge graph API
if d.pgStores != nil && d.pgStores.KnowledgeGraph != nil {
d.server.SetKnowledgeGraphHandler(httpapi.NewKnowledgeGraphHandler(d.pgStores.KnowledgeGraph, d.providerRegistry))
kgHandler := httpapi.NewKnowledgeGraphHandler(d.pgStores.KnowledgeGraph, d.providerRegistry)
kgHandler.SetUsageCapService(d.usageCapSvc)
d.server.SetKnowledgeGraphHandler(kgHandler)
}
// V3: Evolution metrics + suggestions API
cmd/gateway_lifecycle.go
+1 -1
View File
@@ -140,7 +140,7 @@ func (d *gatewayDeps) runLifecycle(
d.channelMgr.SetContactCollector(contactCollector)
}
go consumeInboundMessages(ctx, d.msgBus, d.agentRouter, d.cfg, deps.sched, d.channelMgr, deps.consumerTeamStore, deps.quotaChecker, d.pgStores.Sessions, d.pgStores.Agents, contactCollector, deps.postTurn, deps.subagentMgr)
go consumeInboundMessages(ctx, d.msgBus, d.agentRouter, d.cfg, deps.sched, d.channelMgr, deps.consumerTeamStore, deps.quotaChecker, d.pgStores.Sessions, d.pgStores.Agents, contactCollector, deps.postTurn, deps.subagentMgr, d.usageCapSvc)
// Webhook callback worker — delivers async webhook_calls rows to receiver callback_url.
// Runs in both editions: Standard (PG, concurrency=4) and Lite (SQLite, concurrency=1).
cmd/gateway_managed.go
+16 -6
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")
}
@@ -177,7 +185,7 @@ func wireExtras(
"disabled_count", n, "edition", edition.Current().Name)
}
handlers := buildHookHandlers(stores, providerReg, appCfg.Hooks)
handlers := buildHookHandlers(stores, providerReg, appCfg.Hooks, usageCapSvc)
stdOpts := hooks.StdDispatcherOpts{
Store: hs,
Audit: hooks.NewAuditWriter(hs, ""),
@@ -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,
@@ -295,7 +304,7 @@ func wireExtras(
writeMemIntc = tools.NewMemoryInterceptor(stores.Memory, workspace)
// Hook KG extraction on memory writes if KG store is available
if stores.KnowledgeGraph != nil && stores.BuiltinTools != nil {
writeMemIntc.SetKGExtractFunc(buildKGExtractFunc(stores.KnowledgeGraph, stores.BuiltinTools, providerReg))
writeMemIntc.SetKGExtractFunc(buildKGExtractFunc(stores.KnowledgeGraph, stores.BuiltinTools, providerReg, usageCapSvc))
}
}
if readTool, ok := toolsReg.Get("read_file"); ok {
@@ -703,7 +712,7 @@ type kgSettings struct {
// buildKGExtractFunc returns a callback that extracts entities from memory content.
// Settings are read from the builtin_tools table on each invocation (not cached),
// so changes take effect immediately without restart.
func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinToolStore, providerReg *providers.Registry) tools.KGExtractFunc {
func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinToolStore, providerReg *providers.Registry, usageCapSvc *usagecaps.Service) tools.KGExtractFunc {
return func(ctx context.Context, agentID, userID, content string) {
slog.Info("kg extract: triggered", "agent", agentID, "user", userID, "content_len", len(content))
// Read settings from DB on each call so admin changes take effect immediately
@@ -727,6 +736,7 @@ func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinTool
return
}
extractor := kg.NewExtractor(p, settings.ExtractionModel, settings.MinConfidence)
extractor.SetUsageCapService(usageCapSvc)
result, err := extractor.Extract(ctx, content)
if err != nil {
slog.Warn("kg extract: extraction failed", "agent", agentID, "error", err)
cmd/gateway_methods.go
+3 -1
View File
@@ -12,14 +12,16 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/gateway/methods"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, cronStore store.CronStore, pairingStore store.PairingStore, cfg *config.Config, cfgPath, workspace, dataDir string, msgBus *bus.MessageBus, execApprovalMgr *tools.ExecApprovalManager, agentStore store.AgentStore, skillStore store.SkillStore, configSecretsStore store.ConfigSecretsStore, teamStore store.TeamStore, contextFileInterceptor *tools.ContextFileInterceptor, logTee *gateway.LogTee, heartbeatStore store.HeartbeatStore, configPermStore store.ConfigPermissionStore, sysConfigStore store.SystemConfigStore, tenantStore store.TenantStore, skillTenantCfgStore store.SkillTenantConfigStore, audioMgr *audio.Manager) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) {
func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, cronStore store.CronStore, pairingStore store.PairingStore, cfg *config.Config, cfgPath, workspace, dataDir string, msgBus *bus.MessageBus, execApprovalMgr *tools.ExecApprovalManager, agentStore store.AgentStore, skillStore store.SkillStore, configSecretsStore store.ConfigSecretsStore, teamStore store.TeamStore, contextFileInterceptor *tools.ContextFileInterceptor, logTee *gateway.LogTee, heartbeatStore store.HeartbeatStore, configPermStore store.ConfigPermissionStore, sysConfigStore store.SystemConfigStore, tenantStore store.TenantStore, skillTenantCfgStore store.SkillTenantConfigStore, audioMgr *audio.Manager, usageCapSvc *usagecaps.Service) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) {
router := server.Router()
// Phase 1: Core methods
chatMethods := methods.NewChatMethods(agents, sessStore, cfg, server.RateLimiter(), msgBus)
chatMethods.SetAudioManager(audioMgr) // Wire TTS auto-apply for WS responses
chatMethods.SetUsageCapService(usageCapSvc)
chatMethods.Register(router)
methods.NewAgentsMethods(agents, cfg, cfgPath, workspace, agentStore, contextFileInterceptor, msgBus).Register(router)
methods.NewSessionsMethods(sessStore, msgBus, cfg).Register(router)
docs/02-providers.md
+23
View File
@@ -54,6 +54,29 @@ 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.
The legacy agent-level `budget_monthly_cents` field is treated as a generated monthly agent USD cap. Existing values are backfilled during migration, and later agent budget edits update or remove the generated cap policy.
Supported price units: input, output, cache read, cache write, reasoning, request, image, and web search.
---
## 2. Supported Providers
### Six Core Provider Types
docs/06-store-data-model.md
+28
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,33 @@ 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, `window_key`, and `source`.
- `usage_cap_counters`: current window used and reserved token/cost counters.
- `usage_cap_reservations`: preflight reservations keyed by LLM call attempt.
- `usage_cap_events`: allow/block/reconcile/skip audit events.
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.
Agent `budget_monthly_cents` values are bridged into `usage_cap_policies` with `source = 'agent_budget_monthly_cents'`, an agent scope, `window_key = 'month'`, and `max_cost_micros = budget_monthly_cents * 10000`. Updating or clearing the agent budget keeps that generated policy in sync; manual cap policies continue to use `source = 'manual'`.
Migration versions:
- PostgreSQL: `000070_usage_caps_pricing`, `000071_usage_cap_policies`, `000072_agent_budget_usage_cap_bridge`.
- 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.
docs/18-http-api.md
+53
View File
@@ -390,6 +390,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:
@@ -1271,11 +1301,34 @@ 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
}
```
Policy responses include read-only `source`. `source="agent_budget_monthly_cents"` means the policy is generated from the agent monthly budget field and must be changed there; direct policy updates/deletes return `409 Conflict`.
---
## 24. Activity & Audit
docs/project-changelog.md
+20
View File
@@ -42,6 +42,26 @@ 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.
- Bridged legacy `budget_monthly_cents` into generated monthly agent USD cap policies, including migration backfill and save-time sync.
- Added web dashboard controls on Usage and Provider detail pages.
- Added Usage page editing for manual cap policies, including enable/disable, scope clearing, and token/USD limit clearing while keeping generated agent-budget caps read-only.
- Added usage-cap decision metadata to LLM spans, including allow/skip/block reason, policy IDs, estimates, actuals, and reconcile status.
**Tests**
- Added pricing and cap service coverage.
- Verified Go builds, SQLite build compatibility, full Go test suite, integration race suite, and Web UI production build.
### Messaging debounce hardening
**Features**
internal/agent/intent_classify.go
+11 -1
View File
@@ -8,6 +8,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/i18n"
"github.com/nextlevelbuilder/goclaw/internal/providers"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// IntentType represents the classified intent of a user message.
@@ -108,6 +109,10 @@ func containsWholeWord(s, kw string) bool {
// Uses keyword fast-path first, then falls back to LLM classification.
// Falls back to IntentNewTask on any error.
func ClassifyIntent(ctx context.Context, provider providers.Provider, model, userMessage string) IntentType {
return ClassifyIntentWithUsageCaps(ctx, nil, provider, model, userMessage)
}
func ClassifyIntentWithUsageCaps(ctx context.Context, usageCaps *usagecaps.Service, provider providers.Provider, model, userMessage string) IntentType {
// Fast-path: keyword matching for obvious patterns (no LLM cost).
if intent, ok := quickClassify(userMessage); ok {
return intent
@@ -116,7 +121,7 @@ func ClassifyIntent(ctx context.Context, provider providers.Provider, model, use
ctx, cancel := context.WithTimeout(ctx, intentClassifyTimeout)
defer cancel()
resp, err := provider.Chat(ctx, providers.ChatRequest{
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: intentSystemPrompt},
{Role: "user", Content: userMessage},
@@ -126,6 +131,11 @@ func ClassifyIntent(ctx context.Context, provider providers.Provider, model, use
providers.OptMaxTokens: 20,
providers.OptTemperature: 0.0,
},
}
resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{
ModelID: model,
Purpose: "intent-classify",
MaxOutputTokens: 20,
})
if err != nil {
return IntentNewTask
internal/agent/loop_compact.go
+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
internal/agent/loop_history_sanitize.go
+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
internal/agent/loop_pipeline_callbacks.go
+109 -46
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"
)
@@ -293,31 +295,86 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c
opts = append(opts, withProvider(provider.Name()))
}
spanID := l.emitLLMSpanStart(ctx, start, state.Iteration+1, chatReq.Messages, opts...)
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)
recordUsageCapAttempt := func(reservation *usagecaps.Reservation) {
if reservation != nil {
opts = append(opts, withUsageCapMetadata(reservation.TraceMetadata()))
}
}
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 {
recordUsageCapAttempt(reservation)
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)
}
recordUsageCapAttempt(reservation)
}
}, 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 {
recordUsageCapAttempt(reservation)
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)
recordUsageCapAttempt(reservation)
}
return callResp, callErr
} else {
callResp, callErr = provider.Chat(ctx, request)
}
if reservation != nil {
reservation.Reconcile(ctx, callResp, callErr)
recordUsageCapAttempt(reservation)
}
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 +399,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 +572,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),
})
}
internal/agent/loop_tracing.go
+22 -5
View File
@@ -14,6 +14,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
"github.com/nextlevelbuilder/goclaw/internal/tracing"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
func (l *Loop) emit(event AgentEvent) {
@@ -51,20 +52,33 @@ func (l *Loop) IsRunning() bool { return l.activeRuns.Load() > 0 }
type spanOption func(*spanOverrides)
type spanOverrides struct {
model string
provider string
model string
provider string
usageCapAttempts []usagecaps.TraceMetadata
}
func withModel(m string) spanOption { return func(o *spanOverrides) { o.model = m } }
func withProvider(p string) spanOption { return func(o *spanOverrides) { o.provider = p } }
func withUsageCapMetadata(metadata usagecaps.TraceMetadata) spanOption {
return func(o *spanOverrides) {
if !metadata.Empty() {
o.usageCapAttempts = append(o.usageCapAttempts, metadata)
}
}
}
// resolveSpan returns (model, provider) applying any overrides on top of agent defaults.
func (l *Loop) resolveSpan(opts []spanOption) (string, string) {
o := l.resolveSpanOverrides(opts)
return o.model, o.provider
}
func (l *Loop) resolveSpanOverrides(opts []spanOption) spanOverrides {
o := spanOverrides{model: l.model, provider: l.provider.Name()}
for _, fn := range opts {
fn(&o)
}
return o.model, o.provider
return o
}
// ---------------------------------------------------------------------------
@@ -150,6 +164,7 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.
"status": store.SpanStatusCompleted,
}
var spanMetadata json.RawMessage
spanOpts := l.resolveSpanOverrides(opts)
if callErr != nil {
updates["status"] = store.SpanStatusError
@@ -176,8 +191,7 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.
}
}
// Calculate cost if pricing config is available.
model, providerName := l.resolveSpan(opts)
if pricing := tracing.LookupPricing(l.modelPricing, providerName, model); pricing != nil {
if pricing := tracing.LookupPricing(l.modelPricing, spanOpts.provider, spanOpts.model); pricing != nil {
cost := tracing.CalculateCost(pricing, resp.Usage)
if cost > 0 {
updates["total_cost"] = cost
@@ -203,6 +217,9 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time.
if decision := providers.ReasoningDecisionFromContext(ctx); decision != nil {
spanMetadata = providers.MergeReasoningMetadata(spanMetadata, *decision)
}
if len(spanOpts.usageCapAttempts) > 0 {
spanMetadata = usagecaps.MergeTraceMetadata(spanMetadata, spanOpts.usageCapAttempts)
}
if len(spanMetadata) > 0 {
updates["metadata"] = spanMetadata
}
internal/agent/loop_types.go
+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"`
internal/agent/memoryflush.go
+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")
internal/agent/resolver.go
+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,
internal/agent/title_generate.go
+13 -3
View File
@@ -7,6 +7,7 @@ import (
"time"
"github.com/nextlevelbuilder/goclaw/internal/providers"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
const titleGenerateTimeout = 120 * time.Second
@@ -16,10 +17,14 @@ const titleSystemPrompt = `Generate a short title (max 15 words) for this conver
// GenerateTitle uses a lightweight LLM call to create a short conversation title
// from the user's first message. Returns empty string on error.
func GenerateTitle(ctx context.Context, provider providers.Provider, model, userMessage string) string {
return GenerateTitleWithUsageCaps(ctx, nil, provider, model, userMessage)
}
func GenerateTitleWithUsageCaps(ctx context.Context, usageCaps *usagecaps.Service, provider providers.Provider, model, userMessage string) string {
ctx, cancel := context.WithTimeout(ctx, titleGenerateTimeout)
defer cancel()
resp, err := provider.Chat(ctx, providers.ChatRequest{
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: titleSystemPrompt},
{Role: "user", Content: userMessage},
@@ -29,13 +34,18 @@ func GenerateTitle(ctx context.Context, provider providers.Provider, model, user
// Larger budget: thinking-capable models (Gemini 2.5/3, GPT-5 reasoning)
// can consume output tokens on reasoning traces. 256 leaves room for a
// 15-word title even when the provider allocates some budget to thinking.
providers.OptMaxTokens: 256,
providers.OptTemperature: 0.3,
providers.OptMaxTokens: 256,
providers.OptTemperature: 0.3,
// Disable extended thinking for title generation — it's a trivial task
// that doesn't benefit from reasoning and defaults (esp. Gemini's "high")
// otherwise eat the entire max_tokens budget, truncating the title to 1 word.
providers.OptThinkingLevel: "off",
},
}
resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{
ModelID: model,
Purpose: "session-title",
MaxOutputTokens: 256,
})
if err != nil {
slog.Warn("title generation failed", "error", err)
internal/agent/usage_caps_runtime.go
+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
}
internal/channels/history_compaction.go
+10 -3
View File
@@ -11,6 +11,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// CompactionConfig configures LLM-based history compaction.
@@ -20,6 +21,7 @@ type CompactionConfig struct {
MaxTokens int // max output tokens for summarization (default 4096)
Provider providers.Provider // LLM provider for summarization
Model string // model to use for summarization
UsageCaps *usagecaps.Service
}
// MaybeCompact checks if compaction is needed for a history key and triggers it in background.
@@ -55,7 +57,7 @@ func (ph *PendingHistory) MaybeCompact(historyKey string, currentCount int, cfg
// CompactGroup performs LLM-based compaction on a pending message group.
// Reused by both auto-compact (channel) and HTTP compact endpoint.
// Returns the number of entries remaining after compaction.
func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, historyKey string, provider providers.Provider, model string, keepRecent, maxTokens int) (int, error) {
func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, historyKey string, provider providers.Provider, model string, keepRecent, maxTokens int, usageCaps *usagecaps.Service) (int, error) {
if keepRecent <= 0 {
keepRecent = 40
}
@@ -94,13 +96,18 @@ func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName,
fmt.Fprintf(&sb, "%s%s: %s\n", prefix, ts, e.Body)
}
resp, err := provider.Chat(ctx, providers.ChatRequest{
req := providers.ChatRequest{
Messages: []providers.Message{{
Role: "user",
Content: "Summarize these group chat messages concisely, preserving key topics, decisions, names, and important context:\n\n" + sb.String(),
}},
Model: model,
Options: map[string]any{"max_tokens": maxTokens, "temperature": 0.3},
}
resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{
ModelID: model,
Purpose: "pending-history-compaction",
MaxOutputTokens: maxTokens,
})
if err != nil {
return 0, fmt.Errorf("llm summarize: %w", err)
@@ -185,7 +192,7 @@ func (ph *PendingHistory) runCompaction(historyKey string, cfg *CompactionConfig
return
}
_, err = CompactGroup(ctx, ph.store, ph.channelName, historyKey, cfg.Provider, cfg.Model, cfg.KeepRecent, cfg.MaxTokens)
_, err = CompactGroup(ctx, ph.store, ph.channelName, historyKey, cfg.Provider, cfg.Model, cfg.KeepRecent, cfg.MaxTokens, cfg.UsageCaps)
if err != nil {
slog.Warn("compaction.failed", "channel", ph.channelName, "key", historyKey, "error", err)
return
internal/channels/instance_loader.go
+9 -2
View File
@@ -16,6 +16,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// reloadStartTimeout bounds how long Reload() will wait for a single channel's
@@ -39,6 +40,7 @@ type InstanceLoader struct {
agentStore store.AgentStore
providerReg *providers.Registry
pendingCompactCfg *config.PendingCompactionConfig
usageCaps *usagecaps.Service
factories map[string]ChannelFactory
manager *Manager
msgBus *bus.MessageBus
@@ -78,6 +80,10 @@ func (l *InstanceLoader) SetPendingCompactionConfig(cfg *config.PendingCompactio
l.pendingCompactCfg = cfg
}
func (l *InstanceLoader) SetUsageCapService(s *usagecaps.Service) {
l.usageCaps = s
}
// RegisterFactory registers a factory for a channel type (e.g., "telegram", "discord").
func (l *InstanceLoader) RegisterFactory(channelType string, factory ChannelFactory) {
l.factories[channelType] = factory
@@ -306,8 +312,9 @@ func (l *InstanceLoader) loadInstance(ctx context.Context, inst store.ChannelIns
if p != nil && model != "" {
cc := &CompactionConfig{
Provider: p,
Model: model,
Provider: p,
Model: model,
UsageCaps: l.usageCaps,
}
if l.pendingCompactCfg != nil {
cc.Threshold = l.pendingCompactCfg.Threshold
internal/consolidation/dreaming_worker.go
+14 -2
View File
@@ -11,9 +11,10 @@ import (
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/bgalert"
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
const (
@@ -31,6 +32,7 @@ type dreamingWorker struct {
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
alertDeps bgalert.AlertDeps
usageCaps *usagecaps.Service
// threshold/debounce are the global defaults. Per-agent overrides come
// from resolveConfig which reads the agent's MemoryConfig.Dreaming JSONB.
@@ -95,6 +97,11 @@ func (w *dreamingWorker) Handle(ctx context.Context, event eventbus.DomainEvent)
ctx = store.WithTenantID(ctx, tid)
}
}
if event.AgentID != "" {
if aid, err := uuid.Parse(event.AgentID); err == nil {
ctx = store.WithAgentID(ctx, aid)
}
}
agentID := event.AgentID
userID := event.UserID
@@ -205,7 +212,7 @@ func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Prov
}
body := strings.Join(summaries, "\n---\n")
resp, err := provider.Chat(ctx, providers.ChatRequest{
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: dreamingSystemPrompt},
{Role: "user", Content: "Session summaries:\n---\n" + body + "\n---"},
@@ -214,6 +221,11 @@ func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Prov
Options: map[string]any{
providers.OptMaxTokens: dreamingMaxTokens,
},
}
resp, err := w.usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{
ModelID: model,
Purpose: "dreaming-synthesis",
MaxOutputTokens: dreamingMaxTokens,
})
if err != nil {
return "", fmt.Errorf("dreaming chat: %w", err)
internal/consolidation/episodic_worker.go
+13 -5
View File
@@ -10,19 +10,21 @@ import (
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/bgalert"
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// episodicWorker handles session.completed events → creates episodic summaries.
type episodicWorker struct {
store store.EpisodicStore
sessions store.SessionCoreStore // for reading session messages during summarization
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
sessions store.SessionCoreStore // for reading session messages during summarization
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
eventBus eventbus.DomainEventBus
alertDeps bgalert.AlertDeps
usageCaps *usagecaps.Service
}
// resolveProvider delegates to shared background provider resolution.
@@ -53,6 +55,7 @@ func (w *episodicWorker) Handle(ctx context.Context, event eventbus.DomainEvent)
if err != nil {
return fmt.Errorf("episodic: invalid agent_id %q: %w", event.AgentID, err)
}
ctx = store.WithAgentID(ctx, agentUUID)
// Build source_id for idempotency
sourceID := fmt.Sprintf("%s:%d", payload.SessionKey, payload.CompactionCount)
@@ -168,13 +171,18 @@ func (w *episodicWorker) summarizeFromMessages(ctx context.Context, provider pro
sctx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
resp, err := provider.Chat(sctx, providers.ChatRequest{
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: summarizationPrompt},
{Role: "user", Content: sb.String()},
},
Model: model,
Options: map[string]any{"max_tokens": 1024, "temperature": 0.3},
}
resp, err := w.usageCaps.Chat(sctx, provider, req, usagecaps.ChatOptions{
ModelID: model,
Purpose: "episodic-summary",
MaxOutputTokens: 1024,
})
if err != nil {
return "", err
internal/consolidation/workers.go
+4
View File
@@ -13,6 +13,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// ConsolidationDeps bundles all dependencies for the consolidation pipeline.
@@ -26,6 +27,7 @@ type ConsolidationDeps struct {
Registry *providers.Registry // provider resolution
Extractor EntityExtractor
AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors
UsageCaps *usagecaps.Service
// AgentStore is optional: when present, the dreaming worker reads
// per-agent overrides from MemoryConfig.Dreaming. If nil, the worker
// uses its built-in defaults for every agent.
@@ -42,6 +44,7 @@ func Register(deps ConsolidationDeps) func() {
registry: deps.Registry,
eventBus: deps.EventBus,
alertDeps: deps.AlertDeps,
usageCaps: deps.UsageCaps,
}
semantic := &semanticWorker{
kgStore: deps.KGStore,
@@ -59,6 +62,7 @@ func Register(deps ConsolidationDeps) func() {
systemConfigs: deps.SystemConfigs,
registry: deps.Registry,
alertDeps: deps.AlertDeps,
usageCaps: deps.UsageCaps,
threshold: dreamingDefaultThreshold,
debounce: dreamingDefaultDebounce,
resolveConfig: newAgentStoreResolver(deps.AgentStore),
internal/gateway/methods/agents_update.go
+7 -1
View File
@@ -60,6 +60,10 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client
if req.Params != nil {
json.Unmarshal(req.Params, &params)
}
var rawParams map[string]json.RawMessage
if req.Params != nil {
_ = json.Unmarshal(req.Params, &rawParams)
}
if params.AgentID == "" {
client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId")))
@@ -104,7 +108,9 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client
if params.IsDefault != nil {
updates["is_default"] = *params.IsDefault
}
if params.BudgetCents != nil {
if rawBudget, ok := rawParams["budget_monthly_cents"]; ok && strings.TrimSpace(string(rawBudget)) == "null" {
updates["budget_monthly_cents"] = nil
} else if params.BudgetCents != nil {
updates["budget_monthly_cents"] = *params.BudgetCents
}
// Per-agent JSONB config overrides
internal/gateway/methods/chat.go
+10 -1
View File
@@ -20,6 +20,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/sessions"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
)
@@ -32,6 +33,7 @@ type ChatMethods struct {
eventBus bus.EventPublisher
postTurn tools.PostTurnProcessor
audioMgr *audio.Manager // for TTS auto-apply on WS responses (nil = disabled)
usageCaps *usagecaps.Service
debouncer *chatDebouncer
}
@@ -46,6 +48,10 @@ func (m *ChatMethods) SetAudioManager(mgr *audio.Manager) {
m.audioMgr = mgr
}
func (m *ChatMethods) SetUsageCapService(s *usagecaps.Service) {
m.usageCaps = s
}
// SetPostTurnProcessor sets the post-turn processor for team task dispatch.
func (m *ChatMethods) SetPostTurnProcessor(pt tools.PostTurnProcessor) {
m.postTurn = pt
@@ -352,7 +358,10 @@ func (m *ChatMethods) dispatchChatSends(requests []chatSendRequest) {
// Use runCtxBase (WithoutCancel + tenant-aware) so title save uses correct tenant.
titleCtx := runCtxBase
go func() {
title := agent.GenerateTitle(titleCtx, agentProvider, agentModel, userMsg)
if uid := loop.UUID(); uid != uuid.Nil {
titleCtx = store.WithAgentID(titleCtx, uid)
}
title := agent.GenerateTitleWithUsageCaps(titleCtx, m.usageCaps, agentProvider, agentModel, userMsg)
if title == "" {
return
}
internal/gateway/server.go
+5
View File
@@ -639,6 +639,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) }
internal/hooks/handlers/prompt.go
+33 -1
View File
@@ -17,6 +17,8 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/hooks"
"github.com/nextlevelbuilder/goclaw/internal/hooks/budget"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// ── Public surface ──────────────────────────────────────────────────────────
@@ -46,6 +48,9 @@ type PromptHandler struct {
// budget checks are skipped (Lite edition behavior).
Budget *budget.Store
// UsageCaps enforces provider/model/tenant cost caps before the hook LLM call.
UsageCaps *usagecaps.Service
// DefaultModel is used when a hook config does not specify one.
// Recommended: "haiku" for cheap evaluation.
DefaultModel string
@@ -155,7 +160,17 @@ func (h *PromptHandler) Execute(ctx context.Context, cfg hooks.HookConfig, ev ho
req := h.buildChatRequest(cfg, ev, resolvedModel)
// 6. Call provider.
resp, err := provider.Chat(ctx, req)
callCtx := ctx
if ev.TenantID != uuid.Nil {
callCtx = store.WithTenantID(callCtx, ev.TenantID)
}
resp, err := h.UsageCaps.Chat(callCtx, provider, req, usagecaps.ChatOptions{
TenantID: ev.TenantID,
ProviderName: provider.Name(),
ModelID: resolvedModel,
Purpose: "hook-prompt",
MaxOutputTokens: promptMaxTokens(req),
})
if err != nil {
// Fail-closed on transport/provider error for blocking events.
if ev.HookEvent.IsBlocking() {
@@ -300,6 +315,23 @@ func (h *PromptHandler) buildChatRequest(cfg hooks.HookConfig, ev hooks.Event, m
}
}
func promptMaxTokens(req providers.ChatRequest) int {
if req.Options == nil {
return 512
}
if v, ok := req.Options[providers.OptMaxTokens]; ok {
switch n := v.(type) {
case int:
return n
case int64:
return int(n)
case float64:
return int(n)
}
}
return 512
}
// sanitizeToolInput canonicalizes a tool_input map into stable JSON with
// sorted keys, stripping only structural noise. Actual injection-attack
// detection is delegated to the evaluator LLM (which has the anti-injection
internal/http/knowledge_graph.go
+9 -1
View File
@@ -7,12 +7,14 @@ import (
kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// KnowledgeGraphHandler handles KG entity/relation management endpoints.
type KnowledgeGraphHandler struct {
store store.KnowledgeGraphStore
providerReg *providers.Registry
usageCaps *usagecaps.Service
}
// NewKnowledgeGraphHandler creates a handler for KG management endpoints.
@@ -20,6 +22,10 @@ func NewKnowledgeGraphHandler(s store.KnowledgeGraphStore, providerReg *provider
return &KnowledgeGraphHandler{store: s, providerReg: providerReg}
}
func (h *KnowledgeGraphHandler) SetUsageCapService(s *usagecaps.Service) {
h.usageCaps = s
}
// NewExtractor creates an Extractor from the given provider name and model.
func (h *KnowledgeGraphHandler) NewExtractor(ctx context.Context, providerName, model string, minConfidence float64) *kg.Extractor {
if h.providerReg == nil || providerName == "" || model == "" {
@@ -29,7 +35,9 @@ func (h *KnowledgeGraphHandler) NewExtractor(ctx context.Context, providerName,
if err != nil {
return nil
}
return kg.NewExtractor(p, model, minConfidence)
extractor := kg.NewExtractor(p, model, minConfidence)
extractor.SetUsageCapService(h.usageCaps)
return extractor
}
// RegisterRoutes registers all KG routes on the given mux.
internal/http/knowledge_graph_handlers.go
+11 -6
View File
@@ -5,6 +5,7 @@ import (
"net/http"
"strconv"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/i18n"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
@@ -156,6 +157,10 @@ func (h *KnowledgeGraphHandler) handleTraverse(w http.ResponseWriter, r *http.Re
func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Request) {
locale := extractLocale(r)
agentID := r.PathValue("agentID")
callCtx := r.Context()
if parsedAgentID, err := uuid.Parse(agentID); err == nil {
callCtx = store.WithAgentID(callCtx, parsedAgentID)
}
var body struct {
Text string `json:"text"`
@@ -176,13 +181,13 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req
return
}
extractor := h.NewExtractor(r.Context(), body.Provider, body.Model, body.MinConf)
extractor := h.NewExtractor(callCtx, body.Provider, body.Model, body.MinConf)
if extractor == nil {
writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidProviderOrModel)})
return
}
result, err := extractor.Extract(r.Context(), body.Text)
result, err := extractor.Extract(callCtx, body.Text)
if err != nil {
slog.Warn("kg.extract failed", "error", err)
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
@@ -213,10 +218,10 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req
}
writeJSON(w, http.StatusOK, map[string]any{
"entities": len(result.Entities),
"relations": len(result.Relations),
"dedup_merged": dedupMerged,
"dedup_flagged": dedupFlagged,
"entities": len(result.Entities),
"relations": len(result.Relations),
"dedup_merged": dedupMerged,
"dedup_flagged": dedupFlagged,
})
}
internal/http/pending_messages.go
+7 -1
View File
@@ -11,6 +11,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// PendingMessagesHandler handles pending message HTTP endpoints.
@@ -22,6 +23,7 @@ type PendingMessagesHandler struct {
maxTokens int // max output tokens for LLM summarization (0 = use default)
cfgProvider string // config-level provider override (empty = resolve from agent)
cfgModel string // config-level model override (empty = resolve from agent)
usageCaps *usagecaps.Service
}
func NewPendingMessagesHandler(s store.PendingMessageStore, agentStore store.AgentStore, providerReg *providers.Registry) *PendingMessagesHandler {
@@ -40,6 +42,10 @@ func (h *PendingMessagesHandler) SetProviderModel(provider, model string) {
h.cfgModel = model
}
func (h *PendingMessagesHandler) SetUsageCapService(s *usagecaps.Service) {
h.usageCaps = s
}
func (h *PendingMessagesHandler) RegisterRoutes(mux *http.ServeMux) {
mux.HandleFunc("GET /v1/pending-messages", h.authMiddleware(h.handleListGroups))
mux.HandleFunc("GET /v1/pending-messages/messages", h.authMiddleware(h.handleListMessages))
@@ -148,7 +154,7 @@ func (h *PendingMessagesHandler) handleCompact(w http.ResponseWriter, r *http.Re
go func() {
ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 180*time.Second)
defer cancel()
remaining, err := channels.CompactGroup(ctx, h.store, req.ChannelName, req.HistoryKey, provider, model, keepRecent, h.maxTokens)
remaining, err := channels.CompactGroup(ctx, h.store, req.ChannelName, req.HistoryKey, provider, model, keepRecent, h.maxTokens, h.usageCaps)
if err != nil {
slog.Warn("compact.failed", "channel", req.ChannelName, "key", req.HistoryKey, "error", err)
} else {
internal/http/provider_verify.go
+10 -1
View File
@@ -16,6 +16,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/i18n"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// HandleVerifyProviderForTest invokes the verify handler directly without auth
@@ -126,8 +127,9 @@ func (h *ProvidersHandler) handleVerifyProvider(w http.ResponseWriter, r *http.R
ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
defer cancel()
ctx = store.WithTenantID(ctx, p.TenantID)
_, err = provider.Chat(ctx, providers.ChatRequest{
reqChat := providers.ChatRequest{
Messages: []providers.Message{
{Role: "user", Content: "hi"},
},
@@ -136,6 +138,13 @@ func (h *ProvidersHandler) handleVerifyProvider(w http.ResponseWriter, r *http.R
// Use a small but safe value — reasoning models need headroom beyond 1 token.
"max_tokens": 50,
},
}
_, err = h.usageCaps.Chat(ctx, provider, reqChat, usagecaps.ChatOptions{
TenantID: p.TenantID,
ProviderName: p.Name,
ModelID: req.Model,
Purpose: "provider-verify",
MaxOutputTokens: 50,
})
if err != nil {
writeJSON(w, http.StatusOK, map[string]any{"valid": false, "error": friendlyVerifyError(err)})
internal/http/providers.go
+6
View File
@@ -22,6 +22,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/permissions"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
)
@@ -39,6 +40,7 @@ type ProvidersHandler struct {
tracingStore store.TracingStore // optional: for provider-scoped pool activity
agents store.AgentCRUDStore // optional: for provider pool activity agent lookup
modelReg providers.ModelRegistry // optional: forward-compat model resolver for Anthropic
usageCaps *usagecaps.Service
}
// NewProvidersHandler creates a handler for provider management endpoints.
@@ -85,6 +87,10 @@ func (h *ProvidersHandler) SetModelRegistry(r providers.ModelRegistry) {
h.modelReg = r
}
func (h *ProvidersHandler) SetUsageCapService(s *usagecaps.Service) {
h.usageCaps = s
}
// resolveAPIBase returns the provider's api_base, falling back to config/env if empty.
// For Ollama/OllamaCloud providers, applies a safety-net normalization: if the stored
// value is missing the /v1 suffix (pre-existing record before write-time normalization),
internal/http/summoner.go
+5 -1
View File
@@ -12,6 +12,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/bus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// Summoning event type constants.
@@ -48,14 +49,16 @@ type AgentSummoner struct {
agents store.AgentStore
providerReg *providers.Registry
msgBus *bus.MessageBus
usageCaps *usagecaps.Service
}
// NewAgentSummoner creates a summoner backed by the given stores and provider registry.
func NewAgentSummoner(agents store.AgentStore, providerReg *providers.Registry, msgBus *bus.MessageBus) *AgentSummoner {
func NewAgentSummoner(agents store.AgentStore, providerReg *providers.Registry, msgBus *bus.MessageBus, usageCaps *usagecaps.Service) *AgentSummoner {
return &AgentSummoner{
agents: agents,
providerReg: providerReg,
msgBus: msgBus,
usageCaps: usageCaps,
}
}
@@ -73,6 +76,7 @@ const singleCallTimeout = 300 * time.Second
func (s *AgentSummoner) SummonAgent(agentID uuid.UUID, tenantID uuid.UUID, providerName, model, description string) {
ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 600*time.Second)
defer cancel()
ctx = store.WithAgentID(ctx, agentID)
s.ensureBackfillFiles(ctx, agentID)
s.emitEvent(agentID, tenantID, SummonEventStarted, "", "")
internal/http/summoner_regenerate.go
+9 -1
View File
@@ -14,6 +14,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/bus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
)
@@ -23,6 +24,7 @@ import (
func (s *AgentSummoner) RegenerateAgent(agentID uuid.UUID, tenantID uuid.UUID, providerName, model, editPrompt string) {
ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 300*time.Second)
defer cancel()
ctx = store.WithAgentID(ctx, agentID)
s.ensureBackfillFiles(ctx, agentID)
@@ -121,7 +123,7 @@ func (s *AgentSummoner) generateFiles(ctx context.Context, providerName, model,
slog.Info("summoning: calling LLM", "provider", providerName, "model", model, "prompt_len", len(prompt))
resp, err := provider.Chat(ctx, providers.ChatRequest{
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: "You are a file generator. Output ONLY the requested XML-tagged files. No extra commentary."},
{Role: "user", Content: prompt},
@@ -133,6 +135,12 @@ func (s *AgentSummoner) generateFiles(ctx context.Context, providerName, model,
providers.OptSessionKey: summonSessionKey,
providers.OptDisableTools: true,
},
}
resp, err := s.usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{
ProviderName: providerName,
ModelID: model,
Purpose: "agent-summoner",
MaxOutputTokens: 8192,
})
if err != nil {
return nil, fmt.Errorf("%s: %w", providerName, err)
internal/http/usage_caps.go
+397
View File
@@ -0,0 +1,397 @@
package http
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"strconv"
"strings"
"time"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/i18n"
"github.com/nextlevelbuilder/goclaw/internal/permissions"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/usage/pricing"
)
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 writeUsageCapError(w http.ResponseWriter, r *http.Request, status int, key string, args ...any) {
writeJSON(w, status, map[string]string{"error": i18n.T(store.LocaleFromContext(r.Context()), key, args...)})
}
func (h *UsageCapsHandler) handleListPolicies(w http.ResponseWriter, r *http.Request) {
scope := store.UsageCapScope{TenantID: tenantIDOrMaster(r)}
policies, err := h.store.ListUsageCapPolicies(r.Context(), scope, true)
if err != nil {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsListPoliciesFailed)
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 {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidJSON)
return
}
p, err := body.toPolicy(tenantIDOrMaster(r))
if err != nil {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidRequest, err.Error())
return
}
if err := h.store.CreateUsageCapPolicy(r.Context(), &p); err != nil {
slog.Warn("usage_caps.create_policy_failed", "error", err)
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsageCapPolicyValidationFailed)
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 {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "policy")
return
}
bodyBytes, err := io.ReadAll(r.Body)
if err != nil {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidJSON)
return
}
patch, err := policyPatchFromBody(bodyBytes)
if err != nil {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidRequest, err.Error())
return
}
p, err := h.store.UpdateUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id, patch)
if err != nil {
if errors.Is(err, store.ErrUsageCapPolicyManaged) {
writeUsageCapError(w, r, http.StatusConflict, i18n.MsgUsageCapPolicyManaged)
return
}
slog.Warn("usage_caps.update_policy_failed", "error", err)
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsageCapPolicyValidationFailed)
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 {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "policy")
return
}
if err := h.store.DeleteUsageCapPolicy(r.Context(), tenantIDOrMaster(r), id); err != nil {
if errors.Is(err, store.ErrUsageCapPolicyManaged) {
writeUsageCapError(w, r, http.StatusConflict, i18n.MsgUsageCapPolicyManaged)
return
}
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsDeletePolicyFailed)
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 {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsUtilizationFailed)
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 {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsageCapsEventsFailed)
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)
writeUsageCapError(w, r, http.StatusBadGateway, i18n.MsgUsagePricingSyncOpenRouterFailed, err.Error())
return
}
count, err := h.store.UpsertPricingCatalog(r.Context(), entries)
if err != nil {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingStoreCatalogFailed)
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 {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingListFailed)
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 {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidJSON)
return
}
providerID, err := uuid.Parse(body.ProviderID)
if err != nil || body.ModelID == "" {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsagePricingProviderModelRequired)
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)
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgUsagePricingOverrideValidationFailed)
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 {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "provider")
return
}
}
rows, err := h.store.ListPricingOverrides(r.Context(), store.UsagePricingQuery{TenantID: tenantIDOrMaster(r), ProviderID: providerID})
if err != nil {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingListOverridesFailed)
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 {
writeUsageCapError(w, r, http.StatusBadRequest, i18n.MsgInvalidID, "override")
return
}
if err := h.store.DeletePricingOverride(r.Context(), tenantIDOrMaster(r), id); err != nil {
writeUsageCapError(w, r, http.StatusInternalServerError, i18n.MsgUsagePricingDeleteOverrideFailed)
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 policyPatchFromBody(bodyBytes []byte) (store.UsageCapPolicyPatch, error) {
var body policyBody
if err := json.Unmarshal(bodyBytes, &body); err != nil {
return store.UsageCapPolicyPatch{}, err
}
patch, err := body.toPatch()
if err != nil {
return patch, err
}
var raw map[string]json.RawMessage
if err := json.Unmarshal(bodyBytes, &raw); err == nil {
if isJSONNull(raw["max_tokens"]) {
var v *int64
patch.MaxTokens = &v
}
if isJSONNull(raw["max_cost_micros"]) || isJSONNull(raw["max_cost_usd"]) {
var v *int64
patch.MaxCostMicros = &v
}
}
return patch, nil
}
func isJSONNull(raw json.RawMessage) bool {
return strings.TrimSpace(string(raw)) == "null"
}
func tenantIDOrMaster(r *http.Request) uuid.UUID {
if tid := store.TenantIDFromContext(r.Context()); tid != uuid.Nil {
return tid
}
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
}
internal/http/usage_caps_policy_patch_test.go
+35
View File
@@ -0,0 +1,35 @@
package http
import "testing"
func TestPolicyPatchFromBodyClearsTokenAndCostLimits(t *testing.T) {
patch, err := policyPatchFromBody([]byte(`{"max_tokens":null,"max_cost_usd":null}`))
if err != nil {
t.Fatalf("policyPatchFromBody: %v", err)
}
if patch.MaxTokens == nil {
t.Fatal("MaxTokens patch missing")
}
if *patch.MaxTokens != nil {
t.Fatalf("MaxTokens = %v, want nil clear", **patch.MaxTokens)
}
if patch.MaxCostMicros == nil {
t.Fatal("MaxCostMicros patch missing")
}
if *patch.MaxCostMicros != nil {
t.Fatalf("MaxCostMicros = %v, want nil clear", **patch.MaxCostMicros)
}
}
func TestPolicyPatchFromBodyPreservesProvidedCostUSD(t *testing.T) {
patch, err := policyPatchFromBody([]byte(`{"max_cost_usd":12.5}`))
if err != nil {
t.Fatalf("policyPatchFromBody: %v", err)
}
if patch.MaxCostMicros == nil || *patch.MaxCostMicros == nil {
t.Fatal("MaxCostMicros patch missing")
}
if got := **patch.MaxCostMicros; got != 12_500_000 {
t.Fatalf("MaxCostMicros = %d, want 12500000", got)
}
}
internal/http/usage_caps_test.go
+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")
}
}
internal/i18n/catalog_en.go
+19 -4
View File
@@ -93,6 +93,21 @@ func init() {
// Provider
MsgProviderReqFailed: "%s: request failed: %s",
// Usage caps / pricing
MsgUsageCapsListPoliciesFailed: "failed to list usage cap policies",
MsgUsageCapPolicyValidationFailed: "usage cap policy validation failed",
MsgUsageCapPolicyManaged: "managed usage cap policies cannot be modified",
MsgUsageCapsDeletePolicyFailed: "failed to delete usage cap policy",
MsgUsageCapsUtilizationFailed: "failed to load usage cap utilization",
MsgUsageCapsEventsFailed: "failed to load usage cap events",
MsgUsagePricingSyncOpenRouterFailed: "failed to sync OpenRouter pricing: %s",
MsgUsagePricingStoreCatalogFailed: "failed to store pricing catalog",
MsgUsagePricingListFailed: "failed to list model pricing",
MsgUsagePricingProviderModelRequired: "provider_id and model_id are required",
MsgUsagePricingOverrideValidationFailed: "pricing override validation failed",
MsgUsagePricingListOverridesFailed: "failed to list pricing overrides",
MsgUsagePricingDeleteOverrideFailed: "failed to delete pricing override",
// Unknown method
MsgUnknownMethod: "unknown method: %s",
@@ -206,10 +221,10 @@ func init() {
MsgTenantScopeRequired: "tenant scope is required for this operation",
// TTS / Voices
MsgTtsUnknownModel: "unknown tts model: %s",
MsgVoicesListFailed: "failed to list voices: %s",
MsgTtsGeminiInvalidVoice: "invalid Gemini voice: %s",
MsgTtsGeminiSpeakerLimit: "Gemini TTS supports at most 2 speakers",
MsgTtsUnknownModel: "unknown tts model: %s",
MsgVoicesListFailed: "failed to list voices: %s",
MsgTtsGeminiInvalidVoice: "invalid Gemini voice: %s",
MsgTtsGeminiSpeakerLimit: "Gemini TTS supports at most 2 speakers",
MsgTtsGeminiInvalidModel: "invalid Gemini TTS model: %s",
MsgTtsGeminiTextOnly: "Gemini refused to generate audio. Try simpler text without translation or commentary.",
MsgTtsParamOutOfRange: "TTS param %q value %v is out of range [%v, %v]",
internal/i18n/catalog_vi.go
+19 -4
View File
@@ -93,6 +93,21 @@ func init() {
// Provider
MsgProviderReqFailed: "%s: yêu cầu thất bại: %s",
// Usage caps / pricing
MsgUsageCapsListPoliciesFailed: "không thể liệt kê chính sách usage cap",
MsgUsageCapPolicyValidationFailed: "xác thực chính sách usage cap thất bại",
MsgUsageCapPolicyManaged: "không thể chỉnh sửa chính sách usage cap do hệ thống quản lý",
MsgUsageCapsDeletePolicyFailed: "không thể xóa chính sách usage cap",
MsgUsageCapsUtilizationFailed: "không thể tải mức sử dụng usage cap",
MsgUsageCapsEventsFailed: "không thể tải sự kiện usage cap",
MsgUsagePricingSyncOpenRouterFailed: "không thể đồng bộ giá OpenRouter: %s",
MsgUsagePricingStoreCatalogFailed: "không thể lưu catalog giá",
MsgUsagePricingListFailed: "không thể liệt kê giá model",
MsgUsagePricingProviderModelRequired: "provider_id và model_id là bắt buộc",
MsgUsagePricingOverrideValidationFailed: "xác thực override giá thất bại",
MsgUsagePricingListOverridesFailed: "không thể liệt kê override giá",
MsgUsagePricingDeleteOverrideFailed: "không thể xóa override giá",
// Unknown method
MsgUnknownMethod: "phương thức không xác định: %s",
@@ -206,10 +221,10 @@ func init() {
MsgTenantScopeRequired: "cần xác định tenant để thực hiện thao tác này",
// TTS / Giọng đọc
MsgTtsUnknownModel: "model tts không hỗ trợ: %s",
MsgVoicesListFailed: "không tải được danh sách giọng đọc: %s",
MsgTtsGeminiInvalidVoice: "giọng đọc Gemini không hợp lệ: %s",
MsgTtsGeminiSpeakerLimit: "Gemini TTS hỗ trợ tối đa 2 người nói",
MsgTtsUnknownModel: "model tts không hỗ trợ: %s",
MsgVoicesListFailed: "không tải được danh sách giọng đọc: %s",
MsgTtsGeminiInvalidVoice: "giọng đọc Gemini không hợp lệ: %s",
MsgTtsGeminiSpeakerLimit: "Gemini TTS hỗ trợ tối đa 2 người nói",
MsgTtsGeminiInvalidModel: "mô hình Gemini TTS không hợp lệ: %s",
MsgTtsGeminiTextOnly: "Gemini từ chối tạo âm thanh. Vui lòng thử văn bản đơn giản hơn, không dịch hay bình luận.",
MsgTtsParamOutOfRange: "tham số TTS %q có giá trị %v nằm ngoài phạm vi [%v, %v]",
internal/i18n/catalog_zh.go
+19 -4
View File
@@ -93,6 +93,21 @@ func init() {
// Provider
MsgProviderReqFailed: "%s：请求失败：%s",
// Usage caps / pricing
MsgUsageCapsListPoliciesFailed: "无法列出 usage cap 策略",
MsgUsageCapPolicyValidationFailed: "usage cap 策略验证失败",
MsgUsageCapPolicyManaged: "无法修改系统托管的 usage cap 策略",
MsgUsageCapsDeletePolicyFailed: "无法删除 usage cap 策略",
MsgUsageCapsUtilizationFailed: "无法加载 usage cap 使用量",
MsgUsageCapsEventsFailed: "无法加载 usage cap 事件",
MsgUsagePricingSyncOpenRouterFailed: "无法同步 OpenRouter 价格：%s",
MsgUsagePricingStoreCatalogFailed: "无法保存价格目录",
MsgUsagePricingListFailed: "无法列出模型价格",
MsgUsagePricingProviderModelRequired: "provider_id 和 model_id 是必填项",
MsgUsagePricingOverrideValidationFailed: "价格覆盖验证失败",
MsgUsagePricingListOverridesFailed: "无法列出价格覆盖",
MsgUsagePricingDeleteOverrideFailed: "无法删除价格覆盖",
// Unknown method
MsgUnknownMethod: "未知方法：%s",
@@ -206,10 +221,10 @@ func init() {
MsgTenantScopeRequired: "此操作需要指定租户范围",
// TTS / 声音
MsgTtsUnknownModel: "未知的 tts 模型：%s",
MsgVoicesListFailed: "获取声音列表失败：%s",
MsgTtsGeminiInvalidVoice: "无效的 Gemini 声音：%s",
MsgTtsGeminiSpeakerLimit: "Gemini TTS 最多支持 2 位发言人",
MsgTtsUnknownModel: "未知的 tts 模型：%s",
MsgVoicesListFailed: "获取声音列表失败：%s",
MsgTtsGeminiInvalidVoice: "无效的 Gemini 声音：%s",
MsgTtsGeminiSpeakerLimit: "Gemini TTS 最多支持 2 位发言人",
MsgTtsGeminiInvalidModel: "无效的 Gemini TTS 模型：%s",
MsgTtsGeminiTextOnly: "Gemini 拒绝生成音频。请尝试更简单的文本，不要翻译或添加评论。",
MsgTtsParamOutOfRange: "TTS 参数 %q 的值 %v 超出范围 [%v, %v]",
internal/i18n/keys.go
+63 -48
View File
@@ -94,6 +94,21 @@ const (
// --- Provider ---
MsgProviderReqFailed = "error.provider_request_failed" // "%s: request failed: %s"
// --- Usage caps / pricing ---
MsgUsageCapsListPoliciesFailed = "usage_caps.list_policies_failed"
MsgUsageCapPolicyValidationFailed = "usage_caps.policy_validation_failed"
MsgUsageCapPolicyManaged = "usage_caps.policy_managed"
MsgUsageCapsDeletePolicyFailed = "usage_caps.delete_policy_failed"
MsgUsageCapsUtilizationFailed = "usage_caps.utilization_failed"
MsgUsageCapsEventsFailed = "usage_caps.events_failed"
MsgUsagePricingSyncOpenRouterFailed = "usage_pricing.sync_openrouter_failed"
MsgUsagePricingStoreCatalogFailed = "usage_pricing.store_catalog_failed"
MsgUsagePricingListFailed = "usage_pricing.list_failed"
MsgUsagePricingProviderModelRequired = "usage_pricing.provider_model_required"
MsgUsagePricingOverrideValidationFailed = "usage_pricing.override_validation_failed"
MsgUsagePricingListOverridesFailed = "usage_pricing.list_overrides_failed"
MsgUsagePricingDeleteOverrideFailed = "usage_pricing.delete_override_failed"
// --- Unknown method ---
MsgUnknownMethod = "error.unknown_method" // "unknown method: %s"
@@ -123,14 +138,14 @@ const (
MsgInvalidVisibility = "error.invalid_visibility" // "invalid visibility %q: must be one of private, public"
// --- Package updates (Phase 4+5) ---
MsgPackageNotInstalled = "packages.update.not_installed" // "Package {name} is not installed"
MsgPackageUpdateLocked = "packages.update.locked" // "Package {name} is being updated by another request"
MsgPackageNotInstalled = "packages.update.not_installed" // "Package {name} is not installed"
MsgPackageUpdateLocked = "packages.update.locked" // "Package {name} is being updated by another request"
MsgReleaseNotFound = "packages.update.release_not_found" // "Release {tag} not found for {repo}"
MsgAssetNotFound = "packages.update.asset_not_found" // "No compatible asset for {os}/{arch}"
MsgAssetNotFound = "packages.update.asset_not_found" // "No compatible asset for {os}/{arch}"
MsgChecksumMismatch = "packages.update.checksum_mismatch" // "Checksum mismatch for {name}"
MsgUpdateSwapFailed = "packages.update.swap_failed" // "Failed to install {name}; previous version restored"
MsgUpdateManifestDesync = "packages.update.manifest_desync" // "Binary updated but manifest save failed — manual recovery required for {name}"
MsgUpdateCacheStale = "packages.update.cache_stale" // "Updates cache stale; run refresh before applying an update"
MsgUpdateSwapFailed = "packages.update.swap_failed" // "Failed to install {name}; previous version restored"
MsgUpdateManifestDesync = "packages.update.manifest_desync" // "Binary updated but manifest save failed — manual recovery required for {name}"
MsgUpdateCacheStale = "packages.update.cache_stale" // "Updates cache stale; run refresh before applying an update"
// Package update source labels
MsgPackagesUpdatesSourceGithub = "packages.updates.source.github" // "GitHub"
@@ -240,15 +255,15 @@ const (
MsgInvalidRole = "error.invalid_role" // "invalid role: allowed values are owner, admin, operator, member, viewer"
// --- TTS / Voices ---
MsgTtsUnknownModel = "error.tts_unknown_model" // "unknown tts model: %s"
MsgVoicesListFailed = "error.voices_list_failed" // "failed to list voices: %s"
MsgTtsGeminiInvalidVoice = "error.tts_gemini_invalid_voice" // "invalid Gemini voice: %s"
MsgTtsGeminiSpeakerLimit = "error.tts_gemini_speaker_limit" // "Gemini TTS supports at most 2 speakers"
MsgTtsGeminiInvalidModel = "error.tts_gemini_invalid_model" // "invalid Gemini TTS model: %s"
MsgTtsGeminiTextOnly = "error.tts_gemini_text_only" // "Gemini refused to generate audio; try simpler text without translation or commentary"
MsgTtsParamOutOfRange = "error.tts_param_out_of_range" // "TTS param %q value %v is out of range [%v, %v]"
MsgTtsParamUnknownKey = "error.tts_param_unknown_key" // "TTS param %q is not supported by this provider"
MsgTtsMiniMaxVoicesFailed = "error.tts_minimax_voices_failed" // "failed to fetch MiniMax voices: %s"
MsgTtsUnknownModel = "error.tts_unknown_model" // "unknown tts model: %s"
MsgVoicesListFailed = "error.voices_list_failed" // "failed to list voices: %s"
MsgTtsGeminiInvalidVoice = "error.tts_gemini_invalid_voice" // "invalid Gemini voice: %s"
MsgTtsGeminiSpeakerLimit = "error.tts_gemini_speaker_limit" // "Gemini TTS supports at most 2 speakers"
MsgTtsGeminiInvalidModel = "error.tts_gemini_invalid_model" // "invalid Gemini TTS model: %s"
MsgTtsGeminiTextOnly = "error.tts_gemini_text_only" // "Gemini refused to generate audio; try simpler text without translation or commentary"
MsgTtsParamOutOfRange = "error.tts_param_out_of_range" // "TTS param %q value %v is out of range [%v, %v]"
MsgTtsParamUnknownKey = "error.tts_param_unknown_key" // "TTS param %q is not supported by this provider"
MsgTtsMiniMaxVoicesFailed = "error.tts_minimax_voices_failed" // "failed to fetch MiniMax voices: %s"
// --- STT ---
MsgSTTAllProvidersFailed = "error.stt_all_providers_failed" // "All STT providers failed"
@@ -264,50 +279,50 @@ const (
MsgTenantScopeRequired = "error.tenant_scope_required" // "tenant scope is required for this operation"
// --- Webhooks ---
MsgWebhookAuthFailed = "webhook.auth_failed" // "webhook authentication failed"
MsgWebhookHMACInvalid = "webhook.hmac_invalid" // "HMAC signature is invalid"
MsgWebhookHMACTimestampSkew = "webhook.hmac_timestamp_skew" // "request timestamp outside acceptable window"
MsgWebhookBearerRequiredHMAC = "webhook.bearer_required_hmac" // "this webhook requires HMAC authentication"
MsgWebhookRevoked = "webhook.revoked" // "webhook has been revoked"
MsgWebhookKindMismatch = "webhook.kind_mismatch" // "request kind does not match webhook configuration"
MsgWebhookRateLimited = "webhook.rate_limited" // "webhook rate limit exceeded"
MsgWebhookBodyTooLarge = "webhook.body_too_large" // "request body exceeds size limit"
MsgWebhookIdempotencyConflict = "webhook.idempotency_conflict" // "idempotency key conflict: request body mismatch"
MsgWebhookTenantMismatch = "webhook.tenant_mismatch" // "webhook tenant mismatch"
MsgWebhookAgentNotFound = "webhook.agent_not_found" // "webhook agent not found"
MsgWebhookChannelNotFound = "webhook.channel_not_found" // "webhook channel not found"
MsgWebhookMediaSSRFBlocked = "webhook.media_ssrf_blocked" // "media URL blocked by SSRF policy"
MsgWebhookMediaTooLarge = "webhook.media_too_large" // "media file exceeds size limit"
MsgWebhookMediaMIMEDenied = "webhook.media_mime_denied" // "media MIME type is not allowed"
MsgWebhookCallbackURLInvalid = "webhook.callback_url_invalid" // "callback URL is invalid or blocked"
MsgWebhookLLMTimeout = "webhook.llm_timeout" // "LLM processing timed out"
MsgWebhookLaneSaturated = "webhook.lane_saturated" // "webhook processing lane is at capacity"
MsgWebhookAuthFailed = "webhook.auth_failed" // "webhook authentication failed"
MsgWebhookHMACInvalid = "webhook.hmac_invalid" // "HMAC signature is invalid"
MsgWebhookHMACTimestampSkew = "webhook.hmac_timestamp_skew" // "request timestamp outside acceptable window"
MsgWebhookBearerRequiredHMAC = "webhook.bearer_required_hmac" // "this webhook requires HMAC authentication"
MsgWebhookRevoked = "webhook.revoked" // "webhook has been revoked"
MsgWebhookKindMismatch = "webhook.kind_mismatch" // "request kind does not match webhook configuration"
MsgWebhookRateLimited = "webhook.rate_limited" // "webhook rate limit exceeded"
MsgWebhookBodyTooLarge = "webhook.body_too_large" // "request body exceeds size limit"
MsgWebhookIdempotencyConflict = "webhook.idempotency_conflict" // "idempotency key conflict: request body mismatch"
MsgWebhookTenantMismatch = "webhook.tenant_mismatch" // "webhook tenant mismatch"
MsgWebhookAgentNotFound = "webhook.agent_not_found" // "webhook agent not found"
MsgWebhookChannelNotFound = "webhook.channel_not_found" // "webhook channel not found"
MsgWebhookMediaSSRFBlocked = "webhook.media_ssrf_blocked" // "media URL blocked by SSRF policy"
MsgWebhookMediaTooLarge = "webhook.media_too_large" // "media file exceeds size limit"
MsgWebhookMediaMIMEDenied = "webhook.media_mime_denied" // "media MIME type is not allowed"
MsgWebhookCallbackURLInvalid = "webhook.callback_url_invalid" // "callback URL is invalid or blocked"
MsgWebhookLLMTimeout = "webhook.llm_timeout" // "LLM processing timed out"
MsgWebhookLaneSaturated = "webhook.lane_saturated" // "webhook processing lane is at capacity"
MsgWebhookLocalhostOnlyViolation = "webhook.localhost_only_violation" // "this webhook is restricted to localhost callers"
MsgWebhookMediaChannelUnsupported = "webhook.media_channel_unsupported" // "channel does not support media attachments"
MsgWebhookIPDenied = "webhook.ip_denied" // "request origin is not in the IP allowlist"
MsgWebhookEncryptionUnavailable = "webhook.encryption_unavailable" // "webhook encryption key not configured; set GOCLAW_ENCRYPTION_KEY to enable webhooks"
// --- Workstation permissions ---
MsgWorkstationCmdDenied = "error.workstation_cmd_denied" // "command denied by workstation policy: %s"
MsgWorkstationEnvDenied = "error.workstation_env_denied" // "env var denied by policy: %s"
MsgWorkstationInputInvalid = "error.workstation_input_invalid" // "command contains invalid characters: %s"
MsgWorkstationRateLimit = "error.workstation_rate_limit" // "workstation rate limit exceeded"
MsgWorkstationPermNotFound = "error.workstation_perm_not_found" // "permission entry not found: %s"
MsgWorkstationCmdDenied = "error.workstation_cmd_denied" // "command denied by workstation policy: %s"
MsgWorkstationEnvDenied = "error.workstation_env_denied" // "env var denied by policy: %s"
MsgWorkstationInputInvalid = "error.workstation_input_invalid" // "command contains invalid characters: %s"
MsgWorkstationRateLimit = "error.workstation_rate_limit" // "workstation rate limit exceeded"
MsgWorkstationPermNotFound = "error.workstation_perm_not_found" // "permission entry not found: %s"
// --- Workstation activity (Phase 7) ---
MsgWorkstationActivityTitle = "ui.workstations.activity.title" // "Recent Activity"
MsgWorkstationActionExec = "ui.workstations.activity.action_exec" // "Exec"
MsgWorkstationActionDeny = "ui.workstations.activity.action_deny" // "Denied"
MsgWorkstationActivityTitle = "ui.workstations.activity.title" // "Recent Activity"
MsgWorkstationActionExec = "ui.workstations.activity.action_exec" // "Exec"
MsgWorkstationActionDeny = "ui.workstations.activity.action_deny" // "Denied"
// --- Workstation ---
MsgWorkstationNotFound = "error.workstation_not_found" // "workstation not found: %s"
MsgWorkstationKeyExists = "error.workstation_key_exists" // "workstation key already in use: %s"
MsgInvalidBackend = "error.invalid_backend" // "invalid backend type: %s (must be ssh|docker)"
MsgWorkstationInactive = "error.workstation_inactive" // "workstation is inactive: %s"
MsgInvalidMetadataShape = "error.invalid_metadata_shape" // "invalid metadata for %s backend: %s"
MsgWorkstationRequired = "error.workstation_required" // "no workstation bound to agent; pass workstation_id"
MsgWorkstationNotFound = "error.workstation_not_found" // "workstation not found: %s"
MsgWorkstationKeyExists = "error.workstation_key_exists" // "workstation key already in use: %s"
MsgInvalidBackend = "error.invalid_backend" // "invalid backend type: %s (must be ssh|docker)"
MsgWorkstationInactive = "error.workstation_inactive" // "workstation is inactive: %s"
MsgInvalidMetadataShape = "error.invalid_metadata_shape" // "invalid metadata for %s backend: %s"
MsgWorkstationRequired = "error.workstation_required" // "no workstation bound to agent; pass workstation_id"
MsgWorkstationAccessDenied = "error.workstation_access_denied" // "agent %s not authorized for workstation %s"
MsgBackendNotReady = "error.backend_not_ready" // "workstation backend not ready: %s"
MsgBackendNotReady = "error.backend_not_ready" // "workstation backend not ready: %s"
// --- Hooks ---
MsgHookInvalidMatcher = "hook.invalid_matcher" // "invalid matcher regex: %s"
internal/knowledgegraph/extractor.go
+17 -3
View File
@@ -9,6 +9,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
)
// ExtractionResult holds entities and relations extracted from text.
@@ -22,6 +23,7 @@ type Extractor struct {
provider providers.Provider
model string
minConfidence float64
usageCaps *usagecaps.Service
}
// NewExtractor creates a new Extractor with the given provider, model, and confidence threshold.
@@ -32,6 +34,11 @@ func NewExtractor(provider providers.Provider, model string, minConfidence float
return &Extractor{provider: provider, model: model, minConfidence: minConfidence}
}
// SetUsageCapService enables cost enforcement for LLM extraction calls.
func (e *Extractor) SetUsageCapService(s *usagecaps.Service) {
e.usageCaps = s
}
const maxChunkChars = 12000
// Extract calls the LLM to extract entities and relations from text.
@@ -72,7 +79,11 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR
},
}
resp, err := e.provider.Chat(ctx, req)
resp, err := e.usageCaps.Chat(ctx, e.provider, req, usagecaps.ChatOptions{
ModelID: e.model,
Purpose: "knowledge-graph-extract",
MaxOutputTokens: 8192,
})
if err != nil {
return nil, fmt.Errorf("kg extraction LLM call: %w", err)
}
@@ -85,7 +96,11 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR
text = text[:retryMaxChars] + "\n\n[...truncated]"
}
req.Messages[1].Content = text
resp, err = e.provider.Chat(ctx, req)
resp, err = e.usageCaps.Chat(ctx, e.provider, req, usagecaps.ChatOptions{
ModelID: e.model,
Purpose: "knowledge-graph-extract-retry",
MaxOutputTokens: 8192,
})
if err != nil {
return nil, fmt.Errorf("kg extraction LLM retry: %w", err)
}
@@ -286,4 +301,3 @@ func stripCodeBlock(s string) string {
}
return strings.TrimSpace(s)
}
internal/providers/anthropic_stream.go
+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
internal/providers/codex.go
+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.
internal/providers/model_fallback.go
+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,
internal/providers/model_fallback_test.go
+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",
internal/providers/openai_chat.go
+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
internal/providers/openai_http.go
+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
internal/providers/openai_types.go
+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 {
internal/providers/types.go
+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"`
}
internal/store/pg/agents.go
+95
View File
@@ -133,6 +133,12 @@ func (s *PGAgentStore) Create(ctx context.Context, agent *store.AgentData) error
if err != nil {
return err
}
if agent.BudgetMonthlyCents != nil {
if err := s.syncAgentBudgetUsageCap(ctx, tenantID, agent.ID, agent.BudgetMonthlyCents); err != nil {
_, _ = s.db.ExecContext(ctx, "DELETE FROM agents WHERE id = $1", agent.ID)
return err
}
}
// Generate embedding for new agent with frontmatter
if agent.Frontmatter != "" && s.embProvider != nil {
@@ -189,6 +195,22 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str
if len(updates) == 0 {
return nil
}
var err error
var budgetCents *int
rawBudget, syncBudget := updates["budget_monthly_cents"]
if syncBudget {
budgetCents, err = budgetCentsFromUpdate(rawBudget)
if err != nil {
return err
}
}
var budgetTenantID uuid.UUID
if syncBudget {
budgetTenantID, err = s.agentTenantID(ctx, id)
if err != nil {
return err
}
}
// Coerce NOT NULL columns: null → default to prevent constraint violations.
// Promoted TEXT columns (migration 000037): null → empty string.
@@ -248,6 +270,11 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str
return err
}
}
if syncBudget {
if err := s.syncAgentBudgetUsageCap(ctx, budgetTenantID, id, budgetCents); err != nil {
return err
}
}
// Regenerate embedding when frontmatter changes
if _, hasFrontmatter := updates["frontmatter"]; hasFrontmatter && s.embProvider != nil {
@@ -262,6 +289,74 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str
return nil
}
func (s *PGAgentStore) agentTenantID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) {
if !store.IsCrossTenant(ctx) {
tid := store.TenantIDFromContext(ctx)
if tid != uuid.Nil {
return tid, nil
}
}
var tenantID uuid.UUID
if err := s.db.QueryRowContext(ctx, "SELECT tenant_id FROM agents WHERE id = $1 AND deleted_at IS NULL", id).Scan(&tenantID); err != nil {
return uuid.Nil, fmt.Errorf("agent not found: %s", id)
}
return tenantID, nil
}
func budgetCentsFromUpdate(v any) (*int, error) {
if v == nil {
return nil, nil
}
var cents int
switch n := v.(type) {
case int:
cents = n
case int32:
cents = int(n)
case int64:
cents = int(n)
case float64:
if n != float64(int(n)) {
return nil, fmt.Errorf("budget_monthly_cents must be an integer")
}
cents = int(n)
default:
return nil, fmt.Errorf("budget_monthly_cents must be an integer")
}
if cents < 0 {
return nil, fmt.Errorf("budget_monthly_cents must be non-negative")
}
return &cents, nil
}
func (s *PGAgentStore) syncAgentBudgetUsageCap(ctx context.Context, tenantID, agentID uuid.UUID, budgetCents *int) error {
if budgetCents == nil || *budgetCents <= 0 {
_, err := s.db.ExecContext(ctx,
`DELETE FROM usage_cap_policies
WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`,
tenantID, agentID, store.UsageCapSourceAgentBudget)
return err
}
costMicros := int64(*budgetCents) * 10000
_, err := s.db.ExecContext(ctx, `
INSERT INTO usage_cap_policies (
tenant_id, agent_id, window_key, max_cost_micros, enabled, priority, source
) VALUES ($1,$2,'month',$3,true,90,$4)
ON CONFLICT (tenant_id, agent_id) WHERE source = 'agent_budget_monthly_cents'
DO UPDATE SET
window_key='month',
provider_id=NULL,
provider_type=NULL,
model_id=NULL,
max_tokens=NULL,
max_cost_micros=EXCLUDED.max_cost_micros,
enabled=true,
priority=EXCLUDED.priority,
updated_at=now()`,
tenantID, agentID, costMicros, store.UsageCapSourceAgentBudget)
return err
}
func isEmptyOrNullJSONUpdate(v any) bool {
if v == nil {
return true
internal/store/pg/agents_update_null_coerce_test.go
+96
View File
@@ -1,8 +1,14 @@
package pg
import (
"context"
"database/sql"
"encoding/json"
"errors"
"testing"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
func TestIsEmptyOrNullJSONUpdate(t *testing.T) {
@@ -30,3 +36,93 @@ func TestIsEmptyOrNullJSONUpdate(t *testing.T) {
})
}
}
func TestPGAgentStoreSyncsMonthlyBudgetUsageCap(t *testing.T) {
db := hooksTestDB(t)
tenantID, _ := seedTenantAndAgent(t, db)
ctx := tenantScopedCtx(tenantID)
agentStore := NewPGAgentStore(db)
budgetCents := 123
agentID := uuid.New()
agent := &store.AgentData{
BaseModel: store.BaseModel{ID: agentID},
TenantID: tenantID, AgentKey: "budget-agent-" + agentID.String(),
OwnerID: "owner", Provider: "openai", Model: "gpt-test",
AgentType: store.AgentTypePredefined, Status: store.AgentStatusActive,
BudgetMonthlyCents: &budgetCents,
}
if err := agentStore.Create(ctx, agent); err != nil {
t.Fatalf("Create: %v", err)
}
assertAgentBudgetPolicy(t, db, tenantID, agentID, 1230000)
var generatedPolicyID uuid.UUID
if err := db.QueryRowContext(context.Background(),
`SELECT id FROM usage_cap_policies
WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`,
tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&generatedPolicyID); err != nil {
t.Fatalf("query generated policy id: %v", err)
}
if err := NewPGUsageCapStore(db).DeleteUsageCapPolicy(ctx, tenantID, generatedPolicyID); err == nil {
t.Fatal("DeleteUsageCapPolicy deleted generated agent budget policy")
} else if !errors.Is(err, store.ErrUsageCapPolicyManaged) {
t.Fatalf("DeleteUsageCapPolicy error = %v, want managed policy error", err)
}
disablePolicy := false
if _, err := NewPGUsageCapStore(db).UpdateUsageCapPolicy(ctx, tenantID, generatedPolicyID, store.UsageCapPolicyPatch{
Enabled: &disablePolicy,
}); err == nil {
t.Fatal("UpdateUsageCapPolicy updated generated agent budget policy")
} else if !errors.Is(err, store.ErrUsageCapPolicyManaged) {
t.Fatalf("UpdateUsageCapPolicy error = %v, want managed policy error", err)
}
assertAgentBudgetPolicy(t, db, tenantID, agentID, 1230000)
if err := agentStore.Update(ctx, agentID, map[string]any{"budget_monthly_cents": float64(456)}); err != nil {
t.Fatalf("Update budget: %v", err)
}
assertAgentBudgetPolicy(t, db, tenantID, agentID, 4560000)
if err := agentStore.Update(ctx, agentID, map[string]any{"budget_monthly_cents": nil}); err != nil {
t.Fatalf("Clear budget: %v", err)
}
var count int
if err := db.QueryRowContext(context.Background(),
`SELECT COUNT(*) FROM usage_cap_policies
WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`,
tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&count); err != nil {
t.Fatalf("count cleared policy: %v", err)
}
if count != 0 {
t.Fatalf("agent budget policy count after clear = %d, want 0", count)
}
}
func assertAgentBudgetPolicy(t *testing.T, db rowQueryer, tenantID, agentID uuid.UUID, wantCostMicros int64) {
t.Helper()
var windowKey, source string
var costMicros int64
var count int
if err := db.QueryRowContext(context.Background(),
`SELECT COUNT(*), COALESCE(MAX(window_key), ''), COALESCE(MAX(source), ''), COALESCE(MAX(max_cost_micros), 0)
FROM usage_cap_policies
WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`,
tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&count, &windowKey, &source, &costMicros); err != nil {
t.Fatalf("query agent budget policy: %v", err)
}
if count != 1 {
t.Fatalf("agent budget policy count = %d, want 1", count)
}
if windowKey != store.UsageCapWindowMonth {
t.Fatalf("window_key = %q, want month", windowKey)
}
if source != store.UsageCapSourceAgentBudget {
t.Fatalf("source = %q, want %q", source, store.UsageCapSourceAgentBudget)
}
if costMicros != wantCostMicros {
t.Fatalf("max_cost_micros = %d, want %d", costMicros, wantCostMicros)
}
}
type rowQueryer interface {
QueryRowContext(context.Context, string, ...any) *sql.Row
}
internal/store/pg/factory.go
+38 -37
View File
@@ -24,49 +24,50 @@ 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),
BrowserCookies: NewPGBrowserCookieStore(db, cfg.EncryptionKey),
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),
BrowserCookies: NewPGBrowserCookieStore(db, cfg.EncryptionKey),
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.
internal/store/pg/usage_caps.go
+431
View File
@@ -0,0 +1,431 @@
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
}
if p.Source == "" {
p.Source = store.UsageCapSourceManual
}
const q = `
INSERT INTO usage_cap_policies (
id, tenant_id, agent_id, provider_id, provider_type, model_id, window_key,
max_tokens, max_cost_micros, source, enabled, priority
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
RETURNING created_at, updated_at`
return s.db.QueryRowContext(ctx, q,
p.ID, p.TenantID, uuidPtrVal(p.AgentID), uuidPtrVal(p.ProviderID),
nullEmpty(p.ProviderType), nullEmpty(p.ModelID), p.Window,
intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros), p.Source, 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 p.Source == store.UsageCapSourceAgentBudget {
return nil, fmt.Errorf("%w: agent monthly budget", store.ErrUsageCapPolicyManaged)
}
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_key=$7, max_tokens=$8, max_cost_micros=$9,
enabled=$10, priority=$11, updated_at=now()
WHERE tenant_id=$1 AND id=$2
RETURNING updated_at`
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 {
res, err := s.db.ExecContext(ctx,
`DELETE FROM usage_cap_policies
WHERE tenant_id=$1 AND id=$2 AND COALESCE(source,'manual') <> $3`,
tenantID, id, store.UsageCapSourceAgentBudget)
if err != nil {
return err
}
if n, _ := res.RowsAffected(); n == 0 {
if _, err := s.getPolicy(ctx, tenantID, id); err == nil {
return fmt.Errorf("%w: agent monthly budget", store.ErrUsageCapPolicyManaged)
}
return sql.ErrNoRows
}
return nil
}
func (s *PGUsageCapStore) ReserveUsage(ctx context.Context, req store.UsageReserveRequest, policies []store.UsageCapPolicy) (*store.UsageReservationResult, error) {
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 + $3,
reserved_cost_micros = reserved_cost_micros + $4,
updated_at = now()
WHERE policy_id=$1 AND window_start=$2
AND ($5::bigint IS NULL OR used_tokens + reserved_tokens + $3 <= $5)
AND ($6::bigint IS NULL OR used_cost_micros + reserved_cost_micros + $4 <= $6)
RETURNING true`, p.ID, start, req.EstimatedTokens, req.EstimatedCostMicros, intPtrVal(p.MaxTokens), intPtrVal(p.MaxCostMicros)).Scan(&ok)
if err != nil {
if !errors.Is(err, sql.ErrNoRows) {
return nil, err
}
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}
err := s.db.QueryRowContext(ctx, `
SELECT used_tokens, reserved_tokens, used_cost_micros, reserved_cost_micros
FROM usage_cap_counters WHERE policy_id=$1 AND window_start=$2`, p.ID, start).Scan(&u.UsedTokens, &u.ReservedTokens, &u.UsedCostMicros, &u.ReservedCostMicros)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
out = append(out, u)
continue
}
return nil, err
}
out = append(out, u)
}
return out, nil
}
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_key, max_tokens, max_cost_micros, COALESCE(source,'manual'), enabled, priority, created_at, updated_at FROM usage_cap_policies`
func scanPolicy(row scanner) (store.UsageCapPolicy, error) {
var p store.UsageCapPolicy
var agentID, providerID uuid.NullUUID
var maxTokens, maxCost sql.NullInt64
err := row.Scan(&p.ID, &p.TenantID, &agentID, &providerID, &p.ProviderType, &p.ModelID,
&p.Window, &maxTokens, &maxCost, &p.Source, &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
}
internal/store/pg/usage_caps_test.go
+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, "ucp-"+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, "ucpm-"+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 }
internal/store/pg/usage_pricing.go
+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
}
internal/store/stores.go
+39 -36
View File
@@ -4,42 +4,42 @@ 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
BrowserCookies BrowserCookieStore
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
BrowserCookies BrowserCookieStore
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.
@@ -53,4 +53,7 @@ type Stores struct {
WorkstationLinks AgentWorkstationLinkStore
WorkstationPermissions WorkstationPermissionStore
WorkstationActivity WorkstationActivityStore
// UsageCaps is Standard/PostgreSQL only in the first budget-control rollout.
UsageCaps UsageCapStore
}
internal/store/usage_caps.go
+202
View File
@@ -0,0 +1,202 @@
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"
UsageCapSourceManual = "manual"
UsageCapSourceAgentBudget = "agent_budget_monthly_cents"
)
var (
ErrUsageCapExceeded = errors.New("usage cap exceeded")
ErrUsageCapPolicyManaged = errors.New("usage cap policy is managed by another setting")
)
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"`
Source string `json:"source,omitempty" db:"source"`
Enabled bool `json:"enabled" db:"enabled"`
Priority int `json:"priority" db:"priority"`
CreatedAt time.Time `json:"created_at" db:"created_at"`
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
}
internal/tools/read_audio.go
+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 {
internal/tools/read_audio_resolve.go
+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)
}
internal/tools/read_document.go
+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 {
internal/tools/read_document_resolve.go
+26 -2
View File
@@ -151,7 +151,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)
}
@@ -178,7 +194,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",
@@ -188,7 +204,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)
}
internal/tools/read_image.go
+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)
}
internal/tools/read_video.go
+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 {
internal/tools/read_video_resolve.go
+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)
}
internal/tools/subagent.go
+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 {
internal/tools/subagent_exec.go
+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,
})
}
internal/tools/subagent_spawn.go
+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),
internal/tools/usage_caps.go
+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
}
internal/upgrade/version.go
+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 = 69
const RequiredSchemaVersion uint = 72
internal/usage/caps/chat_call.go
+150
View File
@@ -0,0 +1,150 @@
package caps
import (
"context"
"errors"
"fmt"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// ChatOptions identifies a billable non-agent LLM call for usage-cap enforcement.
type ChatOptions struct {
TenantID uuid.UUID
AgentID uuid.UUID
ProviderName string
ModelID string
ReservationKey string
Purpose string
MaxOutputTokens int
}
// Chat wraps Provider.Chat with the same usage-cap preflight and reconciliation
// used by agent loops. A nil service intentionally falls back to direct calls
// for Lite/subscription-only runtimes.
func (s *Service) Chat(ctx context.Context, provider providers.Provider, req providers.ChatRequest, opts ChatOptions) (*providers.ChatResponse, error) {
if provider == nil {
return nil, errors.New("usage cap chat: provider is nil")
}
if s == nil || s.store == nil {
return provider.Chat(ctx, req)
}
if fallback, ok := provider.(*providers.ModelFallbackProvider); ok {
return fallback.ChatWithHook(ctx, req, func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) {
callOpts := opts
callOpts.ProviderName = entry.ProviderName
if callOpts.ProviderName == "" && entry.Provider != nil {
callOpts.ProviderName = entry.Provider.Name()
}
callOpts.ModelID = actualReq.Model
callOpts.ReservationKey = ""
usageReq := s.chatRequest(callCtx, entry.Provider, actualReq, callOpts)
scopedCtx := scopedRequestContext(callCtx, usageReq)
reservation, err := s.Preflight(scopedCtx, usageReq)
if err != nil {
return nil, err
}
return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) {
if reservation != nil {
reservation.Reconcile(scopedCtx, resp, callErr)
}
}, nil
})
}
usageReq := s.chatRequest(ctx, provider, req, opts)
scopedCtx := scopedRequestContext(ctx, usageReq)
reservation, err := s.Preflight(scopedCtx, usageReq)
if err != nil {
return nil, err
}
resp, err := provider.Chat(scopedCtx, req)
if reservation != nil {
reservation.Reconcile(scopedCtx, resp, err)
}
return resp, err
}
func (s *Service) chatRequest(ctx context.Context, provider providers.Provider, req providers.ChatRequest, opts ChatOptions) Request {
tenantID := opts.TenantID
if tenantID == uuid.Nil {
tenantID = store.TenantIDFromContext(ctx)
}
if tenantID == uuid.Nil {
tenantID = store.MasterTenantID
}
agentID := opts.AgentID
if agentID == uuid.Nil {
agentID = store.AgentIDFromContext(ctx)
}
providerName := opts.ProviderName
if providerName == "" && provider != nil {
providerName = provider.Name()
}
modelID := opts.ModelID
if modelID == "" {
modelID = req.Model
}
if modelID == "" && provider != nil {
modelID = provider.DefaultModel()
}
return Request{
TenantID: tenantID,
AgentID: agentID,
ProviderName: providerName,
ModelID: modelID,
ReservationKey: reservationKey(opts),
Messages: req.Messages,
MaxOutputTokens: maxOutputTokens(req, opts.MaxOutputTokens),
}
}
func scopedRequestContext(ctx context.Context, req Request) context.Context {
if req.TenantID != uuid.Nil && store.TenantIDFromContext(ctx) != req.TenantID {
ctx = store.WithTenantID(ctx, req.TenantID)
}
if req.AgentID != uuid.Nil && store.AgentIDFromContext(ctx) != req.AgentID {
ctx = store.WithAgentID(ctx, req.AgentID)
}
return ctx
}
func reservationKey(opts ChatOptions) string {
if opts.ReservationKey != "" {
return opts.ReservationKey
}
purpose := opts.Purpose
if purpose == "" {
purpose = "llm"
}
return fmt.Sprintf("%s:%s", purpose, uuid.NewString())
}
func maxOutputTokens(req providers.ChatRequest, fallback int) int {
if fallback <= 0 {
fallback = 1024
}
if req.Options == nil {
return fallback
}
v, ok := req.Options[providers.OptMaxTokens]
if !ok {
return fallback
}
switch n := v.(type) {
case int:
return n
case int64:
return int(n)
case int32:
return int(n)
case float64:
return int(n)
case float32:
return int(n)
default:
return fallback
}
}
internal/usage/caps/reservation_metadata.go
+125
View File
@@ -0,0 +1,125 @@
package caps
import (
"encoding/json"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/usage/pricing"
)
const TraceMetadataKey = "usage_caps"
type TraceMetadata struct {
Decision string `json:"decision"`
Reason string `json:"reason,omitempty"`
ReservationKey string `json:"reservation_key,omitempty"`
ProviderName string `json:"provider_name,omitempty"`
ProviderType string `json:"provider_type,omitempty"`
ModelID string `json:"model_id,omitempty"`
EstimatedTokens int64 `json:"estimated_tokens"`
EstimatedCostMicros int64 `json:"estimated_cost_micros"`
ActualTokens int64 `json:"actual_tokens"`
ActualCostMicros int64 `json:"actual_cost_micros"`
ReconcileStatus string `json:"reconcile_status,omitempty"`
PolicyCount int `json:"policy_count"`
PolicyIDs []string `json:"policy_ids,omitempty"`
}
func (r *Reservation) TraceMetadata() TraceMetadata {
if r == nil {
return TraceMetadata{}
}
decision := r.decision
if decision == "" {
if r.skipped {
decision = store.UsageCapEventSkip
} else {
decision = store.UsageCapEventAllow
}
}
policyCount := 0
policyIDs := make([]string, 0)
if r.result != nil {
policyCount = len(r.result.Policies)
for _, policy := range r.result.Policies {
if policy.ID != uuid.Nil {
policyIDs = append(policyIDs, policy.ID.String())
}
}
}
if r.blockedPolicyID != uuid.Nil {
policyIDs = append(policyIDs, r.blockedPolicyID.String())
if policyCount == 0 {
policyCount = 1
}
}
return TraceMetadata{
Decision: decision,
Reason: r.reason,
ReservationKey: r.key,
ProviderName: r.providerName,
ProviderType: r.providerType,
ModelID: r.modelID,
EstimatedTokens: r.usage.TotalTokens(),
EstimatedCostMicros: r.estimatedCostMicros,
ActualTokens: r.actualTokens,
ActualCostMicros: r.actualCostMicros,
ReconcileStatus: r.reconcileStatus,
PolicyCount: policyCount,
PolicyIDs: policyIDs,
}
}
func (m TraceMetadata) Empty() bool {
return m.Decision == "" && m.Reason == "" && m.ReservationKey == "" &&
m.ProviderName == "" && m.ProviderType == "" && m.ModelID == "" &&
m.EstimatedTokens == 0 && m.EstimatedCostMicros == 0 &&
m.ActualTokens == 0 && m.ActualCostMicros == 0 &&
m.ReconcileStatus == "" && m.PolicyCount == 0 && len(m.PolicyIDs) == 0
}
func MergeTraceMetadata(existing json.RawMessage, entries []TraceMetadata) json.RawMessage {
clean := make([]TraceMetadata, 0, len(entries))
for _, entry := range entries {
if !entry.Empty() {
clean = append(clean, entry)
}
}
if len(clean) == 0 {
return existing
}
payload := map[string]any{}
if len(existing) > 0 {
_ = json.Unmarshal(existing, &payload)
}
payload[TraceMetadataKey] = map[string]any{"attempts": clean}
data, err := json.Marshal(payload)
if err != nil {
return existing
}
return json.RawMessage(data)
}
func skippedReservation(req Request, reason string) *Reservation {
return &Reservation{
skipped: true, decision: store.UsageCapEventSkip, reason: reason,
providerName: req.ProviderName, modelID: req.ModelID,
}
}
func skippedScopedReservation(req Request, scope store.UsageCapScope, reason string) *Reservation {
r := skippedReservation(req, reason)
r.providerType = scope.ProviderType
r.modelID = scope.ModelID
return r
}
func blockedReservation(req Request, scope store.UsageCapScope, key string, usage pricing.BillableUsage, costMicros int64, policyID uuid.UUID, reason string) *Reservation {
return &Reservation{
key: key, usage: usage, estimatedCostMicros: costMicros,
decision: store.UsageCapEventBlock, reason: reason,
blockedPolicyID: policyID,
providerName: req.ProviderName, providerType: scope.ProviderType, modelID: scope.ModelID,
}
}
internal/usage/caps/service.go
+298
View File
@@ -0,0 +1,298 @@
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
actualTokens int64
actualCostMicros int64
reconcileStatus string
skipped bool
decision string
reason string
blockedPolicyID uuid.UUID
providerName string
providerType string
modelID string
}
func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, error) {
if s == nil || s.store == nil {
return skippedReservation(req, "service_disabled"), nil
}
ctx = scopedRequestContext(ctx, req)
providerData, err := s.resolveProvider(ctx, req.TenantID, req.ProviderName)
if err != nil {
return skippedReservation(req, "provider_metadata_missing"), nil
}
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 skippedScopedReservation(req, scope, "provider_not_billable_api"), nil
}
policies, err := s.store.ListUsageCapPolicies(ctx, scope, false)
if err != nil {
return nil, err
}
if len(policies) == 0 {
return skippedScopedReservation(req, scope, "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
}
key := req.ReservationKey
if key == "" {
key = uuid.NewString()
}
metadata := map[string]any{"model_id": req.ModelID, "pricing_source": "token_only"}
var prices store.UsagePricingFields
var costMicros int64
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) {
_ = s.store.InsertUsageCapEvent(ctx, &store.UsageCapEvent{
TenantID: req.TenantID,
ReservationKey: key, Decision: store.UsageCapEventBlock, Reason: "pricing_unknown",
EstimatedTokens: usage.TotalTokens(), EstimatedCostMicros: 0,
Metadata: mustJSON(map[string]any{"model_id": req.ModelID, "provider": req.ProviderName}),
})
return blockedReservation(req, scope, key, usage, 0, uuid.Nil, "pricing_unknown"), fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID)
}
return nil, err
}
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}
}
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 blockedReservation(req, scope, key, usage, costMicros, blockedPolicy, reason), ErrCapExceeded
}
return nil, err
}
return &Reservation{
key: key, result: result, svc: s, usage: usage, prices: prices,
estimatedCostMicros: costMicros, decision: store.UsageCapEventAllow,
reason: "reserved", providerName: req.ProviderName,
providerType: scope.ProviderType, modelID: scope.ModelID,
}, nil
}
func (r *Reservation) Reconcile(ctx context.Context, resp *providers.ChatResponse, callErr error) {
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
}
r.actualTokens = actual.TotalTokens()
r.actualCostMicros = cost
r.reconcileStatus = status
reconcileCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
defer cancel()
if err := r.svc.store.ReconcileUsage(reconcileCtx, store.UsageReconcileRequest{
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
}
internal/usage/caps/service_test.go
+527
View File
@@ -0,0 +1,527 @@
package caps
import (
"context"
"database/sql"
"encoding/json"
"errors"
"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)
}
metadata := reservation.TraceMetadata()
if metadata.Decision != store.UsageCapEventAllow {
t.Fatalf("Decision = %q, want allow", metadata.Decision)
}
if metadata.PolicyCount != 1 {
t.Fatalf("PolicyCount = %d, want 1", metadata.PolicyCount)
}
}
func TestPreflightIncludesRequestPricingWhenConfigured(t *testing.T) {
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",
}, requireTenant: policy.TenantID}
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)
}
metadata := reservation.TraceMetadata()
if metadata.ActualTokens != 5 {
t.Fatalf("ActualTokens = %d, want 5", metadata.ActualTokens)
}
if metadata.ReconcileStatus != "reconciled" {
t.Fatalf("ReconcileStatus = %q, want reconciled", metadata.ReconcileStatus)
}
}
func TestPreflightTraceMetadataForCapExceeded(t *testing.T) {
policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(10), Enabled: true}
usageStore := &fakeUsageCapStore{
policies: []store.UsageCapPolicy{policy},
reserveErr: &store.UsageCapExceededError{PolicyID: policy.ID, Reason: "token_cap_exceeded"},
}
providerStore := &fakeProviderStore{provider: &store.LLMProviderData{
BaseModel: store.BaseModel{ID: uuid.New()},
Name: "openrouter",
ProviderType: store.ProviderOpenRouter,
APIKey: "sk-test",
}}
svc := NewService(usageStore, providerStore)
reservation, err := svc.Preflight(context.Background(), Request{
TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "token/model",
ReservationKey: "blocked", Messages: []providers.Message{{Role: "user", Content: "hello"}},
MaxOutputTokens: 10,
})
if !errors.Is(err, ErrCapExceeded) {
t.Fatalf("Preflight error = %v, want ErrCapExceeded", err)
}
metadata := reservation.TraceMetadata()
if metadata.Decision != store.UsageCapEventBlock {
t.Fatalf("Decision = %q, want block", metadata.Decision)
}
if metadata.Reason != "token_cap_exceeded" {
t.Fatalf("Reason = %q, want token_cap_exceeded", metadata.Reason)
}
if metadata.ReservationKey != "blocked" {
t.Fatalf("ReservationKey = %q, want blocked", metadata.ReservationKey)
}
if len(metadata.PolicyIDs) != 1 || metadata.PolicyIDs[0] != policy.ID.String() {
t.Fatalf("PolicyIDs = %v, want [%s]", metadata.PolicyIDs, policy.ID)
}
}
func TestPreflightRecordsPricingUnknownBlockEvent(t *testing.T) {
policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxCostMicros: int64Ptr(1000), Enabled: true}
usageStore := &fakeUsageCapStore{
policies: []store.UsageCapPolicy{policy},
resolveErr: sql.ErrNoRows,
}
providerStore := &fakeProviderStore{provider: &store.LLMProviderData{
BaseModel: store.BaseModel{ID: uuid.New()},
Name: "openrouter",
ProviderType: store.ProviderOpenRouter,
APIKey: "sk-test",
}}
svc := NewService(usageStore, providerStore)
reservation, err := svc.Preflight(context.Background(), Request{
TenantID: policy.TenantID, ProviderName: "openrouter", ModelID: "missing/model",
ReservationKey: "pricing-missing", Messages: []providers.Message{{Role: "user", Content: "hello"}},
MaxOutputTokens: 10,
})
if !errors.Is(err, ErrPricingUnknown) {
t.Fatalf("Preflight error = %v, want ErrPricingUnknown", err)
}
if reservation == nil {
t.Fatal("Preflight returned nil reservation")
}
metadata := reservation.TraceMetadata()
if metadata.Reason != "pricing_unknown" {
t.Fatalf("reservation metadata = %+v, want pricing_unknown", metadata)
}
if len(usageStore.events) != 1 {
t.Fatalf("events = %d, want 1", len(usageStore.events))
}
event := usageStore.events[0]
if event.Decision != store.UsageCapEventBlock || event.Reason != "pricing_unknown" {
t.Fatalf("event decision/reason = %q/%q, want block/pricing_unknown", event.Decision, event.Reason)
}
if event.ReservationKey != "pricing-missing" {
t.Fatalf("event reservation_key = %q, want pricing-missing", event.ReservationKey)
}
}
func TestServiceChatBlocksBeforeProviderCall(t *testing.T) {
policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: uuid.New(), MaxTokens: int64Ptr(10), Enabled: true}
usageStore := &fakeUsageCapStore{
policies: []store.UsageCapPolicy{policy},
reserveErr: &store.UsageCapExceededError{PolicyID: policy.ID, Reason: "token_cap_exceeded"},
}
providerStore := &fakeProviderStore{provider: &store.LLMProviderData{
BaseModel: store.BaseModel{ID: uuid.New()},
Name: "openrouter",
ProviderType: store.ProviderOpenRouter,
APIKey: "sk-test",
}, requireTenant: policy.TenantID}
svc := NewService(usageStore, providerStore)
provider := &fakeChatProvider{name: "openrouter", model: "token/model"}
_, err := svc.Chat(context.Background(), provider, providers.ChatRequest{
Messages: []providers.Message{{Role: "user", Content: "hello"}},
Model: "token/model",
Options: map[string]any{providers.OptMaxTokens: 20},
}, ChatOptions{
TenantID: policy.TenantID,
ProviderName: "openrouter",
Purpose: "test-block",
})
if !errors.Is(err, ErrCapExceeded) {
t.Fatalf("Chat error = %v, want ErrCapExceeded", err)
}
if provider.calls != 0 {
t.Fatalf("provider calls = %d, want 0", provider.calls)
}
}
func TestMergeTraceMetadataPreservesExistingSections(t *testing.T) {
existing := json.RawMessage(`{"thinking":{"effort":"high"}}`)
merged := MergeTraceMetadata(existing, []TraceMetadata{{
Decision: store.UsageCapEventAllow,
Reason: "reserved",
ModelID: "openai/gpt-test",
}})
var payload map[string]json.RawMessage
if err := json.Unmarshal(merged, &payload); err != nil {
t.Fatalf("Unmarshal merged metadata: %v", err)
}
if len(payload["thinking"]) == 0 {
t.Fatal("existing thinking metadata was removed")
}
var usagePayload struct {
Attempts []TraceMetadata `json:"attempts"`
}
if err := json.Unmarshal(payload[TraceMetadataKey], &usagePayload); err != nil {
t.Fatalf("Unmarshal usage caps metadata: %v", err)
}
if len(usagePayload.Attempts) != 1 || usagePayload.Attempts[0].Decision != store.UsageCapEventAllow {
t.Fatalf("usage cap attempts = %+v, want one allow attempt", usagePayload.Attempts)
}
}
func TestCountImagesOnlyCountsImageMIMEs(t *testing.T) {
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 fakeChatProvider struct {
name string
model string
calls int
resp *providers.ChatResponse
err error
}
func (p *fakeChatProvider) Chat(context.Context, providers.ChatRequest) (*providers.ChatResponse, error) {
p.calls++
if p.resp != nil || p.err != nil {
return p.resp, p.err
}
return &providers.ChatResponse{
Content: "ok",
Usage: &providers.Usage{PromptTokens: 1, CompletionTokens: 1, TotalTokens: 2},
}, nil
}
func (p *fakeChatProvider) ChatStream(ctx context.Context, req providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) {
return p.Chat(ctx, req)
}
func (p *fakeChatProvider) DefaultModel() string { return p.model }
func (p *fakeChatProvider) Name() string { return p.name }
type fakeUsageCapStore struct {
policies []store.UsageCapPolicy
resolved *store.ResolvedUsagePricing
resolveErr error
resolveCalls int
reserveErr error
reserved store.UsageReserveRequest
reconciled store.UsageReconcileRequest
reconcileCalls int
reconcileCtxCanceled bool
events []store.UsageCapEvent
}
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
if s.reserveErr != nil {
return nil, s.reserveErr
}
return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Policies: policies}, nil
}
func (s *fakeUsageCapStore) ReconcileUsage(ctx context.Context, req store.UsageReconcileRequest) error {
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, event *store.UsageCapEvent) error {
if event != nil {
s.events = append(s.events, *event)
}
return nil
}
type fakeProviderStore struct {
provider *store.LLMProviderData
masterProvider *store.LLMProviderData
requireTenant uuid.UUID
}
func (s *fakeProviderStore) CreateProvider(context.Context, *store.LLMProviderData) error { return nil }
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 s.requireTenant != uuid.Nil && store.TenantIDFromContext(ctx) != s.requireTenant {
return nil, sql.ErrNoRows
}
if store.TenantIDFromContext(ctx) == store.MasterTenantID && s.masterProvider != nil {
return s.masterProvider, nil
}
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 }
internal/usage/pricing/decimal.go
+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
}
internal/usage/pricing/decimal_test.go
+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())
}
}
internal/usage/pricing/openrouter.go
+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
}
}
internal/vault/enrich_classify.go
+9
View File
@@ -7,6 +7,7 @@ import (
"maps"
"slices"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
@@ -40,6 +41,14 @@ func (w *EnrichWorker) classifyLinks(ctx context.Context, provider providers.Pro
if provider == nil {
return
}
if tid, err := uuid.Parse(tenantID); err == nil {
ctx = store.WithTenantID(ctx, tid)
}
if agentID != "" {
if aid, err := uuid.Parse(agentID); err == nil {
ctx = store.WithAgentID(ctx, aid)
}
}
capped := results
if len(capped) > classifyMaxSourceDocs {
internal/vault/enrich_worker.go
+35 -20
View File
@@ -16,20 +16,21 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/bgalert"
"github.com/nextlevelbuilder/goclaw/internal/bus"
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"golang.org/x/sync/semaphore"
)
const (
enrichMaxDedupEntries = 10000
enrichSimilarityLimit = 10
enrichSimilarityMin = 0.7
enrichMaxConcurrent = 3 // max concurrent batch summarize calls across chunks
enrichBatchSize = 5 // docs per enrichment chunk (1 LLM call per chunk)
enrichBatchItemMaxRunes = 3000 // per-file content limit in batch summarize
enrichMaxRetries = 3 // shared retry count for LLM calls (summarize + classify)
enrichMaxDedupEntries = 10000
enrichSimilarityLimit = 10
enrichSimilarityMin = 0.7
enrichMaxConcurrent = 3 // max concurrent batch summarize calls across chunks
enrichBatchSize = 5 // docs per enrichment chunk (1 LLM call per chunk)
enrichBatchItemMaxRunes = 3000 // per-file content limit in batch summarize
enrichMaxRetries = 3 // shared retry count for LLM calls (summarize + classify)
)
// Shared retry config for all enrichment LLM calls.
@@ -41,12 +42,13 @@ var (
// EnrichWorkerDeps bundles dependencies for the vault enrichment worker.
type EnrichWorkerDeps struct {
VaultStore store.VaultStore
SystemConfigs store.SystemConfigStore // per-tenant provider config
Registry *providers.Registry // provider resolution
SystemConfigs store.SystemConfigStore // per-tenant provider config
Registry *providers.Registry // provider resolution
EventBus eventbus.DomainEventBus
MsgBus bus.EventPublisher // for WS event broadcast
TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe)
AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors
MsgBus bus.EventPublisher // for WS event broadcast
TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe)
AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors
UsageCaps *usagecaps.Service
}
// RegisterEnrichWorker subscribes the enrichment worker to vault doc events.
@@ -60,6 +62,7 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *Enri
registry: deps.Registry,
msgBus: deps.MsgBus,
alertDeps: deps.AlertDeps,
usageCaps: deps.UsageCaps,
dedup: make(map[string]string),
sem: semaphore.NewWeighted(enrichMaxConcurrent),
progress: progress,
@@ -74,11 +77,12 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *Enri
// Exported so HTTP handlers can call Stop/EnqueueUnenriched.
type EnrichWorker struct {
vault store.VaultStore
teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
msgBus bus.EventPublisher // for error event broadcast
alertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors
teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
msgBus bus.EventPublisher // for error event broadcast
alertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors
usageCaps *usagecaps.Service
queue enrichBatchQueue
progress *EnrichProgress
@@ -296,6 +300,14 @@ func (w *EnrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD
// Batch-fetch all existing docs in a single query.
tenantID := pending[0].TenantID
if tid, err := uuid.Parse(tenantID); err == nil {
ctx = store.WithTenantID(ctx, tid)
}
if pending[0].AgentID != "" {
if aid, err := uuid.Parse(pending[0].AgentID); err == nil {
ctx = store.WithAgentID(ctx, aid)
}
}
// Resolve provider once per chunk (all items share tenantID)
provider, model := w.resolveProviderForTenant(ctx, tenantID)
@@ -515,7 +527,11 @@ func (w *EnrichWorker) chatWithRetry(ctx context.Context, provider providers.Pro
}
}
cctx, cancel := context.WithTimeout(ctx, enrichRetryTimeouts[attempt])
resp, err := provider.Chat(cctx, req)
resp, err := w.usageCaps.Chat(cctx, provider, req, usagecaps.ChatOptions{
ModelID: req.Model,
Purpose: logPrefix,
MaxOutputTokens: 4096,
})
cancel()
if err != nil {
lastErr = err
@@ -563,7 +579,6 @@ func (w *EnrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUps
}
}
// recordDedup stores a processed hash and evicts ~25% entries if over capacity.
func (w *EnrichWorker) recordDedup(docID, hash string) {
w.dedupMu.Lock()
migrations/000070_usage_caps_pricing.down.sql
+2
View File
@@ -0,0 +1,2 @@
DROP TABLE IF EXISTS usage_pricing_overrides;
DROP TABLE IF EXISTS usage_pricing_catalog;
migrations/000070_usage_caps_pricing.up.sql
+45
View File
@@ -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;
migrations/000071_usage_cap_policies.down.sql
+4
View File
@@ -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;
migrations/000071_usage_cap_policies.up.sql
+68
View File
@@ -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_key TEXT NOT NULL CHECK (window_key IN ('hour', 'day', 'week', 'month')),
max_tokens BIGINT CHECK (max_tokens IS NULL OR max_tokens >= 0),
max_cost_micros BIGINT CHECK (max_cost_micros IS NULL OR max_cost_micros >= 0),
enabled BOOLEAN NOT NULL DEFAULT true,
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);
migrations/000072_agent_budget_usage_cap_bridge.down.sql
+7
View File
@@ -0,0 +1,7 @@
DELETE FROM usage_cap_policies
WHERE source = 'agent_budget_monthly_cents';
DROP INDEX IF EXISTS idx_usage_cap_policies_agent_budget_source;
ALTER TABLE usage_cap_policies
DROP COLUMN IF EXISTS source;
migrations/000072_agent_budget_usage_cap_bridge.up.sql
+23
View File
@@ -0,0 +1,23 @@
ALTER TABLE usage_cap_policies
ADD COLUMN IF NOT EXISTS source TEXT NOT NULL DEFAULT 'manual';
CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_cap_policies_agent_budget_source
ON usage_cap_policies (tenant_id, agent_id)
WHERE source = 'agent_budget_monthly_cents';
INSERT INTO usage_cap_policies (
tenant_id, agent_id, window_key, max_cost_micros, enabled, priority, source
)
SELECT
tenant_id,
id,
'month',
budget_monthly_cents::BIGINT * 10000,
true,
90,
'agent_budget_monthly_cents'
FROM agents
WHERE deleted_at IS NULL
AND budget_monthly_cents IS NOT NULL
AND budget_monthly_cents > 0
ON CONFLICT DO NOTHING;
ui/web/src/i18n/locales/en/providers.json
+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",
ui/web/src/i18n/locales/en/usage.json
+46 -1
View File
@@ -19,7 +19,52 @@
"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",
"edit": "Edit cap",
"save": "Save cap",
"cancel": "Cancel edit",
"scope": "Scope",
"tokens": "Tokens",
"cost": "Cost",
"empty": "No cap policies yet.",
"tenantScope": "Tenant cap",
"tenantScoped": "Tenant scoped",
"agentScoped": "Agent scoped",
"agentBudgetSource": "Agent budget",
"agentBudgetManaged": "Clear this from the agent monthly budget field.",
"enabled": "Enabled",
"disabled": "Disabled",
"delete": "Delete cap",
"recentBlocks": "Recent blocks",
"blocked": "Blocked",
"toast": {
"created": "Usage cap created",
"updated": "Usage cap updated",
"createFailed": "Could not create usage cap",
"updateFailed": "Could not update usage cap",
"deleted": "Usage cap deleted",
"deleteFailed": "Could not delete usage cap"
}
},
"analytics": {
"title": "Usage Analytics",
ui/web/src/i18n/locales/vi/providers.json
+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",
ui/web/src/i18n/locales/vi/usage.json
+46 -1
View File
@@ -19,7 +19,52 @@
"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",
"edit": "Sửa cap",
"save": "Lưu cap",
"cancel": "Hủy sửa",
"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",
"agentBudgetSource": "Ngân sách agent",
"agentBudgetManaged": "Xóa giới hạn này từ ô ngân sách hàng tháng của agent.",
"enabled": "Đã bật",
"disabled": "Đã tắt",
"delete": "Xóa cap",
"recentBlocks": "Lần chặn gần đây",
"blocked": "Đã chặn",
"toast": {
"created": "Đã tạo usage cap",
"updated": "Đã cập nhật usage cap",
"createFailed": "Không thể tạo usage cap",
"updateFailed": "Không thể cập nhật usage cap",
"deleted": "Đã xóa usage cap",
"deleteFailed": "Không thể xóa usage cap"
}
},
"analytics": {
"title": "Phân tích sử dụng",
ui/web/src/i18n/locales/zh/providers.json
+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",
ui/web/src/i18n/locales/zh/usage.json
+46 -1
View File
@@ -19,7 +19,52 @@
"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": "创建上限",
"edit": "编辑上限",
"save": "保存上限",
"cancel": "取消编辑",
"scope": "范围",
"tokens": "令牌",
"cost": "费用",
"empty": "暂无上限策略。",
"tenantScope": "租户上限",
"tenantScoped": "租户范围",
"agentScoped": "Agent范围",
"agentBudgetSource": "Agent预算",
"agentBudgetManaged": "请在 Agent 月度预算字段中清除此上限。",
"enabled": "已启用",
"disabled": "已禁用",
"delete": "删除上限",
"recentBlocks": "近期拦截",
"blocked": "已拦截",
"toast": {
"created": "用量上限已创建",
"updated": "用量上限已更新",
"createFailed": "无法创建用量上限",
"updateFailed": "无法更新用量上限",
"deleted": "用量上限已删除",
"deleteFailed": "无法删除用量上限"
}
},
"analytics": {
"title": "用量分析",
ui/web/src/lib/query-keys.ts
+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,
Loaded 100 of 108 files, more files were not shown because too many files have changed in this diff. Show more