goclaw/internal/agent/usage_events_test.go

package agent

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"github.com/google/uuid"

	"github.com/nextlevelbuilder/goclaw/internal/providers"
	"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"
)

type fakeUsageEventStore struct {
	events chan store.UsageEvent
}

func newFakeUsageEventStore() *fakeUsageEventStore {
	return &fakeUsageEventStore{events: make(chan store.UsageEvent, 8)}
}

func (s *fakeUsageEventStore) InsertEvent(_ context.Context, event *store.UsageEvent) error {
	if event != nil {
		s.events <- *event
	}
	return nil
}

func (s *fakeUsageEventStore) InsertEvents(ctx context.Context, events []store.UsageEvent) error {
	for i := range events {
		if err := s.InsertEvent(ctx, &events[i]); err != nil {
			return err
		}
	}
	return nil
}

func (s *fakeUsageEventStore) GetEventTimeSeries(context.Context, store.UsageEventQuery) ([]store.UsageEventTimeSeries, error) {
	return nil, nil
}

func (s *fakeUsageEventStore) RefreshEventRollupHour(context.Context, time.Time) error {
	return nil
}

func (s *fakeUsageEventStore) GetLatestEventRollupBucket(context.Context) (*time.Time, error) {
	return nil, nil
}

func (s *fakeUsageEventStore) GetEventBreakdown(context.Context, store.UsageEventQuery) ([]store.UsageEventBreakdown, error) {
	return nil, nil
}

func (s *fakeUsageEventStore) GetEventSummary(context.Context, store.UsageEventQuery) (*store.UsageEventSummary, error) {
	return nil, nil
}

type usageFakeTool struct {
	name string
}

func (t usageFakeTool) Name() string               { return t.name }
func (t usageFakeTool) Description() string        { return "" }
func (t usageFakeTool) Parameters() map[string]any { return nil }
func (t usageFakeTool) Execute(context.Context, map[string]any) *tools.Result {
	return tools.NewResult("ok")
}

func TestRecordToolUsageEvent_RuntimeAliasUsesCanonicalName(t *testing.T) {
	storeSpy := newFakeUsageEventStore()
	registry := tools.NewRegistry()
	registry.Register(usageFakeTool{name: "exec"})
	registry.RegisterAlias("Bash", "exec")
	loop := &Loop{registry: registry, usageEvents: storeSpy, agentUUID: uuid.New(), tenantID: uuid.New()}
	ctx := tracing.WithTraceID(store.WithTenantID(t.Context(), loop.tenantID), uuid.New())

	canonical := loop.canonicalToolName("Bash")
	loop.recordToolUsageEvent(ctx, &RunRequest{RunID: "run-1", SessionKey: "session-1", Channel: "telegram"}, canonical, "Bash", "call-1",
		map[string]any{"command": "echo secret"}, time.Now().Add(-time.Millisecond), tools.NewResult("ok"), uuid.New())

	event := waitUsageEvent(t, storeSpy)
	if event.EventType != store.UsageEventTypeRuntimeToolCall {
		t.Fatalf("event type = %q, want %q", event.EventType, store.UsageEventTypeRuntimeToolCall)
	}
	if event.ResourceName != "exec" {
		t.Fatalf("resource = %q, want canonical exec", event.ResourceName)
	}
	if string(event.Metadata) == "" || json.Valid(event.Metadata) == false {
		t.Fatalf("metadata should contain valid alias metadata, got %q", string(event.Metadata))
	}
	if containsJSONKey(event.Metadata, "command") {
		t.Fatalf("metadata leaked raw command args: %s", string(event.Metadata))
	}
}

func TestRecordToolUsageEvent_UseSkillCountsSkillName(t *testing.T) {
	storeSpy := newFakeUsageEventStore()
	registry := tools.NewRegistry()
	registry.Register(tools.NewUseSkillTool())
	loop := &Loop{registry: registry, usageEvents: storeSpy, agentUUID: uuid.New(), tenantID: uuid.New()}
	ctx := tracing.WithTraceID(store.WithTenantID(t.Context(), loop.tenantID), uuid.New())

	loop.recordToolUsageEvent(ctx, &RunRequest{RunID: "run-1", SessionKey: "session-1"}, "use_skill", "use_skill", "call-1",
		map[string]any{"name": "ck:plan"}, time.Now().Add(-time.Millisecond), tools.NewResult("ok"), uuid.New())

	event := waitUsageEvent(t, storeSpy)
	if event.EventType != store.UsageEventTypeSkillActivation {
		t.Fatalf("event type = %q, want skill activation", event.EventType)
	}
	if event.ResourceType != store.UsageResourceTypeSkill || event.ResourceName != "ck:plan" {
		t.Fatalf("resource = %s/%s, want skill ck:plan", event.ResourceType, event.ResourceName)
	}
	if event.Source != store.UsageSourceUseSkill {
		t.Fatalf("source = %q, want use_skill", event.Source)
	}
}

func TestRecordToolUsageEvent_PreservesProviderCacheUsage(t *testing.T) {
	storeSpy := newFakeUsageEventStore()
	loop := &Loop{usageEvents: storeSpy, agentUUID: uuid.New(), tenantID: uuid.New()}
	ctx := tracing.WithTraceID(store.WithTenantID(t.Context(), loop.tenantID), uuid.New())
	result := tools.NewResult("ok")
	result.Provider = "openai-codex"
	result.Model = "gpt-5.5"
	result.Usage = &providers.Usage{
		PromptTokens:        100,
		CompletionTokens:    20,
		TotalTokens:         120,
		CacheReadTokens:     80,
		CacheCreationTokens: 10,
		ThinkingTokens:      5,
	}

	loop.recordToolUsageEvent(ctx, &RunRequest{RunID: "run-1", SessionKey: "session-1"}, "read_image", "read_image", "call-1",
		nil, time.Now().Add(-time.Millisecond), result, uuid.New())

	event := waitUsageEvent(t, storeSpy)
	if event.InputTokens != 100 || event.OutputTokens != 20 || event.TotalTokens != 120 {
		t.Fatalf("tokens = input %d output %d total %d", event.InputTokens, event.OutputTokens, event.TotalTokens)
	}
	if event.CacheReadTokens != 80 || event.CacheCreateTokens != 10 || event.ThinkingTokens != 5 {
		t.Fatalf("cache/thinking tokens = read %d create %d thinking %d", event.CacheReadTokens, event.CacheCreateTokens, event.ThinkingTokens)
	}
}

func TestRecordToolUsageEvent_ComputesCost(t *testing.T) {
	tenantID := uuid.New()
	providerID := uuid.New()
	storeSpy := newFakeUsageEventStore()
	input := "0.000001"
	output := "0.000002"
	usageStore := &tracePricingUsageStore{resolved: &store.ResolvedUsagePricing{
		ModelID: "openai/gpt-4o-mini",
		Source:  "catalog",
		Pricing: store.UsagePricingFields{Input: &input, Output: &output},
	}}
	providerStore := &tracePricingProviderStore{provider: &store.LLMProviderData{
		BaseModel:    store.BaseModel{ID: providerID},
		TenantID:     tenantID,
		Name:         "openai",
		ProviderType: store.ProviderOpenAICompat,
		APIKey:       "sk-test",
		Enabled:      true,
	}}
	loop := &Loop{
		usageEvents: storeSpy,
		agentUUID:   uuid.New(),
		tenantID:    tenantID,
		usageCaps:   usagecaps.NewService(usageStore, providerStore),
	}
	ctx := tracing.WithTraceID(store.WithTenantID(t.Context(), tenantID), uuid.New())
	result := tools.NewResult("ok")
	result.Provider = "openai"
	result.Model = "gpt-4o-mini"
	result.Usage = &providers.Usage{
		PromptTokens:     1000,
		CompletionTokens: 500,
		TotalTokens:      1500,
	}

	loop.recordToolUsageEvent(ctx, &RunRequest{RunID: "run-1", SessionKey: "session-1"}, "read_image", "read_image", "call-1",
		nil, time.Now().Add(-time.Millisecond), result, uuid.New())

	event := waitUsageEvent(t, storeSpy)
	if !floatClose(event.CostUSD, 0.002) {
		t.Fatalf("cost = %.9f, want 0.002", event.CostUSD)
	}
}

func TestRecordSkillSlashUsageEvent_RequiresTraceContext(t *testing.T) {
	storeSpy := newFakeUsageEventStore()
	loop := &Loop{usageEvents: storeSpy, agentUUID: uuid.New(), tenantID: uuid.New()}
	ctx := store.WithRunContext(store.WithTenantID(t.Context(), loop.tenantID), &store.RunContext{
		RunID:      "run-1",
		SessionKey: "session-1",
		Channel:    "web",
	})

	loop.recordSkillSlashUsageEvent(ctx, "ck:plan")
	select {
	case event := <-storeSpy.events:
		t.Fatalf("unexpected event without trace: %+v", event)
	case <-time.After(50 * time.Millisecond):
	}

	loop.recordSkillSlashUsageEvent(tracing.WithTraceID(ctx, uuid.New()), "ck:plan")
	event := waitUsageEvent(t, storeSpy)
	if event.Source != store.UsageSourceSlashCommand || event.ResourceName != "ck:plan" {
		t.Fatalf("slash event = %s/%s, want slash-command ck:plan", event.Source, event.ResourceName)
	}
	if event.RunID != "run-1" || event.SessionKey != "session-1" || event.Channel != "web" {
		t.Fatalf("run context not copied: run=%q session=%q channel=%q", event.RunID, event.SessionKey, event.Channel)
	}
}

func waitUsageEvent(t *testing.T, storeSpy *fakeUsageEventStore) store.UsageEvent {
	t.Helper()
	select {
	case event := <-storeSpy.events:
		return event
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for usage event")
		return store.UsageEvent{}
	}
}

func containsJSONKey(raw json.RawMessage, key string) bool {
	var values map[string]any
	if err := json.Unmarshal(raw, &values); err != nil {
		return false
	}
	_, ok := values[key]
	return ok
}