Merge pull request #45 from manaflow-ai/fix/fable-headerless-probe

Show headerless Fable probe quota exhaustion

Author: lawrencecchen

cmd/subrouter/sr.go

@@ -1325,7 +1325,9 @@ func fetchClaudeUsageWindows(ctx context.Context, client *http.Client, accessTok
 	if !usageWindowNamed(windows, agentclaude.FableWindowName) {
 		fableWindows, probeErr := agentclaude.FetchFableUsageWindows(ctx, client, accessToken)
 		if probeErr == nil && len(fableWindows) > 0 {
-			windows = mergeUsageWindows(windows, fableWindows)
+			if err == nil || fableProbeHasPrimaryWindows(fableWindows) {
+				windows = mergeUsageWindows(windows, fableWindows)
+			}
 		} else if err != nil {
 			if probeErr != nil {
 				return nil, probeErr
@@ -1339,6 +1341,15 @@ func fetchClaudeUsageWindows(ctx context.Context, client *http.Client, accessTok
 	return windows, nil
 }
 
+func fableProbeHasPrimaryWindows(windows []accounts.UsageWindow) bool {
+	for _, window := range windows {
+		if window.Name == "5h" || window.Name == "7d" {
+			return true
+		}
+	}
+	return false
+}
+
 func mergeUsageWindows(base, extra []accounts.UsageWindow) []accounts.UsageWindow {
 	out := append([]accounts.UsageWindow(nil), base...)
 	index := make(map[string]int, len(out))

internal/agents/claude/store.go

@@ -749,6 +749,14 @@ func FetchFableUsageWindows(ctx context.Context, client *http.Client, accessToke
 	if len(windows) > 0 {
 		return windows, nil
 	}
+	if res.StatusCode == http.StatusTooManyRequests {
+		return []accounts.UsageWindow{{
+			Name:               FableWindowName,
+			UsedPercent:        100,
+			LimitWindowSeconds: int64(7 * 24 * time.Hour / time.Second),
+			Feature:            FableModel,
+		}}, nil
+	}
 	if res.StatusCode == http.StatusUnauthorized {
 		return nil, fmt.Errorf("fable probe failed: %s", res.Status)
 	}

internal/proxy/claude_ratelimit_routing_test.go

@@ -423,6 +423,41 @@ func TestClaudeFableProbeAddsHiddenOAuthAppsWindow(t *testing.T) {
 	}
 }
 
+func TestClaudeFableProbeTreatsHeaderless429AsFableExhausted(t *testing.T) {
+	usageBody := `{"five_hour":{"utilization":10.0,"resets_at":"2030-01-01T00:00:00+00:00"},"seven_day":{"utilization":60.0,"resets_at":"2030-01-02T00:00:00+00:00"}}`
+	client := &http.Client{Transport: proxyRoundTripFunc(func(req *http.Request) (*http.Response, error) {
+		switch req.URL.Path {
+		case "/api/oauth/usage":
+			return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(usageBody)), Header: http.Header{}}, nil
+		case "/v1/messages":
+			return &http.Response{StatusCode: http.StatusTooManyRequests, Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"rate_limit_error","message":"Error"}}`)), Header: http.Header{}}, nil
+		default:
+			t.Fatalf("unexpected path %s", req.URL.Path)
+			return nil, nil
+		}
+	})}
+	windows, err := fetchAccountUsageWindowsLive(context.Background(), client, accounts.Account{
+		ID:       "acct",
+		Provider: accounts.ProviderClaude,
+		AuthMode: accounts.AuthModeOAuth,
+		Token:    "tok",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	score := scoreFromUsageWindows(accounts.ProviderClaude, "acct", windows)
+	if score.Headroom == 0 {
+		t.Fatalf("base Headroom = %.3f, hidden Fable bucket should not exhaust non-Fable Claude models", score.Headroom)
+	}
+	fableScore, ok := score.ModelScores[selectacct.ModelKey(agentclaude.FableModel)]
+	if !ok {
+		t.Fatalf("missing Fable model score: %+v", score.ModelScores)
+	}
+	if fableScore.Headroom != 0 {
+		t.Fatalf("Fable Headroom = %.3f, want 0 from headerless Fable probe 429", fableScore.Headroom)
+	}
+}
+
 // 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

@@ -993,8 +993,10 @@ func (s Server) withRequestTimeExhaustionWindows(statuses []AccountUsageStatus)
 		}
 		name := "request-limit"
 		windowSeconds := int64(7 * 24 * 60 * 60)
+		feature := ""
 		if status.Provider == accounts.ProviderClaude {
-			name = "oauth-apps-weekly"
+			name = agentclaude.FableWindowName
+			feature = agentclaude.FableModel
 		}
 		if usageWindowNamed(status.Windows, name) {
 			continue
@@ -1004,6 +1006,7 @@ func (s Server) withRequestTimeExhaustionWindows(statuses []AccountUsageStatus)
 			UsedPercent:        100,
 			LimitWindowSeconds: windowSeconds,
 			ResetAfterSeconds:  resetAfter,
+			Feature:            feature,
 		})
 	}
 	return out
@@ -1374,7 +1377,9 @@ func fetchAccountUsageWindowsLive(ctx context.Context, client *http.Client, acco
 		windows := claudeUsageWindows(usage)
 		if !usageWindowNamed(windows, agentclaude.FableWindowName) {
 			if fableWindows, probeErr := agentclaude.FetchFableUsageWindows(ctx, client, account.Token); probeErr == nil && len(fableWindows) > 0 {
-				windows = mergeUsageWindows(windows, fableWindows)
+				if err == nil || fableProbeHasPrimaryWindows(fableWindows) {
+					windows = mergeUsageWindows(windows, fableWindows)
+				}
 			} else if err != nil {
 				if probeErr != nil {
 					return nil, probeErr
@@ -1390,6 +1395,15 @@ func fetchAccountUsageWindowsLive(ctx context.Context, client *http.Client, acco
 	return accounts.FetchCodexUsage(ctx, client, account)
 }
 
+func fableProbeHasPrimaryWindows(windows []accounts.UsageWindow) bool {
+	for _, window := range windows {
+		if window.Name == "5h" || window.Name == "7d" {
+			return true
+		}
+	}
+	return false
+}
+
 func mergeUsageWindows(base, extra []accounts.UsageWindow) []accounts.UsageWindow {
 	out := append([]accounts.UsageWindow(nil), base...)
 	index := make(map[string]int, len(out))

internal/proxy/proxy_websocket_test.go

@@ -1096,14 +1096,17 @@ func TestUsageStatusEndpointIncludesClaudeRequestTimeExhaustion(t *testing.T) {
 		t.Fatalf("statuses = %+v, want one", statuses)
 	}
 	for _, window := range statuses[0].Windows {
-		if window.Name == "oauth-apps-weekly" {
+		if window.Name == agentclaude.FableWindowName {
 			if window.UsedPercent != 100 || window.ResetAfterSeconds <= 0 {
-				t.Fatalf("oauth-apps-weekly = %+v, want saturated with reset", window)
+				t.Fatalf("Fable window = %+v, want saturated with reset", window)
+			}
+			if window.Feature != agentclaude.FableModel {
+				t.Fatalf("Fable Feature = %q, want %q", window.Feature, agentclaude.FableModel)
 			}
 			return
 		}
 	}
-	t.Fatalf("missing oauth-apps-weekly window: %+v", statuses[0].Windows)
+	t.Fatalf("missing Fable window: %+v", statuses[0].Windows)
 }
 
 func TestReloadAccountsHotLoadsNewAccountWithoutRestart(t *testing.T) {