diff --git a/docs/project-changelog.md b/docs/project-changelog.md index b965c81a..84b9c87c 100644 --- a/docs/project-changelog.md +++ b/docs/project-changelog.md @@ -18,6 +18,8 @@ Significant changes, features, and fixes in reverse chronological order. - 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** diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index 6befa4d1..9d19e0b6 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -295,6 +295,11 @@ 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...) + recordUsageCapAttempt := func(reservation *usagecaps.Reservation) { + if reservation != nil { + opts = append(opts, withUsageCapMetadata(reservation.TraceMetadata())) + } + } emitChunk := func(chunk providers.StreamChunk) { if chunk.Thinking != "" { @@ -320,6 +325,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c 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) { @@ -329,6 +335,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } else { reservation.Reconcile(callCtx, callResp, callErr) } + recordUsageCapAttempt(reservation) } }, nil } @@ -339,6 +346,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } reservation, reserveErr := l.reserveLLMUsage(ctx, req, state, request, attempt) if reserveErr != nil { + recordUsageCapAttempt(reservation) return nil, reserveErr } var callResp *providers.ChatResponse @@ -353,6 +361,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c }) if reservation != nil { reservation.ReconcileStream(ctx, callResp, callErr, streamed) + recordUsageCapAttempt(reservation) } return callResp, callErr } else { @@ -360,6 +369,7 @@ func (l *Loop) makeCallLLM(req *RunRequest, emitRun func(AgentEvent)) func(ctx c } if reservation != nil { reservation.Reconcile(ctx, callResp, callErr) + recordUsageCapAttempt(reservation) } return callResp, callErr } 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/http/usage_caps.go b/internal/http/usage_caps.go index 1fc64c8f..cf1940ca 100644 --- a/internal/http/usage_caps.go +++ b/internal/http/usage_caps.go @@ -5,9 +5,11 @@ import ( "encoding/json" "errors" "fmt" + "io" "log/slog" "net/http" "strconv" + "strings" "time" "github.com/google/uuid" @@ -97,12 +99,12 @@ func (h *UsageCapsHandler) handleUpdatePolicy(w http.ResponseWriter, r *http.Req writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid policy id"}) return } - var body policyBody - if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + bodyBytes, err := io.ReadAll(r.Body) + if err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"}) return } - patch, err := body.toPatch() + patch, err := policyPatchFromBody(bodyBytes) if err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) return @@ -325,6 +327,33 @@ func (b policyBody) toPatch() (store.UsageCapPolicyPatch, error) { 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 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/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 index 12db8dd9..d3fbbaf3 100644 --- a/internal/usage/caps/service.go +++ b/internal/usage/caps/service.go @@ -50,17 +50,25 @@ type Reservation struct { 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 &Reservation{skipped: true, reason: "service_disabled"}, nil + return skippedReservation(req, "service_disabled"), nil } providerData, err := s.resolveProvider(ctx, req.TenantID, req.ProviderName) if err != nil { - return &Reservation{skipped: true, reason: "provider_metadata_missing"}, nil + return skippedReservation(req, "provider_metadata_missing"), nil } scope := store.UsageCapScope{ TenantID: req.TenantID, AgentID: req.AgentID, ProviderID: providerData.ID, @@ -71,14 +79,14 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err TenantID: req.TenantID, Decision: store.UsageCapEventSkip, Reason: "provider_not_billable_api", Metadata: mustJSON(scope), }) - return &Reservation{skipped: true, reason: "provider_not_billable_api"}, nil + 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 &Reservation{skipped: true, reason: "no_policy"}, nil + return skippedScopedReservation(req, scope, "no_policy"), nil } usage := pricing.BillableUsage{ InputTokens: int64(EstimateInputTokens(req.Messages)), @@ -88,6 +96,10 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err 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 @@ -95,7 +107,7 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err 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) { - return nil, fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID) + return blockedReservation(req, scope, key, usage, 0, uuid.Nil, "pricing_unknown"), fmt.Errorf("%w: %s", ErrPricingUnknown, req.ModelID) } return nil, err } @@ -109,10 +121,6 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err } metadata = map[string]any{"source": resolved.Source, "model_id": resolved.ModelID} } - key := req.ReservationKey - if key == "" { - key = uuid.NewString() - } result, err := s.store.ReserveUsage(ctx, store.UsageReserveRequest{ UsageCapScope: scope, ReservationKey: key, EstimatedTokens: usage.TotalTokens(), EstimatedCostMicros: costMicros, @@ -136,11 +144,16 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err 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 nil, ErrCapExceeded + 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}, nil + 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) { @@ -189,6 +202,9 @@ func (r *Reservation) reconcile(ctx context.Context, resp *providers.ChatRespons 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{ diff --git a/internal/usage/caps/service_test.go b/internal/usage/caps/service_test.go index e1350d6c..7a14a084 100644 --- a/internal/usage/caps/service_test.go +++ b/internal/usage/caps/service_test.go @@ -3,6 +3,8 @@ package caps import ( "context" "database/sql" + "encoding/json" + "errors" "testing" "github.com/google/uuid" @@ -61,6 +63,13 @@ func TestPreflightTokenOnlyCapDoesNotRequirePricing(t *testing.T) { 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) { @@ -224,6 +233,76 @@ func TestReservationReconcileIgnoresUnpricedRequestCount(t *testing.T) { 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 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) { @@ -245,6 +324,7 @@ type fakeUsageCapStore struct { resolved *store.ResolvedUsagePricing resolveErr error resolveCalls int + reserveErr error reserved store.UsageReserveRequest reconciled store.UsageReconcileRequest reconcileCalls int @@ -287,6 +367,9 @@ func (s *fakeUsageCapStore) DeleteUsageCapPolicy(context.Context, uuid.UUID, uui } 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 { diff --git a/ui/web/src/i18n/locales/en/usage.json b/ui/web/src/i18n/locales/en/usage.json index 70a45e09..477b8608 100644 --- a/ui/web/src/i18n/locales/en/usage.json +++ b/ui/web/src/i18n/locales/en/usage.json @@ -40,6 +40,9 @@ "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", @@ -56,7 +59,9 @@ "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" } diff --git a/ui/web/src/i18n/locales/vi/usage.json b/ui/web/src/i18n/locales/vi/usage.json index 28b72586..dd2a9f45 100644 --- a/ui/web/src/i18n/locales/vi/usage.json +++ b/ui/web/src/i18n/locales/vi/usage.json @@ -40,6 +40,9 @@ "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í", @@ -56,7 +59,9 @@ "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" } diff --git a/ui/web/src/i18n/locales/zh/usage.json b/ui/web/src/i18n/locales/zh/usage.json index 58d6b2f4..b200c31f 100644 --- a/ui/web/src/i18n/locales/zh/usage.json +++ b/ui/web/src/i18n/locales/zh/usage.json @@ -40,6 +40,9 @@ "maxTokens": "最大令牌", "maxCost": "最大费用 USD", "create": "创建上限", + "edit": "编辑上限", + "save": "保存上限", + "cancel": "取消编辑", "scope": "范围", "tokens": "令牌", "cost": "费用", @@ -56,7 +59,9 @@ "blocked": "已拦截", "toast": { "created": "用量上限已创建", + "updated": "用量上限已更新", "createFailed": "无法创建用量上限", + "updateFailed": "无法更新用量上限", "deleted": "用量上限已删除", "deleteFailed": "无法删除用量上限" } diff --git a/ui/web/src/pages/usage/components/usage-cap-row.tsx b/ui/web/src/pages/usage/components/usage-cap-row.tsx new file mode 100644 index 00000000..99786c2d --- /dev/null +++ b/ui/web/src/pages/usage/components/usage-cap-row.tsx @@ -0,0 +1,44 @@ +import { useTranslation } from "react-i18next"; +import { Pencil, Trash2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { formatCost, formatTokens } from "@/lib/format"; +import type { UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps"; + +interface UsageCapRowProps { + row: UsageCapUtilization; + onEdit: (policy: UsageCapPolicy) => void; + onDelete: () => void; +} + +export function UsageCapRow({ row, onEdit, onDelete }: UsageCapRowProps) { + const { t } = useTranslation("usage"); + const p = row.policy; + const tokenUsed = row.used_tokens + row.reserved_tokens; + const costUsed = row.used_cost_micros + row.reserved_cost_micros; + const tokenPct = p.max_tokens ? Math.min(100, Math.round((tokenUsed / p.max_tokens) * 100)) : 0; + const costPct = p.max_cost_micros ? Math.min(100, Math.round((costUsed / p.max_cost_micros) * 100)) : 0; + const isAgentBudget = p.source === "agent_budget_monthly_cents"; + + return ( + + +
{p.model_id || p.provider_type || t("caps.tenantScope")}
+
+ {p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")} + {isAgentBudget ? {t("caps.agentBudgetSource")} : null} +
+ + {t(`caps.windows.${p.window}`)} + {p.max_tokens ? `${formatTokens(tokenUsed)} / ${formatTokens(p.max_tokens)} (${tokenPct}%)` : "-"} + {p.max_cost_micros ? `${formatCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "-"} + {p.enabled ? t("caps.enabled") : t("caps.disabled")} + +
+ + +
+ + + ); +} diff --git a/ui/web/src/pages/usage/components/usage-caps-panel.tsx b/ui/web/src/pages/usage/components/usage-caps-panel.tsx index 2371cf4b..972f1c34 100644 --- a/ui/web/src/pages/usage/components/usage-caps-panel.tsx +++ b/ui/web/src/pages/usage/components/usage-caps-panel.tsx @@ -1,17 +1,18 @@ import { useMemo, useState } from "react"; import type { ReactNode } from "react"; import { useTranslation } from "react-i18next"; -import { Plus, RefreshCw, ShieldAlert, Trash2 } from "lucide-react"; +import { Plus, RefreshCw, Save, ShieldAlert, X } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { Badge } from "@/components/ui/badge"; +import { Switch } from "@/components/ui/switch"; import { formatCost, formatDate, formatTokens } from "@/lib/format"; import { useAgents } from "@/pages/agents/hooks/use-agents"; import { useProviders } from "@/pages/providers/hooks/use-providers"; import { useUsageCaps } from "../hooks/use-usage-caps"; -import type { UsageCapPolicy, UsageCapUtilization } from "@/types/usage-caps"; +import type { UsageCapPolicy } from "@/types/usage-caps"; +import { UsageCapRow } from "./usage-cap-row"; const ALL = "__all__"; @@ -19,37 +20,68 @@ export function UsageCapsPanel() { const { t } = useTranslation("usage"); const { agents } = useAgents(); const { providers } = useProviders(); - const { utilization, events, refreshing, refresh, createPolicy, deletePolicy } = useUsageCaps(); + const { utilization, events, refreshing, refresh, createPolicy, updatePolicy, deletePolicy } = useUsageCaps(); const [windowValue, setWindowValue] = useState("day"); const [agentId, setAgentId] = useState(ALL); const [providerId, setProviderId] = useState(ALL); const [modelId, setModelId] = useState(""); const [maxTokens, setMaxTokens] = useState(""); const [maxCost, setMaxCost] = useState(""); + const [enabled, setEnabled] = useState(true); + const [editingPolicy, setEditingPolicy] = useState(null); const [saving, setSaving] = useState(false); const provider = useMemo(() => providers.find((p) => p.id === providerId), [providerId, providers]); + const providerType = provider?.provider_type ?? (editingPolicy && providerId === (editingPolicy.provider_id ?? ALL) ? editingPolicy.provider_type : undefined); const blockedEvents = events.filter((event) => event.decision === "block"); + const isEditing = editingPolicy != null; + + const resetForm = () => { + setEditingPolicy(null); + setWindowValue("day"); + setAgentId(ALL); + setProviderId(ALL); + setModelId(""); + setMaxTokens(""); + setMaxCost(""); + setEnabled(true); + }; + + const startEdit = (policy: UsageCapPolicy) => { + setEditingPolicy(policy); + setWindowValue(policy.window); + setAgentId(policy.agent_id ?? ALL); + setProviderId(policy.provider_id ?? ALL); + setModelId(policy.model_id ?? ""); + setMaxTokens(policy.max_tokens ? String(policy.max_tokens) : ""); + setMaxCost(policy.max_cost_micros ? String(policy.max_cost_micros / 1_000_000) : ""); + setEnabled(policy.enabled); + }; const onSubmit = async () => { - const tokens = Number(maxTokens); - const cost = Number(maxCost); - if ((!Number.isFinite(tokens) || tokens <= 0) && (!Number.isFinite(cost) || cost <= 0)) return; + const tokenValue = maxTokens.trim(); + const costValue = maxCost.trim(); + const tokens = Number(tokenValue); + const cost = Number(costValue); + if ((!tokenValue || !Number.isFinite(tokens) || tokens <= 0) && (!costValue || !Number.isFinite(cost) || cost <= 0)) return; setSaving(true); try { - await createPolicy({ + const input = { window: windowValue, - agent_id: agentId === ALL ? undefined : agentId, - provider_id: providerId === ALL ? undefined : providerId, - provider_type: provider?.provider_type, - model_id: modelId.trim() || undefined, - max_tokens: Number.isFinite(tokens) && tokens > 0 ? Math.floor(tokens) : undefined, - max_cost_usd: Number.isFinite(cost) && cost > 0 ? cost : undefined, - enabled: true, - }); - setMaxTokens(""); - setMaxCost(""); - setModelId(""); + agent_id: agentId === ALL ? (isEditing ? "" : undefined) : agentId, + provider_id: providerId === ALL ? (isEditing ? "" : undefined) : providerId, + provider_type: providerType ?? (isEditing ? "" : undefined), + model_id: modelId.trim() || (isEditing ? "" : undefined), + max_tokens: tokenValue && Number.isFinite(tokens) && tokens > 0 ? Math.floor(tokens) : (isEditing ? null : undefined), + max_cost_usd: costValue && Number.isFinite(cost) && cost > 0 ? cost : (isEditing ? null : undefined), + enabled, + }; + if (editingPolicy) { + await updatePolicy(editingPolicy.id, input); + } else { + await createPolicy(input); + } + resetForm(); } finally { setSaving(false); } @@ -68,7 +100,7 @@ export function UsageCapsPanel() { -
+
setMaxCost(e.target.value)} inputMode="decimal" placeholder="25" className="text-base md:text-sm" /> - + {isEditing ? ( + + ) : null} +
+ + +
+ + {enabled ? t("caps.enabled") : t("caps.disabled")}
@@ -127,7 +170,7 @@ export function UsageCapsPanel() { {utilization.length === 0 ? ( {t("caps.empty")} ) : utilization.map((row) => ( - void deletePolicy(row.policy.id)} /> + void deletePolicy(row.policy.id)} /> ))} @@ -153,31 +196,3 @@ export function UsageCapsPanel() { function Field({ label, children }: { label: string; children: ReactNode }) { return
{children}
; } - -function UsageCapRow({ row, onDelete }: { row: UsageCapUtilization; onDelete: () => void }) { - const { t } = useTranslation("usage"); - const p = row.policy; - const tokenUsed = row.used_tokens + row.reserved_tokens; - const costUsed = row.used_cost_micros + row.reserved_cost_micros; - const tokenPct = p.max_tokens ? Math.min(100, Math.round((tokenUsed / p.max_tokens) * 100)) : 0; - const costPct = p.max_cost_micros ? Math.min(100, Math.round((costUsed / p.max_cost_micros) * 100)) : 0; - const isAgentBudget = p.source === "agent_budget_monthly_cents"; - return ( - - -
{p.model_id || p.provider_type || t("caps.tenantScope")}
-
- {p.agent_id ? t("caps.agentScoped") : t("caps.tenantScoped")} - {isAgentBudget ? {t("caps.agentBudgetSource")} : null} -
- - {t(`caps.windows.${p.window}`)} - {p.max_tokens ? `${formatTokens(tokenUsed)} / ${formatTokens(p.max_tokens)} (${tokenPct}%)` : "—"} - {p.max_cost_micros ? `${formatCost(costUsed / 1_000_000)} / ${formatCost(p.max_cost_micros / 1_000_000)} (${costPct}%)` : "—"} - {p.enabled ? t("caps.enabled") : t("caps.disabled")} - - - - - ); -} diff --git a/ui/web/src/pages/usage/hooks/use-usage-caps.ts b/ui/web/src/pages/usage/hooks/use-usage-caps.ts index 64c0e4c4..9d307e3b 100644 --- a/ui/web/src/pages/usage/hooks/use-usage-caps.ts +++ b/ui/web/src/pages/usage/hooks/use-usage-caps.ts @@ -12,8 +12,8 @@ export interface UsageCapPolicyInput { provider_type?: string; model_id?: string; window: UsageCapPolicy["window"]; - max_tokens?: number; - max_cost_usd?: number; + max_tokens?: number | null; + max_cost_usd?: number | null; enabled?: boolean; } @@ -81,6 +81,20 @@ export function useUsageCaps() { [http, refresh], ); + const updatePolicy = useCallback( + async (id: string, input: UsageCapPolicyInput) => { + try { + await http.patch(`/v1/usage-caps/policies/${id}`, input); + await refresh(); + toast.success(i18next.t("usage:caps.toast.updated")); + } catch (err) { + toast.error(i18next.t("usage:caps.toast.updateFailed"), err instanceof Error ? err.message : ""); + throw err; + } + }, + [http, refresh], + ); + return { policies: policiesQuery.data ?? [], utilization: utilizationQuery.data ?? [], @@ -89,6 +103,7 @@ export function useUsageCaps() { refreshing: policiesQuery.isFetching || utilizationQuery.isFetching || eventsQuery.isFetching, refresh, createPolicy, + updatePolicy, deletePolicy, }; }