feat: Improve usage limit handling and data freshness

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() {
           </Flex>
           <Dialog.Description>
             <Text color="gray" className="text-sm">
-              {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.
             </Text>
           </Dialog.Description>
           <Flex justify="end" gap="3" mt="2">

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 };
 }

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) {

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<SeatStore>()((set) => ({
+export const useSeatStore = create<SeatStore>()((set, get) => ({
   ...initialState,
 
   fetchSeat: async (options?: { autoProvision?: boolean }) => {
@@ -90,10 +92,21 @@ export const useSeatStore = create<SeatStore>()((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);
     }
   },

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);
   });
 });

apps/code/src/renderer/features/billing/stores/usageLimitStore.ts

@@ -1,23 +1,19 @@
 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;
 }
 
 type UsageLimitStore = UsageLimitState & UsageLimitActions;
 
 export const useUsageLimitStore = create<UsageLimitStore>()((set) => ({
   isOpen: false,
-  context: null,
 
-  show: (context) => set({ isOpen: true, context }),
+  show: () => set({ isOpen: true }),
   hide: () => set({ isOpen: false }),
 }));

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,
     ),

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,

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 ? <Spinner size="1" /> : "Subscribe — $200/mo"}
+              {isLoading ? <Spinner size="1" /> : "Subscribe - $200/mo"}
             </Button>
           </Flex>
         </Dialog.Content>

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);
 

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";

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)