diff --git a/cmd/gateway.go b/cmd/gateway.go index debcc1da..8cc54660 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -46,6 +46,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/internal/vault" "github.com/nextlevelbuilder/goclaw/pkg/protocol" @@ -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, cfg.Skills) // 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)) diff --git a/cmd/gateway_agents.go b/cmd/gateway_agents.go index 0dcbda6b..2efec76a 100644 --- a/cmd/gateway_agents.go +++ b/cmd/gateway_agents.go @@ -15,6 +15,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tts" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // resolveEmbeddingProvider selects an embedding provider from DB only. @@ -159,7 +160,7 @@ func buildEmbeddingProvider( return nil } -func setupSubagents(providerReg *providers.Registry, cfg *config.Config, msgBus *bus.MessageBus, toolsReg *tools.Registry, workspace string, sandboxMgr sandbox.Manager, secureCLIStore store.SecureCLIStore) *tools.SubagentManager { +func setupSubagents(providerReg *providers.Registry, cfg *config.Config, msgBus *bus.MessageBus, toolsReg *tools.Registry, workspace string, sandboxMgr sandbox.Manager, secureCLIStore store.SecureCLIStore, usageCapSvc *usagecaps.Service) *tools.SubagentManager { names := providerReg.List(context.Background()) if len(names) == 0 { return nil @@ -207,7 +208,9 @@ func setupSubagents(providerReg *providers.Registry, cfg *config.Config, msgBus return reg } - return tools.NewSubagentManager(provider, providerReg, agentCfg.Model, msgBus, toolsFactory, subCfg) + manager := tools.NewSubagentManager(provider, providerReg, agentCfg.Model, msgBus, toolsFactory, subCfg) + manager.SetUsageCapService(usageCapSvc) + return manager } // buildSubagentToolsRegistry produces a cloned tool registry for a subagent diff --git a/cmd/gateway_consumer.go b/cmd/gateway_consumer.go index aa6d8469..e2759ef1 100644 --- a/cmd/gateway_consumer.go +++ b/cmd/gateway_consumer.go @@ -19,6 +19,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/scheduler" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -26,7 +27,7 @@ import ( // and routes them through the scheduler/agent loop, then publishes the response back. // Also handles subagent announcements: routes them through the parent agent's session // (matching TS subagent-announce.ts pattern) so the agent can reformulate for the user. -func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents *agent.Router, cfg *config.Config, sched *scheduler.Scheduler, channelMgr *channels.Manager, teamStore store.TeamStore, quotaChecker *channels.QuotaChecker, sessStore store.SessionStore, agentStore store.AgentStore, contactCollector *store.ContactCollector, postTurn tools.PostTurnProcessor, subagentMgr *tools.SubagentManager) { +func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents *agent.Router, cfg *config.Config, sched *scheduler.Scheduler, channelMgr *channels.Manager, teamStore store.TeamStore, quotaChecker *channels.QuotaChecker, sessStore store.SessionStore, agentStore store.AgentStore, contactCollector *store.ContactCollector, postTurn tools.PostTurnProcessor, subagentMgr *tools.SubagentManager, usageCapSvc *usagecaps.Service) { slog.Info("inbound message consumer started") // Inbound message deduplication (matching TS src/infra/dedupe.ts + inbound-dedupe.ts). @@ -58,6 +59,7 @@ func consumeInboundMessages(ctx context.Context, msgBus *bus.MessageBus, agents QuotaChecker: quotaChecker, ContactCollector: contactCollector, SubagentMgr: subagentMgr, + UsageCaps: usageCapSvc, GetAnnounceMu: getAnnounceMu, } diff --git a/cmd/gateway_consumer_deps.go b/cmd/gateway_consumer_deps.go index faf3755c..11174f2a 100644 --- a/cmd/gateway_consumer_deps.go +++ b/cmd/gateway_consumer_deps.go @@ -10,6 +10,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/scheduler" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ConsumerDeps bundles shared dependencies for consumer message handlers. @@ -28,6 +29,7 @@ type ConsumerDeps struct { ContactCollector *store.ContactCollector TaskRunSessions sync.Map SubagentMgr *tools.SubagentManager + UsageCaps *usagecaps.Service BgWg sync.WaitGroup GetAnnounceMu func(string) *sync.Mutex } diff --git a/cmd/gateway_consumer_normal.go b/cmd/gateway_consumer_normal.go index 0b243cf8..a719c446 100644 --- a/cmd/gateway_consumer_normal.go +++ b/cmd/gateway_consumer_normal.go @@ -333,7 +333,11 @@ func processNormalMessage( if locale == "" { locale = "en" } - intent := agent.ClassifyIntent(ctx, loop.Provider(), loop.Model(), msg.Content) + classifyCtx := ctx + if uid := loop.UUID(); uid != uuid.Nil { + classifyCtx = store.WithAgentID(classifyCtx, uid) + } + intent := agent.ClassifyIntentWithUsageCaps(classifyCtx, deps.UsageCaps, loop.Provider(), loop.Model(), msg.Content) switch intent { case agent.IntentStatusQuery: status := deps.Agents.GetActivity(sessionKey) diff --git a/cmd/gateway_deps.go b/cmd/gateway_deps.go index 0487f358..c30ca4f2 100644 --- a/cmd/gateway_deps.go +++ b/cmd/gateway_deps.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/internal/vault" ) @@ -28,13 +29,14 @@ type gatewayDeps struct { channelMgr *channels.Manager agentRouter *agent.Router toolsReg *tools.Registry - skillsLoader *skills.Loader // optional: enables skill creation in evolution approval + skillsLoader *skills.Loader // optional: enables skill creation in evolution approval permCache *cache.PermissionCache // nil if no tenant store; closed on shutdown to stop sweep goroutines - enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered - enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue + enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered + enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue workspace string dataDir string domainBus eventbus.DomainEventBus + usageCapSvc *usagecaps.Service audioMgr *audio.Manager // nil if TTS not configured; used by TTSHandler ttsHandler *httpapi.TTSHandler // nil if TTS not configured; for hot-reload } diff --git a/cmd/gateway_hooks.go b/cmd/gateway_hooks.go index 26b35887..29340bc5 100644 --- a/cmd/gateway_hooks.go +++ b/cmd/gateway_hooks.go @@ -7,12 +7,13 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/edition" "github.com/nextlevelbuilder/goclaw/internal/hooks" - hookhandlers "github.com/nextlevelbuilder/goclaw/internal/hooks/handlers" "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" + hookhandlers "github.com/nextlevelbuilder/goclaw/internal/hooks/handlers" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/security" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/store/pg" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // sharedHookHandlers is populated by wireExtras so the gateway.go router @@ -27,7 +28,7 @@ var sharedHookHandlers map[hooks.HandlerType]hooks.Handler // Budget wiring (C1 fix): the PromptHandler receives a budget.Store bound // to pg.NewPGHookBudget so token spend is atomically deducted per tenant. // When the DB handle is unavailable, budget falls back to nil (Lite desktop). -func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, hooksCfg config.HooksConfig) map[hooks.HandlerType]hooks.Handler { +func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, hooksCfg config.HooksConfig, usageCapSvc *usagecaps.Service) map[hooks.HandlerType]hooks.Handler { encryptKey := os.Getenv("GOCLAW_ENCRYPTION_KEY") var budgetStore *budget.Store @@ -38,6 +39,7 @@ func buildHookHandlers(stores *store.Stores, providerReg *providers.Registry, ho promptHandler := &hookhandlers.PromptHandler{ Resolver: hookhandlers.NewRegistryResolver(providerReg, stores.SystemConfigs), Budget: budgetStore, + UsageCaps: usageCapSvc, DefaultModel: "haiku", } diff --git a/cmd/gateway_http_handlers.go b/cmd/gateway_http_handlers.go index 5ad49409..7d7cb532 100644 --- a/cmd/gateway_http_handlers.go +++ b/cmd/gateway_http_handlers.go @@ -2,14 +2,16 @@ package cmd import ( "github.com/nextlevelbuilder/goclaw/internal/bus" + "github.com/nextlevelbuilder/goclaw/internal/config" httpapi "github.com/nextlevelbuilder/goclaw/internal/http" "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, skillUploadConfig config.SkillsConfig) (*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 +26,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) @@ -37,6 +39,10 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir if len(dirs) > 0 { skillsH = httpapi.NewSkillsHandler(manageStore, dirs[0], dataDir, bundledSkillsDir, msgBus, stores.SkillTenantCfgs, stores.Tenants) skillsH.SetDB(stores.DB) + skillsH.SetUploadLimitConfig(skillUploadConfig) + if stores.SystemConfigs != nil { + skillsH.SetSystemConfigStore(stores.SystemConfigs) + } } } } @@ -61,6 +67,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 +97,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir if stores != nil && stores.PendingMessages != nil { pendingMessagesH = httpapi.NewPendingMessagesHandler(stores.PendingMessages, stores.Agents, providerReg) + pendingMessagesH.SetUsageCapService(usageCapSvc) } if stores != nil && stores.SecureCLI != nil { diff --git a/cmd/gateway_http_wiring.go b/cmd/gateway_http_wiring.go index 28efdb47..3479472e 100644 --- a/cmd/gateway_http_wiring.go +++ b/cmd/gateway_http_wiring.go @@ -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 diff --git a/cmd/gateway_lifecycle.go b/cmd/gateway_lifecycle.go index 3a8ef20a..631bdd4e 100644 --- a/cmd/gateway_lifecycle.go +++ b/cmd/gateway_lifecycle.go @@ -140,7 +140,7 @@ func (d *gatewayDeps) runLifecycle( d.channelMgr.SetContactCollector(contactCollector) } - go consumeInboundMessages(ctx, d.msgBus, d.agentRouter, d.cfg, deps.sched, d.channelMgr, deps.consumerTeamStore, deps.quotaChecker, d.pgStores.Sessions, d.pgStores.Agents, contactCollector, deps.postTurn, deps.subagentMgr) + go consumeInboundMessages(ctx, d.msgBus, d.agentRouter, d.cfg, deps.sched, d.channelMgr, deps.consumerTeamStore, deps.quotaChecker, d.pgStores.Sessions, d.pgStores.Agents, contactCollector, deps.postTurn, deps.subagentMgr, d.usageCapSvc) // Webhook callback worker — delivers async webhook_calls rows to receiver callback_url. // Runs in both editions: Standard (PG, concurrency=4) and Lite (SQLite, concurrency=1). diff --git a/cmd/gateway_lifecycle_shell_deny_groups.go b/cmd/gateway_lifecycle_shell_deny_groups.go index 28c8be7f..46e0e0df 100644 --- a/cmd/gateway_lifecycle_shell_deny_groups.go +++ b/cmd/gateway_lifecycle_shell_deny_groups.go @@ -30,6 +30,10 @@ func subscribeShellDenyGroupsReload(msgBus *bus.MessageBus, toolsReg *tools.Regi return } et.SetGlobalShellDenyGroups(updatedCfg.Tools.ShellDenyGroups) - slog.Info("shell deny groups reloaded via pub/sub", "groups", len(updatedCfg.Tools.ShellDenyGroups)) + et.SetCommandKeywordAllowlist(updatedCfg.Tools.CommandKeywordAllowlist) + slog.Info("shell deny groups reloaded via pub/sub", + "groups", len(updatedCfg.Tools.ShellDenyGroups), + "command_keyword_allowlist_rules", len(updatedCfg.Tools.CommandKeywordAllowlist), + ) }) } diff --git a/cmd/gateway_lifecycle_shell_deny_groups_test.go b/cmd/gateway_lifecycle_shell_deny_groups_test.go index 13368dd6..e66be5c2 100644 --- a/cmd/gateway_lifecycle_shell_deny_groups_test.go +++ b/cmd/gateway_lifecycle_shell_deny_groups_test.go @@ -35,6 +35,37 @@ func TestShellDenyGroupsConfigReload_UpdatesGlobal(t *testing.T) { if v, ok := got["package_install"]; !ok || v != true { t.Fatalf("expected pub/sub to set global package_install=true, got %v", got) } + rules := execTool.CommandKeywordAllowlistForTest() + if len(rules) != 0 { + t.Fatalf("expected empty command keyword allowlist, got %v", rules) + } +} + +func TestShellDenyGroupsConfigReload_UpdatesCommandKeywordAllowlist(t *testing.T) { + msgBus := bus.New() + defer msgBus.Unsubscribe("shell-deny-groups-config-reload") + + toolsReg := tools.NewRegistry() + execTool := tools.NewExecTool("/tmp", false) + toolsReg.Register(execTool) + + subscribeShellDenyGroupsReload(msgBus, toolsReg) + + msgBus.Broadcast(bus.Event{ + Name: bus.TopicConfigChanged, + Payload: &config.Config{ + Tools: config.ToolsConfig{ + CommandKeywordAllowlist: []config.CommandKeywordAllowlistRule{ + {ID: "github-content", Command: "gh", Args: []string{"--body"}, Keywords: []string{"secret"}}, + }, + }, + }, + }) + + rules := execTool.CommandKeywordAllowlistForTest() + if len(rules) != 1 || rules[0].ID != "github-content" { + t.Fatalf("expected command keyword allowlist to reload, got %v", rules) + } } // TestShellDenyGroupsConfigReload_IgnoresOtherEvents: subscriber must guard diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index 18411a0f..0ae38647 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -31,6 +31,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store/pg" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -58,6 +59,7 @@ func wireExtras( sandboxMgr sandbox.Manager, redisClient any, // nil when built without -tags redis or when Redis is unconfigured domainBus eventbus.DomainEventBus, + usageCapSvc *usagecaps.Service, ) (*tools.ContextFileInterceptor, *mcpbridge.Pool, *media.Store, tools.PostTurnProcessor) { // 1. Build cache instances (in-memory or Redis depending on build tags) agentCtxCache, userCtxCache := makeCaches(redisClient) @@ -82,9 +84,15 @@ func wireExtras( } } // Register media analysis tools (need mediaStore for file access). - toolsReg.Register(tools.NewReadDocumentTool(providerReg, mediaStore)) - toolsReg.Register(tools.NewReadAudioTool(providerReg, mediaStore)) - toolsReg.Register(tools.NewReadVideoTool(providerReg, mediaStore)) + readDocumentTool := tools.NewReadDocumentTool(providerReg, mediaStore) + readDocumentTool.SetUsageCapService(usageCapSvc) + toolsReg.Register(readDocumentTool) + readAudioTool := tools.NewReadAudioTool(providerReg, mediaStore) + readAudioTool.SetUsageCapService(usageCapSvc) + toolsReg.Register(readAudioTool) + readVideoTool := tools.NewReadVideoTool(providerReg, mediaStore) + readVideoTool.SetUsageCapService(usageCapSvc) + toolsReg.Register(readVideoTool) toolsReg.Register(tools.NewCreateVideoTool(providerReg)) slog.Info("media tools registered", "tools", "read_document,read_audio,read_video,create_video") } @@ -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, ""), @@ -201,6 +209,7 @@ func wireExtras( ToolPolicy: toolPE, Skills: skillsLoader, SkillAccessStore: skillAccessStore, + SkillSlashCommands: appCfg.Skills.SlashCommands, HasMemory: hasMemory, TraceCollector: traceCollector, EnsureUserProfile: ensureUserProfile, @@ -228,6 +237,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 +305,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 +713,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 +737,7 @@ func buildKGExtractFunc(kgStore store.KnowledgeGraphStore, bts store.BuiltinTool return } extractor := kg.NewExtractor(p, settings.ExtractionModel, settings.MinConfidence) + extractor.SetUsageCapService(usageCapSvc) result, err := extractor.Extract(ctx, content) if err != nil { slog.Warn("kg extract: extraction failed", "agent", agentID, "error", err) diff --git a/cmd/gateway_methods.go b/cmd/gateway_methods.go index d994217e..7d7f0664 100644 --- a/cmd/gateway_methods.go +++ b/cmd/gateway_methods.go @@ -12,14 +12,16 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/gateway/methods" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) -func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, cronStore store.CronStore, pairingStore store.PairingStore, cfg *config.Config, cfgPath, workspace, dataDir string, msgBus *bus.MessageBus, execApprovalMgr *tools.ExecApprovalManager, agentStore store.AgentStore, skillStore store.SkillStore, configSecretsStore store.ConfigSecretsStore, teamStore store.TeamStore, contextFileInterceptor *tools.ContextFileInterceptor, logTee *gateway.LogTee, heartbeatStore store.HeartbeatStore, configPermStore store.ConfigPermissionStore, sysConfigStore store.SystemConfigStore, tenantStore store.TenantStore, skillTenantCfgStore store.SkillTenantConfigStore, audioMgr *audio.Manager) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) { +func registerAllMethods(server *gateway.Server, agents *agent.Router, sessStore store.SessionStore, cronStore store.CronStore, pairingStore store.PairingStore, cfg *config.Config, cfgPath, workspace, dataDir string, msgBus *bus.MessageBus, execApprovalMgr *tools.ExecApprovalManager, agentStore store.AgentStore, skillStore store.SkillStore, configSecretsStore store.ConfigSecretsStore, teamStore store.TeamStore, contextFileInterceptor *tools.ContextFileInterceptor, logTee *gateway.LogTee, heartbeatStore store.HeartbeatStore, configPermStore store.ConfigPermissionStore, sysConfigStore store.SystemConfigStore, tenantStore store.TenantStore, skillTenantCfgStore store.SkillTenantConfigStore, audioMgr *audio.Manager, usageCapSvc *usagecaps.Service) (*methods.PairingMethods, *methods.HeartbeatMethods, *methods.ChatMethods, *methods.ConfigPermissionsMethods) { router := server.Router() // Phase 1: Core methods chatMethods := methods.NewChatMethods(agents, sessStore, cfg, server.RateLimiter(), msgBus) chatMethods.SetAudioManager(audioMgr) // Wire TTS auto-apply for WS responses + chatMethods.SetUsageCapService(usageCapSvc) chatMethods.Register(router) methods.NewAgentsMethods(agents, cfg, cfgPath, workspace, agentStore, contextFileInterceptor, msgBus).Register(router) methods.NewSessionsMethods(sessStore, msgBus, cfg).Register(router) diff --git a/cmd/gateway_setup.go b/cmd/gateway_setup.go index 491bd680..aba394db 100644 --- a/cmd/gateway_setup.go +++ b/cmd/gateway_setup.go @@ -14,11 +14,11 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bootstrap" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/edition" mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp" "github.com/nextlevelbuilder/goclaw/internal/permissions" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/sandbox" - "github.com/nextlevelbuilder/goclaw/internal/edition" "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/store/pg" @@ -214,6 +214,7 @@ func setupToolRegistry( // Apply global shell deny-group toggles before any request can arrive. // Per-agent overrides via store.WithShellDenyGroups still win per-key. et.SetGlobalShellDenyGroups(cfg.Tools.ShellDenyGroups) + et.SetCommandKeywordAllowlist(cfg.Tools.CommandKeywordAllowlist) et.DenyPaths(dataDir, ".goclaw/") // Allow skills execution: master-tenant skills-store + all tenant-scoped skills-store dirs. et.AllowPathExemptions( @@ -603,4 +604,3 @@ func setupSkillsSystem( return skillsLoader, skillSearchTool, globalSkillsDir, bundledSkillsDir, builtinSkillsDir } - diff --git a/cmd/gateway_system_config_sync_test.go b/cmd/gateway_system_config_sync_test.go index 3b9e2b2c..be831ca6 100644 --- a/cmd/gateway_system_config_sync_test.go +++ b/cmd/gateway_system_config_sync_test.go @@ -46,3 +46,17 @@ func TestSeedConfigForContextPersistsZeroInboundDebounce(t *testing.T) { t.Fatalf("gateway.inbound_debounce_ms = %q, want 0", got) } } + +func TestSeedConfigForContextDoesNotCreateSkillUploadTenantOverride(t *testing.T) { + t.Parallel() + + sc := &captureSystemConfigStore{data: map[string]string{}} + cfg := config.Default() + cfg.Skills.MaxUploadSizeMB = 64 + + seedConfigForContext(store.WithTenantID(context.Background(), store.MasterTenantID), sc, cfg, false) + + if _, ok := sc.data[config.SkillMaxUploadSizeSystemConfigKey]; ok { + t.Fatalf("%s should not be seeded; missing key lets SKILL.md frontmatter override global config", config.SkillMaxUploadSizeSystemConfigKey) + } +} diff --git a/docs/02-providers.md b/docs/02-providers.md index b9633d0f..642464e7 100644 --- a/docs/02-providers.md +++ b/docs/02-providers.md @@ -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 diff --git a/docs/03-tools-system.md b/docs/03-tools-system.md index ea7a5d2d..d1b3dd8b 100644 --- a/docs/03-tools-system.md +++ b/docs/03-tools-system.md @@ -246,6 +246,32 @@ User-facing parameter schemas for the most commonly configured tools. ``` Available presets: `gh`, `gcloud`, `aws`, `kubectl`, `terraform`. +### Credentialed CLI keyword allowlist + +`config.tools.commandKeywordAllowlist` lets operators allow specific product or security vocabulary inside selected credentialed CLI content arguments without disabling `deny_args`. + +Example: + +```json +{ + "tools": { + "commandKeywordAllowlist": [ + { + "id": "github-content", + "command": "gh", + "subcommands": ["issue create", "issue edit", "pr create", "pr comment"], + "args": ["--body", "--title"], + "argPositions": [], + "keywords": ["secret", "secrets", "token", "credential"], + "reason": "Allow security vocabulary in GitHub issue and PR prose" + } + ] + } +} +``` + +The rule above allows `gh issue create --body "secret rotation notes"` but still blocks command paths like `gh secret set TOKEN`. `argPositions` are 0-based after the matched subcommand. The scanner evaluates command arguments only; it does not read the contents of files passed through arguments such as `--body-file`. + --- ## 6. Interception Layer @@ -346,6 +372,13 @@ Custom tools are shell-based tools defined at runtime via the HTTP API — no re | `env` | no | Encrypted environment variables injected at runtime | | `enabled` | no | Toggle without deleting (default true) | +Credentialed CLI env entries support two API/UI kinds: + +- `sensitive` (default): encrypted at rest, masked in normal API responses, replace-only in UI, and flattened only at credential injection time. +- `value`: encrypted at rest but visible to authorized admins in API/UI for non-secret settings such as public URLs, domains, limits, regions, and feature flags. + +Legacy env JSON like `{"TOKEN":"..."}` is still accepted and treated as `sensitive`. + **Execution:** Template placeholders are rendered with shell-escaped argument values, then run via `sh -c`. The same deny-pattern check as the `exec` tool applies — no reverse shells, no `curl | sh`, etc. **Scope:** @@ -406,7 +439,7 @@ Current adopters: `web_search`, `web_fetch`, `tts`, `create_image`, `read_image` - Per-key: agent value takes precedence over global value - Multi-tenant invariant: each tenant's config is isolated -**Live reload:** Changes to `config.tools.shellDenyGroups` propagate via `bus.TopicConfigChanged` pub/sub. Next agent turn automatically applies new toggles. +**Live reload:** Changes to `config.tools.shellDenyGroups` and `config.tools.commandKeywordAllowlist` propagate via `bus.TopicConfigChanged` pub/sub. Next agent turn automatically applies new toggles. **Deny-group classes** (from `internal/tools/shell_deny_groups.go` — all denied by default): diff --git a/docs/04-gateway-protocol.md b/docs/04-gateway-protocol.md index 61f764bb..91af3125 100644 --- a/docs/04-gateway-protocol.md +++ b/docs/04-gateway-protocol.md @@ -442,7 +442,7 @@ All CRUD endpoints require `Authorization: Bearer ` and `X-GoClaw-User-Id | Method | Path | Description | |--------|------|-------------| | GET | `/v1/skills` | List skills | -| POST | `/v1/skills/upload` | Upload skill ZIP (max 20 MB) | +| POST | `/v1/skills/upload` | Upload skill ZIP (configurable, default 20 MB, max 500 MB) | | DELETE | `/v1/skills/{id}` | Delete a skill | **Traces** (`/v1/traces`): diff --git a/docs/06-store-data-model.md b/docs/06-store-data-model.md index eddc0118..bfbbbc8e 100644 --- a/docs/06-store-data-model.md +++ b/docs/06-store-data-model.md @@ -45,6 +45,7 @@ The `Stores` struct is the top-level container holding all PostgreSQL-backed sto | ContactStore | `PGContactStore` | Channel contacts (auto-collected), cross-channel deduplication, merge | | ActivityStore | `PGActivityStore` | Audit logs, action tracking, compliance | | SnapshotStore | `PGSnapshotStore` | Hourly usage snapshots, cost aggregation, time series queries | +| UsageCapStore | `PGUsageCapStore` | OpenRouter pricing catalog, pricing overrides, cap policies, reservations, counters, events | | SecureCLIStore | `PGSecureCLIStore` | CLI binary configs with encrypted credential injection | | APIKeyStore | `PGAPIKeyStore` | Gateway API keys, scopes, expiration, revocation | | HookStore | `PGHookStore` | Lifecycle hook definitions (event, handler type, matcher, config), execution audit log | @@ -85,6 +86,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. diff --git a/docs/07-bootstrap-skills-memory.md b/docs/07-bootstrap-skills-memory.md index cc2af0fb..62c721a0 100644 --- a/docs/07-bootstrap-skills-memory.md +++ b/docs/07-bootstrap-skills-memory.md @@ -279,6 +279,30 @@ This decision is re-evaluated each time the system prompt is built, so newly hot --- +## 9.5. Explicit Slash Skill Commands + +Users can bypass implicit skill matching by starting a prompt with a slash command: + +| Pattern | Behavior | +|---------|----------| +| `/ prompt` | Activates the skill by slug and treats `prompt` as the skill input | +| `/use prompt` | Activates the skill by slug or display name | +| `/list-skills` | Shows available skills for the current agent context | +| `/help ` | Shows description and usage guidance for one skill | + +Slash detection runs during prompt construction after request context is scoped and before the skills section is built. A matched skill narrows the per-request `SkillFilter` to that skill and injects the full `SKILL.md` instructions into the system prompt for the current turn only. Normal matching remains unchanged for messages that do not start with the configured prefix, path-like strings such as `/home/user/file`, or unresolved commands without suggestions. + +Tenant settings live in `system_configs`: + +| Key | Default | Behavior | +|-----|---------|----------| +| `skills.slash_commands.enabled` | `true` | Enable slash command detection | +| `skills.slash_commands.suggest_not_found` | `true` | Suggest similar skills for unknown commands | +| `skills.slash_commands.partial_matching` | `false` | Allow unique prefixes such as `/frontend` | +| `skills.slash_commands.prefix` | `/` | Single-character command prefix | + +--- + ## 10. Skills -- BM25 Search An in-memory BM25 index provides keyword-based skill search. The index is lazily rebuilt whenever the skill version changes. diff --git a/docs/09-security.md b/docs/09-security.md index 124e4fb8..5fcdc693 100644 --- a/docs/09-security.md +++ b/docs/09-security.md @@ -221,11 +221,14 @@ AES-256-GCM encryption for secrets stored in PostgreSQL. Key provided via `GOCLA | LLM provider API keys | `llm_providers` | `api_key` | | MCP server API keys | `mcp_servers` | `api_key` | | Custom tool env vars | `custom_tools` | `env` | +| Credentialed CLI env vars | `secure_cli_binaries`, `secure_cli_agent_grants`, `secure_cli_user_credentials` | `encrypted_env` | **Format**: `"aes-gcm:" + base64(12-byte nonce + ciphertext + GCM tag)` Backward compatible: values without the `aes-gcm:` prefix are returned as plaintext (for migration from unencrypted data). +Credentialed CLI env entries have a separate visibility kind inside the encrypted JSON blob when `GOCLAW_ENCRYPTION_KEY` is configured. `sensitive` entries are masked in normal API/UI responses and never returned raw except through the explicit audited grant reveal flow. `value` entries use the same at-rest storage path but are returned to authorized admins for operational review. + --- ## 4. Rate Limiting -- Gateway + Tool diff --git a/docs/14-skills-runtime.md b/docs/14-skills-runtime.md index b50798ac..512a459d 100644 --- a/docs/14-skills-runtime.md +++ b/docs/14-skills-runtime.md @@ -29,6 +29,11 @@ How skills access Python, Node.js, and system tools inside Docker containers and └─────────────────────────────────────────────────────────┘ ``` +Explicit skill activation is handled before runtime execution. When a user starts +their prompt with `/` or `/use `, the gateway resolves +the skill, injects its `SKILL.md` into the current turn, and then normal runtime +rules apply to any scripts or package dependencies that skill uses. + --- ## 2. Pre-installed Packages (Option A) @@ -139,7 +144,7 @@ To install additional packages: pip3 install or npm install -g - Run Python/Node scripts via exec tool - Install packages via `pip3 install` / `npm install -g` -- Access files in `/app/workspace/` including `.media/` subdirectory +- Access files in `/app/workspace/`, including `.uploads/` for current user uploads and `.media/` for legacy media refs - Read skill files from `.goclaw/skills-store/` ### What Agents CANNOT Do @@ -156,16 +161,18 @@ To install additional packages: pip3 install or npm install -g Uploaded files (from web chat, Telegram, Discord, etc.) are persisted to: ``` -/app/workspace/.media/{sessionHash}/{uuid}.{ext} +/app/workspace/.uploads/{safe-original-name}-{8hex}.{ext} ``` +Uploads without a usable original filename fall back to `{uuid}.{ext}`. Legacy media refs may still resolve from `.media/{sessionHash}/{uuid}.{ext}`. + The `enrichDocumentPaths()` function injects the full path into `` tags: ``` - + ``` -Agents can read these files directly via exec — no copy to `/tmp` needed. +Agents can read these files directly via exec — no copy to `/tmp` needed. For archive uploads such as `.zip`, inspect or extract with commands like `unzip -l ""` or `unzip -q "" -d `. --- diff --git a/docs/16-skill-publishing.md b/docs/16-skill-publishing.md index 14acb69e..7244c1e5 100644 --- a/docs/16-skill-publishing.md +++ b/docs/16-skill-publishing.md @@ -183,7 +183,7 @@ Unlike the HTTP upload handler, the tool does **not** archive the skill on missi | `..` in relative path | Skip (prevent traversal) | | Symlinks | Skip (prevent escape) | | System artifacts | Skip (`.DS_Store`, `__MACOSX`, `Thumbs.db`, etc.) | -| Total dir size > 20 MB | Reject with error | +| Total upload size exceeds configured limit | Reject with error. Default is 20 MB; configurable via `skills.max_upload_size_mb`, `GOCLAW_SKILLS_MAX_UPLOAD_SIZE_MB`, tenant `skills.max_upload_size_mb`, or SKILL.md frontmatter `max_upload_size_mb`; clamped to 1-500 MB. | --- diff --git a/docs/18-http-api.md b/docs/18-http-api.md index 97aa7001..e5227b98 100644 --- a/docs/18-http-api.md +++ b/docs/18-http-api.md @@ -321,7 +321,7 @@ Use `direct_selection_count` plus the `selected_provider` sequence to verify rea | Method | Path | Description | |--------|------|-------------| | `GET` | `/v1/skills` | List all skills | -| `POST` | `/v1/skills/upload` | Upload ZIP with SKILL.md (20 MB limit) | +| `POST` | `/v1/skills/upload` | Upload ZIP with SKILL.md (configurable 20 MB default, 1-500 MB range) | | `GET` | `/v1/skills/{id}` | Get skill details | | `PUT` | `/v1/skills/{id}` | Update skill metadata | | `DELETE` | `/v1/skills/{id}` | Delete skill (not system skills) | @@ -331,6 +331,18 @@ Use `direct_selection_count` plus the `selected_provider` sequence to verify rea ### Skill Grants +Skill upload size is enforced per ZIP file. The effective limit resolves in this order: +tenant `system_configs["skills.max_upload_size_mb"]`, then `SKILL.md` frontmatter +`max_upload_size_mb`, then config/env `skills.max_upload_size_mb` / +`GOCLAW_SKILLS_MAX_UPLOAD_SIZE_MB`, then the default 20 MB. Values are clamped +to 1-500 MB. + +Skill slash-command behavior is configured through tenant `system_configs`: +`skills.slash_commands.enabled`, `skills.slash_commands.suggest_not_found`, +`skills.slash_commands.partial_matching`, and `skills.slash_commands.prefix`. +The default prefix is `/`; supported prompt forms are `/ prompt`, +`/use prompt`, `/list-skills`, and `/help `. + | Method | Path | Description | |--------|------|-------------| | `POST` | `/v1/skills/{id}/grants/agent` | Grant skill to agent | @@ -390,6 +402,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 +1313,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 diff --git a/docs/21-agent-evolution-and-skill-management.md b/docs/21-agent-evolution-and-skill-management.md index 21be814c..06ddf443 100644 --- a/docs/21-agent-evolution-and-skill-management.md +++ b/docs/21-agent-evolution-and-skill-management.md @@ -450,7 +450,7 @@ System skills (`is_system=true`) cannot be modified through any path. | Symlink detection | `filepath.WalkDir` + `d.Type()&os.ModeSymlink` check | | Path traversal | `strings.Contains(rel, "..")` rejection | | Content size limit | 100KB max for SKILL.md content | -| Companion size limit | 20MB max total for companion files (scripts, assets) | +| Companion size limit | Configurable per ZIP upload; default 20MB, clamped to 1-500MB | | Soft-delete | Files moved to `.trash/`, never hard-deleted | --- diff --git a/docs/project-changelog.md b/docs/project-changelog.md index 8af81eb1..6c3c6664 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -18,6 +18,60 @@ Significant changes, features, and fixes in reverse chronological order. - Added runtime binary discovery, SecureCLI preset, deny-pattern, and Dockerfile contract coverage for Google Workspace CLI. +### Slash skill commands + +**Features** + +- Added explicit slash skill activation for `/`, `/use `, `/list-skills`, and `/help `. +- Added tenant settings for slash command enablement, similar-skill suggestions, partial matching, and custom prefix. + +**Tests** + +- Added backend coverage for parser false positives, exact/partial skill resolution, suggestions, help/list commands, and config overlays. + +### Configurable skill upload limits + +**Features** + +- Added configurable skill ZIP upload limits with config/env, SKILL.md frontmatter, and tenant system setting support. +- Added dashboard settings and dynamic upload validation so the Web UI follows the tenant limit instead of hardcoding 20MB. + +**Tests** + +- Added backend coverage for limit precedence, clamping, oversized rejection, and frontend coverage for parameterized upload validation. + +### CLI environment variable visibility + +**Features** + +- Added `sensitive` and `value` kinds for secure CLI environment variables across binary defaults, agent grant overrides, and user overrides. +- Plain value entries are visible to authorized admins for operational config review, while sensitive entries remain masked and replace-only. + +**Fixes** + +- Stopped per-user credential reads from returning legacy sensitive env values raw. +- Kept legacy `{"KEY":"value"}` env blobs backward-compatible by treating them as sensitive. + +**Tests** + +- Added backend regression coverage for env kind parsing, sanitized API responses, runtime flattening, and invalid kind rejection. +- Verified Web UI build after adding env-kind controls and warnings. + +### Command keyword allowlist + +**Features** + +- Added scoped credentialed CLI keyword allowlist config for content arguments and positional arguments, with runtime reload and web config editing. + +**Fixes** + +- Kept credentialed CLI `deny_args` active for command paths such as `gh secret set` while allowing approved GitHub issue/PR prose to mention security vocabulary. +- Added security audit logging for real allowlisted pass-throughs without logging full argument values. + +**Tests** + +- Added regression coverage for scoped keyword masking, disabled rules, unsafe positional rules, config reload, and race-safe policy snapshots. + ### Browser cookie sync and config UI **Features** @@ -39,6 +93,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** diff --git a/internal/agent/intent_classify.go b/internal/agent/intent_classify.go index 2fd4e660..46201708 100644 --- a/internal/agent/intent_classify.go +++ b/internal/agent/intent_classify.go @@ -8,6 +8,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // IntentType represents the classified intent of a user message. @@ -108,6 +109,10 @@ func containsWholeWord(s, kw string) bool { // Uses keyword fast-path first, then falls back to LLM classification. // Falls back to IntentNewTask on any error. func ClassifyIntent(ctx context.Context, provider providers.Provider, model, userMessage string) IntentType { + return ClassifyIntentWithUsageCaps(ctx, nil, provider, model, userMessage) +} + +func ClassifyIntentWithUsageCaps(ctx context.Context, usageCaps *usagecaps.Service, provider providers.Provider, model, userMessage string) IntentType { // Fast-path: keyword matching for obvious patterns (no LLM cost). if intent, ok := quickClassify(userMessage); ok { return intent @@ -116,7 +121,7 @@ func ClassifyIntent(ctx context.Context, provider providers.Provider, model, use ctx, cancel := context.WithTimeout(ctx, intentClassifyTimeout) defer cancel() - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: intentSystemPrompt}, {Role: "user", Content: userMessage}, @@ -126,6 +131,11 @@ func ClassifyIntent(ctx context.Context, provider providers.Provider, model, use providers.OptMaxTokens: 20, providers.OptTemperature: 0.0, }, + } + resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "intent-classify", + MaxOutputTokens: 20, }) if err != nil { return IntentNewTask diff --git a/internal/agent/loop_compact.go b/internal/agent/loop_compact.go index c6c31753..7adb6382 100644 --- a/internal/agent/loop_compact.go +++ b/internal/agent/loop_compact.go @@ -89,14 +89,15 @@ func (l *Loop) compactMessagesInPlace(ctx context.Context, messages []providers. inTokens := l.estimateSummaryInputTokens(toSummarize) slog.Info("compact_budget", "agent", l.id, "in_tokens", inTokens, "out_tokens", dynamicSummaryMax(inTokens)) - resp, err := l.provider.Chat(sctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{{ Role: "user", Content: compactionSummaryPrompt + sb.String(), }}, Model: l.model, Options: map[string]any{"max_tokens": dynamicSummaryMax(inTokens), "temperature": 0.3}, - }) + } + resp, err := l.callInternalLLMWithUsage(sctx, chatReq, "mid-loop-compaction") if err != nil { slog.Warn("mid_loop_compaction_failed", "agent", l.id, "error", err) return nil diff --git a/internal/agent/loop_history.go b/internal/agent/loop_history.go index da91d3b0..36cf064c 100644 --- a/internal/agent/loop_history.go +++ b/internal/agent/loop_history.go @@ -117,6 +117,8 @@ func (l *Loop) buildMessages(ctx context.Context, history []providers.Message, s } } + userMessage, extraSystemPrompt, skillFilter = l.applySkillSlashCommand(ctx, userMessage, extraSystemPrompt, skillFilter) + // Build tool list, filtering out skill_manage when skill_evolve is off. // Also applies ChannelAware filtering so channel-specific tools don't // appear in ## Tooling when the current channel doesn't support them. diff --git a/internal/agent/loop_history_sanitize.go b/internal/agent/loop_history_sanitize.go index 6e0834b8..706f0dc9 100644 --- a/internal/agent/loop_history_sanitize.go +++ b/internal/agent/loop_history_sanitize.go @@ -284,11 +284,12 @@ func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { inTokens := l.estimateSummaryInputTokens(toSummarize) slog.Info("compact_budget", "agent", l.id, "in_tokens", inTokens, "out_tokens", dynamicSummaryMax(inTokens)) - resp, err := l.provider.Chat(sctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{{Role: "user", Content: prompt.String()}}, Model: l.model, Options: map[string]any{"max_tokens": dynamicSummaryMax(inTokens), "temperature": 0.3}, - }) + } + resp, err := l.callInternalLLMWithUsage(sctx, chatReq, "session-summarization") if err != nil { slog.Warn("summarization failed", "session", sessionKey, "error", err) return diff --git a/internal/agent/loop_input_media.go b/internal/agent/loop_input_media.go index fa7989eb..49cb3394 100644 --- a/internal/agent/loop_input_media.go +++ b/internal/agent/loop_input_media.go @@ -27,10 +27,11 @@ func isTextMime(mime string) bool { } // collectRefsByKind gathers MediaRefs of a given kind from message history -// (reverse order) and current-turn refs. Historical first, current last. +// in chronological order, then appends current-turn refs. The last ref is the +// newest document for read_document's omitted media_id fallback. func collectRefsByKind(messages []providers.Message, currentRefs []providers.MediaRef, kind string) []providers.MediaRef { var refs []providers.MediaRef - for i := len(messages) - 1; i >= 0; i-- { + for i := range messages { for _, ref := range messages[i].MediaRefs { if ref.Kind == kind { refs = append(refs, ref) diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index c8abda9a..9d19e0b6 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -2,6 +2,7 @@ package agent import ( "context" + "fmt" "log/slog" "strings" "time" @@ -15,6 +16,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/internal/workspace" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -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), + }) +} diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index b21c1817..469362a7 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) func (l *Loop) emit(event AgentEvent) { @@ -51,20 +52,33 @@ func (l *Loop) IsRunning() bool { return l.activeRuns.Load() > 0 } type spanOption func(*spanOverrides) type spanOverrides struct { - model string - provider string + model string + provider string + usageCapAttempts []usagecaps.TraceMetadata } func withModel(m string) spanOption { return func(o *spanOverrides) { o.model = m } } func withProvider(p string) spanOption { return func(o *spanOverrides) { o.provider = p } } +func withUsageCapMetadata(metadata usagecaps.TraceMetadata) spanOption { + return func(o *spanOverrides) { + if !metadata.Empty() { + o.usageCapAttempts = append(o.usageCapAttempts, metadata) + } + } +} // resolveSpan returns (model, provider) applying any overrides on top of agent defaults. func (l *Loop) resolveSpan(opts []spanOption) (string, string) { + o := l.resolveSpanOverrides(opts) + return o.model, o.provider +} + +func (l *Loop) resolveSpanOverrides(opts []spanOption) spanOverrides { o := spanOverrides{model: l.model, provider: l.provider.Name()} for _, fn := range opts { fn(&o) } - return o.model, o.provider + return o } // --------------------------------------------------------------------------- @@ -150,6 +164,7 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time. "status": store.SpanStatusCompleted, } var spanMetadata json.RawMessage + spanOpts := l.resolveSpanOverrides(opts) if callErr != nil { updates["status"] = store.SpanStatusError @@ -176,8 +191,7 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time. } } // Calculate cost if pricing config is available. - model, providerName := l.resolveSpan(opts) - if pricing := tracing.LookupPricing(l.modelPricing, providerName, model); pricing != nil { + if pricing := tracing.LookupPricing(l.modelPricing, spanOpts.provider, spanOpts.model); pricing != nil { cost := tracing.CalculateCost(pricing, resp.Usage) if cost > 0 { updates["total_cost"] = cost @@ -203,6 +217,9 @@ func (l *Loop) emitLLMSpanEnd(ctx context.Context, spanID uuid.UUID, start time. if decision := providers.ReasoningDecisionFromContext(ctx); decision != nil { spanMetadata = providers.MergeReasoningMetadata(spanMetadata, *decision) } + if len(spanOpts.usageCapAttempts) > 0 { + spanMetadata = usagecaps.MergeTraceMetadata(spanMetadata, spanOpts.usageCapAttempts) + } if len(spanMetadata) > 0 { updates["metadata"] = spanMetadata } diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index 50c50228..816e06a1 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -23,6 +23,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/tokencount" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // bootstrapAutoCleanupTurns is the number of user messages after which @@ -80,13 +81,13 @@ type Loop struct { // agentUUID is the canonical DB primary key. Use for SQL WHERE/JOIN, // DomainEvent.AgentID, OTel span attributes, and context propagation via // store.WithAgentID. See docs/agent-identity-conventions.md. - agentUUID uuid.UUID - tenantID uuid.UUID // agent's owning tenant + agentUUID uuid.UUID + tenantID uuid.UUID // agent's owning tenant // agentOtherConfig is a defensive byte copy of agents.other_config JSONB. // Copied once at Loop construction; used to build AgentAudioSnapshot at tool dispatch. agentOtherConfig json.RawMessage - agentType string // "open" or "predefined" - defaultTimezone string // system default timezone for bootstrap pre-fill + agentType string // "open" or "predefined" + defaultTimezone string // system default timezone for bootstrap pre-fill provider providers.Provider model string modelRegistry providers.ModelRegistry // resolves per-model context window at run time (nil = use static contextWindow) @@ -108,7 +109,7 @@ type Loop struct { // Memory flush runs if callback != nil; auto-inject runs if AutoInjector != nil. autoInjector memory.AutoInjector // v3 L0 memory auto-inject (nil = disabled) - eventPub bus.EventPublisher // currently unused by Loop; kept for future use + eventPub bus.EventPublisher // currently unused by Loop; kept for future use domainBus eventbus.DomainEventBus // V3 domain event bus for consolidation pipeline sessions store.SessionStore tools tools.ToolExecutor @@ -121,11 +122,12 @@ type Loop struct { summarizeMu sync.Map // sessionKey → *sync.Mutex // Bootstrap/persona context (loaded at startup, injected into system prompt) - ownerIDs []string - skillsLoader *skills.Loader - skillAllowList []string // nil = all, [] = none, ["x","y"] = filter - hasMemory bool - contextFiles []bootstrap.ContextFile + ownerIDs []string + skillsLoader *skills.Loader + skillAllowList []string // nil = all, [] = none, ["x","y"] = filter + skillSlashCommands config.SkillSlashCommandConfig + hasMemory bool + contextFiles []bootstrap.ContextFile // Per-user profile + file seeding + dynamic context loading ensureUserProfile EnsureUserProfileFunc // create/resolve user profile + workspace @@ -137,11 +139,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 @@ -184,6 +186,7 @@ type Loop struct { // Tenant-specific allowed paths beyond workspace (from system_configs['allowed_paths']). // Filesystem tools (read_file, write_file, edit, list_files) check these at execution time. tenantAllowedPaths []string + systemConfigs store.SystemConfigStore // Per-tenant disabled tools (tool name → true means excluded from LLM) disabledTools map[string]bool @@ -239,13 +242,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 @@ -328,11 +332,12 @@ type LoopConfig struct { OnEvent func(AgentEvent) // Bootstrap/persona context - OwnerIDs []string - SkillsLoader *skills.Loader - SkillAllowList []string // nil = all, [] = none, ["x","y"] = filter - HasMemory bool - ContextFiles []bootstrap.ContextFile + OwnerIDs []string + SkillsLoader *skills.Loader + SkillAllowList []string // nil = all, [] = none, ["x","y"] = filter + SkillSlashCommands config.SkillSlashCommandConfig + HasMemory bool + ContextFiles []bootstrap.ContextFile // Compaction config CompactionCfg *config.CompactionConfig @@ -350,11 +355,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 @@ -381,6 +386,7 @@ type LoopConfig struct { // Tenant-specific allowed paths beyond workspace (from system_configs['allowed_paths']). TenantAllowedPaths []string + SystemConfigs store.SystemConfigStore // Per-tenant disabled tools (tool name → true means excluded) DisabledTools map[string]bool @@ -430,19 +436,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 @@ -527,6 +534,7 @@ func NewLoop(cfg LoopConfig) *Loop { ownerIDs: cfg.OwnerIDs, skillsLoader: cfg.SkillsLoader, skillAllowList: cfg.SkillAllowList, + skillSlashCommands: cfg.SkillSlashCommands, hasMemory: cfg.HasMemory, contextFiles: cfg.ContextFiles, defaultTimezone: cfg.DefaultTimezone, @@ -550,6 +558,7 @@ func NewLoop(cfg LoopConfig) *Loop { builtinToolSettings: cfg.BuiltinToolSettings, tenantToolSettings: cfg.TenantToolSettings, tenantAllowedPaths: cfg.TenantAllowedPaths, + systemConfigs: cfg.SystemConfigs, disabledTools: cfg.DisabledTools, reasoningConfig: cfg.ReasoningConfig, promptMode: cfg.PromptMode, @@ -568,6 +577,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 +592,36 @@ func NewLoop(cfg LoopConfig) *Loop { // RunRequest is the input for processing a message through the agent. type RunRequest struct { - SessionKey string // composite key: agent:{agentId}:{channel}:{peerKind}:{chatId} - Message string // user message - Media []bus.MediaFile // local media files with MIME types - ForwardMedia []bus.MediaFile // media files to forward to output (from delegation results) - Channel string // source channel instance name (e.g. "my-telegram-bot") - ChannelType string // platform type (e.g. "zalo_personal", "telegram") — for system prompt context - BitrixPortalDomain string // bitrix24-only: portal domain (e.g. "tamgiac.bitrix24.com") for entity URL construction - ChatTitle string // group chat display name (e.g. Telegram group title) - ChatID string // source chat ID - PeerKind string // "direct" or "group" (for session key building and tool context) - RunID string // unique run identifier - UserID string // external user ID (TEXT, free-form) for multi-tenant scoping - SenderID string // original individual sender ID (preserved in group chats for permission checks) - SenderName string // display name from channel metadata (for bootstrap auto-contact) - Role string // caller's RBAC role (admin/operator/viewer/owner); bypasses per-user grants for authenticated admins (#915) - Stream bool // whether to stream response chunks - ExtraSystemPrompt string // optional: injected into system prompt (skills, subagent context, etc.) - SkillFilter []string // per-request skill override: nil=use agent default, []=no skills, ["x","y"]=whitelist - HistoryLimit int // max user turns to keep in context (0=unlimited, from channel config) - ToolAllow []string // per-group tool allow list (nil = no restriction, supports "group:xxx") - LocalKey string // composite key with topic/thread suffix for routing (e.g. "-100123:topic:42") - ParentTraceID uuid.UUID // if set, reuse parent trace instead of creating new (announce runs) - ParentRootSpanID uuid.UUID // if set, nest announce agent span under this parent span - LinkedTraceID uuid.UUID // if set, create new trace with parent_trace_id pointing to this (team task runs) - TraceName string // override trace name (default: "chat ") - TraceTags []string // additional tags for the trace (e.g. "cron") - MaxIterations int // per-request override (0 = use agent default, must be lower) - ModelOverride string // per-request model override (heartbeat uses cheaper model) - ProviderOverride providers.Provider // per-request provider override (heartbeat uses different provider) - LightContext bool // skip loading context files (only inject ExtraSystemPrompt) + SessionKey string // composite key: agent:{agentId}:{channel}:{peerKind}:{chatId} + Message string // user message + Media []bus.MediaFile // local media files with MIME types + ForwardMedia []bus.MediaFile // media files to forward to output (from delegation results) + Channel string // source channel instance name (e.g. "my-telegram-bot") + ChannelType string // platform type (e.g. "zalo_personal", "telegram") — for system prompt context + BitrixPortalDomain string // bitrix24-only: portal domain (e.g. "tamgiac.bitrix24.com") for entity URL construction + ChatTitle string // group chat display name (e.g. Telegram group title) + ChatID string // source chat ID + PeerKind string // "direct" or "group" (for session key building and tool context) + RunID string // unique run identifier + UserID string // external user ID (TEXT, free-form) for multi-tenant scoping + SenderID string // original individual sender ID (preserved in group chats for permission checks) + SenderName string // display name from channel metadata (for bootstrap auto-contact) + Role string // caller's RBAC role (admin/operator/viewer/owner); bypasses per-user grants for authenticated admins (#915) + Stream bool // whether to stream response chunks + ExtraSystemPrompt string // optional: injected into system prompt (skills, subagent context, etc.) + SkillFilter []string // per-request skill override: nil=use agent default, []=no skills, ["x","y"]=whitelist + HistoryLimit int // max user turns to keep in context (0=unlimited, from channel config) + ToolAllow []string // per-group tool allow list (nil = no restriction, supports "group:xxx") + LocalKey string // composite key with topic/thread suffix for routing (e.g. "-100123:topic:42") + ParentTraceID uuid.UUID // if set, reuse parent trace instead of creating new (announce runs) + ParentRootSpanID uuid.UUID // if set, nest announce agent span under this parent span + LinkedTraceID uuid.UUID // if set, create new trace with parent_trace_id pointing to this (team task runs) + TraceName string // override trace name (default: "chat ") + TraceTags []string // additional tags for the trace (e.g. "cron") + MaxIterations int // per-request override (0 = use agent default, must be lower) + ModelOverride string // per-request model override (heartbeat uses cheaper model) + ProviderOverride providers.Provider // per-request provider override (heartbeat uses different provider) + LightContext bool // skip loading context files (only inject ExtraSystemPrompt) // Run classification RunKind string // "delegation", "announce" — empty for user-initiated runs @@ -646,7 +656,7 @@ type RunRequest struct { // RunResult is the output of a completed agent run. type RunResult struct { Content string `json:"content"` - Thinking string `json:"thinking,omitempty"` // reasoning content from thinking models (Claude, o3, DeepSeek-R1, Kimi) + Thinking string `json:"thinking,omitempty"` // reasoning content from thinking models (Claude, o3, DeepSeek-R1, Kimi) RunID string `json:"runId"` Iterations int `json:"iterations"` Usage *providers.Usage `json:"usage,omitempty"` diff --git a/internal/agent/media_persist_test.go b/internal/agent/media_persist_test.go index 297f4a94..b819d208 100644 --- a/internal/agent/media_persist_test.go +++ b/internal/agent/media_persist_test.go @@ -72,6 +72,14 @@ func TestPersistMedia_NamingScheme(t *testing.T) { wantSlug: "", wantExtPat: `\.pdf`, }, + { + name: "zip_archive_preserves_archive_extension", + filename: "codex.zip", + content: "zip bytes", + mime: "application/zip", + wantSlug: "codex", + wantExtPat: `\.zip`, + }, } var loop Loop diff --git a/internal/agent/media_test.go b/internal/agent/media_test.go index cd66173c..f5fd4144 100644 --- a/internal/agent/media_test.go +++ b/internal/agent/media_test.go @@ -73,6 +73,34 @@ func TestEnrichImageIDs_SkipsAlreadyEnriched(t *testing.T) { } } +func TestCollectRefsByKindOrdersOldestToNewestThenCurrent(t *testing.T) { + messages := []providers.Message{ + { + Role: "user", + Content: "old", + MediaRefs: []providers.MediaRef{{ID: "old-doc", Kind: "document"}}, + }, + { + Role: "user", + Content: "latest-history", + MediaRefs: []providers.MediaRef{{ID: "latest-history-doc", Kind: "document"}}, + }, + } + current := []providers.MediaRef{{ID: "current-doc", Kind: "document"}} + + got := collectRefsByKind(messages, current, "document") + want := []string{"old-doc", "latest-history-doc", "current-doc"} + + if len(got) != len(want) { + t.Fatalf("got %d refs, want %d", len(got), len(want)) + } + for i := range want { + if got[i].ID != want[i] { + t.Fatalf("ref[%d] = %q, want %q", i, got[i].ID, want[i]) + } + } +} + // testMediaStore creates a temporary media.Store for tests. func testMediaStore(t *testing.T) *media.Store { t.Helper() diff --git a/internal/agent/memoryflush.go b/internal/agent/memoryflush.go index d5106100..7db493a0 100644 --- a/internal/agent/memoryflush.go +++ b/internal/agent/memoryflush.go @@ -186,7 +186,7 @@ func (l *Loop) runMemoryFlush(ctx context.Context, sessionKey string, settings * // Run LLM iteration loop (max 5 iterations for flush) maxFlushIter := 5 for range maxFlushIter { - resp, err := l.provider.Chat(flushCtx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: messages, Tools: toolDefs, Model: l.model, @@ -194,7 +194,8 @@ func (l *Loop) runMemoryFlush(ctx context.Context, sessionKey string, settings * "max_tokens": 4096, "temperature": 0.3, }, - }) + } + resp, err := l.callInternalLLMWithUsage(flushCtx, chatReq, "memory-flush") if err != nil { slog.Warn("memory flush: LLM call failed", "error", err) l.extractiveMemoryFallback(flushCtx, sessionKey, history, "LLM error") diff --git a/internal/agent/resolver.go b/internal/agent/resolver.go index afd3a936..823e17fc 100644 --- a/internal/agent/resolver.go +++ b/internal/agent/resolver.go @@ -25,6 +25,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ResolverDeps holds shared dependencies for the agent resolver. @@ -84,7 +85,8 @@ type ResolverDeps struct { MCPGrantChecker mcpbridge.GrantChecker // Skill access store — for per-agent skill visibility filtering - SkillAccessStore store.SkillAccessStore + SkillAccessStore store.SkillAccessStore + SkillSlashCommands config.SkillSlashCommandConfig // Config permission store for group file writer checks ConfigPermStore store.ConfigPermissionStore @@ -97,6 +99,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 @@ -489,6 +492,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { AgentToolPolicy: agentToolPolicyForTeam(agentToolPolicyWithWorkspace(agentToolPolicyWithMCP(ag.ParseToolsConfig(), hasMCPTools), hasTeam), isTeamLead), SkillsLoader: deps.Skills, SkillAllowList: skillAllowList, + SkillSlashCommands: deps.SkillSlashCommands, HasMemory: hasMemory, ContextFiles: contextFiles, EnsureUserProfile: deps.EnsureUserProfile, @@ -509,6 +513,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { BuiltinToolSettings: builtinSettings, TenantToolSettings: tenantToolSettings, TenantAllowedPaths: tenantAllowedPaths, + SystemConfigs: deps.SystemConfigs, DisabledTools: disabledTools, ReasoningConfig: store.ResolveEffectiveReasoningConfig(providerReasoningDefaults, ag.ParseReasoningConfig()), PromptMode: PromptMode(ag.ParsePromptMode()), @@ -528,6 +533,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { ModelPricing: deps.ModelPricing, BudgetMonthlyCents: derefInt(ag.BudgetMonthlyCents), TracingStore: deps.TracingStore, + UsageCaps: deps.UsageCaps, MemoryStore: deps.MemoryStore, MCPStore: deps.MCPStore, MCPPool: deps.MCPPool, diff --git a/internal/agent/skill_slash_command_guidance.go b/internal/agent/skill_slash_command_guidance.go new file mode 100644 index 00000000..6a652e92 --- /dev/null +++ b/internal/agent/skill_slash_command_guidance.go @@ -0,0 +1,65 @@ +package agent + +import ( + "fmt" + "path/filepath" + "sort" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/skills" +) + +func (r skillSlashCommandResult) systemPromptSection() string { + switch r.Kind { + case skillSlashCommandActivate: + return fmt.Sprintf("## Explicit Skill Activation\n\nThe user explicitly requested skill `%s` (%s). Use this skill for the current request and treat the remaining user message as the skill input. Call `use_skill` with name `%s` for observability when available, then follow these instructions:\n\n%s", r.Skill.Slug, r.Skill.Name, r.Skill.Slug, r.SkillContent) + case skillSlashCommandList, skillSlashCommandHelp, skillSlashCommandUnknown: + return r.Guidance + default: + return "" + } +} + +func buildSkillSlashListGuidance(all []skills.Info) string { + if len(all) == 0 { + return "## Skill Slash Command\n\nNo skills are currently available." + } + sort.Slice(all, func(i, j int) bool { return all[i].Slug < all[j].Slug }) + var lines []string + lines = append(lines, "## Skill Slash Command", "", "Available skills:") + for _, skill := range all { + lines = append(lines, fmt.Sprintf("- `%s` - %s", skill.Slug, skillDisplayDescription(skill))) + } + return strings.Join(lines, "\n") +} + +func buildSkillSlashHelpGuidance(skill skills.Info) string { + return fmt.Sprintf("## Skill Slash Command\n\nSkill `%s` (%s)\nDescription: %s\nLocation: %s\nExplain this skill and how to invoke it with the configured slash prefix.", skill.Slug, skill.Name, skillDisplayDescription(skill), filepath.ToSlash(skill.Path)) +} + +func buildSkillSlashUnknownGuidance(target string, suggestions []skills.Info) string { + var lines []string + lines = append(lines, "## Skill Slash Command", "", fmt.Sprintf("Requested skill `%s` was not found. Suggest these available alternatives:", target)) + for _, skill := range suggestions { + lines = append(lines, fmt.Sprintf("- `%s` - %s", skill.Slug, skillDisplayDescription(skill))) + } + return strings.Join(lines, "\n") +} + +func skillDisplayDescription(skill skills.Info) string { + if strings.TrimSpace(skill.Description) != "" { + return strings.TrimSpace(skill.Description) + } + return skill.Name +} + +func appendExtraPrompt(existing, addition string) string { + addition = strings.TrimSpace(addition) + if addition == "" { + return existing + } + if strings.TrimSpace(existing) == "" { + return addition + } + return existing + "\n\n" + addition +} diff --git a/internal/agent/skill_slash_command_matching.go b/internal/agent/skill_slash_command_matching.go new file mode 100644 index 00000000..68a4bb1c --- /dev/null +++ b/internal/agent/skill_slash_command_matching.go @@ -0,0 +1,124 @@ +package agent + +import ( + "sort" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/skills" +) + +type parsedSkillSlashCommand struct { + verb string + target string + rest string +} + +func parseSkillSlashCommand(message, prefix string) (parsedSkillSlashCommand, bool) { + message = strings.TrimSpace(message) + if message == "" || !strings.HasPrefix(message, prefix) { + return parsedSkillSlashCommand{}, false + } + after := strings.TrimSpace(strings.TrimPrefix(message, prefix)) + if after == "" || looksLikePath(after) { + return parsedSkillSlashCommand{}, false + } + first, rest, _ := strings.Cut(after, " ") + first = strings.TrimSpace(first) + rest = strings.TrimSpace(rest) + switch strings.ToLower(first) { + case "list-skills": + return parsedSkillSlashCommand{verb: "list-skills"}, true + case "help": + if rest == "" { + return parsedSkillSlashCommand{}, false + } + return parsedSkillSlashCommand{verb: "help", target: rest}, true + case "use", "activate": + if rest == "" { + return parsedSkillSlashCommand{}, false + } + return parsedSkillSlashCommand{verb: strings.ToLower(first), rest: rest}, true + default: + return parsedSkillSlashCommand{verb: "direct", target: first, rest: rest}, true + } +} + +func looksLikePath(value string) bool { + first, _, _ := strings.Cut(value, " ") + return strings.Contains(first, "/") || strings.Contains(first, "\\") || strings.Contains(first, ".") +} + +func firstWord(value string) (string, string) { + first, rest, ok := strings.Cut(strings.TrimSpace(value), " ") + if !ok { + return strings.TrimSpace(value), "" + } + return strings.TrimSpace(first), strings.TrimSpace(rest) +} + +func matchSkillCommandTarget(all []skills.Info, raw string, partial bool) (skills.Info, bool, string) { + raw = strings.TrimSpace(raw) + if raw == "" { + return skills.Info{}, false, "" + } + type candidate struct { + info skills.Info + matchText string + remainder string + score int + } + var matches []candidate + lowerRaw := strings.ToLower(raw) + partialTarget, partialRemainder := firstWord(raw) + lowerPartialTarget := strings.ToLower(partialTarget) + for _, skill := range all { + for _, value := range []string{skill.Slug, skill.Name} { + value = strings.TrimSpace(value) + if value == "" { + continue + } + lowerValue := strings.ToLower(value) + if lowerRaw == lowerValue { + matches = append(matches, candidate{info: skill, matchText: value, score: len([]rune(value))}) + continue + } + if strings.HasPrefix(lowerRaw, lowerValue+" ") { + remainder := trimMatchedSkillCommandPrefix(raw, value) + matches = append(matches, candidate{info: skill, matchText: value, remainder: remainder, score: len([]rune(value))}) + continue + } + if partial && lowerPartialTarget != "" && strings.HasPrefix(lowerValue, lowerPartialTarget) { + matches = append(matches, candidate{info: skill, matchText: value, remainder: partialRemainder, score: len([]rune(partialTarget))}) + } + } + } + if len(matches) == 0 { + return skills.Info{}, false, "" + } + sort.Slice(matches, func(i, j int) bool { + return matches[i].score > matches[j].score + }) + bestBySlug := make([]candidate, 0, len(matches)) + seenSlug := make(map[string]struct{}, len(matches)) + for _, match := range matches { + if _, ok := seenSlug[match.info.Slug]; ok { + continue + } + seenSlug[match.info.Slug] = struct{}{} + bestBySlug = append(bestBySlug, match) + } + best := bestBySlug[0] + if len(bestBySlug) > 1 && bestBySlug[0].score == bestBySlug[1].score { + return skills.Info{}, false, "" + } + return best.info, true, best.remainder +} + +func trimMatchedSkillCommandPrefix(raw, matched string) string { + rawRunes := []rune(raw) + matchedRunes := []rune(matched) + if len(rawRunes) < len(matchedRunes) { + return "" + } + return strings.TrimSpace(string(rawRunes[len(matchedRunes):])) +} diff --git a/internal/agent/skill_slash_command_similarity.go b/internal/agent/skill_slash_command_similarity.go new file mode 100644 index 00000000..8b3c9fed --- /dev/null +++ b/internal/agent/skill_slash_command_similarity.go @@ -0,0 +1,102 @@ +package agent + +import ( + "sort" + "strings" + "unicode" + + "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/skills" +) + +func unknownSkillSlashResult(all []skills.Info, target string, cfg config.SkillSlashCommandConfig) skillSlashCommandResult { + if !cfg.EffectiveSuggestNotFound() { + return skillSlashCommandResult{Kind: skillSlashCommandNone} + } + suggestions := similarSkills(all, target, 3) + if len(suggestions) == 0 { + return skillSlashCommandResult{Kind: skillSlashCommandNone} + } + return skillSlashCommandResult{ + Kind: skillSlashCommandUnknown, + Guidance: buildSkillSlashUnknownGuidance(target, suggestions), + Suggestions: suggestions, + } +} + +func similarSkills(all []skills.Info, target string, limit int) []skills.Info { + target = strings.ToLower(strings.TrimSpace(target)) + if target == "" { + return nil + } + type scored struct { + info skills.Info + score int + } + var scoredSkills []scored + for _, skill := range all { + best := scoreSimilarSkill(target, skill) + if best <= 3 { + scoredSkills = append(scoredSkills, scored{info: skill, score: best}) + } + } + sort.Slice(scoredSkills, func(i, j int) bool { + if scoredSkills[i].score == scoredSkills[j].score { + return scoredSkills[i].info.Slug < scoredSkills[j].info.Slug + } + return scoredSkills[i].score < scoredSkills[j].score + }) + if len(scoredSkills) > limit { + scoredSkills = scoredSkills[:limit] + } + out := make([]skills.Info, len(scoredSkills)) + for i, scored := range scoredSkills { + out[i] = scored.info + } + return out +} + +func scoreSimilarSkill(target string, skill skills.Info) int { + targetRunes := []rune(target) + slug := strings.ToLower(skill.Slug) + best := slashSkillDistance(target, slug) + if slugRunes := []rune(slug); len(slugRunes) >= len(targetRunes) { + best = min(best, slashSkillDistance(target, string(slugRunes[:len(targetRunes)]))) + } + name := strings.ToLower(skill.Name) + if name != "" { + best = min(best, slashSkillDistance(target, name)) + if nameRunes := []rune(name); len(nameRunes) >= len(targetRunes) { + best = min(best, slashSkillDistance(target, string(nameRunes[:len(targetRunes)]))) + } + } + if strings.HasPrefix(slug, target) || strings.Contains(name, target) { + best = 0 + } + return best +} + +func slashSkillDistance(a, b string) int { + ar := []rune(a) + br := []rune(b) + if len(ar) == 0 { + return len(br) + } + prev := make([]int, len(br)+1) + for j := range prev { + prev[j] = j + } + for i, ca := range ar { + cur := make([]int, len(br)+1) + cur[0] = i + 1 + for j, cb := range br { + cost := 0 + if unicode.ToLower(ca) != unicode.ToLower(cb) { + cost = 1 + } + cur[j+1] = min(min(cur[j]+1, prev[j+1]+1), prev[j]+cost) + } + prev = cur + } + return prev[len(br)] +} diff --git a/internal/agent/skill_slash_commands.go b/internal/agent/skill_slash_commands.go new file mode 100644 index 00000000..c9b1f34e --- /dev/null +++ b/internal/agent/skill_slash_commands.go @@ -0,0 +1,130 @@ +package agent + +import ( + "context" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/skills" +) + +type skillSlashCommandKind int + +const ( + skillSlashCommandNone skillSlashCommandKind = iota + skillSlashCommandActivate + skillSlashCommandList + skillSlashCommandHelp + skillSlashCommandUnknown +) + +type skillSlashCommandResult struct { + Kind skillSlashCommandKind + Skill skills.Info + SkillContent string + RemainingPrompt string + Guidance string + Suggestions []skills.Info +} + +func (l *Loop) applySkillSlashCommand(ctx context.Context, message, extraPrompt string, skillFilter []string) (string, string, []string) { + result := resolveSkillSlashCommand(ctx, l.skillsLoader, l.resolveSkillSlashCommandConfig(ctx), message) + if result.Kind == skillSlashCommandNone { + return message, extraPrompt, skillFilter + } + extraPrompt = appendExtraPrompt(extraPrompt, result.systemPromptSection()) + switch result.Kind { + case skillSlashCommandActivate: + if result.RemainingPrompt == "" { + message = "Use the activated skill to help with the user's request." + } else { + message = result.RemainingPrompt + } + skillFilter = []string{result.Skill.Slug} + case skillSlashCommandList: + message = "List the available skills shown in the system instructions." + case skillSlashCommandHelp: + message = "Explain the requested skill and how it should be used." + case skillSlashCommandUnknown: + message = "Explain that the requested skill was not found and suggest available alternatives." + } + return message, extraPrompt, skillFilter +} + +func (l *Loop) resolveSkillSlashCommandConfig(ctx context.Context) config.SkillSlashCommandConfig { + cfg := l.skillSlashCommands + if l.systemConfigs == nil { + return cfg + } + if raw, err := l.systemConfigs.Get(ctx, config.SkillSlashCommandsEnabledSystemConfigKey); err == nil && strings.TrimSpace(raw) != "" { + v := parseSkillSlashBool(raw) + cfg.Enabled = &v + } + if raw, err := l.systemConfigs.Get(ctx, config.SkillSlashSuggestNotFoundSystemConfigKey); err == nil && strings.TrimSpace(raw) != "" { + v := parseSkillSlashBool(raw) + cfg.SuggestNotFound = &v + } + if raw, err := l.systemConfigs.Get(ctx, config.SkillSlashPartialMatchingSystemConfigKey); err == nil && strings.TrimSpace(raw) != "" { + cfg.PartialMatching = parseSkillSlashBool(raw) + } + if raw, err := l.systemConfigs.Get(ctx, config.SkillSlashCommandPrefixSystemConfigKey); err == nil && strings.TrimSpace(raw) != "" { + cfg.Prefix = raw + } + return cfg +} + +func resolveSkillSlashCommand(ctx context.Context, loader *skills.Loader, cfg config.SkillSlashCommandConfig, message string) skillSlashCommandResult { + if loader == nil || !cfg.EffectiveEnabled() { + return skillSlashCommandResult{Kind: skillSlashCommandNone} + } + parsed, ok := parseSkillSlashCommand(message, cfg.EffectivePrefix()) + if !ok { + return skillSlashCommandResult{Kind: skillSlashCommandNone} + } + all := loader.ListSkills(ctx) + switch parsed.verb { + case "list-skills": + return skillSlashCommandResult{Kind: skillSlashCommandList, Guidance: buildSkillSlashListGuidance(all)} + case "help": + skill, matched, _ := matchSkillCommandTarget(all, parsed.target, cfg.EffectivePartialMatching()) + if !matched { + return unknownSkillSlashResult(all, parsed.target, cfg) + } + return skillSlashCommandResult{Kind: skillSlashCommandHelp, Skill: skill, Guidance: buildSkillSlashHelpGuidance(skill)} + case "use", "activate": + return resolveSkillActivation(ctx, loader, all, parsed.rest, cfg) + default: + return resolveSkillActivation(ctx, loader, all, parsed.target+" "+parsed.rest, cfg) + } +} + +func resolveSkillActivation(ctx context.Context, loader *skills.Loader, all []skills.Info, raw string, cfg config.SkillSlashCommandConfig) skillSlashCommandResult { + skill, matched, remainder := matchSkillCommandTarget(all, raw, cfg.EffectivePartialMatching()) + if !matched { + fields := strings.Fields(raw) + target := strings.TrimSpace(raw) + if len(fields) > 0 { + target = fields[0] + } + return unknownSkillSlashResult(all, target, cfg) + } + content, ok := loader.LoadSkill(ctx, skill.Slug) + if !ok { + return unknownSkillSlashResult(all, skill.Slug, cfg) + } + return skillSlashCommandResult{ + Kind: skillSlashCommandActivate, + Skill: skill, + SkillContent: content, + RemainingPrompt: strings.TrimSpace(remainder), + } +} + +func parseSkillSlashBool(raw string) bool { + switch strings.ToLower(strings.TrimSpace(raw)) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} diff --git a/internal/agent/skill_slash_commands_test.go b/internal/agent/skill_slash_commands_test.go new file mode 100644 index 00000000..d67c0d7b --- /dev/null +++ b/internal/agent/skill_slash_commands_test.go @@ -0,0 +1,138 @@ +package agent + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/skills" +) + +func TestResolveSkillSlashCommandExactSlug(t *testing.T) { + loader := newSlashTestLoader(t) + result := resolveSkillSlashCommand(context.Background(), loader, config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/"}, "/frontend-design build a landing page") + + if result.Kind != skillSlashCommandActivate { + t.Fatalf("kind = %v, want activate", result.Kind) + } + if result.Skill.Slug != "frontend-design" { + t.Fatalf("slug = %q, want frontend-design", result.Skill.Slug) + } + if result.RemainingPrompt != "build a landing page" { + t.Fatalf("remaining = %q", result.RemainingPrompt) + } + if !strings.Contains(result.SkillContent, "Use responsive components.") { + t.Fatal("expected loaded SKILL.md content") + } +} + +func TestResolveSkillSlashCommandExactNameUseSyntax(t *testing.T) { + loader := newSlashTestLoader(t) + result := resolveSkillSlashCommand(context.Background(), loader, config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/"}, "/use Frontend Design build a landing page") + + if result.Kind != skillSlashCommandActivate { + t.Fatalf("kind = %v, want activate", result.Kind) + } + if result.Skill.Slug != "frontend-design" { + t.Fatalf("slug = %q, want frontend-design", result.Skill.Slug) + } + if result.RemainingPrompt != "build a landing page" { + t.Fatalf("remaining = %q", result.RemainingPrompt) + } +} + +func TestResolveSkillSlashCommandPartialMatchRequiresUniqueEnabled(t *testing.T) { + loader := newSlashTestLoader(t) + + disabled := resolveSkillSlashCommand(context.Background(), loader, config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/"}, "/front build") + if disabled.Kind != skillSlashCommandUnknown { + t.Fatalf("disabled partial kind = %v, want unknown", disabled.Kind) + } + + enabled := resolveSkillSlashCommand(context.Background(), loader, config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/", PartialMatching: true}, "/front build") + if enabled.Kind != skillSlashCommandActivate { + t.Fatalf("enabled partial kind = %v, want activate", enabled.Kind) + } + if enabled.Skill.Slug != "frontend-design" { + t.Fatalf("slug = %q, want frontend-design", enabled.Skill.Slug) + } +} + +func TestResolveSkillSlashCommandFalsePositives(t *testing.T) { + loader := newSlashTestLoader(t) + for _, msg := range []string{"/home/user/project", "/etc/config.yaml", "https://example.com/path", "regular prompt"} { + result := resolveSkillSlashCommand(context.Background(), loader, config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/"}, msg) + if result.Kind != skillSlashCommandNone { + t.Fatalf("%q kind = %v, want none", msg, result.Kind) + } + } +} + +func TestResolveSkillSlashCommandListAndHelp(t *testing.T) { + loader := newSlashTestLoader(t) + cfg := config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/"} + + list := resolveSkillSlashCommand(context.Background(), loader, cfg, "/list-skills") + if list.Kind != skillSlashCommandList { + t.Fatalf("list kind = %v, want list", list.Kind) + } + if !strings.Contains(list.Guidance, "frontend-design") || !strings.Contains(list.Guidance, "git-helper") { + t.Fatalf("list guidance missing skills: %s", list.Guidance) + } + + help := resolveSkillSlashCommand(context.Background(), loader, cfg, "/help frontend-design") + if help.Kind != skillSlashCommandHelp { + t.Fatalf("help kind = %v, want help", help.Kind) + } + if help.Skill.Slug != "frontend-design" || !strings.Contains(help.Guidance, "Frontend Design") { + t.Fatalf("unexpected help result: %#v", help) + } + + helpByName := resolveSkillSlashCommand(context.Background(), loader, cfg, "/help Frontend Design") + if helpByName.Kind != skillSlashCommandHelp { + t.Fatalf("help by name kind = %v, want help", helpByName.Kind) + } + if helpByName.Skill.Slug != "frontend-design" { + t.Fatalf("help by name slug = %q, want frontend-design", helpByName.Skill.Slug) + } +} + +func TestResolveSkillSlashCommandSuggestsUnknown(t *testing.T) { + loader := newSlashTestLoader(t) + result := resolveSkillSlashCommand(context.Background(), loader, config.SkillSlashCommandConfig{Enabled: boolPtr(true), Prefix: "/", SuggestNotFound: boolPtr(true)}, "/fronted build") + + if result.Kind != skillSlashCommandUnknown { + t.Fatalf("kind = %v, want unknown", result.Kind) + } + if len(result.Suggestions) == 0 || result.Suggestions[0].Slug != "frontend-design" { + t.Fatalf("suggestions = %#v", result.Suggestions) + } +} + +func boolPtr(v bool) *bool { + return &v +} + +func newSlashTestLoader(t *testing.T) *skills.Loader { + t.Helper() + t.Setenv("GOCLAW_DISABLE_PERSONAL_SKILLS", "1") + root := t.TempDir() + writeSkill(t, root, "frontend-design", "Frontend Design", "Create polished UI layouts.", "Use responsive components.") + writeSkill(t, root, "git-helper", "Git Helper", "Handle git workflows.", "Use clean commits.") + return skills.NewLoader("", root, "") +} + +func writeSkill(t *testing.T, root, slug, name, description, body string) { + t.Helper() + dir := filepath.Join(root, slug) + if err := os.MkdirAll(dir, 0755); err != nil { + t.Fatal(err) + } + content := "---\nname: " + name + "\ndescription: " + description + "\n---\n\n" + body + "\n" + if err := os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte(content), 0644); err != nil { + t.Fatal(err) + } +} diff --git a/internal/agent/title_generate.go b/internal/agent/title_generate.go index eab223cf..67fcd302 100644 --- a/internal/agent/title_generate.go +++ b/internal/agent/title_generate.go @@ -7,6 +7,7 @@ import ( "time" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) const titleGenerateTimeout = 120 * time.Second @@ -16,10 +17,14 @@ const titleSystemPrompt = `Generate a short title (max 15 words) for this conver // GenerateTitle uses a lightweight LLM call to create a short conversation title // from the user's first message. Returns empty string on error. func GenerateTitle(ctx context.Context, provider providers.Provider, model, userMessage string) string { + return GenerateTitleWithUsageCaps(ctx, nil, provider, model, userMessage) +} + +func GenerateTitleWithUsageCaps(ctx context.Context, usageCaps *usagecaps.Service, provider providers.Provider, model, userMessage string) string { ctx, cancel := context.WithTimeout(ctx, titleGenerateTimeout) defer cancel() - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: titleSystemPrompt}, {Role: "user", Content: userMessage}, @@ -29,13 +34,18 @@ func GenerateTitle(ctx context.Context, provider providers.Provider, model, user // Larger budget: thinking-capable models (Gemini 2.5/3, GPT-5 reasoning) // can consume output tokens on reasoning traces. 256 leaves room for a // 15-word title even when the provider allocates some budget to thinking. - providers.OptMaxTokens: 256, - providers.OptTemperature: 0.3, + providers.OptMaxTokens: 256, + providers.OptTemperature: 0.3, // Disable extended thinking for title generation — it's a trivial task // that doesn't benefit from reasoning and defaults (esp. Gemini's "high") // otherwise eat the entire max_tokens budget, truncating the title to 1 word. providers.OptThinkingLevel: "off", }, + } + resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "session-title", + MaxOutputTokens: 256, }) if err != nil { slog.Warn("title generation failed", "error", err) diff --git a/internal/agent/usage_caps_runtime.go b/internal/agent/usage_caps_runtime.go new file mode 100644 index 00000000..e883a9c3 --- /dev/null +++ b/internal/agent/usage_caps_runtime.go @@ -0,0 +1,88 @@ +package agent + +import ( + "context" + "fmt" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" +) + +func (l *Loop) reserveInternalLLMUsage(ctx context.Context, chatReq providers.ChatRequest, purpose string) (*usagecaps.Reservation, error) { + if l.usageCaps == nil || l.provider == nil { + return nil, nil + } + return l.reserveInternalLLMUsageFor(ctx, chatReq, purpose, l.provider.Name(), chatReq.Model) +} + +func (l *Loop) reserveInternalLLMUsageFor(ctx context.Context, chatReq providers.ChatRequest, purpose, providerName, model string) (*usagecaps.Reservation, error) { + if l.usageCaps == nil { + return nil, nil + } + if model == "" { + model = chatReq.Model + } + if model == "" { + model = l.model + } + tenantID := store.TenantIDFromContext(ctx) + if tenantID == uuid.Nil { + tenantID = l.tenantID + } + return l.usageCaps.Preflight(ctx, usagecaps.Request{ + TenantID: tenantID, + AgentID: l.agentUUID, + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("%s:%s:%s", purpose, l.agentUUID.String(), uuid.NewString()), + Messages: chatReq.Messages, + MaxOutputTokens: l.maxOutputTokensFromRequest(chatReq), + }) +} + +func (l *Loop) callInternalLLMWithUsage(ctx context.Context, chatReq providers.ChatRequest, purpose string) (*providers.ChatResponse, error) { + if fallbackProvider, ok := l.provider.(*providers.ModelFallbackProvider); ok { + before := func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + candidatePurpose := fmt.Sprintf("%s:%s:%s", purpose, entry.ProviderName, actualReq.Model) + reservation, err := l.reserveInternalLLMUsageFor(callCtx, actualReq, candidatePurpose, entry.ProviderName, actualReq.Model) + if err != nil { + return nil, err + } + return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) { + if reservation != nil { + reservation.Reconcile(callCtx, resp, callErr) + } + }, nil + } + return fallbackProvider.ChatWithHook(ctx, chatReq, before) + } + reservation, reserveErr := l.reserveInternalLLMUsage(ctx, chatReq, purpose) + if reserveErr != nil { + return nil, reserveErr + } + resp, err := l.provider.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } + return resp, err +} + +func (l *Loop) maxOutputTokensFromRequest(chatReq providers.ChatRequest) int { + maxTokens := l.effectiveMaxTokens() + if chatReq.Options == nil { + return maxTokens + } + if v, ok := chatReq.Options[providers.OptMaxTokens]; ok { + switch n := v.(type) { + case int: + maxTokens = n + case int64: + maxTokens = int(n) + case float64: + maxTokens = int(n) + } + } + return maxTokens +} diff --git a/internal/channels/history.go b/internal/channels/history.go index 36692857..054d61a4 100644 --- a/internal/channels/history.go +++ b/internal/channels/history.go @@ -36,9 +36,11 @@ const ( // Stored in RAM only (not persisted to DB) — used by channels that defer media // download until the bot is actually mentioned (e.g. Telegram). type MediaRef struct { - Type string // "image", "video", "audio", "voice", "document", "animation" - FileID string // platform-specific file ID for lazy download - FileSize int64 // file size in bytes (0 if unknown) — used to skip large files + Type string // "image", "video", "audio", "voice", "document", "animation" + FileID string // platform-specific file ID for lazy download + FileSize int64 // file size in bytes (0 if unknown) — used to skip large files + FileName string // original filename, if the platform exposes one + ContentType string // MIME type, if the platform exposes one } // HistoryEntry represents a single tracked group message. diff --git a/internal/channels/history_compaction.go b/internal/channels/history_compaction.go index 6480a6ad..79b20624 100644 --- a/internal/channels/history_compaction.go +++ b/internal/channels/history_compaction.go @@ -11,6 +11,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // CompactionConfig configures LLM-based history compaction. @@ -20,6 +21,7 @@ type CompactionConfig struct { MaxTokens int // max output tokens for summarization (default 4096) Provider providers.Provider // LLM provider for summarization Model string // model to use for summarization + UsageCaps *usagecaps.Service } // MaybeCompact checks if compaction is needed for a history key and triggers it in background. @@ -55,7 +57,7 @@ func (ph *PendingHistory) MaybeCompact(historyKey string, currentCount int, cfg // CompactGroup performs LLM-based compaction on a pending message group. // Reused by both auto-compact (channel) and HTTP compact endpoint. // Returns the number of entries remaining after compaction. -func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, historyKey string, provider providers.Provider, model string, keepRecent, maxTokens int) (int, error) { +func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, historyKey string, provider providers.Provider, model string, keepRecent, maxTokens int, usageCaps *usagecaps.Service) (int, error) { if keepRecent <= 0 { keepRecent = 40 } @@ -94,13 +96,18 @@ func CompactGroup(ctx context.Context, s store.PendingMessageStore, channelName, fmt.Fprintf(&sb, "%s%s: %s\n", prefix, ts, e.Body) } - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{{ Role: "user", Content: "Summarize these group chat messages concisely, preserving key topics, decisions, names, and important context:\n\n" + sb.String(), }}, Model: model, Options: map[string]any{"max_tokens": maxTokens, "temperature": 0.3}, + } + resp, err := usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "pending-history-compaction", + MaxOutputTokens: maxTokens, }) if err != nil { return 0, fmt.Errorf("llm summarize: %w", err) @@ -185,7 +192,7 @@ func (ph *PendingHistory) runCompaction(historyKey string, cfg *CompactionConfig return } - _, err = CompactGroup(ctx, ph.store, ph.channelName, historyKey, cfg.Provider, cfg.Model, cfg.KeepRecent, cfg.MaxTokens) + _, err = CompactGroup(ctx, ph.store, ph.channelName, historyKey, cfg.Provider, cfg.Model, cfg.KeepRecent, cfg.MaxTokens, cfg.UsageCaps) if err != nil { slog.Warn("compaction.failed", "channel", ph.channelName, "key", historyKey, "error", err) return diff --git a/internal/channels/instance_loader.go b/internal/channels/instance_loader.go index df6d677f..7f845d80 100644 --- a/internal/channels/instance_loader.go +++ b/internal/channels/instance_loader.go @@ -16,6 +16,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providerresolve" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // reloadStartTimeout bounds how long Reload() will wait for a single channel's @@ -39,6 +40,7 @@ type InstanceLoader struct { agentStore store.AgentStore providerReg *providers.Registry pendingCompactCfg *config.PendingCompactionConfig + usageCaps *usagecaps.Service factories map[string]ChannelFactory manager *Manager msgBus *bus.MessageBus @@ -78,6 +80,10 @@ func (l *InstanceLoader) SetPendingCompactionConfig(cfg *config.PendingCompactio l.pendingCompactCfg = cfg } +func (l *InstanceLoader) SetUsageCapService(s *usagecaps.Service) { + l.usageCaps = s +} + // RegisterFactory registers a factory for a channel type (e.g., "telegram", "discord"). func (l *InstanceLoader) RegisterFactory(channelType string, factory ChannelFactory) { l.factories[channelType] = factory @@ -306,8 +312,9 @@ func (l *InstanceLoader) loadInstance(ctx context.Context, inst store.ChannelIns if p != nil && model != "" { cc := &CompactionConfig{ - Provider: p, - Model: model, + Provider: p, + Model: model, + UsageCaps: l.usageCaps, } if l.pendingCompactCfg != nil { cc.Threshold = l.pendingCompactCfg.Threshold diff --git a/internal/channels/telegram/handlers.go b/internal/channels/telegram/handlers.go index 7999103c..685d621f 100644 --- a/internal/channels/telegram/handlers.go +++ b/internal/channels/telegram/handlers.go @@ -513,14 +513,8 @@ func (c *Channel) handleMessage(ctx context.Context, update telego.Update) { // Resolve deferred media from history entries (lazy download). if histRefs := c.GroupHistory().CollectMediaRefs(localKey); len(histRefs) > 0 { histMedia, histErrors := c.resolveMediaRefs(ctx, histRefs) - for _, m := range histMedia { - mediaFiles = append(mediaFiles, bus.MediaFile{ - Path: m.FilePath, - MimeType: m.ContentType, - Filename: m.FileName, - }) - } if len(histMedia) > 0 { + mediaFiles = prependMediaInfoFiles(mediaFiles, histMedia) histTags := buildMediaTags(histMedia) annotated = "[Media from recent group messages — only analyze if user asks about them]\n" + histTags + "\n[/Media]\n\n" + annotated } @@ -577,12 +571,12 @@ func (c *Channel) handleMessage(ctx context.Context, update telego.Update) { // user sees typing indicator → first content appears directly. metadata := map[string]string{ - "message_id": fmt.Sprintf("%d", message.MessageID), - "user_id": fmt.Sprintf("%d", user.ID), + "message_id": fmt.Sprintf("%d", message.MessageID), + "user_id": fmt.Sprintf("%d", user.ID), tools.MetaUsername: user.Username, - "first_name": user.FirstName, - "is_group": fmt.Sprintf("%t", isGroup), - "local_key": localKey, + "first_name": user.FirstName, + "is_group": fmt.Sprintf("%t", isGroup), + "local_key": localKey, } if message.Chat.Title != "" { metadata[tools.MetaChatTitle] = message.Chat.Title diff --git a/internal/channels/telegram/media.go b/internal/channels/telegram/media.go index a7d92d02..1a1faa8c 100644 --- a/internal/channels/telegram/media.go +++ b/internal/channels/telegram/media.go @@ -14,6 +14,7 @@ import ( "github.com/mymmrac/telego" + "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/channels" "github.com/nextlevelbuilder/goclaw/internal/channels/media" "github.com/nextlevelbuilder/goclaw/internal/tools" @@ -359,6 +360,24 @@ func buildMediaTags(mediaList []MediaInfo) string { return media.BuildMediaTags(mediaList) } +func prependMediaInfoFiles(current []bus.MediaFile, history []MediaInfo) []bus.MediaFile { + if len(history) == 0 { + return current + } + ordered := make([]bus.MediaFile, 0, len(history)+len(current)) + for _, m := range history { + if m.FilePath == "" { + continue + } + ordered = append(ordered, bus.MediaFile{ + Path: m.FilePath, + MimeType: m.ContentType, + Filename: m.FileName, + }) + } + return append(ordered, current...) +} + // extractDocumentContent delegates to the shared media package. func extractDocumentContent(filePath, fileName string) (string, error) { return media.ExtractDocumentContent(filePath, fileName) @@ -412,22 +431,51 @@ func extractMediaRefs(msg *telego.Message) []channels.MediaRef { refs = append(refs, channels.MediaRef{Type: "image", FileID: photo.FileID, FileSize: int64(photo.FileSize)}) } if msg.Video != nil { - refs = append(refs, channels.MediaRef{Type: "video", FileID: msg.Video.FileID, FileSize: int64(msg.Video.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "video", + FileID: msg.Video.FileID, + FileSize: int64(msg.Video.FileSize), + FileName: msg.Video.FileName, + ContentType: msg.Video.MimeType, + }) } if msg.VideoNote != nil { refs = append(refs, channels.MediaRef{Type: "video", FileID: msg.VideoNote.FileID, FileSize: int64(msg.VideoNote.FileSize)}) } if msg.Animation != nil { - refs = append(refs, channels.MediaRef{Type: "animation", FileID: msg.Animation.FileID, FileSize: int64(msg.Animation.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "animation", + FileID: msg.Animation.FileID, + FileSize: int64(msg.Animation.FileSize), + FileName: msg.Animation.FileName, + ContentType: msg.Animation.MimeType, + }) } if msg.Audio != nil { - refs = append(refs, channels.MediaRef{Type: "audio", FileID: msg.Audio.FileID, FileSize: int64(msg.Audio.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "audio", + FileID: msg.Audio.FileID, + FileSize: int64(msg.Audio.FileSize), + FileName: msg.Audio.FileName, + ContentType: msg.Audio.MimeType, + }) } if msg.Voice != nil { - refs = append(refs, channels.MediaRef{Type: "voice", FileID: msg.Voice.FileID, FileSize: int64(msg.Voice.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "voice", + FileID: msg.Voice.FileID, + FileSize: int64(msg.Voice.FileSize), + ContentType: msg.Voice.MimeType, + }) } if msg.Document != nil { - refs = append(refs, channels.MediaRef{Type: "document", FileID: msg.Document.FileID, FileSize: int64(msg.Document.FileSize)}) + refs = append(refs, channels.MediaRef{ + Type: "document", + FileID: msg.Document.FileID, + FileSize: int64(msg.Document.FileSize), + FileName: msg.Document.FileName, + ContentType: msg.Document.MimeType, + }) } return refs } @@ -492,9 +540,12 @@ func (c *Channel) resolveMediaRefs(ctx context.Context, refs []channels.MediaRef continue } results = append(results, MediaInfo{ - Type: ref.Type, - FilePath: filePath, - FileID: ref.FileID, + Type: ref.Type, + FilePath: filePath, + FileID: ref.FileID, + ContentType: ref.ContentType, + FileName: ref.FileName, + FileSize: ref.FileSize, }) } return results, errs diff --git a/internal/channels/telegram/media_test.go b/internal/channels/telegram/media_test.go index f568db34..a9822c70 100644 --- a/internal/channels/telegram/media_test.go +++ b/internal/channels/telegram/media_test.go @@ -3,6 +3,11 @@ package telegram import ( "strings" "testing" + + "github.com/mymmrac/telego" + + "github.com/nextlevelbuilder/goclaw/internal/bus" + "github.com/nextlevelbuilder/goclaw/internal/channels" ) // --- buildMediaTags tests --- @@ -148,3 +153,52 @@ func TestBuildMediaTags_UnknownType(t *testing.T) { t.Errorf("expected empty string for unknown type, got: %q", got) } } + +func TestExtractMediaRefsPreservesDocumentMetadata(t *testing.T) { + msg := &telego.Message{ + Document: &telego.Document{ + FileID: "telegram-file-id", + FileName: "codex.zip", + MimeType: "application/zip", + FileSize: 1234, + }, + } + + got := extractMediaRefs(msg) + want := []channels.MediaRef{{ + Type: "document", + FileID: "telegram-file-id", + FileSize: 1234, + FileName: "codex.zip", + ContentType: "application/zip", + }} + + if len(got) != len(want) { + t.Fatalf("got %d refs, want %d", len(got), len(want)) + } + if got[0] != want[0] { + t.Fatalf("ref = %+v, want %+v", got[0], want[0]) + } +} + +func TestPrependMediaInfoFilesKeepsHistoryBeforeCurrent(t *testing.T) { + current := []bus.MediaFile{{ + Path: "/workspace/.uploads/current.pdf", + MimeType: "application/pdf", + Filename: "current.pdf", + }} + history := []MediaInfo{{ + Type: "document", + FilePath: "/workspace/.uploads/codex.zip", + ContentType: "application/zip", + FileName: "codex.zip", + }} + + got := prependMediaInfoFiles(current, history) + if len(got) != 2 { + t.Fatalf("got %d files, want 2", len(got)) + } + if got[0].Path != history[0].FilePath || got[1].Path != current[0].Path { + t.Fatalf("order = [%q, %q], want history before current", got[0].Path, got[1].Path) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index bdc2cd36..86d6f0d3 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "log/slog" + "strings" "sync" "time" @@ -47,6 +48,7 @@ type Config struct { Providers ProvidersConfig `json:"providers"` Gateway GatewayConfig `json:"gateway"` Tools ToolsConfig `json:"tools"` + Skills SkillsConfig `json:"skills"` Sessions SessionsConfig `json:"sessions"` Database DatabaseConfig `json:"database"` Tts TtsConfig `json:"tts"` @@ -129,7 +131,82 @@ type DatabaseConfig struct { // SkillsConfig configures the skills storage system. type SkillsConfig struct { - StorageDir string `json:"storage_dir,omitempty"` // directory for skill content (default: dataDir/skills-store/) + StorageDir string `json:"storage_dir,omitempty"` // directory for skill content (default: dataDir/skills-store/) + MaxUploadSizeMB int `json:"max_upload_size_mb,omitempty"` // per-file ZIP upload limit + SlashCommands SkillSlashCommandConfig `json:"slash_commands,omitempty"` +} + +// SkillSlashCommandConfig controls explicit slash-command skill activation. +type SkillSlashCommandConfig struct { + Enabled *bool `json:"enabled,omitempty"` + SuggestNotFound *bool `json:"suggest_not_found,omitempty"` + PartialMatching bool `json:"partial_matching,omitempty"` + Prefix string `json:"prefix,omitempty"` +} + +const ( + DefaultSkillMaxUploadSizeMB = 20 + MinSkillMaxUploadSizeMB = 1 + MaxSkillMaxUploadSizeMB = 500 + + DefaultSkillSlashCommandPrefix = "/" + + SkillMaxUploadSizeSystemConfigKey = "skills.max_upload_size_mb" + SkillSlashCommandsEnabledSystemConfigKey = "skills.slash_commands.enabled" + SkillSlashSuggestNotFoundSystemConfigKey = "skills.slash_commands.suggest_not_found" + SkillSlashPartialMatchingSystemConfigKey = "skills.slash_commands.partial_matching" + SkillSlashCommandPrefixSystemConfigKey = "skills.slash_commands.prefix" +) + +func ClampSkillMaxUploadSizeMB(value int) int { + if value == 0 { + return DefaultSkillMaxUploadSizeMB + } + if value < MinSkillMaxUploadSizeMB { + return MinSkillMaxUploadSizeMB + } + if value > MaxSkillMaxUploadSizeMB { + return MaxSkillMaxUploadSizeMB + } + return value +} + +func (c SkillsConfig) EffectiveMaxUploadSizeMB() int { + return ClampSkillMaxUploadSizeMB(c.MaxUploadSizeMB) +} + +func (c SkillsConfig) EffectiveMaxUploadSizeBytes() int64 { + return int64(c.EffectiveMaxUploadSizeMB()) << 20 +} + +func (c SkillSlashCommandConfig) EffectiveEnabled() bool { + if c.Enabled == nil { + return true + } + return *c.Enabled +} + +func (c SkillSlashCommandConfig) EffectiveSuggestNotFound() bool { + if c.SuggestNotFound == nil { + return true + } + return *c.SuggestNotFound +} + +func (c SkillSlashCommandConfig) EffectivePartialMatching() bool { + return c.PartialMatching +} + +func (c SkillSlashCommandConfig) EffectivePrefix() string { + prefix := strings.TrimSpace(c.Prefix) + if prefix == "" { + return DefaultSkillSlashCommandPrefix + } + runes := []rune(prefix) + if len(runes) != 1 || runes[0] == ' ' || runes[0] == '\t' || runes[0] == '\n' || runes[0] == '\r' { + return DefaultSkillSlashCommandPrefix + } + return prefix } // AgentBinding maps a channel/peer pattern to a specific agent. @@ -479,6 +556,7 @@ func (c *Config) ReplaceFrom(src *Config) { c.Providers = src.Providers c.Gateway = src.Gateway c.Tools = src.Tools + c.Skills = src.Skills c.Sessions = src.Sessions c.Database = src.Database c.Tts = src.Tts diff --git a/internal/config/config_channels.go b/internal/config/config_channels.go index a125d50e..6a9c5c70 100644 --- a/internal/config/config_channels.go +++ b/internal/config/config_channels.go @@ -24,25 +24,25 @@ type ChannelsConfig struct { } type TelegramConfig struct { - Enabled bool `json:"enabled"` - Token string `json:"token"` - Proxy string `json:"proxy,omitempty"` - APIServer string `json:"api_server,omitempty"` // custom Telegram Bot API server URL (e.g. "http://localhost:8081") - AllowFrom FlexibleStringSlice `json:"allow_from"` - DMPolicy string `json:"dm_policy,omitempty"` // "pairing" (default), "allowlist", "open", "disabled" - GroupPolicy string `json:"group_policy,omitempty"` // "open" (default), "allowlist", "disabled" - RequireMention *bool `json:"require_mention,omitempty"` // require @bot mention in groups (default true) - MentionMode string `json:"mention_mode,omitempty"` // "strict" (default) = only respond when mentioned; "yield" = respond unless another bot is mentioned - HistoryLimit int `json:"history_limit,omitempty"` // max pending group messages for context (default 50, 0=disabled) - DMStream *bool `json:"dm_stream,omitempty"` // enable streaming for DMs (default false) — edits placeholder progressively - GroupStream *bool `json:"group_stream,omitempty"` // enable streaming for groups (default false) — sends new message, edits progressively - DraftTransport *bool `json:"draft_transport,omitempty"` // use sendMessageDraft for DM streaming (default true) — stealth preview, no notifications per edit - ReasoningStream *bool `json:"reasoning_stream,omitempty"` // show reasoning as separate message when provider emits thinking events (default true) - ReactionLevel string `json:"reaction_level,omitempty"` // "off" (default), "minimal", "full" — status emoji reactions - MediaMaxBytes int64 `json:"media_max_bytes,omitempty"` // max media download size in bytes (default 20MB) - LinkPreview *bool `json:"link_preview,omitempty"` // enable URL previews in messages (default true) - BlockReply *bool `json:"block_reply,omitempty"` // override gateway block_reply (nil = inherit) - ForceIPv4 bool `json:"force_ipv4,omitempty"` // force IPv4 for all Telegram API requests (use when IPv6 routing is broken) + Enabled bool `json:"enabled"` + Token string `json:"token"` + Proxy string `json:"proxy,omitempty"` + APIServer string `json:"api_server,omitempty"` // custom Telegram Bot API server URL (e.g. "http://localhost:8081") + AllowFrom FlexibleStringSlice `json:"allow_from"` + DMPolicy string `json:"dm_policy,omitempty"` // "pairing" (default), "allowlist", "open", "disabled" + GroupPolicy string `json:"group_policy,omitempty"` // "open" (default), "allowlist", "disabled" + RequireMention *bool `json:"require_mention,omitempty"` // require @bot mention in groups (default true) + MentionMode string `json:"mention_mode,omitempty"` // "strict" (default) = only respond when mentioned; "yield" = respond unless another bot is mentioned + HistoryLimit int `json:"history_limit,omitempty"` // max pending group messages for context (default 50, 0=disabled) + DMStream *bool `json:"dm_stream,omitempty"` // enable streaming for DMs (default false) — edits placeholder progressively + GroupStream *bool `json:"group_stream,omitempty"` // enable streaming for groups (default false) — sends new message, edits progressively + DraftTransport *bool `json:"draft_transport,omitempty"` // use sendMessageDraft for DM streaming (default true) — stealth preview, no notifications per edit + ReasoningStream *bool `json:"reasoning_stream,omitempty"` // show reasoning as separate message when provider emits thinking events (default true) + ReactionLevel string `json:"reaction_level,omitempty"` // "off" (default), "minimal", "full" — status emoji reactions + MediaMaxBytes int64 `json:"media_max_bytes,omitempty"` // max media download size in bytes (default 20MB) + LinkPreview *bool `json:"link_preview,omitempty"` // enable URL previews in messages (default true) + BlockReply *bool `json:"block_reply,omitempty"` // override gateway block_reply (nil = inherit) + ForceIPv4 bool `json:"force_ipv4,omitempty"` // force IPv4 for all Telegram API requests (use when IPv6 routing is broken) // Optional STT (Speech-to-Text) pipeline for voice/audio inbound messages. // When stt_proxy_url is set, audio/voice messages are transcribed before being forwarded to the agent. @@ -133,7 +133,7 @@ type SlackConfig struct { type WhatsAppConfig struct { Enabled bool `json:"enabled"` - AuthDir string `json:"auth_dir,omitempty"` // optional: SQLite auth dir override (desktop) + AuthDir string `json:"auth_dir,omitempty"` // optional: SQLite auth dir override (desktop) AllowFrom FlexibleStringSlice `json:"allow_from"` DMPolicy string `json:"dm_policy,omitempty"` // "pairing" (default for DB instances), "open", "allowlist", "disabled" GroupPolicy string `json:"group_policy,omitempty"` // "pairing" (default for DB instances), "open" (default for config), "allowlist", "disabled" @@ -196,25 +196,25 @@ type FeishuConfig struct { // ProvidersConfig maps provider name to its config. type ProvidersConfig struct { - Anthropic ProviderConfig `json:"anthropic"` - OpenAI ProviderConfig `json:"openai"` - OpenRouter ProviderConfig `json:"openrouter"` - Groq ProviderConfig `json:"groq"` - Gemini ProviderConfig `json:"gemini"` - DeepSeek ProviderConfig `json:"deepseek"` - Mistral ProviderConfig `json:"mistral"` - XAI ProviderConfig `json:"xai"` - MiniMax ProviderConfig `json:"minimax"` - Cohere ProviderConfig `json:"cohere"` - Perplexity ProviderConfig `json:"perplexity"` - DashScope ProviderConfig `json:"dashscope"` - Bailian ProviderConfig `json:"bailian"` - Zai ProviderConfig `json:"zai"` - ZaiCoding ProviderConfig `json:"zai_coding"` - Ollama OllamaConfig `json:"ollama"` // local Ollama instance (no API key needed) - OllamaCloud ProviderConfig `json:"ollama_cloud"` // Ollama Cloud (API key required) - ClaudeCLI ClaudeCLIConfig `json:"claude_cli"` - ACP ACPConfig `json:"acp"` + Anthropic ProviderConfig `json:"anthropic"` + OpenAI ProviderConfig `json:"openai"` + OpenRouter ProviderConfig `json:"openrouter"` + Groq ProviderConfig `json:"groq"` + Gemini ProviderConfig `json:"gemini"` + DeepSeek ProviderConfig `json:"deepseek"` + Mistral ProviderConfig `json:"mistral"` + XAI ProviderConfig `json:"xai"` + MiniMax ProviderConfig `json:"minimax"` + Cohere ProviderConfig `json:"cohere"` + Perplexity ProviderConfig `json:"perplexity"` + DashScope ProviderConfig `json:"dashscope"` + Bailian ProviderConfig `json:"bailian"` + Zai ProviderConfig `json:"zai"` + ZaiCoding ProviderConfig `json:"zai_coding"` + Ollama OllamaConfig `json:"ollama"` // local Ollama instance (no API key needed) + OllamaCloud ProviderConfig `json:"ollama_cloud"` // Ollama Cloud (API key required) + ClaudeCLI ClaudeCLIConfig `json:"claude_cli"` + ACP ACPConfig `json:"acp"` Novita ProviderConfig `json:"novita"` // Novita AI (OpenAI-compatible endpoint) BytePlus ProviderConfig `json:"byteplus"` // BytePlus ModelArk (Seed 2.0) BytePlusCoding ProviderConfig `json:"byteplus_coding"` // BytePlus ModelArk Coding Plan @@ -362,16 +362,16 @@ type QuotaConfig struct { // GatewayConfig controls the gateway server. type GatewayConfig struct { - Host string `json:"host"` - Port int `json:"port"` - Token string `json:"token,omitempty"` // bearer token for WS/HTTP auth - OwnerIDs []string `json:"owner_ids,omitempty"` // sender IDs considered "owner" - AllowedOrigins []string `json:"allowed_origins,omitempty"` // WebSocket CORS whitelist (empty = allow all) - MaxMessageChars int `json:"max_message_chars,omitempty"` // max user message characters (default 32000) - RateLimitRPM int `json:"rate_limit_rpm,omitempty"` // rate limit: requests per minute per user (default 20, 0 = disabled) - InjectionAction string `json:"injection_action,omitempty"` // prompt injection action: "log", "warn" (default), "block", "off" - InboundDebounceMs int `json:"inbound_debounce_ms,omitempty"` // merge rapid channel/Web Chat messages from same sender/session (0 = no wait) - Quota *QuotaConfig `json:"quota,omitempty"` // per-user/group request quotas + Host string `json:"host"` + Port int `json:"port"` + Token string `json:"token,omitempty"` // bearer token for WS/HTTP auth + OwnerIDs []string `json:"owner_ids,omitempty"` // sender IDs considered "owner" + AllowedOrigins []string `json:"allowed_origins,omitempty"` // WebSocket CORS whitelist (empty = allow all) + MaxMessageChars int `json:"max_message_chars,omitempty"` // max user message characters (default 32000) + RateLimitRPM int `json:"rate_limit_rpm,omitempty"` // rate limit: requests per minute per user (default 20, 0 = disabled) + InjectionAction string `json:"injection_action,omitempty"` // prompt injection action: "log", "warn" (default), "block", "off" + InboundDebounceMs int `json:"inbound_debounce_ms,omitempty"` // merge rapid channel/Web Chat messages from same sender/session (0 = no wait) + Quota *QuotaConfig `json:"quota,omitempty"` // per-user/group request quotas BlockReply *bool `json:"block_reply,omitempty"` // deliver intermediate text during tool iterations (default false) ToolStatus *bool `json:"tool_status,omitempty"` // show tool name in streaming preview during tool execution (default true) TaskRecoveryIntervalSec int `json:"task_recovery_interval_sec,omitempty"` // team task recovery ticker interval in seconds (default 300 = 5min) @@ -381,18 +381,32 @@ type GatewayConfig struct { // ToolsConfig controls tool availability, policy, and web search. type ToolsConfig struct { - Profile string `json:"profile,omitempty"` // global profile: "minimal", "coding", "messaging", "full" - Allow []string `json:"allow,omitempty"` // global allow list (tool names or "group:xxx") - Deny []string `json:"deny,omitempty"` // global deny list - AlsoAllow []string `json:"alsoAllow,omitempty"` // additive: adds without removing existing - ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` // per-provider overrides - ShellDenyGroups map[string]bool `json:"shellDenyGroups,omitempty"` // global shell deny-group toggles (group name -> denied); per-agent overrides win per-key - ExecApproval ExecApprovalCfg `json:"execApproval"` // exec command approval settings - WebFetch WebFetchPolicyConfig `json:"web_fetch"` // domain policy for URL fetching - Browser BrowserToolConfig `json:"browser"` - RateLimitPerHour int `json:"rate_limit_per_hour,omitempty"` // max tool executions per hour per session (0 = disabled) - ScrubCredentials *bool `json:"scrub_credentials,omitempty"` // auto-redact API keys/tokens in tool output (default true) - McpServers map[string]*MCPServerConfig `json:"mcp_servers,omitempty"` // external MCP server connections + Profile string `json:"profile,omitempty"` // global profile: "minimal", "coding", "messaging", "full" + Allow []string `json:"allow,omitempty"` // global allow list (tool names or "group:xxx") + Deny []string `json:"deny,omitempty"` // global deny list + AlsoAllow []string `json:"alsoAllow,omitempty"` // additive: adds without removing existing + ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` // per-provider overrides + ShellDenyGroups map[string]bool `json:"shellDenyGroups,omitempty"` // global shell deny-group toggles (group name -> denied); per-agent overrides win per-key + CommandKeywordAllowlist []CommandKeywordAllowlistRule `json:"commandKeywordAllowlist,omitempty"` // scoped bypass for credentialed CLI content keywords + ExecApproval ExecApprovalCfg `json:"execApproval"` // exec command approval settings + WebFetch WebFetchPolicyConfig `json:"web_fetch"` // domain policy for URL fetching + Browser BrowserToolConfig `json:"browser"` + RateLimitPerHour int `json:"rate_limit_per_hour,omitempty"` // max tool executions per hour per session (0 = disabled) + ScrubCredentials *bool `json:"scrub_credentials,omitempty"` // auto-redact API keys/tokens in tool output (default true) + McpServers map[string]*MCPServerConfig `json:"mcp_servers,omitempty"` // external MCP server connections +} + +// CommandKeywordAllowlistRule scopes product/security vocabulary that may appear +// in credentialed CLI content args without disabling command-path deny patterns. +type CommandKeywordAllowlistRule struct { + ID string `json:"id,omitempty"` + Command string `json:"command"` + Subcommands []string `json:"subcommands,omitempty"` + Args []string `json:"args,omitempty"` + ArgPositions []int `json:"argPositions,omitempty"` // 0-based after matched subcommand + Keywords []string `json:"keywords,omitempty"` + Reason string `json:"reason,omitempty"` + Enabled *bool `json:"enabled,omitempty"` // nil means enabled } // MCPServerConfig configures a single external MCP server connection. @@ -429,24 +443,24 @@ type WebFetchPolicyConfig struct { // BrowserToolConfig controls the browser automation tool. type BrowserToolConfig struct { - Enabled bool `json:"enabled"` // enable the browser tool (default false) - Headless bool `json:"headless,omitempty"` // run Chrome in headless mode (ignored when RemoteURL is set) - RemoteURL string `json:"remote_url,omitempty"` // CDP endpoint for remote Chrome sidecar, e.g. "ws://chrome:9222" - ActionTimeoutMs int `json:"action_timeout_ms,omitempty"` // per-action timeout in ms (default 30000) - IdleTimeoutMs int `json:"idle_timeout_ms,omitempty"` // idle page auto-close in ms (default 600000, 0=disabled) - MaxPages int `json:"max_pages,omitempty"` // max open pages per tenant (default 5) - CookieSyncEnabled bool `json:"cookie_sync_enabled"` // apply selected synced cookies to scoped browser sessions + Enabled bool `json:"enabled"` // enable the browser tool (default false) + Headless bool `json:"headless,omitempty"` // run Chrome in headless mode (ignored when RemoteURL is set) + RemoteURL string `json:"remote_url,omitempty"` // CDP endpoint for remote Chrome sidecar, e.g. "ws://chrome:9222" + ActionTimeoutMs int `json:"action_timeout_ms,omitempty"` // per-action timeout in ms (default 30000) + IdleTimeoutMs int `json:"idle_timeout_ms,omitempty"` // idle page auto-close in ms (default 600000, 0=disabled) + MaxPages int `json:"max_pages,omitempty"` // max open pages per tenant (default 5) + CookieSyncEnabled bool `json:"cookie_sync_enabled"` // apply selected synced cookies to scoped browser sessions } // ToolPolicySpec defines a tool policy at any level (global, per-agent, per-provider). type ToolPolicySpec struct { - Profile string `json:"profile,omitempty"` - Allow []string `json:"allow,omitempty"` - Deny []string `json:"deny,omitempty"` - AlsoAllow []string `json:"alsoAllow,omitempty"` - ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` - Wait *WaitToolPolicy `json:"wait,omitempty"` - ToolCallPrefix string `json:"toolCallPrefix,omitempty"` // prefix to strip from model's tool call names before registry lookup + Profile string `json:"profile,omitempty"` + Allow []string `json:"allow,omitempty"` + Deny []string `json:"deny,omitempty"` + AlsoAllow []string `json:"alsoAllow,omitempty"` + ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` + Wait *WaitToolPolicy `json:"wait,omitempty"` + ToolCallPrefix string `json:"toolCallPrefix,omitempty"` // prefix to strip from model's tool call names before registry lookup } // WaitToolPolicy configures per-agent safety bounds for the wait tool. diff --git a/internal/config/config_load.go b/internal/config/config_load.go index 205dd511..e215bad3 100644 --- a/internal/config/config_load.go +++ b/internal/config/config_load.go @@ -63,6 +63,15 @@ func isLoopbackGatewayHost(host string) bool { return err == nil && addr.IsLoopback() } +func parseEnvBool(value string) bool { + switch strings.ToLower(strings.TrimSpace(value)) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + // Default returns a Config with sensible defaults. func Default() *Config { return &Config{ @@ -107,6 +116,9 @@ func Default() *Config { }, RateLimitPerHour: 150, }, + Skills: SkillsConfig{ + MaxUploadSizeMB: DefaultSkillMaxUploadSizeMB, + }, Sessions: SessionsConfig{}, } } @@ -234,6 +246,26 @@ func (c *Config) applyEnvOverrides() { c.Gateway.Port = port } } + if v := os.Getenv("GOCLAW_SKILLS_MAX_UPLOAD_SIZE_MB"); v != "" { + if mb, err := strconv.Atoi(v); err == nil { + c.Skills.MaxUploadSizeMB = ClampSkillMaxUploadSizeMB(mb) + } + } + envBoolPtr := func(key string, dst **bool) { + if v := os.Getenv(key); v != "" { + b := parseEnvBool(v) + *dst = &b + } + } + envBool := func(key string, dst *bool) { + if v := os.Getenv(key); v != "" { + *dst = parseEnvBool(v) + } + } + envBoolPtr("GOCLAW_SKILLS_SLASH_COMMANDS_ENABLED", &c.Skills.SlashCommands.Enabled) + envBoolPtr("GOCLAW_SKILLS_SLASH_COMMANDS_SUGGEST_NOT_FOUND", &c.Skills.SlashCommands.SuggestNotFound) + envBool("GOCLAW_SKILLS_SLASH_COMMANDS_PARTIAL_MATCHING", &c.Skills.SlashCommands.PartialMatching) + envStr("GOCLAW_SKILLS_SLASH_COMMANDS_PREFIX", &c.Skills.SlashCommands.Prefix) // Database envStr("GOCLAW_POSTGRES_DSN", &c.Database.PostgresDSN) diff --git a/internal/config/config_load_test.go b/internal/config/config_load_test.go index 14923e73..912c04bc 100644 --- a/internal/config/config_load_test.go +++ b/internal/config/config_load_test.go @@ -24,6 +24,21 @@ func TestDefault_SensibleDefaults(t *testing.T) { if cfg.Agents.Defaults.MaxToolIterations != DefaultMaxIterations { t.Fatalf("default max iterations: got %d", cfg.Agents.Defaults.MaxToolIterations) } + if cfg.Skills.EffectiveMaxUploadSizeMB() != DefaultSkillMaxUploadSizeMB { + t.Fatalf("default skill upload max: got %d, want %d", cfg.Skills.EffectiveMaxUploadSizeMB(), DefaultSkillMaxUploadSizeMB) + } + if !cfg.Skills.SlashCommands.EffectiveEnabled() { + t.Fatal("slash commands should default enabled") + } + if !cfg.Skills.SlashCommands.EffectiveSuggestNotFound() { + t.Fatal("slash command suggestions should default enabled") + } + if cfg.Skills.SlashCommands.EffectivePartialMatching() { + t.Fatal("slash command partial matching should default disabled") + } + if cfg.Skills.SlashCommands.EffectivePrefix() != "/" { + t.Fatalf("slash command prefix = %q, want /", cfg.Skills.SlashCommands.EffectivePrefix()) + } } @@ -103,6 +118,124 @@ func TestLoad_EnvVarOverrides(t *testing.T) { } } +func TestLoad_SkillsMaxUploadSizeFromFileAndEnv(t *testing.T) { + dir := t.TempDir() + cfgPath := filepath.Join(dir, "config.json5") + os.WriteFile(cfgPath, []byte(`{"skills":{"max_upload_size_mb":64}}`), 0644) + + cfg, err := Load(cfgPath) + if err != nil { + t.Fatalf("load error: %v", err) + } + if cfg.Skills.EffectiveMaxUploadSizeMB() != 64 { + t.Fatalf("file skill upload max: got %d, want 64", cfg.Skills.EffectiveMaxUploadSizeMB()) + } + + t.Setenv("GOCLAW_SKILLS_MAX_UPLOAD_SIZE_MB", "128") + cfg, err = Load(cfgPath) + if err != nil { + t.Fatalf("load with env error: %v", err) + } + if cfg.Skills.EffectiveMaxUploadSizeMB() != 128 { + t.Fatalf("env skill upload max: got %d, want 128", cfg.Skills.EffectiveMaxUploadSizeMB()) + } +} + +func TestLoad_SkillSlashCommandsFromFileEnvAndSystemConfig(t *testing.T) { + dir := t.TempDir() + cfgPath := filepath.Join(dir, "config.json5") + os.WriteFile(cfgPath, []byte(`{ + "skills": { + "slash_commands": { + "enabled": false, + "suggest_not_found": false, + "partial_matching": true, + "prefix": "!" + } + } + }`), 0644) + + cfg, err := Load(cfgPath) + if err != nil { + t.Fatalf("load error: %v", err) + } + if cfg.Skills.SlashCommands.EffectiveEnabled() { + t.Fatal("file enabled override should be false") + } + if cfg.Skills.SlashCommands.EffectiveSuggestNotFound() { + t.Fatal("file suggestion override should be false") + } + if !cfg.Skills.SlashCommands.EffectivePartialMatching() { + t.Fatal("file partial matching override should be true") + } + if cfg.Skills.SlashCommands.EffectivePrefix() != "!" { + t.Fatalf("file prefix = %q, want !", cfg.Skills.SlashCommands.EffectivePrefix()) + } + + t.Setenv("GOCLAW_SKILLS_SLASH_COMMANDS_ENABLED", "true") + t.Setenv("GOCLAW_SKILLS_SLASH_COMMANDS_SUGGEST_NOT_FOUND", "true") + t.Setenv("GOCLAW_SKILLS_SLASH_COMMANDS_PARTIAL_MATCHING", "false") + t.Setenv("GOCLAW_SKILLS_SLASH_COMMANDS_PREFIX", "#") + cfg, err = Load(cfgPath) + if err != nil { + t.Fatalf("load with env error: %v", err) + } + if !cfg.Skills.SlashCommands.EffectiveEnabled() { + t.Fatal("env enabled override should be true") + } + if !cfg.Skills.SlashCommands.EffectiveSuggestNotFound() { + t.Fatal("env suggestion override should be true") + } + if cfg.Skills.SlashCommands.EffectivePartialMatching() { + t.Fatal("env partial matching override should be false") + } + if cfg.Skills.SlashCommands.EffectivePrefix() != "#" { + t.Fatalf("env prefix = %q, want #", cfg.Skills.SlashCommands.EffectivePrefix()) + } + + cfg.ApplySystemConfigs(map[string]string{ + "skills.slash_commands.enabled": "false", + "skills.slash_commands.suggest_not_found": "false", + "skills.slash_commands.partial_matching": "true", + "skills.slash_commands.prefix": "%", + }) + if cfg.Skills.SlashCommands.EffectiveEnabled() { + t.Fatal("system enabled override should be false") + } + if cfg.Skills.SlashCommands.EffectiveSuggestNotFound() { + t.Fatal("system suggestion override should be false") + } + if !cfg.Skills.SlashCommands.EffectivePartialMatching() { + t.Fatal("system partial matching override should be true") + } + if cfg.Skills.SlashCommands.EffectivePrefix() != "%" { + t.Fatalf("system prefix = %q, want %%", cfg.Skills.SlashCommands.EffectivePrefix()) + } +} + +func TestSkillsMaxUploadSizeClampAndSystemConfigOverlay(t *testing.T) { + cfg := Default() + cfg.Skills.MaxUploadSizeMB = 0 + if got := cfg.Skills.EffectiveMaxUploadSizeMB(); got != DefaultSkillMaxUploadSizeMB { + t.Fatalf("zero upload max: got %d, want %d", got, DefaultSkillMaxUploadSizeMB) + } + + cfg.Skills.MaxUploadSizeMB = -10 + if got := cfg.Skills.EffectiveMaxUploadSizeMB(); got != MinSkillMaxUploadSizeMB { + t.Fatalf("negative upload max: got %d, want %d", got, MinSkillMaxUploadSizeMB) + } + + cfg.Skills.MaxUploadSizeMB = 999 + if got := cfg.Skills.EffectiveMaxUploadSizeMB(); got != MaxSkillMaxUploadSizeMB { + t.Fatalf("high upload max: got %d, want %d", got, MaxSkillMaxUploadSizeMB) + } + + cfg.ApplySystemConfigs(map[string]string{"skills.max_upload_size_mb": "77"}) + if got := cfg.Skills.EffectiveMaxUploadSizeMB(); got != 77 { + t.Fatalf("system config upload max: got %d, want 77", got) + } +} + func TestLoad_EnvVarOverrides_InvalidPort(t *testing.T) { t.Setenv("GOCLAW_PORT", "not-a-number") diff --git a/internal/config/config_system.go b/internal/config/config_system.go index 4789327c..3f046809 100644 --- a/internal/config/config_system.go +++ b/internal/config/config_system.go @@ -74,6 +74,14 @@ func (c *Config) ApplySystemConfigs(configs map[string]string) { integer("tools.browser.max_pages", &c.Tools.Browser.MaxPages) boolValue("tools.browser.cookie_sync_enabled", &c.Tools.Browser.CookieSyncEnabled) + // Skills + integer(SkillMaxUploadSizeSystemConfigKey, &c.Skills.MaxUploadSizeMB) + c.Skills.MaxUploadSizeMB = ClampSkillMaxUploadSizeMB(c.Skills.MaxUploadSizeMB) + boolean(SkillSlashCommandsEnabledSystemConfigKey, &c.Skills.SlashCommands.Enabled) + boolean(SkillSlashSuggestNotFoundSystemConfigKey, &c.Skills.SlashCommands.SuggestNotFound) + boolValue(SkillSlashPartialMatchingSystemConfigKey, &c.Skills.SlashCommands.PartialMatching) + str(SkillSlashCommandPrefixSystemConfigKey, &c.Skills.SlashCommands.Prefix) + // TTS str("tts.provider", &c.Tts.Provider) str("tts.auto", &c.Tts.Auto) diff --git a/internal/consolidation/dreaming_worker.go b/internal/consolidation/dreaming_worker.go index 4a97bf76..4a981793 100644 --- a/internal/consolidation/dreaming_worker.go +++ b/internal/consolidation/dreaming_worker.go @@ -11,9 +11,10 @@ import ( "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/bgalert" "github.com/nextlevelbuilder/goclaw/internal/eventbus" - "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/providerresolve" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) const ( @@ -31,6 +32,7 @@ type dreamingWorker struct { systemConfigs store.SystemConfigStore // per-tenant provider config registry *providers.Registry // provider resolution alertDeps bgalert.AlertDeps + usageCaps *usagecaps.Service // threshold/debounce are the global defaults. Per-agent overrides come // from resolveConfig which reads the agent's MemoryConfig.Dreaming JSONB. @@ -95,6 +97,11 @@ func (w *dreamingWorker) Handle(ctx context.Context, event eventbus.DomainEvent) ctx = store.WithTenantID(ctx, tid) } } + if event.AgentID != "" { + if aid, err := uuid.Parse(event.AgentID); err == nil { + ctx = store.WithAgentID(ctx, aid) + } + } agentID := event.AgentID userID := event.UserID @@ -205,7 +212,7 @@ func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Prov } body := strings.Join(summaries, "\n---\n") - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: dreamingSystemPrompt}, {Role: "user", Content: "Session summaries:\n---\n" + body + "\n---"}, @@ -214,6 +221,11 @@ func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Prov Options: map[string]any{ providers.OptMaxTokens: dreamingMaxTokens, }, + } + resp, err := w.usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "dreaming-synthesis", + MaxOutputTokens: dreamingMaxTokens, }) if err != nil { return "", fmt.Errorf("dreaming chat: %w", err) diff --git a/internal/consolidation/episodic_worker.go b/internal/consolidation/episodic_worker.go index 8a5f8cf5..b93a5383 100644 --- a/internal/consolidation/episodic_worker.go +++ b/internal/consolidation/episodic_worker.go @@ -10,19 +10,21 @@ import ( "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/bgalert" "github.com/nextlevelbuilder/goclaw/internal/eventbus" - "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/providerresolve" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // episodicWorker handles session.completed events → creates episodic summaries. type episodicWorker struct { store store.EpisodicStore - sessions store.SessionCoreStore // for reading session messages during summarization - systemConfigs store.SystemConfigStore // per-tenant provider config - registry *providers.Registry // provider resolution + sessions store.SessionCoreStore // for reading session messages during summarization + systemConfigs store.SystemConfigStore // per-tenant provider config + registry *providers.Registry // provider resolution eventBus eventbus.DomainEventBus alertDeps bgalert.AlertDeps + usageCaps *usagecaps.Service } // resolveProvider delegates to shared background provider resolution. @@ -53,6 +55,7 @@ func (w *episodicWorker) Handle(ctx context.Context, event eventbus.DomainEvent) if err != nil { return fmt.Errorf("episodic: invalid agent_id %q: %w", event.AgentID, err) } + ctx = store.WithAgentID(ctx, agentUUID) // Build source_id for idempotency sourceID := fmt.Sprintf("%s:%d", payload.SessionKey, payload.CompactionCount) @@ -168,13 +171,18 @@ func (w *episodicWorker) summarizeFromMessages(ctx context.Context, provider pro sctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - resp, err := provider.Chat(sctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: summarizationPrompt}, {Role: "user", Content: sb.String()}, }, Model: model, Options: map[string]any{"max_tokens": 1024, "temperature": 0.3}, + } + resp, err := w.usageCaps.Chat(sctx, provider, req, usagecaps.ChatOptions{ + ModelID: model, + Purpose: "episodic-summary", + MaxOutputTokens: 1024, }) if err != nil { return "", err diff --git a/internal/consolidation/workers.go b/internal/consolidation/workers.go index 1c5fb82a..2c3be24e 100644 --- a/internal/consolidation/workers.go +++ b/internal/consolidation/workers.go @@ -13,6 +13,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/eventbus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ConsolidationDeps bundles all dependencies for the consolidation pipeline. @@ -26,6 +27,7 @@ type ConsolidationDeps struct { Registry *providers.Registry // provider resolution Extractor EntityExtractor AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + UsageCaps *usagecaps.Service // AgentStore is optional: when present, the dreaming worker reads // per-agent overrides from MemoryConfig.Dreaming. If nil, the worker // uses its built-in defaults for every agent. @@ -42,6 +44,7 @@ func Register(deps ConsolidationDeps) func() { registry: deps.Registry, eventBus: deps.EventBus, alertDeps: deps.AlertDeps, + usageCaps: deps.UsageCaps, } semantic := &semanticWorker{ kgStore: deps.KGStore, @@ -59,6 +62,7 @@ func Register(deps ConsolidationDeps) func() { systemConfigs: deps.SystemConfigs, registry: deps.Registry, alertDeps: deps.AlertDeps, + usageCaps: deps.UsageCaps, threshold: dreamingDefaultThreshold, debounce: dreamingDefaultDebounce, resolveConfig: newAgentStoreResolver(deps.AgentStore), diff --git a/internal/gateway/methods/agents_update.go b/internal/gateway/methods/agents_update.go index e1f96651..d0d76d69 100644 --- a/internal/gateway/methods/agents_update.go +++ b/internal/gateway/methods/agents_update.go @@ -60,6 +60,10 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client if req.Params != nil { json.Unmarshal(req.Params, ¶ms) } + var rawParams map[string]json.RawMessage + if req.Params != nil { + _ = json.Unmarshal(req.Params, &rawParams) + } if params.AgentID == "" { client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInvalidRequest, i18n.T(locale, i18n.MsgRequired, "agentId"))) @@ -104,7 +108,9 @@ func (m *AgentsMethods) handleUpdate(ctx context.Context, client *gateway.Client if params.IsDefault != nil { updates["is_default"] = *params.IsDefault } - if params.BudgetCents != nil { + if rawBudget, ok := rawParams["budget_monthly_cents"]; ok && strings.TrimSpace(string(rawBudget)) == "null" { + updates["budget_monthly_cents"] = nil + } else if params.BudgetCents != nil { updates["budget_monthly_cents"] = *params.BudgetCents } // Per-agent JSONB config overrides diff --git a/internal/gateway/methods/chat.go b/internal/gateway/methods/chat.go index 3ed1f2a1..3e02e70e 100644 --- a/internal/gateway/methods/chat.go +++ b/internal/gateway/methods/chat.go @@ -20,6 +20,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/sessions" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -32,6 +33,7 @@ type ChatMethods struct { eventBus bus.EventPublisher postTurn tools.PostTurnProcessor audioMgr *audio.Manager // for TTS auto-apply on WS responses (nil = disabled) + usageCaps *usagecaps.Service debouncer *chatDebouncer } @@ -46,6 +48,10 @@ func (m *ChatMethods) SetAudioManager(mgr *audio.Manager) { m.audioMgr = mgr } +func (m *ChatMethods) SetUsageCapService(s *usagecaps.Service) { + m.usageCaps = s +} + // SetPostTurnProcessor sets the post-turn processor for team task dispatch. func (m *ChatMethods) SetPostTurnProcessor(pt tools.PostTurnProcessor) { m.postTurn = pt @@ -352,7 +358,10 @@ func (m *ChatMethods) dispatchChatSends(requests []chatSendRequest) { // Use runCtxBase (WithoutCancel + tenant-aware) so title save uses correct tenant. titleCtx := runCtxBase go func() { - title := agent.GenerateTitle(titleCtx, agentProvider, agentModel, userMsg) + if uid := loop.UUID(); uid != uuid.Nil { + titleCtx = store.WithAgentID(titleCtx, uid) + } + title := agent.GenerateTitleWithUsageCaps(titleCtx, m.usageCaps, agentProvider, agentModel, userMsg) if title == "" { return } diff --git a/internal/gateway/methods/config.go b/internal/gateway/methods/config.go index 590d66a8..481181b3 100644 --- a/internal/gateway/methods/config.go +++ b/internal/gateway/methods/config.go @@ -22,7 +22,7 @@ type ConfigMethods struct { cfgPath string secretsStore store.ConfigSecretsStore syncFn func(ctx context.Context, cfg *config.Config) // nil-safe; syncs non-secret settings to system_configs - eventBus bus.EventPublisher // nil-safe; broadcasts config change events + eventBus bus.EventPublisher // nil-safe; broadcasts config change events } func NewConfigMethods(cfg *config.Config, cfgPath string, secretsStore store.ConfigSecretsStore, eventBus bus.EventPublisher) *ConfigMethods { @@ -269,6 +269,45 @@ func (m *ConfigMethods) handleSchema(_ context.Context, client *gateway.Client, "type": "object", "description": "Tool configuration (browser, exec, web search)", }, + "skills": map[string]any{ + "type": "object", + "description": "Skill storage and upload settings", + "properties": map[string]any{ + "max_upload_size_mb": map[string]any{ + "type": "integer", + "minimum": config.MinSkillMaxUploadSizeMB, + "maximum": config.MaxSkillMaxUploadSizeMB, + "default": config.DefaultSkillMaxUploadSizeMB, + "description": "Maximum skill ZIP upload size in MB", + }, + "slash_commands": map[string]any{ + "type": "object", + "description": "Explicit slash command skill activation settings", + "properties": map[string]any{ + "enabled": map[string]any{ + "type": "boolean", + "default": true, + "description": "Enable slash command detection in user prompts", + }, + "suggest_not_found": map[string]any{ + "type": "boolean", + "default": true, + "description": "Suggest similar skills when a requested skill is not found", + }, + "partial_matching": map[string]any{ + "type": "boolean", + "default": false, + "description": "Allow unique skill slug/name prefixes", + }, + "prefix": map[string]any{ + "type": "string", + "default": config.DefaultSkillSlashCommandPrefix, + "description": "Single-character slash command prefix", + }, + }, + }, + }, + }, "sessions": map[string]any{ "type": "object", "description": "Session storage configuration", diff --git a/internal/gateway/server.go b/internal/gateway/server.go index b742ef87..6cb1b0cf 100644 --- a/internal/gateway/server.go +++ b/internal/gateway/server.go @@ -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) } diff --git a/internal/hooks/handlers/prompt.go b/internal/hooks/handlers/prompt.go index 2ec9fb1f..062c3c1f 100644 --- a/internal/hooks/handlers/prompt.go +++ b/internal/hooks/handlers/prompt.go @@ -17,6 +17,8 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/hooks" "github.com/nextlevelbuilder/goclaw/internal/hooks/budget" "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ── Public surface ────────────────────────────────────────────────────────── @@ -46,6 +48,9 @@ type PromptHandler struct { // budget checks are skipped (Lite edition behavior). Budget *budget.Store + // UsageCaps enforces provider/model/tenant cost caps before the hook LLM call. + UsageCaps *usagecaps.Service + // DefaultModel is used when a hook config does not specify one. // Recommended: "haiku" for cheap evaluation. DefaultModel string @@ -155,7 +160,17 @@ func (h *PromptHandler) Execute(ctx context.Context, cfg hooks.HookConfig, ev ho req := h.buildChatRequest(cfg, ev, resolvedModel) // 6. Call provider. - resp, err := provider.Chat(ctx, req) + callCtx := ctx + if ev.TenantID != uuid.Nil { + callCtx = store.WithTenantID(callCtx, ev.TenantID) + } + resp, err := h.UsageCaps.Chat(callCtx, provider, req, usagecaps.ChatOptions{ + TenantID: ev.TenantID, + ProviderName: provider.Name(), + ModelID: resolvedModel, + Purpose: "hook-prompt", + MaxOutputTokens: promptMaxTokens(req), + }) if err != nil { // Fail-closed on transport/provider error for blocking events. if ev.HookEvent.IsBlocking() { @@ -300,6 +315,23 @@ func (h *PromptHandler) buildChatRequest(cfg hooks.HookConfig, ev hooks.Event, m } } +func promptMaxTokens(req providers.ChatRequest) int { + if req.Options == nil { + return 512 + } + if v, ok := req.Options[providers.OptMaxTokens]; ok { + switch n := v.(type) { + case int: + return n + case int64: + return int(n) + case float64: + return int(n) + } + } + return 512 +} + // sanitizeToolInput canonicalizes a tool_input map into stable JSON with // sorted keys, stripping only structural noise. Actual injection-attack // detection is delegated to the evaluator LLM (which has the anti-injection diff --git a/internal/http/knowledge_graph.go b/internal/http/knowledge_graph.go index 7333f851..fcc61b90 100644 --- a/internal/http/knowledge_graph.go +++ b/internal/http/knowledge_graph.go @@ -7,12 +7,14 @@ import ( kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // KnowledgeGraphHandler handles KG entity/relation management endpoints. type KnowledgeGraphHandler struct { store store.KnowledgeGraphStore providerReg *providers.Registry + usageCaps *usagecaps.Service } // NewKnowledgeGraphHandler creates a handler for KG management endpoints. @@ -20,6 +22,10 @@ func NewKnowledgeGraphHandler(s store.KnowledgeGraphStore, providerReg *provider return &KnowledgeGraphHandler{store: s, providerReg: providerReg} } +func (h *KnowledgeGraphHandler) SetUsageCapService(s *usagecaps.Service) { + h.usageCaps = s +} + // NewExtractor creates an Extractor from the given provider name and model. func (h *KnowledgeGraphHandler) NewExtractor(ctx context.Context, providerName, model string, minConfidence float64) *kg.Extractor { if h.providerReg == nil || providerName == "" || model == "" { @@ -29,7 +35,9 @@ func (h *KnowledgeGraphHandler) NewExtractor(ctx context.Context, providerName, if err != nil { return nil } - return kg.NewExtractor(p, model, minConfidence) + extractor := kg.NewExtractor(p, model, minConfidence) + extractor.SetUsageCapService(h.usageCaps) + return extractor } // RegisterRoutes registers all KG routes on the given mux. diff --git a/internal/http/knowledge_graph_handlers.go b/internal/http/knowledge_graph_handlers.go index 80cb420f..8e31b319 100644 --- a/internal/http/knowledge_graph_handlers.go +++ b/internal/http/knowledge_graph_handlers.go @@ -5,6 +5,7 @@ import ( "net/http" "strconv" + "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -156,6 +157,10 @@ func (h *KnowledgeGraphHandler) handleTraverse(w http.ResponseWriter, r *http.Re func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Request) { locale := extractLocale(r) agentID := r.PathValue("agentID") + callCtx := r.Context() + if parsedAgentID, err := uuid.Parse(agentID); err == nil { + callCtx = store.WithAgentID(callCtx, parsedAgentID) + } var body struct { Text string `json:"text"` @@ -176,13 +181,13 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req return } - extractor := h.NewExtractor(r.Context(), body.Provider, body.Model, body.MinConf) + extractor := h.NewExtractor(callCtx, body.Provider, body.Model, body.MinConf) if extractor == nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidProviderOrModel)}) return } - result, err := extractor.Extract(r.Context(), body.Text) + result, err := extractor.Extract(callCtx, body.Text) if err != nil { slog.Warn("kg.extract failed", "error", err) writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()}) @@ -213,10 +218,10 @@ func (h *KnowledgeGraphHandler) handleExtract(w http.ResponseWriter, r *http.Req } writeJSON(w, http.StatusOK, map[string]any{ - "entities": len(result.Entities), - "relations": len(result.Relations), - "dedup_merged": dedupMerged, - "dedup_flagged": dedupFlagged, + "entities": len(result.Entities), + "relations": len(result.Relations), + "dedup_merged": dedupMerged, + "dedup_flagged": dedupFlagged, }) } diff --git a/internal/http/pending_messages.go b/internal/http/pending_messages.go index 6bc7d0a6..3f144fa7 100644 --- a/internal/http/pending_messages.go +++ b/internal/http/pending_messages.go @@ -11,6 +11,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providerresolve" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // PendingMessagesHandler handles pending message HTTP endpoints. @@ -22,6 +23,7 @@ type PendingMessagesHandler struct { maxTokens int // max output tokens for LLM summarization (0 = use default) cfgProvider string // config-level provider override (empty = resolve from agent) cfgModel string // config-level model override (empty = resolve from agent) + usageCaps *usagecaps.Service } func NewPendingMessagesHandler(s store.PendingMessageStore, agentStore store.AgentStore, providerReg *providers.Registry) *PendingMessagesHandler { @@ -40,6 +42,10 @@ func (h *PendingMessagesHandler) SetProviderModel(provider, model string) { h.cfgModel = model } +func (h *PendingMessagesHandler) SetUsageCapService(s *usagecaps.Service) { + h.usageCaps = s +} + func (h *PendingMessagesHandler) RegisterRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /v1/pending-messages", h.authMiddleware(h.handleListGroups)) mux.HandleFunc("GET /v1/pending-messages/messages", h.authMiddleware(h.handleListMessages)) @@ -148,7 +154,7 @@ func (h *PendingMessagesHandler) handleCompact(w http.ResponseWriter, r *http.Re go func() { ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 180*time.Second) defer cancel() - remaining, err := channels.CompactGroup(ctx, h.store, req.ChannelName, req.HistoryKey, provider, model, keepRecent, h.maxTokens) + remaining, err := channels.CompactGroup(ctx, h.store, req.ChannelName, req.HistoryKey, provider, model, keepRecent, h.maxTokens, h.usageCaps) if err != nil { slog.Warn("compact.failed", "channel", req.ChannelName, "key", req.HistoryKey, "error", err) } else { diff --git a/internal/http/provider_verify.go b/internal/http/provider_verify.go index 48e4e591..8b3e402e 100644 --- a/internal/http/provider_verify.go +++ b/internal/http/provider_verify.go @@ -16,6 +16,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // HandleVerifyProviderForTest invokes the verify handler directly without auth @@ -126,8 +127,9 @@ func (h *ProvidersHandler) handleVerifyProvider(w http.ResponseWriter, r *http.R ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) defer cancel() + ctx = store.WithTenantID(ctx, p.TenantID) - _, err = provider.Chat(ctx, providers.ChatRequest{ + reqChat := providers.ChatRequest{ Messages: []providers.Message{ {Role: "user", Content: "hi"}, }, @@ -136,6 +138,13 @@ func (h *ProvidersHandler) handleVerifyProvider(w http.ResponseWriter, r *http.R // Use a small but safe value — reasoning models need headroom beyond 1 token. "max_tokens": 50, }, + } + _, err = h.usageCaps.Chat(ctx, provider, reqChat, usagecaps.ChatOptions{ + TenantID: p.TenantID, + ProviderName: p.Name, + ModelID: req.Model, + Purpose: "provider-verify", + MaxOutputTokens: 50, }) if err != nil { writeJSON(w, http.StatusOK, map[string]any{"valid": false, "error": friendlyVerifyError(err)}) diff --git a/internal/http/providers.go b/internal/http/providers.go index bee754f0..9127982f 100644 --- a/internal/http/providers.go +++ b/internal/http/providers.go @@ -22,6 +22,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/permissions" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -39,6 +40,7 @@ type ProvidersHandler struct { tracingStore store.TracingStore // optional: for provider-scoped pool activity agents store.AgentCRUDStore // optional: for provider pool activity agent lookup modelReg providers.ModelRegistry // optional: forward-compat model resolver for Anthropic + usageCaps *usagecaps.Service } // NewProvidersHandler creates a handler for provider management endpoints. @@ -85,6 +87,10 @@ func (h *ProvidersHandler) SetModelRegistry(r providers.ModelRegistry) { h.modelReg = r } +func (h *ProvidersHandler) SetUsageCapService(s *usagecaps.Service) { + h.usageCaps = s +} + // resolveAPIBase returns the provider's api_base, falling back to config/env if empty. // For Ollama/OllamaCloud providers, applies a safety-net normalization: if the stored // value is missing the /v1 suffix (pre-existing record before write-time normalization), diff --git a/internal/http/secure_cli.go b/internal/http/secure_cli.go index 5aa05a5c..ee4d0ad8 100644 --- a/internal/http/secure_cli.go +++ b/internal/http/secure_cli.go @@ -7,7 +7,6 @@ import ( "net/http" "os/exec" "regexp" - "sort" "strings" "github.com/google/uuid" @@ -70,68 +69,13 @@ func (h *SecureCLIHandler) emitCacheInvalidate(key string) { // envKeysFromDecryptedJSON returns sorted env variable names from plaintext env JSON (decrypted blob). func envKeysFromDecryptedJSON(env []byte) []string { - empty := []string{} - if len(env) == 0 { - return empty - } - var m map[string]any - if err := json.Unmarshal(env, &m); err != nil { - return empty - } - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - sort.Strings(keys) - return keys + return store.SecureCLIEnvKeys(env) } -// mergeSecureCLIEnv merges incoming env from the UI with existing stored env. -// Incoming defines the full set of keys shown in the form: keys omitted were removed. -// Empty string means "keep existing value" for that key when it already exists. -func mergeSecureCLIEnv(existingJSON []byte, incoming map[string]any) (map[string]string, error) { - existing := map[string]string{} - if len(existingJSON) > 0 { - if err := json.Unmarshal(existingJSON, &existing); err != nil { - return nil, fmt.Errorf("parse existing env: %w", err) - } - } - out := make(map[string]string) - for k, v := range incoming { - if k == "" { - continue - } - sv, err := envValueAsString(v) - if err != nil { - return nil, fmt.Errorf("invalid environment variable value") - } - if sv != "" { - out[k] = sv - continue - } - if ev, ok := existing[k]; ok { - out[k] = ev - } - } - return out, nil -} - -func envValueAsString(v any) (string, error) { - switch t := v.(type) { - case string: - return t, nil - case float64: - return fmt.Sprint(t), nil - case bool: - if t { - return "true", nil - } - return "false", nil - case nil: - return "", nil - default: - return "", fmt.Errorf("value must be a string") - } +func populateBinaryEnvResponse(b *store.SecureCLIBinary) { + b.EnvKeys = store.SecureCLIEnvKeys(b.EncryptedEnv) + b.Env = store.SanitizeSecureCLIEnvJSON(b.EncryptedEnv) + b.EncryptedEnv = nil } func (h *SecureCLIHandler) handleList(w http.ResponseWriter, r *http.Request) { @@ -142,27 +86,26 @@ func (h *SecureCLIHandler) handleList(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgFailedToList, "CLI credentials")}) return } - // Never send env values; only variable names for editing. + // Sensitive env values stay masked; value-kind entries are returned for editing. for i := range result { - result[i].EnvKeys = envKeysFromDecryptedJSON(result[i].EncryptedEnv) - result[i].EncryptedEnv = nil + populateBinaryEnvResponse(&result[i]) } writeJSON(w, http.StatusOK, map[string]any{"items": result}) } // secureCLICreateRequest supports both preset-based and custom creation. type secureCLICreateRequest struct { - Preset string `json:"preset,omitempty"` // auto-fill from preset - BinaryName string `json:"binary_name"` - BinaryPath *string `json:"binary_path,omitempty"` - Description string `json:"description"` - Env map[string]string `json:"env"` // plaintext env vars (encrypted by store) - DenyArgs json.RawMessage `json:"deny_args,omitempty"` - DenyVerbose json.RawMessage `json:"deny_verbose,omitempty"` - TimeoutSeconds int `json:"timeout_seconds,omitempty"` - Tips string `json:"tips,omitempty"` - IsGlobal *bool `json:"is_global,omitempty"` - Enabled bool `json:"enabled"` + Preset string `json:"preset,omitempty"` // auto-fill from preset + BinaryName string `json:"binary_name"` + BinaryPath *string `json:"binary_path,omitempty"` + Description string `json:"description"` + Env json.RawMessage `json:"env"` // plaintext env vars or env entry objects (encrypted by store) + DenyArgs json.RawMessage `json:"deny_args,omitempty"` + DenyVerbose json.RawMessage `json:"deny_verbose,omitempty"` + TimeoutSeconds int `json:"timeout_seconds,omitempty"` + Tips string `json:"tips,omitempty"` + IsGlobal *bool `json:"is_global,omitempty"` + Enabled bool `json:"enabled"` } func (h *SecureCLIHandler) handleCreate(w http.ResponseWriter, r *http.Request) { @@ -209,10 +152,14 @@ func (h *SecureCLIHandler) handleCreate(w http.ResponseWriter, r *http.Request) return } - // Serialize env as JSON bytes (store layer encrypts) - envJSON, err := json.Marshal(req.Env) + envEntries, err := store.ParseSecureCLIEnv(req.Env) if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid env"}) + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, err.Error())}) + return + } + envJSON, err := store.SerializeSecureCLIEnv(envEntries) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, err.Error())}) return } @@ -240,7 +187,7 @@ func (h *SecureCLIHandler) handleCreate(w http.ResponseWriter, r *http.Request) } tools.ResetCredentialScrubValues() // clear stale scrub values - b.EncryptedEnv = nil // don't return credentials + populateBinaryEnvResponse(b) emitAudit(h.msgBus, r, "secure_cli.created", "secure_cli", b.ID.String()) h.emitCacheInvalidate(b.ID.String()) writeJSON(w, http.StatusCreated, b) @@ -260,8 +207,7 @@ func (h *SecureCLIHandler) handleGet(w http.ResponseWriter, r *http.Request) { return } - b.EnvKeys = envKeysFromDecryptedJSON(b.EncryptedEnv) - b.EncryptedEnv = nil // don't expose credential values + populateBinaryEnvResponse(b) writeJSON(w, http.StatusOK, b) } @@ -293,25 +239,23 @@ func (h *SecureCLIHandler) handleUpdate(w http.ResponseWriter, r *http.Request) // If env is updated, merge with stored env so empty values mean "keep existing secret". if envVal, ok := updates["env"]; ok { - if envMap, isMap := envVal.(map[string]any); isMap { - cur, err := h.store.Get(r.Context(), id) - if err != nil { - writeJSON(w, http.StatusNotFound, map[string]string{"error": i18n.T(locale, i18n.MsgNotFound, "credential", id.String())}) - return - } - merged, err := mergeSecureCLIEnv(cur.EncryptedEnv, envMap) - if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) - return - } - envJSON, err := json.Marshal(merged) - if err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid env"}) - return - } - updates["encrypted_env"] = string(envJSON) - delete(updates, "env") + envJSON, err := json.Marshal(envVal) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, err.Error())}) + return } + cur, err := h.store.Get(r.Context(), id) + if err != nil { + writeJSON(w, http.StatusNotFound, map[string]string{"error": i18n.T(locale, i18n.MsgNotFound, "credential", id.String())}) + return + } + merged, err := store.MergeSecureCLIEnv(cur.EncryptedEnv, envJSON) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, err.Error())}) + return + } + updates["encrypted_env"] = string(merged) + delete(updates, "env") } if err := h.store.Update(r.Context(), id, updates); err != nil { diff --git a/internal/http/secure_cli_agent_grants.go b/internal/http/secure_cli_agent_grants.go index 11f87289..342a9b2e 100644 --- a/internal/http/secure_cli_agent_grants.go +++ b/internal/http/secure_cli_agent_grants.go @@ -68,7 +68,7 @@ func (h *SecureCLIGrantHandler) RegisterRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /v1/cli-credentials/{id}/agent-grants/{grantId}", auth(h.handleGet)) mux.HandleFunc("PUT /v1/cli-credentials/{id}/agent-grants/{grantId}", auth(h.handleUpdate)) mux.HandleFunc("DELETE /v1/cli-credentials/{id}/agent-grants/{grantId}", auth(h.handleDelete)) - // POST (not GET) to prevent caching and satisfy CSRF semantics per Red Team C1. + // POST keeps revealed secret material out of URL/history and avoids query caching. mux.HandleFunc("POST /v1/cli-credentials/{id}/agent-grants/{grantId}/env:reveal", auth(h.handleRevealEnv)) } @@ -76,45 +76,41 @@ func (h *SecureCLIGrantHandler) RegisterRoutes(mux *http.ServeMux) { // EnvVars is optional; plaintext values are encrypted by the store layer. // Clients MUST NOT send encrypted_env — that field is never accepted from the wire. type grantCreateRequest struct { - AgentID uuid.UUID `json:"agent_id"` - EnvVars map[string]string `json:"env_vars,omitempty"` - DenyArgs *json.RawMessage `json:"deny_args,omitempty"` - DenyVerbose *json.RawMessage `json:"deny_verbose,omitempty"` - TimeoutSeconds *int `json:"timeout_seconds,omitempty"` - Tips *string `json:"tips,omitempty"` - Enabled *bool `json:"enabled,omitempty"` + AgentID uuid.UUID `json:"agent_id"` + EnvVars json.RawMessage `json:"env_vars,omitempty"` + DenyArgs *json.RawMessage `json:"deny_args,omitempty"` + DenyVerbose *json.RawMessage `json:"deny_verbose,omitempty"` + TimeoutSeconds *int `json:"timeout_seconds,omitempty"` + Tips *string `json:"tips,omitempty"` + Enabled *bool `json:"enabled,omitempty"` } -// populateGrantEnvFields sets EnvKeys (sorted) and EnvSet from the grant's decrypted env bytes. -// Plaintext values are never exposed — only key names. +// populateGrantEnvFields sets sorted key names, env presence, and sanitized entries. func populateGrantEnvFields(g *store.SecureCLIAgentGrant) { if len(g.EncryptedEnv) == 0 { g.EnvKeys = []string{} + g.Env = nil g.EnvSet = false return } - var m map[string]any - if err := json.Unmarshal(g.EncryptedEnv, &m); err != nil { - g.EnvKeys = []string{} - g.EnvSet = false - return - } - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - sort.Strings(keys) + keys := store.SecureCLIEnvKeys(g.EncryptedEnv) g.EnvKeys = keys + g.Env = store.SanitizeSecureCLIEnvJSON(g.EncryptedEnv) g.EnvSet = len(keys) > 0 } // validateAndSerializeEnvVars validates env keys/values via denylist and returns serialized JSON. // Returns (nil, 400 error response written) on denial, (jsonBytes, nil) on success. // Never logs env values or keys in error paths. -func validateAndSerializeEnvVars(w http.ResponseWriter, locale string, envVars map[string]string) ([]byte, bool) { - if len(envVars) == 0 { - b, _ := json.Marshal(envVars) - return b, true +func validateAndSerializeEnvVars(w http.ResponseWriter, locale string, raw json.RawMessage) ([]byte, bool) { + envEntries, err := store.ParseSecureCLIEnv(raw) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, err.Error())}) + return nil, false + } + envVars := make(map[string]string, len(envEntries)) + for key, entry := range envEntries { + envVars[key] = entry.Value } denied, valErr := crypto.ValidateGrantEnvVars(envVars) if valErr != nil { @@ -129,7 +125,7 @@ func validateAndSerializeEnvVars(w http.ResponseWriter, locale string, envVars m }) return nil, false } - b, err := json.Marshal(envVars) + b, err := store.SerializeSecureCLIEnv(envEntries) if err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, "serialization failed")}) return nil, false @@ -249,7 +245,7 @@ func (h *SecureCLIGrantHandler) handleCreate(w http.ResponseWriter, r *http.Requ envJSON, ok := validateAndSerializeEnvVars(w, locale, req.EnvVars) if !ok { // Grant was created but env validation failed; clean it up to avoid orphan row. - // Finding #13: log rollback-delete failures for ops visibility. + // Log rollback-delete failures so operators can clean up orphan rows. if delErr := h.grants.Delete(r.Context(), g.ID); delErr != nil { slog.Error("secure_cli_grants.create.rollback_delete", "grant_id", g.ID, @@ -261,7 +257,7 @@ func (h *SecureCLIGrantHandler) handleCreate(w http.ResponseWriter, r *http.Requ } if err := h.grants.UpdateGrantEnv(r.Context(), g.ID, envJSON); err != nil { slog.Error("secure_cli_grants.create.set_env", "grant_id", g.ID, "error", err) - // Finding #13: log rollback-delete failures for ops visibility. + // Log rollback-delete failures so operators can clean up orphan rows. if delErr := h.grants.Delete(r.Context(), g.ID); delErr != nil { slog.Error("secure_cli_grants.create.rollback_delete", "grant_id", g.ID, @@ -323,7 +319,7 @@ func (h *SecureCLIGrantHandler) handleUpdate(w http.ResponseWriter, r *http.Requ } if allowedScalar[k] { var decoded any - // Finding #3: return 400 on Unmarshal failure — silent discard means admin + // Return 400 on unmarshal failure; silent discard means admin // thinks they applied a change (e.g. enabled: "false") but the grant is unchanged. if err := json.Unmarshal(v, &decoded); err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{ @@ -335,7 +331,7 @@ func (h *SecureCLIGrantHandler) handleUpdate(w http.ResponseWriter, r *http.Requ } } // 3-state env_vars semantics: absent=skip, null=clear, {...}=replace. - // Finding #15: {} (empty map) is treated as clear — same as null. + // Empty map is treated as clear, same as null. // TS type: absent | null | Record — see ui/web/src/types/cli-credential.ts. var envJSON []byte envPresent := false @@ -343,18 +339,22 @@ func (h *SecureCLIGrantHandler) handleUpdate(w http.ResponseWriter, r *http.Requ envPresent = true var envPtr *map[string]string if string(envRaw) != "null" { - var m map[string]string - if err := json.Unmarshal(envRaw, &m); err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, "env_vars must be a string map")}) + envEntries, err := store.ParseSecureCLIEnv(envRaw) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, "env_vars must be a string map or env entries")}) return } + m := make(map[string]string, len(envEntries)) + for key, entry := range envEntries { + m[key] = entry.Value + } envPtr = &m } // envPtr == nil → clear; envPtr != nil → replace. // Note: envPtr pointing to an empty map ({}) is treated as clear (same as null) — // envJSON stays nil and UpdateGrantEnv(nil) removes the override. if envPtr != nil && len(*envPtr) > 0 { - j, ok := validateAndSerializeEnvVars(w, locale, *envPtr) + j, ok := validateAndSerializeEnvVars(w, locale, envRaw) if !ok { return } @@ -430,8 +430,8 @@ func (h *SecureCLIGrantHandler) handleRevealEnv(w http.ResponseWriter, r *http.R locale := store.LocaleFromContext(ctx) // Rate limit: 10 reveals/min per authenticated caller (context UserID). - // Finding #2: require non-empty UserID from authenticated context. - // If UserID is empty, the auth middleware failed to populate it — reject rather + // Require non-empty UserID from authenticated context. + // If UserID is empty, auth middleware failed to populate it — reject rather // than fall back to a spoofable header or IP address. callerID := store.UserIDFromContext(ctx) if callerID == "" { @@ -476,8 +476,8 @@ func (h *SecureCLIGrantHandler) handleRevealEnv(w http.ResponseWriter, r *http.R writeJSON(w, http.StatusOK, map[string]any{"env_vars": map[string]string{}}) return } - var envVars map[string]string - if err := json.Unmarshal(g.EncryptedEnv, &envVars); err != nil { + envVars, err := store.FlattenSecureCLIEnv(g.EncryptedEnv) + if err != nil { slog.Error("secure_cli_grants.reveal.parse", "grant_id", g.ID, "error", err) writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgInternalError, "parse grant env")}) return diff --git a/internal/http/secure_cli_agent_grants_test.go b/internal/http/secure_cli_agent_grants_test.go index fd77450e..9a4f322c 100644 --- a/internal/http/secure_cli_agent_grants_test.go +++ b/internal/http/secure_cli_agent_grants_test.go @@ -3,6 +3,7 @@ package http import ( "context" "database/sql" + "encoding/json" "io" "net/http" "net/http/httptest" @@ -59,8 +60,16 @@ func (s *fakeSecureCLIGrantStore) Delete(context.Context, uuid.UUID) error { return nil } -func (s *fakeSecureCLIGrantStore) ListByBinary(context.Context, uuid.UUID) ([]store.SecureCLIAgentGrant, error) { - return nil, nil +func (s *fakeSecureCLIGrantStore) ListByBinary(_ context.Context, binaryID uuid.UUID) ([]store.SecureCLIAgentGrant, error) { + grants := make([]store.SecureCLIAgentGrant, 0, len(s.grants)) + for _, grant := range s.grants { + if grant == nil || grant.BinaryID != binaryID { + continue + } + cp := *grant + grants = append(grants, cp) + } + return grants, nil } func (s *fakeSecureCLIGrantStore) ListByAgent(context.Context, uuid.UUID) ([]store.SecureCLIAgentGrant, error) { @@ -202,3 +211,42 @@ func TestSecureCLIGrantUpdateRejectsInvalidEnvVarsBeforeScalarUpdate(t *testing. t.Fatal("invalid env_vars request must not persist scalar grant updates") } } + +func TestSecureCLIGrantGetSanitizesMixedEnv(t *testing.T) { + binaryID := uuid.New() + grantID := uuid.New() + fake := &fakeSecureCLIGrantStore{ + grants: map[uuid.UUID]*store.SecureCLIAgentGrant{ + grantID: { + BaseModel: store.BaseModel{ID: grantID}, + BinaryID: binaryID, + AgentID: uuid.New(), + Enabled: true, + EncryptedEnv: []byte(`{"TOKEN":"secret-token","PUBLIC_BASE_URL":{"kind":"value","value":"https://goclaw.sh"}}`), + }, + }, + } + h := NewSecureCLIGrantHandler(fake, nil, nil) + rr, req := requestWithGrantPath(http.MethodGet, nil, binaryID, grantID) + + h.handleGet(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", rr.Code, rr.Body.String()) + } + if strings.Contains(rr.Body.String(), "secret-token") { + t.Fatalf("sensitive grant env leaked in response: %s", rr.Body.String()) + } + var got struct { + Env map[string]store.SecureCLIEnvResponseEntry `json:"env"` + } + if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Env["TOKEN"].Masked || got.Env["TOKEN"].Value != nil { + t.Fatalf("TOKEN not masked: %#v", got.Env["TOKEN"]) + } + if got.Env["PUBLIC_BASE_URL"].Value == nil || *got.Env["PUBLIC_BASE_URL"].Value != "https://goclaw.sh" { + t.Fatalf("PUBLIC_BASE_URL not returned: %#v", got.Env["PUBLIC_BASE_URL"]) + } +} diff --git a/internal/http/secure_cli_env_response_test.go b/internal/http/secure_cli_env_response_test.go new file mode 100644 index 00000000..4e08bcd0 --- /dev/null +++ b/internal/http/secure_cli_env_response_test.go @@ -0,0 +1,129 @@ +package http + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +type fakeSecureCLIStore struct { + binary *store.SecureCLIBinary + user *store.SecureCLIUserCredential +} + +func (s *fakeSecureCLIStore) Create(context.Context, *store.SecureCLIBinary) error { return nil } +func (s *fakeSecureCLIStore) Get(context.Context, uuid.UUID) (*store.SecureCLIBinary, error) { + cp := *s.binary + return &cp, nil +} +func (s *fakeSecureCLIStore) Update(context.Context, uuid.UUID, map[string]any) error { return nil } +func (s *fakeSecureCLIStore) Delete(context.Context, uuid.UUID) error { return nil } +func (s *fakeSecureCLIStore) List(context.Context) ([]store.SecureCLIBinary, error) { + cp := *s.binary + return []store.SecureCLIBinary{cp}, nil +} +func (s *fakeSecureCLIStore) LookupByBinary(context.Context, string, *uuid.UUID, string) (*store.SecureCLIBinary, error) { + return nil, nil +} +func (s *fakeSecureCLIStore) ListEnabled(context.Context) ([]store.SecureCLIBinary, error) { + return nil, nil +} +func (s *fakeSecureCLIStore) ListForAgent(context.Context, uuid.UUID) ([]store.SecureCLIBinary, error) { + return nil, nil +} +func (s *fakeSecureCLIStore) IsRegisteredBinary(context.Context, string) (bool, error) { + return false, nil +} +func (s *fakeSecureCLIStore) GetUserCredentials(context.Context, uuid.UUID, string) (*store.SecureCLIUserCredential, error) { + cp := *s.user + return &cp, nil +} +func (s *fakeSecureCLIStore) SetUserCredentials(context.Context, uuid.UUID, string, []byte) error { + return nil +} +func (s *fakeSecureCLIStore) DeleteUserCredentials(context.Context, uuid.UUID, string) error { + return nil +} +func (s *fakeSecureCLIStore) ListUserCredentials(context.Context, uuid.UUID) ([]store.SecureCLIUserCredential, error) { + cp := *s.user + return []store.SecureCLIUserCredential{cp}, nil +} + +func TestSecureCLIGetSanitizesMixedEnv(t *testing.T) { + id := uuid.New() + h := NewSecureCLIHandler(&fakeSecureCLIStore{ + binary: &store.SecureCLIBinary{ + BaseModel: store.BaseModel{ID: id}, + BinaryName: "gh", + EncryptedEnv: []byte(`{"TOKEN":"secret-token","PUBLIC_BASE_URL":{"kind":"value","value":"https://goclaw.sh"}}`), + }, + }, nil) + req := httptest.NewRequest(http.MethodGet, "/v1/cli-credentials/"+id.String(), nil) + req.SetPathValue("id", id.String()) + rec := httptest.NewRecorder() + + h.handleGet(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + if strings.Contains(rec.Body.String(), "secret-token") { + t.Fatalf("sensitive env leaked in response: %s", rec.Body.String()) + } + var got struct { + Env map[string]store.SecureCLIEnvResponseEntry `json:"env"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Env["TOKEN"].Masked || got.Env["TOKEN"].Value != nil { + t.Fatalf("TOKEN not masked: %#v", got.Env["TOKEN"]) + } + if got.Env["PUBLIC_BASE_URL"].Value == nil || *got.Env["PUBLIC_BASE_URL"].Value != "https://goclaw.sh" { + t.Fatalf("value env not returned: %#v", got.Env["PUBLIC_BASE_URL"]) + } +} + +func TestSecureCLIUserCredentialsGetDoesNotReturnLegacySensitiveRaw(t *testing.T) { + binaryID := uuid.New() + h := NewSecureCLIHandler(&fakeSecureCLIStore{ + user: &store.SecureCLIUserCredential{ + ID: uuid.New(), + BinaryID: binaryID, + UserID: "user-1", + EncryptedEnv: []byte(`{"TOKEN":"secret-token","REGION":{"kind":"value","value":"asia-southeast1"}}`), + }, + }, nil) + req := httptest.NewRequest(http.MethodGet, "/v1/cli-credentials/"+binaryID.String()+"/user-credentials/user-1", nil) + req.SetPathValue("id", binaryID.String()) + req.SetPathValue("userId", "user-1") + rec := httptest.NewRecorder() + + h.handleGetUserCredentials(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + if strings.Contains(rec.Body.String(), "secret-token") { + t.Fatalf("legacy sensitive env leaked in response: %s", rec.Body.String()) + } + var got struct { + Env map[string]store.SecureCLIEnvResponseEntry `json:"env"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Env["TOKEN"].Masked || got.Env["TOKEN"].Value != nil { + t.Fatalf("TOKEN not masked: %#v", got.Env["TOKEN"]) + } + if got.Env["REGION"].Value == nil || *got.Env["REGION"].Value != "asia-southeast1" { + t.Fatalf("REGION not returned: %#v", got.Env["REGION"]) + } +} diff --git a/internal/http/secure_cli_user_credentials.go b/internal/http/secure_cli_user_credentials.go index 7f810594..5f089acf 100644 --- a/internal/http/secure_cli_user_credentials.go +++ b/internal/http/secure_cli_user_credentials.go @@ -25,13 +25,14 @@ func (h *SecureCLIHandler) handleListUserCredentials(w http.ResponseWriter, r *h } // Return without env values for listing (names only + timestamps) type entry struct { - ID uuid.UUID `json:"id"` - BinaryID uuid.UUID `json:"binary_id"` - UserID string `json:"user_id"` - HasEnv bool `json:"has_env"` - EnvKeys []string `json:"env_keys,omitempty"` - CreatedAt string `json:"created_at"` - UpdatedAt string `json:"updated_at"` + ID uuid.UUID `json:"id"` + BinaryID uuid.UUID `json:"binary_id"` + UserID string `json:"user_id"` + HasEnv bool `json:"has_env"` + EnvKeys []string `json:"env_keys,omitempty"` + Env map[string]store.SecureCLIEnvResponseEntry `json:"env,omitempty"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` } entries := make([]entry, 0, len(creds)) for _, c := range creds { @@ -42,6 +43,7 @@ func (h *SecureCLIHandler) handleListUserCredentials(w http.ResponseWriter, r *h UserID: c.UserID, HasEnv: len(c.EncryptedEnv) > 0, EnvKeys: envKeys, + Env: store.SanitizeSecureCLIEnvJSON(c.EncryptedEnv), CreatedAt: c.CreatedAt, UpdatedAt: c.UpdatedAt, }) @@ -73,14 +75,9 @@ func (h *SecureCLIHandler) handleGetUserCredentials(w http.ResponseWriter, r *ht return } - // Return decrypted env as JSON object (admin-only endpoint) - var envObj any - if len(cred.EncryptedEnv) > 0 { - _ = json.Unmarshal(cred.EncryptedEnv, &envObj) - } writeJSON(w, http.StatusOK, map[string]any{ "user_id": cred.UserID, - "env": envObj, + "env": store.SanitizeSecureCLIEnvJSON(cred.EncryptedEnv), }) } @@ -108,14 +105,23 @@ func (h *SecureCLIHandler) handleSetUserCredentials(w http.ResponseWriter, r *ht return } - // Validate env is a JSON object - var envCheck map[string]string - if err := json.Unmarshal(body.Env, &envCheck); err != nil { - writeJSON(w, http.StatusBadRequest, map[string]string{"error": "env must be a JSON object with string values"}) + existing, err := h.store.GetUserCredentials(r.Context(), binaryID, userID) + if err != nil { + locale := store.LocaleFromContext(r.Context()) + writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgInternalError, err.Error())}) + return + } + var existingEnv []byte + if existing != nil { + existingEnv = existing.EncryptedEnv + } + envJSON, err := store.MergeSecureCLIEnv(existingEnv, body.Env) + if err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgGrantEnvValueInvalid, err.Error())}) return } - if err := h.store.SetUserCredentials(r.Context(), binaryID, userID, body.Env); err != nil { + if err := h.store.SetUserCredentials(r.Context(), binaryID, userID, envJSON); err != nil { locale := store.LocaleFromContext(r.Context()) writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgInternalError, err.Error())}) return diff --git a/internal/http/skills.go b/internal/http/skills.go index 7457cbe8..7d04a4e5 100644 --- a/internal/http/skills.go +++ b/internal/http/skills.go @@ -21,8 +21,6 @@ import ( "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) -const maxSkillUploadSize = 20 << 20 // 20 MB - var ( aggregateInstallDeps = skills.AggregateMissingDeps installManagedDeps = skills.InstallDeps @@ -40,11 +38,13 @@ type SkillsHandler struct { tenantStore store.TenantStore db *sql.DB // for export/import direct queries uploadLocks sync.Map // per-slug mutex; bounded by validated slug set, entries are tiny (*sync.Mutex) + uploadLimitCfg config.SkillsConfig + systemConfigs store.SystemConfigStore } // NewSkillsHandler creates a handler for skill management endpoints. func NewSkillsHandler(skills store.SkillManageStore, baseDir, dataDir, bundledDir string, msgBus *bus.MessageBus, tenantCfgStore store.SkillTenantConfigStore, tenantStore store.TenantStore) *SkillsHandler { - return &SkillsHandler{skills: skills, baseDir: baseDir, dataDir: dataDir, bundledDir: bundledDir, msgBus: msgBus, tenantCfgStore: tenantCfgStore, tenantStore: tenantStore} + return &SkillsHandler{skills: skills, baseDir: baseDir, dataDir: dataDir, bundledDir: bundledDir, msgBus: msgBus, tenantCfgStore: tenantCfgStore, tenantStore: tenantStore, uploadLimitCfg: config.SkillsConfig{MaxUploadSizeMB: config.DefaultSkillMaxUploadSizeMB}} } // tenantSkillsDir returns the skills-store directory scoped to the requesting tenant. diff --git a/internal/http/skills_grants.go b/internal/http/skills_grants.go index eccebcad..60ce877b 100644 --- a/internal/http/skills_grants.go +++ b/internal/http/skills_grants.go @@ -1,8 +1,6 @@ package http import ( - "archive/zip" - "io" "log/slog" "net/http" @@ -225,16 +223,3 @@ func (h *SkillsHandler) handleRevokeUser(w http.ResponseWriter, r *http.Request) } // --- Helpers --- - -func readZipFile(f *zip.File) (string, error) { - rc, err := f.Open() - if err != nil { - return "", err - } - defer rc.Close() - data, err := io.ReadAll(rc) - if err != nil { - return "", err - } - return string(data), nil -} diff --git a/internal/http/skills_upload.go b/internal/http/skills_upload.go index dbd9e14d..4f69171f 100644 --- a/internal/http/skills_upload.go +++ b/internal/http/skills_upload.go @@ -5,6 +5,7 @@ import ( "context" "crypto/sha256" "encoding/json" + "errors" "fmt" "io" "log/slog" @@ -43,7 +44,8 @@ func (h *SkillsHandler) handleUpload(w http.ResponseWriter, r *http.Request) { return } - r.Body = http.MaxBytesReader(w, r.Body, maxSkillUploadSize) + preParseLimitMB := h.resolvePreParseUploadLimitMB(r.Context()) + r.Body = http.MaxBytesReader(w, r.Body, skillUploadBodyBytes(preParseLimitMB)) file, header, err := r.FormFile("file") if err != nil { @@ -70,6 +72,11 @@ func (h *SkillsHandler) handleUpload(w http.ResponseWriter, r *http.Request) { size, err := io.Copy(tmp, file) if err != nil { tmp.Close() + var maxBytesErr *http.MaxBytesError + if errors.As(err, &maxBytesErr) { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidRequest, skillUploadTooLargeMessage(maxBytesErr.Limit, preParseLimitMB))}) + return + } writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgInternalError, "failed to save upload")}) return } @@ -141,6 +148,11 @@ func (h *SkillsHandler) handleUpload(w http.ResponseWriter, r *http.Request) { } name, description, slug, frontmatter := skills.ParseSkillFrontmatter(skillContent) + uploadLimitMB := h.resolveSkillUploadLimitMB(r.Context(), frontmatter) + if size > skillUploadLimitBytes(uploadLimitMB) { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidRequest, skillUploadTooLargeMessage(size, uploadLimitMB))}) + return + } if name == "" { writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgRequired, "name in SKILL.md frontmatter")}) return @@ -239,17 +251,11 @@ func (h *SkillsHandler) handleUpload(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgInternalError, "failed to create skill file directory")}) return } - data, err := readZipFile(f) - if err != nil { + if err := copyZipFileToPath(f, destPath); err != nil { os.RemoveAll(destDir) writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgInvalidRequest, "failed to read ZIP entry")}) return } - if err := os.WriteFile(destPath, []byte(data), 0644); err != nil { - os.RemoveAll(destDir) - writeJSON(w, http.StatusInternalServerError, map[string]string{"error": i18n.T(locale, i18n.MsgInternalError, "failed to write skill files")}) - return - } if cleanName == "SKILL.md" { wroteSkillMD = true } diff --git a/internal/http/skills_upload_limits.go b/internal/http/skills_upload_limits.go new file mode 100644 index 00000000..75a31a09 --- /dev/null +++ b/internal/http/skills_upload_limits.go @@ -0,0 +1,88 @@ +package http + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/config" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +const ( + skillUploadMaxSizeConfigKey = config.SkillMaxUploadSizeSystemConfigKey + skillUploadMultipartOverheadBytes = int64(1 << 20) +) + +func (h *SkillsHandler) SetUploadLimitConfig(cfg config.SkillsConfig) { + h.uploadLimitCfg = cfg +} + +func (h *SkillsHandler) SetSystemConfigStore(s store.SystemConfigStore) { + h.systemConfigs = s +} + +func (h *SkillsHandler) resolveSkillUploadLimitMB(ctx context.Context, frontmatter map[string]string) int { + if mb, ok := h.resolveTenantSkillUploadLimitMB(ctx); ok { + return mb + } + if mb, ok := parseSkillUploadLimitMB(frontmatter["max_upload_size_mb"]); ok { + return config.ClampSkillMaxUploadSizeMB(mb) + } + return h.uploadLimitCfg.EffectiveMaxUploadSizeMB() +} + +func (h *SkillsHandler) resolvePreParseUploadLimitMB(ctx context.Context) int { + if mb, ok := h.resolveTenantSkillUploadLimitMB(ctx); ok { + return mb + } + return config.MaxSkillMaxUploadSizeMB +} + +func (h *SkillsHandler) resolveTenantSkillUploadLimitMB(ctx context.Context) (int, bool) { + if h.systemConfigs == nil { + return 0, false + } + raw, err := h.systemConfigs.Get(ctx, skillUploadMaxSizeConfigKey) + if err != nil { + return 0, false + } + mb, ok := parseSkillUploadLimitMB(raw) + if !ok { + return 0, false + } + return config.ClampSkillMaxUploadSizeMB(mb), true +} + +func parseSkillUploadLimitMB(raw string) (int, bool) { + raw = strings.TrimSpace(raw) + if raw == "" { + return 0, false + } + mb, err := strconv.Atoi(raw) + if err != nil { + return 0, false + } + return mb, true +} + +func skillUploadLimitBytes(limitMB int) int64 { + return int64(config.ClampSkillMaxUploadSizeMB(limitMB)) << 20 +} + +func skillUploadBodyBytes(limitMB int) int64 { + return skillUploadLimitBytes(limitMB) + skillUploadMultipartOverheadBytes +} + +func skillUploadTooLargeMessage(size int64, limitMB int) string { + return fmt.Sprintf("skill ZIP size %s exceeds %d MB limit", formatUploadBytes(size), limitMB) +} + +func formatUploadBytes(size int64) string { + const mb = int64(1 << 20) + if size >= mb { + return fmt.Sprintf("%.1f MB", float64(size)/float64(mb)) + } + return fmt.Sprintf("%d bytes", size) +} diff --git a/internal/http/skills_upload_test.go b/internal/http/skills_upload_test.go index c2bb6a0a..de45592c 100644 --- a/internal/http/skills_upload_test.go +++ b/internal/http/skills_upload_test.go @@ -12,11 +12,13 @@ import ( "net/http/httptest" "path/filepath" "reflect" + "strings" "testing" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/bus" + "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/pkg/protocol" @@ -670,6 +672,65 @@ func TestHandleUpload_ReturnsGrantErrors(t *testing.T) { } } +func TestResolveSkillUploadLimitMBPrecedenceAndClamp(t *testing.T) { + ctx := store.WithTenantID(context.Background(), store.MasterTenantID) + handler, _, _, _ := newTestUploadHandler(t) + handler.SetUploadLimitConfig(config.SkillsConfig{MaxUploadSizeMB: 30}) + + if got := handler.resolveSkillUploadLimitMB(ctx, nil); got != 30 { + t.Fatalf("global limit = %d, want 30", got) + } + if got := handler.resolveSkillUploadLimitMB(ctx, map[string]string{"max_upload_size_mb": "100"}); got != 100 { + t.Fatalf("frontmatter limit = %d, want 100", got) + } + + handler.SetSystemConfigStore(&skillUploadSystemConfigStore{data: map[string]string{skillUploadMaxSizeConfigKey: "40"}}) + if got := handler.resolveSkillUploadLimitMB(ctx, map[string]string{"max_upload_size_mb": "100"}); got != 40 { + t.Fatalf("tenant limit = %d, want 40", got) + } + + handler.SetSystemConfigStore(&skillUploadSystemConfigStore{data: map[string]string{skillUploadMaxSizeConfigKey: "999"}}) + if got := handler.resolveSkillUploadLimitMB(ctx, nil); got != config.MaxSkillMaxUploadSizeMB { + t.Fatalf("clamped tenant limit = %d, want %d", got, config.MaxSkillMaxUploadSizeMB) + } +} + +func TestUploadRejectsZipAboveConfiguredLimit(t *testing.T) { + handler, _, ctx, _ := newTestUploadHandler(t) + handler.SetUploadLimitConfig(config.SkillsConfig{MaxUploadSizeMB: 1}) + + req := newZipUploadRequestWithBinary(t, ctx, map[string][]byte{ + "SKILL.md": []byte(skillMarkdown("Large Skill", "large-skill")), + "large.bin": bytes.Repeat([]byte("x"), (1<<20)+1), + }) + w := httptest.NewRecorder() + handler.handleUpload(w, req) + + if w.Code != http.StatusBadRequest { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } + if !strings.Contains(w.Body.String(), "exceeds") || !strings.Contains(w.Body.String(), "1 MB") { + t.Fatalf("body = %s, want configured upload limit error", w.Body.String()) + } +} + +func TestUploadAllowsFrontmatterLimitAboveGlobalDefault(t *testing.T) { + handler, _, ctx, _ := newTestUploadHandler(t) + handler.SetUploadLimitConfig(config.SkillsConfig{MaxUploadSizeMB: 1}) + skillMD := "---\nname: Video Skill\nslug: video-skill\nmax_upload_size_mb: 2\n---\nSkill body\n" + + req := newZipUploadRequestWithBinary(t, ctx, map[string][]byte{ + "SKILL.md": []byte(skillMD), + "large.bin": bytes.Repeat([]byte("x"), (1<<20)+(128<<10)), + }) + w := httptest.NewRecorder() + handler.handleUpload(w, req) + + if w.Code != http.StatusCreated { + t.Fatalf("status = %d, body = %s", w.Code, w.Body.String()) + } +} + func newTestUploadHandler(t *testing.T) (*SkillsHandler, *skillManageStoreStub, context.Context, string) { t.Helper() @@ -738,6 +799,43 @@ func newZipUploadRequestWithManagers(t *testing.T, ctx context.Context, files ma return req.WithContext(ctx) } +func newZipUploadRequestWithBinary(t *testing.T, ctx context.Context, files map[string][]byte) *http.Request { + t.Helper() + + var zipBuf bytes.Buffer + zw := zip.NewWriter(&zipBuf) + for name, content := range files { + header := &zip.FileHeader{Name: name, Method: zip.Store} + w, err := zw.CreateHeader(header) + if err != nil { + t.Fatalf("zip create %s: %v", name, err) + } + if _, err := w.Write(content); err != nil { + t.Fatalf("zip write %s: %v", name, err) + } + } + if err := zw.Close(); err != nil { + t.Fatalf("zip close: %v", err) + } + + var body bytes.Buffer + mw := multipart.NewWriter(&body) + part, err := mw.CreateFormFile("file", "skill.zip") + if err != nil { + t.Fatalf("multipart file: %v", err) + } + if _, err := part.Write(zipBuf.Bytes()); err != nil { + t.Fatalf("multipart write: %v", err) + } + if err := mw.Close(); err != nil { + t.Fatalf("multipart close: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/skills/upload", &body) + req.Header.Set("Content-Type", mw.FormDataContentType()) + return req.WithContext(ctx) +} + func newUploadManagerIDsFormRequest(t *testing.T, raw string) *http.Request { t.Helper() @@ -770,6 +868,34 @@ type skillManageStoreStub struct { lastUpdates map[uuid.UUID]map[string]any } +type skillUploadSystemConfigStore struct { + data map[string]string +} + +func (s *skillUploadSystemConfigStore) Get(_ context.Context, key string) (string, error) { + if v, ok := s.data[key]; ok { + return v, nil + } + return "", errors.New("not found") +} + +func (s *skillUploadSystemConfigStore) Set(_ context.Context, key, value string) error { + if s.data == nil { + s.data = map[string]string{} + } + s.data[key] = value + return nil +} + +func (s *skillUploadSystemConfigStore) Delete(_ context.Context, key string) error { + delete(s.data, key) + return nil +} + +func (s *skillUploadSystemConfigStore) List(_ context.Context) (map[string]string, error) { + return s.data, nil +} + type skillGrantCall struct { SkillID uuid.UUID AgentID uuid.UUID diff --git a/internal/http/skills_zip_helpers.go b/internal/http/skills_zip_helpers.go new file mode 100644 index 00000000..6a798400 --- /dev/null +++ b/internal/http/skills_zip_helpers.go @@ -0,0 +1,46 @@ +package http + +import ( + "archive/zip" + "errors" + "io" + "os" +) + +const maxSkillMarkdownBytes = 100 << 10 + +var errSkillMarkdownTooLarge = errors.New("SKILL.md exceeds 100 KB limit") + +func readZipFile(f *zip.File) (string, error) { + rc, err := f.Open() + if err != nil { + return "", err + } + defer rc.Close() + + data, err := io.ReadAll(io.LimitReader(rc, maxSkillMarkdownBytes+1)) + if err != nil { + return "", err + } + if len(data) > maxSkillMarkdownBytes { + return "", errSkillMarkdownTooLarge + } + return string(data), nil +} + +func copyZipFileToPath(f *zip.File, destPath string) error { + rc, err := f.Open() + if err != nil { + return err + } + defer rc.Close() + + dst, err := os.OpenFile(destPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) + if err != nil { + return err + } + defer dst.Close() + + _, err = io.Copy(dst, rc) + return err +} diff --git a/internal/http/summoner.go b/internal/http/summoner.go index 540257c1..64f2effb 100644 --- a/internal/http/summoner.go +++ b/internal/http/summoner.go @@ -12,6 +12,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // Summoning event type constants. @@ -48,14 +49,16 @@ type AgentSummoner struct { agents store.AgentStore providerReg *providers.Registry msgBus *bus.MessageBus + usageCaps *usagecaps.Service } // NewAgentSummoner creates a summoner backed by the given stores and provider registry. -func NewAgentSummoner(agents store.AgentStore, providerReg *providers.Registry, msgBus *bus.MessageBus) *AgentSummoner { +func NewAgentSummoner(agents store.AgentStore, providerReg *providers.Registry, msgBus *bus.MessageBus, usageCaps *usagecaps.Service) *AgentSummoner { return &AgentSummoner{ agents: agents, providerReg: providerReg, msgBus: msgBus, + usageCaps: usageCaps, } } @@ -73,6 +76,7 @@ const singleCallTimeout = 300 * time.Second func (s *AgentSummoner) SummonAgent(agentID uuid.UUID, tenantID uuid.UUID, providerName, model, description string) { ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 600*time.Second) defer cancel() + ctx = store.WithAgentID(ctx, agentID) s.ensureBackfillFiles(ctx, agentID) s.emitEvent(agentID, tenantID, SummonEventStarted, "", "") diff --git a/internal/http/summoner_regenerate.go b/internal/http/summoner_regenerate.go index e3237719..126f87db 100644 --- a/internal/http/summoner_regenerate.go +++ b/internal/http/summoner_regenerate.go @@ -14,6 +14,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "github.com/nextlevelbuilder/goclaw/pkg/protocol" ) @@ -23,6 +24,7 @@ import ( func (s *AgentSummoner) RegenerateAgent(agentID uuid.UUID, tenantID uuid.UUID, providerName, model, editPrompt string) { ctx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 300*time.Second) defer cancel() + ctx = store.WithAgentID(ctx, agentID) s.ensureBackfillFiles(ctx, agentID) @@ -121,7 +123,7 @@ func (s *AgentSummoner) generateFiles(ctx context.Context, providerName, model, slog.Info("summoning: calling LLM", "provider", providerName, "model", model, "prompt_len", len(prompt)) - resp, err := provider.Chat(ctx, providers.ChatRequest{ + req := providers.ChatRequest{ Messages: []providers.Message{ {Role: "system", Content: "You are a file generator. Output ONLY the requested XML-tagged files. No extra commentary."}, {Role: "user", Content: prompt}, @@ -133,6 +135,12 @@ func (s *AgentSummoner) generateFiles(ctx context.Context, providerName, model, providers.OptSessionKey: summonSessionKey, providers.OptDisableTools: true, }, + } + resp, err := s.usageCaps.Chat(ctx, provider, req, usagecaps.ChatOptions{ + ProviderName: providerName, + ModelID: model, + Purpose: "agent-summoner", + MaxOutputTokens: 8192, }) if err != nil { return nil, fmt.Errorf("%s: %w", providerName, err) diff --git a/internal/http/usage_caps.go b/internal/http/usage_caps.go new file mode 100644 index 00000000..bf2c3718 --- /dev/null +++ b/internal/http/usage_caps.go @@ -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 +} diff --git a/internal/http/usage_caps_policy_patch_test.go b/internal/http/usage_caps_policy_patch_test.go new file mode 100644 index 00000000..6038370e --- /dev/null +++ b/internal/http/usage_caps_policy_patch_test.go @@ -0,0 +1,35 @@ +package http + +import "testing" + +func TestPolicyPatchFromBodyClearsTokenAndCostLimits(t *testing.T) { + patch, err := policyPatchFromBody([]byte(`{"max_tokens":null,"max_cost_usd":null}`)) + if err != nil { + t.Fatalf("policyPatchFromBody: %v", err) + } + if patch.MaxTokens == nil { + t.Fatal("MaxTokens patch missing") + } + if *patch.MaxTokens != nil { + t.Fatalf("MaxTokens = %v, want nil clear", **patch.MaxTokens) + } + if patch.MaxCostMicros == nil { + t.Fatal("MaxCostMicros patch missing") + } + if *patch.MaxCostMicros != nil { + t.Fatalf("MaxCostMicros = %v, want nil clear", **patch.MaxCostMicros) + } +} + +func TestPolicyPatchFromBodyPreservesProvidedCostUSD(t *testing.T) { + patch, err := policyPatchFromBody([]byte(`{"max_cost_usd":12.5}`)) + if err != nil { + t.Fatalf("policyPatchFromBody: %v", err) + } + if patch.MaxCostMicros == nil || *patch.MaxCostMicros == nil { + t.Fatal("MaxCostMicros patch missing") + } + if got := **patch.MaxCostMicros; got != 12_500_000 { + t.Fatalf("MaxCostMicros = %d, want 12500000", got) + } +} diff --git a/internal/http/usage_caps_test.go b/internal/http/usage_caps_test.go new file mode 100644 index 00000000..3923972a --- /dev/null +++ b/internal/http/usage_caps_test.go @@ -0,0 +1,17 @@ +package http + +import ( + "testing" + + "github.com/google/uuid" +) + +func TestUsageCapPolicyBodyRejectsInvalidUUID(t *testing.T) { + invalid := "not-a-uuid" + if _, err := (policyBody{AgentID: &invalid}).toPolicy(uuid.New()); err == nil { + t.Fatal("toPolicy accepted invalid agent_id") + } + if _, err := (policyBody{ProviderID: &invalid}).toPatch(); err == nil { + t.Fatal("toPatch accepted invalid provider_id") + } +} diff --git a/internal/i18n/catalog_en.go b/internal/i18n/catalog_en.go index c88fd7e5..9ec12d1e 100644 --- a/internal/i18n/catalog_en.go +++ b/internal/i18n/catalog_en.go @@ -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]", diff --git a/internal/i18n/catalog_vi.go b/internal/i18n/catalog_vi.go index fd3e62af..4dc62ea8 100644 --- a/internal/i18n/catalog_vi.go +++ b/internal/i18n/catalog_vi.go @@ -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]", diff --git a/internal/i18n/catalog_zh.go b/internal/i18n/catalog_zh.go index ba5f9c77..85d88923 100644 --- a/internal/i18n/catalog_zh.go +++ b/internal/i18n/catalog_zh.go @@ -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]", diff --git a/internal/i18n/keys.go b/internal/i18n/keys.go index b2b80122..6d31d667 100644 --- a/internal/i18n/keys.go +++ b/internal/i18n/keys.go @@ -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" diff --git a/internal/knowledgegraph/extractor.go b/internal/knowledgegraph/extractor.go index 491c6c2c..addb80df 100644 --- a/internal/knowledgegraph/extractor.go +++ b/internal/knowledgegraph/extractor.go @@ -9,6 +9,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // ExtractionResult holds entities and relations extracted from text. @@ -22,6 +23,7 @@ type Extractor struct { provider providers.Provider model string minConfidence float64 + usageCaps *usagecaps.Service } // NewExtractor creates a new Extractor with the given provider, model, and confidence threshold. @@ -32,6 +34,11 @@ func NewExtractor(provider providers.Provider, model string, minConfidence float return &Extractor{provider: provider, model: model, minConfidence: minConfidence} } +// SetUsageCapService enables cost enforcement for LLM extraction calls. +func (e *Extractor) SetUsageCapService(s *usagecaps.Service) { + e.usageCaps = s +} + const maxChunkChars = 12000 // Extract calls the LLM to extract entities and relations from text. @@ -72,7 +79,11 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR }, } - resp, err := e.provider.Chat(ctx, req) + resp, err := e.usageCaps.Chat(ctx, e.provider, req, usagecaps.ChatOptions{ + ModelID: e.model, + Purpose: "knowledge-graph-extract", + MaxOutputTokens: 8192, + }) if err != nil { return nil, fmt.Errorf("kg extraction LLM call: %w", err) } @@ -85,7 +96,11 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR text = text[:retryMaxChars] + "\n\n[...truncated]" } req.Messages[1].Content = text - resp, err = e.provider.Chat(ctx, req) + resp, err = e.usageCaps.Chat(ctx, e.provider, req, usagecaps.ChatOptions{ + ModelID: e.model, + Purpose: "knowledge-graph-extract-retry", + MaxOutputTokens: 8192, + }) if err != nil { return nil, fmt.Errorf("kg extraction LLM retry: %w", err) } @@ -286,4 +301,3 @@ func stripCodeBlock(s string) string { } return strings.TrimSpace(s) } - diff --git a/internal/media/store.go b/internal/media/store.go index 162871b6..f9cada68 100644 --- a/internal/media/store.go +++ b/internal/media/store.go @@ -109,6 +109,12 @@ func ExtFromMime(mime string) string { return ".wav" case strings.HasPrefix(mime, "application/pdf"): return ".pdf" + case strings.HasPrefix(mime, "application/zip"), strings.HasPrefix(mime, "application/x-zip-compressed"): + return ".zip" + case strings.HasPrefix(mime, "application/x-tar"): + return ".tar" + case strings.HasPrefix(mime, "application/gzip"), strings.HasPrefix(mime, "application/x-gzip"): + return ".gz" case mime == "application/vnd.openxmlformats-officedocument.wordprocessingml.document": return ".docx" case mime == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": diff --git a/internal/media/store_test.go b/internal/media/store_test.go new file mode 100644 index 00000000..e8d1a12d --- /dev/null +++ b/internal/media/store_test.go @@ -0,0 +1,23 @@ +package media + +import "testing" + +func TestExtFromMimeArchiveTypes(t *testing.T) { + tests := []struct { + mime string + want string + }{ + {mime: "application/zip", want: ".zip"}, + {mime: "application/x-zip-compressed", want: ".zip"}, + {mime: "application/x-tar", want: ".tar"}, + {mime: "application/gzip", want: ".gz"}, + } + + for _, tt := range tests { + t.Run(tt.mime, func(t *testing.T) { + if got := ExtFromMime(tt.mime); got != tt.want { + t.Fatalf("ExtFromMime(%q) = %q, want %q", tt.mime, got, tt.want) + } + }) + } +} diff --git a/internal/providers/anthropic_stream.go b/internal/providers/anthropic_stream.go index f9936c32..0d2451c0 100644 --- a/internal/providers/anthropic_stream.go +++ b/internal/providers/anthropic_stream.go @@ -150,7 +150,7 @@ func (p *AnthropicProvider) ChatStream(ctx context.Context, req ChatRequest, onC } if err := sse.Err(); err != nil { - return nil, fmt.Errorf("anthropic stream read error: %w", err) + return result, fmt.Errorf("anthropic stream read error: %w", err) } // Parse accumulated tool call JSON arguments diff --git a/internal/providers/codex.go b/internal/providers/codex.go index e876ff20..8687cbb5 100644 --- a/internal/providers/codex.go +++ b/internal/providers/codex.go @@ -168,7 +168,7 @@ func (p *CodexProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk } if err := sse.Err(); err != nil { - return nil, fmt.Errorf("%s: stream read error: %w", p.name, err) + return result, fmt.Errorf("%s: stream read error: %w", p.name, err) } // Assemble generated images from image accumulator into ChatResponse. diff --git a/internal/providers/model_fallback.go b/internal/providers/model_fallback.go index 9df97e03..2f9310c7 100644 --- a/internal/providers/model_fallback.go +++ b/internal/providers/model_fallback.go @@ -21,6 +21,13 @@ type ModelFallbackProvider struct { maxAttempts int } +type FallbackCallInfo struct { + Streamed bool +} + +type FallbackAfterCall func(*ChatResponse, error, FallbackCallInfo) +type FallbackBeforeCall func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (after FallbackAfterCall, err error) + func NewModelFallbackProvider(primary FallbackCandidate, fallbacks []FallbackCandidate, maxAttempts int, cooldownEnabled bool) *ModelFallbackProvider { var tracker *CooldownTracker if cooldownEnabled { @@ -64,6 +71,22 @@ func (p *ModelFallbackProvider) Chat(ctx context.Context, req ChatRequest) (*Cha }) } +func (p *ModelFallbackProvider) ChatWithHook(ctx context.Context, req ChatRequest, before FallbackBeforeCall) (*ChatResponse, error) { + return p.runOrdered(ctx, req, func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (*ChatResponse, error) { + nextReq := req + nextReq.Model = entry.Model + after, err := before(ctx, entry, nextReq) + if err != nil { + return nil, err + } + resp, err := entry.Provider.Chat(ctx, nextReq) + if after != nil { + after(resp, err, FallbackCallInfo{}) + } + return resp, err + }) +} + func (p *ModelFallbackProvider) ChatStream(ctx context.Context, req ChatRequest, onChunk func(StreamChunk)) (*ChatResponse, error) { return p.runOrdered(ctx, req, func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (*ChatResponse, error) { nextReq := req @@ -82,6 +105,31 @@ func (p *ModelFallbackProvider) ChatStream(ctx context.Context, req ChatRequest, }) } +func (p *ModelFallbackProvider) ChatStreamWithHook(ctx context.Context, req ChatRequest, onChunk func(StreamChunk), before FallbackBeforeCall) (*ChatResponse, error) { + return p.runOrdered(ctx, req, func(ctx context.Context, entry FallbackCandidate, req ChatRequest) (*ChatResponse, error) { + nextReq := req + nextReq.Model = entry.Model + after, err := before(ctx, entry, nextReq) + if err != nil { + return nil, err + } + streamed := false + resp, err := entry.Provider.ChatStream(ctx, nextReq, func(chunk StreamChunk) { + if chunk.Content != "" || chunk.Thinking != "" || len(chunk.Images) > 0 { + streamed = true + } + onChunk(chunk) + }) + if after != nil { + after(resp, err, FallbackCallInfo{Streamed: streamed}) + } + if streamed && err != nil { + return nil, noFallbackAfterStreamError{err: err} + } + return resp, err + }) +} + func (p *ModelFallbackProvider) runOrdered( ctx context.Context, req ChatRequest, diff --git a/internal/providers/model_fallback_test.go b/internal/providers/model_fallback_test.go index 5bea45bb..a891374d 100644 --- a/internal/providers/model_fallback_test.go +++ b/internal/providers/model_fallback_test.go @@ -94,6 +94,33 @@ func TestModelFallbackProviderDoesNotFallbackAfterStreamChunk(t *testing.T) { } } +func TestModelFallbackProviderChatStreamWithHookReportsStreamedChunks(t *testing.T) { + streamErr := &HTTPError{Status: 429, Body: "rate limited"} + primary := &testFallbackProvider{ + name: "primary", + model: "primary-model", + streamErr: streamErr, + } + provider := NewModelFallbackProvider(FallbackCandidate{ + ProviderName: "primary", + Provider: primary, + Model: "primary-model", + }, nil, 1, false) + + var streamed bool + _, err := provider.ChatStreamWithHook(context.Background(), ChatRequest{}, func(StreamChunk) {}, func(context.Context, FallbackCandidate, ChatRequest) (FallbackAfterCall, error) { + return func(_ *ChatResponse, _ error, info FallbackCallInfo) { + streamed = info.Streamed + }, nil + }) + if err == nil { + t.Fatal("ChatStreamWithHook() error = nil, want stream error") + } + if !streamed { + t.Fatal("FallbackCallInfo.Streamed = false, want true after partial stream") + } +} + func TestModelFallbackProviderFallsBackToSameModelOnDifferentProvider(t *testing.T) { primary := &testFallbackProvider{ name: "primary", diff --git a/internal/providers/openai_chat.go b/internal/providers/openai_chat.go index b3991990..c5b1f960 100644 --- a/internal/providers/openai_chat.go +++ b/internal/providers/openai_chat.go @@ -108,13 +108,19 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun PromptTokens: chunk.Usage.PromptTokens, CompletionTokens: chunk.Usage.CompletionTokens, TotalTokens: chunk.Usage.TotalTokens, + RequestCount: 1, } if chunk.Usage.PromptTokensDetails != nil { result.Usage.CacheReadTokens = chunk.Usage.PromptTokensDetails.CachedTokens + result.Usage.CacheCreationTokens = chunk.Usage.PromptTokensDetails.CacheWriteTokens + result.Usage.PromptTokensIncludeCachedSegments = true } if chunk.Usage.CompletionTokensDetails != nil && chunk.Usage.CompletionTokensDetails.ReasoningTokens > 0 { result.Usage.ThinkingTokens = chunk.Usage.CompletionTokensDetails.ReasoningTokens } + if chunk.Usage.ServerToolUse != nil { + result.Usage.WebSearchCount = chunk.Usage.ServerToolUse.WebSearchRequests + } } if len(chunk.Choices) == 0 { @@ -181,7 +187,7 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun // Check for scanner errors (timeout, connection reset, etc.) if err := sse.Err(); err != nil { - return nil, fmt.Errorf("%s: stream read error: %w", p.name, err) + return result, fmt.Errorf("%s: stream read error: %w", p.name, err) } // Parse accumulated tool call arguments diff --git a/internal/providers/openai_http.go b/internal/providers/openai_http.go index 896021e0..99b269a8 100644 --- a/internal/providers/openai_http.go +++ b/internal/providers/openai_http.go @@ -125,13 +125,19 @@ func (p *OpenAIProvider) parseResponse(resp *openAIResponse) *ChatResponse { PromptTokens: resp.Usage.PromptTokens, CompletionTokens: resp.Usage.CompletionTokens, TotalTokens: resp.Usage.TotalTokens, + RequestCount: 1, } if resp.Usage.PromptTokensDetails != nil { result.Usage.CacheReadTokens = resp.Usage.PromptTokensDetails.CachedTokens + result.Usage.CacheCreationTokens = resp.Usage.PromptTokensDetails.CacheWriteTokens + result.Usage.PromptTokensIncludeCachedSegments = true } if resp.Usage.CompletionTokensDetails != nil && resp.Usage.CompletionTokensDetails.ReasoningTokens > 0 { result.Usage.ThinkingTokens = resp.Usage.CompletionTokensDetails.ReasoningTokens } + if resp.Usage.ServerToolUse != nil { + result.Usage.WebSearchCount = resp.Usage.ServerToolUse.WebSearchRequests + } } return result diff --git a/internal/providers/openai_types.go b/internal/providers/openai_types.go index 57361e1c..48389cfd 100644 --- a/internal/providers/openai_types.go +++ b/internal/providers/openai_types.go @@ -49,18 +49,25 @@ type openAIUsage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` TotalTokens int `json:"total_tokens"` + Cost float64 `json:"cost,omitempty"` PromptTokensDetails *openAIPromptDetails `json:"prompt_tokens_details,omitempty"` CompletionTokensDetails *openAICompletionDetails `json:"completion_tokens_details,omitempty"` + ServerToolUse *openAIServerToolUse `json:"server_tool_use,omitempty"` } type openAIPromptDetails struct { - CachedTokens int `json:"cached_tokens"` + CachedTokens int `json:"cached_tokens"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` } type openAICompletionDetails struct { ReasoningTokens int `json:"reasoning_tokens"` } +type openAIServerToolUse struct { + WebSearchRequests int `json:"web_search_requests,omitempty"` +} + // Streaming types type openAIStreamChunk struct { diff --git a/internal/providers/types.go b/internal/providers/types.go index 555a2bf0..d7f03413 100644 --- a/internal/providers/types.go +++ b/internal/providers/types.go @@ -186,10 +186,14 @@ type ToolFunctionSchema struct { // Usage tracks token consumption. type Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` - CacheCreationTokens int `json:"cache_creation_input_tokens,omitempty"` - CacheReadTokens int `json:"cache_read_input_tokens,omitempty"` - ThinkingTokens int `json:"thinking_tokens,omitempty"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + CacheCreationTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadTokens int `json:"cache_read_input_tokens,omitempty"` + PromptTokensIncludeCachedSegments bool `json:"prompt_tokens_include_cached_segments,omitempty"` + ThinkingTokens int `json:"thinking_tokens,omitempty"` + RequestCount int `json:"request_count,omitempty"` + ImageCount int `json:"image_count,omitempty"` + WebSearchCount int `json:"web_search_count,omitempty"` } diff --git a/internal/store/pg/agents.go b/internal/store/pg/agents.go index 234cf1a6..e5a19437 100644 --- a/internal/store/pg/agents.go +++ b/internal/store/pg/agents.go @@ -133,6 +133,12 @@ func (s *PGAgentStore) Create(ctx context.Context, agent *store.AgentData) error if err != nil { return err } + if agent.BudgetMonthlyCents != nil { + if err := s.syncAgentBudgetUsageCap(ctx, tenantID, agent.ID, agent.BudgetMonthlyCents); err != nil { + _, _ = s.db.ExecContext(ctx, "DELETE FROM agents WHERE id = $1", agent.ID) + return err + } + } // Generate embedding for new agent with frontmatter if agent.Frontmatter != "" && s.embProvider != nil { @@ -189,6 +195,22 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str if len(updates) == 0 { return nil } + var err error + var budgetCents *int + rawBudget, syncBudget := updates["budget_monthly_cents"] + if syncBudget { + budgetCents, err = budgetCentsFromUpdate(rawBudget) + if err != nil { + return err + } + } + var budgetTenantID uuid.UUID + if syncBudget { + budgetTenantID, err = s.agentTenantID(ctx, id) + if err != nil { + return err + } + } // Coerce NOT NULL columns: null → default to prevent constraint violations. // Promoted TEXT columns (migration 000037): null → empty string. @@ -248,6 +270,11 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str return err } } + if syncBudget { + if err := s.syncAgentBudgetUsageCap(ctx, budgetTenantID, id, budgetCents); err != nil { + return err + } + } // Regenerate embedding when frontmatter changes if _, hasFrontmatter := updates["frontmatter"]; hasFrontmatter && s.embProvider != nil { @@ -262,6 +289,74 @@ func (s *PGAgentStore) Update(ctx context.Context, id uuid.UUID, updates map[str return nil } +func (s *PGAgentStore) agentTenantID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { + if !store.IsCrossTenant(ctx) { + tid := store.TenantIDFromContext(ctx) + if tid != uuid.Nil { + return tid, nil + } + } + var tenantID uuid.UUID + if err := s.db.QueryRowContext(ctx, "SELECT tenant_id FROM agents WHERE id = $1 AND deleted_at IS NULL", id).Scan(&tenantID); err != nil { + return uuid.Nil, fmt.Errorf("agent not found: %s", id) + } + return tenantID, nil +} + +func budgetCentsFromUpdate(v any) (*int, error) { + if v == nil { + return nil, nil + } + var cents int + switch n := v.(type) { + case int: + cents = n + case int32: + cents = int(n) + case int64: + cents = int(n) + case float64: + if n != float64(int(n)) { + return nil, fmt.Errorf("budget_monthly_cents must be an integer") + } + cents = int(n) + default: + return nil, fmt.Errorf("budget_monthly_cents must be an integer") + } + if cents < 0 { + return nil, fmt.Errorf("budget_monthly_cents must be non-negative") + } + return ¢s, nil +} + +func (s *PGAgentStore) syncAgentBudgetUsageCap(ctx context.Context, tenantID, agentID uuid.UUID, budgetCents *int) error { + if budgetCents == nil || *budgetCents <= 0 { + _, err := s.db.ExecContext(ctx, + `DELETE FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget) + return err + } + costMicros := int64(*budgetCents) * 10000 + _, err := s.db.ExecContext(ctx, ` +INSERT INTO usage_cap_policies ( + tenant_id, agent_id, window_key, max_cost_micros, enabled, priority, source +) VALUES ($1,$2,'month',$3,true,90,$4) +ON CONFLICT (tenant_id, agent_id) WHERE source = 'agent_budget_monthly_cents' +DO UPDATE SET + window_key='month', + provider_id=NULL, + provider_type=NULL, + model_id=NULL, + max_tokens=NULL, + max_cost_micros=EXCLUDED.max_cost_micros, + enabled=true, + priority=EXCLUDED.priority, + updated_at=now()`, + tenantID, agentID, costMicros, store.UsageCapSourceAgentBudget) + return err +} + func isEmptyOrNullJSONUpdate(v any) bool { if v == nil { return true diff --git a/internal/store/pg/agents_update_null_coerce_test.go b/internal/store/pg/agents_update_null_coerce_test.go index 3148b68e..5782d545 100644 --- a/internal/store/pg/agents_update_null_coerce_test.go +++ b/internal/store/pg/agents_update_null_coerce_test.go @@ -1,8 +1,14 @@ package pg import ( + "context" + "database/sql" "encoding/json" + "errors" "testing" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" ) func TestIsEmptyOrNullJSONUpdate(t *testing.T) { @@ -30,3 +36,93 @@ func TestIsEmptyOrNullJSONUpdate(t *testing.T) { }) } } + +func TestPGAgentStoreSyncsMonthlyBudgetUsageCap(t *testing.T) { + db := hooksTestDB(t) + tenantID, _ := seedTenantAndAgent(t, db) + ctx := tenantScopedCtx(tenantID) + agentStore := NewPGAgentStore(db) + budgetCents := 123 + agentID := uuid.New() + agent := &store.AgentData{ + BaseModel: store.BaseModel{ID: agentID}, + TenantID: tenantID, AgentKey: "budget-agent-" + agentID.String(), + OwnerID: "owner", Provider: "openai", Model: "gpt-test", + AgentType: store.AgentTypePredefined, Status: store.AgentStatusActive, + BudgetMonthlyCents: &budgetCents, + } + if err := agentStore.Create(ctx, agent); err != nil { + t.Fatalf("Create: %v", err) + } + assertAgentBudgetPolicy(t, db, tenantID, agentID, 1230000) + var generatedPolicyID uuid.UUID + if err := db.QueryRowContext(context.Background(), + `SELECT id FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&generatedPolicyID); err != nil { + t.Fatalf("query generated policy id: %v", err) + } + if err := NewPGUsageCapStore(db).DeleteUsageCapPolicy(ctx, tenantID, generatedPolicyID); err == nil { + t.Fatal("DeleteUsageCapPolicy deleted generated agent budget policy") + } else if !errors.Is(err, store.ErrUsageCapPolicyManaged) { + t.Fatalf("DeleteUsageCapPolicy error = %v, want managed policy error", err) + } + disablePolicy := false + if _, err := NewPGUsageCapStore(db).UpdateUsageCapPolicy(ctx, tenantID, generatedPolicyID, store.UsageCapPolicyPatch{ + Enabled: &disablePolicy, + }); err == nil { + t.Fatal("UpdateUsageCapPolicy updated generated agent budget policy") + } else if !errors.Is(err, store.ErrUsageCapPolicyManaged) { + t.Fatalf("UpdateUsageCapPolicy error = %v, want managed policy error", err) + } + assertAgentBudgetPolicy(t, db, tenantID, agentID, 1230000) + + if err := agentStore.Update(ctx, agentID, map[string]any{"budget_monthly_cents": float64(456)}); err != nil { + t.Fatalf("Update budget: %v", err) + } + assertAgentBudgetPolicy(t, db, tenantID, agentID, 4560000) + + if err := agentStore.Update(ctx, agentID, map[string]any{"budget_monthly_cents": nil}); err != nil { + t.Fatalf("Clear budget: %v", err) + } + var count int + if err := db.QueryRowContext(context.Background(), + `SELECT COUNT(*) FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&count); err != nil { + t.Fatalf("count cleared policy: %v", err) + } + if count != 0 { + t.Fatalf("agent budget policy count after clear = %d, want 0", count) + } +} + +func assertAgentBudgetPolicy(t *testing.T, db rowQueryer, tenantID, agentID uuid.UUID, wantCostMicros int64) { + t.Helper() + var windowKey, source string + var costMicros int64 + var count int + if err := db.QueryRowContext(context.Background(), + `SELECT COUNT(*), COALESCE(MAX(window_key), ''), COALESCE(MAX(source), ''), COALESCE(MAX(max_cost_micros), 0) + FROM usage_cap_policies + WHERE tenant_id=$1 AND agent_id=$2 AND source=$3`, + tenantID, agentID, store.UsageCapSourceAgentBudget).Scan(&count, &windowKey, &source, &costMicros); err != nil { + t.Fatalf("query agent budget policy: %v", err) + } + if count != 1 { + t.Fatalf("agent budget policy count = %d, want 1", count) + } + if windowKey != store.UsageCapWindowMonth { + t.Fatalf("window_key = %q, want month", windowKey) + } + if source != store.UsageCapSourceAgentBudget { + t.Fatalf("source = %q, want %q", source, store.UsageCapSourceAgentBudget) + } + if costMicros != wantCostMicros { + t.Fatalf("max_cost_micros = %d, want %d", costMicros, wantCostMicros) + } +} + +type rowQueryer interface { + QueryRowContext(context.Context, string, ...any) *sql.Row +} diff --git a/internal/store/pg/factory.go b/internal/store/pg/factory.go index 6027e483..bde42a89 100644 --- a/internal/store/pg/factory.go +++ b/internal/store/pg/factory.go @@ -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. diff --git a/internal/store/pg/usage_caps.go b/internal/store/pg/usage_caps.go new file mode 100644 index 00000000..f693cb63 --- /dev/null +++ b/internal/store/pg/usage_caps.go @@ -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 +} diff --git a/internal/store/pg/usage_caps_test.go b/internal/store/pg/usage_caps_test.go new file mode 100644 index 00000000..216ad6de --- /dev/null +++ b/internal/store/pg/usage_caps_test.go @@ -0,0 +1,180 @@ +package pg + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func TestPGUsageCapStoreReserveUsageIdempotent(t *testing.T) { + db := hooksTestDB(t) + tenantID, _ := seedTenantAndAgent(t, db) + usageStore := NewPGUsageCapStore(db) + maxTokens := int64(100) + policy := &store.UsageCapPolicy{ + TenantID: tenantID, Window: store.UsageCapWindowHour, + MaxTokens: &maxTokens, Enabled: true, Priority: 100, + } + if err := usageStore.CreateUsageCapPolicy(context.Background(), policy); err != nil { + t.Fatalf("CreateUsageCapPolicy: %v", err) + } + req := store.UsageReserveRequest{ + UsageCapScope: store.UsageCapScope{TenantID: tenantID}, + ReservationKey: "duplicate-reservation", + EstimatedTokens: 10, + } + for i := range 2 { + if _, err := usageStore.ReserveUsage(context.Background(), req, []store.UsageCapPolicy{*policy}); err != nil { + t.Fatalf("ReserveUsage call %d: %v", i+1, err) + } + } + + var reservedTokens, reservationRows int64 + if err := db.QueryRow(`SELECT COALESCE(SUM(reserved_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&reservedTokens); err != nil { + t.Fatalf("query reserved tokens: %v", err) + } + if err := db.QueryRow(`SELECT COUNT(*) FROM usage_cap_reservations WHERE policy_id=$1 AND reservation_key=$2`, policy.ID, req.ReservationKey).Scan(&reservationRows); err != nil { + t.Fatalf("query reservation rows: %v", err) + } + if reservedTokens != 10 { + t.Fatalf("reserved_tokens = %d, want 10", reservedTokens) + } + if reservationRows != 1 { + t.Fatalf("reservation rows = %d, want 1", reservationRows) + } + + var wg sync.WaitGroup + errs := make(chan error, 2) + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + errs <- usageStore.ReconcileUsage(context.Background(), store.UsageReconcileRequest{ + ReservationKey: req.ReservationKey, + ActualTokens: 7, + Status: "reconciled", + }) + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("ReconcileUsage: %v", err) + } + } + var usedTokens int64 + if err := db.QueryRow(`SELECT COALESCE(SUM(used_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&usedTokens); err != nil { + t.Fatalf("query used tokens: %v", err) + } + if err := db.QueryRow(`SELECT COALESCE(SUM(reserved_tokens),0) FROM usage_cap_counters WHERE policy_id=$1`, policy.ID).Scan(&reservedTokens); err != nil { + t.Fatalf("query reserved tokens after reconcile: %v", err) + } + if usedTokens != 7 { + t.Fatalf("used_tokens = %d, want 7", usedTokens) + } + if reservedTokens != 0 { + t.Fatalf("reserved_tokens after reconcile = %d, want 0", reservedTokens) + } +} + +func TestPGUsageCapStoreRejectsCrossTenantRefs(t *testing.T) { + db := hooksTestDB(t) + tenantA, agentA := seedTenantAndAgent(t, db) + tenantB, _ := seedTenantAndAgent(t, db) + usageStore := NewPGUsageCapStore(db) + + policy := &store.UsageCapPolicy{ + TenantID: tenantB, AgentID: &agentA, Window: store.UsageCapWindowDay, + MaxTokens: int64PtrPG(100), Enabled: true, + } + if err := usageStore.CreateUsageCapPolicy(context.Background(), policy); err == nil { + t.Fatal("CreateUsageCapPolicy accepted agent_id from another tenant") + } + + providerID := uuid.New() + if _, err := db.Exec( + `INSERT INTO llm_providers (id, tenant_id, name, provider_type, api_key, enabled) + VALUES ($1,$2,$3,'openrouter','sk-test',true)`, + providerID, tenantA, "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 } diff --git a/internal/store/pg/usage_pricing.go b/internal/store/pg/usage_pricing.go new file mode 100644 index 00000000..64ef138e --- /dev/null +++ b/internal/store/pg/usage_pricing.go @@ -0,0 +1,362 @@ +package pg + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "math/big" + "strconv" + "strings" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +type PGUsageCapStore struct { + db *sql.DB +} + +func NewPGUsageCapStore(db *sql.DB) *PGUsageCapStore { + return &PGUsageCapStore{db: db} +} + +func (s *PGUsageCapStore) UpsertPricingCatalog(ctx context.Context, entries []store.UsagePricingCatalogEntry) (int, error) { + const q = ` +INSERT INTO usage_pricing_catalog ( + model_id, canonical_model_id, raw_pricing, raw_model, + input_price, output_price, cache_read_price, cache_write_price, + reasoning_price, request_price, image_price, web_search_price, synced_at +) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13) +ON CONFLICT (model_id) DO UPDATE SET + canonical_model_id = EXCLUDED.canonical_model_id, + raw_pricing = EXCLUDED.raw_pricing, + raw_model = EXCLUDED.raw_model, + input_price = EXCLUDED.input_price, + output_price = EXCLUDED.output_price, + cache_read_price = EXCLUDED.cache_read_price, + cache_write_price = EXCLUDED.cache_write_price, + reasoning_price = EXCLUDED.reasoning_price, + request_price = EXCLUDED.request_price, + image_price = EXCLUDED.image_price, + web_search_price = EXCLUDED.web_search_price, + synced_at = EXCLUDED.synced_at, + updated_at = now()` + for _, e := range entries { + if strings.TrimSpace(e.ModelID) == "" { + continue + } + if err := validateUsagePricingFields(e.Pricing); err != nil { + return 0, err + } + if len(e.RawPricing) == 0 { + e.RawPricing = json.RawMessage(`{}`) + } + if len(e.RawModel) == 0 { + e.RawModel = json.RawMessage(`{}`) + } + if _, err := s.db.ExecContext(ctx, q, + e.ModelID, nullEmpty(e.CanonicalModelID), e.RawPricing, e.RawModel, + priceVal(e.Pricing.Input), priceVal(e.Pricing.Output), + priceVal(e.Pricing.CacheRead), priceVal(e.Pricing.CacheWrite), + priceVal(e.Pricing.Reasoning), priceVal(e.Pricing.Request), + priceVal(e.Pricing.Image), priceVal(e.Pricing.WebSearch), e.SyncedAt, + ); err != nil { + return 0, err + } + } + return len(entries), nil +} + +func (s *PGUsageCapStore) ListPricingCatalog(ctx context.Context, q store.UsagePricingQuery) ([]store.UsagePricingCatalogEntry, error) { + limit := q.Limit + if limit <= 0 || limit > 500 { + limit = 100 + } + args := []any{} + where := "TRUE" + if q.ModelID != "" { + args = append(args, "%"+q.ModelID+"%") + where = "model_id ILIKE $1" + } + args = append(args, limit) + rows, err := s.db.QueryContext(ctx, ` +SELECT id, model_id, COALESCE(canonical_model_id,''), raw_pricing, raw_model, + input_price::text, output_price::text, cache_read_price::text, cache_write_price::text, + reasoning_price::text, request_price::text, image_price::text, web_search_price::text, + synced_at, created_at, updated_at +FROM usage_pricing_catalog +WHERE `+where+` +ORDER BY model_id +LIMIT $`+strconv.Itoa(len(args))+``, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []store.UsagePricingCatalogEntry + for rows.Next() { + e, err := scanCatalog(rows) + if err != nil { + return nil, err + } + out = append(out, e) + } + return out, rows.Err() +} + +func (s *PGUsageCapStore) PutPricingOverride(ctx context.Context, o *store.UsagePricingOverride) error { + if o.ID == uuid.Nil { + o.ID = uuid.New() + } + if o.ProviderID == uuid.Nil || o.TenantID == uuid.Nil || o.ModelID == "" { + return errors.New("tenant_id, provider_id, and model_id are required") + } + if err := validateUsagePricingFields(o.Pricing); err != nil { + return err + } + if err := s.validateUsageCapRefs(ctx, o.TenantID, nil, &o.ProviderID); err != nil { + return err + } + const q = ` +INSERT INTO usage_pricing_overrides ( + id, tenant_id, provider_id, provider_type, model_id, + input_price, output_price, cache_read_price, cache_write_price, + reasoning_price, request_price, image_price, web_search_price, enabled +) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) +ON CONFLICT (tenant_id, provider_id, model_id) DO UPDATE SET + provider_type = EXCLUDED.provider_type, + input_price = EXCLUDED.input_price, + output_price = EXCLUDED.output_price, + cache_read_price = EXCLUDED.cache_read_price, + cache_write_price = EXCLUDED.cache_write_price, + reasoning_price = EXCLUDED.reasoning_price, + request_price = EXCLUDED.request_price, + image_price = EXCLUDED.image_price, + web_search_price = EXCLUDED.web_search_price, + enabled = EXCLUDED.enabled, + updated_at = now() +RETURNING id, created_at, updated_at` + return s.db.QueryRowContext(ctx, q, + o.ID, o.TenantID, o.ProviderID, o.ProviderType, o.ModelID, + priceVal(o.Pricing.Input), priceVal(o.Pricing.Output), + priceVal(o.Pricing.CacheRead), priceVal(o.Pricing.CacheWrite), + priceVal(o.Pricing.Reasoning), priceVal(o.Pricing.Request), + priceVal(o.Pricing.Image), priceVal(o.Pricing.WebSearch), o.Enabled, + ).Scan(&o.ID, &o.CreatedAt, &o.UpdatedAt) +} + +func (s *PGUsageCapStore) ListPricingOverrides(ctx context.Context, q store.UsagePricingQuery) ([]store.UsagePricingOverride, error) { + args := []any{q.TenantID} + where := "tenant_id = $1" + if q.ProviderID != uuid.Nil { + args = append(args, q.ProviderID) + where += " AND provider_id = $" + itoa(len(args)) + } + rows, err := s.db.QueryContext(ctx, overrideSelectSQL+" WHERE "+where+" ORDER BY updated_at DESC", args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []store.UsagePricingOverride + for rows.Next() { + o, err := scanOverride(rows) + if err != nil { + return nil, err + } + out = append(out, o) + } + return out, rows.Err() +} + +func (s *PGUsageCapStore) DeletePricingOverride(ctx context.Context, tenantID, id uuid.UUID) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM usage_pricing_overrides WHERE tenant_id=$1 AND id=$2`, tenantID, id) + return err +} + +func (s *PGUsageCapStore) ResolvePricing(ctx context.Context, tenantID, providerID uuid.UUID, providerName, providerType, modelID string) (*store.ResolvedUsagePricing, error) { + candidates := usagePricingModelCandidates(providerName, providerType, modelID) + if tenantID != uuid.Nil && providerID != uuid.Nil { + for _, candidate := range candidates { + row := s.db.QueryRowContext(ctx, overrideSelectSQL+` WHERE tenant_id=$1 AND provider_id=$2 AND model_id=$3 AND enabled=true`, tenantID, providerID, candidate) + if o, err := scanOverride(row); err == nil { + return &store.ResolvedUsagePricing{ModelID: o.ModelID, ProviderID: providerID, ProviderType: providerType, Source: "override", Pricing: o.Pricing, OverrideID: o.ID}, nil + } else if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + } + } + for _, candidate := range candidates { + row := s.db.QueryRowContext(ctx, ` +SELECT id, model_id, COALESCE(canonical_model_id,''), raw_pricing, raw_model, + input_price::text, output_price::text, cache_read_price::text, cache_write_price::text, + reasoning_price::text, request_price::text, image_price::text, web_search_price::text, + synced_at, created_at, updated_at +FROM usage_pricing_catalog WHERE model_id=$1 OR canonical_model_id=$1 LIMIT 1`, candidate) + e, err := scanCatalog(row) + if err == nil { + return &store.ResolvedUsagePricing{ModelID: e.ModelID, ProviderID: providerID, ProviderType: providerType, Source: "catalog", Pricing: e.Pricing, CatalogSynced: &e.SyncedAt}, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + } + return nil, sql.ErrNoRows +} + +func usagePricingModelCandidates(providerName, providerType, modelID string) []string { + modelID = strings.TrimSpace(modelID) + if modelID == "" { + return nil + } + out := []string{modelID} + if strings.Contains(modelID, "/") { + return out + } + for _, prefix := range openRouterProviderPrefixes(providerName, providerType) { + out = appendUniqueString(out, prefix+"/"+modelID) + } + return out +} + +func openRouterProviderPrefixes(providerName, providerType string) []string { + switch providerType { + case store.ProviderAnthropicNative: + return []string{"anthropic"} + case store.ProviderGeminiNative, store.ProviderVertex: + return []string{"google"} + case store.ProviderOpenAICompat: + switch normalizeProviderAlias(providerName) { + case "openai", "azure", "azure-openai", "azure_openai": + return []string{"openai"} + case "anthropic": + return []string{"anthropic"} + case "gemini", "google", "vertex": + return []string{"google"} + } + return nil + case store.ProviderOpenRouter: + return nil + case store.ProviderGroq: + return []string{"groq"} + case store.ProviderDeepSeek: + return []string{"deepseek"} + case store.ProviderMistral: + return []string{"mistralai"} + case store.ProviderXAI: + return []string{"x-ai"} + case store.ProviderMiniMax: + return []string{"minimax"} + case store.ProviderCohere: + return []string{"cohere"} + case store.ProviderPerplexity: + return []string{"perplexity"} + case store.ProviderDashScope: + return []string{"qwen"} + default: + return nil + } +} + +func normalizeProviderAlias(s string) string { + s = strings.ToLower(strings.TrimSpace(s)) + s = strings.ReplaceAll(s, " ", "-") + return s +} + +func appendUniqueString(values []string, next string) []string { + next = strings.TrimSpace(next) + if next == "" { + return values + } + for _, existing := range values { + if existing == next { + return values + } + } + return append(values, next) +} + +const overrideSelectSQL = `SELECT id, tenant_id, provider_id, provider_type, model_id, + input_price::text, output_price::text, cache_read_price::text, cache_write_price::text, + reasoning_price::text, request_price::text, image_price::text, web_search_price::text, + enabled, created_at, updated_at FROM usage_pricing_overrides` + +type scanner interface{ Scan(dest ...any) error } + +func scanOverride(row scanner) (store.UsagePricingOverride, error) { + var o store.UsagePricingOverride + var prices [8]sql.NullString + err := row.Scan(&o.ID, &o.TenantID, &o.ProviderID, &o.ProviderType, &o.ModelID, + &prices[0], &prices[1], &prices[2], &prices[3], &prices[4], &prices[5], &prices[6], &prices[7], + &o.Enabled, &o.CreatedAt, &o.UpdatedAt) + o.Pricing = pricingFromNulls(prices) + return o, err +} + +func scanCatalog(row scanner) (store.UsagePricingCatalogEntry, error) { + var e store.UsagePricingCatalogEntry + var prices [8]sql.NullString + err := row.Scan(&e.ID, &e.ModelID, &e.CanonicalModelID, &e.RawPricing, &e.RawModel, + &prices[0], &prices[1], &prices[2], &prices[3], &prices[4], &prices[5], &prices[6], &prices[7], + &e.SyncedAt, &e.CreatedAt, &e.UpdatedAt) + e.Pricing = pricingFromNulls(prices) + return e, err +} + +func pricingFromNulls(p [8]sql.NullString) store.UsagePricingFields { + return store.UsagePricingFields{ + Input: pricePtr(p[0]), Output: pricePtr(p[1]), + CacheRead: pricePtr(p[2]), CacheWrite: pricePtr(p[3]), + Reasoning: pricePtr(p[4]), Request: pricePtr(p[5]), + Image: pricePtr(p[6]), WebSearch: pricePtr(p[7]), + } +} + +func pricePtr(v sql.NullString) *string { + if !v.Valid { + return nil + } + s := v.String + return &s +} + +func priceVal(v *string) any { + if v == nil || strings.TrimSpace(*v) == "" { + return nil + } + return strings.TrimSpace(*v) +} + +func validateUsagePricingFields(fields store.UsagePricingFields) error { + values := map[string]*string{ + "input": fields.Input, + "output": fields.Output, + "cache_read": fields.CacheRead, + "cache_write": fields.CacheWrite, + "reasoning": fields.Reasoning, + "request": fields.Request, + "image": fields.Image, + "web_search": fields.WebSearch, + } + for name, raw := range values { + if raw == nil || strings.TrimSpace(*raw) == "" { + continue + } + rat, ok := new(big.Rat).SetString(strings.TrimSpace(*raw)) + if !ok { + return fmt.Errorf("invalid %s price", name) + } + if rat.Sign() < 0 { + return fmt.Errorf("%s price must be non-negative", name) + } + } + return nil +} + +func nullEmpty(s string) any { + if strings.TrimSpace(s) == "" { + return nil + } + return s +} diff --git a/internal/store/secure_cli_env.go b/internal/store/secure_cli_env.go new file mode 100644 index 00000000..fc8ded35 --- /dev/null +++ b/internal/store/secure_cli_env.go @@ -0,0 +1,200 @@ +package store + +import ( + "bytes" + "encoding/json" + "fmt" + "sort" + "strings" +) + +const ( + SecureCLIEnvKindSensitive = "sensitive" + SecureCLIEnvKindValue = "value" +) + +// SecureCLIEnvEntry is the stored per-key env representation; legacy KEY:string maps decode as sensitive. +type SecureCLIEnvEntry struct { + Kind string `json:"kind"` + Value string `json:"value"` +} + +// SecureCLIEnvResponseEntry is safe to serialize in admin API responses. +type SecureCLIEnvResponseEntry struct { + Kind string `json:"kind"` + Value *string `json:"value"` + Masked bool `json:"masked"` +} + +func ParseSecureCLIEnv(raw []byte) (map[string]SecureCLIEnvEntry, error) { + if len(bytes.TrimSpace(raw)) == 0 { + return map[string]SecureCLIEnvEntry{}, nil + } + + var payload map[string]json.RawMessage + if err := json.Unmarshal(raw, &payload); err != nil { + return nil, err + } + + env := make(map[string]SecureCLIEnvEntry, len(payload)) + for key, item := range payload { + key = strings.TrimSpace(key) + if key == "" { + continue + } + entry, err := parseSecureCLIEnvEntry(item) + if err != nil { + return nil, fmt.Errorf("%s: %w", key, err) + } + env[key] = entry + } + return env, nil +} + +func parseSecureCLIEnvEntry(raw json.RawMessage) (SecureCLIEnvEntry, error) { + var legacy string + if err := json.Unmarshal(raw, &legacy); err == nil { + return SecureCLIEnvEntry{Kind: SecureCLIEnvKindSensitive, Value: legacy}, nil + } + trimmed := bytes.TrimSpace(raw) + if len(trimmed) > 0 && trimmed[0] != '{' { + value, err := secureCLIEnvValueAsString(raw) + if err != nil { + return SecureCLIEnvEntry{}, err + } + return SecureCLIEnvEntry{Kind: SecureCLIEnvKindSensitive, Value: value}, nil + } + + var obj struct { + Kind string `json:"kind"` + Value json.RawMessage `json:"value"` + } + if err := json.Unmarshal(raw, &obj); err != nil { + return SecureCLIEnvEntry{}, err + } + + kind := strings.ToLower(strings.TrimSpace(obj.Kind)) + if kind == "" { + kind = SecureCLIEnvKindSensitive + } + if kind != SecureCLIEnvKindSensitive && kind != SecureCLIEnvKindValue { + return SecureCLIEnvEntry{}, fmt.Errorf("invalid env kind %q", obj.Kind) + } + + value, err := secureCLIEnvValueAsString(obj.Value) + if err != nil { + return SecureCLIEnvEntry{}, err + } + return SecureCLIEnvEntry{Kind: kind, Value: value}, nil +} + +func secureCLIEnvValueAsString(raw json.RawMessage) (string, error) { + if len(bytes.TrimSpace(raw)) == 0 || bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + return "", nil + } + var s string + if err := json.Unmarshal(raw, &s); err == nil { + return s, nil + } + var b bool + if err := json.Unmarshal(raw, &b); err == nil { + if b { + return "true", nil + } + return "false", nil + } + var f float64 + if err := json.Unmarshal(raw, &f); err == nil { + return fmt.Sprint(f), nil + } + return "", fmt.Errorf("env value must be string, number, bool, or null") +} + +func SerializeSecureCLIEnv(env map[string]SecureCLIEnvEntry) ([]byte, error) { + normalized := make(map[string]SecureCLIEnvEntry, len(env)) + for key, entry := range env { + key = strings.TrimSpace(key) + if key == "" { + continue + } + kind := strings.ToLower(strings.TrimSpace(entry.Kind)) + if kind == "" { + kind = SecureCLIEnvKindSensitive + } + if kind != SecureCLIEnvKindSensitive && kind != SecureCLIEnvKindValue { + return nil, fmt.Errorf("%s: invalid env kind %q", key, entry.Kind) + } + normalized[key] = SecureCLIEnvEntry{Kind: kind, Value: entry.Value} + } + return json.Marshal(normalized) +} + +func FlattenSecureCLIEnv(raw []byte) (map[string]string, error) { + entries, err := ParseSecureCLIEnv(raw) + if err != nil { + return nil, err + } + flat := make(map[string]string, len(entries)) + for key, entry := range entries { + flat[key] = entry.Value + } + return flat, nil +} + +func MergeSecureCLIEnv(existingJSON []byte, incoming json.RawMessage) ([]byte, error) { + existing, err := ParseSecureCLIEnv(existingJSON) + if err != nil { + return nil, fmt.Errorf("parse existing env: %w", err) + } + incomingEntries, err := ParseSecureCLIEnv(incoming) + if err != nil { + return nil, fmt.Errorf("parse incoming env: %w", err) + } + + out := make(map[string]SecureCLIEnvEntry, len(incomingEntries)) + for key, entry := range incomingEntries { + if entry.Kind == SecureCLIEnvKindSensitive && entry.Value == "" { + if prev, ok := existing[key]; ok { + prev.Kind = SecureCLIEnvKindSensitive + out[key] = prev + continue + } + } + out[key] = entry + } + return SerializeSecureCLIEnv(out) +} + +func SecureCLIEnvKeys(raw []byte) []string { + env, err := ParseSecureCLIEnv(raw) + if err != nil { + return []string{} + } + keys := make([]string, 0, len(env)) + for key := range env { + keys = append(keys, key) + } + sort.Strings(keys) + return keys +} + +func SanitizeSecureCLIEnv(env map[string]SecureCLIEnvEntry) map[string]SecureCLIEnvResponseEntry { + out := make(map[string]SecureCLIEnvResponseEntry, len(env)) + for key, entry := range env { + if entry.Kind == SecureCLIEnvKindValue { + value := entry.Value + out[key] = SecureCLIEnvResponseEntry{Kind: SecureCLIEnvKindValue, Value: &value, Masked: false} + continue + } + out[key] = SecureCLIEnvResponseEntry{Kind: SecureCLIEnvKindSensitive, Value: nil, Masked: true} + } + return out +} + +func SanitizeSecureCLIEnvJSON(raw []byte) map[string]SecureCLIEnvResponseEntry { + env, err := ParseSecureCLIEnv(raw) + if err != nil { + return map[string]SecureCLIEnvResponseEntry{} + } + return SanitizeSecureCLIEnv(env) +} diff --git a/internal/store/secure_cli_env_test.go b/internal/store/secure_cli_env_test.go new file mode 100644 index 00000000..2a678b84 --- /dev/null +++ b/internal/store/secure_cli_env_test.go @@ -0,0 +1,102 @@ +package store + +import ( + "encoding/json" + "testing" +) + +func TestParseSecureCLIEnvLegacyMapDefaultsSensitive(t *testing.T) { + env, err := ParseSecureCLIEnv([]byte(`{"TOKEN":"secret","PUBLIC_BASE_URL":"https://goclaw.sh"}`)) + if err != nil { + t.Fatalf("ParseSecureCLIEnv() error = %v", err) + } + if got := env["TOKEN"].Kind; got != SecureCLIEnvKindSensitive { + t.Fatalf("TOKEN kind = %q, want %q", got, SecureCLIEnvKindSensitive) + } + if got := env["TOKEN"].Value; got != "secret" { + t.Fatalf("TOKEN value = %q", got) + } + if got := env["PUBLIC_BASE_URL"].Kind; got != SecureCLIEnvKindSensitive { + t.Fatalf("PUBLIC_BASE_URL kind = %q, want default sensitive", got) + } +} + +func TestParseSecureCLIEnvLegacyScalarsDefaultSensitive(t *testing.T) { + env, err := ParseSecureCLIEnv([]byte(`{"MAX_UPLOAD_SIZE_MB":100,"DEBUG":true}`)) + if err != nil { + t.Fatalf("ParseSecureCLIEnv() error = %v", err) + } + if got := env["MAX_UPLOAD_SIZE_MB"]; got.Kind != SecureCLIEnvKindSensitive || got.Value != "100" { + t.Fatalf("MAX_UPLOAD_SIZE_MB = %#v, want sensitive 100", got) + } + if got := env["DEBUG"]; got.Kind != SecureCLIEnvKindSensitive || got.Value != "true" { + t.Fatalf("DEBUG = %#v, want sensitive true", got) + } +} + +func TestSanitizeSecureCLIEnvMasksSensitiveAndReturnsValues(t *testing.T) { + env := map[string]SecureCLIEnvEntry{ + "TOKEN": {Kind: SecureCLIEnvKindSensitive, Value: "secret"}, + "PUBLIC_BASE_URL": {Kind: SecureCLIEnvKindValue, Value: "https://goclaw.sh"}, + } + got := SanitizeSecureCLIEnv(env) + + if got["TOKEN"].Value != nil { + t.Fatalf("sensitive value leaked: %q", *got["TOKEN"].Value) + } + if !got["TOKEN"].Masked { + t.Fatalf("sensitive masked = false") + } + if got["PUBLIC_BASE_URL"].Value == nil || *got["PUBLIC_BASE_URL"].Value != "https://goclaw.sh" { + t.Fatalf("value entry not returned: %#v", got["PUBLIC_BASE_URL"]) + } + if got["PUBLIC_BASE_URL"].Masked { + t.Fatalf("value entry masked = true") + } +} + +func TestMergeSecureCLIEnvPreservesExistingSensitiveOnEmptyValue(t *testing.T) { + existing := []byte(`{"TOKEN":{"kind":"sensitive","value":"old"},"PUBLIC_BASE_URL":{"kind":"value","value":"https://old.example"}}`) + incoming := json.RawMessage(`{"TOKEN":{"kind":"sensitive","value":""},"PUBLIC_BASE_URL":{"kind":"value","value":"https://new.example"}}`) + + merged, err := MergeSecureCLIEnv(existing, incoming) + if err != nil { + t.Fatalf("MergeSecureCLIEnv() error = %v", err) + } + env, err := ParseSecureCLIEnv(merged) + if err != nil { + t.Fatalf("ParseSecureCLIEnv(merged) error = %v", err) + } + if got := env["TOKEN"].Value; got != "old" { + t.Fatalf("TOKEN value = %q, want preserved old", got) + } + if got := env["PUBLIC_BASE_URL"].Value; got != "https://new.example" { + t.Fatalf("PUBLIC_BASE_URL = %q", got) + } + if got := env["PUBLIC_BASE_URL"].Kind; got != SecureCLIEnvKindValue { + t.Fatalf("PUBLIC_BASE_URL kind = %q", got) + } +} + +func TestFlattenSecureCLIEnvSupportsEntryShape(t *testing.T) { + got, err := FlattenSecureCLIEnv([]byte(`{ + "TOKEN":{"kind":"sensitive","value":"secret"}, + "PUBLIC_BASE_URL":{"kind":"value","value":"https://goclaw.sh"} + }`)) + if err != nil { + t.Fatalf("FlattenSecureCLIEnv() error = %v", err) + } + if got["TOKEN"] != "secret" { + t.Fatalf("TOKEN = %q", got["TOKEN"]) + } + if got["PUBLIC_BASE_URL"] != "https://goclaw.sh" { + t.Fatalf("PUBLIC_BASE_URL = %q", got["PUBLIC_BASE_URL"]) + } +} + +func TestParseSecureCLIEnvRejectsInvalidKind(t *testing.T) { + _, err := ParseSecureCLIEnv([]byte(`{"TOKEN":{"kind":"plain","value":"secret"}}`)) + if err == nil { + t.Fatalf("ParseSecureCLIEnv() error = nil, want invalid kind error") + } +} diff --git a/internal/store/secure_cli_store.go b/internal/store/secure_cli_store.go index 2c8f117c..bac2a08d 100644 --- a/internal/store/secure_cli_store.go +++ b/internal/store/secure_cli_store.go @@ -37,6 +37,8 @@ type SecureCLIBinary struct { UserEnv []byte `json:"-" db:"-"` // per-user encrypted env (populated by LookupByBinary LEFT JOIN) // EnvKeys is set by HTTP handlers only (names from decrypted env, no values); not a DB column. EnvKeys []string `json:"env_keys,omitempty" db:"-"` + // Env is set by HTTP handlers only. Sensitive values are masked; value entries are visible. + Env map[string]SecureCLIEnvResponseEntry `json:"env,omitempty" db:"-"` // AgentGrantsSummary is populated by List only — lightweight per-grant summary (no env bytes). AgentGrantsSummary []AgentGrantSummary `json:"agent_grants_summary" db:"-"` } @@ -92,6 +94,8 @@ type SecureCLIAgentGrant struct { EncryptedEnv []byte `json:"-" db:"encrypted_env"` // EnvKeys is populated by HTTP handlers only (sorted key names, no values). Not a DB column. EnvKeys []string `json:"env_keys,omitempty" db:"-"` + // Env is populated by HTTP handlers only for sanitized responses. + Env map[string]SecureCLIEnvResponseEntry `json:"env,omitempty" db:"-"` // EnvSet indicates whether this grant has an env override. Not a DB column. EnvSet bool `json:"env_set" db:"-"` CreatedAt time.Time `json:"created_at" db:"created_at"` diff --git a/internal/store/stores.go b/internal/store/stores.go index 2f7e7112..7a1a9887 100644 --- a/internal/store/stores.go +++ b/internal/store/stores.go @@ -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 } diff --git a/internal/store/usage_caps.go b/internal/store/usage_caps.go new file mode 100644 index 00000000..6848a4c6 --- /dev/null +++ b/internal/store/usage_caps.go @@ -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 +} diff --git a/internal/tools/command_keyword_allowlist.go b/internal/tools/command_keyword_allowlist.go new file mode 100644 index 00000000..b03490eb --- /dev/null +++ b/internal/tools/command_keyword_allowlist.go @@ -0,0 +1,187 @@ +package tools + +import ( + "regexp" + "strconv" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/config" +) + +type commandKeywordAllowAudit struct { + RuleID string + Command string + Subcommand string + Arg string + Keyword string + Reason string +} + +func applyCommandKeywordAllowlist(command string, args []string, rules []config.CommandKeywordAllowlistRule) ([]string, []commandKeywordAllowAudit) { + if len(args) == 0 || len(rules) == 0 { + return args, nil + } + + var out []string + var audits []commandKeywordAllowAudit + for _, rule := range rules { + if !commandKeywordAllowlistRuleEnabled(rule) { + continue + } + if normalizeBinaryName(rule.Command) != normalizeBinaryName(command) { + continue + } + subcommand, argStart, ok := matchCommandKeywordSubcommand(args, rule.Subcommands) + if !ok { + continue + } + argNames := commandKeywordSet(rule.Args) + argPositions := commandKeywordPositionSet(rule.ArgPositions) + keywords := commandKeywordSet(rule.Keywords) + if (len(argNames) == 0 && len(argPositions) == 0) || len(keywords) == 0 { + continue + } + if len(argPositions) > 0 && subcommand == "" { + continue + } + if out == nil { + out = slicesClone(args) + } + for i := argStart; i < len(out); i++ { + argName, valueIndex, inlineValue, ok := commandKeywordArgValue(out, i, argNames) + if !ok { + if _, positional := argPositions[i-argStart]; positional { + argName = "argPositions[" + strconv.Itoa(i-argStart) + "]" + valueIndex = i + inlineValue = out[i] + ok = true + } + } + if !ok { + continue + } + value := inlineValue + if valueIndex != i { + value = out[valueIndex] + } + masked, hits := maskCommandKeywords(value, keywords) + if len(hits) == 0 { + continue + } + if valueIndex == i { + prefix := out[i][:strings.Index(out[i], "=")+1] + out[i] = prefix + masked + } else { + out[valueIndex] = masked + i = valueIndex + } + for _, keyword := range hits { + audits = append(audits, commandKeywordAllowAudit{ + RuleID: rule.ID, + Command: normalizeBinaryName(command), + Subcommand: subcommand, + Arg: argName, + Keyword: keyword, + Reason: rule.Reason, + }) + } + } + } + if out == nil { + return args, audits + } + return out, audits +} + +func commandKeywordAllowlistRuleEnabled(rule config.CommandKeywordAllowlistRule) bool { + return rule.Enabled == nil || *rule.Enabled +} + +func matchCommandKeywordSubcommand(args []string, subcommands []string) (string, int, bool) { + if len(subcommands) == 0 { + return "", 0, true + } + for _, subcommand := range subcommands { + parts := strings.Fields(strings.ToLower(strings.TrimSpace(subcommand))) + if len(parts) == 0 || len(parts) > len(args) { + continue + } + matched := true + for i, part := range parts { + if strings.ToLower(args[i]) != part { + matched = false + break + } + } + if matched { + return strings.Join(parts, " "), len(parts), true + } + } + return "", 0, false +} + +func commandKeywordSet(values []string) map[string]struct{} { + if len(values) == 0 { + return nil + } + set := make(map[string]struct{}, len(values)) + for _, value := range values { + value = strings.ToLower(strings.TrimSpace(value)) + if value != "" { + set[value] = struct{}{} + } + } + return set +} + +func commandKeywordPositionSet(values []int) map[int]struct{} { + if len(values) == 0 { + return nil + } + set := make(map[int]struct{}, len(values)) + for _, value := range values { + if value >= 0 { + set[value] = struct{}{} + } + } + return set +} + +func commandKeywordArgValue(args []string, i int, argNames map[string]struct{}) (string, int, string, bool) { + arg := args[i] + if eq := strings.Index(arg, "="); eq > 0 { + name := strings.ToLower(arg[:eq]) + if _, ok := argNames[name]; ok { + return name, i, arg[eq+1:], true + } + return "", 0, "", false + } + name := strings.ToLower(arg) + if _, ok := argNames[name]; !ok || i+1 >= len(args) { + return "", 0, "", false + } + return name, i + 1, "", true +} + +func maskCommandKeywords(value string, keywords map[string]struct{}) (string, []string) { + if value == "" || len(keywords) == 0 { + return value, nil + } + var hits []string + masked := value + for keyword := range keywords { + re := regexp.MustCompile(`(?i)\b` + regexp.QuoteMeta(keyword) + `\b`) + if !re.MatchString(masked) { + continue + } + hits = append(hits, keyword) + masked = re.ReplaceAllString(masked, "__allowlisted_keyword__") + } + return masked, hits +} + +func slicesClone[T any](src []T) []T { + dst := make([]T, len(src)) + copy(dst, src) + return dst +} diff --git a/internal/tools/credentialed_exec.go b/internal/tools/credentialed_exec.go index b811ca22..44efd721 100644 --- a/internal/tools/credentialed_exec.go +++ b/internal/tools/credentialed_exec.go @@ -39,8 +39,8 @@ var wrapperBinaries = map[string]bool{ // normalizeBinaryName returns the lowercased file base of a binary reference. // Examples: "/usr/bin/gh" → "gh", "./GH" → "gh", " Gh " → "gh". -// Applied at BOTH the gate lookup and lookupCredentialedBinary so the two -// layers agree on identity. (Red Team F5) +// Applied at both the gate lookup and lookupCredentialedBinary so the two +// layers agree on identity. func normalizeBinaryName(s string) string { return filepath.Base(strings.TrimSpace(strings.ToLower(s))) } @@ -381,7 +381,8 @@ func (t *ExecTool) executeCredentialed(ctx context.Context, cred *store.SecureCL } // Step 3: Per-binary deny check (deny_args) - if p := matchesBinaryDeny(args, cred.DenyArgs); p != "" { + denyArgs, allowAudits := applyCommandKeywordAllowlist(binary, args, t.commandKeywordAllowlistSnapshot()) + if p := matchesBinaryDeny(denyArgs, cred.DenyArgs); p != "" { return credentialedDenyError(binary, args, p) } // Per-binary verbose deny check (deny_verbose) — per-arg start-anchored match @@ -389,6 +390,20 @@ func (t *ExecTool) executeCredentialed(ctx context.Context, cred *store.SecureCL if p := matchesBinaryVerbose(args, cred.DenyVerbose); p != "" { return credentialedDenyError(binary, args, p) } + for _, audit := range allowAudits { + slog.Info("security.command_keyword_allowlist", + "binary", audit.Command, + "subcommand", audit.Subcommand, + "arg", audit.Arg, + "keyword", audit.Keyword, + "rule_id", audit.RuleID, + "reason", audit.Reason, + "agent_id", store.AgentIDFromContext(ctx), + "user_id", store.UserIDFromContext(ctx), + "credential_user_id", store.CredentialUserIDFromContext(ctx), + "tenant_id", store.TenantIDFromContext(ctx), + ) + } // Step 4: Decrypt env vars from store (already decrypted by store layer). // Per-user env overrides take priority over binary/grant env. @@ -421,13 +436,15 @@ func mergeCredentialedEnv(cred *store.SecureCLIBinary) (map[string]string, error return envMap, nil } if len(cred.EncryptedEnv) > 0 { - if err := json.Unmarshal(cred.EncryptedEnv, &envMap); err != nil { + baseEnv, err := store.FlattenSecureCLIEnv(cred.EncryptedEnv) + if err != nil { return nil, err } + maps.Copy(envMap, baseEnv) } if len(cred.UserEnv) > 0 { - var userEnvMap map[string]string - if err := json.Unmarshal(cred.UserEnv, &userEnvMap); err != nil { + userEnvMap, err := store.FlattenSecureCLIEnv(cred.UserEnv) + if err != nil { return nil, err } maps.Copy(envMap, userEnvMap) @@ -613,7 +630,7 @@ func (t *ExecTool) lookupCredentialedBinary(ctx context.Context, command string) } // Normalize lookup key so path/case variants (/usr/bin/gh, ./gh, GH) all // resolve to the same registry row. Same helper is used by the gate - // branch in Execute — identity must agree at both layers. (Red Team F5) + // branch in Execute because identity must agree at both layers. normBinary := normalizeBinaryName(binary) // Get agent ID from context for scoped lookup agentID := store.AgentIDFromContext(ctx) diff --git a/internal/tools/credentialed_exec_env_test.go b/internal/tools/credentialed_exec_env_test.go index e164a8ae..91e48cb5 100644 --- a/internal/tools/credentialed_exec_env_test.go +++ b/internal/tools/credentialed_exec_env_test.go @@ -43,3 +43,24 @@ func TestMergeCredentialedEnvFailsClosedOnInvalidUserEnv(t *testing.T) { t.Fatal("expected invalid per-user env JSON to fail closed") } } + +func TestMergeCredentialedEnvFlattensSensitiveValueEntries(t *testing.T) { + binary := &store.SecureCLIBinary{ + EncryptedEnv: []byte(`{ + "TOKEN":{"kind":"sensitive","value":"secret"}, + "PUBLIC_BASE_URL":{"kind":"value","value":"https://goclaw.sh"} + }`), + UserEnv: []byte(`{"PUBLIC_BASE_URL":{"kind":"value","value":"https://user.example"}}`), + } + + env, err := mergeCredentialedEnv(binary) + if err != nil { + t.Fatalf("mergeCredentialedEnv() error = %v", err) + } + if env["TOKEN"] != "secret" { + t.Fatalf("TOKEN = %q", env["TOKEN"]) + } + if env["PUBLIC_BASE_URL"] != "https://user.example" { + t.Fatalf("PUBLIC_BASE_URL = %q", env["PUBLIC_BASE_URL"]) + } +} diff --git a/internal/tools/credentialed_exec_test.go b/internal/tools/credentialed_exec_test.go index 293c5f75..1b2886b0 100644 --- a/internal/tools/credentialed_exec_test.go +++ b/internal/tools/credentialed_exec_test.go @@ -5,6 +5,8 @@ import ( "os" "path/filepath" "testing" + + "github.com/nextlevelbuilder/goclaw/internal/config" ) func TestDetectShellOperators(t *testing.T) { @@ -208,6 +210,113 @@ func TestMatchesBinaryDenyJoinedArgs(t *testing.T) { } } +func TestApplyCommandKeywordAllowlistScopesContentArgs(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`auth\s+`, `repo\s+delete`, `secret\s+`, `token\s+`}) + rules := []config.CommandKeywordAllowlistRule{ + { + ID: "github-content", + Command: "gh", + Subcommands: []string{"issue create", "pr create"}, + Args: []string{"--body", "--title"}, + ArgPositions: []int{0}, + Keywords: []string{"secret", "token"}, + Reason: "GitHub issue and PR prose may discuss security terms.", + }, + } + + tests := []struct { + name string + args []string + wantHit bool + wantAudit int + }{ + { + name: "issue body content allowed", + args: []string{"issue", "create", "--body", "secret rotation details"}, + wantHit: false, + wantAudit: 1, + }, + { + name: "pr title content allowed", + args: []string{"pr", "create", "--title=token handling notes"}, + wantHit: false, + wantAudit: 1, + }, + { + name: "positional content allowed", + args: []string{"issue", "create", "token handling notes"}, + wantHit: false, + wantAudit: 1, + }, + { + name: "command path stays blocked", + args: []string{"secret", "set", "TOKEN"}, + wantHit: true, + }, + { + name: "non-allowlisted arg stays blocked", + args: []string{"issue", "create", "--label", "secret incident"}, + wantHit: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sanitized, audits := applyCommandKeywordAllowlist("gh", tt.args, rules) + got := matchesBinaryDeny(sanitized, ghPatterns) + if (got != "") != tt.wantHit { + t.Fatalf("matchesBinaryDeny(%v) after allowlist = %q, wantHit=%v", sanitized, got, tt.wantHit) + } + if len(audits) != tt.wantAudit { + t.Fatalf("audit count = %d, want %d", len(audits), tt.wantAudit) + } + }) + } +} + +func TestApplyCommandKeywordAllowlistIgnoresDisabledRules(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`secret\s+`}) + enabled := false + rules := []config.CommandKeywordAllowlistRule{ + { + ID: "disabled-github-content", + Command: "gh", + Subcommands: []string{"issue create"}, + Args: []string{"--body"}, + Keywords: []string{"secret"}, + Enabled: &enabled, + }, + } + + sanitized, audits := applyCommandKeywordAllowlist("gh", []string{"issue", "create", "--body", "secret notes"}, rules) + if got := matchesBinaryDeny(sanitized, ghPatterns); got == "" { + t.Fatalf("disabled rule bypassed deny_args; sanitized args = %v", sanitized) + } + if len(audits) != 0 { + t.Fatalf("disabled rule emitted audit records: %v", audits) + } +} + +func TestApplyCommandKeywordAllowlistRequiresSubcommandForPositions(t *testing.T) { + ghPatterns, _ := json.Marshal([]string{`secret\s+`}) + rules := []config.CommandKeywordAllowlistRule{ + { + ID: "unsafe-position", + Command: "gh", + ArgPositions: []int{0}, + Keywords: []string{"secret"}, + }, + } + + sanitized, audits := applyCommandKeywordAllowlist("gh", []string{"secret", "set", "TOKEN"}, rules) + if got := matchesBinaryDeny(sanitized, ghPatterns); got == "" { + t.Fatalf("position rule without subcommand bypassed command-path deny; sanitized args = %v", sanitized) + } + if len(audits) != 0 { + t.Fatalf("position rule without subcommand emitted audit records: %v", audits) + } +} + func TestResolveAndMatchBinaryUsesConfiguredExecutablePath(t *testing.T) { t.Setenv("PATH", "/usr/bin") binDir := t.TempDir() diff --git a/internal/tools/read_audio.go b/internal/tools/read_audio.go index e0a92283..a329dd26 100644 --- a/internal/tools/read_audio.go +++ b/internal/tools/read_audio.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // --- Context helpers for media audio --- @@ -45,12 +46,17 @@ var audioModelDefaults = map[string]string{ type ReadAudioTool struct { registry *providers.Registry mediaLoader MediaPathLoader + usageCaps *usagecaps.Service } func NewReadAudioTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadAudioTool { return &ReadAudioTool{registry: registry, mediaLoader: mediaLoader} } +func (t *ReadAudioTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadAudioTool) Name() string { return "read_audio" } func (t *ReadAudioTool) Description() string { diff --git a/internal/tools/read_audio_resolve.go b/internal/tools/read_audio_resolve.go index 65e27d74..f7f62a1e 100644 --- a/internal/tools/read_audio_resolve.go +++ b/internal/tools/read_audio_resolve.go @@ -93,7 +93,19 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, // providers exposing a /v1/audio/transcriptions endpoint. if isTranscriptionModel(model) { slog.Info("read_audio: using openai transcription API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := openaiTranscriptionCall(ctx, cp.APIKey(), cp.APIBase(), model, data, mime) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("openai transcription call: %w", err) } @@ -103,7 +115,19 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, // Gemini: use File API (inlineData doesn't work for audio). if ptype == "gemini" { slog.Info("read_audio: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 120*time.Second) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("gemini file API: %w", err) } @@ -113,7 +137,19 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, // Native OpenAI chat-audio (gpt-4o-audio-preview etc.): input_audio content part. if ptype == "openai" { slog.Info("read_audio: using openai input_audio API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := openaiAudioCall(ctx, cp.APIKey(), cp.APIBase(), model, prompt, data, mime) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("openai audio call: %w", err) } @@ -128,7 +164,7 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, } slog.Info("read_audio: using chat API fallback", "provider", providerName, "model", model, "size", len(data)) - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -141,7 +177,15 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, "max_tokens": 16384, "temperature": 0.2, }, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("chat API: %w", err) } diff --git a/internal/tools/read_document.go b/internal/tools/read_document.go index 7b3952df..4167d76a 100644 --- a/internal/tools/read_document.go +++ b/internal/tools/read_document.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // textReadableMIMEs are MIME types whose content can be returned directly without LLM analysis. @@ -68,12 +69,17 @@ var documentModelDefaults = map[string]string{ type ReadDocumentTool struct { registry *providers.Registry mediaLoader MediaPathLoader + usageCaps *usagecaps.Service } func NewReadDocumentTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadDocumentTool { return &ReadDocumentTool{registry: registry, mediaLoader: mediaLoader} } +func (t *ReadDocumentTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadDocumentTool) Name() string { return "read_document" } func (t *ReadDocumentTool) Description() string { diff --git a/internal/tools/read_document_resolve.go b/internal/tools/read_document_resolve.go index 47d72cf9..24de1584 100644 --- a/internal/tools/read_document_resolve.go +++ b/internal/tools/read_document_resolve.go @@ -43,8 +43,7 @@ func (t *ReadDocumentTool) resolveDocumentFile(ctx context.Context, mediaID, doc } } if ref == nil { - slog.Warn("read_document: media_id not found, falling back to most recent", "media_id", mediaID) - ref = &refs[len(refs)-1] + return "", "", fmt.Errorf("document media_id %q not found in this conversation", mediaID) } } else { // Use the last (most recent) document ref. @@ -152,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) } @@ -179,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", @@ -189,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) } diff --git a/internal/tools/read_document_resolve_test.go b/internal/tools/read_document_resolve_test.go index bc2ac937..172ac87f 100644 --- a/internal/tools/read_document_resolve_test.go +++ b/internal/tools/read_document_resolve_test.go @@ -66,7 +66,7 @@ func TestResolveDocumentFileMatchesUploadedFilenameAlias(t *testing.T) { } } -func TestResolveDocumentFileInvalidMediaIDFallsBackToMostRecent(t *testing.T) { +func TestResolveDocumentFileInvalidMediaIDReturnsError(t *testing.T) { refs := []providers.MediaRef{ {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/old.pdf", MimeType: "application/pdf"}, {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/latest.pdf", MimeType: "application/pdf"}, @@ -75,6 +75,23 @@ func TestResolveDocumentFileInvalidMediaIDFallsBackToMostRecent(t *testing.T) { tool := NewReadDocumentTool(nil, nil) ctx := WithMediaDocRefs(context.Background(), refs) gotPath, _, err := tool.resolveDocumentFile(ctx, "not-a-real-media-id", "") + if err == nil { + t.Fatalf("resolveDocumentFile returned path %q, want explicit media_id error", gotPath) + } + if !strings.Contains(err.Error(), "not-a-real-media-id") { + t.Fatalf("error = %q, want requested media_id", err.Error()) + } +} + +func TestResolveDocumentFileOmittedMediaIDUsesLastRef(t *testing.T) { + refs := []providers.MediaRef{ + {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/old.pdf", MimeType: "application/pdf"}, + {ID: uuid.NewString(), Kind: "document", Path: "/workspace/.uploads/latest.pdf", MimeType: "application/pdf"}, + } + + tool := NewReadDocumentTool(nil, nil) + ctx := WithMediaDocRefs(context.Background(), refs) + gotPath, _, err := tool.resolveDocumentFile(ctx, "", "") if err != nil { t.Fatalf("resolveDocumentFile returned error: %v", err) } diff --git a/internal/tools/read_image.go b/internal/tools/read_image.go index 9d754d5c..892a6dfb 100644 --- a/internal/tools/read_image.go +++ b/internal/tools/read_image.go @@ -10,6 +10,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // --- Context helpers for media images --- @@ -46,13 +47,18 @@ var visionModelDefaults = map[string]string{ // ReadImageTool uses a vision-capable provider to describe images attached to the current message. type ReadImageTool struct { - registry *providers.Registry + registry *providers.Registry + usageCaps *usagecaps.Service } func NewReadImageTool(registry *providers.Registry) *ReadImageTool { return &ReadImageTool{registry: registry} } +func (t *ReadImageTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadImageTool) Name() string { return "read_image" } func (t *ReadImageTool) Description() string { @@ -152,7 +158,7 @@ func (t *ReadImageTool) callProvider(ctx context.Context, cp credentialProvider, opts["disable_tools"] = true } - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -162,7 +168,15 @@ func (t *ReadImageTool) callProvider(ctx context.Context, cp credentialProvider, }, Model: model, Options: opts, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("vision provider error: %w", err) } diff --git a/internal/tools/read_video.go b/internal/tools/read_video.go index b6be69fb..46df93c4 100644 --- a/internal/tools/read_video.go +++ b/internal/tools/read_video.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/nextlevelbuilder/goclaw/internal/providers" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // --- Context helpers for media video --- @@ -45,12 +46,17 @@ var videoModelDefaults = map[string]string{ type ReadVideoTool struct { registry *providers.Registry mediaLoader MediaPathLoader + usageCaps *usagecaps.Service } func NewReadVideoTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadVideoTool { return &ReadVideoTool{registry: registry, mediaLoader: mediaLoader} } +func (t *ReadVideoTool) SetUsageCapService(svc *usagecaps.Service) { + t.usageCaps = svc +} + func (t *ReadVideoTool) Name() string { return "read_video" } func (t *ReadVideoTool) Description() string { diff --git a/internal/tools/read_video_resolve.go b/internal/tools/read_video_resolve.go index 8a67d925..f912fb96 100644 --- a/internal/tools/read_video_resolve.go +++ b/internal/tools/read_video_resolve.go @@ -70,7 +70,19 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, ptype := GetParamString(params, "_provider_type", providerTypeFromName(providerName)) if cp != nil && ptype == "gemini" { slog.Info("read_video: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + chatReq := providers.ChatRequest{ + Messages: []providers.Message{{Role: "user", Content: prompt}}, + Model: model, + Options: map[string]any{"max_tokens": 16384}, + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } resp, err := geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 180*time.Second) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("gemini file API: %w", err) } @@ -84,7 +96,7 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, } slog.Info("read_video: using chat API fallback", "provider", providerName, "model", model, "size", len(data)) - resp, err := p.Chat(ctx, providers.ChatRequest{ + chatReq := providers.ChatRequest{ Messages: []providers.Message{ { Role: "user", @@ -97,7 +109,15 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, "max_tokens": 16384, "temperature": 0.2, }, - }) + } + reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + if reserveErr != nil { + return nil, nil, reserveErr + } + resp, err := p.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } if err != nil { return nil, nil, fmt.Errorf("chat API: %w", err) } diff --git a/internal/tools/shell.go b/internal/tools/shell.go index 6be167cd..4e07182f 100644 --- a/internal/tools/shell.go +++ b/internal/tools/shell.go @@ -13,8 +13,10 @@ import ( "regexp" "runtime" "strings" + "sync" "time" + "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/sandbox" "github.com/nextlevelbuilder/goclaw/internal/store" "golang.org/x/text/unicode/norm" @@ -35,27 +37,56 @@ func DefaultDenyPatterns() []*regexp.Regexp { // ExecTool executes shell commands, optionally inside a sandbox container. type ExecTool struct { - workspace string - timeout time.Duration - pathDenyPatterns []*regexp.Regexp // always-on path-based denials (DenyPaths) - pathDenyRoots []string // raw deny roots for nested workspace exemptions - denyExemptions []string // substrings that exempt a command from deny - restrict bool - sandboxMgr sandbox.Manager // nil = no sandbox, execute on host - approvalMgr *ExecApprovalManager // nil = no approval needed - agentID string // for approval request context - secureCLIStore store.SecureCLIStore // nil = no credentialed exec + workspace string + timeout time.Duration + pathDenyPatterns []*regexp.Regexp // always-on path-based denials (DenyPaths) + pathDenyRoots []string // raw deny roots for nested workspace exemptions + denyExemptions []string // substrings that exempt a command from deny + restrict bool + sandboxMgr sandbox.Manager // nil = no sandbox, execute on host + approvalMgr *ExecApprovalManager // nil = no approval needed + agentID string // for approval request context + secureCLIStore store.SecureCLIStore // nil = no credentialed exec + policyMu sync.RWMutex + commandKeywordAllowlist []config.CommandKeywordAllowlistRule // globalDenyGroups holds global shell deny-group toggles from config.tools. // Per-agent overrides from context (store.WithShellDenyGroups) win per-key. // Updated at startup and via TopicConfigChanged pub/sub for runtime reload. globalDenyGroups map[string]bool } +// SetCommandKeywordAllowlist replaces the scoped credentialed CLI keyword +// allowlist. The slice is defensively copied so config reload callers cannot +// mutate the running tool policy after assignment. +func (t *ExecTool) SetCommandKeywordAllowlist(rules []config.CommandKeywordAllowlistRule) { + t.policyMu.Lock() + defer t.policyMu.Unlock() + if len(rules) == 0 { + t.commandKeywordAllowlist = nil + return + } + t.commandKeywordAllowlist = slicesClone(rules) +} + +// CommandKeywordAllowlistForTest exposes the effective global allowlist for +// cross-package config reload tests. Not for production callers. +func (t *ExecTool) CommandKeywordAllowlistForTest() []config.CommandKeywordAllowlistRule { + return t.commandKeywordAllowlistSnapshot() +} + +func (t *ExecTool) commandKeywordAllowlistSnapshot() []config.CommandKeywordAllowlistRule { + t.policyMu.RLock() + defer t.policyMu.RUnlock() + return slicesClone(t.commandKeywordAllowlist) +} + // SetGlobalShellDenyGroups replaces the global shell deny-group toggles. The // caller's map is defensively copied so later mutations cannot leak into the // tool's internal state. Passing nil or an empty map clears the global config // (per-agent context overrides, if any, still apply on their own). func (t *ExecTool) SetGlobalShellDenyGroups(groups map[string]bool) { + t.policyMu.Lock() + defer t.policyMu.Unlock() if len(groups) == 0 { t.globalDenyGroups = nil return @@ -70,14 +101,20 @@ func (t *ExecTool) SetGlobalShellDenyGroups(groups map[string]bool) { // empty, the other is returned directly (no allocation). func (t *ExecTool) effectiveDenyGroups(ctx context.Context) map[string]bool { agent := store.ShellDenyGroupsFromContext(ctx) - if len(t.globalDenyGroups) == 0 { + t.policyMu.RLock() + global := t.globalDenyGroups + if len(global) > 0 { + global = maps.Clone(global) + } + t.policyMu.RUnlock() + if len(global) == 0 { return agent } if len(agent) == 0 { - return t.globalDenyGroups + return global } - merged := make(map[string]bool, len(t.globalDenyGroups)+len(agent)) - maps.Copy(merged, t.globalDenyGroups) + merged := make(map[string]bool, len(global)+len(agent)) + maps.Copy(merged, global) // agent wins per-key maps.Copy(merged, agent) return merged diff --git a/internal/tools/subagent.go b/internal/tools/subagent.go index 865ee178..b642e1d9 100644 --- a/internal/tools/subagent.go +++ b/internal/tools/subagent.go @@ -18,6 +18,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // SubagentConfig configures the subagent system. @@ -40,33 +41,34 @@ const ( // SubagentTask tracks a running or completed subagent. type SubagentTask struct { - ID string `json:"id"` - ParentID string `json:"parentId"` - Task string `json:"task"` - Label string `json:"label"` - Status string `json:"status"` // "running", "completed", "failed", "cancelled" - Result string `json:"result,omitempty"` - Depth int `json:"depth"` - Model string `json:"model,omitempty"` // model override for this subagent - TotalInputTokens int64 `json:"totalInputTokens,omitempty"` - TotalOutputTokens int64 `json:"totalOutputTokens,omitempty"` - OriginChannel string `json:"originChannel,omitempty"` - OriginChatID string `json:"originChatId,omitempty"` - OriginPeerKind string `json:"originPeerKind,omitempty"` // "direct" or "group" (for session key building) - OriginLocalKey string `json:"originLocalKey,omitempty"` // composite key with topic/thread suffix for routing - OriginUserID string `json:"originUserId,omitempty"` // parent's userID for per-user scoping propagation - OriginSenderID string `json:"originSenderId,omitempty"` // real acting sender; preserves permission attribution in announce re-ingress (#915) - OriginRole string `json:"originRole,omitempty"` // parent's RBAC role; bypasses per-user grants for admin/operator/owner in re-ingress (#915) - OriginSessionKey string `json:"originSessionKey,omitempty"` // exact parent session key for announce routing (WS uses non-standard format) - CreatedAt int64 `json:"createdAt"` - CompletedAt int64 `json:"completedAt,omitempty"` - Media []bus.MediaFile `json:"-"` // media files from tool results - OriginTenantID uuid.UUID `json:"-"` // parent's tenant for announce routing - OriginTraceID uuid.UUID `json:"-"` // parent trace for announce linking - OriginRootSpanID uuid.UUID `json:"-"` // parent agent's root span ID - cancelFunc context.CancelFunc `json:"-"` // per-task context cancel - spawnConfig SubagentConfig `json:"-"` // resolved config at spawn time (per-agent override merged) - dbID uuid.UUID `json:"-"` // persistent DB UUID (zero if not persisted) + ID string `json:"id"` + ParentID string `json:"parentId"` + Task string `json:"task"` + Label string `json:"label"` + Status string `json:"status"` // "running", "completed", "failed", "cancelled" + Result string `json:"result,omitempty"` + Depth int `json:"depth"` + Model string `json:"model,omitempty"` // model override for this subagent + TotalInputTokens int64 `json:"totalInputTokens,omitempty"` + TotalOutputTokens int64 `json:"totalOutputTokens,omitempty"` + OriginChannel string `json:"originChannel,omitempty"` + OriginChatID string `json:"originChatId,omitempty"` + OriginPeerKind string `json:"originPeerKind,omitempty"` // "direct" or "group" (for session key building) + OriginLocalKey string `json:"originLocalKey,omitempty"` // composite key with topic/thread suffix for routing + OriginUserID string `json:"originUserId,omitempty"` // parent's userID for per-user scoping propagation + OriginSenderID string `json:"originSenderId,omitempty"` // real acting sender; preserves permission attribution in announce re-ingress (#915) + OriginRole string `json:"originRole,omitempty"` // parent's RBAC role; bypasses per-user grants for admin/operator/owner in re-ingress (#915) + OriginSessionKey string `json:"originSessionKey,omitempty"` // exact parent session key for announce routing (WS uses non-standard format) + CreatedAt int64 `json:"createdAt"` + CompletedAt int64 `json:"completedAt,omitempty"` + Media []bus.MediaFile `json:"-"` // media files from tool results + OriginAgentID uuid.UUID `json:"-"` // parent agent UUID for usage caps and scoped tools + OriginTenantID uuid.UUID `json:"-"` // parent's tenant for announce routing + OriginTraceID uuid.UUID `json:"-"` // parent trace for announce linking + OriginRootSpanID uuid.UUID `json:"-"` // parent agent's root span ID + cancelFunc context.CancelFunc `json:"-"` // per-task context cancel + spawnConfig SubagentConfig `json:"-"` // resolved config at spawn time (per-agent override merged) + dbID uuid.UUID `json:"-"` // persistent DB UUID (zero if not persisted) } // SubagentManager manages the lifecycle of spawned subagents. @@ -74,8 +76,8 @@ type SubagentManager struct { mu sync.RWMutex tasks map[string]*SubagentTask config SubagentConfig - provider providers.Provider // default provider (fallback) - providerReg *providers.Registry // registry for resolving parent's provider + provider providers.Provider // default provider (fallback) + providerReg *providers.Registry // registry for resolving parent's provider model string msgBus *bus.MessageBus @@ -83,6 +85,7 @@ type SubagentManager struct { createTools func() *Registry announceQueue *AnnounceQueue // optional: batches announces with debounce taskStore store.SubagentTaskStore // optional: persists tasks to DB (fire-and-forget) + usageCaps *usagecaps.Service } // NewSubagentManager creates a new subagent manager. @@ -116,6 +119,10 @@ func (sm *SubagentManager) SetTaskStore(s store.SubagentTaskStore) { sm.taskStore = s } +func (sm *SubagentManager) SetUsageCapService(s *usagecaps.Service) { + sm.usageCaps = s +} + // effectiveConfig returns the per-agent context override merged with defaults, // or falls back to sm.config when no override is present. func (sm *SubagentManager) effectiveConfig(ctx context.Context) SubagentConfig { diff --git a/internal/tools/subagent_exec.go b/internal/tools/subagent_exec.go index cfd38de1..15cc7c49 100644 --- a/internal/tools/subagent_exec.go +++ b/internal/tools/subagent_exec.go @@ -13,6 +13,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tracing" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) // runTask executes the subagent in a goroutine. @@ -253,10 +254,7 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) } slog.Info("subagent LLM retry", "id", task.ID, "iteration", iteration, "attempt", attempt+1) } - // ctx is the parent agent's run context — cancelling the parent (e.g. agent abort) - // cascades here and to all subsequent tool calls in this iteration. - // Do NOT replace ctx with context.Background() here; that would detach abort propagation. - resp, err = activeProvider.Chat(ctx, chatReq) + resp, err = sm.chatSubagentWithUsageCap(ctx, task, activeProvider, model, chatReq, iteration, attempt+1) if err == nil { break } @@ -344,3 +342,44 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask) return iteration } + +func (sm *SubagentManager) chatSubagentWithUsageCap(ctx context.Context, task *SubagentTask, activeProvider providers.Provider, model string, chatReq providers.ChatRequest, iteration, attempt int) (*providers.ChatResponse, error) { + if fallbackProvider, ok := activeProvider.(*providers.ModelFallbackProvider); ok { + before := func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + reservation, err := sm.reserveSubagentUsage(callCtx, task, entry.ProviderName, actualReq.Model, actualReq, fmt.Sprintf("%d:%d:%s:%s", iteration, attempt, entry.ProviderName, actualReq.Model)) + if err != nil { + return nil, err + } + return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) { + if reservation != nil { + reservation.Reconcile(callCtx, resp, callErr) + } + }, nil + } + return fallbackProvider.ChatWithHook(ctx, chatReq, before) + } + reservation, err := sm.reserveSubagentUsage(ctx, task, activeProvider.Name(), model, chatReq, fmt.Sprintf("%d:%d", iteration, attempt)) + if err != nil { + return nil, err + } + resp, err := activeProvider.Chat(ctx, chatReq) + if reservation != nil { + reservation.Reconcile(ctx, resp, err) + } + return resp, err +} + +func (sm *SubagentManager) reserveSubagentUsage(ctx context.Context, task *SubagentTask, providerName, model string, chatReq providers.ChatRequest, suffix string) (*usagecaps.Reservation, error) { + if sm.usageCaps == nil { + return nil, nil + } + return sm.usageCaps.Preflight(ctx, usagecaps.Request{ + TenantID: task.OriginTenantID, + AgentID: task.OriginAgentID, + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("%s:%s", task.ID, suffix), + Messages: chatReq.Messages, + MaxOutputTokens: 4096, + }) +} diff --git a/internal/tools/subagent_spawn.go b/internal/tools/subagent_spawn.go index c9e43d05..0f44430b 100644 --- a/internal/tools/subagent_spawn.go +++ b/internal/tools/subagent_spawn.go @@ -82,16 +82,17 @@ func (sm *SubagentManager) Spawn( OriginChannel: channel, OriginChatID: chatID, OriginPeerKind: peerKind, - OriginLocalKey: ToolLocalKeyFromCtx(ctx), - OriginUserID: store.UserIDFromContext(ctx), - OriginSenderID: store.SenderIDFromContext(ctx), - OriginRole: store.RoleFromContext(ctx), - OriginSessionKey: ToolSessionKeyFromCtx(ctx), - OriginTenantID: store.TenantIDFromContext(ctx), - OriginTraceID: tracing.TraceIDFromContext(ctx), - OriginRootSpanID: tracing.ParentSpanIDFromContext(ctx), - CreatedAt: time.Now().UnixMilli(), - spawnConfig: cfg, + OriginLocalKey: ToolLocalKeyFromCtx(ctx), + OriginUserID: store.UserIDFromContext(ctx), + OriginSenderID: store.SenderIDFromContext(ctx), + OriginRole: store.RoleFromContext(ctx), + OriginSessionKey: ToolSessionKeyFromCtx(ctx), + OriginAgentID: store.AgentIDFromContext(ctx), + OriginTenantID: store.TenantIDFromContext(ctx), + OriginTraceID: tracing.TraceIDFromContext(ctx), + OriginRootSpanID: tracing.ParentSpanIDFromContext(ctx), + CreatedAt: time.Now().UnixMilli(), + spawnConfig: cfg, } // Detach from parent's cancellation chain so subagent survives after parent run completes. // WithoutCancel preserves all context values (agent ID, workspace, trace info, etc.) @@ -166,6 +167,7 @@ func (sm *SubagentManager) RunSync( OriginSenderID: store.SenderIDFromContext(ctx), OriginRole: store.RoleFromContext(ctx), OriginSessionKey: ToolSessionKeyFromCtx(ctx), + OriginAgentID: store.AgentIDFromContext(ctx), OriginTenantID: store.TenantIDFromContext(ctx), OriginTraceID: tracing.TraceIDFromContext(ctx), OriginRootSpanID: tracing.ParentSpanIDFromContext(ctx), diff --git a/internal/tools/usage_caps.go b/internal/tools/usage_caps.go new file mode 100644 index 00000000..0dc59e94 --- /dev/null +++ b/internal/tools/usage_caps.go @@ -0,0 +1,44 @@ +package tools + +import ( + "context" + "fmt" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" +) + +func reserveToolLLMUsage(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest) (*usagecaps.Reservation, error) { + if svc == nil { + return nil, nil + } + return svc.Preflight(ctx, usagecaps.Request{ + TenantID: store.TenantIDFromContext(ctx), + AgentID: store.AgentIDFromContext(ctx), + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("tool:%s:%s", toolName, uuid.NewString()), + Messages: req.Messages, + MaxOutputTokens: maxOutputTokensFromOptions(req.Options), + }) +} + +func maxOutputTokensFromOptions(options map[string]any) int { + maxTokens := 1024 + if options == nil { + return maxTokens + } + if v, ok := options[providers.OptMaxTokens]; ok { + switch n := v.(type) { + case int: + maxTokens = n + case int64: + maxTokens = int(n) + case float64: + maxTokens = int(n) + } + } + return maxTokens +} diff --git a/internal/upgrade/version.go b/internal/upgrade/version.go index 7ba6bf92..10ad72c6 100644 --- a/internal/upgrade/version.go +++ b/internal/upgrade/version.go @@ -2,4 +2,4 @@ package upgrade // RequiredSchemaVersion is the schema migration version this binary requires. // Bump this whenever adding a new SQL migration file. -const RequiredSchemaVersion uint = 69 +const RequiredSchemaVersion uint = 72 diff --git a/internal/usage/caps/chat_call.go b/internal/usage/caps/chat_call.go new file mode 100644 index 00000000..5883d263 --- /dev/null +++ b/internal/usage/caps/chat_call.go @@ -0,0 +1,150 @@ +package caps + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +// ChatOptions identifies a billable non-agent LLM call for usage-cap enforcement. +type ChatOptions struct { + TenantID uuid.UUID + AgentID uuid.UUID + ProviderName string + ModelID string + ReservationKey string + Purpose string + MaxOutputTokens int +} + +// Chat wraps Provider.Chat with the same usage-cap preflight and reconciliation +// used by agent loops. A nil service intentionally falls back to direct calls +// for Lite/subscription-only runtimes. +func (s *Service) Chat(ctx context.Context, provider providers.Provider, req providers.ChatRequest, opts ChatOptions) (*providers.ChatResponse, error) { + if provider == nil { + return nil, errors.New("usage cap chat: provider is nil") + } + if s == nil || s.store == nil { + return provider.Chat(ctx, req) + } + if fallback, ok := provider.(*providers.ModelFallbackProvider); ok { + return fallback.ChatWithHook(ctx, req, func(callCtx context.Context, entry providers.FallbackCandidate, actualReq providers.ChatRequest) (providers.FallbackAfterCall, error) { + callOpts := opts + callOpts.ProviderName = entry.ProviderName + if callOpts.ProviderName == "" && entry.Provider != nil { + callOpts.ProviderName = entry.Provider.Name() + } + callOpts.ModelID = actualReq.Model + callOpts.ReservationKey = "" + usageReq := s.chatRequest(callCtx, entry.Provider, actualReq, callOpts) + scopedCtx := scopedRequestContext(callCtx, usageReq) + reservation, err := s.Preflight(scopedCtx, usageReq) + if err != nil { + return nil, err + } + return func(resp *providers.ChatResponse, callErr error, _ providers.FallbackCallInfo) { + if reservation != nil { + reservation.Reconcile(scopedCtx, resp, callErr) + } + }, nil + }) + } + + usageReq := s.chatRequest(ctx, provider, req, opts) + scopedCtx := scopedRequestContext(ctx, usageReq) + reservation, err := s.Preflight(scopedCtx, usageReq) + if err != nil { + return nil, err + } + resp, err := provider.Chat(scopedCtx, req) + if reservation != nil { + reservation.Reconcile(scopedCtx, resp, err) + } + return resp, err +} + +func (s *Service) chatRequest(ctx context.Context, provider providers.Provider, req providers.ChatRequest, opts ChatOptions) Request { + tenantID := opts.TenantID + if tenantID == uuid.Nil { + tenantID = store.TenantIDFromContext(ctx) + } + if tenantID == uuid.Nil { + tenantID = store.MasterTenantID + } + agentID := opts.AgentID + if agentID == uuid.Nil { + agentID = store.AgentIDFromContext(ctx) + } + providerName := opts.ProviderName + if providerName == "" && provider != nil { + providerName = provider.Name() + } + modelID := opts.ModelID + if modelID == "" { + modelID = req.Model + } + if modelID == "" && provider != nil { + modelID = provider.DefaultModel() + } + return Request{ + TenantID: tenantID, + AgentID: agentID, + ProviderName: providerName, + ModelID: modelID, + ReservationKey: reservationKey(opts), + Messages: req.Messages, + MaxOutputTokens: maxOutputTokens(req, opts.MaxOutputTokens), + } +} + +func scopedRequestContext(ctx context.Context, req Request) context.Context { + if req.TenantID != uuid.Nil && store.TenantIDFromContext(ctx) != req.TenantID { + ctx = store.WithTenantID(ctx, req.TenantID) + } + if req.AgentID != uuid.Nil && store.AgentIDFromContext(ctx) != req.AgentID { + ctx = store.WithAgentID(ctx, req.AgentID) + } + return ctx +} + +func reservationKey(opts ChatOptions) string { + if opts.ReservationKey != "" { + return opts.ReservationKey + } + purpose := opts.Purpose + if purpose == "" { + purpose = "llm" + } + return fmt.Sprintf("%s:%s", purpose, uuid.NewString()) +} + +func maxOutputTokens(req providers.ChatRequest, fallback int) int { + if fallback <= 0 { + fallback = 1024 + } + if req.Options == nil { + return fallback + } + v, ok := req.Options[providers.OptMaxTokens] + if !ok { + return fallback + } + switch n := v.(type) { + case int: + return n + case int64: + return int(n) + case int32: + return int(n) + case float64: + return int(n) + case float32: + return int(n) + default: + return fallback + } +} diff --git a/internal/usage/caps/reservation_metadata.go b/internal/usage/caps/reservation_metadata.go new file mode 100644 index 00000000..ad60a88d --- /dev/null +++ b/internal/usage/caps/reservation_metadata.go @@ -0,0 +1,125 @@ +package caps + +import ( + "encoding/json" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/usage/pricing" +) + +const TraceMetadataKey = "usage_caps" + +type TraceMetadata struct { + Decision string `json:"decision"` + Reason string `json:"reason,omitempty"` + ReservationKey string `json:"reservation_key,omitempty"` + ProviderName string `json:"provider_name,omitempty"` + ProviderType string `json:"provider_type,omitempty"` + ModelID string `json:"model_id,omitempty"` + EstimatedTokens int64 `json:"estimated_tokens"` + EstimatedCostMicros int64 `json:"estimated_cost_micros"` + ActualTokens int64 `json:"actual_tokens"` + ActualCostMicros int64 `json:"actual_cost_micros"` + ReconcileStatus string `json:"reconcile_status,omitempty"` + PolicyCount int `json:"policy_count"` + PolicyIDs []string `json:"policy_ids,omitempty"` +} + +func (r *Reservation) TraceMetadata() TraceMetadata { + if r == nil { + return TraceMetadata{} + } + decision := r.decision + if decision == "" { + if r.skipped { + decision = store.UsageCapEventSkip + } else { + decision = store.UsageCapEventAllow + } + } + policyCount := 0 + policyIDs := make([]string, 0) + if r.result != nil { + policyCount = len(r.result.Policies) + for _, policy := range r.result.Policies { + if policy.ID != uuid.Nil { + policyIDs = append(policyIDs, policy.ID.String()) + } + } + } + if r.blockedPolicyID != uuid.Nil { + policyIDs = append(policyIDs, r.blockedPolicyID.String()) + if policyCount == 0 { + policyCount = 1 + } + } + return TraceMetadata{ + Decision: decision, + Reason: r.reason, + ReservationKey: r.key, + ProviderName: r.providerName, + ProviderType: r.providerType, + ModelID: r.modelID, + EstimatedTokens: r.usage.TotalTokens(), + EstimatedCostMicros: r.estimatedCostMicros, + ActualTokens: r.actualTokens, + ActualCostMicros: r.actualCostMicros, + ReconcileStatus: r.reconcileStatus, + PolicyCount: policyCount, + PolicyIDs: policyIDs, + } +} + +func (m TraceMetadata) Empty() bool { + return m.Decision == "" && m.Reason == "" && m.ReservationKey == "" && + m.ProviderName == "" && m.ProviderType == "" && m.ModelID == "" && + m.EstimatedTokens == 0 && m.EstimatedCostMicros == 0 && + m.ActualTokens == 0 && m.ActualCostMicros == 0 && + m.ReconcileStatus == "" && m.PolicyCount == 0 && len(m.PolicyIDs) == 0 +} + +func MergeTraceMetadata(existing json.RawMessage, entries []TraceMetadata) json.RawMessage { + clean := make([]TraceMetadata, 0, len(entries)) + for _, entry := range entries { + if !entry.Empty() { + clean = append(clean, entry) + } + } + if len(clean) == 0 { + return existing + } + payload := map[string]any{} + if len(existing) > 0 { + _ = json.Unmarshal(existing, &payload) + } + payload[TraceMetadataKey] = map[string]any{"attempts": clean} + data, err := json.Marshal(payload) + if err != nil { + return existing + } + return json.RawMessage(data) +} + +func skippedReservation(req Request, reason string) *Reservation { + return &Reservation{ + skipped: true, decision: store.UsageCapEventSkip, reason: reason, + providerName: req.ProviderName, modelID: req.ModelID, + } +} + +func skippedScopedReservation(req Request, scope store.UsageCapScope, reason string) *Reservation { + r := skippedReservation(req, reason) + r.providerType = scope.ProviderType + r.modelID = scope.ModelID + return r +} + +func blockedReservation(req Request, scope store.UsageCapScope, key string, usage pricing.BillableUsage, costMicros int64, policyID uuid.UUID, reason string) *Reservation { + return &Reservation{ + key: key, usage: usage, estimatedCostMicros: costMicros, + decision: store.UsageCapEventBlock, reason: reason, + blockedPolicyID: policyID, + providerName: req.ProviderName, providerType: scope.ProviderType, modelID: scope.ModelID, + } +} diff --git a/internal/usage/caps/service.go b/internal/usage/caps/service.go new file mode 100644 index 00000000..30061bec --- /dev/null +++ b/internal/usage/caps/service.go @@ -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 +} diff --git a/internal/usage/caps/service_test.go b/internal/usage/caps/service_test.go new file mode 100644 index 00000000..6a597404 --- /dev/null +++ b/internal/usage/caps/service_test.go @@ -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 } diff --git a/internal/usage/pricing/decimal.go b/internal/usage/pricing/decimal.go new file mode 100644 index 00000000..21976542 --- /dev/null +++ b/internal/usage/pricing/decimal.go @@ -0,0 +1,119 @@ +package pricing + +import ( + "errors" + "fmt" + "math/big" + "strings" + + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +var ErrUnknownPricing = errors.New("usage pricing unknown") + +type BillableUsage struct { + InputTokens int64 + OutputTokens int64 + CacheReadTokens int64 + CacheWriteTokens int64 + ReasoningTokens int64 + RequestCount int64 + ImageCount int64 + WebSearchCount int64 +} + +func FromProviderUsage(u *providers.Usage) BillableUsage { + if u == nil { + return BillableUsage{} + } + inputTokens := int64(u.PromptTokens) + cacheReadTokens := int64(u.CacheReadTokens) + cacheWriteTokens := int64(u.CacheCreationTokens) + if u.PromptTokensIncludeCachedSegments { + inputTokens -= cacheReadTokens + cacheWriteTokens + if inputTokens < 0 { + inputTokens = 0 + } + } + return BillableUsage{ + InputTokens: inputTokens, + OutputTokens: int64(u.CompletionTokens), + CacheReadTokens: cacheReadTokens, + CacheWriteTokens: cacheWriteTokens, + ReasoningTokens: int64(u.ThinkingTokens), + RequestCount: int64(u.RequestCount), + ImageCount: int64(u.ImageCount), + WebSearchCount: int64(u.WebSearchCount), + } +} + +func (u BillableUsage) TotalTokens() int64 { + return u.InputTokens + u.OutputTokens + u.CacheReadTokens + u.CacheWriteTokens +} + +func CostMicros(fields store.UsagePricingFields, usage BillableUsage) (int64, error) { + var total int64 + add := func(name string, price *string, units int64) error { + if units <= 0 { + return nil + } + if price == nil || strings.TrimSpace(*price) == "" { + return fmt.Errorf("%w: %s", ErrUnknownPricing, name) + } + micros, err := microsForUnits(*price, units) + if err != nil { + return fmt.Errorf("%s pricing: %w", name, err) + } + total += micros + return nil + } + if err := add("input", fields.Input, usage.InputTokens); err != nil { + return 0, err + } + if usage.ReasoningTokens > 0 && fields.Reasoning != nil { + visible := max(usage.OutputTokens-usage.ReasoningTokens, 0) + if err := add("output", fields.Output, visible); err != nil { + return 0, err + } + if err := add("reasoning", fields.Reasoning, usage.ReasoningTokens); err != nil { + return 0, err + } + } else if err := add("output", fields.Output, usage.OutputTokens); err != nil { + return 0, err + } + if err := add("cache_read", fields.CacheRead, usage.CacheReadTokens); err != nil { + return 0, err + } + if err := add("cache_write", fields.CacheWrite, usage.CacheWriteTokens); err != nil { + return 0, err + } + if err := add("request", fields.Request, usage.RequestCount); err != nil { + return 0, err + } + if err := add("image", fields.Image, usage.ImageCount); err != nil { + return 0, err + } + if err := add("web_search", fields.WebSearch, usage.WebSearchCount); err != nil { + return 0, err + } + return total, nil +} + +func microsForUnits(price string, units int64) (int64, error) { + r, ok := new(big.Rat).SetString(strings.TrimSpace(price)) + if !ok { + return 0, fmt.Errorf("invalid decimal %q", price) + } + r.Mul(r, big.NewRat(units, 1)) + r.Mul(r, big.NewRat(1_000_000, 1)) + q := new(big.Int).Quo(r.Num(), r.Denom()) + rem := new(big.Int).Rem(r.Num(), r.Denom()) + if rem.Sign() > 0 { + q.Add(q, big.NewInt(1)) + } + if !q.IsInt64() { + return 0, errors.New("cost overflows int64 micros") + } + return q.Int64(), nil +} diff --git a/internal/usage/pricing/decimal_test.go b/internal/usage/pricing/decimal_test.go new file mode 100644 index 00000000..515330fb --- /dev/null +++ b/internal/usage/pricing/decimal_test.go @@ -0,0 +1,66 @@ +package pricing + +import ( + "errors" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func strp(s string) *string { return &s } + +func TestCostMicrosTokenAndFlatDimensions(t *testing.T) { + got, err := CostMicros(store.UsagePricingFields{ + Input: strp("0.000001"), + Output: strp("0.000002"), + CacheRead: strp("0.0000001"), + CacheWrite: strp("0.0000002"), + Request: strp("0.01"), + Image: strp("0.02"), + WebSearch: strp("0.03"), + }, BillableUsage{ + InputTokens: 1000, + OutputTokens: 500, + CacheReadTokens: 100, + CacheWriteTokens: 50, + RequestCount: 1, + ImageCount: 2, + WebSearchCount: 1, + }) + if err != nil { + t.Fatal(err) + } + want := int64(1000 + 1000 + 10 + 10 + 10_000 + 40_000 + 30_000) + if got != want { + t.Fatalf("cost micros = %d, want %d", got, want) + } +} + +func TestCostMicrosMissingRequiredPrice(t *testing.T) { + _, err := CostMicros(store.UsagePricingFields{Input: strp("0.000001")}, BillableUsage{ + InputTokens: 1, + OutputTokens: 1, + }) + if !errors.Is(err, ErrUnknownPricing) { + t.Fatalf("err = %v, want ErrUnknownPricing", err) + } +} + +func TestFromProviderUsageSeparatesIncludedCacheTokens(t *testing.T) { + got := FromProviderUsage(&providers.Usage{ + PromptTokens: 2006, + CompletionTokens: 100, + CacheReadTokens: 1920, + PromptTokensIncludeCachedSegments: true, + }) + if got.InputTokens != 86 { + t.Fatalf("InputTokens = %d, want 86", got.InputTokens) + } + if got.CacheReadTokens != 1920 { + t.Fatalf("CacheReadTokens = %d, want 1920", got.CacheReadTokens) + } + if got.TotalTokens() != 2106 { + t.Fatalf("TotalTokens = %d, want 2106", got.TotalTokens()) + } +} diff --git a/internal/usage/pricing/openrouter.go b/internal/usage/pricing/openrouter.go new file mode 100644 index 00000000..98cb84e9 --- /dev/null +++ b/internal/usage/pricing/openrouter.go @@ -0,0 +1,91 @@ +package pricing + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +const OpenRouterModelsURL = "https://openrouter.ai/api/v1/models" + +type openRouterModelsResponse struct { + Data []json.RawMessage `json:"data"` +} + +type openRouterModel struct { + ID string `json:"id"` + Name string `json:"name"` + Pricing map[string]any `json:"pricing"` +} + +func FetchOpenRouterCatalog(ctx context.Context, client *http.Client) ([]store.UsagePricingCatalogEntry, error) { + if client == nil { + client = http.DefaultClient + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, OpenRouterModelsURL, nil) + if err != nil { + return nil, err + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("openrouter models status %d", resp.StatusCode) + } + var payload openRouterModelsResponse + if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil { + return nil, err + } + now := time.Now().UTC() + entries := make([]store.UsagePricingCatalogEntry, 0, len(payload.Data)) + for _, raw := range payload.Data { + var model openRouterModel + if err := json.Unmarshal(raw, &model); err != nil || model.ID == "" { + continue + } + rawPricing, _ := json.Marshal(model.Pricing) + entries = append(entries, store.UsagePricingCatalogEntry{ + ModelID: model.ID, + CanonicalModelID: model.ID, + Pricing: mapOpenRouterPricing(model.Pricing), + RawPricing: rawPricing, + RawModel: raw, + SyncedAt: now, + }) + } + return entries, nil +} + +func mapOpenRouterPricing(raw map[string]any) store.UsagePricingFields { + return store.UsagePricingFields{ + Input: decimalStringPtr(raw["prompt"]), + Output: decimalStringPtr(raw["completion"]), + CacheRead: decimalStringPtr(raw["input_cache_read"]), + CacheWrite: decimalStringPtr(raw["input_cache_write"]), + Reasoning: decimalStringPtr(raw["internal_reasoning"]), + Request: decimalStringPtr(raw["request"]), + Image: decimalStringPtr(raw["image"]), + WebSearch: decimalStringPtr(raw["web_search"]), + } +} + +func decimalStringPtr(v any) *string { + switch x := v.(type) { + case string: + if x == "" { + return nil + } + return &x + case float64: + s := fmt.Sprintf("%.18g", x) + return &s + default: + return nil + } +} diff --git a/internal/vault/enrich_classify.go b/internal/vault/enrich_classify.go index b8aaf857..227a27cc 100644 --- a/internal/vault/enrich_classify.go +++ b/internal/vault/enrich_classify.go @@ -7,6 +7,7 @@ import ( "maps" "slices" + "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -40,6 +41,14 @@ func (w *EnrichWorker) classifyLinks(ctx context.Context, provider providers.Pro if provider == nil { return } + if tid, err := uuid.Parse(tenantID); err == nil { + ctx = store.WithTenantID(ctx, tid) + } + if agentID != "" { + if aid, err := uuid.Parse(agentID); err == nil { + ctx = store.WithAgentID(ctx, aid) + } + } capped := results if len(capped) > classifyMaxSourceDocs { diff --git a/internal/vault/enrich_worker.go b/internal/vault/enrich_worker.go index 18d9f49e..2fd4cb1c 100644 --- a/internal/vault/enrich_worker.go +++ b/internal/vault/enrich_worker.go @@ -16,20 +16,21 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bgalert" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/eventbus" - "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/providerresolve" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" "golang.org/x/sync/semaphore" ) const ( - enrichMaxDedupEntries = 10000 - enrichSimilarityLimit = 10 - enrichSimilarityMin = 0.7 - enrichMaxConcurrent = 3 // max concurrent batch summarize calls across chunks - enrichBatchSize = 5 // docs per enrichment chunk (1 LLM call per chunk) - enrichBatchItemMaxRunes = 3000 // per-file content limit in batch summarize - enrichMaxRetries = 3 // shared retry count for LLM calls (summarize + classify) + enrichMaxDedupEntries = 10000 + enrichSimilarityLimit = 10 + enrichSimilarityMin = 0.7 + enrichMaxConcurrent = 3 // max concurrent batch summarize calls across chunks + enrichBatchSize = 5 // docs per enrichment chunk (1 LLM call per chunk) + enrichBatchItemMaxRunes = 3000 // per-file content limit in batch summarize + enrichMaxRetries = 3 // shared retry count for LLM calls (summarize + classify) ) // Shared retry config for all enrichment LLM calls. @@ -41,12 +42,13 @@ var ( // EnrichWorkerDeps bundles dependencies for the vault enrichment worker. type EnrichWorkerDeps struct { VaultStore store.VaultStore - SystemConfigs store.SystemConfigStore // per-tenant provider config - Registry *providers.Registry // provider resolution + SystemConfigs store.SystemConfigStore // per-tenant provider config + Registry *providers.Registry // provider resolution EventBus eventbus.DomainEventBus - MsgBus bus.EventPublisher // for WS event broadcast - TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe) - AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + MsgBus bus.EventPublisher // for WS event broadcast + TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe) + AlertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + UsageCaps *usagecaps.Service } // RegisterEnrichWorker subscribes the enrichment worker to vault doc events. @@ -60,6 +62,7 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *Enri registry: deps.Registry, msgBus: deps.MsgBus, alertDeps: deps.AlertDeps, + usageCaps: deps.UsageCaps, dedup: make(map[string]string), sem: semaphore.NewWeighted(enrichMaxConcurrent), progress: progress, @@ -74,11 +77,12 @@ func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *Enri // Exported so HTTP handlers can call Stop/EnqueueUnenriched. type EnrichWorker struct { vault store.VaultStore - teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil - systemConfigs store.SystemConfigStore // per-tenant provider config - registry *providers.Registry // provider resolution - msgBus bus.EventPublisher // for error event broadcast - alertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil + systemConfigs store.SystemConfigStore // per-tenant provider config + registry *providers.Registry // provider resolution + msgBus bus.EventPublisher // for error event broadcast + alertDeps bgalert.AlertDeps // for reporting non-retryable LLM errors + usageCaps *usagecaps.Service queue enrichBatchQueue progress *EnrichProgress @@ -296,6 +300,14 @@ func (w *EnrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD // Batch-fetch all existing docs in a single query. tenantID := pending[0].TenantID + if tid, err := uuid.Parse(tenantID); err == nil { + ctx = store.WithTenantID(ctx, tid) + } + if pending[0].AgentID != "" { + if aid, err := uuid.Parse(pending[0].AgentID); err == nil { + ctx = store.WithAgentID(ctx, aid) + } + } // Resolve provider once per chunk (all items share tenantID) provider, model := w.resolveProviderForTenant(ctx, tenantID) @@ -515,7 +527,11 @@ func (w *EnrichWorker) chatWithRetry(ctx context.Context, provider providers.Pro } } cctx, cancel := context.WithTimeout(ctx, enrichRetryTimeouts[attempt]) - resp, err := provider.Chat(cctx, req) + resp, err := w.usageCaps.Chat(cctx, provider, req, usagecaps.ChatOptions{ + ModelID: req.Model, + Purpose: logPrefix, + MaxOutputTokens: 4096, + }) cancel() if err != nil { lastErr = err @@ -563,7 +579,6 @@ func (w *EnrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUps } } - // recordDedup stores a processed hash and evicts ~25% entries if over capacity. func (w *EnrichWorker) recordDedup(docID, hash string) { w.dedupMu.Lock() diff --git a/migrations/000070_usage_caps_pricing.down.sql b/migrations/000070_usage_caps_pricing.down.sql new file mode 100644 index 00000000..18c41ec8 --- /dev/null +++ b/migrations/000070_usage_caps_pricing.down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS usage_pricing_overrides; +DROP TABLE IF EXISTS usage_pricing_catalog; diff --git a/migrations/000070_usage_caps_pricing.up.sql b/migrations/000070_usage_caps_pricing.up.sql new file mode 100644 index 00000000..d69e35d1 --- /dev/null +++ b/migrations/000070_usage_caps_pricing.up.sql @@ -0,0 +1,45 @@ +CREATE TABLE IF NOT EXISTS usage_pricing_catalog ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + model_id TEXT NOT NULL UNIQUE, + canonical_model_id TEXT, + raw_pricing JSONB NOT NULL DEFAULT '{}'::jsonb, + raw_model JSONB NOT NULL DEFAULT '{}'::jsonb, + input_price NUMERIC(30, 18) CHECK (input_price IS NULL OR input_price >= 0), + output_price NUMERIC(30, 18) CHECK (output_price IS NULL OR output_price >= 0), + cache_read_price NUMERIC(30, 18) CHECK (cache_read_price IS NULL OR cache_read_price >= 0), + cache_write_price NUMERIC(30, 18) CHECK (cache_write_price IS NULL OR cache_write_price >= 0), + reasoning_price NUMERIC(30, 18) CHECK (reasoning_price IS NULL OR reasoning_price >= 0), + request_price NUMERIC(30, 18) CHECK (request_price IS NULL OR request_price >= 0), + image_price NUMERIC(30, 18) CHECK (image_price IS NULL OR image_price >= 0), + web_search_price NUMERIC(30, 18) CHECK (web_search_price IS NULL OR web_search_price >= 0), + synced_at TIMESTAMPTZ NOT NULL DEFAULT now(), + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX IF NOT EXISTS idx_usage_pricing_catalog_synced_at + ON usage_pricing_catalog (synced_at DESC); + +CREATE TABLE IF NOT EXISTS usage_pricing_overrides ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, + provider_id UUID NOT NULL REFERENCES llm_providers(id) ON DELETE CASCADE, + provider_type TEXT NOT NULL, + model_id TEXT NOT NULL, + input_price NUMERIC(30, 18) CHECK (input_price IS NULL OR input_price >= 0), + output_price NUMERIC(30, 18) CHECK (output_price IS NULL OR output_price >= 0), + cache_read_price NUMERIC(30, 18) CHECK (cache_read_price IS NULL OR cache_read_price >= 0), + cache_write_price NUMERIC(30, 18) CHECK (cache_write_price IS NULL OR cache_write_price >= 0), + reasoning_price NUMERIC(30, 18) CHECK (reasoning_price IS NULL OR reasoning_price >= 0), + request_price NUMERIC(30, 18) CHECK (request_price IS NULL OR request_price >= 0), + image_price NUMERIC(30, 18) CHECK (image_price IS NULL OR image_price >= 0), + web_search_price NUMERIC(30, 18) CHECK (web_search_price IS NULL OR web_search_price >= 0), + enabled BOOLEAN NOT NULL DEFAULT true, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (tenant_id, provider_id, model_id) +); + +CREATE INDEX IF NOT EXISTS idx_usage_pricing_overrides_tenant_provider + ON usage_pricing_overrides (tenant_id, provider_id, model_id) + WHERE enabled; diff --git a/migrations/000071_usage_cap_policies.down.sql b/migrations/000071_usage_cap_policies.down.sql new file mode 100644 index 00000000..49e817d5 --- /dev/null +++ b/migrations/000071_usage_cap_policies.down.sql @@ -0,0 +1,4 @@ +DROP TABLE IF EXISTS usage_cap_events; +DROP TABLE IF EXISTS usage_cap_reservations; +DROP TABLE IF EXISTS usage_cap_counters; +DROP TABLE IF EXISTS usage_cap_policies; diff --git a/migrations/000071_usage_cap_policies.up.sql b/migrations/000071_usage_cap_policies.up.sql new file mode 100644 index 00000000..09a8a578 --- /dev/null +++ b/migrations/000071_usage_cap_policies.up.sql @@ -0,0 +1,68 @@ +CREATE TABLE IF NOT EXISTS usage_cap_policies ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE, + agent_id UUID REFERENCES agents(id) ON DELETE CASCADE, + provider_id UUID REFERENCES llm_providers(id) ON DELETE CASCADE, + provider_type TEXT, + model_id TEXT, + window_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); diff --git a/migrations/000072_agent_budget_usage_cap_bridge.down.sql b/migrations/000072_agent_budget_usage_cap_bridge.down.sql new file mode 100644 index 00000000..e26c4b4a --- /dev/null +++ b/migrations/000072_agent_budget_usage_cap_bridge.down.sql @@ -0,0 +1,7 @@ +DELETE FROM usage_cap_policies +WHERE source = 'agent_budget_monthly_cents'; + +DROP INDEX IF EXISTS idx_usage_cap_policies_agent_budget_source; + +ALTER TABLE usage_cap_policies + DROP COLUMN IF EXISTS source; diff --git a/migrations/000072_agent_budget_usage_cap_bridge.up.sql b/migrations/000072_agent_budget_usage_cap_bridge.up.sql new file mode 100644 index 00000000..700d1b9d --- /dev/null +++ b/migrations/000072_agent_budget_usage_cap_bridge.up.sql @@ -0,0 +1,23 @@ +ALTER TABLE usage_cap_policies + ADD COLUMN IF NOT EXISTS source TEXT NOT NULL DEFAULT 'manual'; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_cap_policies_agent_budget_source + ON usage_cap_policies (tenant_id, agent_id) + WHERE source = 'agent_budget_monthly_cents'; + +INSERT INTO usage_cap_policies ( + tenant_id, agent_id, window_key, max_cost_micros, enabled, priority, source +) +SELECT + tenant_id, + id, + 'month', + budget_monthly_cents::BIGINT * 10000, + true, + 90, + 'agent_budget_monthly_cents' +FROM agents +WHERE deleted_at IS NULL + AND budget_monthly_cents IS NOT NULL + AND budget_monthly_cents > 0 +ON CONFLICT DO NOTHING; diff --git a/ui/web/src/components/layout/system-settings-constants.ts b/ui/web/src/components/layout/system-settings-constants.ts index 3bb1226c..9cbcfc5d 100644 --- a/ui/web/src/components/layout/system-settings-constants.ts +++ b/ui/web/src/components/layout/system-settings-constants.ts @@ -44,6 +44,11 @@ export interface InitState { kgMinConfidence: string; bgProvider: string; bgModel: string; + skillUploadMaxSize: string; + skillSlashEnabled: boolean; + skillSlashSuggest: boolean; + skillSlashPartial: boolean; + skillSlashPrefix: string; } export const DEFAULTS: InitState = { @@ -54,6 +59,11 @@ export const DEFAULTS: InitState = { compThreshold: "", compKeepRecent: "", compMaxTokens: "", kgProvider: "", kgModel: "", kgMinConfidence: "0.75", bgProvider: "", bgModel: "", + skillUploadMaxSize: "20", + skillSlashEnabled: true, + skillSlashSuggest: true, + skillSlashPartial: false, + skillSlashPrefix: "/", }; export function parseBool(v: string | undefined, fallback: boolean): boolean { diff --git a/ui/web/src/components/layout/system-settings-modal.tsx b/ui/web/src/components/layout/system-settings-modal.tsx index 01cc4815..4bc0aca9 100644 --- a/ui/web/src/components/layout/system-settings-modal.tsx +++ b/ui/web/src/components/layout/system-settings-modal.tsx @@ -17,6 +17,7 @@ import { toast } from "@/stores/use-toast-store"; import { EMBEDDING_MODELS, DEFAULT_EMBEDDING_MODELS, DEFAULTS, parseBool, type InitState } from "./system-settings-constants"; import { SystemSettingsEmbeddingCard } from "./system-settings-embedding-card"; import { SystemSettingsCompactionCard } from "./system-settings-compaction-card"; +import { SystemSettingsSkillsCard } from "./system-settings-skills-card"; import { Eye, MessageSquareText, Brain } from "lucide-react"; interface SystemSettingsModalProps { @@ -60,6 +61,11 @@ export function SystemSettingsModal({ open, onOpenChange }: SystemSettingsModalP // Background Workers const [bgProvider, setBgProvider] = useState(""); const [bgModel, setBgModel] = useState(""); + const [skillUploadMaxSize, setSkillUploadMaxSize] = useState("20"); + const [skillSlashEnabled, setSkillSlashEnabled] = useState(true); + const [skillSlashSuggest, setSkillSlashSuggest] = useState(true); + const [skillSlashPartial, setSkillSlashPartial] = useState(false); + const [skillSlashPrefix, setSkillSlashPrefix] = useState("/"); const applyConfigs = useCallback(( configs: Record, @@ -76,6 +82,11 @@ export function SystemSettingsModal({ open, onOpenChange }: SystemSettingsModalP kgProvider: kgSettings?.extraction_provider ?? "", kgModel: kgSettings?.extraction_model ?? "", kgMinConfidence: String(kgSettings?.min_confidence ?? 0.75), bgProvider: configs["background.provider"] ?? "", bgModel: configs["background.model"] ?? "", + skillUploadMaxSize: configs["skills.max_upload_size_mb"] ?? "20", + skillSlashEnabled: parseBool(configs["skills.slash_commands.enabled"], true), + skillSlashSuggest: parseBool(configs["skills.slash_commands.suggest_not_found"], true), + skillSlashPartial: parseBool(configs["skills.slash_commands.partial_matching"], false), + skillSlashPrefix: configs["skills.slash_commands.prefix"] ?? "/", }; setInit(s); setEmbProvider(s.embProvider); setEmbModel(s.embModel); setEmbMaxChunkLen(s.embMaxChunkLen); setEmbChunkOverlap(s.embChunkOverlap); @@ -83,6 +94,11 @@ export function SystemSettingsModal({ open, onOpenChange }: SystemSettingsModalP setCompProvider(s.compProvider); setCompModel(s.compModel); setCompThreshold(s.compThreshold); setCompKeepRecent(s.compKeepRecent); setCompMaxTokens(s.compMaxTokens); setKgProvider(s.kgProvider); setKgModel(s.kgModel); setKgMinConfidence(s.kgMinConfidence); setBgProvider(s.bgProvider); setBgModel(s.bgModel); + setSkillUploadMaxSize(s.skillUploadMaxSize); + setSkillSlashEnabled(s.skillSlashEnabled); + setSkillSlashSuggest(s.skillSlashSuggest); + setSkillSlashPartial(s.skillSlashPartial); + setSkillSlashPrefix(s.skillSlashPrefix); resetEmb(); }, [resetEmb]); @@ -126,6 +142,11 @@ export function SystemSettingsModal({ open, onOpenChange }: SystemSettingsModalP if (compMaxTokens !== init.compMaxTokens) updates["compaction.max_tokens"] = compMaxTokens; if (bgProvider !== init.bgProvider) updates["background.provider"] = bgProvider; if (bgModel !== init.bgModel) updates["background.model"] = bgModel; + if (skillUploadMaxSize !== init.skillUploadMaxSize) updates["skills.max_upload_size_mb"] = skillUploadMaxSize; + if (skillSlashEnabled !== init.skillSlashEnabled) updates["skills.slash_commands.enabled"] = String(skillSlashEnabled); + if (skillSlashSuggest !== init.skillSlashSuggest) updates["skills.slash_commands.suggest_not_found"] = String(skillSlashSuggest); + if (skillSlashPartial !== init.skillSlashPartial) updates["skills.slash_commands.partial_matching"] = String(skillSlashPartial); + if (skillSlashPrefix !== init.skillSlashPrefix) updates["skills.slash_commands.prefix"] = skillSlashPrefix.trim() || "/"; for (const [key, value] of Object.entries(updates)) await http.put(`/v1/system-configs/${key}`, { value }); const kgChanged = kgProvider !== init.kgProvider || kgModel !== init.kgModel || kgMinConfidence !== init.kgMinConfidence; if (kgChanged) { @@ -203,6 +224,19 @@ export function SystemSettingsModal({ open, onOpenChange }: SystemSettingsModalP + + void; + slashEnabled: boolean; + setSlashEnabled: (value: boolean) => void; + slashSuggest: boolean; + setSlashSuggest: (value: boolean) => void; + slashPartial: boolean; + setSlashPartial: (value: boolean) => void; + slashPrefix: string; + setSlashPrefix: (value: string) => void; +} + +export function SystemSettingsSkillsCard({ + uploadMaxSize, + setUploadMaxSize, + slashEnabled, + setSlashEnabled, + slashSuggest, + setSlashSuggest, + slashPartial, + setSlashPartial, + slashPrefix, + setSlashPrefix, +}: SystemSettingsSkillsCardProps) { + const { t } = useTranslation("system-settings"); + + return ( + + + + + {t("skills.title")} + + {t("skills.description")} + + +
+
+ +

{t("skills.maxUploadSizeHint")}

+
+ setUploadMaxSize(e.target.value)} + className="w-24 shrink-0 text-base md:text-sm" + /> +
+
+ + + +
+
+ +

{t("skills.slashPrefixHint")}

+
+ setSlashPrefix(e.target.value.slice(0, 1))} + className="w-16 shrink-0 text-center text-base md:text-sm" + /> +
+
+
+
+ ); +} + +function SkillSwitchRow({ + id, + label, + hint, + checked, + onCheckedChange, +}: { + id: string; + label: string; + hint: string; + checked: boolean; + onCheckedChange: (value: boolean) => void; +}) { + return ( +
+
+ +

{hint}

+
+ +
+ ); +} diff --git a/ui/web/src/i18n/locales/en/cli-credentials.json b/ui/web/src/i18n/locales/en/cli-credentials.json index a668eb86..c2a15e1e 100644 --- a/ui/web/src/i18n/locales/en/cli-credentials.json +++ b/ui/web/src/i18n/locales/en/cli-credentials.json @@ -39,6 +39,10 @@ "noEnvVarsHint": "Click \"Add Variable\" to define environment variables for this CLI tool.", "envKeyPlaceholder": "ENV_VAR_NAME", "envValuePlaceholder": "value", + "envKindSensitive": "Sensitive", + "envKindValue": "Value", + "envValueWarning": "Plaintext values are visible to authorized users.", + "envValueSuspicious": "This plaintext value looks like a token, password, or key.", "invalidEnvKey": "\"{{key}}\" is not a valid env variable name (use A-Z, 0-9, underscore).", "binaryPathHint": "Leave blank to auto-detect from PATH", "checkBinary": "Check", diff --git a/ui/web/src/i18n/locales/en/config.json b/ui/web/src/i18n/locales/en/config.json index a1059e77..f22f94bc 100644 --- a/ui/web/src/i18n/locales/en/config.json +++ b/ui/web/src/i18n/locales/en/config.json @@ -208,6 +208,10 @@ "tools.execAskMode": "Ask Mode", "tools.execAskModeTip": "When to prompt the user for approval before executing a command.", "tools.execAllowlistLabel": "Allowed Commands (one pattern per line)", + "tools.commandKeywordAllowlist": "Command Keyword Allowlist", + "tools.commandKeywordAllowlistTip": "JSON rules that allow configured keywords only inside selected credentialed CLI content arguments. Command-path deny rules still apply.", + "tools.commandKeywordAllowlistJsonError": "Invalid JSON.", + "tools.commandKeywordAllowlistArrayError": "Use a JSON array of rules.", "tools.scrubCredentials": "Scrub Credentials", "browser.title": "Browser Automation", diff --git a/ui/web/src/i18n/locales/en/providers.json b/ui/web/src/i18n/locales/en/providers.json index fb4fbbca..e9c38633 100644 --- a/ui/web/src/i18n/locales/en/providers.json +++ b/ui/web/src/i18n/locales/en/providers.json @@ -258,6 +258,38 @@ "verifyFailed": "Embedding verification failed", "dimensionsMismatch": "{{count}} dimensions — does not match required 1536" }, + "pricing": { + "title": "Model Pricing", + "description": "Sync OpenRouter prices and override this provider's model costs for usage caps.", + "skippedDescription": "Usage caps skip subscription and local providers, but prices can still be prepared for later API-based routing.", + "sync": "Sync OpenRouter", + "model": "Model", + "saveOverride": "Save Override", + "catalogMatches": "OpenRouter catalog matches", + "source": "Source", + "configuredFields": "Configured fields", + "override": "Override", + "emptyOverrides": "No pricing overrides for this provider.", + "deleteOverride": "Delete override", + "fields": { + "input": "Input", + "output": "Output", + "cache_read": "Cache read", + "cache_write": "Cache write", + "reasoning": "Reasoning", + "request": "Request", + "image": "Image", + "web_search": "Web search" + }, + "toast": { + "synced": "Synced {{count}} OpenRouter models", + "syncFailed": "Could not sync OpenRouter pricing", + "saved": "Pricing override saved", + "saveFailed": "Could not save pricing override", + "deleted": "Pricing override deleted", + "deleteFailed": "Could not delete pricing override" + } + }, "reasoning": { "off": "Off", "offDesc": "No default extended thinking", diff --git a/ui/web/src/i18n/locales/en/skills.json b/ui/web/src/i18n/locales/en/skills.json index 69801a19..e8f01366 100644 --- a/ui/web/src/i18n/locales/en/skills.json +++ b/ui/web/src/i18n/locales/en/skills.json @@ -106,7 +106,7 @@ }, "upload": { "title": "Upload Skills", - "description": "Upload ZIP files containing SKILL.md with YAML frontmatter (name, description, slug).", + "description": "Upload ZIP files containing SKILL.md with YAML frontmatter (name, description, slug). Max {{max}} MB per file.", "button": "Upload", "uploading": "Uploading...", "dropOrClick": "Click or drag & drop .zip files", @@ -114,7 +114,7 @@ "onlyZip": "Only .zip files are accepted", "cancel": "Cancel", "done": "Done", - "tooLarge": "File exceeds 20MB limit", + "tooLarge": "File exceeds the configured upload limit", "invalidZip": "Invalid ZIP file", "noSkillMd": "ZIP must contain SKILL.md at root", "emptySkillMd": "SKILL.md is empty", diff --git a/ui/web/src/i18n/locales/en/system-settings.json b/ui/web/src/i18n/locales/en/system-settings.json index fafc76c7..911e74b1 100644 --- a/ui/web/src/i18n/locales/en/system-settings.json +++ b/ui/web/src/i18n/locales/en/system-settings.json @@ -64,6 +64,20 @@ "modelPlaceholder": "(default)", "info": "Used for vault enrichment (document summarization), consolidation (session summaries), and dreaming. Leave empty to fall back to the agent default provider." }, + "skills": { + "title": "Skills", + "description": "Tenant skill package and slash command settings.", + "maxUploadSize": "Max upload size", + "maxUploadSizeHint": "Per-file ZIP limit in MB. Allowed range: 1-500.", + "slashEnabled": "Enable slash commands", + "slashEnabledHint": "Detect /skill-name and /use skill-name at the start of user prompts.", + "slashSuggest": "Suggest similar skills", + "slashSuggestHint": "Show close matches when a requested skill is not found.", + "slashPartial": "Allow partial matching", + "slashPartialHint": "Let unique prefixes like /frontend activate matching skills.", + "slashPrefix": "Command prefix", + "slashPrefixHint": "Single character used before skill commands." + }, "compaction": { "title": "Pending Message Compaction", "description": "Summarize long pending message history to stay within context limits.", diff --git a/ui/web/src/i18n/locales/en/usage.json b/ui/web/src/i18n/locales/en/usage.json index 740e92e2..477b8608 100644 --- a/ui/web/src/i18n/locales/en/usage.json +++ b/ui/web/src/i18n/locales/en/usage.json @@ -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", diff --git a/ui/web/src/i18n/locales/vi/cli-credentials.json b/ui/web/src/i18n/locales/vi/cli-credentials.json index 9cd7f738..9d2a5889 100644 --- a/ui/web/src/i18n/locales/vi/cli-credentials.json +++ b/ui/web/src/i18n/locales/vi/cli-credentials.json @@ -39,6 +39,10 @@ "noEnvVarsHint": "Nhấn \"Thêm biến\" để khai báo biến môi trường cho công cụ CLI này.", "envKeyPlaceholder": "TÊN_BIẾN", "envValuePlaceholder": "giá trị", + "envKindSensitive": "Nhạy cảm", + "envKindValue": "Giá trị", + "envValueWarning": "Giá trị plaintext hiển thị với người dùng có quyền.", + "envValueSuspicious": "Giá trị plaintext này giống token, mật khẩu hoặc key.", "invalidEnvKey": "\"{{key}}\" không phải tên biến môi trường hợp lệ (dùng A-Z, 0-9, gạch dưới).", "binaryPathHint": "Để trống để tự động tìm từ PATH", "checkBinary": "Kiểm tra", diff --git a/ui/web/src/i18n/locales/vi/config.json b/ui/web/src/i18n/locales/vi/config.json index b0c6a77d..aaace75b 100644 --- a/ui/web/src/i18n/locales/vi/config.json +++ b/ui/web/src/i18n/locales/vi/config.json @@ -208,6 +208,10 @@ "tools.execAskMode": "Chế độ hỏi", "tools.execAskModeTip": "Khi nào nhắc người dùng phê duyệt trước khi thực thi lệnh.", "tools.execAllowlistLabel": "Lệnh được phép (mỗi dòng một mẫu)", + "tools.commandKeywordAllowlist": "Allowlist từ khóa theo lệnh", + "tools.commandKeywordAllowlistTip": "Rule JSON cho phép từ khóa đã cấu hình chỉ trong argument nội dung của CLI có credential. Deny rule theo command path vẫn áp dụng.", + "tools.commandKeywordAllowlistJsonError": "JSON không hợp lệ.", + "tools.commandKeywordAllowlistArrayError": "Dùng mảng JSON các rule.", "tools.scrubCredentials": "Ẩn thông tin xác thực", "browser.title": "Tự động hóa trình duyệt", diff --git a/ui/web/src/i18n/locales/vi/providers.json b/ui/web/src/i18n/locales/vi/providers.json index a50b7edf..dc898ade 100644 --- a/ui/web/src/i18n/locales/vi/providers.json +++ b/ui/web/src/i18n/locales/vi/providers.json @@ -293,6 +293,38 @@ "verifyFailed": "Kiểm tra embedding thất bại", "dimensionsMismatch": "{{count}} chiều — không khớp yêu cầu 1536" }, + "pricing": { + "title": "Giá Model", + "description": "Đồng bộ giá OpenRouter và ghi đè chi phí model của provider này cho usage cap.", + "skippedDescription": "Usage cap bỏ qua provider subscription và local, nhưng vẫn có thể chuẩn bị giá cho route API-based sau này.", + "sync": "Đồng bộ OpenRouter", + "model": "Model", + "saveOverride": "Lưu override", + "catalogMatches": "Kết quả trong catalog OpenRouter", + "source": "Nguồn", + "configuredFields": "Trường đã cấu hình", + "override": "Override", + "emptyOverrides": "Provider này chưa có pricing override.", + "deleteOverride": "Xóa override", + "fields": { + "input": "Input", + "output": "Output", + "cache_read": "Cache read", + "cache_write": "Cache write", + "reasoning": "Reasoning", + "request": "Request", + "image": "Image", + "web_search": "Web search" + }, + "toast": { + "synced": "Đã đồng bộ {{count}} model OpenRouter", + "syncFailed": "Không thể đồng bộ giá OpenRouter", + "saved": "Đã lưu pricing override", + "saveFailed": "Không thể lưu pricing override", + "deleted": "Đã xóa pricing override", + "deleteFailed": "Không thể xóa pricing override" + } + }, "reasoning": { "off": "Off", "offDesc": "No default extended thinking", diff --git a/ui/web/src/i18n/locales/vi/skills.json b/ui/web/src/i18n/locales/vi/skills.json index 2ea7605b..329d18e4 100644 --- a/ui/web/src/i18n/locales/vi/skills.json +++ b/ui/web/src/i18n/locales/vi/skills.json @@ -64,7 +64,7 @@ }, "upload": { "title": "Tải lên Skill", - "description": "Tải lên tệp ZIP chứa SKILL.md với YAML frontmatter (name, description, slug).", + "description": "Tải lên tệp ZIP chứa SKILL.md với YAML frontmatter (name, description, slug). Tối đa {{max}} MB mỗi tệp.", "button": "Tải lên", "uploading": "Đang tải lên...", "dropOrClick": "Nhấp hoặc kéo thả tệp .zip", @@ -72,7 +72,7 @@ "onlyZip": "Chỉ chấp nhận tệp .zip", "cancel": "Hủy", "done": "Xong", - "tooLarge": "Tệp vượt quá giới hạn 20MB", + "tooLarge": "Tệp vượt quá giới hạn tải lên đã cấu hình", "invalidZip": "Tệp ZIP không hợp lệ", "noSkillMd": "ZIP phải chứa SKILL.md ở thư mục gốc", "emptySkillMd": "SKILL.md trống", diff --git a/ui/web/src/i18n/locales/vi/system-settings.json b/ui/web/src/i18n/locales/vi/system-settings.json index 994ffc63..4100d6ef 100644 --- a/ui/web/src/i18n/locales/vi/system-settings.json +++ b/ui/web/src/i18n/locales/vi/system-settings.json @@ -64,6 +64,20 @@ "modelPlaceholder": "(mặc định)", "info": "Dùng cho vault enrichment (tóm tắt tài liệu), consolidation (tóm tắt phiên) và dreaming. Để trống sẽ fallback về agent default provider." }, + "skills": { + "title": "Skills", + "description": "Cài đặt gói skill và slash command theo tenant.", + "maxUploadSize": "Dung lượng tải lên tối đa", + "maxUploadSizeHint": "Giới hạn ZIP theo từng tệp, tính bằng MB. Khoảng cho phép: 1-500.", + "slashEnabled": "Bật slash command", + "slashEnabledHint": "Nhận diện /skill-name và /use skill-name ở đầu prompt.", + "slashSuggest": "Gợi ý skill gần đúng", + "slashSuggestHint": "Hiển thị các skill gần giống khi không tìm thấy skill được gọi.", + "slashPartial": "Cho phép khớp một phần", + "slashPartialHint": "Cho phép prefix duy nhất như /frontend kích hoạt skill tương ứng.", + "slashPrefix": "Prefix command", + "slashPrefixHint": "Một ký tự dùng trước skill command." + }, "compaction": { "title": "Nén tin nhắn chờ", "description": "Tóm tắt lịch sử tin nhắn chờ dài để nằm trong giới hạn context.", diff --git a/ui/web/src/i18n/locales/vi/usage.json b/ui/web/src/i18n/locales/vi/usage.json index 22aefdb2..dd2a9f45 100644 --- a/ui/web/src/i18n/locales/vi/usage.json +++ b/ui/web/src/i18n/locales/vi/usage.json @@ -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", diff --git a/ui/web/src/i18n/locales/zh/cli-credentials.json b/ui/web/src/i18n/locales/zh/cli-credentials.json index b0e4d929..ceca7f4e 100644 --- a/ui/web/src/i18n/locales/zh/cli-credentials.json +++ b/ui/web/src/i18n/locales/zh/cli-credentials.json @@ -39,6 +39,10 @@ "noEnvVarsHint": "点击\"添加变量\"为此 CLI 工具定义环境变量。", "envKeyPlaceholder": "变量名", "envValuePlaceholder": "值", + "envKindSensitive": "敏感", + "envKindValue": "值", + "envValueWarning": "明文值对有权限的用户可见。", + "envValueSuspicious": "此明文值看起来像令牌、密码或密钥。", "invalidEnvKey": "\"{{key}}\" 不是有效的环境变量名(使用 A-Z、0-9、下划线)。", "binaryPathHint": "留空自动从 PATH 检测", "checkBinary": "检查", diff --git a/ui/web/src/i18n/locales/zh/config.json b/ui/web/src/i18n/locales/zh/config.json index 1994f1df..21fe8e1d 100644 --- a/ui/web/src/i18n/locales/zh/config.json +++ b/ui/web/src/i18n/locales/zh/config.json @@ -208,6 +208,10 @@ "tools.execAskMode": "询问模式", "tools.execAskModeTip": "在执行命令之前何时提示用户审批。", "tools.execAllowlistLabel": "允许的命令(每行一个模式)", + "tools.commandKeywordAllowlist": "命令关键词允许列表", + "tools.commandKeywordAllowlistTip": "JSON 规则只允许配置的关键词出现在指定的凭证 CLI 内容参数中。命令路径拒绝规则仍然生效。", + "tools.commandKeywordAllowlistJsonError": "JSON 无效。", + "tools.commandKeywordAllowlistArrayError": "请使用规则 JSON 数组。", "tools.scrubCredentials": "擦除凭证", "browser.title": "浏览器自动化", diff --git a/ui/web/src/i18n/locales/zh/providers.json b/ui/web/src/i18n/locales/zh/providers.json index 282dcb05..c4f99013 100644 --- a/ui/web/src/i18n/locales/zh/providers.json +++ b/ui/web/src/i18n/locales/zh/providers.json @@ -311,6 +311,38 @@ "verifyFailed": "Embedding验证失败", "dimensionsMismatch": "{{count}}维 — 不符合要求的1536" }, + "pricing": { + "title": "模型定价", + "description": "同步 OpenRouter 价格,并为此 Provider 的模型费用设置覆盖以供用量上限使用。", + "skippedDescription": "用量上限会跳过订阅和本地 Provider,但仍可为之后的 API 路由预先配置价格。", + "sync": "同步 OpenRouter", + "model": "模型", + "saveOverride": "保存覆盖", + "catalogMatches": "OpenRouter 目录匹配", + "source": "来源", + "configuredFields": "已配置字段", + "override": "覆盖", + "emptyOverrides": "此 Provider 暂无定价覆盖。", + "deleteOverride": "删除覆盖", + "fields": { + "input": "输入", + "output": "输出", + "cache_read": "缓存读取", + "cache_write": "缓存写入", + "reasoning": "推理", + "request": "请求", + "image": "图片", + "web_search": "网页搜索" + }, + "toast": { + "synced": "已同步 {{count}} 个 OpenRouter 模型", + "syncFailed": "无法同步 OpenRouter 定价", + "saved": "定价覆盖已保存", + "saveFailed": "无法保存定价覆盖", + "deleted": "定价覆盖已删除", + "deleteFailed": "无法删除定价覆盖" + } + }, "reasoning": { "off": "Off", "offDesc": "No default extended thinking", diff --git a/ui/web/src/i18n/locales/zh/skills.json b/ui/web/src/i18n/locales/zh/skills.json index a423d2ab..e5db5acb 100644 --- a/ui/web/src/i18n/locales/zh/skills.json +++ b/ui/web/src/i18n/locales/zh/skills.json @@ -64,7 +64,7 @@ }, "upload": { "title": "上传Skill", - "description": "上传包含带 YAML frontmatter (name, description, slug) 的 SKILL.md 的 ZIP 文件。", + "description": "上传包含带 YAML frontmatter (name, description, slug) 的 SKILL.md 的 ZIP 文件。每个文件最大 {{max}} MB。", "button": "上传", "uploading": "上传中...", "dropOrClick": "点击或拖放 .zip 文件", @@ -72,7 +72,7 @@ "onlyZip": "只接受 .zip 文件", "cancel": "取消", "done": "完成", - "tooLarge": "文件超过 20MB 限制", + "tooLarge": "文件超过配置的上传限制", "invalidZip": "无效的 ZIP 文件", "noSkillMd": "ZIP 必须包含根目录的 SKILL.md", "emptySkillMd": "SKILL.md 为空", diff --git a/ui/web/src/i18n/locales/zh/system-settings.json b/ui/web/src/i18n/locales/zh/system-settings.json index b87a3aaa..7f951fdc 100644 --- a/ui/web/src/i18n/locales/zh/system-settings.json +++ b/ui/web/src/i18n/locales/zh/system-settings.json @@ -64,6 +64,20 @@ "modelPlaceholder": "(默认)", "info": "用于知识库文档摘要、会话整合和梦境。留空则回退到Agent默认Provider。" }, + "skills": { + "title": "技能", + "description": "租户技能包和斜杠命令设置。", + "maxUploadSize": "最大上传大小", + "maxUploadSizeHint": "单个 ZIP 文件限制,单位 MB。允许范围:1-500。", + "slashEnabled": "启用斜杠命令", + "slashEnabledHint": "在用户提示开头检测 /skill-name 和 /use skill-name。", + "slashSuggest": "建议相似技能", + "slashSuggestHint": "找不到请求的技能时显示接近匹配项。", + "slashPartial": "允许部分匹配", + "slashPartialHint": "允许像 /frontend 这样的唯一前缀激活匹配技能。", + "slashPrefix": "命令前缀", + "slashPrefixHint": "技能命令前使用的单个字符。" + }, "compaction": { "title": "待处理消息压缩", "description": "压缩长待处理消息历史以保持在上下文限制内。", diff --git a/ui/web/src/i18n/locales/zh/usage.json b/ui/web/src/i18n/locales/zh/usage.json index 8dd29a7d..b200c31f 100644 --- a/ui/web/src/i18n/locales/zh/usage.json +++ b/ui/web/src/i18n/locales/zh/usage.json @@ -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": "用量分析", diff --git a/ui/web/src/lib/query-keys.ts b/ui/web/src/lib/query-keys.ts index e58a2f88..3ce184e1 100644 --- a/ui/web/src/lib/query-keys.ts +++ b/ui/web/src/lib/query-keys.ts @@ -5,6 +5,8 @@ export const queryKeys = { providers: { all: ["providers"] as const, models: (providerId: string) => ["providers", providerId, "models"] as const, + pricing: (providerId: string) => ["providers", providerId, "pricing"] as const, + pricingCatalog: (model: string) => ["providers", "pricing-catalog", model] as const, chatgptOAuthStatuses: (providerKeys: string[]) => ["providers", "chatgpt-oauth-statuses", ...providerKeys] as const, chatgptOAuthQuotas: (providerNames: string[]) => ["providers", "chatgpt-oauth-quotas", ...providerNames] as const, codexPoolActivity: (providerId: string, limit: number) => ["providers", providerId, "codex-pool-activity", limit] as const, @@ -69,6 +71,11 @@ export const queryKeys = { usage: { all: ["usage"] as const, records: (params: Record) => ["usage", "records", params] as const, + caps: { + policies: ["usage", "caps", "policies"] as const, + utilization: ["usage", "caps", "utilization"] as const, + events: ["usage", "caps", "events"] as const, + }, }, teams: { all: ["teams"] as const, diff --git a/ui/web/src/pages/cli-credentials/__tests__/cli-credential-grants-dialog-helpers.test.ts b/ui/web/src/pages/cli-credentials/__tests__/cli-credential-grants-dialog-helpers.test.ts index 7f290f11..ce54fe45 100644 --- a/ui/web/src/pages/cli-credentials/__tests__/cli-credential-grants-dialog-helpers.test.ts +++ b/ui/web/src/pages/cli-credentials/__tests__/cli-credential-grants-dialog-helpers.test.ts @@ -9,7 +9,7 @@ import type { CLIAgentGrant } from "../hooks/use-cli-credentials"; describe("cli credential grant env helpers", () => { it("omits env_vars when existing masked values are not revealed", () => { const payload = buildEnvVarsPayload( - { overrideEnabled: true, entries: [{ key: "TOKEN", value: "", masked: true }] }, + { overrideEnabled: true, entries: [{ key: "TOKEN", value: "", kind: "sensitive", masked: true }] }, true, ); expect(payload).toBeUndefined(); @@ -20,14 +20,14 @@ describe("cli credential grant env helpers", () => { { overrideEnabled: true, entries: [ - { key: " CLI_ENV ", value: "agent-value", masked: false }, - { key: "", value: "ignored", masked: false }, - { key: "MASKED", value: "", masked: true }, + { key: " CLI_ENV ", value: "agent-value", kind: "value", masked: false }, + { key: "", value: "ignored", kind: "sensitive", masked: false }, + { key: "MASKED", value: "", kind: "sensitive", masked: true }, ], }, false, ); - expect(payload).toEqual({ CLI_ENV: "agent-value" }); + expect(payload).toEqual({ CLI_ENV: { kind: "value", value: "agent-value" } }); }); it("clears existing env override when override is disabled", () => { @@ -39,13 +39,31 @@ describe("cli credential grant env helpers", () => { const state = envStateFromGrant({ env_set: true, env_keys: ["API_KEY", "TOKEN"], - } as CLIAgentGrant); + } as unknown as CLIAgentGrant); expect(state).toEqual({ overrideEnabled: true, entries: [ - { key: "API_KEY", value: "", masked: true }, - { key: "TOKEN", value: "", masked: true }, + { key: "API_KEY", value: "", kind: "sensitive", masked: true }, + { key: "TOKEN", value: "", kind: "sensitive", masked: true }, + ], + }); + }); + + it("derives visible value entries from sanitized grant env metadata", () => { + const state = envStateFromGrant({ + env_set: true, + env: { + PUBLIC_BASE_URL: { kind: "value", value: "https://goclaw.sh", masked: false }, + TOKEN: { kind: "sensitive", value: null, masked: true }, + }, + } as unknown as CLIAgentGrant); + + expect(state).toEqual({ + overrideEnabled: true, + entries: [ + { key: "PUBLIC_BASE_URL", value: "https://goclaw.sh", kind: "value", masked: false }, + { key: "TOKEN", value: "", kind: "sensitive", masked: true }, ], }); }); diff --git a/ui/web/src/pages/cli-credentials/cli-credential-env-vars-section.tsx b/ui/web/src/pages/cli-credentials/cli-credential-env-vars-section.tsx index 89fa83ab..528bf6fe 100644 --- a/ui/web/src/pages/cli-credentials/cli-credential-env-vars-section.tsx +++ b/ui/web/src/pages/cli-credentials/cli-credential-env-vars-section.tsx @@ -4,11 +4,14 @@ import { Plus, X } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import type { CLIPreset } from "./hooks/use-cli-credentials"; +import type { CLIEnvEntryKind } from "@/types/cli-credential"; export interface ManualEnvEntry { key: string; value: string; + kind: CLIEnvEntryKind; } interface CliCredentialEnvVarsSectionProps { @@ -20,6 +23,10 @@ interface CliCredentialEnvVarsSectionProps { setManualEnvEntries: (updater: (prev: ManualEnvEntry[]) => ManualEnvEntry[]) => void; } +const SUSPICIOUS_VALUE_RE = /(api[_-]?key|token|secret|password|credential|bearer\s+[a-z0-9._-]+|sk-[a-z0-9_-]{12,}|gh[pousr]_[a-z0-9_]{20,})/i; +export const isSuspiciousPlaintextEnv = (key: string, value: string) => + SUSPICIOUS_VALUE_RE.test(`${key}=${value}`); + /** Env var inputs: preset-driven fields or free-form key/value pairs in manual mode. */ export function CliCredentialEnvVarsSection({ isManualMode, @@ -33,14 +40,14 @@ export function CliCredentialEnvVarsSection({ const { t: tc } = useTranslation("common"); const addEntry = useCallback(() => { - setManualEnvEntries((prev) => [...prev, { key: "", value: "" }]); + setManualEnvEntries((prev) => [...prev, { key: "", value: "", kind: "sensitive" }]); }, [setManualEnvEntries]); const removeEntry = useCallback((index: number) => { setManualEnvEntries((prev) => prev.filter((_, i) => i !== index)); }, [setManualEnvEntries]); - const updateEntry = useCallback((index: number, field: "key" | "value", val: string) => { + const updateEntry = useCallback((index: number, field: "key" | "value" | "kind", val: string) => { setManualEnvEntries((prev) => prev.map((entry, i) => (i === index ? { ...entry, [field]: val } : entry)), ); @@ -88,34 +95,57 @@ export function CliCredentialEnvVarsSection({

{t("form.noEnvVarsHint")}

)} {manualEnvEntries.map((entry, idx) => ( -
-
- updateEntry(idx, "key", e.target.value)} - className="text-base md:text-sm font-mono" - /> +
+
+
+ updateEntry(idx, "key", e.target.value)} + className="text-base md:text-sm font-mono" + /> +
+
+ updateEntry(idx, "value", e.target.value)} + className="text-base md:text-sm" + /> +
+
-
- updateEntry(idx, "value", e.target.value)} - className="text-base md:text-sm" - /> -
- + + + + {entry.kind === "value" && ( +

+ {isSuspiciousPlaintextEnv(entry.key, entry.value) + ? t("form.envValueSuspicious") + : t("form.envValueWarning")} +

+ )}
))}
diff --git a/ui/web/src/pages/cli-credentials/cli-credential-form-dialog.tsx b/ui/web/src/pages/cli-credentials/cli-credential-form-dialog.tsx index e7cfc08e..edeb906d 100644 --- a/ui/web/src/pages/cli-credentials/cli-credential-form-dialog.tsx +++ b/ui/web/src/pages/cli-credentials/cli-credential-form-dialog.tsx @@ -17,6 +17,7 @@ import { CliCredentialEnvVarsSection } from "./cli-credential-env-vars-section"; import { CliCredentialBinaryFields } from "./cli-credential-binary-fields"; import { CliCredentialScopeFields } from "./cli-credential-scope-fields"; import { cliCredentialSchema, type CliCredentialFormData } from "@/schemas/credential.schema"; +import type { CLIEnvEntryResponse, CLIEnvPayload } from "@/types/cli-credential"; interface Props { open: boolean; @@ -29,6 +30,20 @@ interface Props { const NONE_PRESET = "__none__"; const ENV_KEY_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/; +function manualEntriesFromEnv( + env: Record | undefined, + fallbackKeys: string[], +): ManualEnvEntry[] { + if (env && Object.keys(env).length > 0) { + return Object.entries(env).map(([key, entry]) => ({ + key, + value: entry.value ?? "", + kind: entry.kind ?? "sensitive", + })); + } + return fallbackKeys.map((key) => ({ key, value: "", kind: "sensitive" })); +} + export function CliCredentialFormDialog({ open, onOpenChange, credential, presets, onSubmit }: Props) { const { t } = useTranslation("cli-credentials"); const { t: tc } = useTranslation("common"); @@ -91,13 +106,13 @@ export function CliCredentialFormDialog({ open, onOpenChange, credential, preset return; } - const applyEnvKeys = (keys: string[]) => { + const applyEnvState = (env: Record | undefined, keys: string[]) => { setInitialEnvKeys(keys); - setManualEnvEntries(keys.length > 0 ? keys.map((k) => ({ key: k, value: "" })) : []); + setManualEnvEntries(manualEntriesFromEnv(env, keys)); }; if (credential.env_keys !== undefined) { - applyEnvKeys(credential.env_keys ?? []); + applyEnvState(credential.env, credential.env_keys ?? []); return; } @@ -106,9 +121,9 @@ export function CliCredentialFormDialog({ open, onOpenChange, credential, preset try { const full = await http.get(`/v1/cli-credentials/${credential.id}`); if (cancelled) return; - applyEnvKeys(full.env_keys ?? []); + applyEnvState(full.env, full.env_keys ?? []); } catch { - if (!cancelled) applyEnvKeys([]); + if (!cancelled) applyEnvState(undefined, []); } })(); return () => { cancelled = true; }; @@ -156,16 +171,22 @@ export function CliCredentialFormDialog({ open, onOpenChange, credential, preset const splitCommaList = (v: string): string[] => v.split(",").map((s) => s.trim()).filter(Boolean); - const buildEnvPayload = (): Record | null => { - if (!isManualMode) return envValues; - const env: Record = {}; + const buildEnvPayload = (): CLIEnvPayload | null => { + if (!isManualMode) { + const presetEnv: CLIEnvPayload = {}; + for (const [key, value] of Object.entries(envValues)) { + presetEnv[key] = { kind: "sensitive", value }; + } + return presetEnv; + } + const env: CLIEnvPayload = {}; for (const entry of manualEnvEntries) { const k = entry.key.trim(); if (k && !ENV_KEY_PATTERN.test(k)) { setError(t("form.invalidEnvKey", { key: k })); return null; } - if (k) env[k] = entry.value; + if (k) env[k] = { kind: entry.kind, value: entry.value }; } return env; }; diff --git a/ui/web/src/pages/cli-credentials/cli-credential-grant-env-row.tsx b/ui/web/src/pages/cli-credentials/cli-credential-grant-env-row.tsx new file mode 100644 index 00000000..1d14eada --- /dev/null +++ b/ui/web/src/pages/cli-credentials/cli-credential-grant-env-row.tsx @@ -0,0 +1,73 @@ +import { X } from "lucide-react"; +import { useTranslation } from "react-i18next"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { isSuspiciousPlaintextEnv } from "./cli-credential-env-vars-section"; +import type { GrantEnvEntry } from "./cli-credential-grant-env-section"; + +interface Props { + entry: GrantEnvEntry; + hasError: boolean; + onRemove: () => void; + onUpdate: (field: "key" | "value" | "kind", value: string) => void; +} + +export function CliCredentialGrantEnvRow({ entry, hasError, onRemove, onUpdate }: Props) { + const { t } = useTranslation("cli-credentials"); + + return ( +
+
+ onUpdate("key", e.target.value)} + className={`text-base md:text-sm font-mono${hasError ? " border-destructive" : ""}`} + /> + {hasError && ( +

+ {t("grants.envVars.deniedKey", { key: entry.key })} +

+ )} +
+
+ {entry.masked ? ( + + ) : ( + onUpdate("value", e.target.value)} + className="text-base md:text-sm" + /> + )} +
+ {!entry.masked && ( + + )} + + {!entry.masked && entry.kind === "value" && ( +

+ {isSuspiciousPlaintextEnv(entry.key, entry.value) + ? t("form.envValueSuspicious") + : t("form.envValueWarning")} +

+ )} +
+ ); +} diff --git a/ui/web/src/pages/cli-credentials/cli-credential-grant-env-section.tsx b/ui/web/src/pages/cli-credentials/cli-credential-grant-env-section.tsx index ebff9a70..5e546540 100644 --- a/ui/web/src/pages/cli-credentials/cli-credential-grant-env-section.tsx +++ b/ui/web/src/pages/cli-credentials/cli-credential-grant-env-section.tsx @@ -6,13 +6,14 @@ */ import { useState, useCallback, useEffect, useRef } from "react"; import { useTranslation } from "react-i18next"; -import { Plus, X, Eye } from "lucide-react"; +import { Plus, Eye } from "lucide-react"; import { Button } from "@/components/ui/button"; -import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Switch } from "@/components/ui/switch"; import { toast } from "@/stores/use-toast-store"; import { useHttp } from "@/hooks/use-ws"; +import type { CLIEnvEntryKind } from "@/types/cli-credential"; +import { CliCredentialGrantEnvRow } from "./cli-credential-grant-env-row"; // Keep in sync with internal/crypto/env_denylist.go. // Backend is authoritative; this list drives inline UX warnings only. @@ -23,7 +24,7 @@ const ENV_DENYLIST_EXACT = new Set([ "PYTHONPATH", "PYTHONHOME", "PYTHONSTARTUP", "GIT_SSH_COMMAND", "GIT_SSH", "GIT_EXEC_PATH", "GIT_CONFIG_SYSTEM", "SSH_AUTH_SOCK", - // Finding #6 additions — keep in sync with internal/crypto/env_denylist.go + // Shell startup and proxy/certificate variables can alter command behavior. "BASH_ENV", "ENV", "PROMPT_COMMAND", "PERL5LIB", "RUBYOPT", "HTTPS_PROXY", "HTTP_PROXY", "NO_PROXY", @@ -36,6 +37,7 @@ const ENV_DENYLIST_PREFIXES = ["DYLD_", "GOCLAW_", "LD_", "NPM_CONFIG_"]; export interface GrantEnvEntry { key: string; value: string; + kind: CLIEnvEntryKind; masked: boolean; // true = not yet revealed from server } @@ -63,10 +65,10 @@ export function CliCredentialGrantEnvSection({ const [revealing, setRevealing] = useState(false); const [revealed, setRevealed] = useState(false); const { overrideEnabled, entries } = state; - // Finding #10: track blur timeout so we can cancel it on reveal/unmount. + // Track reveal timeout so plaintext can be cleared on reveal refresh/unmount. const blurTimeoutRef = useRef | null>(null); - // Finding #10: clear revealed plaintext from entries on component unmount. + // Clear revealed plaintext from entries on component unmount. // This is defense-in-depth — plaintext should not persist in React state beyond use. useEffect(() => { return () => { @@ -77,7 +79,6 @@ export function CliCredentialGrantEnvSection({ entries: state.entries.map((e) => ({ ...e, value: "", masked: e.masked })), }); }; - // eslint-disable-next-line react-hooks/exhaustive-deps }, []); const setEntries = useCallback( @@ -89,10 +90,10 @@ export function CliCredentialGrantEnvSection({ const handleToggle = useCallback((checked: boolean) => { if (checked) { if (initialEnvSet && !revealed && entries.every((e) => e.masked)) { - const masked: GrantEnvEntry[] = initialEnvKeys.map((k) => ({ key: k, value: "", masked: true })); - onChange({ overrideEnabled: true, entries: masked.length > 0 ? masked : [{ key: "", value: "", masked: false }] }); + const masked: GrantEnvEntry[] = initialEnvKeys.map((k) => ({ key: k, value: "", kind: "sensitive", masked: true })); + onChange({ overrideEnabled: true, entries: masked.length > 0 ? masked : [{ key: "", value: "", kind: "sensitive", masked: false }] }); } else if (entries.length === 0) { - onChange({ overrideEnabled: true, entries: [{ key: "", value: "", masked: false }] }); + onChange({ overrideEnabled: true, entries: [{ key: "", value: "", kind: "sensitive", masked: false }] }); } else { onChange({ overrideEnabled: true, entries }); } @@ -105,16 +106,19 @@ export function CliCredentialGrantEnvSection({ if (!grantId) return; setRevealing(true); try { - // POST — not GET (C1 red-team). Direct call, not cached by TanStack Query. + // POST keeps reveal out of URL/history and avoids query caching. const res = await http.post<{ env_vars: Record }>( `/v1/cli-credentials/${binaryId}/agent-grants/${grantId}/env:reveal`, ); const filled: GrantEnvEntry[] = Object.entries(res.env_vars).map(([k, v]) => ({ - key: k, value: v, masked: false, + key: k, + value: v, + kind: entries.find((entry) => entry.key === k)?.kind ?? "sensitive", + masked: false, })); onChange({ overrideEnabled: true, entries: filled.length > 0 ? filled : entries }); setRevealed(true); - // Finding #10: wipe plaintext after 30s of inactivity (defense-in-depth). + // Wipe plaintext after 30s of inactivity. if (blurTimeoutRef.current) clearTimeout(blurTimeoutRef.current); blurTimeoutRef.current = setTimeout(() => { onChange({ @@ -133,9 +137,9 @@ export function CliCredentialGrantEnvSection({ } }, [grantId, binaryId, http, onChange, entries, t]); - const addEntry = useCallback(() => setEntries((p) => [...p, { key: "", value: "", masked: false }]), [setEntries]); + const addEntry = useCallback(() => setEntries((p) => [...p, { key: "", value: "", kind: "sensitive", masked: false }]), [setEntries]); const removeEntry = useCallback((i: number) => setEntries((p) => p.filter((_, j) => j !== i)), [setEntries]); - const updateEntry = useCallback((i: number, f: "key" | "value", v: string) => + const updateEntry = useCallback((i: number, f: "key" | "value" | "kind", v: string) => setEntries((p) => p.map((e, j) => j === i ? { ...e, [f]: v, masked: false } : e)), [setEntries]); const isDenied = (k: string) => { @@ -174,32 +178,13 @@ export function CliCredentialGrantEnvSection({ {entries.map((entry, idx) => { const hasError = isDenied(entry.key) || isRejected(entry.key); return ( -
-
- updateEntry(idx, "key", e.target.value)} - className={`text-base md:text-sm font-mono${hasError ? " border-destructive" : ""}`} /> - {hasError && ( -

- {t("grants.envVars.deniedKey", { key: entry.key })} -

- )} -
-
- {entry.masked ? ( - - ) : ( - updateEntry(idx, "value", e.target.value)} - className="text-base md:text-sm" /> - )} -
- -
+ removeEntry(idx)} + onUpdate={(field, value) => updateEntry(idx, field, value)} + /> ); })} {entries.length === 0 && ( diff --git a/ui/web/src/pages/cli-credentials/cli-credential-grants-dialog-helpers.ts b/ui/web/src/pages/cli-credentials/cli-credential-grants-dialog-helpers.ts index 3d5e0ac2..c59d33ea 100644 --- a/ui/web/src/pages/cli-credentials/cli-credential-grants-dialog-helpers.ts +++ b/ui/web/src/pages/cli-credentials/cli-credential-grants-dialog-helpers.ts @@ -3,6 +3,7 @@ */ import type { CLIAgentGrant } from "./hooks/use-cli-credentials"; import type { GrantEnvState, GrantEnvEntry } from "./cli-credential-grant-env-section"; +import type { CLIEnvPayload } from "@/types/cli-credential"; export const EMPTY_ENV_STATE: GrantEnvState = { overrideEnabled: false, entries: [] }; @@ -15,14 +16,14 @@ export const EMPTY_ENV_STATE: GrantEnvState = { overrideEnabled: false, entries: export function buildEnvVarsPayload( envState: GrantEnvState, originalEnvSet: boolean, -): Record | null | undefined { +): CLIEnvPayload | null | undefined { const { overrideEnabled, entries } = envState; if (overrideEnabled) { const allMasked = entries.length > 0 && entries.every((e: GrantEnvEntry) => e.masked); if (allMasked) return undefined; // not revealed; don't overwrite - const result: Record = {}; + const result: CLIEnvPayload = {}; for (const e of entries) { - if (!e.masked && e.key.trim()) result[e.key.trim()] = e.value; + if (!e.masked && e.key.trim()) result[e.key.trim()] = { kind: e.kind, value: e.value }; } return result; } @@ -31,10 +32,23 @@ export function buildEnvVarsPayload( /** Derive initial GrantEnvState from an existing grant. */ export function envStateFromGrant(grant: CLIAgentGrant): GrantEnvState { + if (grant.env && Object.keys(grant.env).length > 0) { + return { + overrideEnabled: true, + entries: Object.entries(grant.env) + .sort(([a], [b]) => a.localeCompare(b)) + .map(([key, entry]) => ({ + key, + value: entry.kind === "value" && entry.value !== null ? entry.value : "", + kind: entry.kind, + masked: entry.masked, + })), + }; + } if (grant.env_set && grant.env_keys && grant.env_keys.length > 0) { return { overrideEnabled: true, - entries: grant.env_keys.map((k) => ({ key: k, value: "", masked: true })), + entries: grant.env_keys.map((k) => ({ key: k, value: "", kind: "sensitive", masked: true })), }; } return EMPTY_ENV_STATE; diff --git a/ui/web/src/pages/cli-credentials/cli-user-credentials-dialog.tsx b/ui/web/src/pages/cli-credentials/cli-user-credentials-dialog.tsx index 7cfe6c57..79a95a9a 100644 --- a/ui/web/src/pages/cli-credentials/cli-user-credentials-dialog.tsx +++ b/ui/web/src/pages/cli-credentials/cli-user-credentials-dialog.tsx @@ -13,11 +13,12 @@ import { Button } from "@/components/ui/button"; import { Label } from "@/components/ui/label"; import { Badge } from "@/components/ui/badge"; import { UserPickerCombobox } from "@/components/shared/user-picker-combobox"; -import { KeyValueEditor } from "@/components/shared/key-value-editor"; import { toast } from "@/stores/use-toast-store"; import { useHttp } from "@/hooks/use-ws"; import i18next from "i18next"; +import { CliCredentialEnvVarsSection, type ManualEnvEntry } from "./cli-credential-env-vars-section"; import type { SecureCLIBinary } from "./hooks/use-cli-credentials"; +import type { CLIEnvEntryResponse, CLIEnvPayload } from "@/types/cli-credential"; interface UserCredEntry { id: string; @@ -36,11 +37,26 @@ interface CLIUserCredentialsDialogProps { binary: SecureCLIBinary; } -const SENSITIVE_ENV_RE = /^.*(key|secret|token|password|credential).*$/i; -const isSensitiveEnv = (key: string) => SENSITIVE_ENV_RE.test(key.trim()); - type ViewState = "list" | "form"; +function entriesFromEnv(env: Record | null | undefined): ManualEnvEntry[] { + if (!env || Object.keys(env).length === 0) return []; + return Object.entries(env).map(([key, entry]) => ({ + key, + value: entry.value ?? "", + kind: entry.kind ?? "sensitive", + })); +} + +function envPayloadFromEntries(entries: ManualEnvEntry[]): CLIEnvPayload { + const env: CLIEnvPayload = {}; + for (const entry of entries) { + const key = entry.key.trim(); + if (key) env[key] = { kind: entry.kind, value: entry.value }; + } + return env; +} + export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUserCredentialsDialogProps) { const { t } = useTranslation("cli-credentials"); const http = useHttp(); @@ -54,7 +70,7 @@ export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUser const [userId, setUserId] = useState(""); // Separate search text from selected value (onChange fires on every keystroke) const [userSearchText, setUserSearchText] = useState(""); - const [env, setEnv] = useState>({}); + const [envEntries, setEnvEntries] = useState([]); const [saving, setSaving] = useState(false); const [deleting, setDeletingId] = useState(null); @@ -80,7 +96,7 @@ export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUser setEditEntry(null); setUserId(""); setUserSearchText(""); - setEnv({}); + setEnvEntries([]); loadList(); }, [open, loadList]); @@ -88,7 +104,7 @@ export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUser setEditEntry(null); setUserId(""); setUserSearchText(""); - setEnv({}); + setEnvEntries([]); setView("form"); }; @@ -96,14 +112,14 @@ export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUser setEditEntry(entry); setUserId(entry.user_id); setUserSearchText(entry.user_id); - setEnv({}); + setEnvEntries([]); setView("form"); // Load existing env for edit try { - const res = await http.get<{ user_id: string; env: Record | null }>( + const res = await http.get<{ user_id: string; env: Record | null }>( `/v1/cli-credentials/${binary.id}/user-credentials/${entry.user_id}`, ); - setEnv(res.env ?? {}); + setEnvEntries(entriesFromEnv(res.env)); } catch { // leave env empty — user can re-enter } @@ -112,6 +128,7 @@ export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUser const handleSave = async () => { const uid = userId.trim(); if (!uid) return; + const env = envPayloadFromEntries(envEntries); // New entry needs at least one variable; edits may clear all keys (empty object). if (!editEntry && Object.keys(env).length === 0) { toast.error(i18next.t("cli-credentials:userCredentials.envRequired")); @@ -251,13 +268,13 @@ export function CLIUserCredentialsDialog({ open, onOpenChange, binary }: CLIUser
- undefined} + manualEnvEntries={envEntries} + setManualEnvEntries={setEnvEntries} />
diff --git a/ui/web/src/pages/config/sections/tools-exec-section.tsx b/ui/web/src/pages/config/sections/tools-exec-section.tsx index a0caf55a..7486f094 100644 --- a/ui/web/src/pages/config/sections/tools-exec-section.tsx +++ b/ui/web/src/pages/config/sections/tools-exec-section.tsx @@ -19,7 +19,6 @@ import { } from "@/components/ui/select"; import { InfoLabel } from "@/components/shared/info-label"; - type ToolsData = Record; interface Props { @@ -31,10 +30,14 @@ interface Props { export function ToolsExecSection({ data, onSave, saving }: Props) { const { t } = useTranslation("config"); const [draft, setDraft] = useState(data ?? {}); + const [allowlistText, setAllowlistText] = useState(""); + const [allowlistError, setAllowlistError] = useState(null); const [dirty, setDirty] = useState(false); useEffect(() => { setDraft(data ?? {}); + setAllowlistText(JSON.stringify(data?.commandKeywordAllowlist ?? [], null, 2)); + setAllowlistError(null); setDirty(false); }, [data]); @@ -46,6 +49,30 @@ export function ToolsExecSection({ data, onSave, saving }: Props) { setDirty(true); }; + const updateCommandKeywordAllowlist = (value: string) => { + setAllowlistText(value); + setDirty(true); + + const trimmed = value.trim(); + if (!trimmed) { + setAllowlistError(null); + setDraft((prev) => ({ ...prev, commandKeywordAllowlist: [] })); + return; + } + + try { + const parsed = JSON.parse(trimmed); + if (!Array.isArray(parsed)) { + setAllowlistError(t("tools.commandKeywordAllowlistArrayError")); + return; + } + setAllowlistError(null); + setDraft((prev) => ({ ...prev, commandKeywordAllowlist: parsed })); + } catch { + setAllowlistError(t("tools.commandKeywordAllowlistJsonError")); + } + }; + if (!data) return null; const exec = draft.execApproval ?? {}; @@ -91,15 +118,41 @@ export function ToolsExecSection({ data, onSave, saving }: Props) { allowlist: e.target.value.split("\n").filter(Boolean), }) } - className="min-h-[80px] font-mono text-xs" + className="min-h-[80px] font-mono text-base md:text-sm" placeholder="git * npm * ls *" /> )} +
+ + {t("tools.commandKeywordAllowlist")} + +