Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cmd/gateway_http_wiring.go
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer(

// Usage analytics API
if d.pgStores.Snapshots != nil {
d.server.SetUsageHandler(httpapi.NewUsageHandler(d.pgStores.Snapshots, d.pgStores.DB))
d.server.SetUsageHandler(httpapi.NewUsageHandler(d.pgStores.Snapshots, d.pgStores.UsageEvents, d.pgStores.DB))
}
if d.pgStores.UsageCaps != nil {
d.server.SetUsageCapsHandler(httpapi.NewUsageCapsHandler(d.pgStores.UsageCaps, d.pgStores.Tenants))
Expand Down
1 change: 1 addition & 0 deletions cmd/gateway_managed.go
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,7 @@ func wireExtras(
ModelPricing: appCfg.Telemetry.ModelPricing,
TracingStore: stores.Tracing,
UsageCaps: usageCapSvc,
UsageEvents: stores.UsageEvents,
MemoryStore: stores.Memory,
ContactStore: stores.Contacts,
TenantStore: stores.Tenants,
Expand Down
2 changes: 1 addition & 1 deletion cmd/gateway_setup.go
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ func wireTracingAndCron(
// Start snapshot worker for hourly usage aggregation
var snapshotWorker *tracing.SnapshotWorker
if stores.Snapshots != nil {
snapshotWorker = tracing.NewSnapshotWorker(stores.DB, stores.Snapshots)
snapshotWorker = tracing.NewSnapshotWorker(stores.DB, stores.Snapshots, stores.UsageEvents)
snapshotWorker.Start()

// Backfill historical data in background
Expand Down
3 changes: 3 additions & 0 deletions internal/agent/loop_context.go
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,8 @@ func (l *Loop) injectContext(ctx context.Context, req *RunRequest) (contextSetup
AgentKey: l.id,
TenantID: l.tenantID,
UserID: req.UserID,
RunID: req.RunID,
SessionKey: req.SessionKey,
CredentialUserID: credUserID,
AgentType: l.agentType,
SenderID: req.SenderID,
Expand All @@ -388,6 +390,7 @@ func (l *Loop) injectContext(ctx context.Context, req *RunRequest) (contextSetup
SharedContext: store.IsSharedContext(ctx),
RestrictToWorkspace: l.restrictToWs != nil && *l.restrictToWs,
BuiltinToolSettings: l.builtinToolSettings,
Channel: req.Channel,
ChannelType: req.ChannelType,
SubagentsCfg: l.subagentsCfg,
ParentModel: l.model,
Expand Down
30 changes: 24 additions & 6 deletions internal/agent/loop_pipeline_tool_callbacks.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx
emitRun := makeToolEmitRun(l, req)
return func(ctx context.Context, state *pipeline.RunState, tc providers.ToolCall) ([]providers.Message, error) {
tc = l.normalizeToolCall(tc)
registryName := l.resolveToolCallName(tc.Name)
registryName := l.canonicalToolName(l.resolveToolCallName(tc.Name))
argsJSON, _ := json.Marshal(tc.Arguments)
slog.Info("tool call", "agent", l.id, "tool", tc.Name, "args_len", len(argsJSON))

Expand All @@ -35,7 +35,7 @@ func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx

// Emit tool span start for tracing.
toolStart := time.Now().UTC()
toolSpanID := l.emitToolSpanStart(ctx, toolStart, tc.Name, tc.ID, string(argsJSON))
toolSpanID := l.emitToolSpanStart(ctx, toolStart, registryName, tc.ID, string(argsJSON))

// Inject agent audio snapshot so TTS tool (and any future audio consumers)
// can read agent-level voice/model config without an extra DB lookup.
Expand All @@ -56,6 +56,7 @@ func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx
toolDuration := time.Since(toolStart)

l.emitToolSpanEnd(ctx, toolSpanID, toolStart, result)
l.recordToolUsageEvent(ctx, req, registryName, tc.Name, tc.ID, tc.Arguments, toolStart, result, toolSpanID)

// v3 evolution metrics: record tool execution non-blocking (best-effort).
l.recordToolMetric(ctx, req.SessionKey, registryName, !result.IsError, toolDuration)
Expand All @@ -77,6 +78,10 @@ func (l *Loop) makeExecuteToolCall(req *RunRequest, bridgeRS *runState) func(ctx
type toolRawResult struct {
result *tools.Result
duration time.Duration
start time.Time
spanID uuid.UUID
toolName string
rawName string
}

// makeExecuteToolRaw wraps tool I/O only (parallel-safe, no state mutation).
Expand All @@ -85,7 +90,7 @@ func (l *Loop) makeExecuteToolRaw(req *RunRequest) func(ctx context.Context, tc
emitRun := makeToolEmitRun(l, req)
return func(ctx context.Context, tc providers.ToolCall) (providers.Message, any, error) {
tc = l.normalizeToolCall(tc)
registryName := l.resolveToolCallName(tc.Name)
registryName := l.canonicalToolName(l.resolveToolCallName(tc.Name))
argsJSON, _ := json.Marshal(tc.Arguments)
slog.Info("tool call", "agent", l.id, "tool", tc.Name, "args_len", len(argsJSON))

Expand All @@ -102,7 +107,7 @@ func (l *Loop) makeExecuteToolRaw(req *RunRequest) func(ctx context.Context, tc

// Emit tool span start (goroutine-safe: channel send only).
start := time.Now().UTC()
spanID := l.emitToolSpanStart(ctx, start, tc.Name, tc.ID, string(argsJSON))
spanID := l.emitToolSpanStart(ctx, start, registryName, tc.ID, string(argsJSON))

// Inject agent audio snapshot (parallel path — same as sequential makeExecuteToolCall).
if l.agentUUID != uuid.Nil {
Expand All @@ -129,7 +134,7 @@ func (l *Loop) makeExecuteToolRaw(req *RunRequest) func(ctx context.Context, tc
ToolCallID: tc.ID,
IsError: result.IsError,
}
return msg, &toolRawResult{result: result, duration: dur}, nil
return msg, &toolRawResult{result: result, duration: dur, start: start, spanID: spanID, toolName: registryName, rawName: tc.Name}, nil
}
}

Expand All @@ -155,20 +160,33 @@ func (l *Loop) makeProcessToolResult(req *RunRequest, bridgeRS *runState) func(c
emitRun := makeToolEmitRun(l, req)
return func(ctx context.Context, state *pipeline.RunState, tc providers.ToolCall, rawMsg providers.Message, rawData any) []providers.Message {
tc = l.normalizeToolCall(tc)
registryName := l.resolveToolCallName(tc.Name)
registryName := l.canonicalToolName(l.resolveToolCallName(tc.Name))

// Extract result and timing from toolRawResult wrapper.
var result *tools.Result
var dur time.Duration
var start time.Time
var spanID uuid.UUID
var rawName string
if raw, ok := rawData.(*toolRawResult); ok && raw != nil {
result = raw.result
dur = raw.duration
start = raw.start
spanID = raw.spanID
registryName = raw.toolName
rawName = raw.rawName
} else if r, ok := rawData.(*tools.Result); ok {
result = r // backward compat
}
if result == nil {
return []providers.Message{rawMsg}
}
if rawName == "" {
rawName = tc.Name
}
if !start.IsZero() {
l.recordToolUsageEvent(ctx, req, registryName, rawName, tc.ID, tc.Arguments, start, result, spanID)
}

// Record tool metrics (non-blocking, best-effort).
l.recordToolMetric(ctx, req.SessionKey, registryName, !result.IsError, dur)
Expand Down
3 changes: 3 additions & 0 deletions internal/agent/loop_types.go
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,7 @@ type Loop struct {
budgetMonthlyCents int
tracingStore store.TracingStore
usageCaps *usagecaps.Service
usageEvents store.UsageEventStore

// Memory store for extractive memory fallback (writes directly when LLM flush fails)
memStore store.MemoryStore
Expand Down Expand Up @@ -441,6 +442,7 @@ type LoopConfig struct {
BudgetMonthlyCents int
TracingStore store.TracingStore
UsageCaps *usagecaps.Service
UsageEvents store.UsageEventStore

// Memory store for extractive memory fallback (writes directly when LLM flush fails)
MemoryStore store.MemoryStore
Expand Down Expand Up @@ -586,6 +588,7 @@ func NewLoop(cfg LoopConfig) *Loop {
budgetMonthlyCents: cfg.BudgetMonthlyCents,
tracingStore: cfg.TracingStore,
usageCaps: cfg.UsageCaps,
usageEvents: cfg.UsageEvents,
memStore: cfg.MemoryStore,
mcpStore: cfg.MCPStore,
mcpPool: cfg.MCPPool,
Expand Down
2 changes: 2 additions & 0 deletions internal/agent/resolver.go
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ type ResolverDeps struct {
// Tracing store for budget enforcement queries
TracingStore store.TracingStore
UsageCaps *usagecaps.Service
UsageEvents store.UsageEventStore

// Memory store for extractive memory fallback
MemoryStore store.MemoryStore
Expand Down Expand Up @@ -536,6 +537,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc {
BudgetMonthlyCents: derefInt(ag.BudgetMonthlyCents),
TracingStore: deps.TracingStore,
UsageCaps: deps.UsageCaps,
UsageEvents: deps.UsageEvents,
MemoryStore: deps.MemoryStore,
MCPStore: deps.MCPStore,
MCPPool: deps.MCPPool,
Expand Down
1 change: 1 addition & 0 deletions internal/agent/skill_slash_commands.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ func (l *Loop) applySkillSlashCommand(ctx context.Context, req *RunRequest, mess
message = result.RemainingPrompt
}
skillFilter = []string{result.Skill.Slug}
l.recordSkillSlashUsageEvent(ctx, result.Skill.Slug)
l.recordSkillUsage(ctx, req, result.Skill.Slug, "", "slash", store.SkillUsageStatusStarted, "", 0)
case skillSlashCommandList:
message = "List the available skills shown in the system instructions."
Expand Down
200 changes: 200 additions & 0 deletions internal/agent/usage_events.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
package agent

import (
"context"
"encoding/json"
"log/slog"
"slices"
"time"

"github.com/google/uuid"

"github.com/nextlevelbuilder/goclaw/internal/store"
"github.com/nextlevelbuilder/goclaw/internal/tools"
"github.com/nextlevelbuilder/goclaw/internal/tracing"
)

type mcpUsageTool interface {
ServerName() string
OriginalName() string
}

func (l *Loop) canonicalToolName(name string) string {
if l.registry == nil {
return name
}
if tool, ok := l.registry.Get(name); ok && tool != nil {
return tool.Name()
}
return name
}

func (l *Loop) recordToolUsageEvent(ctx context.Context, req *RunRequest, canonicalName, rawName, toolCallID string, args map[string]any, start time.Time, result *tools.Result, spanID uuid.UUID) {
if l.usageEvents == nil || result == nil {
return
}
if tracing.TraceIDFromContext(ctx) == uuid.Nil {
return
}

resourceName := canonicalName
resourceID := canonicalName
eventType := store.UsageEventTypeToolCall
resourceType := store.UsageResourceTypeTool
source := store.UsageSourceToolCall
metadata := map[string]any{}

if canonicalName == "use_skill" {
eventType = store.UsageEventTypeSkillActivation
resourceType = store.UsageResourceTypeSkill
source = store.UsageSourceUseSkill
if skill, _ := args["name"].(string); skill != "" {
resourceName = skill
resourceID = skill
}
} else if l.registry != nil {
tool, ok := l.registry.Get(canonicalName)
if !ok {
tool = nil
}
if mcpTool, ok := tool.(mcpUsageTool); ok {
eventType = store.UsageEventTypeMCPToolCall
resourceType = store.UsageResourceTypeMCPTool
resourceName = mcpTool.ServerName() + "/" + mcpTool.OriginalName()
resourceID = mcpTool.OriginalName()
metadata["server"] = mcpTool.ServerName()
metadata["tool"] = mcpTool.OriginalName()
} else if l.isRuntimeTool(canonicalName) {
eventType = store.UsageEventTypeRuntimeToolCall
resourceType = store.UsageResourceTypeRuntimeTool
}
}

if rawName != "" && rawName != canonicalName {
metadata["raw_tool_name"] = rawName
}

event := l.baseUsageEvent(ctx, req, start, eventType, resourceType, resourceName, resourceID, source)
event.SpanID = uuidPtr(spanID)
event.Status = "completed"
if result.IsError {
event.Status = "error"
event.ErrorCount = 1
}
event.DurationMS = int(time.Since(start).Milliseconds())
if result.Usage != nil {
event.InputTokens = int64(result.Usage.PromptTokens)
event.OutputTokens = int64(result.Usage.CompletionTokens)
event.TotalTokens = int64(result.Usage.TotalTokens)
if event.TotalTokens == 0 {
event.TotalTokens = event.InputTokens + event.OutputTokens
}
}
event.Provider = result.Provider
event.Model = result.Model
event.Metadata = usageMetadata(metadata)
l.insertUsageEventBestEffort(ctx, event)
}

func (l *Loop) recordSkillSlashUsageEvent(ctx context.Context, skillSlug string) {
if l.usageEvents == nil || skillSlug == "" {
return
}
traceID := tracing.TraceIDFromContext(ctx)
if traceID == uuid.Nil {
return
}
rc := store.RunContextFromCtx(ctx)
event := l.baseUsageEvent(ctx, nil, time.Now().UTC(),
store.UsageEventTypeSkillActivation,
store.UsageResourceTypeSkill,
skillSlug,
skillSlug,
store.UsageSourceSlashCommand,
)
event.TraceID = uuidPtr(traceID)
if rc != nil {
event.RunID = rc.RunID
event.SessionKey = rc.SessionKey
event.Channel = rc.Channel
if rc.TeamID != "" {
if teamID, err := uuid.Parse(rc.TeamID); err == nil {
event.TeamID = &teamID
}
}
}
event.Metadata = usageMetadata(map[string]any{"activation_source": store.UsageSourceSlashCommand})
l.insertUsageEventBestEffort(ctx, event)
}

func (l *Loop) baseUsageEvent(ctx context.Context, req *RunRequest, eventTime time.Time, eventType, resourceType, resourceName, resourceID, source string) store.UsageEvent {
tenantID := store.TenantIDFromContext(ctx)
if tenantID == uuid.Nil {
tenantID = l.tenantID
}
traceID := tracing.TraceIDFromContext(ctx)
event := store.UsageEvent{
ID: uuid.New(),
TenantID: tenantID,
EventTime: eventTime.UTC(),
BucketHour: eventTime.UTC().Truncate(time.Hour),
EventType: eventType,
ResourceType: resourceType,
ResourceName: resourceName,
ResourceID: resourceID,
Source: source,
AgentID: uuidPtr(l.agentUUID),
TeamID: tracing.TraceTeamIDPtrFromContext(ctx),
TraceID: uuidPtr(traceID),
Status: "completed",
CallCount: 1,
}
if req != nil {
event.RunID = req.RunID
event.SessionKey = req.SessionKey
event.Channel = req.Channel
if req.TeamID != "" && event.TeamID == nil {
if teamID, err := uuid.Parse(req.TeamID); err == nil {
event.TeamID = &teamID
}
}
}
return event
}

func (l *Loop) isRuntimeTool(toolName string) bool {
if l.registry == nil {
return false
}
members, ok := l.registry.GetToolGroup("runtime")
return ok && slices.Contains(members, toolName)
}

func (l *Loop) insertUsageEventBestEffort(ctx context.Context, event store.UsageEvent) {
tenantID := event.TenantID
go func() {
bgCtx, cancel := context.WithTimeout(store.WithTenantID(context.Background(), tenantID), 5*time.Second)
defer cancel()
if err := l.usageEvents.InsertEvent(bgCtx, &event); err != nil {
slog.Debug("usage.event.record_failed", "resource", event.ResourceName, "error", err)
}
}()
}

func uuidPtr(id uuid.UUID) *uuid.UUID {
if id == uuid.Nil {
return nil
}
return &id
}

func usageMetadata(values map[string]any) json.RawMessage {
if len(values) == 0 {
return nil
}
data, err := json.Marshal(values)
if err != nil {
return nil
}
return data
}
Loading
Loading