From 84bf37ed767cc1b7743cedf51c793800a392ed65 Mon Sep 17 00:00:00 2001 From: Max Beizer Date: Mon, 9 Mar 2026 15:33:45 -0500 Subject: [PATCH] feat: add PRU-aware budget config with Copilot plan presets - Add MonthlyLimit and CopilotPlan fields to ProviderConfig - Add CopilotPlanLimits map: free=50, pro=300, pro_plus=1500, business=300, enterprise=1000 PRUs/month - Add GetCopilotMonthlyLimit() with resolution order: MonthlyLimit > CopilotPlan preset > weeklyTokens*4 fallback - Update budget.Manager to use GetCopilotMonthlyLimit() instead of hardcoded weeklyBudget*4 approximation - Update doctor command to use config-driven monthly limit - Table-driven tests for all plan presets and fallback behavior Config example: [providers.copilot] copilot_plan = "pro" # auto-sets 300 PRUs/month # or: monthly_limit = 500 # explicit override Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cmd/nightshift/commands/doctor.go | 2 +- internal/budget/budget.go | 6 +- internal/config/config.go | 39 +++++++++++++ internal/config/config_test.go | 93 +++++++++++++++++++++++++++++++ 4 files changed, 135 insertions(+), 5 deletions(-) diff --git a/cmd/nightshift/commands/doctor.go b/cmd/nightshift/commands/doctor.go index 1a0f0ee..2c41869 100644 --- a/cmd/nightshift/commands/doctor.go +++ b/cmd/nightshift/commands/doctor.go @@ -246,7 +246,7 @@ func checkProviders(cfg *config.Config, add func(string, checkStatus, string)) ( add("copilot.data_path", statusOK, path) } copilotProvider = providers.NewCopilotWithPath(path) - monthlyLimit := int64(cfg.GetProviderBudget("copilot")) + monthlyLimit := cfg.GetCopilotMonthlyLimit() if pct, err := copilotProvider.GetUsedPercent(mode, monthlyLimit); err != nil { add("copilot.usage", statusFail, err.Error()) } else { diff --git a/internal/budget/budget.go b/internal/budget/budget.go index 31a718c..207bd4b 100644 --- a/internal/budget/budget.go +++ b/internal/budget/budget.go @@ -312,10 +312,8 @@ func (m *Manager) GetUsedPercent(provider string) (float64, error) { if m.copilot == nil { return 0, fmt.Errorf("copilot provider not configured") } - // Copilot uses monthly request limits, not weekly token budgets - // Convert weekly budget to monthly limit for consistency - // Note: This is a simplification; actual monthly limits should be configured separately - monthlyLimit := weeklyBudget * 4 // Approximate: 4 weeks per month + // Use config-driven monthly limit (plan preset, explicit, or fallback) + monthlyLimit := m.cfg.GetCopilotMonthlyLimit() return m.copilot.GetUsedPercent(mode, monthlyLimit) default: diff --git a/internal/config/config.go b/internal/config/config.go index f0d7347..92495f0 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -76,6 +76,17 @@ type ProviderConfig struct { DangerouslySkipPermissions bool `mapstructure:"dangerously_skip_permissions"` // DangerouslyBypassApprovalsAndSandbox tells the CLI to bypass approvals and sandboxing. DangerouslyBypassApprovalsAndSandbox bool `mapstructure:"dangerously_bypass_approvals_and_sandbox"` + + // MonthlyLimit sets the monthly premium request (PRU) limit for Copilot. + // If zero, falls back to CopilotPlan preset or weeklyTokens * 4. + // Only meaningful for the copilot provider. + MonthlyLimit int64 `mapstructure:"monthly_limit"` + + // CopilotPlan sets the GitHub Copilot plan for automatic PRU limit detection. + // Valid values: "free" (50), "pro" (300), "pro_plus" (1500), + // "business" (300), "enterprise" (1000). + // If empty, uses MonthlyLimit or falls back to weeklyTokens * 4. + CopilotPlan string `mapstructure:"copilot_plan"` } // ProjectConfig defines a project to manage. @@ -498,6 +509,34 @@ func (c *Config) GetProviderBudget(provider string) int { return c.Budget.WeeklyTokens } +// CopilotPlanLimits maps Copilot plan names to monthly PRU limits. +var CopilotPlanLimits = map[string]int64{ + "free": 50, + "pro": 300, + "pro_plus": 1500, + "business": 300, + "enterprise": 1000, +} + +// GetCopilotMonthlyLimit returns the monthly PRU limit for the Copilot provider. +// Resolution order: MonthlyLimit field > CopilotPlan preset > weeklyTokens * 4 fallback. +func (c *Config) GetCopilotMonthlyLimit() int64 { + // Explicit monthly limit takes priority + if c.Providers.Copilot.MonthlyLimit > 0 { + return c.Providers.Copilot.MonthlyLimit + } + + // Plan-based preset + if c.Providers.Copilot.CopilotPlan != "" { + if limit, ok := CopilotPlanLimits[c.Providers.Copilot.CopilotPlan]; ok { + return limit + } + } + + // Fallback: derive from weekly token budget + return int64(c.GetProviderBudget("copilot")) * 4 +} + // IsTaskEnabled checks if a task type is enabled. func (c *Config) IsTaskEnabled(task string) bool { // Check if explicitly disabled diff --git a/internal/config/config_test.go b/internal/config/config_test.go index df34cdc..73f3386 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -154,6 +154,99 @@ func TestGetProviderBudget(t *testing.T) { } } +func TestGetCopilotMonthlyLimit(t *testing.T) { + tests := []struct { + name string + cfg Config + expected int64 + }{ + { + name: "explicit monthly limit takes priority", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{ + MonthlyLimit: 500, + CopilotPlan: "pro", // ignored when MonthlyLimit is set + }, + }, + Budget: BudgetConfig{WeeklyTokens: 100}, + }, + expected: 500, + }, + { + name: "plan preset: free", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{CopilotPlan: "free"}, + }, + }, + expected: 50, + }, + { + name: "plan preset: pro", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{CopilotPlan: "pro"}, + }, + }, + expected: 300, + }, + { + name: "plan preset: pro_plus", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{CopilotPlan: "pro_plus"}, + }, + }, + expected: 1500, + }, + { + name: "plan preset: business", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{CopilotPlan: "business"}, + }, + }, + expected: 300, + }, + { + name: "plan preset: enterprise", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{CopilotPlan: "enterprise"}, + }, + }, + expected: 1000, + }, + { + name: "fallback to weekly tokens * 4", + cfg: Config{ + Budget: BudgetConfig{WeeklyTokens: 200}, + }, + expected: 800, + }, + { + name: "unknown plan falls back to weekly tokens * 4", + cfg: Config{ + Providers: ProvidersConfig{ + Copilot: ProviderConfig{CopilotPlan: "unknown_plan"}, + }, + Budget: BudgetConfig{WeeklyTokens: 150}, + }, + expected: 600, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := tt.cfg.GetCopilotMonthlyLimit() + if got != tt.expected { + t.Errorf("GetCopilotMonthlyLimit() = %d, want %d", got, tt.expected) + } + }) + } +} + func TestNormalizeBudgetConfig(t *testing.T) { cfg := &Config{ Budget: BudgetConfig{