diff --git a/apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx b/apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx index 25250a817..386687dc5 100644 --- a/apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx +++ b/apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx @@ -5,7 +5,6 @@ import { Button, Dialog, Flex, Text } from "@radix-ui/themes"; export function UsageLimitModal() { const isOpen = useUsageLimitStore((s) => s.isOpen); - const context = useUsageLimitStore((s) => s.context); const hide = useUsageLimitStore((s) => s.hide); const handleUpgrade = () => { @@ -27,9 +26,8 @@ export function UsageLimitModal() { - {context === "mid-task" - ? "You've hit your free plan usage limit. Your current task can't continue until usage resets or you upgrade to Pro." - : "You've reached your free plan usage limit. Upgrade to Pro for unlimited usage."} + You've reached your free plan usage limit. Upgrade to Pro for + unlimited usage. diff --git a/apps/code/src/renderer/features/billing/hooks/useUsage.ts b/apps/code/src/renderer/features/billing/hooks/useUsage.ts index 4cd77bf0e..a48e426af 100644 --- a/apps/code/src/renderer/features/billing/hooks/useUsage.ts +++ b/apps/code/src/renderer/features/billing/hooks/useUsage.ts @@ -2,16 +2,20 @@ import { useTRPC } from "@renderer/trpc"; import { useRendererWindowFocusStore } from "@stores/rendererWindowFocusStore"; import { useQuery } from "@tanstack/react-query"; -const USAGE_REFETCH_INTERVAL_MS = 60_000; +const USAGE_REFETCH_INTERVAL_MS = 30_000; export function useUsage({ enabled = true }: { enabled?: boolean } = {}) { const trpc = useTRPC(); const focused = useRendererWindowFocusStore((s) => s.focused); - const { data: usage, isLoading } = useQuery({ + const { + data: usage, + isLoading, + refetch, + } = useQuery({ ...trpc.llmGateway.usage.queryOptions(), enabled, refetchInterval: focused && enabled ? USAGE_REFETCH_INTERVAL_MS : false, refetchIntervalInBackground: false, }); - return { usage: usage ?? null, isLoading }; + return { usage: usage ?? null, isLoading, refetch }; } diff --git a/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts b/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts index 00f842eab..ac08dae2d 100644 --- a/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts +++ b/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts @@ -20,16 +20,15 @@ export function useUsageLimitDetection(billingEnabled: boolean) { const exceeded = isUsageExceeded(usage); if (exceeded && !hasAlertedRef.current) { - hasAlertedRef.current = true; - const sessions = useSessionStore.getState().sessions; const hasActiveSession = Object.values(sessions).some( (s) => s.status === "connected" && s.isPromptPending, ); - useUsageLimitStore - .getState() - .show(hasActiveSession ? "mid-task" : "idle"); + if (hasActiveSession) { + hasAlertedRef.current = true; + useUsageLimitStore.getState().show(); + } } if (!exceeded) { diff --git a/apps/code/src/renderer/features/billing/stores/seatStore.ts b/apps/code/src/renderer/features/billing/stores/seatStore.ts index fb3d13acb..46a1b92bf 100644 --- a/apps/code/src/renderer/features/billing/stores/seatStore.ts +++ b/apps/code/src/renderer/features/billing/stores/seatStore.ts @@ -7,6 +7,7 @@ import { trpcClient } from "@renderer/trpc"; import type { SeatData } from "@shared/types/seat"; import { PLAN_FREE, PLAN_PRO } from "@shared/types/seat"; import { logger } from "@utils/logger"; +import { queryClient } from "@utils/queryClient"; import { getPostHogUrl } from "@utils/urls"; import { create } from "zustand"; @@ -71,6 +72,7 @@ function invalidatePlanCache(): void { trpcClient.llmGateway.invalidatePlanCache.mutate().catch((err) => { log.warn("Failed to invalidate plan cache", err); }); + void queryClient.invalidateQueries({ queryKey: [["llmGateway"]] }); } const initialState: SeatStoreState = { @@ -80,7 +82,7 @@ const initialState: SeatStoreState = { redirectUrl: null, }; -export const useSeatStore = create()((set) => ({ +export const useSeatStore = create()((set, get) => ({ ...initialState, fetchSeat: async (options?: { autoProvision?: boolean }) => { @@ -90,10 +92,21 @@ export const useSeatStore = create()((set) => ({ let seat = await client.getMySeat(); if (!seat && options?.autoProvision) { log.info("No seat found, auto-provisioning free plan"); - seat = await client.createSeat(PLAN_FREE); + try { + seat = await client.createSeat(PLAN_FREE); + } catch { + log.info("Auto-provision failed, re-fetching seat"); + seat = await client.getMySeat(); + } } set({ seat, isLoading: false }); } catch (error) { + const { seat: existingSeat } = get(); + if (existingSeat) { + log.warn("fetchSeat failed but seat already loaded, keeping it", error); + set({ isLoading: false }); + return; + } handleSeatError(error, set); } }, diff --git a/apps/code/src/renderer/features/billing/stores/usageLimitStore.test.ts b/apps/code/src/renderer/features/billing/stores/usageLimitStore.test.ts index f97834b76..23dc20f0b 100644 --- a/apps/code/src/renderer/features/billing/stores/usageLimitStore.test.ts +++ b/apps/code/src/renderer/features/billing/stores/usageLimitStore.test.ts @@ -3,34 +3,22 @@ import { useUsageLimitStore } from "./usageLimitStore"; describe("usageLimitStore", () => { beforeEach(() => { - useUsageLimitStore.setState({ isOpen: false, context: null }); + useUsageLimitStore.setState({ isOpen: false }); }); - it("starts closed with no context", () => { + it("starts closed", () => { const state = useUsageLimitStore.getState(); expect(state.isOpen).toBe(false); - expect(state.context).toBeNull(); }); - it("show opens with mid-task context", () => { - useUsageLimitStore.getState().show("mid-task"); - const state = useUsageLimitStore.getState(); - expect(state.isOpen).toBe(true); - expect(state.context).toBe("mid-task"); - }); - - it("show opens with idle context", () => { - useUsageLimitStore.getState().show("idle"); - const state = useUsageLimitStore.getState(); - expect(state.isOpen).toBe(true); - expect(state.context).toBe("idle"); + it("show opens the modal", () => { + useUsageLimitStore.getState().show(); + expect(useUsageLimitStore.getState().isOpen).toBe(true); }); - it("hide closes but preserves context for exit animation", () => { - useUsageLimitStore.getState().show("mid-task"); + it("hide closes the modal", () => { + useUsageLimitStore.getState().show(); useUsageLimitStore.getState().hide(); - const state = useUsageLimitStore.getState(); - expect(state.isOpen).toBe(false); - expect(state.context).toBe("mid-task"); + expect(useUsageLimitStore.getState().isOpen).toBe(false); }); }); diff --git a/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts b/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts index 7fd54234a..c587db0f3 100644 --- a/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts +++ b/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts @@ -1,14 +1,11 @@ import { create } from "zustand"; -type UsageLimitContext = "mid-task" | "idle"; - interface UsageLimitState { isOpen: boolean; - context: UsageLimitContext | null; } interface UsageLimitActions { - show: (context: UsageLimitContext) => void; + show: () => void; hide: () => void; } @@ -16,8 +13,7 @@ type UsageLimitStore = UsageLimitState & UsageLimitActions; export const useUsageLimitStore = create()((set) => ({ isOpen: false, - context: null, - show: (context) => set({ isOpen: true, context }), + show: () => set({ isOpen: true }), hide: () => set({ isOpen: false }), })); diff --git a/apps/code/src/renderer/features/sessions/service/service.test.ts b/apps/code/src/renderer/features/sessions/service/service.test.ts index 83fbe1c40..836e2bd20 100644 --- a/apps/code/src/renderer/features/sessions/service/service.test.ts +++ b/apps/code/src/renderer/features/sessions/service/service.test.ts @@ -270,6 +270,7 @@ vi.mock("@utils/session", async () => { extractPromptText: vi.fn((p) => (typeof p === "string" ? p : "text")), getUserShellExecutesSinceLastPrompt: vi.fn(() => []), isFatalSessionError: actual.isFatalSessionError, + isRateLimitError: actual.isRateLimitError, normalizePromptToBlocks: vi.fn((p) => typeof p === "string" ? [{ type: "text", text: p }] : p, ), diff --git a/apps/code/src/renderer/features/sessions/service/service.ts b/apps/code/src/renderer/features/sessions/service/service.ts index 4e34683b5..23e99e951 100644 --- a/apps/code/src/renderer/features/sessions/service/service.ts +++ b/apps/code/src/renderer/features/sessions/service/service.ts @@ -8,6 +8,7 @@ import { getAuthenticatedClient, } from "@features/auth/hooks/authClient"; import { fetchAuthState } from "@features/auth/hooks/authQueries"; +import { useUsageLimitStore } from "@features/billing/stores/usageLimitStore"; import { useSessionAdapterStore } from "@features/sessions/stores/sessionAdapterStore"; import { getPersistedConfigOptions, @@ -70,6 +71,7 @@ import { extractPromptText, getUserShellExecutesSinceLastPrompt, isFatalSessionError, + isRateLimitError, normalizePromptToBlocks, shellExecutesToContextBlocks, } from "@utils/session"; @@ -1414,6 +1416,18 @@ export class SessionService { sessionStoreSetters.clearOptimisticItems(session.taskRunId); + if (isRateLimitError(errorMessage, errorDetails)) { + log.warn("Rate limit exceeded, showing usage limit modal", { + taskRunId: session.taskRunId, + }); + sessionStoreSetters.updateSession(session.taskRunId, { + isPromptPending: false, + promptStartedAt: null, + }); + useUsageLimitStore.getState().show(); + return { stopReason: "rate_limited" }; + } + if (isFatalSessionError(errorMessage, errorDetails)) { log.error("Fatal prompt error, attempting recovery", { taskRunId: session.taskRunId, diff --git a/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx b/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx index 4d00c7a22..3b40ba1a1 100644 --- a/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx +++ b/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx @@ -20,7 +20,7 @@ import { import { Tooltip } from "@renderer/components/ui/Tooltip"; import { PLAN_PRO_ALPHA } from "@shared/types/seat"; import { getPostHogUrl } from "@utils/urls"; -import { useState } from "react"; +import { useEffect, useState } from "react"; function formatResetTime(seconds: number): string { if (seconds < 3600) return "less than 1 hour"; @@ -43,14 +43,24 @@ export function PlanUsageSettings() { error, redirectUrl, } = useSeat(); - const { upgradeToPro, cancelSeat, reactivateSeat, clearError } = + const { fetchSeat, upgradeToPro, cancelSeat, reactivateSeat, clearError } = useSeatStore(); const [showUpgradeDialog, setShowUpgradeDialog] = useState(false); + const isAlpha = seat?.plan_key === PLAN_PRO_ALPHA; - const { usage, isLoading: usageLoading } = useUsage({ + const { + usage, + isLoading: usageLoading, + refetch: refetchUsage, + } = useUsage({ enabled: seat !== null, }); + useEffect(() => { + void fetchSeat(); + void refetchUsage(); + }, [fetchSeat, refetchUsage]); + const formattedActiveUntil = activeUntil ? activeUntil.toLocaleDateString(undefined, { month: "short", @@ -310,7 +320,7 @@ export function PlanUsageSettings() { }} disabled={isLoading} > - {isLoading ? : "Subscribe — $200/mo"} + {isLoading ? : "Subscribe - $200/mo"} diff --git a/apps/code/src/renderer/hooks/useFeatureFlag.ts b/apps/code/src/renderer/hooks/useFeatureFlag.ts index b1263f93d..de841080a 100644 --- a/apps/code/src/renderer/hooks/useFeatureFlag.ts +++ b/apps/code/src/renderer/hooks/useFeatureFlag.ts @@ -1,23 +1,15 @@ import { isFeatureFlagEnabled, onFeatureFlagsLoaded } from "@utils/analytics"; import { useEffect, useState } from "react"; -// only if in dev -const IS_DEV = import.meta.env.DEV; - export function useFeatureFlag( flagKey: string, defaultValue: boolean = false, ): boolean { const [enabled, setEnabled] = useState( - () => IS_DEV || isFeatureFlagEnabled(flagKey) || defaultValue, + () => isFeatureFlagEnabled(flagKey) || defaultValue, ); useEffect(() => { - if (IS_DEV) { - setEnabled(true); - return; - } - // Update immediately in case flags loaded between render and effect setEnabled(isFeatureFlagEnabled(flagKey) || defaultValue); diff --git a/apps/code/src/renderer/utils/session.ts b/apps/code/src/renderer/utils/session.ts index 550b24f04..ec99a997d 100644 --- a/apps/code/src/renderer/utils/session.ts +++ b/apps/code/src/renderer/utils/session.ts @@ -221,4 +221,4 @@ export function normalizePromptToBlocks( return typeof prompt === "string" ? [{ type: "text", text: prompt }] : prompt; } -export { isFatalSessionError } from "@shared/errors"; +export { isFatalSessionError, isRateLimitError } from "@shared/errors"; diff --git a/apps/code/src/shared/errors.ts b/apps/code/src/shared/errors.ts index 36e05a48b..ec6687b04 100644 --- a/apps/code/src/shared/errors.ts +++ b/apps/code/src/shared/errors.ts @@ -20,6 +20,12 @@ export function isAuthError(error: unknown): boolean { return AUTH_ERROR_PATTERNS.some((pattern) => message.includes(pattern)); } +const RATE_LIMIT_PATTERNS = [ + "rate limit exceeded", + "rate_limit", + "[429]", +] as const; + const FATAL_SESSION_ERROR_PATTERNS = [ "internal error", "process exited", @@ -37,10 +43,21 @@ function includesAny( return patterns.some((pattern) => lower.includes(pattern)); } +export function isRateLimitError( + errorMessage: string, + errorDetails?: string, +): boolean { + return ( + includesAny(errorMessage, RATE_LIMIT_PATTERNS) || + includesAny(errorDetails, RATE_LIMIT_PATTERNS) + ); +} + export function isFatalSessionError( errorMessage: string, errorDetails?: string, ): boolean { + if (isRateLimitError(errorMessage, errorDetails)) return false; return ( includesAny(errorMessage, FATAL_SESSION_ERROR_PATTERNS) || includesAny(errorDetails, FATAL_SESSION_ERROR_PATTERNS)