feat: Add usage limit detection and upgrade prompt

apps/code/src/main/services/llm-gateway/service.ts

@@ -1,4 +1,5 @@
 import {
+  getGatewayInvalidatePlanCacheUrl,
   getGatewayUsageUrl,
   getLlmGatewayUrl,
 } from "@posthog/agent/posthog-api";
@@ -158,4 +159,24 @@ export class LlmGatewayService {
 
     return usageOutput.parse(await response.json());
   }
+
+  async invalidatePlanCache(): Promise<void> {
+    const auth = await this.authService.getValidAccessToken();
+    const url = getGatewayInvalidatePlanCacheUrl(auth.apiHost);
+
+    log.debug("Invalidating plan cache", { url });
+
+    const response = await this.authService.authenticatedFetch(fetch, url, {
+      method: "POST",
+    });
+
+    if (!response.ok) {
+      throw new LlmGatewayError(
+        `Failed to invalidate plan cache: HTTP ${response.status}`,
+        "plan_cache_error",
+        undefined,
+        response.status,
+      );
+    }
+  }
 }

apps/code/src/main/trpc/routers/llm-gateway.ts

@@ -26,4 +26,8 @@ export const llmGatewayRouter = router({
   usage: publicProcedure
     .output(usageOutput)
     .query(() => getService().fetchUsage()),
+
+  invalidatePlanCache: publicProcedure.mutation(() =>
+    getService().invalidatePlanCache(),
+  ),
 });

apps/code/src/renderer/components/MainLayout.tsx

@@ -4,6 +4,8 @@ import { HedgehogMode } from "@components/HedgehogMode";
 import { KeyboardShortcutsSheet } from "@components/KeyboardShortcutsSheet";
 
 import { ArchivedTasksView } from "@features/archive/components/ArchivedTasksView";
+import { UsageLimitModal } from "@features/billing/components/UsageLimitModal";
+import { useUsageLimitDetection } from "@features/billing/hooks/useUsageLimitDetection";
 import { CommandMenu } from "@features/command/components/CommandMenu";
 import { CommandCenterView } from "@features/command-center/components/CommandCenterView";
 import { InboxView } from "@features/inbox/components/InboxView";
@@ -20,8 +22,10 @@ import { TourOverlay } from "@features/tour/components/TourOverlay";
 import { useTourStore } from "@features/tour/stores/tourStore";
 import { createFirstTaskTour } from "@features/tour/tours/createFirstTaskTour";
 import { useConnectivity } from "@hooks/useConnectivity";
+import { useFeatureFlag } from "@hooks/useFeatureFlag";
 import { useIntegrations } from "@hooks/useIntegrations";
 import { Box, Flex } from "@radix-ui/themes";
+import { BILLING_FLAG } from "@shared/constants";
 import { useCommandMenuStore } from "@stores/commandMenuStore";
 import { useNavigationStore } from "@stores/navigationStore";
 import { useShortcutsSheetStore } from "@stores/shortcutsSheetStore";
@@ -43,12 +47,14 @@ export function MainLayout() {
   } = useShortcutsSheetStore();
   const { data: tasks } = useTasks();
   const { showPrompt, isChecking, check, dismiss } = useConnectivity();
+  const billingEnabled = useFeatureFlag(BILLING_FLAG);
 
   const startTour = useTourStore((s) => s.startTour);
   const isFirstTaskTourDone = useTourStore((s) =>
     s.completedTourIds.includes(createFirstTaskTour.id),
   );
 
+  useUsageLimitDetection(billingEnabled);
   useIntegrations();
   useTaskDeepLink();
   useInboxDeepLink();
@@ -119,6 +125,7 @@ export function MainLayout() {
       />
       <SettingsDialog />
       <TourOverlay />
+      {billingEnabled && <UsageLimitModal />}
       <HedgehogMode />
     </Flex>
   );

apps/code/src/renderer/features/auth/hooks/useAuthSession.ts

@@ -11,6 +11,7 @@ import { useAuthUiStateStore } from "@features/auth/stores/authUiStateStore";
 import { useSeatStore } from "@features/billing/stores/seatStore";
 import { useFeatureFlag } from "@hooks/useFeatureFlag";
 import { trpcClient } from "@renderer/trpc/client";
+import { BILLING_FLAG } from "@shared/constants";
 import { identifyUser, resetUser } from "@utils/analytics";
 import { logger } from "@utils/logger";
 import { useEffect } from "react";
@@ -104,7 +105,7 @@ export function useAuthSession() {
   const { data: currentUser } = useCurrentUser({ client });
   const authIdentity = getAuthIdentity(authState);
 
-  const billingEnabled = useFeatureFlag("posthog-code-billing");
+  const billingEnabled = useFeatureFlag(BILLING_FLAG);
 
   useAuthSubscriptionSync();
   useAuthIdentitySync(authIdentity, authState.cloudRegion);

apps/code/src/renderer/features/billing/components/SidebarUsageBar.tsx

@@ -0,0 +1,59 @@
+import { useUsage } from "@features/billing/hooks/useUsage";
+import { isUsageExceeded } from "@features/billing/utils";
+import { useSettingsDialogStore } from "@features/settings/stores/settingsDialogStore";
+import { useFeatureFlag } from "@hooks/useFeatureFlag";
+import { useSeat } from "@hooks/useSeat";
+import { Circle } from "@phosphor-icons/react";
+import { BILLING_FLAG } from "@shared/constants";
+
+export function SidebarUsageBar() {
+  const billingEnabled = useFeatureFlag(BILLING_FLAG);
+  const { seat, isPro } = useSeat();
+  const seatLoaded = seat !== null;
+  const { usage } = useUsage({
+    enabled: billingEnabled && seatLoaded && !isPro,
+  });
+
+  if (!billingEnabled || !seatLoaded || isPro || !usage) return null;
+
+  const usagePercent = Math.max(
+    usage.sustained.used_percent,
+    usage.burst.used_percent,
+  );
+  const exceeded = isUsageExceeded(usage);
+
+  const handleUpgrade = () => {
+    useSettingsDialogStore.getState().open("plan-usage");
+  };
+
+  return (
+    <div className="shrink-0 border-gray-6 border-t px-3 py-3">
+      <div className="flex items-center justify-between">
+        <span className="font-medium text-gray-11 text-xs">
+          Free plan
+          <Circle
+            size={4}
+            weight="fill"
+            className="mx-1.5 inline text-gray-8"
+          />
+          <span className="font-normal text-gray-10">
+            {exceeded ? "Limit reached" : `${Math.round(usagePercent)}% used`}
+          </span>
+        </span>
+        <button
+          type="button"
+          className="bg-transparent font-medium text-accent-11 text-xs transition-colors hover:text-accent-12"
+          onClick={handleUpgrade}
+        >
+          Upgrade
+        </button>
+      </div>
+      <div className="mt-2 h-2.5 w-full overflow-hidden rounded-full bg-gray-4">
+        <div
+          className={`h-full rounded-full transition-all ${exceeded ? "bg-red-9" : "bg-accent-9"}`}
+          style={{ width: `${Math.min(Math.ceil(usagePercent), 100)}%` }}
+        />
+      </div>
+    </div>
+  );
+}

apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx

@@ -0,0 +1,47 @@
+import { useUsageLimitStore } from "@features/billing/stores/usageLimitStore";
+import { useSettingsDialogStore } from "@features/settings/stores/settingsDialogStore";
+import { WarningCircle } from "@phosphor-icons/react";
+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 = () => {
+    hide();
+    useSettingsDialogStore.getState().open("plan-usage");
+  };
+
+  return (
+    <Dialog.Root open={isOpen}>
+      <Dialog.Content
+        maxWidth="400px"
+        onInteractOutside={(e) => e.preventDefault()}
+        onEscapeKeyDown={hide}
+      >
+        <Flex direction="column" gap="3">
+          <Flex align="center" gap="2">
+            <WarningCircle size={20} weight="bold" color="var(--red-9)" />
+            <Dialog.Title className="mb-0">Usage limit reached</Dialog.Title>
+          </Flex>
+          <Dialog.Description>
+            <Text size="2" color="gray">
+              {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."}
+            </Text>
+          </Dialog.Description>
+          <Flex justify="end" gap="3" mt="2">
+            <Button type="button" variant="soft" color="gray" onClick={hide}>
+              Not now
+            </Button>
+            <Button type="button" onClick={handleUpgrade}>
+              View upgrade options
+            </Button>
+          </Flex>
+        </Flex>
+      </Dialog.Content>
+    </Dialog.Root>
+  );
+}

apps/code/src/renderer/features/billing/hooks/useUsage.ts

@@ -0,0 +1,17 @@
+import { useTRPC } from "@renderer/trpc";
+import { useRendererWindowFocusStore } from "@stores/rendererWindowFocusStore";
+import { useQuery } from "@tanstack/react-query";
+
+const USAGE_REFETCH_INTERVAL_MS = 60_000;
+
+export function useUsage({ enabled = true }: { enabled?: boolean } = {}) {
+  const trpc = useTRPC();
+  const focused = useRendererWindowFocusStore((s) => s.focused);
+  const { data: usage, isLoading } = useQuery({
+    ...trpc.llmGateway.usage.queryOptions(),
+    enabled,
+    refetchInterval: focused && enabled ? USAGE_REFETCH_INTERVAL_MS : false,
+    refetchIntervalInBackground: false,
+  });
+  return { usage: usage ?? null, isLoading };
+}

apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts

@@ -0,0 +1,38 @@
+import { useUsageLimitStore } from "@features/billing/stores/usageLimitStore";
+import { isUsageExceeded } from "@features/billing/utils";
+import { useSessionStore } from "@features/sessions/stores/sessionStore";
+import { useSeat } from "@hooks/useSeat";
+import { useEffect, useRef } from "react";
+import { useUsage } from "./useUsage";
+
+export function useUsageLimitDetection(billingEnabled: boolean) {
+  const { seat, isPro } = useSeat();
+  const seatLoaded = seat !== null;
+  const { usage } = useUsage({
+    enabled: billingEnabled && seatLoaded && !isPro,
+  });
+  const hasAlertedRef = useRef(false);
+
+  useEffect(() => {
+    if (!billingEnabled || !seatLoaded || isPro || !usage) return;
+
+    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 (!exceeded) {
+      hasAlertedRef.current = false;
+    }
+  }, [billingEnabled, seatLoaded, isPro, usage]);
+}