tiennm99/goclaw
1
0 Fork 0
mirror of https://github.com/tiennm99/goclaw.git synced 2026-10-20 02:17:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): add AI budget usage caps

Browse Source
This commit is contained in:
Goon committed 2026-05-23 00:16:03 +07:00
1 parent 43049d3b3f
commit 9ff5a629d2
67 files changed
+4195 -249
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cmd/gateway.go
+10 -2
View File
@@ -46,6 +46,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/skills"
"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/internal/vault"
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
@@ -282,8 +283,15 @@ func runGateway() {
slog.Info("bootstrap: capabilities backfill complete", "agents", count)
}
usageCapSvc := usagecaps.NewService(pgStores.UsageCaps, pgStores.Providers)
if readImage, ok := toolsReg.Get("read_image"); ok {
if t, ok := readImage.(*tools.ReadImageTool); ok {
t.SetUsageCapService(usageCapSvc)
}
}
// Subagent system (secureCLI store wired so subagent ExecTools enforce the gate)
subagentMgr := setupSubagents(providerRegistry, cfg, msgBus, toolsReg, workspace, sandboxMgr, pgStores.SecureCLI)
subagentMgr := setupSubagents(providerRegistry, cfg, msgBus, toolsReg, workspace, sandboxMgr, pgStores.SecureCLI, usageCapSvc)
if subagentMgr != nil {
// Wire announce queue for batched subagent result delivery (matching TS debounce pattern).
announceQueue := tools.NewAnnounceQueue(1000, 20, makeDelegateAnnounceCallback(subagentMgr, msgBus))
@@ -336,7 +344,7 @@ func runGateway() {
var mcpPool *mcpbridge.Pool
var mediaStore *media.Store
var postTurn tools.PostTurnProcessor
contextFileInterceptor, mcpPool, mediaStore, postTurn = wireExtras(pgStores, agentRouter, providerRegistry, modelReg, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, redisClient, domainBus)
contextFileInterceptor, mcpPool, mediaStore, postTurn = wireExtras(pgStores, agentRouter, providerRegistry, modelReg, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, redisClient, domainBus, usageCapSvc)
if mcpPool != nil {
defer mcpPool.Stop()
}
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_http_wiring.go
+3
View File
@@ -134,6 +134,9 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer(
if d.pgStores.Snapshots != nil {
d.server.SetUsageHandler(httpapi.NewUsageHandler(d.pgStores.Snapshots, d.pgStores.DB))
}
if d.pgStores.UsageCaps != nil {
d.server.SetUsageCapsHandler(httpapi.NewUsageCapsHandler(d.pgStores.UsageCaps, d.pgStores.Tenants))
}
// Runtime package management (install/uninstall system/pip/npm/github packages)
// Wire the update registry AFTER initGitHubInstaller so DefaultGitHubInstaller() is set.
cmd/gateway_managed.go
+12 -3
View File
@@ -31,6 +31,7 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/store/pg"
"github.com/nextlevelbuilder/goclaw/internal/tools"
"github.com/nextlevelbuilder/goclaw/internal/tracing"
usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps"
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
)
@@ -58,6 +59,7 @@ func wireExtras(
sandboxMgr sandbox.Manager,
redisClient any, // nil when built without -tags redis or when Redis is unconfigured
domainBus eventbus.DomainEventBus,
usageCapSvc *usagecaps.Service,
) (*tools.ContextFileInterceptor, *mcpbridge.Pool, *media.Store, tools.PostTurnProcessor) {
// 1. Build cache instances (in-memory or Redis depending on build tags)
agentCtxCache, userCtxCache := makeCaches(redisClient)
@@ -82,9 +84,15 @@ func wireExtras(
}
}
// Register media analysis tools (need mediaStore for file access).
toolsReg.Register(tools.NewReadDocumentTool(providerReg, mediaStore))
toolsReg.Register(tools.NewReadAudioTool(providerReg, mediaStore))
toolsReg.Register(tools.NewReadVideoTool(providerReg, mediaStore))
readDocumentTool := tools.NewReadDocumentTool(providerReg, mediaStore)
readDocumentTool.SetUsageCapService(usageCapSvc)
toolsReg.Register(readDocumentTool)
readAudioTool := tools.NewReadAudioTool(providerReg, mediaStore)
readAudioTool.SetUsageCapService(usageCapSvc)
toolsReg.Register(readAudioTool)
readVideoTool := tools.NewReadVideoTool(providerReg, mediaStore)
readVideoTool.SetUsageCapService(usageCapSvc)
toolsReg.Register(readVideoTool)
toolsReg.Register(tools.NewCreateVideoTool(providerReg))
slog.Info("media tools registered", "tools", "read_document,read_audio,read_video,create_video")
}
@@ -228,6 +236,7 @@ func wireExtras(
MediaStore: mediaStore,
ModelPricing: appCfg.Telemetry.ModelPricing,
TracingStore: stores.Tracing,
UsageCaps: usageCapSvc,
MemoryStore: stores.Memory,
ContactStore: stores.Contacts,
TenantStore: stores.Tenants,