diff --git a/apps/code/src/main/services/llm-gateway/service.ts b/apps/code/src/main/services/llm-gateway/service.ts index df79a2eb9..94f22e19a 100644 --- a/apps/code/src/main/services/llm-gateway/service.ts +++ b/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 { + 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, + ); + } + } } diff --git a/apps/code/src/main/trpc/routers/llm-gateway.ts b/apps/code/src/main/trpc/routers/llm-gateway.ts index a2dafcea7..fe9862e1a 100644 --- a/apps/code/src/main/trpc/routers/llm-gateway.ts +++ b/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(), + ), }); diff --git a/apps/code/src/renderer/components/MainLayout.tsx b/apps/code/src/renderer/components/MainLayout.tsx index 37cce59fe..b05f1b28c 100644 --- a/apps/code/src/renderer/components/MainLayout.tsx +++ b/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() { /> + {billingEnabled && } ); diff --git a/apps/code/src/renderer/features/auth/hooks/useAuthSession.ts b/apps/code/src/renderer/features/auth/hooks/useAuthSession.ts index 56673fcbd..03d489290 100644 --- a/apps/code/src/renderer/features/auth/hooks/useAuthSession.ts +++ b/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); diff --git a/apps/code/src/renderer/features/billing/components/SidebarUsageBar.tsx b/apps/code/src/renderer/features/billing/components/SidebarUsageBar.tsx new file mode 100644 index 000000000..027d6cc2e --- /dev/null +++ b/apps/code/src/renderer/features/billing/components/SidebarUsageBar.tsx @@ -0,0 +1,54 @@ +import { useFreeUsage } from "@features/billing/hooks/useFreeUsage"; +import { isUsageExceeded } from "@features/billing/utils"; +import { useSettingsDialogStore } from "@features/settings/stores/settingsDialogStore"; +import { useFeatureFlag } from "@hooks/useFeatureFlag"; +import { Circle } from "@phosphor-icons/react"; +import { BILLING_FLAG } from "@shared/constants"; + +export function SidebarUsageBar() { + const billingEnabled = useFeatureFlag(BILLING_FLAG); + const usage = useFreeUsage(billingEnabled); + + if (!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 ( +
+
+ + Free plan + + + {exceeded ? "Limit reached" : `${Math.round(usagePercent)}% used`} + + + +
+
+
+
+
+ ); +} diff --git a/apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx b/apps/code/src/renderer/features/billing/components/UsageLimitModal.tsx new file mode 100644 index 000000000..17fe34686 --- /dev/null +++ b/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 ( + + e.preventDefault()} + onEscapeKeyDown={hide} + > + + + + Usage limit reached + + + + {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."} + + + + + + + + + + ); +} diff --git a/apps/code/src/renderer/features/billing/hooks/useFreeUsage.ts b/apps/code/src/renderer/features/billing/hooks/useFreeUsage.ts new file mode 100644 index 000000000..fb20fa6f1 --- /dev/null +++ b/apps/code/src/renderer/features/billing/hooks/useFreeUsage.ts @@ -0,0 +1,14 @@ +import { useSeat } from "@hooks/useSeat"; +import type { UsageOutput } from "@main/services/llm-gateway/schemas"; +import { useUsage } from "./useUsage"; + +export function useFreeUsage(billingEnabled: boolean): UsageOutput | null { + const { seat, isPro } = useSeat(); + const seatLoaded = seat !== null; + const { usage } = useUsage({ + enabled: billingEnabled && seatLoaded && !isPro, + }); + + if (!billingEnabled || !seatLoaded || isPro || !usage) return null; + return usage; +} diff --git a/apps/code/src/renderer/features/billing/hooks/useUsage.ts b/apps/code/src/renderer/features/billing/hooks/useUsage.ts new file mode 100644 index 000000000..4cd77bf0e --- /dev/null +++ b/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 }; +} diff --git a/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts b/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts new file mode 100644 index 000000000..0a3ef978f --- /dev/null +++ b/apps/code/src/renderer/features/billing/hooks/useUsageLimitDetection.ts @@ -0,0 +1,33 @@ +import { useUsageLimitStore } from "@features/billing/stores/usageLimitStore"; +import { isUsageExceeded } from "@features/billing/utils"; +import { useSessionStore } from "@features/sessions/stores/sessionStore"; +import { useEffect, useRef } from "react"; +import { useFreeUsage } from "./useFreeUsage"; + +export function useUsageLimitDetection(billingEnabled: boolean) { + const usage = useFreeUsage(billingEnabled); + const hasAlertedRef = useRef(false); + + useEffect(() => { + if (!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; + } + }, [usage]); +} diff --git a/apps/code/src/renderer/features/billing/stores/seatStore.test.ts b/apps/code/src/renderer/features/billing/stores/seatStore.test.ts index 974f93404..84c707e47 100644 --- a/apps/code/src/renderer/features/billing/stores/seatStore.test.ts +++ b/apps/code/src/renderer/features/billing/stores/seatStore.test.ts @@ -2,13 +2,8 @@ import type { SeatData } from "@shared/types/seat"; import { PLAN_FREE, PLAN_PRO } from "@shared/types/seat"; import { beforeEach, describe, expect, it, vi } from "vitest"; -const mockIsFeatureFlagEnabled = vi.hoisted(() => vi.fn()); const mockGetAuthenticatedClient = vi.hoisted(() => vi.fn()); -vi.mock("@utils/analytics", () => ({ - isFeatureFlagEnabled: mockIsFeatureFlagEnabled, -})); - vi.mock("@features/auth/hooks/authClient", () => ({ getAuthenticatedClient: mockGetAuthenticatedClient, })); @@ -45,8 +40,21 @@ vi.mock("@utils/urls", () => ({ getPostHogUrl: (path: string) => `https://posthog.com${path}`, })); +vi.mock("@renderer/trpc", () => ({ + trpcClient: { + llmGateway: { + invalidatePlanCache: { mutate: vi.fn().mockResolvedValue(undefined) }, + }, + }, +})); + +import { trpcClient } from "@renderer/trpc"; import { useSeatStore } from "./seatStore"; +const mockInvalidatePlanCache = vi.mocked( + trpcClient.llmGateway.invalidatePlanCache.mutate, +); + function makeSeat(overrides: Partial = {}): SeatData { return { id: 1, @@ -86,61 +94,8 @@ describe("seatStore", () => { }); }); - describe("billing flag gate", () => { - it("fetchSeat does not call API when billing is disabled", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(false); - const client = mockClient(); - - await useSeatStore.getState().fetchSeat({ autoProvision: true }); - - expect(client.getMySeat).not.toHaveBeenCalled(); - expect(client.createSeat).not.toHaveBeenCalled(); - expect(useSeatStore.getState().seat).toBeNull(); - expect(useSeatStore.getState().error).toBe("Billing is not enabled"); - }); - - it("provisionFreeSeat does not call API when billing is disabled", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(false); - const client = mockClient(); - - await useSeatStore.getState().provisionFreeSeat(); - - expect(client.getMySeat).not.toHaveBeenCalled(); - expect(client.createSeat).not.toHaveBeenCalled(); - }); - - it("upgradeToPro does not call API when billing is disabled", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(false); - const client = mockClient(); - - await useSeatStore.getState().upgradeToPro(); - - expect(client.getMySeat).not.toHaveBeenCalled(); - expect(client.upgradeSeat).not.toHaveBeenCalled(); - }); - - it("cancelSeat does not call API when billing is disabled", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(false); - const client = mockClient(); - - await useSeatStore.getState().cancelSeat(); - - expect(client.cancelSeat).not.toHaveBeenCalled(); - }); - - it("reactivateSeat does not call API when billing is disabled", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(false); - const client = mockClient(); - - await useSeatStore.getState().reactivateSeat(); - - expect(client.reactivateSeat).not.toHaveBeenCalled(); - }); - }); - describe("fetchSeat", () => { it("fetches existing seat", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const seat = makeSeat(); mockClient({ getMySeat: vi.fn().mockResolvedValue(seat) }); @@ -152,7 +107,6 @@ describe("seatStore", () => { }); it("auto-provisions free seat when none exists", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const seat = makeSeat(); const client = mockClient({ getMySeat: vi.fn().mockResolvedValue(null), @@ -166,7 +120,6 @@ describe("seatStore", () => { }); it("does not auto-provision when option is false", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const client = mockClient(); await useSeatStore.getState().fetchSeat(); @@ -178,7 +131,6 @@ describe("seatStore", () => { describe("provisionFreeSeat", () => { it("creates free seat when none exists", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const seat = makeSeat(); const client = mockClient({ createSeat: vi.fn().mockResolvedValue(seat), @@ -188,10 +140,10 @@ describe("seatStore", () => { expect(client.createSeat).toHaveBeenCalledWith(PLAN_FREE); expect(useSeatStore.getState().seat).toEqual(seat); + expect(mockInvalidatePlanCache).toHaveBeenCalled(); }); it("uses existing seat instead of creating", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const existing = makeSeat(); const client = mockClient({ getMySeat: vi.fn().mockResolvedValue(existing), @@ -201,12 +153,12 @@ describe("seatStore", () => { expect(client.createSeat).not.toHaveBeenCalled(); expect(useSeatStore.getState().seat).toEqual(existing); + expect(mockInvalidatePlanCache).not.toHaveBeenCalled(); }); }); describe("upgradeToPro", () => { it("upgrades existing free seat to pro", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const freeSeat = makeSeat({ plan_key: PLAN_FREE }); const proSeat = makeSeat({ plan_key: PLAN_PRO }); const client = mockClient({ @@ -218,10 +170,10 @@ describe("seatStore", () => { expect(client.upgradeSeat).toHaveBeenCalledWith(PLAN_PRO); expect(useSeatStore.getState().seat).toEqual(proSeat); + expect(mockInvalidatePlanCache).toHaveBeenCalled(); }); it("no-ops when already on pro", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const proSeat = makeSeat({ plan_key: PLAN_PRO }); const client = mockClient({ getMySeat: vi.fn().mockResolvedValue(proSeat), @@ -235,7 +187,6 @@ describe("seatStore", () => { }); it("creates pro seat when none exists", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const proSeat = makeSeat({ plan_key: PLAN_PRO }); const client = mockClient({ createSeat: vi.fn().mockResolvedValue(proSeat), @@ -244,12 +195,12 @@ describe("seatStore", () => { await useSeatStore.getState().upgradeToPro(); expect(client.createSeat).toHaveBeenCalledWith(PLAN_PRO); + expect(mockInvalidatePlanCache).toHaveBeenCalled(); }); }); describe("cancelSeat", () => { it("cancels and re-fetches seat", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const canceledSeat = makeSeat({ status: "canceling" }); const client = mockClient({ getMySeat: vi.fn().mockResolvedValue(canceledSeat), @@ -259,12 +210,12 @@ describe("seatStore", () => { expect(client.cancelSeat).toHaveBeenCalled(); expect(useSeatStore.getState().seat).toEqual(canceledSeat); + expect(mockInvalidatePlanCache).toHaveBeenCalled(); }); }); describe("reactivateSeat", () => { it("reactivates seat", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const seat = makeSeat({ status: "active" }); mockClient({ reactivateSeat: vi.fn().mockResolvedValue(seat), @@ -273,12 +224,12 @@ describe("seatStore", () => { await useSeatStore.getState().reactivateSeat(); expect(useSeatStore.getState().seat).toEqual(seat); + expect(mockInvalidatePlanCache).toHaveBeenCalled(); }); }); describe("error handling", () => { it("sets redirect URL on subscription required error", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const { SeatSubscriptionRequiredError } = await import( "@renderer/api/posthogClient" ); @@ -300,7 +251,6 @@ describe("seatStore", () => { }); it("sets error on payment failure", async () => { - mockIsFeatureFlagEnabled.mockReturnValue(true); const { SeatPaymentFailedError } = await import( "@renderer/api/posthogClient" ); @@ -314,6 +264,16 @@ describe("seatStore", () => { expect(useSeatStore.getState().error).toBe("Card declined"); }); + + it("does not invalidate plan cache on failure", async () => { + mockClient({ + getMySeat: vi.fn().mockRejectedValue(new Error("Network error")), + }); + + await useSeatStore.getState().upgradeToPro(); + + expect(mockInvalidatePlanCache).not.toHaveBeenCalled(); + }); }); describe("reset", () => { diff --git a/apps/code/src/renderer/features/billing/stores/seatStore.ts b/apps/code/src/renderer/features/billing/stores/seatStore.ts index 17a38590b..fb3d13acb 100644 --- a/apps/code/src/renderer/features/billing/stores/seatStore.ts +++ b/apps/code/src/renderer/features/billing/stores/seatStore.ts @@ -3,9 +3,9 @@ import { SeatPaymentFailedError, SeatSubscriptionRequiredError, } from "@renderer/api/posthogClient"; +import { trpcClient } from "@renderer/trpc"; import type { SeatData } from "@shared/types/seat"; import { PLAN_FREE, PLAN_PRO } from "@shared/types/seat"; -import { isFeatureFlagEnabled } from "@utils/analytics"; import { logger } from "@utils/logger"; import { getPostHogUrl } from "@utils/urls"; import { create } from "zustand"; @@ -31,14 +31,6 @@ interface SeatStoreActions { type SeatStore = SeatStoreState & SeatStoreActions; -const BILLING_FLAG = "posthog-code-billing"; - -function assertBillingEnabled(): void { - if (!isFeatureFlagEnabled(BILLING_FLAG)) { - throw new Error("Billing is not enabled"); - } -} - async function getClient() { const client = await getAuthenticatedClient(); if (!client) { @@ -75,6 +67,12 @@ function handleSeatError( set({ isLoading: false, error: error.message }); } +function invalidatePlanCache(): void { + trpcClient.llmGateway.invalidatePlanCache.mutate().catch((err) => { + log.warn("Failed to invalidate plan cache", err); + }); +} + const initialState: SeatStoreState = { seat: null, isLoading: false, @@ -88,7 +86,6 @@ export const useSeatStore = create()((set) => ({ fetchSeat: async (options?: { autoProvision?: boolean }) => { set({ isLoading: true, error: null, redirectUrl: null }); try { - assertBillingEnabled(); const client = await getClient(); let seat = await client.getMySeat(); if (!seat && options?.autoProvision) { @@ -105,7 +102,6 @@ export const useSeatStore = create()((set) => ({ log.info("Provisioning free seat"); set({ isLoading: true, error: null, redirectUrl: null }); try { - assertBillingEnabled(); const client = await getClient(); const existing = await client.getMySeat(); if (existing) { @@ -119,6 +115,7 @@ export const useSeatStore = create()((set) => ({ const seat = await client.createSeat(PLAN_FREE); log.info("Free seat created", { id: seat.id, plan: seat.plan_key }); set({ seat, isLoading: false }); + invalidatePlanCache(); } catch (error) { log.error("provisionFreeSeat failed", error); handleSeatError(error, set); @@ -128,7 +125,6 @@ export const useSeatStore = create()((set) => ({ upgradeToPro: async () => { set({ isLoading: true, error: null, redirectUrl: null }); try { - assertBillingEnabled(); const client = await getClient(); const existing = await client.getMySeat(); if (existing) { @@ -138,10 +134,12 @@ export const useSeatStore = create()((set) => ({ } const seat = await client.upgradeSeat(PLAN_PRO); set({ seat, isLoading: false }); + invalidatePlanCache(); return; } const seat = await client.createSeat(PLAN_PRO); set({ seat, isLoading: false }); + invalidatePlanCache(); } catch (error) { handleSeatError(error, set); } @@ -150,11 +148,11 @@ export const useSeatStore = create()((set) => ({ cancelSeat: async () => { set({ isLoading: true, error: null, redirectUrl: null }); try { - assertBillingEnabled(); const client = await getClient(); await client.cancelSeat(); const seat = await client.getMySeat(); set({ seat, isLoading: false }); + invalidatePlanCache(); } catch (error) { handleSeatError(error, set); } @@ -163,10 +161,10 @@ export const useSeatStore = create()((set) => ({ reactivateSeat: async () => { set({ isLoading: true, error: null, redirectUrl: null }); try { - assertBillingEnabled(); const client = await getClient(); const seat = await client.reactivateSeat(); set({ seat, isLoading: false }); + invalidatePlanCache(); } catch (error) { 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 new file mode 100644 index 000000000..f97834b76 --- /dev/null +++ b/apps/code/src/renderer/features/billing/stores/usageLimitStore.test.ts @@ -0,0 +1,36 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { useUsageLimitStore } from "./usageLimitStore"; + +describe("usageLimitStore", () => { + beforeEach(() => { + useUsageLimitStore.setState({ isOpen: false, context: null }); + }); + + it("starts closed with no context", () => { + 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("hide closes but preserves context for exit animation", () => { + useUsageLimitStore.getState().show("mid-task"); + useUsageLimitStore.getState().hide(); + const state = useUsageLimitStore.getState(); + expect(state.isOpen).toBe(false); + expect(state.context).toBe("mid-task"); + }); +}); diff --git a/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts b/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts new file mode 100644 index 000000000..7fd54234a --- /dev/null +++ b/apps/code/src/renderer/features/billing/stores/usageLimitStore.ts @@ -0,0 +1,23 @@ +import { create } from "zustand"; + +type UsageLimitContext = "mid-task" | "idle"; + +interface UsageLimitState { + isOpen: boolean; + context: UsageLimitContext | null; +} + +interface UsageLimitActions { + show: (context: UsageLimitContext) => void; + hide: () => void; +} + +type UsageLimitStore = UsageLimitState & UsageLimitActions; + +export const useUsageLimitStore = create()((set) => ({ + isOpen: false, + context: null, + + show: (context) => set({ isOpen: true, context }), + hide: () => set({ isOpen: false }), +})); diff --git a/apps/code/src/renderer/features/billing/utils.test.ts b/apps/code/src/renderer/features/billing/utils.test.ts new file mode 100644 index 000000000..8e8db3c3f --- /dev/null +++ b/apps/code/src/renderer/features/billing/utils.test.ts @@ -0,0 +1,53 @@ +import type { UsageOutput } from "@main/services/llm-gateway/schemas"; +import { describe, expect, it } from "vitest"; +import { isUsageExceeded } from "./utils"; + +function makeUsage( + overrides: Partial<{ + sustained: boolean; + burst: boolean; + isRateLimited: boolean; + }> = {}, +): UsageOutput { + return { + product: "posthog_code", + user_id: 1, + sustained: { + used_percent: 50, + resets_in_seconds: 3600, + exceeded: overrides.sustained ?? false, + }, + burst: { + used_percent: 30, + resets_in_seconds: 600, + exceeded: overrides.burst ?? false, + }, + is_rate_limited: overrides.isRateLimited ?? false, + }; +} + +describe("isUsageExceeded", () => { + it("returns false when nothing is exceeded", () => { + expect(isUsageExceeded(makeUsage())).toBe(false); + }); + + it("returns true when sustained is exceeded", () => { + expect(isUsageExceeded(makeUsage({ sustained: true }))).toBe(true); + }); + + it("returns true when burst is exceeded", () => { + expect(isUsageExceeded(makeUsage({ burst: true }))).toBe(true); + }); + + it("returns true when rate limited", () => { + expect(isUsageExceeded(makeUsage({ isRateLimited: true }))).toBe(true); + }); + + it("returns true when all flags are set", () => { + expect( + isUsageExceeded( + makeUsage({ sustained: true, burst: true, isRateLimited: true }), + ), + ).toBe(true); + }); +}); diff --git a/apps/code/src/renderer/features/billing/utils.ts b/apps/code/src/renderer/features/billing/utils.ts new file mode 100644 index 000000000..f0ad86830 --- /dev/null +++ b/apps/code/src/renderer/features/billing/utils.ts @@ -0,0 +1,7 @@ +import type { UsageOutput } from "@main/services/llm-gateway/schemas"; + +export function isUsageExceeded(usage: UsageOutput): boolean { + return ( + usage.is_rate_limited || usage.sustained.exceeded || usage.burst.exceeded + ); +} diff --git a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx index 643585e0d..25d0e5e3d 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx @@ -5,12 +5,8 @@ import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { useTRPC } from "@renderer/trpc/client"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; -import { logger } from "@utils/logger"; -import { useEffect } from "react"; import type { ToolViewProps } from "./toolCallUtils"; -const log = logger.scope("mcp-tool-block"); - interface McpToolBlockProps extends ToolViewProps { mcpToolName: string; } @@ -35,15 +31,6 @@ export function McpToolBlock(props: McpToolBlockProps) { ), ); - // TODO: Remove this, used for local debugging only - useEffect(() => { - log.debug("McpToolBlock render", { - mcpToolName, - hasUi, - isDisabledForServer, - }); - }, [mcpToolName, hasUi, isDisabledForServer]); - // When MCP Apps discovery completes (possibly after this component mounted), // invalidate the hasUiForTool query so we pick up newly-discovered UIs. useSubscription( diff --git a/apps/code/src/renderer/features/sessions/service/service.ts b/apps/code/src/renderer/features/sessions/service/service.ts index 2146a877a..41e1ba26a 100644 --- a/apps/code/src/renderer/features/sessions/service/service.ts +++ b/apps/code/src/renderer/features/sessions/service/service.ts @@ -2126,35 +2126,6 @@ export class SessionService { await this.appendAndPersist(taskId, session, event, storedEntry); } - /** - * Append a user shell execute event (synchronous version for backwards compatibility). - */ - async appendUserShellExecute( - taskId: string, - command: string, - cwd: string, - result: { stdout: string; stderr: string; exitCode: number }, - ): Promise { - const id = `user-shell-${Date.now()}-${Math.random() - .toString(36) - .slice(2, 9)}`; - const session = sessionStoreSetters.getSessionByTaskId(taskId); - if (!session) return; - - const storedEntry: StoredLogEntry = { - type: "notification", - timestamp: new Date().toISOString(), - notification: { - method: "_array/user_shell_execute", - params: { id, command, cwd, result }, - }, - }; - - const event = createUserShellExecuteEvent(command, cwd, result, id); - - await this.appendAndPersist(taskId, session, event, storedEntry); - } - /** * Retry connecting to the existing session (resume attempt using * the sessionId from logs). Does NOT tear down — avoids the connect diff --git a/apps/code/src/renderer/features/settings/components/SettingsDialog.tsx b/apps/code/src/renderer/features/settings/components/SettingsDialog.tsx index 10ea83948..b16165af0 100644 --- a/apps/code/src/renderer/features/settings/components/SettingsDialog.tsx +++ b/apps/code/src/renderer/features/settings/components/SettingsDialog.tsx @@ -29,6 +29,7 @@ import { Wrench, } from "@phosphor-icons/react"; import { Avatar, Box, Flex, ScrollArea, Text } from "@radix-ui/themes"; +import { BILLING_FLAG } from "@shared/constants"; import { type ReactNode, useEffect, useMemo } from "react"; import { useHotkeys } from "react-hotkeys-hook"; import { AdvancedSettings } from "./sections/AdvancedSettings"; @@ -128,7 +129,7 @@ export function SettingsDialog() { const client = useOptionalAuthenticatedClient(); const { data: user } = useCurrentUser({ client }); const { seat, planLabel } = useSeat(); - const billingEnabled = useFeatureFlag("posthog-code-billing"); + const billingEnabled = useFeatureFlag(BILLING_FLAG); const logoutMutation = useLogoutMutation(); const sidebarItems = useMemo( 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 438dfe384..10f00e916 100644 --- a/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx +++ b/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx @@ -1,4 +1,6 @@ +import { useUsage } from "@features/billing/hooks/useUsage"; import { useSeatStore } from "@features/billing/stores/seatStore"; +import { useFeatureFlag } from "@hooks/useFeatureFlag"; import { useSeat } from "@hooks/useSeat"; import { ArrowSquareOut, @@ -16,8 +18,7 @@ import { Text, } from "@radix-ui/themes"; import { Tooltip } from "@renderer/components/ui/Tooltip"; -import { useTRPC } from "@renderer/trpc"; -import { useQuery } from "@tanstack/react-query"; +import { BILLING_FLAG } from "@shared/constants"; import { getPostHogUrl } from "@utils/urls"; import { useState } from "react"; @@ -38,14 +39,6 @@ function formatResetTime(seconds: number): string { return `${days} days`; } -function useUsage() { - const trpc = useTRPC(); - const { data: usage, isLoading } = useQuery( - trpc.llmGateway.usage.queryOptions(), - ); - return { usage: usage ?? null, isLoading }; -} - export function PlanUsageSettings() { const { seat, @@ -56,10 +49,13 @@ export function PlanUsageSettings() { error, redirectUrl, } = useSeat(); + const billingEnabled = useFeatureFlag(BILLING_FLAG); const { upgradeToPro, cancelSeat, reactivateSeat, clearError } = useSeatStore(); const [showUpgradeDialog, setShowUpgradeDialog] = useState(false); - const { usage, isLoading: usageLoading } = useUsage(); + const { usage, isLoading: usageLoading } = useUsage({ + enabled: seat !== null && !isPro, + }); const formattedActiveUntil = activeUntil ? activeUntil.toLocaleDateString(undefined, { @@ -195,7 +191,7 @@ export function PlanUsageSettings() { borderRadius: "var(--radius-3)", }} > - {isLoading ? ( + {isLoading || !billingEnabled ? ( ) : ( diff --git a/apps/code/src/renderer/features/sidebar/components/SidebarContent.tsx b/apps/code/src/renderer/features/sidebar/components/SidebarContent.tsx index 42cabf8c5..81dc03740 100644 --- a/apps/code/src/renderer/features/sidebar/components/SidebarContent.tsx +++ b/apps/code/src/renderer/features/sidebar/components/SidebarContent.tsx @@ -1,4 +1,5 @@ import { useArchivedTaskIds } from "@features/archive/hooks/useArchivedTaskIds"; +import { SidebarUsageBar } from "@features/billing/components/SidebarUsageBar"; import { ArchiveIcon } from "@phosphor-icons/react"; import { Box, Flex } from "@radix-ui/themes"; import { useNavigationStore } from "@stores/navigationStore"; @@ -12,13 +13,13 @@ export const SidebarContent: React.FC = () => { const navigateToArchived = useNavigationStore( (state) => state.navigateToArchived, ); - return ( + {archivedTaskIds.size > 0 && (