tiennm99/goclaw
1
0 Fork 0
mirror of https://github.com/tiennm99/goclaw.git synced 2026-10-15 08:17: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)