Merge pull request #46 from manaflow-ai/fix/status-refreshes-fable-scheduler

Update scheduler from fresh usage status

Author: lawrencecchen

internal/proxy/claude_ratelimit_routing_test.go

@@ -458,6 +458,41 @@ func TestClaudeFableProbeTreatsHeaderless429AsFableExhausted(t *testing.T) {
 	}
 }
 
+func TestUsageStatusFreshFableWindowUpdatesSchedulerModelPool(t *testing.T) {
+	schedulerRef := selectacct.NewSchedulerRef(selectacct.NewScheduler([]selectacct.Score{
+		{AccountID: "acct", Provider: accounts.ProviderClaude, Headroom: 1, ShortHeadroom: 1},
+	}))
+	server := Server{
+		Accounts: []accounts.Account{{
+			ID:       "acct",
+			Provider: accounts.ProviderClaude,
+			AuthMode: accounts.AuthModeOAuth,
+		}},
+		SchedulerRef: schedulerRef,
+	}
+	server.updateSchedulerFromUsageStatuses([]AccountUsageStatus{{
+		AccountStatus: AccountStatus{
+			ID:       "acct",
+			Provider: accounts.ProviderClaude,
+			AuthMode: accounts.AuthModeOAuth,
+		},
+		UsageFresh: true,
+		Windows: []accounts.UsageWindow{
+			{Name: "5h", UsedPercent: 10, LimitWindowSeconds: 5 * 60 * 60},
+			{Name: "7d", UsedPercent: 40, LimitWindowSeconds: 7 * 24 * 60 * 60},
+			{Name: agentclaude.FableWindowName, Feature: agentclaude.FableModel, UsedPercent: 100, LimitWindowSeconds: 7 * 24 * 60 * 60},
+		},
+	}})
+	base := schedulerRef.Get()
+	if base.Exhausted(accounts.ProviderClaude, "acct") {
+		t.Fatalf("base Claude score should remain usable when only Fable is exhausted")
+	}
+	fable := base.ForModel(agentclaude.FableModel)
+	if !fable.Exhausted(accounts.ProviderClaude, "acct") {
+		t.Fatalf("Fable model score should be exhausted after fresh status update")
+	}
+}
+
 // TestClaudeShortWindowDrivesRouting proves the routing calculation now prefers
 // the account with more 5h headroom even when both accounts have identical
 // account-wide (7d) headroom. Before the fix both scored equal short headroom.

internal/proxy/proxy.go

@@ -277,6 +277,7 @@ type AccountUsageStatus struct {
 	Windows            []accounts.UsageWindow           `json:"windows,omitempty"`
 	Credits            *accounts.CreditsInfo            `json:"credits,omitempty"`
 	ComplimentaryReset *accounts.ComplimentaryResetInfo `json:"complimentary_reset,omitempty"`
+	UsageFresh         bool                             `json:"-"`
 }
 
 func NewAccountRef(store accounts.CodexStore, initial []accounts.Account, client *http.Client) *AccountRef {
@@ -488,6 +489,7 @@ func (r *AccountRef) UsageStatuses(ctx context.Context) []AccountUsageStatus {
 		restored.Active = status.Active
 		restored.Refreshed = status.Refreshed
 		restored.Error = ""
+		restored.UsageFresh = false
 		out[i] = restored
 	}
 	r.usageStatusCache = append([]AccountUsageStatus(nil), out...)
@@ -592,6 +594,7 @@ func (r *AccountRef) usageStatusesLive(ctx context.Context) []AccountUsageStatus
 			next.Windows = details.Windows
 			next.Credits = details.Credits
 			next.ComplimentaryReset = details.ComplimentaryReset
+			next.UsageFresh = true
 			out[i] = next
 		}()
 	}
@@ -628,14 +631,15 @@ func (r *AccountRef) usageStatusesLive(ctx context.Context) []AccountUsageStatus
 			}
 			next.AuthValid = true
 			r.replace(account)
-			windows, _, err := r.FetchUsageWindowsCached(ctx, r.client, account)
+			windows, fresh, err := r.FetchUsageWindowsCached(ctx, r.client, account)
 			if err != nil {
 				next.Error = err.Error()
 				out[i] = next
 				return
 			}
 			next.PlanType = "claude"
 			next.Windows = windows
+			next.UsageFresh = fresh
 			out[i] = next
 		}()
 	}
@@ -953,7 +957,9 @@ func (s Server) handleUsageStatus(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	if s.AccountRef != nil {
-		writeJSON(w, s.withRequestTimeExhaustionWindows(s.AccountRef.UsageStatuses(r.Context())))
+		statuses := s.AccountRef.UsageStatuses(r.Context())
+		s.updateSchedulerFromUsageStatuses(statuses)
+		writeJSON(w, s.withRequestTimeExhaustionWindows(statuses))
 		return
 	}
 	accounts := s.accountList()
@@ -972,6 +978,47 @@ func (s Server) handleUsageStatus(w http.ResponseWriter, r *http.Request) {
 	writeJSON(w, s.withRequestTimeExhaustionWindows(out))
 }
 
+func (s Server) updateSchedulerFromUsageStatuses(statuses []AccountUsageStatus) {
+	if s.SchedulerRef == nil {
+		return
+	}
+	available := oauthAccounts(s.accountList())
+	if len(available) == 0 {
+		return
+	}
+	current := s.scheduler()
+	scores := make([]selectacct.Score, 0, len(available))
+	scoreByID := make(map[string]int, len(available))
+	for _, account := range available {
+		seed := current.ScoreFor(account.Provider, account.ID)
+		seed.AccountID = account.ID
+		seed.Provider = account.Provider
+		seed.Fresh = false
+		scoreByID[selectacct.ScoreKey(account.Provider, account.ID)] = len(scores)
+		scores = append(scores, seed)
+	}
+	scored := 0
+	for _, status := range statuses {
+		if !status.UsageFresh || status.AuthMode != accounts.AuthModeOAuth || len(status.Windows) == 0 {
+			continue
+		}
+		if idx, ok := scoreByID[selectacct.ScoreKey(status.Provider, status.ID)]; ok {
+			next := scoreFromUsageWindows(status.Provider, status.ID, status.Windows)
+			next.Fresh = true
+			scores[idx] = next
+			scored++
+		}
+	}
+	if scored == 0 {
+		return
+	}
+	scheduler := selectacct.NewScheduler(scores)
+	if s.Sessions != nil {
+		scheduler = scheduler.WithSessionCounts(s.Sessions.CountByAccount())
+	}
+	s.SchedulerRef.Set(scheduler)
+}
+
 func (s Server) withRequestTimeExhaustionWindows(statuses []AccountUsageStatus) []AccountUsageStatus {
 	if s.SchedulerRef == nil {
 		return statuses