From e9363b655eb071d255f183c72d0bfecde4354855 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sat, 25 Apr 2026 20:08:26 +0200 Subject: [PATCH 1/9] feat(budget): add user path budget management --- config/config.example.yaml | 12 + config/config.go | 192 ++++++++++ config/config_test.go | 84 +++++ .../admin/dashboard/static/css/dashboard.css | 58 +++ .../admin/dashboard/static/js/dashboard.js | 10 + .../dashboard/static/js/modules/budgets.js | 175 +++++++++ .../static/js/modules/budgets.test.js | 80 ++++ .../dashboard/static/js/modules/workflows.js | 1 + .../admin/dashboard/templates/layout.html | 1 + .../dashboard/templates/page-settings.html | 119 ++++++ internal/admin/handler.go | 79 ++++ internal/app/app.go | 66 +++- internal/budget/factory.go | 132 +++++++ internal/budget/service.go | 185 +++++++++ internal/budget/service_test.go | 138 +++++++ internal/budget/settings_helpers.go | 83 ++++ internal/budget/store.go | 78 ++++ internal/budget/store_mongodb.go | 245 ++++++++++++ internal/budget/store_postgresql.go | 290 ++++++++++++++ internal/budget/store_sqlite.go | 353 ++++++++++++++++++ internal/budget/store_sqlite_test.go | 66 ++++ internal/budget/types.go | 238 ++++++++++++ internal/budget/types_test.go | 44 +++ internal/server/budget_support.go | 45 +++ internal/server/handlers.go | 4 + internal/server/http.go | 7 + internal/server/native_batch_service.go | 4 + internal/server/passthrough_service.go | 4 + .../server/translated_inference_service.go | 13 + 29 files changed, 2791 insertions(+), 15 deletions(-) create mode 100644 internal/admin/dashboard/static/js/modules/budgets.js create mode 100644 internal/admin/dashboard/static/js/modules/budgets.test.js create mode 100644 internal/budget/factory.go create mode 100644 internal/budget/service.go create mode 100644 internal/budget/service_test.go create mode 100644 internal/budget/settings_helpers.go create mode 100644 internal/budget/store.go create mode 100644 internal/budget/store_mongodb.go create mode 100644 internal/budget/store_postgresql.go create mode 100644 internal/budget/store_sqlite.go create mode 100644 internal/budget/store_sqlite_test.go create mode 100644 internal/budget/types.go create mode 100644 internal/budget/types_test.go create mode 100644 internal/server/budget_support.go diff --git a/config/config.example.yaml b/config/config.example.yaml index e6e6a312..2b2070d5 100644 --- a/config/config.example.yaml +++ b/config/config.example.yaml @@ -86,6 +86,18 @@ usage: flush_interval: 5 retention_days: 90 +budgets: + enabled: true # env: BUDGETS_ENABLED; with no configured budgets this has no effect + user_paths: + # Env equivalent: + # SET_BUDGET_USER_PATH_EXAMPLE="daily=10,weekly=50" + - path: "/user/path/example" + limits: + - period: "daily" # hourly, daily, weekly, monthly; stored in DB as period_seconds + amount: 10.00 + - period: "weekly" + amount: 50.00 + metrics: enabled: false endpoint: "/metrics" diff --git a/config/config.go b/config/config.go index e762bf23..5579b0f2 100644 --- a/config/config.go +++ b/config/config.go @@ -36,6 +36,7 @@ type Config struct { Storage StorageConfig `yaml:"storage"` Logging LogConfig `yaml:"logging"` Usage UsageConfig `yaml:"usage"` + Budgets BudgetsConfig `yaml:"budgets"` Metrics MetricsConfig `yaml:"metrics"` HTTP HTTPConfig `yaml:"http"` Admin AdminConfig `yaml:"admin"` @@ -373,6 +374,36 @@ type UsageConfig struct { RetentionDays int `yaml:"retention_days" env:"USAGE_RETENTION_DAYS"` } +// BudgetsConfig holds per-user-path spend limits. +type BudgetsConfig struct { + // Enabled controls whether budget checks are active. + // Default: true. With no configured budgets, this has no effect. + Enabled bool `yaml:"enabled" env:"BUDGETS_ENABLED"` + + // UserPaths declares budget limits by tracked user path. + UserPaths []BudgetUserPathConfig `yaml:"user_paths"` +} + +// BudgetUserPathConfig declares one or more budget limits for a user path. +type BudgetUserPathConfig struct { + Path string `yaml:"path"` + Limits []BudgetLimitConfig `yaml:"limits"` +} + +// BudgetLimitConfig declares one spend limit for a reset period. +type BudgetLimitConfig struct { + // Period accepts hourly, daily, weekly, or monthly. The resolved period is + // persisted as PeriodSeconds in the database. + Period string `yaml:"period"` + + // PeriodSeconds can be set directly instead of Period. Standard values are + // 3600, 86400, 604800, and 2592000. + PeriodSeconds int64 `yaml:"period_seconds"` + + // Amount is the maximum allowed tracked provider spend for the period. + Amount float64 `yaml:"amount"` +} + // StorageConfig holds database storage configuration (used by audit logging, usage tracking, future IAM, etc.) type StorageConfig struct { // Type specifies the storage backend: "sqlite" (default), "postgresql", or "mongodb" @@ -968,6 +999,9 @@ func buildDefaultConfig() *Config { FlushInterval: 5, RetentionDays: 90, }, + Budgets: BudgetsConfig{ + Enabled: true, + }, Metrics: MetricsConfig{ Endpoint: "/metrics", }, @@ -1016,6 +1050,15 @@ func Load() (*LoadResult, error) { if err := applyEnvOverrides(cfg); err != nil { return nil, err } + if err := applyBudgetEnv(cfg); err != nil { + return nil, err + } + if err := validateBudgetConfig(&cfg.Budgets); err != nil { + return nil, err + } + if err := validateBudgetDependencies(cfg); err != nil { + return nil, err + } cfg.Models.ConfiguredProviderModelsMode = ResolveConfiguredProviderModelsMode(cfg.Models.ConfiguredProviderModelsMode) if !cfg.Models.ConfiguredProviderModelsMode.Valid() { return nil, fmt.Errorf("models.configured_provider_models_mode must be one of: fallback, allowlist") @@ -1216,6 +1259,155 @@ func loadFallbackConfig(cfg *FallbackConfig) error { return nil } +func applyBudgetEnv(cfg *Config) error { + if cfg == nil { + return nil + } + if !cfg.Budgets.Enabled { + return nil + } + + const prefix = "SET_BUDGET_" + for _, item := range os.Environ() { + key, value, ok := strings.Cut(item, "=") + if !ok || !strings.HasPrefix(key, prefix) || strings.TrimSpace(value) == "" { + continue + } + path := budgetEnvPath(key[len(prefix):]) + limits, err := parseBudgetEnvLimits(value) + if err != nil { + return fmt.Errorf("invalid value for %s: %w", key, err) + } + if len(limits) == 0 { + continue + } + cfg.Budgets.UserPaths = append(cfg.Budgets.UserPaths, BudgetUserPathConfig{ + Path: path, + Limits: limits, + }) + } + return nil +} + +func budgetEnvPath(suffix string) string { + suffix = strings.Trim(strings.ToLower(strings.TrimSpace(suffix)), "_") + if suffix == "" { + return "/" + } + segments := make([]string, 0) + for _, part := range strings.Split(suffix, "_") { + part = strings.TrimSpace(part) + if part != "" { + segments = append(segments, part) + } + } + if len(segments) == 0 { + return "/" + } + return "/" + strings.Join(segments, "/") +} + +func parseBudgetEnvLimits(raw string) ([]BudgetLimitConfig, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, nil + } + if strings.HasPrefix(raw, "{") { + values := map[string]float64{} + if err := json.Unmarshal([]byte(raw), &values); err != nil { + return nil, err + } + limits := make([]BudgetLimitConfig, 0, len(values)) + for period, amount := range values { + limits = append(limits, BudgetLimitConfig{Period: period, Amount: amount}) + } + return limits, nil + } + if strings.HasPrefix(raw, "[") { + var limits []BudgetLimitConfig + if err := json.Unmarshal([]byte(raw), &limits); err != nil { + return nil, err + } + return limits, nil + } + + fields := strings.FieldsFunc(raw, func(r rune) bool { + return r == ',' || r == ';' || r == '\n' + }) + limits := make([]BudgetLimitConfig, 0, len(fields)) + for _, field := range fields { + field = strings.TrimSpace(field) + if field == "" { + continue + } + period, amountText, ok := strings.Cut(field, "=") + if !ok { + period, amountText, ok = strings.Cut(field, ":") + } + if !ok { + return nil, fmt.Errorf("budget limit %q must use period=amount", field) + } + amount, err := strconv.ParseFloat(strings.TrimSpace(amountText), 64) + if err != nil { + return nil, fmt.Errorf("budget amount %q is not a valid number", amountText) + } + limits = append(limits, BudgetLimitConfig{ + Period: strings.TrimSpace(period), + Amount: amount, + }) + } + return limits, nil +} + +func validateBudgetConfig(cfg *BudgetsConfig) error { + if cfg == nil { + return nil + } + if !cfg.Enabled { + return nil + } + for pathIdx, entry := range cfg.UserPaths { + if strings.TrimSpace(entry.Path) == "" { + return fmt.Errorf("budgets.user_paths[%d].path is required", pathIdx) + } + for limitIdx, limit := range entry.Limits { + if limit.Amount <= 0 { + return fmt.Errorf("budgets.user_paths[%d].limits[%d].amount must be greater than 0", pathIdx, limitIdx) + } + if limit.PeriodSeconds <= 0 { + seconds, ok := budgetPeriodSeconds(limit.Period) + if !ok { + return fmt.Errorf("budgets.user_paths[%d].limits[%d].period must be one of hourly, daily, weekly, monthly or period_seconds must be set", pathIdx, limitIdx) + } + cfg.UserPaths[pathIdx].Limits[limitIdx].PeriodSeconds = seconds + } + } + } + return nil +} + +func validateBudgetDependencies(cfg *Config) error { + if cfg == nil || !cfg.Budgets.Enabled || len(cfg.Budgets.UserPaths) == 0 || cfg.Usage.Enabled { + return nil + } + return fmt.Errorf("budgets require usage tracking to be enabled because spend limits are evaluated from usage cost records") +} + +func budgetPeriodSeconds(period string) (int64, bool) { + switch strings.ToLower(strings.TrimSpace(period)) { + case "hour", "hourly", "hours": + return 3600, true + case "day", "daily", "days": + return 86400, true + case "week", "weekly", "weeks": + return 604800, true + case "month", "monthly", "months": + return 2592000, true + default: + return 0, false + } +} + // applyEnvOverrides walks cfg's struct fields and applies env var overrides // based on `env` struct tags. Maps are skipped. func applyEnvOverrides(cfg *Config) error { diff --git a/config/config_test.go b/config/config_test.go index a9682924..6d2cb287 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -55,6 +55,7 @@ func clearAllConfigEnvVars(t *testing.T) { "LOGGING_FLUSH_INTERVAL", "LOGGING_RETENTION_DAYS", "USAGE_ENABLED", "ENFORCE_RETURNING_USAGE_DATA", "USAGE_BUFFER_SIZE", "USAGE_FLUSH_INTERVAL", "USAGE_RETENTION_DAYS", + "BUDGETS_ENABLED", "GUARDRAILS_ENABLED", "ENABLE_GUARDRAILS_FOR_BATCH_PROCESSING", "FEATURE_FALLBACK_MODE", "FALLBACK_MANUAL_RULES_PATH", "MODEL_OVERRIDES_ENABLED", "MODELS_ENABLED_BY_DEFAULT", "KEEP_ONLY_ALIASES_AT_MODELS_ENDPOINT", "CONFIGURED_PROVIDER_MODELS_MODE", @@ -64,6 +65,13 @@ func clearAllConfigEnvVars(t *testing.T) { t.Setenv(key, "") os.Unsetenv(key) } + for _, item := range os.Environ() { + key, _, _ := strings.Cut(item, "=") + if strings.HasPrefix(key, "SET_BUDGET_") { + t.Setenv(key, "") + os.Unsetenv(key) + } + } clearProviderEnvVars(t) } @@ -157,6 +165,9 @@ func TestBuildDefaultConfig(t *testing.T) { if cfg.Usage.RetentionDays != 90 { t.Errorf("expected Usage.RetentionDays=90, got %d", cfg.Usage.RetentionDays) } + if !cfg.Budgets.Enabled { + t.Error("expected Budgets.Enabled=true") + } if cfg.Metrics.Endpoint != "/metrics" { t.Errorf("expected Metrics.Endpoint=/metrics, got %s", cfg.Metrics.Endpoint) } @@ -205,6 +216,79 @@ func TestBuildDefaultConfig(t *testing.T) { } } +func TestLoadBudgetEnvUserPath(t *testing.T) { + clearAllConfigEnvVars(t) + + withTempDir(t, func(string) { + t.Setenv("SET_BUDGET_USER_PATH_EXAMPLE", "daily=12.5,weekly=50") + + result, err := Load() + if err != nil { + t.Fatalf("Load() failed: %v", err) + } + + entries := result.Config.Budgets.UserPaths + if len(entries) != 1 { + t.Fatalf("expected 1 budget user path, got %d", len(entries)) + } + if got, want := entries[0].Path, "/user/path/example"; got != want { + t.Fatalf("budget env path = %q, want %q", got, want) + } + if len(entries[0].Limits) != 2 { + t.Fatalf("expected 2 budget limits, got %d", len(entries[0].Limits)) + } + if got, want := entries[0].Limits[0].PeriodSeconds, int64(86400); got != want { + t.Fatalf("daily period seconds = %d, want %d", got, want) + } + if got, want := entries[0].Limits[0].Amount, 12.5; got != want { + t.Fatalf("daily amount = %v, want %v", got, want) + } + if got, want := entries[0].Limits[1].PeriodSeconds, int64(604800); got != want { + t.Fatalf("weekly period seconds = %d, want %d", got, want) + } + if got, want := entries[0].Limits[1].Amount, 50.0; got != want { + t.Fatalf("weekly amount = %v, want %v", got, want) + } + }) +} + +func TestLoadBudgetEnvRequiresUsageTracking(t *testing.T) { + clearAllConfigEnvVars(t) + + withTempDir(t, func(string) { + t.Setenv("USAGE_ENABLED", "false") + t.Setenv("SET_BUDGET_USER_PATH_EXAMPLE", "daily=12.5") + + _, err := Load() + if err == nil { + t.Fatal("expected Load() to fail when budgets are configured while usage tracking is disabled") + } + if !strings.Contains(err.Error(), "budgets require usage tracking") { + t.Fatalf("Load() error = %v, want budget usage dependency error", err) + } + }) +} + +func TestLoadDisabledBudgetsIgnoreMalformedBudgetEnv(t *testing.T) { + clearAllConfigEnvVars(t) + + withTempDir(t, func(string) { + t.Setenv("BUDGETS_ENABLED", "false") + t.Setenv("SET_BUDGET_USER_PATH_EXAMPLE", "not-a-budget-limit") + + result, err := Load() + if err != nil { + t.Fatalf("Load() failed: %v", err) + } + if result.Config.Budgets.Enabled { + t.Fatal("expected budgets to be disabled") + } + if len(result.Config.Budgets.UserPaths) != 0 { + t.Fatalf("expected disabled budgets to ignore env user paths, got %d", len(result.Config.Budgets.UserPaths)) + } + }) +} + func TestLoad_ZeroConfig(t *testing.T) { clearAllConfigEnvVars(t) diff --git a/internal/admin/dashboard/static/css/dashboard.css b/internal/admin/dashboard/static/css/dashboard.css index ea27766f..2beb25ac 100644 --- a/internal/admin/dashboard/static/css/dashboard.css +++ b/internal/admin/dashboard/static/css/dashboard.css @@ -1875,6 +1875,23 @@ textarea:focus { border-color: var(--accent); } +.form-input { + width: 100%; + min-height: 38px; + padding: 8px 10px; + background: var(--bg-surface); + border: 1px solid var(--border); + border-radius: var(--radius); + color: var(--text); + font-size: 13px; + font-family: inherit; + outline: none; +} + +.form-input:focus { + border-color: var(--accent); +} + .alias-checkbox { display: inline-flex; align-items: center; @@ -3072,6 +3089,19 @@ body.conversation-drawer-open { border-color: color-mix(in srgb, var(--danger) 75%, var(--border)); } +.pagination-btn-danger { + color: #fff; + border-color: color-mix(in srgb, var(--danger) 76%, #000 12%); + background: var(--danger); + font-weight: 600; + box-shadow: 0 10px 22px color-mix(in srgb, var(--danger) 20%, transparent); +} + +.pagination-btn-danger:hover:not(:disabled) { + background: color-mix(in srgb, var(--danger) 90%, #fff 10%); + border-color: color-mix(in srgb, var(--danger) 82%, #000 14%); +} + .pagination-btn-with-icon { display: inline-flex; align-items: center; @@ -3384,6 +3414,25 @@ body.conversation-drawer-open { line-height: 1.55; } +.budget-settings-section { + width: 100%; +} + +.budget-settings-grid { + width: min(100%, 780px); + grid-template-columns: repeat(3, minmax(140px, 1fr)); +} + +.budget-settings-actions { + display: flex; + flex-wrap: wrap; + gap: 10px; +} + +.budget-reset-dialog { + max-width: 460px; +} + .settings-refresh-icon, .alias-create-icon, .form-action-icon { @@ -3831,6 +3880,15 @@ body.conversation-drawer-open { .settings-form-grid { grid-template-columns: 1fr; } + .budget-settings-grid { + grid-template-columns: 1fr; + } + .budget-settings-actions { + width: 100%; + } + .budget-settings-actions .pagination-btn { + width: 100%; + } /* Date picker mobile */ .date-picker-dropdown { diff --git a/internal/admin/dashboard/static/js/dashboard.js b/internal/admin/dashboard/static/js/dashboard.js index 563a8bda..2dbfe32c 100644 --- a/internal/admin/dashboard/static/js/dashboard.js +++ b/internal/admin/dashboard/static/js/dashboard.js @@ -196,6 +196,9 @@ function dashboard() { if (typeof this.ensureTimezoneOptions === "function") { this.ensureTimezoneOptions(); } + if (typeof this.fetchBudgetSettings === "function") { + this.fetchBudgetSettings(); + } } if (page === "overview") this.renderChart(); if (page === "usage") this.fetchUsagePage(); @@ -356,6 +359,7 @@ function dashboard() { (this.page === "workflows" && this.workflowFormOpen) || (this.page === "guardrails" && this.guardrailFormOpen) || (this.page === "auth-keys" && this.authKeyFormOpen) + || this.budgetResetDialogOpen ); }, @@ -922,6 +926,12 @@ function dashboard() { : null, "dashboardGuardrailsModule", ), + resolveModuleFactory( + typeof dashboardBudgetsModule === "function" + ? dashboardBudgetsModule + : null, + "dashboardBudgetsModule", + ), resolveModuleFactory( typeof dashboardWorkflowsModule === "function" ? dashboardWorkflowsModule diff --git a/internal/admin/dashboard/static/js/modules/budgets.js b/internal/admin/dashboard/static/js/modules/budgets.js new file mode 100644 index 00000000..1c3dcbea --- /dev/null +++ b/internal/admin/dashboard/static/js/modules/budgets.js @@ -0,0 +1,175 @@ +(function(global) { + function dashboardBudgetsModule() { + return { + budgetSettings: { + daily_reset_hour: 0, + daily_reset_minute: 0, + weekly_reset_weekday: 1, + weekly_reset_hour: 0, + weekly_reset_minute: 0, + monthly_reset_day: 1, + monthly_reset_hour: 0, + monthly_reset_minute: 0 + }, + budgetSettingsLoading: false, + budgetSettingsSaving: false, + budgetSettingsNotice: '', + budgetSettingsError: '', + budgetResetDialogOpen: false, + budgetResetConfirmation: '', + budgetResetLoading: false, + + budgetManagementEnabled() { + return typeof this.workflowRuntimeBooleanFlag === 'function' + ? this.workflowRuntimeBooleanFlag('BUDGETS_ENABLED', true) + : true; + }, + + budgetWeekdays() { + return [ + { value: 0, label: 'Sunday' }, + { value: 1, label: 'Monday' }, + { value: 2, label: 'Tuesday' }, + { value: 3, label: 'Wednesday' }, + { value: 4, label: 'Thursday' }, + { value: 5, label: 'Friday' }, + { value: 6, label: 'Saturday' } + ]; + }, + + normalizeBudgetSettings(payload) { + const current = this.budgetSettings || {}; + const numberValue = (key, fallback) => { + const parsed = Number(payload && payload[key]); + return Number.isFinite(parsed) ? parsed : fallback; + }; + return { + daily_reset_hour: numberValue('daily_reset_hour', Number(current.daily_reset_hour || 0)), + daily_reset_minute: numberValue('daily_reset_minute', Number(current.daily_reset_minute || 0)), + weekly_reset_weekday: numberValue('weekly_reset_weekday', Number(current.weekly_reset_weekday || 1)), + weekly_reset_hour: numberValue('weekly_reset_hour', Number(current.weekly_reset_hour || 0)), + weekly_reset_minute: numberValue('weekly_reset_minute', Number(current.weekly_reset_minute || 0)), + monthly_reset_day: numberValue('monthly_reset_day', Number(current.monthly_reset_day || 1)), + monthly_reset_hour: numberValue('monthly_reset_hour', Number(current.monthly_reset_hour || 0)), + monthly_reset_minute: numberValue('monthly_reset_minute', Number(current.monthly_reset_minute || 0)) + }; + }, + + budgetSettingsPayload() { + return this.normalizeBudgetSettings(this.budgetSettings); + }, + + async fetchBudgetSettings() { + if (!this.budgetManagementEnabled()) { + this.budgetSettingsError = ''; + return; + } + this.budgetSettingsLoading = true; + this.budgetSettingsError = ''; + try { + const request = this.requestOptions(); + const res = await fetch('/admin/api/v1/budgets/settings', request); + const handled = this.handleFetchResponse(res, 'budget settings', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + if (!handled) { + this.budgetSettingsError = 'Unable to load budget settings.'; + return; + } + this.budgetSettings = this.normalizeBudgetSettings(await res.json()); + } catch (e) { + console.error('Failed to fetch budget settings:', e); + this.budgetSettingsError = 'Unable to load budget settings.'; + } finally { + this.budgetSettingsLoading = false; + } + }, + + async saveBudgetSettings() { + if (this.budgetSettingsSaving) { + return; + } + this.budgetSettingsSaving = true; + this.budgetSettingsNotice = ''; + this.budgetSettingsError = ''; + try { + const request = this.requestOptions({ + method: 'PUT', + body: JSON.stringify(this.budgetSettingsPayload()) + }); + const res = await fetch('/admin/api/v1/budgets/settings', request); + const handled = this.handleFetchResponse(res, 'budget settings', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + if (!handled) { + this.budgetSettingsError = 'Unable to save budget settings.'; + return; + } + this.budgetSettings = this.normalizeBudgetSettings(await res.json()); + this.budgetSettingsNotice = 'Budget settings saved.'; + } catch (e) { + console.error('Failed to save budget settings:', e); + this.budgetSettingsError = 'Unable to save budget settings.'; + } finally { + this.budgetSettingsSaving = false; + } + }, + + openBudgetResetDialog() { + this.budgetResetConfirmation = ''; + this.budgetSettingsError = ''; + this.budgetResetDialogOpen = true; + setTimeout(() => { + const input = document.getElementById('budget-reset-confirmation'); + if (input && typeof input.focus === 'function') { + input.focus(); + } + }, 0); + }, + + closeBudgetResetDialog() { + this.budgetResetDialogOpen = false; + this.budgetResetConfirmation = ''; + }, + + async resetBudgets() { + if (this.budgetResetLoading) { + return; + } + if (String(this.budgetResetConfirmation || '').trim().toLowerCase() !== 'reset') { + this.budgetSettingsError = 'Type reset to confirm.'; + return; + } + this.budgetResetLoading = true; + this.budgetSettingsNotice = ''; + this.budgetSettingsError = ''; + try { + const request = this.requestOptions({ + method: 'POST', + body: JSON.stringify({ confirmation: 'reset' }) + }); + const res = await fetch('/admin/api/v1/budgets/reset', request); + const handled = this.handleFetchResponse(res, 'budget reset', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + if (!handled) { + this.budgetSettingsError = 'Unable to reset budgets.'; + return; + } + this.closeBudgetResetDialog(); + this.budgetSettingsNotice = 'Budgets reset.'; + } catch (e) { + console.error('Failed to reset budgets:', e); + this.budgetSettingsError = 'Unable to reset budgets.'; + } finally { + this.budgetResetLoading = false; + } + } + }; + } + + global.dashboardBudgetsModule = dashboardBudgetsModule; +})(typeof window !== 'undefined' ? window : globalThis); diff --git a/internal/admin/dashboard/static/js/modules/budgets.test.js b/internal/admin/dashboard/static/js/modules/budgets.test.js new file mode 100644 index 00000000..69df1f44 --- /dev/null +++ b/internal/admin/dashboard/static/js/modules/budgets.test.js @@ -0,0 +1,80 @@ +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const path = require('node:path'); +const vm = require('node:vm'); + +function loadBudgetsModuleFactory(overrides = {}) { + const source = fs.readFileSync(path.join(__dirname, 'budgets.js'), 'utf8'); + const window = { + ...(overrides.window || {}) + }; + const context = { + console, + setTimeout, + clearTimeout, + ...overrides, + window + }; + vm.createContext(context); + vm.runInContext(source, context); + return context.window.dashboardBudgetsModule; +} + +function createBudgetsModule(overrides) { + const factory = loadBudgetsModuleFactory(overrides); + return factory(); +} + +test('budgetManagementEnabled defaults on and respects the runtime flag', () => { + const module = createBudgetsModule(); + assert.equal(module.budgetManagementEnabled(), true); + + module.workflowRuntimeBooleanFlag = (key, fallback) => { + assert.equal(key, 'BUDGETS_ENABLED'); + assert.equal(fallback, true); + return false; + }; + assert.equal(module.budgetManagementEnabled(), false); +}); + +test('budgetSettingsPayload normalizes numeric values before saving', () => { + const module = createBudgetsModule(); + module.budgetSettings = { + daily_reset_hour: '7', + daily_reset_minute: '15', + weekly_reset_weekday: '3', + weekly_reset_hour: '8', + weekly_reset_minute: '45', + monthly_reset_day: '31', + monthly_reset_hour: '9', + monthly_reset_minute: '30' + }; + + assert.equal(JSON.stringify(module.budgetSettingsPayload()), JSON.stringify({ + daily_reset_hour: 7, + daily_reset_minute: 15, + weekly_reset_weekday: 3, + weekly_reset_hour: 8, + weekly_reset_minute: 45, + monthly_reset_day: 31, + monthly_reset_hour: 9, + monthly_reset_minute: 30 + })); +}); + +test('resetBudgets requires the typed reset confirmation before posting', async () => { + const module = createBudgetsModule({ + fetch() { + throw new Error('fetch should not be called'); + } + }); + module.requestOptions = (options) => options || {}; + module.handleFetchResponse = () => true; + module.budgetResetConfirmation = 'nope'; + + await module.resetBudgets(); + + assert.equal(module.budgetSettingsError, 'Type reset to confirm.'); + assert.equal(module.budgetResetLoading, false); +}); diff --git a/internal/admin/dashboard/static/js/modules/workflows.js b/internal/admin/dashboard/static/js/modules/workflows.js index 49e61998..f330cbc2 100644 --- a/internal/admin/dashboard/static/js/modules/workflows.js +++ b/internal/admin/dashboard/static/js/modules/workflows.js @@ -82,6 +82,7 @@ 'FEATURE_FALLBACK_MODE', 'LOGGING_ENABLED', 'USAGE_ENABLED', + 'BUDGETS_ENABLED', 'GUARDRAILS_ENABLED', 'CACHE_ENABLED', 'REDIS_URL', diff --git a/internal/admin/dashboard/templates/layout.html b/internal/admin/dashboard/templates/layout.html index 499826f2..1f8acc0e 100644 --- a/internal/admin/dashboard/templates/layout.html +++ b/internal/admin/dashboard/templates/layout.html @@ -81,6 +81,7 @@

+ diff --git a/internal/admin/dashboard/templates/page-settings.html b/internal/admin/dashboard/templates/page-settings.html index b4a96b32..db7de69e 100644 --- a/internal/admin/dashboard/templates/page-settings.html +++ b/internal/admin/dashboard/templates/page-settings.html @@ -36,6 +36,82 @@

Timezone

+
+
+
+

Budget Resets

+ {{template "inline-help-toggle" .}} +
+

+
+
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+
+ + +
+
+ +
+
@@ -80,6 +156,49 @@

Runtime Refresh

+ +
+ +
{{end}} diff --git a/internal/admin/handler.go b/internal/admin/handler.go index dfa2fe23..8ae2130f 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -20,6 +20,7 @@ import ( "gomodel/internal/aliases" "gomodel/internal/auditlog" "gomodel/internal/authkeys" + "gomodel/internal/budget" "gomodel/internal/core" "gomodel/internal/guardrails" "gomodel/internal/modeloverrides" @@ -37,6 +38,7 @@ type Handler struct { aliases *aliases.Service modelOverrides *modeloverrides.Service workflows *workflows.Service + budgets *budget.Service guardrails guardrails.Catalog guardrailDefs *guardrails.Service runtimeConfig DashboardConfigResponse @@ -53,6 +55,7 @@ const ( DashboardConfigFeatureFallbackMode = "FEATURE_FALLBACK_MODE" DashboardConfigLoggingEnabled = "LOGGING_ENABLED" DashboardConfigUsageEnabled = "USAGE_ENABLED" + DashboardConfigBudgetsEnabled = "BUDGETS_ENABLED" DashboardConfigGuardrailsEnabled = "GUARDRAILS_ENABLED" DashboardConfigCacheEnabled = "CACHE_ENABLED" DashboardConfigRedisURL = "REDIS_URL" @@ -64,6 +67,7 @@ type DashboardConfigResponse struct { FeatureFallbackMode string `json:"FEATURE_FALLBACK_MODE,omitempty"` LoggingEnabled string `json:"LOGGING_ENABLED,omitempty"` UsageEnabled string `json:"USAGE_ENABLED,omitempty"` + BudgetsEnabled string `json:"BUDGETS_ENABLED,omitempty"` GuardrailsEnabled string `json:"GUARDRAILS_ENABLED,omitempty"` CacheEnabled string `json:"CACHE_ENABLED,omitempty"` RedisURL string `json:"REDIS_URL,omitempty"` @@ -173,6 +177,13 @@ func WithWorkflows(service *workflows.Service) Option { } } +// WithBudgets enables budget administration endpoints. +func WithBudgets(service *budget.Service) Option { + return func(h *Handler) { + h.budgets = service + } +} + // WithGuardrailsRegistry enables listing valid guardrail references for workflow authoring. func WithGuardrailsRegistry(registry guardrails.Catalog) Option { return func(h *Handler) { @@ -232,6 +243,7 @@ func normalizeDashboardRuntimeConfig(values DashboardConfigResponse) DashboardCo FeatureFallbackMode: strings.TrimSpace(values.FeatureFallbackMode), LoggingEnabled: strings.TrimSpace(values.LoggingEnabled), UsageEnabled: strings.TrimSpace(values.UsageEnabled), + BudgetsEnabled: strings.TrimSpace(values.BudgetsEnabled), GuardrailsEnabled: strings.TrimSpace(values.GuardrailsEnabled), CacheEnabled: strings.TrimSpace(values.CacheEnabled), RedisURL: strings.TrimSpace(values.RedisURL), @@ -985,6 +997,58 @@ func (h *Handler) DashboardConfig(c *echo.Context) error { return c.JSON(http.StatusOK, cloneDashboardRuntimeConfig(h.runtimeConfig)) } +// BudgetSettings handles GET /admin/api/v1/budgets/settings. +func (h *Handler) BudgetSettings(c *echo.Context) error { + if h.budgets == nil { + return handleError(c, featureUnavailableError("budgets feature is unavailable")) + } + return c.JSON(http.StatusOK, h.budgets.Settings()) +} + +// UpdateBudgetSettings handles PUT /admin/api/v1/budgets/settings. +func (h *Handler) UpdateBudgetSettings(c *echo.Context) error { + if h.budgets == nil { + return handleError(c, featureUnavailableError("budgets feature is unavailable")) + } + var req updateBudgetSettingsRequest + if err := json.NewDecoder(c.Request().Body).Decode(&req); err != nil { + return handleError(c, core.NewInvalidRequestError("invalid request body: "+err.Error(), err)) + } + settings := budget.Settings{ + DailyResetHour: req.DailyResetHour, + DailyResetMinute: req.DailyResetMinute, + WeeklyResetWeekday: req.WeeklyResetWeekday, + WeeklyResetHour: req.WeeklyResetHour, + WeeklyResetMinute: req.WeeklyResetMinute, + MonthlyResetDay: req.MonthlyResetDay, + MonthlyResetHour: req.MonthlyResetHour, + MonthlyResetMinute: req.MonthlyResetMinute, + } + saved, err := h.budgets.SaveSettings(c.Request().Context(), settings) + if err != nil { + return handleError(c, core.NewInvalidRequestError(err.Error(), err)) + } + return c.JSON(http.StatusOK, saved) +} + +// ResetBudgets handles POST /admin/api/v1/budgets/reset. +func (h *Handler) ResetBudgets(c *echo.Context) error { + if h.budgets == nil { + return handleError(c, featureUnavailableError("budgets feature is unavailable")) + } + var req resetBudgetsRequest + if err := json.NewDecoder(c.Request().Body).Decode(&req); err != nil { + return handleError(c, core.NewInvalidRequestError("invalid request body: "+err.Error(), err)) + } + if strings.TrimSpace(strings.ToLower(req.Confirmation)) != "reset" { + return handleError(c, core.NewInvalidRequestError("confirmation must be reset", nil)) + } + if err := h.budgets.ResetAll(c.Request().Context(), time.Now().UTC()); err != nil { + return handleError(c, core.NewProviderError("budgets", http.StatusServiceUnavailable, "failed to reset budgets", err)) + } + return c.JSON(http.StatusOK, map[string]any{"status": "ok"}) +} + // ProviderStatus handles GET /admin/api/v1/providers/status func (h *Handler) ProviderStatus(c *echo.Context) error { return c.JSON(http.StatusOK, h.buildProviderStatusResponse()) @@ -1177,6 +1241,21 @@ type createAuthKeyRequest struct { ExpiresAt *time.Time `json:"expires_at,omitempty"` } +type updateBudgetSettingsRequest struct { + DailyResetHour int `json:"daily_reset_hour"` + DailyResetMinute int `json:"daily_reset_minute"` + WeeklyResetWeekday int `json:"weekly_reset_weekday"` + WeeklyResetHour int `json:"weekly_reset_hour"` + WeeklyResetMinute int `json:"weekly_reset_minute"` + MonthlyResetDay int `json:"monthly_reset_day"` + MonthlyResetHour int `json:"monthly_reset_hour"` + MonthlyResetMinute int `json:"monthly_reset_minute"` +} + +type resetBudgetsRequest struct { + Confirmation string `json:"confirmation"` +} + func featureUnavailableError(message string) error { return core.NewInvalidRequestErrorWithStatus(http.StatusServiceUnavailable, message, nil). WithCode("feature_unavailable") diff --git a/internal/app/app.go b/internal/app/app.go index bd5a6ea8..4efeb94f 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -21,6 +21,7 @@ import ( "gomodel/internal/auditlog" "gomodel/internal/authkeys" "gomodel/internal/batch" + "gomodel/internal/budget" "gomodel/internal/core" "gomodel/internal/fallback" "gomodel/internal/guardrails" @@ -40,6 +41,7 @@ type App struct { providers *providers.InitResult audit *auditlog.Result usage *usage.Result + budgets *budget.Result batch *batch.Result aliases *aliases.Result modelOverrides *modeloverrides.Result @@ -135,6 +137,27 @@ func New(ctx context.Context, cfg Config) (*App, error) { } app.usage = usageResult + var budgetResult *budget.Result + if appCfg.Budgets.Enabled { + sharedBudgetStorage := firstSharedStorage(auditResult.Storage, usageResult.Storage) + if sharedBudgetStorage != nil { + budgetResult, err = budget.NewWithSharedStorage(ctx, appCfg, sharedBudgetStorage) + } else { + budgetResult, err = budget.New(ctx, appCfg) + } + if err != nil { + closeErr := errors.Join(app.usage.Close(), app.audit.Close(), app.providers.Close()) + if closeErr != nil { + return nil, fmt.Errorf("failed to initialize budgets: %w (also: close error: %v)", err, closeErr) + } + return nil, fmt.Errorf("failed to initialize budgets: %w", err) + } + } else { + budgetResult = &budget.Result{} + slog.Info("budgets disabled") + } + app.budgets = budgetResult + // Initialize batch lifecycle storage. var batchResult *batch.Result if auditResult.Storage != nil { @@ -145,7 +168,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { batchResult, err = batch.New(ctx, appCfg) } if err != nil { - closeErr := errors.Join(app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize batch storage: %w (also: close error: %v)", err, closeErr) } @@ -165,7 +188,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { aliasResult, err = aliases.New(ctx, appCfg, providerResult.Registry) } if err != nil { - closeErr := errors.Join(app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize aliases: %w (also: close error: %v)", err, closeErr) } @@ -182,7 +205,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { modelOverrideResult, err = modeloverrides.New(ctx, appCfg, providerResult.Registry) } if err != nil { - closeErr := errors.Join(app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize model overrides: %w (also: close error: %v)", err, closeErr) } @@ -209,7 +232,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { guardrailResult, err = guardrails.New(ctx, appCfg, refreshInterval, guardrailExecutor) } if err != nil { - closeErr := errors.Join(app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize guardrails: %w (also: close error: %v)", err, closeErr) } @@ -219,14 +242,14 @@ func New(ctx context.Context, cfg Config) (*App, error) { seedGuardrails, err := configGuardrailDefinitions(appCfg.Guardrails) if err != nil { - closeErr := errors.Join(app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to prepare guardrail definitions: %w (also: close error: %v)", err, closeErr) } return nil, fmt.Errorf("failed to prepare guardrail definitions: %w", err) } if err := guardrailResult.Service.UpsertDefinitions(ctx, seedGuardrails); err != nil { - closeErr := errors.Join(app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to upsert guardrails: %w (also: close error: %v)", err, closeErr) } @@ -249,7 +272,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { workflowResult, err = workflows.New(ctx, appCfg, workflowCompiler, refreshInterval) } if err != nil { - closeErr := errors.Join(app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize workflows: %w (also: close error: %v)", err, closeErr) } @@ -257,14 +280,14 @@ func New(ctx context.Context, cfg Config) (*App, error) { } defaultWorkflow := defaultWorkflowInput(appCfg, guardrailResult.Service.Names(), seedGuardrails) if err := workflowResult.Service.EnsureDefaultGlobal(ctx, defaultWorkflow); err != nil { - closeErr := errors.Join(workflowResult.Close(), app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(workflowResult.Close(), app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to seed workflows: %w (also: close error: %v)", err, closeErr) } return nil, fmt.Errorf("failed to seed workflows: %w", err) } if err := workflowResult.Service.Refresh(ctx); err != nil { - closeErr := errors.Join(workflowResult.Close(), app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(workflowResult.Close(), app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to load workflows: %w (also: close error: %v)", err, closeErr) } @@ -288,7 +311,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { authKeyResult, err = authkeys.New(ctx, appCfg) } if err != nil { - closeErr := errors.Join(workflowResult.Close(), app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(workflowResult.Close(), app.guardrails.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize auth keys: %w (also: close error: %v)", err, closeErr) } @@ -334,6 +357,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { PprofEnabled: appCfg.Server.PprofEnabled, AuditLogger: auditResult.Logger, UsageLogger: usageResult.Logger, + BudgetChecker: budgetResult.Service, PricingResolver: providerResult.Registry, ModelResolver: app.aliases.Service, ModelAuthorizer: app.modelOverrides.Service, @@ -370,6 +394,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { app.modelOverrides.Service, workflowResult.Service, app.guardrails.Service, + budgetResult.Service, app, dashboardRuntimeConfig(appCfg, usageEnabledForDashboard), adminCfg.UIEnabled, @@ -430,7 +455,7 @@ func New(ctx context.Context, cfg Config) (*App, error) { if app.batch != nil { batchCloseErr = app.batch.Close() } - closeErr := errors.Join(workflowsCloseErr, guardrailsCloseErr, authKeysCloseErr, aliasCloseErr, modelOverridesCloseErr, batchCloseErr, app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(workflowsCloseErr, guardrailsCloseErr, authKeysCloseErr, aliasCloseErr, modelOverridesCloseErr, batchCloseErr, app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to initialize response cache: %w (also: close error: %v)", err, closeErr) } @@ -449,14 +474,14 @@ func New(ctx context.Context, cfg Config) (*App, error) { ResponseCache: rcm, }) if err := guardrailResult.Service.SetExecutor(ctx, internalGuardrailExecutor); err != nil { - closeErr := errors.Join(rcm.Close(), app.workflows.Close(), app.guardrails.Close(), app.authKeys.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(rcm.Close(), app.workflows.Close(), app.guardrails.Close(), app.authKeys.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to wire internal guardrail executor: %w (also: close error: %v)", err, closeErr) } return nil, fmt.Errorf("failed to wire internal guardrail executor: %w", err) } if err := workflowResult.Service.Refresh(ctx); err != nil { - closeErr := errors.Join(rcm.Close(), app.workflows.Close(), app.guardrails.Close(), app.authKeys.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) + closeErr := errors.Join(rcm.Close(), app.workflows.Close(), app.guardrails.Close(), app.authKeys.Close(), app.modelOverrides.Close(), app.aliases.Close(), app.batch.Close(), app.budgets.Close(), app.usage.Close(), app.audit.Close(), app.providers.Close()) if closeErr != nil { return nil, fmt.Errorf("failed to refresh workflows after wiring internal guardrail executor: %w (also: close error: %v)", err, closeErr) } @@ -661,7 +686,15 @@ func (a *App) Shutdown(ctx context.Context) error { } } - // 9. Close usage tracking (flushes pending entries) + // 9. Close budget subsystem. + if a.budgets != nil { + if err := a.budgets.Close(); err != nil { + slog.Error("budgets close error", "error", err) + errs = append(errs, fmt.Errorf("budgets close: %w", err)) + } + } + + // 10. Close usage tracking (flushes pending entries) if a.usage != nil { if err := a.usage.Close(); err != nil { slog.Error("usage logger close error", "error", err) @@ -669,7 +702,7 @@ func (a *App) Shutdown(ctx context.Context) error { } } - // 10. Close audit logging (flushes pending logs) + // 11. Close audit logging (flushes pending logs) if a.audit != nil { if err := a.audit.Close(); err != nil { slog.Error("audit logger close error", "error", err) @@ -749,6 +782,7 @@ func initAdmin( modelOverrideService *modeloverrides.Service, workflowService *workflows.Service, guardrailService *guardrails.Service, + budgetService *budget.Service, runtimeRefresher admin.RuntimeRefresher, runtimeConfig admin.DashboardConfigResponse, uiEnabled bool, @@ -792,6 +826,7 @@ func initAdmin( admin.WithModelOverrides(modelOverrideService), admin.WithWorkflows(workflowService), admin.WithGuardrailService(guardrailService), + admin.WithBudgets(budgetService), admin.WithRuntimeRefresher(runtimeRefresher), admin.WithDashboardRuntimeConfig(runtimeConfig), ) @@ -917,6 +952,7 @@ func dashboardRuntimeConfig(cfg *config.Config, usageEnabled bool) admin.Dashboa FeatureFallbackMode: dashboardFallbackModeValue(cfg), LoggingEnabled: dashboardEnabledValue(cfg != nil && cfg.Logging.Enabled), UsageEnabled: dashboardEnabledValue(cfg != nil && cfg.Usage.Enabled), + BudgetsEnabled: dashboardEnabledValue(cfg != nil && cfg.Budgets.Enabled), GuardrailsEnabled: dashboardEnabledValue(cfg != nil && cfg.Guardrails.Enabled), CacheEnabled: dashboardEnabledValue(cacheAnalyticsConfigured(cfg, usageEnabled)), RedisURL: dashboardEnabledValue(simpleResponseCacheConfigured(cfg)), diff --git a/internal/budget/factory.go b/internal/budget/factory.go new file mode 100644 index 00000000..edb3bab5 --- /dev/null +++ b/internal/budget/factory.go @@ -0,0 +1,132 @@ +package budget + +import ( + "context" + "database/sql" + "errors" + "fmt" + "sync" + + "github.com/jackc/pgx/v5/pgxpool" + "go.mongodb.org/mongo-driver/v2/mongo" + + "gomodel/config" + "gomodel/internal/storage" +) + +type Result struct { + Service *Service + Store Store + Storage storage.Storage + + closeOnce sync.Once + closeErr error +} + +func (r *Result) Close() error { + if r == nil { + return nil + } + r.closeOnce.Do(func() { + var errs []error + if r.Store != nil { + if err := r.Store.Close(); err != nil { + errs = append(errs, fmt.Errorf("store close: %w", err)) + } + } + if r.Storage != nil { + if err := r.Storage.Close(); err != nil { + errs = append(errs, fmt.Errorf("storage close: %w", err)) + } + } + if len(errs) > 0 { + r.closeErr = fmt.Errorf("close errors: %w", errors.Join(errs...)) + } + }) + return r.closeErr +} + +func New(ctx context.Context, cfg *config.Config) (*Result, error) { + if cfg == nil { + return nil, fmt.Errorf("config is required") + } + if !cfg.Budgets.Enabled { + return &Result{}, nil + } + storeConn, err := storage.New(ctx, cfg.Storage.BackendConfig()) + if err != nil { + return nil, fmt.Errorf("failed to create storage: %w", err) + } + result, err := newResult(ctx, cfg, storeConn) + if err != nil { + _ = storeConn.Close() + return nil, err + } + result.Storage = storeConn + return result, nil +} + +func NewWithSharedStorage(ctx context.Context, cfg *config.Config, shared storage.Storage) (*Result, error) { + if cfg == nil { + return nil, fmt.Errorf("config is required") + } + if !cfg.Budgets.Enabled { + return &Result{}, nil + } + if shared == nil { + return nil, fmt.Errorf("shared storage is required") + } + return newResult(ctx, cfg, shared) +} + +func newResult(ctx context.Context, cfg *config.Config, storeConn storage.Storage) (*Result, error) { + store, err := createStore(ctx, storeConn) + if err != nil { + return nil, err + } + service, err := NewService(ctx, store) + if err != nil { + return nil, err + } + if err := seedConfiguredBudgets(ctx, service, cfg.Budgets); err != nil { + return nil, err + } + return &Result{Service: service, Store: store}, nil +} + +func createStore(ctx context.Context, store storage.Storage) (Store, error) { + return storage.ResolveBackend[Store]( + store, + func(db *sql.DB) (Store, error) { return NewSQLiteStore(db) }, + func(pool *pgxpool.Pool) (Store, error) { return NewPostgreSQLStore(ctx, pool) }, + func(db *mongo.Database) (Store, error) { return NewMongoDBStore(ctx, db) }, + ) +} + +func seedConfiguredBudgets(ctx context.Context, service *Service, cfg config.BudgetsConfig) error { + if service == nil || len(cfg.UserPaths) == 0 { + return nil + } + budgets := make([]Budget, 0) + for _, entry := range cfg.UserPaths { + userPath, err := NormalizeUserPath(entry.Path) + if err != nil { + return fmt.Errorf("invalid budget user path %q: %w", entry.Path, err) + } + for _, limit := range entry.Limits { + seconds := limit.PeriodSeconds + if seconds <= 0 { + if parsed, ok := PeriodSeconds(limit.Period); ok { + seconds = parsed + } + } + budgets = append(budgets, Budget{ + UserPath: userPath, + PeriodSeconds: seconds, + Amount: limit.Amount, + Source: "config", + }) + } + } + return service.ReplaceConfigBudgets(ctx, budgets) +} diff --git a/internal/budget/service.go b/internal/budget/service.go new file mode 100644 index 00000000..c690831b --- /dev/null +++ b/internal/budget/service.go @@ -0,0 +1,185 @@ +package budget + +import ( + "context" + "fmt" + "sort" + "strings" + "sync" + "time" +) + +type Service struct { + store Store + mu sync.RWMutex + + budgets []Budget + settings Settings +} + +func NewService(ctx context.Context, store Store) (*Service, error) { + if store == nil { + return nil, fmt.Errorf("budget store is required") + } + service := &Service{ + store: store, + settings: DefaultSettings(), + } + if err := service.Refresh(ctx); err != nil { + return nil, err + } + return service, nil +} + +func (s *Service) Refresh(ctx context.Context) error { + if s == nil || s.store == nil { + return nil + } + budgets, err := s.store.ListBudgets(ctx) + if err != nil { + return err + } + settings, err := s.store.GetSettings(ctx) + if err != nil { + return err + } + sort.SliceStable(budgets, func(i, j int) bool { + if budgets[i].UserPath == budgets[j].UserPath { + return budgets[i].PeriodSeconds < budgets[j].PeriodSeconds + } + return budgets[i].UserPath < budgets[j].UserPath + }) + s.mu.Lock() + s.budgets = budgets + s.settings = settings + s.mu.Unlock() + return nil +} + +func (s *Service) UpsertBudgets(ctx context.Context, budgets []Budget) error { + if s == nil || s.store == nil { + return nil + } + if err := s.store.UpsertBudgets(ctx, budgets); err != nil { + return err + } + return s.Refresh(ctx) +} + +func (s *Service) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error { + if s == nil || s.store == nil { + return nil + } + if err := s.store.ReplaceConfigBudgets(ctx, budgets); err != nil { + return err + } + return s.Refresh(ctx) +} + +func (s *Service) Budgets() []Budget { + if s == nil { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + return append([]Budget(nil), s.budgets...) +} + +func (s *Service) Settings() Settings { + if s == nil { + return DefaultSettings() + } + s.mu.RLock() + defer s.mu.RUnlock() + return s.settings +} + +func (s *Service) SaveSettings(ctx context.Context, settings Settings) (Settings, error) { + if s == nil || s.store == nil { + return Settings{}, fmt.Errorf("budget service is unavailable") + } + saved, err := s.store.SaveSettings(ctx, settings) + if err != nil { + return Settings{}, err + } + if err := s.Refresh(ctx); err != nil { + return Settings{}, err + } + return saved, nil +} + +func (s *Service) ResetAll(ctx context.Context, at time.Time) error { + if s == nil || s.store == nil { + return fmt.Errorf("budget service is unavailable") + } + if at.IsZero() { + at = time.Now().UTC() + } + if err := s.store.ResetAllBudgets(ctx, at.UTC()); err != nil { + return err + } + return s.Refresh(ctx) +} + +func (s *Service) Check(ctx context.Context, userPath string, now time.Time) error { + _, err := s.CheckWithResults(ctx, userPath, now) + return err +} + +func (s *Service) CheckWithResults(ctx context.Context, userPath string, now time.Time) ([]CheckResult, error) { + if s == nil || s.store == nil { + return nil, nil + } + userPath, err := NormalizeUserPath(userPath) + if err != nil { + return nil, err + } + if now.IsZero() { + now = time.Now().UTC() + } + now = now.UTC() + + s.mu.RLock() + budgets := append([]Budget(nil), s.budgets...) + settings := s.settings + s.mu.RUnlock() + if len(budgets) == 0 { + return nil, nil + } + + results := make([]CheckResult, 0) + for _, budget := range budgets { + if !budgetAppliesToPath(budget.UserPath, userPath) { + continue + } + start, end := PeriodBounds(now, budget.PeriodSeconds, settings) + if budget.LastResetAt != nil && budget.LastResetAt.After(start) { + start = budget.LastResetAt.UTC() + } + spent, _, err := s.store.SumUsageCost(ctx, budget.UserPath, start, now) + if err != nil { + return results, err + } + result := CheckResult{ + Budget: budget, + PeriodStart: start, + PeriodEnd: end, + Spent: spent, + Remaining: budget.Amount - spent, + } + results = append(results, result) + if spent >= budget.Amount { + return results, &ExceededError{Result: result} + } + } + return results, nil +} + +func budgetAppliesToPath(budgetPath, requestPath string) bool { + budgetPath = strings.TrimSpace(budgetPath) + requestPath = strings.TrimSpace(requestPath) + if budgetPath == "/" { + return true + } + return requestPath == budgetPath || strings.HasPrefix(requestPath, budgetPath+"/") +} diff --git a/internal/budget/service_test.go b/internal/budget/service_test.go new file mode 100644 index 00000000..665a4ff3 --- /dev/null +++ b/internal/budget/service_test.go @@ -0,0 +1,138 @@ +package budget + +import ( + "context" + "errors" + "testing" + "time" +) + +type fakeStore struct { + budgets []Budget + settings Settings + sum func(userPath string, start, end time.Time) (float64, bool, error) + + lastSumUserPath string + lastSumStart time.Time + lastResetAt time.Time +} + +func (s *fakeStore) ListBudgets(context.Context) ([]Budget, error) { + return append([]Budget(nil), s.budgets...), nil +} + +func (s *fakeStore) UpsertBudgets(context.Context, []Budget) error { + return nil +} + +func (s *fakeStore) ReplaceConfigBudgets(context.Context, []Budget) error { + return nil +} + +func (s *fakeStore) GetSettings(context.Context) (Settings, error) { + if s.settings == (Settings{}) { + return DefaultSettings(), nil + } + return s.settings, nil +} + +func (s *fakeStore) SaveSettings(context.Context, Settings) (Settings, error) { + return Settings{}, nil +} + +func (s *fakeStore) ResetAllBudgets(_ context.Context, at time.Time) error { + s.lastResetAt = at + return nil +} + +func (s *fakeStore) SumUsageCost(_ context.Context, userPath string, start, end time.Time) (float64, bool, error) { + s.lastSumUserPath = userPath + s.lastSumStart = start + if s.sum == nil { + return 0, false, nil + } + return s.sum(userPath, start, end) +} + +func (s *fakeStore) Close() error { + return nil +} + +func TestServiceCheckRejectsExceededBudgetForMatchingUserPath(t *testing.T) { + ctx := context.Background() + store := &fakeStore{ + budgets: []Budget{ + {UserPath: "/team", PeriodSeconds: PeriodDailySeconds, Amount: 10}, + }, + sum: func(userPath string, start, end time.Time) (float64, bool, error) { + if userPath != "/team" { + t.Fatalf("sum user path = %q, want /team", userPath) + } + return 10, true, nil + }, + } + service, err := NewService(ctx, store) + if err != nil { + t.Fatalf("NewService() failed: %v", err) + } + + err = service.Check(ctx, "/team/app", time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)) + var exceeded *ExceededError + if !errors.As(err, &exceeded) { + t.Fatalf("Check() error = %v, want ExceededError", err) + } + if got := exceeded.Result.Budget.UserPath; got != "/team" { + t.Fatalf("exceeded budget path = %q, want /team", got) + } +} + +func TestServiceCheckIgnoresSiblingUserPath(t *testing.T) { + ctx := context.Background() + called := false + store := &fakeStore{ + budgets: []Budget{ + {UserPath: "/team", PeriodSeconds: PeriodDailySeconds, Amount: 10}, + }, + sum: func(userPath string, start, end time.Time) (float64, bool, error) { + called = true + return 0, false, nil + }, + } + service, err := NewService(ctx, store) + if err != nil { + t.Fatalf("NewService() failed: %v", err) + } + + results, err := service.CheckWithResults(ctx, "/team-alpha", time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)) + if err != nil { + t.Fatalf("CheckWithResults() error = %v", err) + } + if len(results) != 0 { + t.Fatalf("expected no matching budgets, got %d", len(results)) + } + if called { + t.Fatal("sum should not be called for a sibling path") + } +} + +func TestServiceCheckStartsAtManualResetWhenNewerThanPeriodStart(t *testing.T) { + ctx := context.Background() + resetAt := time.Date(2026, time.April, 25, 9, 0, 0, 0, time.UTC) + store := &fakeStore{ + budgets: []Budget{ + {UserPath: "/team", PeriodSeconds: PeriodDailySeconds, Amount: 10, LastResetAt: &resetAt}, + }, + } + service, err := NewService(ctx, store) + if err != nil { + t.Fatalf("NewService() failed: %v", err) + } + + _, err = service.CheckWithResults(ctx, "/team", time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC)) + if err != nil { + t.Fatalf("CheckWithResults() error = %v", err) + } + if !store.lastSumStart.Equal(resetAt) { + t.Fatalf("sum start = %s, want reset time %s", store.lastSumStart, resetAt) + } +} diff --git a/internal/budget/settings_helpers.go b/internal/budget/settings_helpers.go new file mode 100644 index 00000000..bcff6b34 --- /dev/null +++ b/internal/budget/settings_helpers.go @@ -0,0 +1,83 @@ +package budget + +import ( + "fmt" + "strconv" + "strings" + "time" +) + +type settingsRowScanner interface { + Next() bool + Scan(dest ...any) error + Err() error +} + +func settingsKeyValues(settings Settings) map[string]int { + return map[string]int{ + settingDailyResetHour: settings.DailyResetHour, + settingDailyResetMinute: settings.DailyResetMinute, + settingWeeklyResetWeekday: settings.WeeklyResetWeekday, + settingWeeklyResetHour: settings.WeeklyResetHour, + settingWeeklyResetMinute: settings.WeeklyResetMinute, + settingMonthlyResetDay: settings.MonthlyResetDay, + settingMonthlyResetHour: settings.MonthlyResetHour, + settingMonthlyResetMinute: settings.MonthlyResetMinute, + } +} + +func applySettingValue(settings *Settings, key, value string) error { + if settings == nil { + return nil + } + parsed, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return fmt.Errorf("budget setting %s must be an integer", key) + } + switch strings.TrimSpace(key) { + case settingDailyResetHour: + settings.DailyResetHour = parsed + case settingDailyResetMinute: + settings.DailyResetMinute = parsed + case settingWeeklyResetWeekday: + settings.WeeklyResetWeekday = parsed + case settingWeeklyResetHour: + settings.WeeklyResetHour = parsed + case settingWeeklyResetMinute: + settings.WeeklyResetMinute = parsed + case settingMonthlyResetDay: + settings.MonthlyResetDay = parsed + case settingMonthlyResetHour: + settings.MonthlyResetHour = parsed + case settingMonthlyResetMinute: + settings.MonthlyResetMinute = parsed + default: + return nil + } + return nil +} + +func scanSettingsRows(rows settingsRowScanner) (Settings, error) { + settings := DefaultSettings() + var latest int64 + for rows.Next() { + var key, value string + var updatedAt int64 + if err := rows.Scan(&key, &value, &updatedAt); err != nil { + return Settings{}, fmt.Errorf("scan budget setting: %w", err) + } + if err := applySettingValue(&settings, key, value); err != nil { + return Settings{}, err + } + if updatedAt > latest { + latest = updatedAt + } + } + if err := rows.Err(); err != nil { + return Settings{}, fmt.Errorf("iterate budget settings: %w", err) + } + if latest > 0 { + settings.UpdatedAt = time.Unix(latest, 0).UTC() + } + return normalizeLoadedSettings(settings) +} diff --git a/internal/budget/store.go b/internal/budget/store.go new file mode 100644 index 00000000..3ab3ced5 --- /dev/null +++ b/internal/budget/store.go @@ -0,0 +1,78 @@ +package budget + +import ( + "context" + "fmt" + "strings" + "time" +) + +// Store persists budget definitions, reset settings, and spend lookups. +type Store interface { + ListBudgets(ctx context.Context) ([]Budget, error) + UpsertBudgets(ctx context.Context, budgets []Budget) error + ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error + GetSettings(ctx context.Context) (Settings, error) + SaveSettings(ctx context.Context, settings Settings) (Settings, error) + ResetAllBudgets(ctx context.Context, at time.Time) error + SumUsageCost(ctx context.Context, userPath string, start, end time.Time) (float64, bool, error) + Close() error +} + +func normalizeBudgetsForUpsert(budgets []Budget) ([]Budget, error) { + if len(budgets) == 0 { + return nil, nil + } + normalized := make([]Budget, 0, len(budgets)) + seen := make(map[string]int, len(budgets)) + for _, budget := range budgets { + item, err := NormalizeBudget(budget) + if err != nil { + return nil, err + } + key := budgetKey(item.UserPath, item.PeriodSeconds) + if existing, ok := seen[key]; ok { + normalized[existing] = item + continue + } + seen[key] = len(normalized) + normalized = append(normalized, item) + } + return normalized, nil +} + +func budgetKey(userPath string, periodSeconds int64) string { + return strings.TrimSpace(userPath) + ":" + fmt.Sprint(periodSeconds) +} + +func normalizeLoadedSettings(settings Settings) (Settings, error) { + defaults := DefaultSettings() + if settings.MonthlyResetDay == 0 { + settings.MonthlyResetDay = defaults.MonthlyResetDay + } + if settings.UpdatedAt.IsZero() { + settings.UpdatedAt = time.Now().UTC() + } + if err := ValidateSettings(settings); err != nil { + return Settings{}, err + } + return settings, nil +} + +func usagePathMatchesBudgetExpr(column string) string { + return "COALESCE(NULLIF(TRIM(" + column + "), ''), '/')" +} + +func usagePathLikePattern(userPath string) string { + if userPath == "/" { + return "/%" + } + return escapeLikeWildcards(userPath) + "/%" +} + +func escapeLikeWildcards(s string) string { + s = strings.ReplaceAll(s, `\`, `\\`) + s = strings.ReplaceAll(s, `%`, `\%`) + s = strings.ReplaceAll(s, `_`, `\_`) + return s +} diff --git a/internal/budget/store_mongodb.go b/internal/budget/store_mongodb.go new file mode 100644 index 00000000..1b3ba6fe --- /dev/null +++ b/internal/budget/store_mongodb.go @@ -0,0 +1,245 @@ +package budget + +import ( + "context" + "fmt" + "regexp" + "strconv" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/mongo/options" +) + +type MongoDBStore struct { + budgets *mongo.Collection + settings *mongo.Collection + usage *mongo.Collection +} + +func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBStore, error) { + if database == nil { + return nil, fmt.Errorf("database is required") + } + store := &MongoDBStore{ + budgets: database.Collection("budgets"), + settings: database.Collection("budget_settings"), + usage: database.Collection("usage"), + } + _, err := store.budgets.Indexes().CreateMany(ctx, []mongo.IndexModel{ + { + Keys: bson.D{{Key: "user_path", Value: 1}, {Key: "period_seconds", Value: 1}}, + Options: options.Index().SetUnique(true), + }, + {Keys: bson.D{{Key: "user_path", Value: 1}}}, + {Keys: bson.D{{Key: "period_seconds", Value: 1}}}, + }) + if err != nil { + return nil, fmt.Errorf("create budget indexes: %w", err) + } + _, err = store.settings.Indexes().CreateOne(ctx, mongo.IndexModel{ + Keys: bson.D{{Key: "key", Value: 1}}, + Options: options.Index().SetUnique(true), + }) + if err != nil { + return nil, fmt.Errorf("create budget settings indexes: %w", err) + } + return store, nil +} + +func (s *MongoDBStore) ListBudgets(ctx context.Context) ([]Budget, error) { + cursor, err := s.budgets.Find(ctx, bson.D{}, options.Find().SetSort(bson.D{{Key: "user_path", Value: 1}, {Key: "period_seconds", Value: 1}})) + if err != nil { + return nil, fmt.Errorf("list budgets: %w", err) + } + defer cursor.Close(ctx) + + budgets := make([]Budget, 0) + for cursor.Next(ctx) { + var budget Budget + if err := cursor.Decode(&budget); err != nil { + return nil, fmt.Errorf("decode budget: %w", err) + } + budgets = append(budgets, budget) + } + if err := cursor.Err(); err != nil { + return nil, fmt.Errorf("iterate budgets: %w", err) + } + return budgets, nil +} + +func (s *MongoDBStore) UpsertBudgets(ctx context.Context, budgets []Budget) error { + budgets, err := normalizeBudgetsForUpsert(budgets) + if err != nil { + return err + } + if len(budgets) == 0 { + return nil + } + for _, budget := range budgets { + filter := bson.D{{Key: "user_path", Value: budget.UserPath}, {Key: "period_seconds", Value: budget.PeriodSeconds}} + update := bson.D{{Key: "$set", Value: bson.D{ + {Key: "user_path", Value: budget.UserPath}, + {Key: "period_seconds", Value: budget.PeriodSeconds}, + {Key: "amount", Value: budget.Amount}, + {Key: "source", Value: budget.Source}, + {Key: "updated_at", Value: budget.UpdatedAt}, + }}, {Key: "$setOnInsert", Value: bson.D{ + {Key: "created_at", Value: budget.CreatedAt}, + {Key: "last_reset_at", Value: budget.LastResetAt}, + }}} + if _, err := s.budgets.UpdateOne(ctx, filter, update, options.UpdateOne().SetUpsert(true)); err != nil { + return fmt.Errorf("upsert budget %s/%d: %w", budget.UserPath, budget.PeriodSeconds, err) + } + } + return nil +} + +func (s *MongoDBStore) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error { + budgets, err := normalizeBudgetsForUpsert(budgets) + if err != nil { + return err + } + for i := range budgets { + budgets[i].Source = "config" + } + + filter := bson.D{{Key: "source", Value: "config"}} + if len(budgets) > 0 { + keep := make(bson.A, 0, len(budgets)) + for _, budget := range budgets { + keep = append(keep, bson.D{ + {Key: "user_path", Value: budget.UserPath}, + {Key: "period_seconds", Value: budget.PeriodSeconds}, + }) + } + filter = append(filter, bson.E{Key: "$nor", Value: keep}) + } + if _, err := s.budgets.DeleteMany(ctx, filter); err != nil { + return fmt.Errorf("delete old config budgets: %w", err) + } + return s.UpsertBudgets(ctx, budgets) +} + +func (s *MongoDBStore) GetSettings(ctx context.Context) (Settings, error) { + cursor, err := s.settings.Find(ctx, bson.D{}) + if err != nil { + return Settings{}, fmt.Errorf("get budget settings: %w", err) + } + defer cursor.Close(ctx) + + settings := DefaultSettings() + var latest time.Time + for cursor.Next(ctx) { + var row struct { + Key string `bson:"key"` + Value string `bson:"value"` + UpdatedAt time.Time `bson:"updated_at"` + } + if err := cursor.Decode(&row); err != nil { + return Settings{}, fmt.Errorf("decode budget setting: %w", err) + } + if err := applySettingValue(&settings, row.Key, row.Value); err != nil { + return Settings{}, err + } + if row.UpdatedAt.After(latest) { + latest = row.UpdatedAt + } + } + if err := cursor.Err(); err != nil { + return Settings{}, fmt.Errorf("iterate budget settings: %w", err) + } + if !latest.IsZero() { + settings.UpdatedAt = latest.UTC() + } + return normalizeLoadedSettings(settings) +} + +func (s *MongoDBStore) SaveSettings(ctx context.Context, settings Settings) (Settings, error) { + if err := ValidateSettings(settings); err != nil { + return Settings{}, err + } + settings.UpdatedAt = time.Now().UTC() + for key, value := range settingsKeyValues(settings) { + filter := bson.D{{Key: "key", Value: key}} + update := bson.D{{Key: "$set", Value: bson.D{ + {Key: "key", Value: key}, + {Key: "value", Value: strconv.Itoa(value)}, + {Key: "updated_at", Value: settings.UpdatedAt}, + }}} + if _, err := s.settings.UpdateOne(ctx, filter, update, options.UpdateOne().SetUpsert(true)); err != nil { + return Settings{}, fmt.Errorf("save budget setting %s: %w", key, err) + } + } + return settings, nil +} + +func (s *MongoDBStore) ResetAllBudgets(ctx context.Context, at time.Time) error { + _, err := s.budgets.UpdateMany(ctx, bson.D{}, bson.D{{Key: "$set", Value: bson.D{ + {Key: "last_reset_at", Value: at.UTC()}, + {Key: "updated_at", Value: at.UTC()}, + }}}) + if err != nil { + return fmt.Errorf("reset all budgets: %w", err) + } + return nil +} + +func (s *MongoDBStore) SumUsageCost(ctx context.Context, userPath string, start, end time.Time) (float64, bool, error) { + userPath, err := NormalizeUserPath(userPath) + if err != nil { + return 0, false, err + } + pathPattern := usagePathRegex(userPath) + pipeline := bson.A{ + bson.D{{Key: "$addFields", Value: bson.D{ + {Key: "_gomodel_budget_user_path", Value: bson.D{{Key: "$cond", Value: bson.A{ + bson.D{{Key: "$ne", Value: bson.A{bson.D{{Key: "$trim", Value: bson.D{{Key: "input", Value: bson.D{{Key: "$ifNull", Value: bson.A{"$user_path", ""}}}}}}}, ""}}}, + bson.D{{Key: "$trim", Value: bson.D{{Key: "input", Value: "$user_path"}}}}, + "/", + }}}}, + }}}, + bson.D{{Key: "$match", Value: bson.D{ + {Key: "timestamp", Value: bson.D{{Key: "$gte", Value: start.UTC()}, {Key: "$lt", Value: end.UTC()}}}, + {Key: "_gomodel_budget_user_path", Value: bson.D{{Key: "$regex", Value: pathPattern}}}, + {Key: "$or", Value: bson.A{ + bson.D{{Key: "cache_type", Value: bson.D{{Key: "$exists", Value: false}}}}, + bson.D{{Key: "cache_type", Value: ""}}, + }}, + }}}, + bson.D{{Key: "$group", Value: bson.D{ + {Key: "_id", Value: nil}, + {Key: "total", Value: bson.D{{Key: "$sum", Value: bson.D{{Key: "$ifNull", Value: bson.A{"$total_cost", 0}}}}}}, + {Key: "has_costs", Value: bson.D{{Key: "$sum", Value: bson.D{{Key: "$cond", Value: bson.A{bson.D{{Key: "$gt", Value: bson.A{"$total_cost", nil}}}, 1, 0}}}}}}, + }}}, + } + cursor, err := s.usage.Aggregate(ctx, pipeline) + if err != nil { + return 0, false, fmt.Errorf("sum usage cost: %w", err) + } + defer cursor.Close(ctx) + + if !cursor.Next(ctx) { + return 0, false, cursor.Err() + } + var row struct { + Total float64 `bson:"total"` + HasCosts int `bson:"has_costs"` + } + if err := cursor.Decode(&row); err != nil { + return 0, false, fmt.Errorf("decode usage cost sum: %w", err) + } + return row.Total, row.HasCosts > 0, nil +} + +func (s *MongoDBStore) Close() error { + return nil +} + +func usagePathRegex(userPath string) string { + if userPath == "/" { + return "^/" + } + return "^" + regexp.QuoteMeta(userPath) + "(?:/|$)" +} diff --git a/internal/budget/store_postgresql.go b/internal/budget/store_postgresql.go new file mode 100644 index 00000000..835b035d --- /dev/null +++ b/internal/budget/store_postgresql.go @@ -0,0 +1,290 @@ +package budget + +import ( + "context" + "fmt" + "strconv" + "strings" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +type PostgreSQLStore struct { + pool *pgxpool.Pool +} + +func NewPostgreSQLStore(ctx context.Context, pool *pgxpool.Pool) (*PostgreSQLStore, error) { + if pool == nil { + return nil, fmt.Errorf("connection pool is required") + } + if _, err := pool.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS budgets ( + user_path TEXT NOT NULL, + period_seconds BIGINT NOT NULL, + amount DOUBLE PRECISION NOT NULL, + source TEXT NOT NULL DEFAULT '', + last_reset_at BIGINT, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (user_path, period_seconds) + ) + `); err != nil { + return nil, fmt.Errorf("failed to create budgets table: %w", err) + } + for _, migration := range []string{ + `ALTER TABLE budgets ADD COLUMN IF NOT EXISTS source TEXT NOT NULL DEFAULT ''`, + `ALTER TABLE budgets ADD COLUMN IF NOT EXISTS last_reset_at BIGINT`, + } { + if _, err := pool.Exec(ctx, migration); err != nil { + return nil, fmt.Errorf("failed to migrate budgets table: %w", err) + } + } + if _, err := pool.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS budget_settings ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at BIGINT NOT NULL + ) + `); err != nil { + return nil, fmt.Errorf("failed to create budget_settings table: %w", err) + } + for _, index := range []string{ + `CREATE INDEX IF NOT EXISTS idx_budgets_user_path ON budgets(user_path)`, + `CREATE INDEX IF NOT EXISTS idx_budgets_period_seconds ON budgets(period_seconds)`, + } { + if _, err := pool.Exec(ctx, index); err != nil { + return nil, fmt.Errorf("failed to create budget index: %w", err) + } + } + return &PostgreSQLStore{pool: pool}, nil +} + +func (s *PostgreSQLStore) ListBudgets(ctx context.Context) ([]Budget, error) { + rows, err := s.pool.Query(ctx, ` + SELECT user_path, period_seconds, amount, source, last_reset_at, created_at, updated_at + FROM budgets + ORDER BY user_path ASC, period_seconds ASC + `) + if err != nil { + return nil, fmt.Errorf("list budgets: %w", err) + } + defer rows.Close() + + budgets := make([]Budget, 0) + for rows.Next() { + budget, err := scanPostgreSQLBudget(rows) + if err != nil { + return nil, err + } + budgets = append(budgets, budget) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate budgets: %w", err) + } + return budgets, nil +} + +func (s *PostgreSQLStore) UpsertBudgets(ctx context.Context, budgets []Budget) error { + budgets, err := normalizeBudgetsForUpsert(budgets) + if err != nil { + return err + } + if len(budgets) == 0 { + return nil + } + tx, err := s.pool.Begin(ctx) + if err != nil { + return fmt.Errorf("begin budget upsert: %w", err) + } + defer tx.Rollback(ctx) //nolint:errcheck + + for _, budget := range budgets { + _, err := tx.Exec(ctx, ` + INSERT INTO budgets (user_path, period_seconds, amount, source, last_reset_at, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ON CONFLICT (user_path, period_seconds) DO UPDATE SET + amount = excluded.amount, + source = excluded.source, + updated_at = excluded.updated_at + `, + budget.UserPath, + budget.PeriodSeconds, + budget.Amount, + budget.Source, + unixOrNil(budget.LastResetAt), + budget.CreatedAt.Unix(), + budget.UpdatedAt.Unix(), + ) + if err != nil { + return fmt.Errorf("upsert budget %s/%d: %w", budget.UserPath, budget.PeriodSeconds, err) + } + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit budget upsert: %w", err) + } + return nil +} + +func (s *PostgreSQLStore) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error { + budgets, err := normalizeBudgetsForUpsert(budgets) + if err != nil { + return err + } + for i := range budgets { + budgets[i].Source = "config" + } + + tx, err := s.pool.Begin(ctx) + if err != nil { + return fmt.Errorf("begin config budget replace: %w", err) + } + defer tx.Rollback(ctx) //nolint:errcheck + + if len(budgets) == 0 { + if _, err := tx.Exec(ctx, `DELETE FROM budgets WHERE source = 'config'`); err != nil { + return fmt.Errorf("delete old config budgets: %w", err) + } + } else { + conditions := make([]string, 0, len(budgets)) + args := make([]any, 0, len(budgets)*2) + for i, budget := range budgets { + base := i*2 + 1 + conditions = append(conditions, fmt.Sprintf(`(user_path = $%d AND period_seconds = $%d)`, base, base+1)) + args = append(args, budget.UserPath, budget.PeriodSeconds) + } + query := `DELETE FROM budgets WHERE source = 'config' AND NOT (` + strings.Join(conditions, " OR ") + `)` + if _, err := tx.Exec(ctx, query, args...); err != nil { + return fmt.Errorf("delete old config budgets: %w", err) + } + } + if err := upsertPostgreSQLBudgets(ctx, tx, budgets); err != nil { + return err + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit config budget replace: %w", err) + } + return nil +} + +func (s *PostgreSQLStore) GetSettings(ctx context.Context) (Settings, error) { + rows, err := s.pool.Query(ctx, `SELECT key, value, updated_at FROM budget_settings`) + if err != nil { + return Settings{}, fmt.Errorf("get budget settings: %w", err) + } + defer rows.Close() + + return scanSettingsRows(rows) +} + +func (s *PostgreSQLStore) SaveSettings(ctx context.Context, settings Settings) (Settings, error) { + if err := ValidateSettings(settings); err != nil { + return Settings{}, err + } + settings.UpdatedAt = time.Now().UTC() + tx, err := s.pool.Begin(ctx) + if err != nil { + return Settings{}, fmt.Errorf("begin budget settings save: %w", err) + } + defer tx.Rollback(ctx) //nolint:errcheck + + for key, value := range settingsKeyValues(settings) { + if _, err := tx.Exec(ctx, ` + INSERT INTO budget_settings (key, value, updated_at) + VALUES ($1, $2, $3) + ON CONFLICT (key) DO UPDATE SET + value = excluded.value, + updated_at = excluded.updated_at + `, key, strconv.Itoa(value), settings.UpdatedAt.Unix()); err != nil { + return Settings{}, fmt.Errorf("save budget setting %s: %w", key, err) + } + } + if err := tx.Commit(ctx); err != nil { + return Settings{}, fmt.Errorf("commit budget settings save: %w", err) + } + return settings, nil +} + +func (s *PostgreSQLStore) ResetAllBudgets(ctx context.Context, at time.Time) error { + _, err := s.pool.Exec(ctx, `UPDATE budgets SET last_reset_at = $1, updated_at = $2`, at.UTC().Unix(), at.UTC().Unix()) + if err != nil { + return fmt.Errorf("reset all budgets: %w", err) + } + return nil +} + +func (s *PostgreSQLStore) SumUsageCost(ctx context.Context, userPath string, start, end time.Time) (float64, bool, error) { + userPath, err := NormalizeUserPath(userPath) + if err != nil { + return 0, false, err + } + userPathExpr := usagePathMatchesBudgetExpr("user_path") + query := `SELECT SUM(total_cost) FROM "usage" + WHERE timestamp >= $1 + AND timestamp < $2 + AND (` + userPathExpr + ` = $3 OR ` + userPathExpr + ` LIKE $4 ESCAPE '\') + AND (cache_type IS NULL OR cache_type = '')` + var total *float64 + if err := s.pool.QueryRow(ctx, query, start.UTC(), end.UTC(), userPath, usagePathLikePattern(userPath)).Scan(&total); err != nil { + return 0, false, fmt.Errorf("sum usage cost: %w", err) + } + if total == nil { + return 0, false, nil + } + return *total, true, nil +} + +func (s *PostgreSQLStore) Close() error { + return nil +} + +func upsertPostgreSQLBudgets(ctx context.Context, tx pgx.Tx, budgets []Budget) error { + for _, budget := range budgets { + _, err := tx.Exec(ctx, ` + INSERT INTO budgets (user_path, period_seconds, amount, source, last_reset_at, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ON CONFLICT (user_path, period_seconds) DO UPDATE SET + amount = excluded.amount, + source = excluded.source, + updated_at = excluded.updated_at + `, + budget.UserPath, + budget.PeriodSeconds, + budget.Amount, + budget.Source, + unixOrNil(budget.LastResetAt), + budget.CreatedAt.Unix(), + budget.UpdatedAt.Unix(), + ) + if err != nil { + return fmt.Errorf("upsert budget %s/%d: %w", budget.UserPath, budget.PeriodSeconds, err) + } + } + return nil +} + +func scanPostgreSQLBudget(row pgx.Row) (Budget, error) { + var budget Budget + var lastResetAt *int64 + var createdAt int64 + var updatedAt int64 + if err := row.Scan( + &budget.UserPath, + &budget.PeriodSeconds, + &budget.Amount, + &budget.Source, + &lastResetAt, + &createdAt, + &updatedAt, + ); err != nil { + return Budget{}, fmt.Errorf("scan budget: %w", err) + } + if lastResetAt != nil { + t := time.Unix(*lastResetAt, 0).UTC() + budget.LastResetAt = &t + } + budget.CreatedAt = time.Unix(createdAt, 0).UTC() + budget.UpdatedAt = time.Unix(updatedAt, 0).UTC() + return budget, nil +} diff --git a/internal/budget/store_sqlite.go b/internal/budget/store_sqlite.go new file mode 100644 index 00000000..70597939 --- /dev/null +++ b/internal/budget/store_sqlite.go @@ -0,0 +1,353 @@ +package budget + +import ( + "context" + "database/sql" + "fmt" + "strconv" + "strings" + "time" +) + +const ( + settingDailyResetHour = "daily_reset_hour" + settingDailyResetMinute = "daily_reset_minute" + settingWeeklyResetWeekday = "weekly_reset_weekday" + settingWeeklyResetHour = "weekly_reset_hour" + settingWeeklyResetMinute = "weekly_reset_minute" + settingMonthlyResetDay = "monthly_reset_day" + settingMonthlyResetHour = "monthly_reset_hour" + settingMonthlyResetMinute = "monthly_reset_minute" +) + +type SQLiteStore struct { + db *sql.DB +} + +func NewSQLiteStore(db *sql.DB) (*SQLiteStore, error) { + if db == nil { + return nil, fmt.Errorf("database connection is required") + } + if _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS budgets ( + user_path TEXT NOT NULL, + period_seconds INTEGER NOT NULL, + amount REAL NOT NULL, + source TEXT NOT NULL DEFAULT '', + last_reset_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (user_path, period_seconds) + ) + `); err != nil { + return nil, fmt.Errorf("failed to create budgets table: %w", err) + } + for _, migration := range []string{ + `ALTER TABLE budgets ADD COLUMN source TEXT NOT NULL DEFAULT ''`, + `ALTER TABLE budgets ADD COLUMN last_reset_at INTEGER`, + } { + if _, err := db.Exec(migration); err != nil && !isSQLiteDuplicateColumnError(err) { + return nil, fmt.Errorf("failed to migrate budgets table: %w", err) + } + } + if _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS budget_settings ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at INTEGER NOT NULL + ) + `); err != nil { + return nil, fmt.Errorf("failed to create budget_settings table: %w", err) + } + for _, index := range []string{ + `CREATE INDEX IF NOT EXISTS idx_budgets_user_path ON budgets(user_path)`, + `CREATE INDEX IF NOT EXISTS idx_budgets_period_seconds ON budgets(period_seconds)`, + } { + if _, err := db.Exec(index); err != nil { + return nil, fmt.Errorf("failed to create budget index: %w", err) + } + } + return &SQLiteStore{db: db}, nil +} + +func (s *SQLiteStore) ListBudgets(ctx context.Context) ([]Budget, error) { + rows, err := s.db.QueryContext(ctx, ` + SELECT user_path, period_seconds, amount, source, last_reset_at, created_at, updated_at + FROM budgets + ORDER BY user_path ASC, period_seconds ASC + `) + if err != nil { + return nil, fmt.Errorf("list budgets: %w", err) + } + defer rows.Close() + + var budgets []Budget + for rows.Next() { + budget, err := scanSQLiteBudget(rows) + if err != nil { + return nil, err + } + budgets = append(budgets, budget) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate budgets: %w", err) + } + return budgets, nil +} + +func (s *SQLiteStore) UpsertBudgets(ctx context.Context, budgets []Budget) error { + budgets, err := normalizeBudgetsForUpsert(budgets) + if err != nil { + return err + } + if len(budgets) == 0 { + return nil + } + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin budget upsert: %w", err) + } + defer tx.Rollback() //nolint:errcheck + + stmt, err := tx.PrepareContext(ctx, ` + INSERT INTO budgets (user_path, period_seconds, amount, source, last_reset_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(user_path, period_seconds) DO UPDATE SET + amount = excluded.amount, + source = excluded.source, + updated_at = excluded.updated_at + `) + if err != nil { + return fmt.Errorf("prepare budget upsert: %w", err) + } + defer stmt.Close() + + for _, budget := range budgets { + if _, err := stmt.ExecContext( + ctx, + budget.UserPath, + budget.PeriodSeconds, + budget.Amount, + budget.Source, + unixOrNil(budget.LastResetAt), + budget.CreatedAt.Unix(), + budget.UpdatedAt.Unix(), + ); err != nil { + return fmt.Errorf("upsert budget %s/%d: %w", budget.UserPath, budget.PeriodSeconds, err) + } + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit budget upsert: %w", err) + } + return nil +} + +func (s *SQLiteStore) ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error { + budgets, err := normalizeBudgetsForUpsert(budgets) + if err != nil { + return err + } + for i := range budgets { + budgets[i].Source = "config" + } + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin config budget replace: %w", err) + } + defer tx.Rollback() //nolint:errcheck + + if len(budgets) == 0 { + if _, err := tx.ExecContext(ctx, `DELETE FROM budgets WHERE source = 'config'`); err != nil { + return fmt.Errorf("delete old config budgets: %w", err) + } + } else { + conditions := make([]string, 0, len(budgets)) + args := make([]any, 0, len(budgets)*2) + for _, budget := range budgets { + conditions = append(conditions, `(user_path = ? AND period_seconds = ?)`) + args = append(args, budget.UserPath, budget.PeriodSeconds) + } + query := `DELETE FROM budgets WHERE source = 'config' AND NOT (` + strings.Join(conditions, " OR ") + `)` + if _, err := tx.ExecContext(ctx, query, args...); err != nil { + return fmt.Errorf("delete old config budgets: %w", err) + } + } + if err := upsertSQLiteBudgets(ctx, tx, budgets); err != nil { + return err + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit config budget replace: %w", err) + } + return nil +} + +func (s *SQLiteStore) GetSettings(ctx context.Context) (Settings, error) { + rows, err := s.db.QueryContext(ctx, `SELECT key, value, updated_at FROM budget_settings`) + if err != nil { + return Settings{}, fmt.Errorf("get budget settings: %w", err) + } + defer rows.Close() + + return scanSettingsRows(rows) +} + +func (s *SQLiteStore) SaveSettings(ctx context.Context, settings Settings) (Settings, error) { + if err := ValidateSettings(settings); err != nil { + return Settings{}, err + } + settings.UpdatedAt = time.Now().UTC() + values := settingsKeyValues(settings) + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return Settings{}, fmt.Errorf("begin budget settings save: %w", err) + } + defer tx.Rollback() //nolint:errcheck + + stmt, err := tx.PrepareContext(ctx, ` + INSERT INTO budget_settings (key, value, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value = excluded.value, + updated_at = excluded.updated_at + `) + if err != nil { + return Settings{}, fmt.Errorf("prepare budget settings save: %w", err) + } + defer stmt.Close() + + for key, value := range values { + if _, err := stmt.ExecContext(ctx, key, strconv.Itoa(value), settings.UpdatedAt.Unix()); err != nil { + return Settings{}, fmt.Errorf("save budget setting %s: %w", key, err) + } + } + if err := tx.Commit(); err != nil { + return Settings{}, fmt.Errorf("commit budget settings save: %w", err) + } + return settings, nil +} + +func (s *SQLiteStore) ResetAllBudgets(ctx context.Context, at time.Time) error { + _, err := s.db.ExecContext(ctx, `UPDATE budgets SET last_reset_at = ?, updated_at = ?`, at.UTC().Unix(), at.UTC().Unix()) + if err != nil { + return fmt.Errorf("reset all budgets: %w", err) + } + return nil +} + +func (s *SQLiteStore) SumUsageCost(ctx context.Context, userPath string, start, end time.Time) (float64, bool, error) { + userPath, err := NormalizeUserPath(userPath) + if err != nil { + return 0, false, err + } + userPathExpr := usagePathMatchesBudgetExpr("user_path") + query := `SELECT SUM(total_cost) FROM usage + WHERE ` + sqliteTimestampEpochExpr() + ` >= unixepoch(?) + AND ` + sqliteTimestampEpochExpr() + ` < unixepoch(?) + AND (` + userPathExpr + ` = ? OR ` + userPathExpr + ` LIKE ? ESCAPE '\') + AND (cache_type IS NULL OR cache_type = '')` + var total sql.NullFloat64 + if err := s.db.QueryRowContext( + ctx, + query, + start.UTC().Format(time.RFC3339Nano), + end.UTC().Format(time.RFC3339Nano), + userPath, + usagePathLikePattern(userPath), + ).Scan(&total); err != nil { + return 0, false, fmt.Errorf("sum usage cost: %w", err) + } + if !total.Valid { + return 0, false, nil + } + return total.Float64, true, nil +} + +func (s *SQLiteStore) Close() error { + return nil +} + +func upsertSQLiteBudgets(ctx context.Context, tx *sql.Tx, budgets []Budget) error { + if len(budgets) == 0 { + return nil + } + stmt, err := tx.PrepareContext(ctx, ` + INSERT INTO budgets (user_path, period_seconds, amount, source, last_reset_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(user_path, period_seconds) DO UPDATE SET + amount = excluded.amount, + source = excluded.source, + updated_at = excluded.updated_at + `) + if err != nil { + return fmt.Errorf("prepare budget upsert: %w", err) + } + defer stmt.Close() + + for _, budget := range budgets { + if _, err := stmt.ExecContext( + ctx, + budget.UserPath, + budget.PeriodSeconds, + budget.Amount, + budget.Source, + unixOrNil(budget.LastResetAt), + budget.CreatedAt.Unix(), + budget.UpdatedAt.Unix(), + ); err != nil { + return fmt.Errorf("upsert budget %s/%d: %w", budget.UserPath, budget.PeriodSeconds, err) + } + } + return nil +} + +func scanSQLiteBudget(scanner interface{ Scan(dest ...any) error }) (Budget, error) { + var budget Budget + var lastResetAt sql.NullInt64 + var createdAt int64 + var updatedAt int64 + if err := scanner.Scan( + &budget.UserPath, + &budget.PeriodSeconds, + &budget.Amount, + &budget.Source, + &lastResetAt, + &createdAt, + &updatedAt, + ); err != nil { + return Budget{}, fmt.Errorf("scan budget: %w", err) + } + budget.LastResetAt = unixPtr(lastResetAt) + budget.CreatedAt = time.Unix(createdAt, 0).UTC() + budget.UpdatedAt = time.Unix(updatedAt, 0).UTC() + return budget, nil +} + +func sqliteTimestampEpochExpr() string { + return "unixepoch(REPLACE(timestamp, ' ', 'T'))" +} + +func isSQLiteDuplicateColumnError(err error) bool { + if err == nil { + return false + } + message := strings.ToLower(err.Error()) + return strings.Contains(message, "duplicate column") || strings.Contains(message, "already exists") +} + +func unixOrNil(value *time.Time) any { + if value == nil { + return nil + } + return value.UTC().Unix() +} + +func unixPtr(value sql.NullInt64) *time.Time { + if !value.Valid { + return nil + } + t := time.Unix(value.Int64, 0).UTC() + return &t +} diff --git a/internal/budget/store_sqlite_test.go b/internal/budget/store_sqlite_test.go new file mode 100644 index 00000000..5b90161d --- /dev/null +++ b/internal/budget/store_sqlite_test.go @@ -0,0 +1,66 @@ +package budget + +import ( + "context" + "database/sql" + "testing" + "time" + + _ "modernc.org/sqlite" +) + +func TestSQLiteStoreReplaceConfigBudgetsRemovesStaleConfigRowsOnly(t *testing.T) { + ctx := context.Background() + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("sql.Open() failed: %v", err) + } + defer db.Close() + + store, err := NewSQLiteStore(db) + if err != nil { + t.Fatalf("NewSQLiteStore() failed: %v", err) + } + resetAt := time.Date(2026, time.April, 25, 9, 0, 0, 0, time.UTC) + if err := store.UpsertBudgets(ctx, []Budget{ + {UserPath: "/team", PeriodSeconds: PeriodDailySeconds, Amount: 10, Source: "config"}, + {UserPath: "/team", PeriodSeconds: PeriodWeeklySeconds, Amount: 50, Source: "config", LastResetAt: &resetAt}, + {UserPath: "/manual", PeriodSeconds: PeriodDailySeconds, Amount: 5, Source: "manual"}, + }); err != nil { + t.Fatalf("UpsertBudgets() failed: %v", err) + } + + if err := store.ReplaceConfigBudgets(ctx, []Budget{ + {UserPath: "/team", PeriodSeconds: PeriodWeeklySeconds, Amount: 75}, + }); err != nil { + t.Fatalf("ReplaceConfigBudgets() failed: %v", err) + } + + got, err := store.ListBudgets(ctx) + if err != nil { + t.Fatalf("ListBudgets() failed: %v", err) + } + if len(got) != 2 { + t.Fatalf("expected 2 budgets after replacement, got %d: %+v", len(got), got) + } + byKey := make(map[string]Budget, len(got)) + for _, budget := range got { + byKey[budgetKey(budget.UserPath, budget.PeriodSeconds)] = budget + } + if _, ok := byKey[budgetKey("/team", PeriodDailySeconds)]; ok { + t.Fatal("stale config daily budget was not removed") + } + weekly := byKey[budgetKey("/team", PeriodWeeklySeconds)] + if weekly.Amount != 75 { + t.Fatalf("weekly amount = %v, want 75", weekly.Amount) + } + if weekly.Source != "config" { + t.Fatalf("weekly source = %q, want config", weekly.Source) + } + if weekly.LastResetAt == nil || !weekly.LastResetAt.Equal(resetAt) { + t.Fatalf("weekly last_reset_at = %v, want %s", weekly.LastResetAt, resetAt) + } + if _, ok := byKey[budgetKey("/manual", PeriodDailySeconds)]; !ok { + t.Fatal("manual budget was removed by config replacement") + } +} diff --git a/internal/budget/types.go b/internal/budget/types.go new file mode 100644 index 00000000..003f3d79 --- /dev/null +++ b/internal/budget/types.go @@ -0,0 +1,238 @@ +package budget + +import ( + "fmt" + "strings" + "time" + + "gomodel/internal/core" +) + +const ( + PeriodHourlySeconds int64 = 3600 + PeriodDailySeconds int64 = 86400 + PeriodWeeklySeconds int64 = 604800 + PeriodMonthlySeconds int64 = 2592000 +) + +// Budget stores one spend limit for one user path and reset period. +type Budget struct { + UserPath string `json:"user_path" bson:"user_path"` + PeriodSeconds int64 `json:"period_seconds" bson:"period_seconds"` + Amount float64 `json:"amount" bson:"amount"` + Source string `json:"source,omitempty" bson:"source,omitempty"` + LastResetAt *time.Time `json:"last_reset_at,omitempty" bson:"last_reset_at,omitempty"` + CreatedAt time.Time `json:"created_at" bson:"created_at"` + UpdatedAt time.Time `json:"updated_at" bson:"updated_at"` +} + +// Settings controls the calendar anchors used to find the active budget period. +// Values are interpreted in UTC. +type Settings struct { + DailyResetHour int `json:"daily_reset_hour" bson:"daily_reset_hour"` + DailyResetMinute int `json:"daily_reset_minute" bson:"daily_reset_minute"` + WeeklyResetWeekday int `json:"weekly_reset_weekday" bson:"weekly_reset_weekday"` + WeeklyResetHour int `json:"weekly_reset_hour" bson:"weekly_reset_hour"` + WeeklyResetMinute int `json:"weekly_reset_minute" bson:"weekly_reset_minute"` + MonthlyResetDay int `json:"monthly_reset_day" bson:"monthly_reset_day"` + MonthlyResetHour int `json:"monthly_reset_hour" bson:"monthly_reset_hour"` + MonthlyResetMinute int `json:"monthly_reset_minute" bson:"monthly_reset_minute"` + UpdatedAt time.Time `json:"updated_at" bson:"updated_at"` +} + +// CheckResult describes one evaluated budget limit. +type CheckResult struct { + Budget Budget `json:"budget"` + PeriodStart time.Time `json:"period_start"` + PeriodEnd time.Time `json:"period_end"` + Spent float64 `json:"spent"` + Remaining float64 `json:"remaining"` +} + +// ExceededError indicates a budget has already been exhausted. +type ExceededError struct { + Result CheckResult +} + +func (e *ExceededError) Error() string { + if e == nil { + return "" + } + return fmt.Sprintf( + "budget exceeded for %s %s limit: spent %.6f of %.6f", + e.Result.Budget.UserPath, + PeriodLabel(e.Result.Budget.PeriodSeconds), + e.Result.Spent, + e.Result.Budget.Amount, + ) +} + +// DefaultSettings returns the reset anchors used when no DB setting exists. +func DefaultSettings() Settings { + return Settings{ + DailyResetHour: 0, + DailyResetMinute: 0, + WeeklyResetWeekday: int(time.Monday), + WeeklyResetHour: 0, + WeeklyResetMinute: 0, + MonthlyResetDay: 1, + MonthlyResetHour: 0, + MonthlyResetMinute: 0, + } +} + +func NormalizeUserPath(raw string) (string, error) { + path, err := core.NormalizeUserPath(raw) + if err != nil { + return "", err + } + if path == "" { + return "/", nil + } + return path, nil +} + +func NormalizeBudget(b Budget) (Budget, error) { + path, err := NormalizeUserPath(b.UserPath) + if err != nil { + return Budget{}, err + } + b.UserPath = path + if b.PeriodSeconds <= 0 { + return Budget{}, fmt.Errorf("period_seconds must be greater than 0") + } + if b.Amount <= 0 { + return Budget{}, fmt.Errorf("amount must be greater than 0") + } + b.Source = strings.TrimSpace(b.Source) + if b.LastResetAt != nil { + t := b.LastResetAt.UTC() + b.LastResetAt = &t + } + now := time.Now().UTC() + if b.CreatedAt.IsZero() { + b.CreatedAt = now + } + b.UpdatedAt = now + return b, nil +} + +func ValidateSettings(settings Settings) error { + if settings.DailyResetHour < 0 || settings.DailyResetHour > 23 { + return fmt.Errorf("daily_reset_hour must be between 0 and 23") + } + if settings.DailyResetMinute < 0 || settings.DailyResetMinute > 59 { + return fmt.Errorf("daily_reset_minute must be between 0 and 59") + } + if settings.WeeklyResetWeekday < int(time.Sunday) || settings.WeeklyResetWeekday > int(time.Saturday) { + return fmt.Errorf("weekly_reset_weekday must be between 0 and 6") + } + if settings.WeeklyResetHour < 0 || settings.WeeklyResetHour > 23 { + return fmt.Errorf("weekly_reset_hour must be between 0 and 23") + } + if settings.WeeklyResetMinute < 0 || settings.WeeklyResetMinute > 59 { + return fmt.Errorf("weekly_reset_minute must be between 0 and 59") + } + if settings.MonthlyResetDay < 1 || settings.MonthlyResetDay > 31 { + return fmt.Errorf("monthly_reset_day must be between 1 and 31") + } + if settings.MonthlyResetHour < 0 || settings.MonthlyResetHour > 23 { + return fmt.Errorf("monthly_reset_hour must be between 0 and 23") + } + if settings.MonthlyResetMinute < 0 || settings.MonthlyResetMinute > 59 { + return fmt.Errorf("monthly_reset_minute must be between 0 and 59") + } + return nil +} + +func PeriodSeconds(period string) (int64, bool) { + switch strings.ToLower(strings.TrimSpace(period)) { + case "hour", "hourly", "hours": + return PeriodHourlySeconds, true + case "day", "daily", "days": + return PeriodDailySeconds, true + case "week", "weekly", "weeks": + return PeriodWeeklySeconds, true + case "month", "monthly", "months": + return PeriodMonthlySeconds, true + default: + return 0, false + } +} + +func PeriodLabel(seconds int64) string { + switch seconds { + case PeriodHourlySeconds: + return "hourly" + case PeriodDailySeconds: + return "daily" + case PeriodWeeklySeconds: + return "weekly" + case PeriodMonthlySeconds: + return "monthly" + default: + return fmt.Sprintf("%ds", seconds) + } +} + +func PeriodBounds(now time.Time, seconds int64, settings Settings) (time.Time, time.Time) { + now = now.UTC() + switch seconds { + case PeriodHourlySeconds: + start := now.Truncate(time.Hour) + return start, start.Add(time.Hour) + case PeriodDailySeconds: + start := anchoredDayStart(now, settings.DailyResetHour, settings.DailyResetMinute) + return start, start.AddDate(0, 0, 1) + case PeriodWeeklySeconds: + start := anchoredWeekStart(now, time.Weekday(settings.WeeklyResetWeekday), settings.WeeklyResetHour, settings.WeeklyResetMinute) + return start, start.AddDate(0, 0, 7) + case PeriodMonthlySeconds: + start := anchoredMonthStart(now, settings.MonthlyResetDay, settings.MonthlyResetHour, settings.MonthlyResetMinute) + nextMonth := time.Date(start.Year(), start.Month()+1, 1, 0, 0, 0, 0, time.UTC) + return start, monthAnchor(nextMonth.Year(), nextMonth.Month(), settings.MonthlyResetDay, settings.MonthlyResetHour, settings.MonthlyResetMinute) + default: + if seconds <= 0 { + return now, now + } + startUnix := now.Unix() - (now.Unix() % seconds) + start := time.Unix(startUnix, 0).UTC() + return start, start.Add(time.Duration(seconds) * time.Second) + } +} + +func anchoredDayStart(now time.Time, hour, minute int) time.Time { + start := time.Date(now.Year(), now.Month(), now.Day(), hour, minute, 0, 0, time.UTC) + if now.Before(start) { + start = start.AddDate(0, 0, -1) + } + return start +} + +func anchoredWeekStart(now time.Time, weekday time.Weekday, hour, minute int) time.Time { + todayAnchor := time.Date(now.Year(), now.Month(), now.Day(), hour, minute, 0, 0, time.UTC) + daysBack := (int(todayAnchor.Weekday()) - int(weekday) + 7) % 7 + start := todayAnchor.AddDate(0, 0, -daysBack) + if now.Before(start) { + start = start.AddDate(0, 0, -7) + } + return start +} + +func anchoredMonthStart(now time.Time, day, hour, minute int) time.Time { + start := monthAnchor(now.Year(), now.Month(), day, hour, minute) + if now.Before(start) { + prev := now.AddDate(0, -1, 0) + start = monthAnchor(prev.Year(), prev.Month(), day, hour, minute) + } + return start +} + +func monthAnchor(year int, month time.Month, day, hour, minute int) time.Time { + anchorDay := min(day, daysInMonth(year, month)) + return time.Date(year, month, anchorDay, hour, minute, 0, 0, time.UTC) +} + +func daysInMonth(year int, month time.Month) int { + return time.Date(year, month+1, 0, 0, 0, 0, 0, time.UTC).Day() +} diff --git a/internal/budget/types_test.go b/internal/budget/types_test.go new file mode 100644 index 00000000..fbbc2b4d --- /dev/null +++ b/internal/budget/types_test.go @@ -0,0 +1,44 @@ +package budget + +import ( + "testing" + "time" +) + +func TestPeriodBoundsUsesConfiguredAnchors(t *testing.T) { + settings := Settings{ + DailyResetHour: 6, + DailyResetMinute: 30, + WeeklyResetWeekday: int(time.Wednesday), + WeeklyResetHour: 9, + WeeklyResetMinute: 15, + MonthlyResetDay: 31, + MonthlyResetHour: 2, + MonthlyResetMinute: 45, + } + now := time.Date(2026, time.April, 25, 12, 0, 0, 0, time.UTC) + + dailyStart, dailyEnd := PeriodBounds(now, PeriodDailySeconds, settings) + if want := time.Date(2026, time.April, 25, 6, 30, 0, 0, time.UTC); !dailyStart.Equal(want) { + t.Fatalf("daily start = %s, want %s", dailyStart, want) + } + if want := time.Date(2026, time.April, 26, 6, 30, 0, 0, time.UTC); !dailyEnd.Equal(want) { + t.Fatalf("daily end = %s, want %s", dailyEnd, want) + } + + weeklyStart, weeklyEnd := PeriodBounds(now, PeriodWeeklySeconds, settings) + if want := time.Date(2026, time.April, 22, 9, 15, 0, 0, time.UTC); !weeklyStart.Equal(want) { + t.Fatalf("weekly start = %s, want %s", weeklyStart, want) + } + if want := time.Date(2026, time.April, 29, 9, 15, 0, 0, time.UTC); !weeklyEnd.Equal(want) { + t.Fatalf("weekly end = %s, want %s", weeklyEnd, want) + } + + monthlyStart, monthlyEnd := PeriodBounds(now, PeriodMonthlySeconds, settings) + if want := time.Date(2026, time.March, 31, 2, 45, 0, 0, time.UTC); !monthlyStart.Equal(want) { + t.Fatalf("monthly start = %s, want %s", monthlyStart, want) + } + if want := time.Date(2026, time.April, 30, 2, 45, 0, 0, time.UTC); !monthlyEnd.Equal(want) { + t.Fatalf("monthly end = %s, want %s", monthlyEnd, want) + } +} diff --git a/internal/server/budget_support.go b/internal/server/budget_support.go new file mode 100644 index 00000000..16338a17 --- /dev/null +++ b/internal/server/budget_support.go @@ -0,0 +1,45 @@ +package server + +import ( + "context" + "errors" + "fmt" + "net/http" + "time" + + "github.com/labstack/echo/v5" + + "gomodel/internal/budget" + "gomodel/internal/core" +) + +type BudgetChecker interface { + Check(ctx context.Context, userPath string, now time.Time) error +} + +func enforceBudget(c *echo.Context, checker BudgetChecker) error { + if checker == nil || c == nil || c.Request() == nil { + return nil + } + userPath := core.UserPathFromContext(c.Request().Context()) + if userPath == "" { + userPath = "/" + } + if err := checker.Check(c.Request().Context(), userPath, time.Now().UTC()); err != nil { + return budgetCheckError(err) + } + return nil +} + +func budgetCheckError(err error) error { + var exceeded *budget.ExceededError + if errors.As(err, &exceeded) { + message := exceeded.Error() + if message == "" { + message = "budget exceeded" + } + return core.NewRateLimitError("budget", message).WithCode("budget_exceeded") + } + return core.NewProviderError("budget", http.StatusServiceUnavailable, fmt.Sprintf("budget check failed: %v", err), err). + WithCode("budget_check_failed") +} diff --git a/internal/server/handlers.go b/internal/server/handlers.go index 5f387675..6ebf98e1 100644 --- a/internal/server/handlers.go +++ b/internal/server/handlers.go @@ -28,6 +28,7 @@ type Handler struct { keepOnlyAliasesAtModelsEndpoint bool logger auditlog.LoggerInterface usageLogger usage.LoggerInterface + budgetChecker BudgetChecker pricingResolver usage.PricingResolver batchStore batchstore.Store responseStore responsestore.Store @@ -134,6 +135,7 @@ func (h *Handler) translatedInference() *translatedInferenceService { translatedRequestPatcher: h.translatedRequestPatcher, logger: h.logger, usageLogger: h.usageLogger, + budgetChecker: h.budgetChecker, pricingResolver: h.pricingResolver, responseCache: h.responseCache, guardrailsHash: h.guardrailsHash, @@ -161,6 +163,7 @@ func (h *Handler) nativeBatch() *nativeBatchService { cleanupPreparedBatchInputFile: h.cleanupPreparedBatchInputFile, cleanupStoredBatchRewrittenInputFile: h.cleanupStoredBatchRewrittenInputFile, usageLogger: h.usageLogger, + budgetChecker: h.budgetChecker, pricingResolver: h.pricingResolver, } } @@ -192,6 +195,7 @@ func (h *Handler) passthrough() *passthroughService { modelAuthorizer: h.modelAuthorizer, logger: h.logger, usageLogger: h.usageLogger, + budgetChecker: h.budgetChecker, pricingResolver: h.pricingResolver, normalizePassthroughV1Prefix: h.normalizePassthroughV1Prefix, enabledPassthroughProviders: h.enabledPassthroughProviders, diff --git a/internal/server/http.go b/internal/server/http.go index 99c0d934..c9a5ca9c 100644 --- a/internal/server/http.go +++ b/internal/server/http.go @@ -53,6 +53,7 @@ type Config struct { PprofEnabled bool // Whether to expose debug profiling routes at /debug/pprof/* AuditLogger auditlog.LoggerInterface // Optional: Audit logger for request/response logging UsageLogger usage.LoggerInterface // Optional: Usage logger for token tracking + BudgetChecker BudgetChecker // Optional: per-user-path budget checker PricingResolver usage.PricingResolver // Optional: Resolves pricing for cost calculation ModelResolver RequestModelResolver // Optional: explicit model resolver used during workflow resolution ModelAuthorizer RequestModelAuthorizer // Optional: request-scoped concrete model access controller @@ -94,10 +95,12 @@ func New(provider core.RoutableProvider, cfg *Config) *Server { // Get loggers from config (may be nil) var auditLogger auditlog.LoggerInterface var usageLogger usage.LoggerInterface + var budgetChecker BudgetChecker var pricingResolver usage.PricingResolver if cfg != nil { auditLogger = cfg.AuditLogger usageLogger = cfg.UsageLogger + budgetChecker = cfg.BudgetChecker pricingResolver = cfg.PricingResolver } @@ -115,6 +118,7 @@ func New(provider core.RoutableProvider, cfg *Config) *Server { } handler := newHandlerWithAuthorizer(provider, auditLogger, usageLogger, pricingResolver, modelResolver, modelAuthorizer, workflowPolicyResolver, fallbackResolver, translatedRequestPatcher) + handler.budgetChecker = budgetChecker if cfg != nil { handler.batchRequestPreparer = cfg.BatchRequestPreparer handler.exposedModelLister = cfg.ExposedModelLister @@ -322,6 +326,9 @@ func New(provider core.RoutableProvider, cfg *Config) *Server { adminAPI.GET("/audit/conversation", cfg.AdminHandler.AuditConversation) adminAPI.GET("/providers/status", cfg.AdminHandler.ProviderStatus) adminAPI.POST("/runtime/refresh", cfg.AdminHandler.RefreshRuntime) + adminAPI.GET("/budgets/settings", cfg.AdminHandler.BudgetSettings) + adminAPI.PUT("/budgets/settings", cfg.AdminHandler.UpdateBudgetSettings) + adminAPI.POST("/budgets/reset", cfg.AdminHandler.ResetBudgets) adminAPI.GET("/models", cfg.AdminHandler.ListModels) adminAPI.GET("/models/categories", cfg.AdminHandler.ListCategories) adminAPI.GET("/model-overrides", cfg.AdminHandler.ListModelOverrides) diff --git a/internal/server/native_batch_service.go b/internal/server/native_batch_service.go index 033b8c09..bdc85d07 100644 --- a/internal/server/native_batch_service.go +++ b/internal/server/native_batch_service.go @@ -26,6 +26,7 @@ type nativeBatchService struct { cleanupPreparedBatchInputFile func(context.Context, string, string) cleanupStoredBatchRewrittenInputFile func(context.Context, *batchstore.StoredBatch) bool usageLogger usage.LoggerInterface + budgetChecker BudgetChecker pricingResolver usage.PricingResolver orchestrator *gateway.BatchOrchestrator @@ -55,6 +56,9 @@ func (s *nativeBatchService) Batches(c *echo.Context) error { if err != nil { return handleError(c, core.NewInvalidRequestError("invalid request body: "+err.Error(), err)) } + if err := enforceBudget(c, s.budgetChecker); err != nil { + return handleError(c, err) + } ctx, requestID := requestContextWithRequestID(c.Request()) result, err := s.batch().Create(ctx, req, batchRequestMeta(c, requestID)) diff --git a/internal/server/passthrough_service.go b/internal/server/passthrough_service.go index 73449f68..0d2e6fda 100644 --- a/internal/server/passthrough_service.go +++ b/internal/server/passthrough_service.go @@ -13,6 +13,7 @@ type passthroughService struct { modelAuthorizer RequestModelAuthorizer logger auditlog.LoggerInterface usageLogger usage.LoggerInterface + budgetChecker BudgetChecker pricingResolver usage.PricingResolver normalizePassthroughV1Prefix bool enabledPassthroughProviders map[string]struct{} @@ -38,6 +39,9 @@ func (s *passthroughService) ProviderPassthrough(c *echo.Context) error { } } } + if err := enforceBudget(c, s.budgetChecker); err != nil { + return handleError(c, err) + } ctx, _ := requestContextWithRequestID(c.Request()) c.SetRequest(c.Request().WithContext(ctx)) diff --git a/internal/server/translated_inference_service.go b/internal/server/translated_inference_service.go index 63537016..36672f68 100644 --- a/internal/server/translated_inference_service.go +++ b/internal/server/translated_inference_service.go @@ -33,6 +33,7 @@ type translatedInferenceService struct { translatedRequestPatcher TranslatedRequestPatcher logger auditlog.LoggerInterface usageLogger usage.LoggerInterface + budgetChecker BudgetChecker pricingResolver usage.PricingResolver responseCache *responsecache.ResponseCacheMiddleware guardrailsHash string @@ -81,6 +82,10 @@ func (s *translatedInferenceService) dispatchChatCompletion(c *echo.Context, req ctx := c.Request().Context() requestID := requestIDFromContextOrHeader(c.Request()) + if err := enforceBudget(c, s.budgetChecker); err != nil { + return handleError(c, err) + } + if req.Stream { if len(s.inference().FallbackSelectors(workflow)) == 0 { if handled, err := s.tryFastPathStreamingChatPassthrough(c, workflow, req); handled { @@ -221,6 +226,10 @@ func (s *translatedInferenceService) dispatchResponses(c *echo.Context, req *cor ctx := c.Request().Context() requestID := requestIDFromContextOrHeader(c.Request()) + if err := enforceBudget(c, s.budgetChecker); err != nil { + return handleError(c, err) + } + if req.Stream { result, err := s.inference().StreamResponses(ctx, workflow, req) if err != nil { @@ -376,6 +385,10 @@ func (s *translatedInferenceService) Embeddings(c *echo.Context) error { } attachPreparedWorkflow(c, prepared.Context, prepared.Workflow) + if err := enforceBudget(c, s.budgetChecker); err != nil { + return handleError(c, err) + } + requestID := requestIDFromContextOrHeader(c.Request()) result, err := s.inference().ExecuteEmbeddings(c.Request().Context(), prepared.Workflow, prepared.Request, requestID, "/v1/embeddings") if err != nil { From 98eb2047a082a9759d6cf35560f0c97ac52bf175 Mon Sep 17 00:00:00 2001 From: "Jakub A. W" Date: Sun, 26 Apr 2026 18:24:17 +0200 Subject: [PATCH 2/9] feat(budget): expand budget management dashboard --- docs/advanced/admin-endpoints.mdx | 24 +- docs/advanced/config-yaml.mdx | 1 + docs/advanced/configuration.mdx | 26 + docs/advanced/workflows.mdx | 3 + docs/docs.json | 1 + docs/features/budgets.mdx | 203 +++ docs/features/user-path.mdx | 14 +- docs/getting-started/quickstart.mdx | 1 + docs/openapi.json | 1086 ++++++++++++++--- .../admin/dashboard/static/css/dashboard.css | 311 ++++- .../admin/dashboard/static/js/dashboard.js | 21 +- .../dashboard/static/js/modules/budgets.js | 525 ++++++++ .../static/js/modules/budgets.test.js | 130 ++ .../js/modules/dashboard-layout.test.js | 140 ++- .../dashboard/static/js/modules/timezone.js | 13 + .../static/js/modules/timezone.test.js | 11 + .../js/modules/workflows-layout.test.js | 19 +- .../dashboard/static/js/modules/workflows.js | 80 +- .../static/js/modules/workflows.test.js | 100 +- internal/admin/dashboard/templates/index.html | 1 + .../dashboard/templates/page-budgets.html | 233 ++++ .../dashboard/templates/page-settings.html | 100 +- .../dashboard/templates/page-workflows.html | 4 + .../admin/dashboard/templates/sidebar.html | 4 + .../dashboard/templates/workflow-chart.html | 27 +- internal/admin/handler.go | 253 ++++ internal/admin/handler_budgets_test.go | 214 ++++ internal/admin/handler_test.go | 4 + internal/app/app.go | 3 + internal/app/app_test.go | 6 + internal/auditlog/auditlog.go | 2 + internal/auditlog/middleware.go | 6 +- internal/budget/service.go | 99 +- internal/budget/service_test.go | 9 + internal/budget/store.go | 2 + internal/budget/store_mongodb.go | 42 + internal/budget/store_postgresql.go | 43 + internal/budget/store_sqlite.go | 45 + internal/budget/types.go | 1 + internal/core/workflow.go | 8 + internal/server/budget_support.go | 3 + internal/server/budget_support_test.go | 63 + internal/server/error_support.go | 9 +- internal/server/error_support_test.go | 28 + internal/server/http.go | 4 + .../internal_chat_completion_executor.go | 3 + internal/workflows/types.go | 6 + 47 files changed, 3668 insertions(+), 263 deletions(-) create mode 100644 docs/features/budgets.mdx create mode 100644 internal/admin/dashboard/templates/page-budgets.html create mode 100644 internal/admin/handler_budgets_test.go create mode 100644 internal/server/budget_support_test.go diff --git a/docs/advanced/admin-endpoints.mdx b/docs/advanced/admin-endpoints.mdx index d15f7a13..ec74aa17 100644 --- a/docs/advanced/admin-endpoints.mdx +++ b/docs/advanced/admin-endpoints.mdx @@ -1,6 +1,6 @@ --- title: "Admin Endpoints" -description: "Built-in REST API and dashboard for monitoring usage, models, and gateway health." +description: "Built-in REST API and dashboard for monitoring usage, budgets, models, and gateway health." icon: "server-cog" --- @@ -10,7 +10,7 @@ GoModel ships with admin endpoints **enabled by default**. The goal is simple: y The admin layer is split into two independently controllable pieces: -1. **Admin REST API** (`/admin/api/v1/*`) — machine-readable JSON endpoints for usage data and model inventory. Protected by `GOMODEL_MASTER_KEY` like all other API routes. +1. **Admin REST API** (`/admin/api/v1/*`) — machine-readable JSON endpoints for usage data, budgets, and model inventory. Protected by `GOMODEL_MASTER_KEY` like all other API routes. 2. **Admin Dashboard UI** (`/admin/dashboard`) — a lightweight, embedded HTML dashboard that visualizes the same data. No external dependencies, no JavaScript frameworks to install — it's compiled into the binary. Both are on by default because observability shouldn't be opt-in. If you don't need them, turn them off with a single environment variable. @@ -122,6 +122,26 @@ The `date` field in the response changes format based on the interval: `YYYY-MM- Returns an empty array if usage tracking is disabled or no data exists for the period. +### Budget endpoints + +Budgets are managed under `/admin/api/v1/budgets`. These endpoints are +available when budget management is enabled. + +| Method | Path | Description | +| -------- | ------------------------------------- | ------------------------------------ | +| `GET` | `/admin/api/v1/budgets` | List budgets with current status | +| `PUT` | `/admin/api/v1/budgets` | Create or update one budget | +| `DELETE` | `/admin/api/v1/budgets` | Delete one budget | +| `GET` | `/admin/api/v1/budgets/settings` | Read budget reset settings | +| `PUT` | `/admin/api/v1/budgets/settings` | Update budget reset settings | +| `POST` | `/admin/api/v1/budgets/reset-one` | Reset one budget period | +| `POST` | `/admin/api/v1/budgets/reset` | Reset all budget periods | + +`PUT`, `DELETE`, and `reset-one` identify periods by either `period` (`hourly`, +`daily`, `weekly`, `monthly`) or `period_seconds`. + +See [Budgets](/features/budgets) for request examples and enforcement behavior. + ### GET /admin/api/v1/models Returns all registered models with both provider type and configured provider name. diff --git a/docs/advanced/config-yaml.mdx b/docs/advanced/config-yaml.mdx index e39a64c7..4ba6f5f4 100644 --- a/docs/advanced/config-yaml.mdx +++ b/docs/advanced/config-yaml.mdx @@ -15,6 +15,7 @@ especially: - custom provider instance names that do not fit the generated `-` env naming - richer reviewable provider model lists, especially when using allowlist mode +- reviewable budget limits by `user_path` - larger nested config that is easier to review in one file For multiple provider instances, env vars support diff --git a/docs/advanced/configuration.mdx b/docs/advanced/configuration.mdx index 6a954698..713b0c7c 100644 --- a/docs/advanced/configuration.mdx +++ b/docs/advanced/configuration.mdx @@ -112,6 +112,22 @@ Storage is shared by audit logging, usage tracking, and future features like IAM | `USAGE_FLUSH_INTERVAL` | Flush interval in seconds | `5` | | `USAGE_RETENTION_DAYS` | Auto-delete after N days (0 = forever) | `90` | +#### Budgets + +Budgets use tracked usage cost records, so budget limits require +`USAGE_ENABLED=true`. + +| Variable | Description | Default | +| ------------------- | ------------------------------------------------------- | ------- | +| `BUDGETS_ENABLED` | Enable budget management and workflow budget checks | `true` | +| `SET_BUDGET_` | Seed budget limits for a user path, such as `daily=10` | _(empty)_ | + +For example, `SET_BUDGET_TEAM_ALPHA="daily=10,weekly=50"` configures limits +for `/team/alpha`. `SET_BUDGET_="monthly=500"` configures the root path `/`. + +See [Budgets](/features/budgets) for YAML examples, periods, matching, and +workflow enforcement. + #### Metrics @@ -222,6 +238,16 @@ cache: redis: url: "redis://my-redis:6379" +budgets: + enabled: true + user_paths: + - path: "/team/alpha" + limits: + - period: "daily" + amount: 10.00 + - period: "weekly" + amount: 50.00 + providers: openai: type: openai diff --git a/docs/advanced/workflows.mdx b/docs/advanced/workflows.mdx index d3ae13b6..f374252a 100644 --- a/docs/advanced/workflows.mdx +++ b/docs/advanced/workflows.mdx @@ -11,6 +11,7 @@ Workflows are immutable workflow-policy versions stored in the gateway and match They currently control gateway-owned behavior such as: - cache +- budgets - audit logging - usage tracking - guardrails @@ -108,6 +109,7 @@ curl -X POST http://localhost:8080/admin/api/v1/workflows \ "schema_version": 1, "features": { "cache": false, + "budget": true, "audit": true, "usage": true, "guardrails": false, @@ -123,3 +125,4 @@ curl -X POST http://localhost:8080/admin/api/v1/workflows \ - `scope_model` requires `scope_provider_name` - `scope_user_path` is normalized to canonical slash form - managed API keys can override the request `X-GoModel-User-Path`; workflow matching uses the effective request user path +- budget enforcement runs only when the global budget feature and the matched workflow's `budget` feature are both enabled; see [Budgets](/features/budgets) diff --git a/docs/docs.json b/docs/docs.json index 0e000b39..c681c215 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -48,6 +48,7 @@ "features/aliases", "features/user-path", "features/passthrough-api", + "features/budgets", "features/cache", "features/failover" ] diff --git a/docs/features/budgets.mdx b/docs/features/budgets.mdx new file mode 100644 index 00000000..b47e5b66 --- /dev/null +++ b/docs/features/budgets.mdx @@ -0,0 +1,203 @@ +--- +title: "Budgets" +description: "Set spend limits per user path and enforce them through workflow budget controls." +icon: "wallet" +keywords: ["budgets", "spend limits", "cost controls", "user path"] +--- + +## Overview + +Budgets let you set spend limits for a `user_path` subtree. GoModel evaluates +them from tracked usage cost records and blocks matching requests when a limit +has already been spent. + +Use budgets when you want limits such as: + +- `/team/alpha` can spend `$10` per day +- `/team/alpha` can spend `$50` per week +- `/` has a global monthly limit + + + Budget enforcement runs only when budgets are globally enabled and the active + workflow has Budget enabled. The Budget workflow control is enabled by default + when the global budget feature is on. + + +## Enable budgets + +Budgets are enabled by default: + +```env +BUDGETS_ENABLED=true +``` + +Budgets depend on usage tracking because spend is calculated from usage cost +records: + +```env +USAGE_ENABLED=true +``` + +If you configure budget limits while usage tracking is disabled, GoModel rejects +the configuration. + +## Create budgets + +You can create budgets in the dashboard: + +```text +Budgets -> Create Budget +``` + +Dashboard-created budgets are marked as `manual`. Budgets loaded from YAML or +environment variables are marked as `config`. + +You can also seed budgets from YAML: + +```yaml +budgets: + enabled: true + user_paths: + - path: "/team/alpha" + limits: + - period: "daily" + amount: 10.00 + - period: "weekly" + amount: 50.00 + - path: "/" + limits: + - period: "monthly" + amount: 500.00 +``` + +Or use environment variables: + +```env +SET_BUDGET_TEAM_ALPHA="daily=10,weekly=50" +SET_BUDGET_="monthly=500" +``` + +The suffix after `SET_BUDGET_` becomes a slash-separated user path: + +- `SET_BUDGET_TEAM_ALPHA` -> `/team/alpha` +- `SET_BUDGET_` -> `/` + +Supported standard periods are: + +| Period | Seconds | +| --------- | --------- | +| `hourly` | `3600` | +| `daily` | `86400` | +| `weekly` | `604800` | +| `monthly` | `2592000` | + +For custom windows, set `period_seconds` in YAML: + +```yaml +budgets: + user_paths: + - path: "/team/alpha" + limits: + - period_seconds: 7200 + amount: 5.00 +``` + +## How matching works + +Budget paths apply to the configured path and its descendants: + +- budget path `/team` +- request path `/team/app` +- result: budget applies + +Sibling paths do not match: + +- budget path `/team` +- request path `/team-alpha` +- result: budget does not apply + +If multiple budgets match a request, GoModel checks all matching limits. The +request is rejected if any matching limit is exhausted. + +## Workflow enforcement + +The active workflow controls whether budget checks run for a request. + +In a workflow payload: + +```json +{ + "schema_version": 1, + "features": { + "budget": true, + "cache": true, + "audit": true, + "usage": true, + "guardrails": true, + "fallback": true + } +} +``` + +Use scoped workflows to turn budget enforcement on or off for a provider, model, +or `user_path` subtree. See [Workflows](/advanced/workflows) for matching +precedence. + + + Response cache hits can return before budget enforcement. If a cached response + exists, GoModel can serve it without spending additional provider cost. + + +## Reset windows + +Budget reset anchors are configured in the dashboard under: + +```text +Settings -> Budget Resets +``` + +The reset settings are evaluated in UTC. + +For monthly budgets, if the configured day does not exist in a month, the reset +runs on the last day of that month. For example, a monthly reset day of `31` +runs on April 30 and on February 28 or 29. + +Editing a budget changes the limit but does not reset the current period. Use +the row-level Reset action to start a new period for one budget, or +`Settings -> Reset All Budgets` to reset all budgets. + +## Admin API + +Budget management is also available through the admin API: + +```bash +curl -H "Authorization: Bearer $GOMODEL_MASTER_KEY" \ + http://localhost:8080/admin/api/v1/budgets +``` + +Create or update a budget: + +```bash +curl -X PUT http://localhost:8080/admin/api/v1/budgets \ + -H "Authorization: Bearer $GOMODEL_MASTER_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "user_path": "/team/alpha", + "period": "daily", + "amount": 10.00 + }' +``` + +Delete a budget: + +```bash +curl -X DELETE http://localhost:8080/admin/api/v1/budgets \ + -H "Authorization: Bearer $GOMODEL_MASTER_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "user_path": "/team/alpha", + "period": "daily" + }' +``` + +See [Admin Endpoints](/advanced/admin-endpoints) for the endpoint list. diff --git a/docs/features/user-path.mdx b/docs/features/user-path.mdx index 5a91fd18..0d950c46 100644 --- a/docs/features/user-path.mdx +++ b/docs/features/user-path.mdx @@ -9,8 +9,8 @@ icon: "route" `user_path` is a normalized hierarchy for the caller, for example `/team/alpha` or `/team/alpha/service`. -GoModel uses it to keep model access, workflows, usage, and audit data scoped to -the right team, tenant, service, or customer. +GoModel uses it to keep model access, workflows, budgets, usage, and audit data +scoped to the right team, tenant, service, or customer. ## API keys @@ -63,7 +63,7 @@ different model access rules. ## Workflows Workflows can also include `scope_user_path`, so different teams or services can -use different cache, audit, usage, guardrail, and fallback settings. +use different budget, cache, audit, usage, guardrail, and fallback settings. You can combine user path with provider and model scope, for example: @@ -72,3 +72,11 @@ You can combine user path with provider and model scope, for example: - `openai_primary` + `gpt-5` + `/team/alpha` See [Workflows](/advanced/workflows) for the full matching order. + +## Budgets + +Budgets are also scoped by `user_path`. A budget for `/team/alpha` applies to +`/team/alpha` and descendants such as `/team/alpha/service`, but not to sibling +paths such as `/team-alpha`. + +See [Budgets](/features/budgets) for spend limits and workflow enforcement. diff --git a/docs/getting-started/quickstart.mdx b/docs/getting-started/quickstart.mdx index c7f14430..377918de 100644 --- a/docs/getting-started/quickstart.mdx +++ b/docs/getting-started/quickstart.mdx @@ -120,6 +120,7 @@ Use one of those model IDs in your requests. ## Next Steps - Understand response caching: [Cache](/features/cache) +- Add spend limits: [Budgets](/features/budgets) - Configure production settings: [Configuration](/advanced/configuration) - Add request policies: [Guardrails](/advanced/guardrails) - Connect OpenClaw: [Using GoModel with OpenClaw](/guides/openclaw) diff --git a/docs/openapi.json b/docs/openapi.json index 507a5bc5..1879bb8c 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -1,7 +1,7 @@ { "openapi": "3.0.0", "info": { - "description": "High-performance AI gateway routing requests to multiple LLM providers (OpenAI, Anthropic, Gemini, Groq, OpenRouter, Z.ai, xAI, Oracle, Ollama). Drop-in OpenAI-compatible API.", + "description": "High-performance AI gateway routing requests to multiple LLM providers (OpenAI, Anthropic, Gemini, Groq, OpenRouter, Z.ai, xAI, MiniMax, Oracle, Ollama). Drop-in OpenAI-compatible API.", "title": "GoModel API", "contact": {}, "version": "1.0" @@ -197,7 +197,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/auditlog.LogListResult" + "$ref": "#/components/schemas/admin.auditLogListResponse" } } } @@ -230,6 +230,421 @@ ] } }, + "/admin/api/v1/budgets": { + "get": { + "tags": [ + "admin" + ], + "summary": "List budgets with current status", + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.budgetListResponse" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "put": { + "tags": [ + "admin" + ], + "summary": "Create or update one budget", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.upsertBudgetRequest" + } + } + }, + "description": "Budget definition", + "required": true + }, + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.budgetListResponse" + } + } + } + }, + "400": { + "description": "Bad Request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "delete": { + "tags": [ + "admin" + ], + "summary": "Delete one budget", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.deleteBudgetRequest" + } + } + }, + "description": "Budget key", + "required": true + }, + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.budgetListResponse" + } + } + } + }, + "400": { + "description": "Bad Request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/admin/api/v1/budgets/reset": { + "post": { + "tags": [ + "admin" + ], + "summary": "Reset all budget periods", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.resetBudgetsRequest" + } + } + }, + "description": "Reset confirmation", + "required": true + }, + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "type": "object", + "additionalProperties": true + } + } + } + }, + "400": { + "description": "Bad Request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/admin/api/v1/budgets/reset-one": { + "post": { + "tags": [ + "admin" + ], + "summary": "Reset one budget period", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.resetBudgetRequest" + } + } + }, + "description": "Budget key", + "required": true + }, + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.budgetListResponse" + } + } + } + }, + "400": { + "description": "Bad Request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, + "/admin/api/v1/budgets/settings": { + "get": { + "tags": [ + "admin" + ], + "summary": "Get budget reset settings", + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/budget.Settings" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "put": { + "tags": [ + "admin" + ], + "summary": "Update budget reset settings", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/admin.updateBudgetSettingsRequest" + } + } + }, + "description": "Budget reset settings", + "required": true + }, + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/budget.Settings" + } + } + } + }, + "400": { + "description": "Bad Request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + }, + "503": { + "description": "Service Unavailable", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.GatewayError" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + } + }, "/admin/api/v1/cache/overview": { "get": { "tags": [ @@ -3570,7 +3985,99 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/core.ResponsesResponse" + "$ref": "#/components/schemas/core.ResponsesResponse" + } + } + } + }, + "400": { + "description": "Bad Request", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" + } + } + } + }, + "401": { + "description": "Unauthorized", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" + } + } + } + }, + "404": { + "description": "Not Found", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" + } + } + } + }, + "501": { + "description": "Not Implemented", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" + } + } + } + }, + "502": { + "description": "Bad Gateway", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" + } + } + } + } + }, + "security": [ + { + "BearerAuth": [] + } + ] + }, + "delete": { + "tags": [ + "responses" + ], + "summary": "Delete a response", + "parameters": [ + { + "description": "Response ID", + "name": "id", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + { + "description": "Provider override for native deletion", + "name": "provider", + "in": "query", + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/core.ResponseDeleteResponse" } } } @@ -3631,12 +4138,14 @@ "BearerAuth": [] } ] - }, - "delete": { + } + }, + "/v1/responses/{id}/cancel": { + "post": { "tags": [ "responses" ], - "summary": "Delete a response", + "summary": "Cancel a response", "parameters": [ { "description": "Response ID", @@ -3648,7 +4157,7 @@ } }, { - "description": "Provider override for native deletion", + "description": "Provider override for native cancellation", "name": "provider", "in": "query", "schema": { @@ -3662,7 +4171,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/core.ResponseDeleteResponse" + "$ref": "#/components/schemas/core.ResponsesResponse" } } } @@ -3725,12 +4234,12 @@ ] } }, - "/v1/responses/{id}/cancel": { - "post": { + "/v1/responses/{id}/input_items": { + "get": { "tags": [ "responses" ], - "summary": "Cancel a response", + "summary": "List response input items", "parameters": [ { "description": "Response ID", @@ -3742,12 +4251,64 @@ } }, { - "description": "Provider override for native cancellation", + "description": "Provider override for native lookups", "name": "provider", "in": "query", "schema": { "type": "string" } + }, + { + "description": "Pagination cursor", + "name": "after", + "in": "query", + "schema": { + "type": "string" + } + }, + { + "description": "Fields to include in listed input items", + "name": "include", + "in": "query", + "explode": true, + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "description": "Fields to include in listed input items", + "name": "include[]", + "in": "query", + "explode": true, + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + }, + { + "description": "Maximum items to return (1-100, default 20)", + "name": "limit", + "in": "query", + "schema": { + "type": "integer" + } + }, + { + "description": "Sort order: asc or desc", + "name": "order", + "in": "query", + "schema": { + "type": "string", + "enum": [ + "asc", + "desc" + ] + } } ], "responses": { @@ -3756,7 +4317,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/core.ResponsesResponse" + "$ref": "#/components/schemas/core.ResponseInputItemListResponse" } } } @@ -3818,173 +4379,278 @@ } ] } + } + }, + "servers": [ + { + "url": "https://gomodel.example.com", + "description": "GoModel HTTPS deployment" }, - "/v1/responses/{id}/input_items": { - "get": { - "tags": [ - "responses" - ], - "summary": "List response input items", - "parameters": [ - { - "description": "Response ID", - "name": "id", - "in": "path", - "required": true, - "schema": { - "type": "string" + { + "url": "http://localhost:8080", + "description": "Local GoModel" + } + ], + "components": { + "securitySchemes": { + "BearerAuth": { + "type": "apiKey", + "name": "Authorization", + "in": "header" + } + }, + "schemas": { + "admin.auditLogEntryResponse": { + "type": "object", + "properties": { + "alias_used": { + "type": "boolean" + }, + "auth_key_id": { + "type": "string" + }, + "auth_method": { + "type": "string" + }, + "cache_type": { + "type": "string" + }, + "client_ip": { + "type": "string" + }, + "data": { + "description": "Data contains flexible request/response information as JSON", + "allOf": [ + { + "$ref": "#/components/schemas/auditlog.LogData" + } + ] + }, + "duration_ns": { + "description": "DurationNs is the request duration in nanoseconds", + "type": "integer" + }, + "error_type": { + "type": "string" + }, + "id": { + "description": "ID is a unique identifier for this log entry (UUID)", + "type": "string" + }, + "method": { + "type": "string" + }, + "path": { + "type": "string" + }, + "provider": { + "description": "canonical provider type used for routing and filters", + "type": "string" + }, + "provider_name": { + "type": "string" + }, + "request_id": { + "description": "Extracted fields for efficient filtering (indexed in relational DBs)", + "type": "string" + }, + "requested_model": { + "description": "Core fields (indexed for queries)", + "type": "string" + }, + "resolved_model": { + "type": "string" + }, + "status_code": { + "type": "integer" + }, + "stream": { + "type": "boolean" + }, + "timestamp": { + "description": "Timestamp is when the request started", + "type": "string" + }, + "usage": { + "$ref": "#/components/schemas/usage.RequestUsageSummary" + }, + "user_path": { + "type": "string" + }, + "workflow_version_id": { + "type": "string" + } + } + }, + "admin.auditLogListResponse": { + "type": "object", + "properties": { + "entries": { + "type": "array", + "items": { + "$ref": "#/components/schemas/admin.auditLogEntryResponse" + } + }, + "limit": { + "type": "integer" + }, + "offset": { + "type": "integer" + }, + "total": { + "type": "integer" + } + } + }, + "admin.budgetListResponse": { + "type": "object", + "properties": { + "budgets": { + "type": "array", + "items": { + "$ref": "#/components/schemas/admin.budgetStatusResponse" } }, - { - "description": "Provider override for native lookups", - "name": "provider", - "in": "query", - "schema": { - "type": "string" - } + "server_time": { + "type": "string" + } + } + }, + "admin.budgetStatusResponse": { + "type": "object", + "properties": { + "amount": { + "type": "number" + }, + "created_at": { + "type": "string" + }, + "has_usage": { + "type": "boolean" + }, + "last_reset_at": { + "type": "string" + }, + "period_end": { + "type": "string" + }, + "period_label": { + "type": "string" + }, + "period_ratio": { + "type": "number" + }, + "period_seconds": { + "type": "integer" + }, + "period_start": { + "type": "string" + }, + "remaining": { + "type": "number" + }, + "source": { + "type": "string" + }, + "spent": { + "type": "number" + }, + "updated_at": { + "type": "string" + }, + "usage_ratio": { + "type": "number" + }, + "user_path": { + "type": "string" + } + } + }, + "admin.deleteBudgetRequest": { + "type": "object", + "properties": { + "period": { + "type": "string" + }, + "period_seconds": { + "type": "integer" }, - { - "description": "Pagination cursor", - "name": "after", - "in": "query", - "schema": { - "type": "string" - } + "user_path": { + "type": "string" + } + } + }, + "admin.resetBudgetRequest": { + "type": "object", + "properties": { + "period": { + "type": "string" }, - { - "description": "Fields to include in listed input items", - "name": "include", - "in": "query", - "explode": true, - "schema": { - "type": "array", - "items": { - "type": "string" - } - } + "period_seconds": { + "type": "integer" }, - { - "description": "Fields to include in listed input items", - "name": "include[]", - "in": "query", - "explode": true, - "schema": { - "type": "array", - "items": { - "type": "string" - } - } + "user_path": { + "type": "string" + } + } + }, + "admin.resetBudgetsRequest": { + "type": "object", + "properties": { + "confirmation": { + "type": "string" + } + } + }, + "admin.updateBudgetSettingsRequest": { + "type": "object", + "properties": { + "daily_reset_hour": { + "type": "integer" }, - { - "description": "Maximum items to return (1-100, default 20)", - "name": "limit", - "in": "query", - "schema": { - "type": "integer" - } + "daily_reset_minute": { + "type": "integer" }, - { - "description": "Sort order: asc or desc", - "name": "order", - "in": "query", - "schema": { - "type": "string", - "enum": [ - "asc", - "desc" - ] - } - } - ], - "responses": { - "200": { - "description": "OK", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/core.ResponseInputItemListResponse" - } - } - } + "monthly_reset_day": { + "type": "integer" }, - "400": { - "description": "Bad Request", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" - } - } - } + "monthly_reset_hour": { + "type": "integer" }, - "401": { - "description": "Unauthorized", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" - } - } - } + "monthly_reset_minute": { + "type": "integer" }, - "404": { - "description": "Not Found", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" - } - } - } + "weekly_reset_hour": { + "type": "integer" }, - "501": { - "description": "Not Implemented", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" - } - } - } + "weekly_reset_minute": { + "type": "integer" }, - "502": { - "description": "Bad Gateway", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/core.OpenAIErrorEnvelope" - } - } - } + "weekly_reset_weekday": { + "type": "integer" } - }, - "security": [ - { - "BearerAuth": [] + } + }, + "admin.upsertBudgetRequest": { + "type": "object", + "properties": { + "amount": { + "type": "number" + }, + "period": { + "type": "string" + }, + "period_seconds": { + "type": "integer" + }, + "source": { + "type": "string" + }, + "user_path": { + "type": "string" } - ] - } - } - }, - "servers": [ - { - "url": "https://gomodel.example.com", - "description": "GoModel HTTPS deployment" - }, - { - "url": "http://localhost:8080", - "description": "Local GoModel" - } - ], - "components": { - "securitySchemes": { - "BearerAuth": { - "type": "apiKey", - "name": "Authorization", - "in": "header" - } - }, - "schemas": { + } + }, "auditlog.ConversationResult": { "type": "object", "properties": { @@ -4013,6 +4679,9 @@ "api_key_hash": { "type": "string" }, + "error_code": { + "type": "string" + }, "error_message": { "description": "Error details (message can be long, so kept in JSON)", "type": "string" @@ -4149,32 +4818,15 @@ } } }, - "auditlog.LogListResult": { - "type": "object", - "properties": { - "entries": { - "type": "array", - "items": { - "$ref": "#/components/schemas/auditlog.LogEntry" - } - }, - "limit": { - "type": "integer" - }, - "offset": { - "type": "integer" - }, - "total": { - "type": "integer" - } - } - }, "auditlog.WorkflowFeaturesSnapshot": { "type": "object", "properties": { "audit": { "type": "boolean" }, + "budget": { + "type": "boolean" + }, "cache": { "type": "boolean" }, @@ -4189,6 +4841,38 @@ } } }, + "budget.Settings": { + "type": "object", + "properties": { + "daily_reset_hour": { + "type": "integer" + }, + "daily_reset_minute": { + "type": "integer" + }, + "monthly_reset_day": { + "type": "integer" + }, + "monthly_reset_hour": { + "type": "integer" + }, + "monthly_reset_minute": { + "type": "integer" + }, + "updated_at": { + "type": "string" + }, + "weekly_reset_hour": { + "type": "integer" + }, + "weekly_reset_minute": { + "type": "integer" + }, + "weekly_reset_weekday": { + "type": "integer" + } + } + }, "core.BatchError": { "type": "object", "properties": { @@ -5606,6 +6290,38 @@ } } }, + "usage.RequestUsageSummary": { + "type": "object", + "properties": { + "cache_write_input_tokens": { + "type": "integer" + }, + "cached_input_ratio": { + "type": "number" + }, + "cached_input_tokens": { + "type": "integer" + }, + "entries": { + "type": "integer" + }, + "estimated_cached_characters": { + "type": "integer" + }, + "input_tokens": { + "type": "integer" + }, + "output_tokens": { + "type": "integer" + }, + "total_tokens": { + "type": "integer" + }, + "uncached_input_tokens": { + "type": "integer" + } + } + }, "usage.UsageLogEntry": { "type": "object", "properties": { diff --git a/internal/admin/dashboard/static/css/dashboard.css b/internal/admin/dashboard/static/css/dashboard.css index 2beb25ac..d3f70011 100644 --- a/internal/admin/dashboard/static/css/dashboard.css +++ b/internal/admin/dashboard/static/css/dashboard.css @@ -3418,9 +3418,291 @@ body.conversation-drawer-open { width: 100%; } +.budget-help-notice { + margin-bottom: 16px; +} + +.budget-list { + display: grid; + gap: 10px; +} + +.budget-row { + display: grid; + grid-template-columns: minmax(0, 1fr) auto; + align-items: center; + gap: 16px; + padding: 12px 14px; + background: var(--bg-surface); + border: 1px solid var(--border); + border-radius: var(--radius); +} + +.budget-row-main { + min-width: 0; +} + +.budget-row-head { + display: flex; + align-items: center; + gap: 10px; + min-width: 0; +} + +.budget-user-path { + display: inline-flex; + align-items: center; + max-width: 100%; + min-height: 24px; + padding: 2px 8px; + border: 1px solid color-mix(in srgb, var(--accent) 32%, var(--border)); + border-radius: 6px; + background: color-mix(in srgb, var(--accent) 9%, var(--bg)); + color: var(--text); + font-family: "SF Mono", Menlo, Consolas, monospace; + font-size: 12px; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.budget-row-meta { + display: inline-flex; + align-items: center; + gap: 8px; + min-width: 0; + color: var(--text-muted); + font-size: 12px; +} + +.budget-source { + padding: 2px 7px; + border: 1px solid var(--border); + border-radius: 999px; + background: var(--bg); + color: var(--text-muted); + font-size: 11px; +} + +.budget-bars { + display: grid; + gap: 6px; + margin-top: 10px; +} + +.budget-bar-line { + display: grid; + grid-template-columns: 118px minmax(0, 1fr); + align-items: center; + gap: 10px; +} + +.budget-bar-label { + display: flex; + align-items: center; + justify-content: space-between; + gap: 6px; + color: var(--text-muted); + font-size: 11px; + line-height: 1; +} + +.budget-bar-track { + position: relative; + height: 14px; + overflow: hidden; + border-radius: 999px; + background: color-mix(in srgb, var(--border) 75%, var(--bg)); +} + +.budget-bar-fill { + height: 100%; + min-width: 0; + border-radius: inherit; + transition: width 0.2s ease; +} + +.budget-bar-fill-usage { + background: color-mix(in srgb, var(--success) 82%, var(--accent)); +} + +.budget-bar-fill-period { + background: color-mix(in srgb, var(--accent) 75%, var(--warning)); +} + +.budget-bar-fill-danger { + background: var(--danger); +} + +.budget-bar-text-row { + position: absolute; + inset: 0; + z-index: 1; + pointer-events: none; + color: var(--text); +} + +.budget-bar-text-row-on-fill { + color: #fff; + clip-path: inset(0 calc(100% - var(--budget-progress, 0%)) 0 0); +} + +.budget-bar-text { + position: absolute; + top: 50%; + max-width: min(44%, 190px); + overflow: hidden; + font-size: 10px; + font-weight: 700; + line-height: 1; + text-overflow: ellipsis; + text-shadow: 0 1px 1px color-mix(in srgb, var(--bg) 72%, transparent); + white-space: nowrap; +} + +.budget-bar-text-start { + left: 6px; + transform: translateY(-50%); +} + +.budget-bar-text-center { + left: 50%; + max-width: min(46%, 240px); + transform: translate(-50%, -50%); +} + +.budget-bar-text-end { + right: 6px; + max-width: min(34%, 180px); + text-align: right; + transform: translateY(-50%); +} + +.budget-period-label { + display: inline-flex; + align-items: center; + justify-content: center; + gap: 5px; + padding: 2px 7px; + border: 1px solid var(--border); + border-radius: 999px; + color: var(--text); + font-weight: 600; + white-space: nowrap; +} + +.budget-period-label::before { + content: ""; + width: 6px; + height: 6px; + flex: 0 0 6px; + border-radius: 999px; + background: currentColor; +} + +.budget-period-label-hourly { + border-color: color-mix(in srgb, var(--accent) 48%, var(--border)); + background: color-mix(in srgb, var(--accent) 10%, var(--bg)); + color: var(--accent); +} + +.budget-period-label-daily { + border-color: color-mix(in srgb, var(--success) 42%, var(--border)); + background: color-mix(in srgb, var(--success) 9%, var(--bg)); + color: var(--success); +} + +.budget-period-label-weekly { + border-color: color-mix(in srgb, var(--warning) 45%, var(--border)); + background: color-mix(in srgb, var(--warning) 10%, var(--bg)); + color: var(--warning); +} + +.budget-period-label-monthly { + border-color: color-mix(in srgb, var(--accent-hover) 58%, var(--border)); + background: color-mix(in srgb, var(--accent-hover) 16%, var(--bg)); + color: var(--accent-hover); + box-shadow: inset 0 -2px 0 color-mix(in srgb, var(--accent-hover) 32%, transparent); +} + +.budget-period-label-custom { + border-style: dashed; + border-color: color-mix(in srgb, var(--text-muted) 58%, var(--border)); + background: color-mix(in srgb, var(--text-muted) 9%, var(--bg)); + color: var(--text-muted); +} + +.budget-period-label-custom::before { + background: transparent; + border: 1px solid currentColor; +} + +.budget-row-actions { + display: flex; + flex-direction: column; + flex-wrap: wrap; + justify-content: flex-end; + gap: 8px; +} + +.budget-action-btn { + min-width: 82px; + justify-content: center; +} + +.budget-action-btn-warning { + color: var(--warning); + border-color: color-mix(in srgb, var(--warning) 50%, var(--border)); +} + +.budget-action-icon { + width: 14px; + height: 14px; + flex: 0 0 14px; +} + +.budget-editor { + background: color-mix(in srgb, var(--bg-surface) 86%, var(--bg) 14%); +} + .budget-settings-grid { - width: min(100%, 780px); - grid-template-columns: repeat(3, minmax(140px, 1fr)); + display: grid; + gap: 12px; +} + +.budget-settings-row { + display: grid; + grid-template-columns: 96px minmax(170px, 1fr) minmax(110px, 140px) minmax(110px, 140px) minmax(220px, 280px); + align-items: start; + gap: 12px; +} + +.budget-settings-period { + align-self: end; + color: var(--text); + font-size: 12px; + font-weight: 700; + letter-spacing: 0; + min-height: 35px; + padding-bottom: 9px; + text-transform: uppercase; +} + +.budget-settings-spacer { + min-height: 1px; +} + +.budget-settings-help-cell { + align-self: start; + min-width: 0; + min-height: 35px; +} + +.budget-settings-help-cell .inline-help-copy { + margin: 22px 0 0; + max-width: 280px; + font-size: 12px; + line-height: 1.35; } .budget-settings-actions { @@ -3757,6 +4039,25 @@ body.conversation-drawer-open { font-size: 11px; } + .budget-row { + grid-template-columns: 1fr; + } + .budget-row-head { + align-items: flex-start; + flex-direction: column; + } + .budget-row-actions { + justify-content: flex-start; + } + .budget-bar-line { + grid-template-columns: 1fr; + gap: 5px; + } + .budget-period-window { + flex-direction: column; + gap: 4px; + } + .form-grid { grid-template-columns: 1fr; } @@ -3883,6 +4184,12 @@ body.conversation-drawer-open { .budget-settings-grid { grid-template-columns: 1fr; } + .budget-settings-row { + grid-template-columns: 1fr; + } + .budget-settings-spacer { + display: none; + } .budget-settings-actions { width: 100%; } diff --git a/internal/admin/dashboard/static/js/dashboard.js b/internal/admin/dashboard/static/js/dashboard.js index 2dbfe32c..219743a3 100644 --- a/internal/admin/dashboard/static/js/dashboard.js +++ b/internal/admin/dashboard/static/js/dashboard.js @@ -160,6 +160,7 @@ function dashboard() { page = [ "overview", "usage", + "budgets", "models", "workflows", "audit-logs", @@ -192,6 +193,9 @@ function dashboard() { ) { this.fetchGuardrailsPage(); } + if (page === "budgets" && typeof this.fetchBudgetsPage === "function") { + this.fetchBudgetsPage(); + } if (page === "settings") { if (typeof this.ensureTimezoneOptions === "function") { this.ensureTimezoneOptions(); @@ -358,8 +362,9 @@ function dashboard() { (this.aliasFormOpen || this.modelOverrideFormOpen)) || (this.page === "workflows" && this.workflowFormOpen) || (this.page === "guardrails" && this.guardrailFormOpen) || - (this.page === "auth-keys" && this.authKeyFormOpen) - || this.budgetResetDialogOpen + (this.page === "auth-keys" && this.authKeyFormOpen) || + (this.page === "budgets" && this.budgetFormOpen) || + this.budgetResetDialogOpen ); }, @@ -454,6 +459,12 @@ function dashboard() { if (typeof this.fetchWorkflowsPage === "function") { requests.push(this.fetchWorkflowsPage()); } + if ( + this.page === "budgets" && + typeof this.fetchBudgetsPage === "function" + ) { + requests.push(this.fetchBudgetsPage()); + } if ( this.hasCalendarModule && typeof this.fetchCalendarData === "function" @@ -544,6 +555,12 @@ function dashboard() { ) { requests.push(this.fetchGuardrailsPage()); } + if ( + this.page === "budgets" && + typeof this.fetchBudgetsPage === "function" + ) { + requests.push(this.fetchBudgetsPage()); + } if (this.page === "usage" && typeof this.fetchUsagePage === "function") { requests.push(this.fetchUsagePage()); } diff --git a/internal/admin/dashboard/static/js/modules/budgets.js b/internal/admin/dashboard/static/js/modules/budgets.js index 1c3dcbea..e42b2c48 100644 --- a/internal/admin/dashboard/static/js/modules/budgets.js +++ b/internal/admin/dashboard/static/js/modules/budgets.js @@ -18,6 +18,25 @@ budgetResetDialogOpen: false, budgetResetConfirmation: '', budgetResetLoading: false, + budgets: [], + budgetsAvailable: true, + budgetsLoading: false, + budgetFilter: '', + budgetError: '', + budgetNotice: '', + budgetFormOpen: false, + budgetFormSubmitting: false, + budgetFormError: '', + budgetEditing: false, + budgetResettingKey: '', + budgetDeletingKey: '', + budgetForm: { + user_path: '', + period: 'daily', + period_seconds: 86400, + amount: '', + source: 'manual' + }, budgetManagementEnabled() { return typeof this.workflowRuntimeBooleanFlag === 'function' @@ -25,6 +44,509 @@ : true; }, + defaultBudgetForm() { + return { + user_path: '', + period: 'daily', + period_seconds: 86400, + amount: '', + source: 'manual' + }; + }, + + budgetPeriodOptions() { + return [ + { value: 'hourly', label: 'Hourly' }, + { value: 'daily', label: 'Daily' }, + { value: 'weekly', label: 'Weekly' }, + { value: 'monthly', label: 'Monthly' }, + { value: 'custom', label: 'Custom seconds' } + ]; + }, + + budgetPeriodSeconds(period) { + switch (String(period || '').trim().toLowerCase()) { + case 'hourly': + return 3600; + case 'daily': + return 86400; + case 'weekly': + return 604800; + case 'monthly': + return 2592000; + default: + return 0; + } + }, + + budgetPeriodFromSeconds(seconds) { + switch (Number(seconds || 0)) { + case 3600: + return 'hourly'; + case 86400: + return 'daily'; + case 604800: + return 'weekly'; + case 2592000: + return 'monthly'; + default: + return 'custom'; + } + }, + + syncBudgetPeriodSeconds() { + const period = String(this.budgetForm && this.budgetForm.period || '').trim(); + const seconds = this.budgetPeriodSeconds(period); + if (seconds > 0) { + this.budgetForm.period_seconds = seconds; + } + }, + + budgetKey(item) { + return String(item && item.user_path || '') + ':' + String(item && item.period_seconds || ''); + }, + + budgetUserPathValidationError(value) { + const trimmed = String(value || '').trim(); + if (!trimmed) { + return 'User path is required.'; + } + const raw = trimmed.startsWith('/') ? trimmed : '/' + trimmed; + const segments = raw.split('/'); + for (const part of segments) { + const segment = String(part || '').trim(); + if (!segment) { + continue; + } + if (segment === '.' || segment === '..') { + return 'User path cannot contain "." or ".." segments.'; + } + if (segment.includes(':')) { + return 'User path cannot contain ":" segments.'; + } + } + return ''; + }, + + normalizeBudgetUserPath(value) { + if (this.budgetUserPathValidationError(value)) { + return ''; + } + const trimmed = String(value || '').trim(); + const raw = trimmed.startsWith('/') ? trimmed : '/' + trimmed; + const segments = raw.split('/'); + const canonical = []; + for (const part of segments) { + const segment = String(part || '').trim(); + if (segment) { + canonical.push(segment); + } + } + return canonical.length ? '/' + canonical.join('/') : '/'; + }, + + normalizeBudgetListPayload(payload) { + if (Array.isArray(payload)) { + return payload; + } + if (payload && Array.isArray(payload.budgets)) { + return payload.budgets; + } + return []; + }, + + filteredBudgets() { + const filter = String(this.budgetFilter || '').trim().toLowerCase(); + if (!filter) { + return this.budgets; + } + return this.budgets.filter((item) => { + const userPath = String(item && item.user_path || '').toLowerCase(); + return userPath.includes(filter); + }); + }, + + async budgetResponseMessage(res, fallback) { + try { + const payload = await res.json(); + if (payload && payload.error && payload.error.message) { + return payload.error.message; + } + } catch (_) { + // Ignore invalid or empty responses and return the fallback message. + } + return fallback; + }, + + async fetchBudgetsPage() { + if (!this.budgetManagementEnabled()) { + this.budgets = []; + this.budgetsAvailable = false; + this.budgetError = ''; + return; + } + await this.fetchBudgets(); + }, + + async fetchBudgets() { + this.budgetsLoading = true; + this.budgetError = ''; + try { + const request = typeof this.requestOptions === 'function' ? this.requestOptions() : { headers: this.headers() }; + const res = await fetch('/admin/api/v1/budgets', request); + if (res.status === 503) { + this.budgetsAvailable = false; + this.budgets = []; + return; + } + const handled = this.handleFetchResponse(res, 'budgets', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + this.budgetsAvailable = true; + if (!handled) { + this.budgetError = 'Unable to load budgets.'; + return; + } + this.budgets = this.normalizeBudgetListPayload(await res.json()); + if (typeof this.renderIconsAfterUpdate === 'function') { + this.renderIconsAfterUpdate(); + } + } catch (e) { + console.error('Failed to fetch budgets:', e); + this.budgets = []; + this.budgetError = 'Unable to load budgets.'; + } finally { + this.budgetsLoading = false; + } + }, + + openBudgetForm(item) { + this.budgetEditing = !!item; + this.budgetFormError = ''; + this.budgetError = ''; + this.budgetNotice = ''; + if (item) { + const periodSeconds = Number(item.period_seconds || 0); + this.budgetForm = { + user_path: String(item.user_path || ''), + period: this.budgetPeriodFromSeconds(periodSeconds), + period_seconds: periodSeconds, + amount: String(item.amount || ''), + source: String(item.source || 'manual') + }; + } else { + this.budgetForm = this.defaultBudgetForm(); + } + this.budgetFormOpen = true; + if (typeof this.renderIconsAfterUpdate === 'function') { + this.renderIconsAfterUpdate(); + } + if (typeof this.$nextTick === 'function') { + this.$nextTick(() => { + const refs = this.$refs || {}; + const input = this.budgetEditing ? refs.budgetAmountInput : refs.budgetUserPathInput; + if (input && typeof input.focus === 'function') { + input.focus({ preventScroll: true }); + } + }); + } + }, + + closeBudgetForm() { + this.budgetFormOpen = false; + this.budgetFormSubmitting = false; + this.budgetFormError = ''; + this.budgetEditing = false; + this.budgetForm = this.defaultBudgetForm(); + }, + + budgetFormPayload() { + const userPathError = this.budgetUserPathValidationError(this.budgetForm.user_path); + if (userPathError) { + this.budgetFormError = userPathError; + return null; + } + const amount = Number(this.budgetForm.amount); + if (!Number.isFinite(amount) || amount <= 0) { + this.budgetFormError = 'Amount must be greater than 0.'; + return null; + } + const period = String(this.budgetForm.period || '').trim(); + let periodSeconds = this.budgetPeriodSeconds(period); + if (period === 'custom') { + periodSeconds = Number(this.budgetForm.period_seconds); + } + if (!Number.isFinite(periodSeconds) || periodSeconds <= 0) { + this.budgetFormError = 'Period seconds must be greater than 0.'; + return null; + } + return { + user_path: this.normalizeBudgetUserPath(this.budgetForm.user_path), + period_seconds: Math.trunc(periodSeconds), + amount, + source: String(this.budgetForm.source || 'manual').trim() || 'manual' + }; + }, + + async submitBudgetForm() { + if (this.budgetFormSubmitting) { + return; + } + const payload = this.budgetFormPayload(); + if (!payload) { + return; + } + this.budgetFormSubmitting = true; + this.budgetFormError = ''; + this.budgetError = ''; + this.budgetNotice = ''; + try { + const request = typeof this.requestOptions === 'function' + ? this.requestOptions({ + method: 'PUT', + body: JSON.stringify(payload) + }) + : { + method: 'PUT', + headers: this.headers(), + body: JSON.stringify(payload) + }; + const res = await fetch('/admin/api/v1/budgets', request); + if (res.status === 503) { + this.budgetsAvailable = false; + this.budgetFormError = 'Budget management is unavailable.'; + return; + } + const handled = this.handleFetchResponse(res, 'budget', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + if (!handled) { + this.budgetFormError = await this.budgetResponseMessage(res, 'Unable to save budget.'); + return; + } + this.closeBudgetForm(); + await this.fetchBudgets(); + this.budgetNotice = 'Budget saved.'; + } catch (e) { + console.error('Failed to save budget:', e); + this.budgetFormError = 'Unable to save budget.'; + } finally { + this.budgetFormSubmitting = false; + } + }, + + async resetBudget(item) { + if (!item) { + return; + } + const key = this.budgetKey(item); + if (this.budgetResettingKey === key) { + return; + } + const label = String(item.user_path || '') + ' ' + this.budgetPeriodLabel(item); + if (global.confirm && !global.confirm('Reset budget "' + label + '"?')) { + return; + } + this.budgetResettingKey = key; + this.budgetError = ''; + this.budgetNotice = ''; + try { + const request = typeof this.requestOptions === 'function' + ? this.requestOptions({ + method: 'POST', + body: JSON.stringify({ + user_path: item.user_path, + period_seconds: item.period_seconds + }) + }) + : { + method: 'POST', + headers: this.headers(), + body: JSON.stringify({ + user_path: item.user_path, + period_seconds: item.period_seconds + }) + }; + const res = await fetch('/admin/api/v1/budgets/reset-one', request); + if (res.status === 503) { + this.budgetsAvailable = false; + this.budgetError = 'Budget management is unavailable.'; + return; + } + const handled = this.handleFetchResponse(res, 'budget reset', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + if (!handled) { + this.budgetError = await this.budgetResponseMessage(res, 'Unable to reset budget.'); + return; + } + await this.fetchBudgets(); + this.budgetNotice = 'Budget reset.'; + } catch (e) { + console.error('Failed to reset budget:', e); + this.budgetError = 'Unable to reset budget.'; + } finally { + this.budgetResettingKey = ''; + } + }, + + async deleteBudget(item) { + if (!item) { + return; + } + const key = this.budgetKey(item); + if (this.budgetDeletingKey === key) { + return; + } + const label = String(item.user_path || '') + ' ' + this.budgetPeriodLabel(item); + if (global.confirm && !global.confirm('Delete budget "' + label + '"? This cannot be undone.')) { + return; + } + this.budgetDeletingKey = key; + this.budgetError = ''; + this.budgetNotice = ''; + try { + const request = typeof this.requestOptions === 'function' + ? this.requestOptions({ + method: 'DELETE', + body: JSON.stringify({ + user_path: item.user_path, + period_seconds: item.period_seconds + }) + }) + : { + method: 'DELETE', + headers: this.headers(), + body: JSON.stringify({ + user_path: item.user_path, + period_seconds: item.period_seconds + }) + }; + const res = await fetch('/admin/api/v1/budgets', request); + if (res.status === 503) { + this.budgetsAvailable = false; + this.budgetError = 'Budget management is unavailable.'; + return; + } + const handled = this.handleFetchResponse(res, 'budget delete', request); + if (typeof this.isStaleAuthFetchResult === 'function' && this.isStaleAuthFetchResult(handled)) { + return; + } + if (!handled) { + this.budgetError = await this.budgetResponseMessage(res, 'Unable to delete budget.'); + return; + } + this.budgets = this.normalizeBudgetListPayload(await res.json()); + if (typeof this.renderIconsAfterUpdate === 'function') { + this.renderIconsAfterUpdate(); + } + this.budgetNotice = 'Budget deleted.'; + } catch (e) { + console.error('Failed to delete budget:', e); + this.budgetError = 'Unable to delete budget.'; + } finally { + this.budgetDeletingKey = ''; + } + }, + + budgetRatio(value) { + const parsed = Number(value); + if (!Number.isFinite(parsed) || parsed < 0) { + return 0; + } + return parsed; + }, + + budgetPercentFromRatio(value, clamp) { + const ratio = this.budgetRatio(value); + const bounded = clamp ? Math.min(ratio, 1) : ratio; + return Math.round(bounded * 1000) / 10; + }, + + budgetUsageRatio(item) { + return this.budgetRatio(item && item.usage_ratio); + }, + + budgetUsagePercent(item) { + return this.budgetPercentFromRatio(this.budgetUsageRatio(item), true); + }, + + budgetPeriodPercent(item) { + return this.budgetPercentFromRatio(item && item.period_ratio, true); + }, + + budgetUsagePercentLabel(item) { + return this.budgetPercentFromRatio(this.budgetUsageRatio(item), false).toFixed(1).replace(/\.0$/, '') + '%'; + }, + + budgetPeriodPercentLabel(item) { + return this.budgetPeriodPercent(item).toFixed(1).replace(/\.0$/, '') + '%'; + }, + + budgetPeriodLabel(item) { + const seconds = Number(item && item.period_seconds || 0); + switch (seconds) { + case 3600: + return 'Hourly'; + case 86400: + return 'Daily'; + case 604800: + return 'Weekly'; + case 2592000: + return 'Monthly'; + default: { + const label = String(item && item.period_label || '').trim(); + return label ? 'Custom ' + label : 'Custom ' + String(seconds || '') + 's'; + } + } + }, + + budgetPeriodClass(item) { + const seconds = Number(item && item.period_seconds || 0); + switch (seconds) { + case 3600: + return 'budget-period-label-hourly'; + case 86400: + return 'budget-period-label-daily'; + case 604800: + return 'budget-period-label-weekly'; + case 2592000: + return 'budget-period-label-monthly'; + default: + return 'budget-period-label-custom'; + } + }, + + budgetSourceLabel(item) { + const source = String(item && item.source || '').trim(); + return source || 'manual'; + }, + + budgetSourceTitle(item) { + const source = this.budgetSourceLabel(item).toLowerCase(); + if (source === 'manual') { + return 'Created from the dashboard.'; + } + if (source === 'config') { + return 'Loaded from configuration.'; + } + return 'Budget source: ' + source; + }, + + budgetRemainingLabel(item) { + const remaining = Number(item && item.remaining); + if (!Number.isFinite(remaining)) { + return ''; + } + if (remaining < 0) { + return this.formatCost(Math.abs(remaining)) + ' over'; + } + return this.formatCost(remaining) + ' remaining'; + }, + budgetWeekdays() { return [ { value: 0, label: 'Sunday' }, @@ -160,6 +682,9 @@ return; } this.closeBudgetResetDialog(); + if (this.page === 'budgets' && typeof this.fetchBudgets === 'function') { + await this.fetchBudgets(); + } this.budgetSettingsNotice = 'Budgets reset.'; } catch (e) { console.error('Failed to reset budgets:', e); diff --git a/internal/admin/dashboard/static/js/modules/budgets.test.js b/internal/admin/dashboard/static/js/modules/budgets.test.js index 69df1f44..9e40fad1 100644 --- a/internal/admin/dashboard/static/js/modules/budgets.test.js +++ b/internal/admin/dashboard/static/js/modules/budgets.test.js @@ -63,6 +63,136 @@ test('budgetSettingsPayload normalizes numeric values before saving', () => { })); }); +test('budgetFormPayload normalizes user path and standard periods', () => { + const module = createBudgetsModule(); + module.budgetForm = { + user_path: 'team/alpha', + period: 'weekly', + period_seconds: 0, + amount: '12.3456', + source: '' + }; + + assert.equal(JSON.stringify(module.budgetFormPayload()), JSON.stringify({ + user_path: '/team/alpha', + period_seconds: 604800, + amount: 12.3456, + source: 'manual' + })); +}); + +test('fetchBudgets reads budget rows from the list envelope', async () => { + let renderIconsCalls = 0; + const module = createBudgetsModule({ + fetch(url) { + assert.equal(url, '/admin/api/v1/budgets'); + return Promise.resolve({ + status: 200, + ok: true, + json: () => Promise.resolve({ + budgets: [ + { user_path: '/team', period_seconds: 86400, amount: 10 } + ] + }) + }); + } + }); + module.requestOptions = () => ({}); + module.handleFetchResponse = () => true; + module.renderIconsAfterUpdate = () => { + renderIconsCalls++; + }; + + await module.fetchBudgets(); + + assert.equal(module.budgetsAvailable, true); + assert.equal(module.budgets.length, 1); + assert.equal(module.budgets[0].user_path, '/team'); + assert.equal(renderIconsCalls, 1); +}); + +test('filteredBudgets filters only by user path', () => { + const module = createBudgetsModule(); + module.budgets = [ + { user_path: '/team/alpha', period_label: 'daily' }, + { user_path: '/team/beta', period_label: 'weekly' }, + { user_path: '/platform', period_label: 'team' } + ]; + + module.budgetFilter = 'TEAM/A'; + assert.equal(JSON.stringify(module.filteredBudgets()), JSON.stringify([ + { user_path: '/team/alpha', period_label: 'daily' } + ])); + + module.budgetFilter = 'weekly'; + assert.equal(JSON.stringify(module.filteredBudgets()), JSON.stringify([])); + + module.budgetFilter = ''; + assert.equal(module.filteredBudgets(), module.budgets); +}); + +test('budgetSourceTitle explains manual and config sources', () => { + const module = createBudgetsModule(); + + assert.equal(module.budgetSourceTitle({ source: 'manual' }), 'Created from the dashboard.'); + assert.equal(module.budgetSourceTitle({ source: 'config' }), 'Loaded from configuration.'); + assert.equal(module.budgetSourceTitle({ source: 'import' }), 'Budget source: import'); +}); + +test('budgetPeriodLabel and class distinguish standard and custom periods', () => { + const module = createBudgetsModule(); + + assert.equal(module.budgetPeriodLabel({ period_seconds: 3600 }), 'Hourly'); + assert.equal(module.budgetPeriodClass({ period_seconds: 3600 }), 'budget-period-label-hourly'); + assert.equal(module.budgetPeriodLabel({ period_seconds: 86400 }), 'Daily'); + assert.equal(module.budgetPeriodClass({ period_seconds: 86400 }), 'budget-period-label-daily'); + assert.equal(module.budgetPeriodLabel({ period_seconds: 604800 }), 'Weekly'); + assert.equal(module.budgetPeriodClass({ period_seconds: 604800 }), 'budget-period-label-weekly'); + assert.equal(module.budgetPeriodLabel({ period_seconds: 2592000 }), 'Monthly'); + assert.equal(module.budgetPeriodClass({ period_seconds: 2592000 }), 'budget-period-label-monthly'); + assert.equal(module.budgetPeriodLabel({ period_seconds: 7200, period_label: '7200s' }), 'Custom 7200s'); + assert.equal(module.budgetPeriodClass({ period_seconds: 7200 }), 'budget-period-label-custom'); +}); + +test('deleteBudget posts the selected budget key and refreshes from the response envelope', async () => { + const requests = []; + const module = createBudgetsModule({ + confirm(message) { + assert.match(message, /Delete budget "\/team Daily"\? This cannot be undone\./); + return true; + }, + fetch(url, request) { + requests.push({ url, request }); + return Promise.resolve({ + status: 200, + ok: true, + json: () => Promise.resolve({ + budgets: [ + { user_path: '/other', period_seconds: 86400, amount: 5 } + ] + }) + }); + } + }); + module.requestOptions = (options) => options || {}; + module.handleFetchResponse = () => true; + + await module.deleteBudget({ user_path: '/team', period_seconds: 86400, period_label: 'daily' }); + + assert.equal(requests.length, 1); + assert.equal(requests[0].url, '/admin/api/v1/budgets'); + assert.equal(requests[0].request.method, 'DELETE'); + assert.equal(requests[0].request.body, JSON.stringify({ + user_path: '/team', + period_seconds: 86400 + })); + assert.equal(module.budgetDeletingKey, ''); + assert.equal(module.budgetNotice, 'Budget deleted.'); + assert.equal(JSON.stringify(module.budgets), JSON.stringify([ + { user_path: '/other', period_seconds: 86400, amount: 5 } + ])); +}); + test('resetBudgets requires the typed reset confirmation before posting', async () => { const module = createBudgetsModule({ fetch() { diff --git a/internal/admin/dashboard/static/js/modules/dashboard-layout.test.js b/internal/admin/dashboard/static/js/modules/dashboard-layout.test.js index 8b2b6021..acd5496e 100644 --- a/internal/admin/dashboard/static/js/modules/dashboard-layout.test.js +++ b/internal/admin/dashboard/static/js/modules/dashboard-layout.test.js @@ -19,6 +19,7 @@ function readDashboardTemplateSource() { readFixture("../../../templates/index.html"), readFixture("../../../templates/page-overview.html"), readFixture("../../../templates/page-usage.html"), + readFixture("../../../templates/page-budgets.html"), readFixture("../../../templates/page-settings.html"), readFixture("../../../templates/page-guardrails.html"), readFixture("../../../templates/page-auth-keys.html"), @@ -62,7 +63,7 @@ test("sidebar and main content share the flex layout without manual content offs assert.doesNotMatch(template, /content-collapsed/); assert.match( template, - /href="\/admin\/dashboard\/overview"[\s\S]*Overview<\/span>[\s\S]*href="\/admin\/dashboard\/models"[\s\S]*Models<\/span>[\s\S]*href="\/admin\/dashboard\/audit-logs"[\s\S]*Audit Logs<\/span>[\s\S]*href="\/admin\/dashboard\/usage"[\s\S]*Usage<\/span>[\s\S]*href="\/admin\/dashboard\/auth-keys"[\s\S]*API Keys<\/span>[\s\S]*href="\/admin\/dashboard\/workflows"[\s\S]*Workflows<\/span>[\s\S]*href="\/admin\/dashboard\/guardrails"[\s\S]*x-show="guardrailsPageVisible\(\)"[\s\S]*Guardrails \(experimental\)<\/span>[\s\S]*href="\/admin\/dashboard\/settings"[\s\S]*Settings<\/span>/, + /href="\/admin\/dashboard\/overview"[\s\S]*Overview<\/span>[\s\S]*href="\/admin\/dashboard\/models"[\s\S]*Models<\/span>[\s\S]*href="\/admin\/dashboard\/audit-logs"[\s\S]*Audit Logs<\/span>[\s\S]*href="\/admin\/dashboard\/usage"[\s\S]*Usage<\/span>[\s\S]*href="\/admin\/dashboard\/budgets"[\s\S]*x-show="budgetManagementEnabled\(\)"[\s\S]*Budgets<\/span>[\s\S]*href="\/admin\/dashboard\/auth-keys"[\s\S]*API Keys<\/span>[\s\S]*href="\/admin\/dashboard\/workflows"[\s\S]*Workflows<\/span>[\s\S]*href="\/admin\/dashboard\/guardrails"[\s\S]*x-show="guardrailsPageVisible\(\)"[\s\S]*Guardrails \(experimental\)<\/span>[\s\S]*href="\/admin\/dashboard\/settings"[\s\S]*Settings<\/span>/, ); const sidebarRule = readCSSRule(css, ".sidebar"); @@ -178,6 +179,10 @@ test("dashboard chrome uses Lucide icons for stable navigation and auth controls template, /data-lucide="chart-column" class="nav-icon"[\s\S]*Usage<\/span>/, ); + assert.match( + template, + /data-lucide="wallet" class="nav-icon"[\s\S]*Budgets<\/span>/, + ); assert.match( template, /data-lucide="key-round" class="nav-icon"[\s\S]*API Keys<\/span>/, @@ -300,7 +305,7 @@ test("dashboard pages reuse a shared auth banner template", () => { const authBannerCalls = indexTemplate.match(/{{template "auth-banner" \.}}/g) || []; - assert.equal(authBannerCalls.length, 8); + assert.equal(authBannerCalls.length, 9); assert.match( indexTemplate, /