fix(billing): suppress budget banner when all chat workloads route to a user-supplied provider (#2040, #2041) (#2053)

Co-authored-by: 李冠辰 <liguanchen@xiaomi.com>
Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>

app/src/hooks/useUsageState.test.ts

@@ -3,6 +3,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 const mockGetCurrentPlan = vi.fn();
 const mockGetTeamUsage = vi.fn();
+const mockLoadAISettings = vi.fn();
 
 vi.mock('../services/api/billingApi', () => ({
   billingApi: { getCurrentPlan: () => mockGetCurrentPlan() },
@@ -12,6 +13,30 @@ vi.mock('../services/api/creditsApi', () => ({
   creditsApi: { getTeamUsage: () => mockGetTeamUsage() },
 }));
 
+vi.mock('../services/api/aiSettingsApi', async () => {
+  const actual = await vi.importActual<typeof import('../services/api/aiSettingsApi')>(
+    '../services/api/aiSettingsApi'
+  );
+  return { ...actual, loadAISettings: () => mockLoadAISettings() };
+});
+
+// All chat workloads routed to OpenHuman — the default for every existing
+// test case (matches the legacy "you have a hosted-backend budget" world).
+const ALL_OPENHUMAN_AI_SETTINGS = {
+  cloudProviders: [],
+  routing: {
+    chat: { kind: 'openhuman' as const },
+    reasoning: { kind: 'openhuman' as const },
+    agentic: { kind: 'openhuman' as const },
+    coding: { kind: 'openhuman' as const },
+    memory: { kind: 'openhuman' as const },
+    embeddings: { kind: 'openhuman' as const },
+    heartbeat: { kind: 'openhuman' as const },
+    learning: { kind: 'openhuman' as const },
+    subconscious: { kind: 'openhuman' as const },
+  },
+};
+
 interface BuildUsageOpts {
   remainingUsd?: number;
   cycleBudgetUsd?: number;
@@ -82,6 +107,11 @@ describe('useUsageState', () => {
     vi.resetModules();
     mockGetCurrentPlan.mockReset();
     mockGetTeamUsage.mockReset();
+    mockLoadAISettings.mockReset();
+    // Default: keep the OpenHuman-routed world so every legacy assertion
+    // about budget gating stays identical until a test opts into the
+    // routed-away scenarios below.
+    mockLoadAISettings.mockResolvedValue(ALL_OPENHUMAN_AI_SETTINGS);
   });
 
   it('does not treat free users with zero recurring budget as exhausted', async () => {
@@ -204,4 +234,300 @@ describe('useUsageState', () => {
       expect(result.current.teamUsage?.remainingUsd).toBe(7);
     });
   });
+
+  // -- #2040 / #2041 — budget banner is suppressed when chat routes away ----
+
+  it('suppresses the budget banner when every chat workload routes to a user-supplied cloud provider (#2040, #2041)', async () => {
+    const { useUsageState } = await import('./useUsageState');
+
+    // Plan + usage say "budget exhausted" — but the user has saved an
+    // OpenRouter key and routed reasoning/agentic/coding away from
+    // openhuman. The banner that previously said "Your included budget is
+    // complete" should NOT show, because the user is paying OpenRouter,
+    // not OpenHuman, for chat inference.
+    mockGetCurrentPlan.mockResolvedValue({
+      plan: 'BASIC',
+      hasActiveSubscription: true,
+      planExpiry: '2026-05-01T00:00:00.000Z',
+      subscription: {
+        id: 'sub_123',
+        status: 'active',
+        currentPeriodEnd: '2026-05-01T00:00:00.000Z',
+        quantity: 1,
+      },
+      monthlyBudgetUsd: 20,
+      weeklyBudgetUsd: 10,
+      fiveHourCapUsd: 3,
+    });
+    mockGetTeamUsage.mockResolvedValue({
+      remainingUsd: 0,
+      cycleBudgetUsd: 10,
+      cycleLimit5hr: 3, // at-the-cap to also exercise the rate-limit gate
+      cycleLimit7day: 10,
+      fiveHourCapUsd: 3,
+      fiveHourResetsAt: null,
+      cycleStartDate: '2026-04-09T00:00:00.000Z',
+      cycleEndsAt: '2026-04-16T00:00:00.000Z',
+      bypassCycleLimit: false,
+    });
+    mockLoadAISettings.mockResolvedValue({
+      cloudProviders: [],
+      routing: {
+        chat: { kind: 'cloud', providerSlug: 'openrouter', model: 'anthropic/claude-sonnet-4.6' },
+        reasoning: {
+          kind: 'cloud',
+          providerSlug: 'openrouter',
+          model: 'google/gemini-3-flash-preview',
+        },
+        agentic: {
+          kind: 'cloud',
+          providerSlug: 'openrouter',
+          model: 'anthropic/claude-sonnet-4.6',
+        },
+        coding: { kind: 'cloud', providerSlug: 'openrouter', model: 'anthropic/claude-sonnet-4.6' },
+        // Background workloads may still route to openhuman — the suppression
+        // logic only consults CHAT_WORKLOADS (chat/reasoning/agentic/coding).
+        memory: { kind: 'openhuman' },
+        embeddings: { kind: 'openhuman' },
+        heartbeat: { kind: 'openhuman' },
+        learning: { kind: 'openhuman' },
+        subconscious: { kind: 'openhuman' },
+      },
+    });
+
+    const { result } = renderHook(() => useUsageState());
+
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(result.current.isFullyRoutedAway).toBe(true);
+    expect(result.current.shouldShowBudgetCompletedMessage).toBe(false);
+    expect(result.current.isBudgetExhausted).toBe(false);
+    expect(result.current.isAtLimit).toBe(false);
+  });
+
+  it('still shows the budget banner when at least one chat workload remains on OpenHuman', async () => {
+    const { useUsageState } = await import('./useUsageState');
+
+    // User has saved an OpenRouter key for agentic+coding but left reasoning
+    // on openhuman — they're still partially dependent on the included
+    // budget, so the banner must keep showing.
+    mockGetCurrentPlan.mockResolvedValue({
+      plan: 'BASIC',
+      hasActiveSubscription: true,
+      planExpiry: '2026-05-01T00:00:00.000Z',
+      subscription: {
+        id: 'sub_123',
+        status: 'active',
+        currentPeriodEnd: '2026-05-01T00:00:00.000Z',
+        quantity: 1,
+      },
+      monthlyBudgetUsd: 20,
+      weeklyBudgetUsd: 10,
+      fiveHourCapUsd: 3,
+    });
+    mockGetTeamUsage.mockResolvedValue({
+      remainingUsd: 0,
+      cycleBudgetUsd: 10,
+      cycleLimit5hr: 1,
+      cycleLimit7day: 10,
+      fiveHourCapUsd: 3,
+      fiveHourResetsAt: null,
+      cycleStartDate: '2026-04-09T00:00:00.000Z',
+      cycleEndsAt: '2026-04-16T00:00:00.000Z',
+      bypassCycleLimit: false,
+    });
+    mockLoadAISettings.mockResolvedValue({
+      cloudProviders: [],
+      routing: {
+        reasoning: { kind: 'openhuman' }, // still on the hosted backend
+        agentic: {
+          kind: 'cloud',
+          providerSlug: 'openrouter',
+          model: 'anthropic/claude-sonnet-4.6',
+        },
+        coding: { kind: 'cloud', providerSlug: 'openrouter', model: 'anthropic/claude-sonnet-4.6' },
+        memory: { kind: 'openhuman' },
+        embeddings: { kind: 'openhuman' },
+        heartbeat: { kind: 'openhuman' },
+        learning: { kind: 'openhuman' },
+        subconscious: { kind: 'openhuman' },
+      },
+    });
+
+    const { result } = renderHook(() => useUsageState());
+
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(result.current.isFullyRoutedAway).toBe(false);
+    expect(result.current.shouldShowBudgetCompletedMessage).toBe(true);
+    expect(result.current.isBudgetExhausted).toBe(true);
+    expect(result.current.isAtLimit).toBe(true);
+  });
+
+  it('treats missing aiSettings (fetch failure) as conservative — banner still shows when budget is otherwise exhausted', async () => {
+    const { useUsageState } = await import('./useUsageState');
+
+    mockGetCurrentPlan.mockResolvedValue({
+      plan: 'BASIC',
+      hasActiveSubscription: true,
+      planExpiry: '2026-05-01T00:00:00.000Z',
+      subscription: {
+        id: 'sub_123',
+        status: 'active',
+        currentPeriodEnd: '2026-05-01T00:00:00.000Z',
+        quantity: 1,
+      },
+      monthlyBudgetUsd: 20,
+      weeklyBudgetUsd: 10,
+      fiveHourCapUsd: 3,
+    });
+    mockGetTeamUsage.mockResolvedValue({
+      remainingUsd: 0,
+      cycleBudgetUsd: 10,
+      cycleLimit5hr: 1,
+      cycleLimit7day: 10,
+      fiveHourCapUsd: 3,
+      fiveHourResetsAt: null,
+      cycleStartDate: '2026-04-09T00:00:00.000Z',
+      cycleEndsAt: '2026-04-16T00:00:00.000Z',
+      bypassCycleLimit: false,
+    });
+    // Simulate a transient failure in the AI-settings fetch — the budget
+    // gate must NOT silently disable itself just because we couldn't read
+    // the routing config.
+    mockLoadAISettings.mockRejectedValue(new Error('network down'));
+
+    const { result } = renderHook(() => useUsageState());
+
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(result.current.isFullyRoutedAway).toBe(false);
+    expect(result.current.shouldShowBudgetCompletedMessage).toBe(true);
+    expect(result.current.isBudgetExhausted).toBe(true);
+    expect(result.current.isAtLimit).toBe(true);
+  });
+
+  it('suppresses the budget banner when every chat workload routes to a local Ollama provider (#2040, #2041, ProviderRef kind=local)', async () => {
+    // Companion to the cloud-provider case above. The PR description claims
+    // suppression also applies to ProviderRef kind='local' (Ollama on-device
+    // inference). graycyrus review on #2053 flagged that none of the new
+    // tests pinned that path. This one does.
+    const { useUsageState } = await import('./useUsageState');
+
+    mockGetCurrentPlan.mockResolvedValue({
+      plan: 'BASIC',
+      hasActiveSubscription: true,
+      planExpiry: '2026-05-01T00:00:00.000Z',
+      subscription: {
+        id: 'sub_123',
+        status: 'active',
+        currentPeriodEnd: '2026-05-01T00:00:00.000Z',
+        quantity: 1,
+      },
+      monthlyBudgetUsd: 20,
+      weeklyBudgetUsd: 10,
+      fiveHourCapUsd: 3,
+    });
+    mockGetTeamUsage.mockResolvedValue({
+      remainingUsd: 0,
+      cycleBudgetUsd: 10,
+      cycleLimit5hr: 3, // at the cap to also exercise rate-limit gating
+      cycleLimit7day: 10,
+      fiveHourCapUsd: 3,
+      fiveHourResetsAt: null,
+      cycleStartDate: '2026-04-09T00:00:00.000Z',
+      cycleEndsAt: '2026-04-16T00:00:00.000Z',
+      bypassCycleLimit: false,
+    });
+    mockLoadAISettings.mockResolvedValue({
+      cloudProviders: [],
+      routing: {
+        // All chat workloads on local Ollama models.
+        chat: { kind: 'local', model: 'qwen3:8b' },
+        reasoning: { kind: 'local', model: 'qwen3:8b' },
+        agentic: { kind: 'local', model: 'qwen3:8b' },
+        coding: { kind: 'local', model: 'qwen3:8b' },
+        // Background workloads are intentionally left on openhuman to
+        // prove the gate is keyed on chat workloads only.
+        memory: { kind: 'openhuman' },
+        embeddings: { kind: 'openhuman' },
+        heartbeat: { kind: 'openhuman' },
+        learning: { kind: 'openhuman' },
+        subconscious: { kind: 'openhuman' },
+      },
+    });
+
+    const { result } = renderHook(() => useUsageState());
+
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    expect(result.current.isFullyRoutedAway).toBe(true);
+    expect(result.current.shouldShowBudgetCompletedMessage).toBe(false);
+    expect(result.current.isBudgetExhausted).toBe(false);
+    expect(result.current.isAtLimit).toBe(false);
+  });
+
+  it('rethrows CoreRpcError(kind=auth_expired) from loadAISettings instead of swallowing it (graycyrus review on #2053)', async () => {
+    // The two sibling fetches (getTeamUsage, getCurrentPlan) explicitly
+    // re-throw auth_expired so coreRpcClient's global re-auth event fires.
+    // loadAISettings goes through the same RPC layer and must follow the
+    // same contract — otherwise a session-expired user gets stale data
+    // instead of a re-auth prompt.
+    const { useUsageState } = await import('./useUsageState');
+    const { CoreRpcError } = await import('../services/coreRpcClient');
+
+    mockGetCurrentPlan.mockResolvedValue({
+      plan: 'FREE',
+      hasActiveSubscription: false,
+      planExpiry: null,
+      subscription: null,
+      monthlyBudgetUsd: 0,
+      weeklyBudgetUsd: 0,
+      fiveHourCapUsd: 0,
+    });
+    mockGetTeamUsage.mockResolvedValue({
+      remainingUsd: 0,
+      cycleBudgetUsd: 0,
+      cycleLimit5hr: 0,
+      cycleLimit7day: 0,
+      fiveHourCapUsd: 0,
+      fiveHourResetsAt: null,
+      cycleStartDate: '2026-04-09T00:00:00.000Z',
+      cycleEndsAt: '2026-04-16T00:00:00.000Z',
+      bypassCycleLimit: false,
+    });
+    mockLoadAISettings.mockRejectedValue(
+      new CoreRpcError(
+        'GET /ai/settings failed (401 Unauthorized): Session expired.',
+        'auth_expired',
+        401
+      )
+    );
+
+    const unhandled = vi.fn();
+    window.addEventListener('unhandledrejection', unhandled);
+    try {
+      const { result } = renderHook(() => useUsageState());
+      await waitFor(() => {
+        expect(result.current.isLoading).toBe(false);
+      });
+      // The hook's outer .catch swallows auth_expired silently (matching
+      // the existing #1472 contract). The rejection must NOT have leaked
+      // to window.unhandledrejection.
+      expect(result.current.teamUsage).toBeNull();
+      expect(result.current.currentPlan).toBeNull();
+      expect(result.current.isFullyRoutedAway).toBe(false);
+      expect(unhandled).not.toHaveBeenCalled();
+    } finally {
+      window.removeEventListener('unhandledrejection', unhandled);
+    }
+  });
 });

app/src/hooks/useUsageState.ts

@@ -1,5 +1,6 @@
 import { useCallback, useEffect, useState } from 'react';
 
+import { type AISettings, CHAT_WORKLOADS, loadAISettings } from '../services/api/aiSettingsApi';
 import { billingApi } from '../services/api/billingApi';
 import { creditsApi, type TeamUsage } from '../services/api/creditsApi';
 import { CoreRpcError } from '../services/coreRpcClient';
@@ -16,6 +17,14 @@ export interface UsageState {
   isAtLimit: boolean;
   isBudgetExhausted: boolean;
   shouldShowBudgetCompletedMessage: boolean;
+  /**
+   * True when every chat workload (reasoning/agentic/coding) is routed to a
+   * non-openhuman provider (a user-configured cloud provider or local Ollama).
+   * Used to suppress the OpenHuman-included-budget banner / modal: when the
+   * user has explicitly bypassed the hosted backend for chat, the included
+   * budget cycle no longer gates them. See #2040 and #2041.
+   */
+  isFullyRoutedAway: boolean;
   isLoading: boolean;
   refresh: () => void;
 }
@@ -23,7 +32,7 @@ export interface UsageState {
 const CACHE_TTL_MS = 60_000;
 
 let _cache: {
-  data: { teamUsage: TeamUsage; currentPlan: CurrentPlanData };
+  data: { teamUsage: TeamUsage; currentPlan: CurrentPlanData; aiSettings: AISettings | null };
   fetchedAt: number;
 } | null = null;
 
@@ -32,6 +41,7 @@ const USAGE_UNAVAILABLE = Symbol('usage-unavailable');
 async function fetchUsageData(): Promise<{
   teamUsage: TeamUsage | null;
   currentPlan: CurrentPlanData | null;
+  aiSettings: AISettings | null;
 } | null> {
   if (_cache && Date.now() - _cache.fetchedAt < CACHE_TTL_MS) {
     return _cache.data;
@@ -40,7 +50,7 @@ async function fetchUsageData(): Promise<{
   // session expiry) cannot reject the Promise.all microtask before the
   // sibling resolves — that race let the unhandled rejection leak to the
   // window's unhandledrejection trap and onward to Sentry (#1472).
-  const [teamUsage, currentPlan] = await Promise.all([
+  const [teamUsage, currentPlan, aiSettings] = await Promise.all([
     creditsApi.getTeamUsage().catch(err => {
       if (err instanceof CoreRpcError && err.kind === 'auth_expired') {
         throw err;
@@ -53,14 +63,32 @@ async function fetchUsageData(): Promise<{
       }
       return USAGE_UNAVAILABLE;
     }),
+    // AI settings drive the "routed away from openhuman" detection used to
+    // suppress the budget banner when the user supplied their own provider
+    // key (#2040 / #2041). Mirror the sibling fetches: re-throw
+    // CoreRpcError(kind='auth_expired') so the documented session-expired
+    // signal still reaches the global re-auth handler (graycyrus review on
+    // #2053). Other failures are treated as "unknown" — the budget gate
+    // stays in its conservative (banner-on) state.
+    loadAISettings().catch(err => {
+      if (err instanceof CoreRpcError && err.kind === 'auth_expired') {
+        throw err;
+      }
+      return USAGE_UNAVAILABLE;
+    }),
   ]);
   const data = {
     teamUsage: teamUsage === USAGE_UNAVAILABLE ? null : (teamUsage as TeamUsage),
     currentPlan: currentPlan === USAGE_UNAVAILABLE ? null : (currentPlan as CurrentPlanData),
+    aiSettings: aiSettings === USAGE_UNAVAILABLE ? null : (aiSettings as AISettings),
   };
   if (data.teamUsage && data.currentPlan) {
     _cache = {
-      data: { teamUsage: data.teamUsage, currentPlan: data.currentPlan },
+      data: {
+        teamUsage: data.teamUsage,
+        currentPlan: data.currentPlan,
+        aiSettings: data.aiSettings,
+      },
       fetchedAt: Date.now(),
     };
   }
@@ -70,6 +98,7 @@ async function fetchUsageData(): Promise<{
 export function useUsageState(): UsageState {
   const [teamUsage, setTeamUsage] = useState<TeamUsage | null>(null);
   const [currentPlan, setCurrentPlan] = useState<CurrentPlanData | null>(null);
+  const [aiSettings, setAiSettings] = useState<AISettings | null>(null);
   const [isLoading, setIsLoading] = useState(false);
   const [fetchCount, setFetchCount] = useState(0);
 
@@ -88,6 +117,7 @@ export function useUsageState(): UsageState {
         if (cancelled || !data) return;
         setTeamUsage(data.teamUsage);
         setCurrentPlan(data.currentPlan);
+        setAiSettings(data.aiSettings);
       })
       .catch((err: unknown) => {
         // CoreRpcError(kind=auth_expired) is the documented signal that the
@@ -119,17 +149,33 @@ export function useUsageState(): UsageState {
         )
       : 0;
 
-  const isBudgetExhausted = teamUsage
+  // When every chat workload routes to a user-supplied provider (cloud or
+  // local Ollama), the OpenHuman included-budget cycle does not gate the
+  // user. Conservative on missing aiSettings (treat as still using
+  // openhuman) so we never silently disable the gate after a transient
+  // fetch failure (#2040, #2041).
+  const isFullyRoutedAway = aiSettings
+    ? CHAT_WORKLOADS.every(w => {
+        const ref = aiSettings.routing[w];
+        return ref !== undefined && ref.kind !== 'openhuman';
+      })
+    : false;
+
+  const rawBudgetExhausted = teamUsage
     ? teamUsage.cycleBudgetUsd > 0.01 && teamUsage.remainingUsd <= 0.01
     : false;
 
   // Some users have no included recurring budget at all. They still need the
   // completed-budget warning in chat even though they are not in an exhausted
-  // paid cycle.
-  const shouldShowBudgetCompletedMessage = teamUsage
-    ? isBudgetExhausted || (teamUsage.cycleBudgetUsd <= 0.01 && teamUsage.remainingUsd <= 0.01)
+  // paid cycle — but only when their chat actually flows through OpenHuman.
+  const rawShouldShowBudgetCompletedMessage = teamUsage
+    ? rawBudgetExhausted || (teamUsage.cycleBudgetUsd <= 0.01 && teamUsage.remainingUsd <= 0.01)
     : false;
 
+  const isBudgetExhausted = !isFullyRoutedAway && rawBudgetExhausted;
+  const shouldShowBudgetCompletedMessage =
+    !isFullyRoutedAway && rawShouldShowBudgetCompletedMessage;
+
   const isAtLimit = isBudgetExhausted;
 
   const isNearLimit = !isAtLimit && teamUsage !== null && usagePct >= 0.8;
@@ -144,6 +190,7 @@ export function useUsageState(): UsageState {
     isAtLimit,
     isBudgetExhausted,
     shouldShowBudgetCompletedMessage,
+    isFullyRoutedAway,
     isLoading,
     refresh,
   };