diff --git a/apps/code/src/renderer/api/posthogClient.ts b/apps/code/src/renderer/api/posthogClient.ts index b564aafda4..6fbb382b5f 100644 --- a/apps/code/src/renderer/api/posthogClient.ts +++ b/apps/code/src/renderer/api/posthogClient.ts @@ -56,6 +56,29 @@ export class SeatPaymentFailedError extends Error { } } +export type UsageLimitType = "burst" | "sustained" | null; + +// Stable message so callers recognize this after a saga reduces the error to a string. +export const CLOUD_USAGE_LIMIT_ERROR_MESSAGE = "Cloud usage limit reached"; + +/** Thrown when the backend rejects a cloud run with a 429 usage-limit error. */ +export class CloudUsageLimitError extends Error { + limitType: UsageLimitType; + resetAt: string | null; + isPro: boolean; + constructor(params: { + limitType: UsageLimitType; + resetAt: string | null; + isPro: boolean; + }) { + super(CLOUD_USAGE_LIMIT_ERROR_MESSAGE); + this.name = "CloudUsageLimitError"; + this.limitType = params.limitType; + this.resetAt = params.resetAt; + this.isPro = params.isPro; + } +} + const log = logger.scope("posthog-client"); export const MCP_CATEGORIES = [ @@ -1106,12 +1129,11 @@ export class PostHogAPIClient { mode: "interactive", }); - const data = await this.api.post( - `/api/projects/{project_id}/tasks/{id}/run/`, - { + const data = await this.withCloudUsageLimitCheck(() => + this.api.post(`/api/projects/{project_id}/tasks/{id}/run/`, { path: { project_id: teamId.toString(), id: taskId }, body, - }, + }), ); return data as unknown as Task; @@ -1328,20 +1350,22 @@ export class PostHogAPIClient { const url = new URL( `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/`, ); - const response = await this.api.fetcher.fetch({ - method: "post", - url, - path: `/api/projects/${teamId}/tasks/${taskId}/runs/`, - overrides: { - body: JSON.stringify({ - ...buildCloudRunRequestBody({ - ...options, - mode: options?.mode ?? "background", + const response = await this.withCloudUsageLimitCheck(() => + this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/runs/`, + overrides: { + body: JSON.stringify({ + ...buildCloudRunRequestBody({ + ...options, + mode: options?.mode ?? "background", + }), + environment: options?.environment ?? "local", }), - environment: options?.environment ?? "local", - }), - }, - }); + }, + }), + ); if (!response.ok) { throw new Error(`Failed to create task run: ${response.statusText}`); @@ -1359,17 +1383,19 @@ export class PostHogAPIClient { const url = new URL( `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/start/`, ); - const response = await this.api.fetcher.fetch({ - method: "post", - url, - path: `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/start/`, - overrides: { - body: JSON.stringify({ - pending_user_message: options?.pendingUserMessage, - pending_user_artifact_ids: options?.pendingUserArtifactIds, - }), - }, - }); + const response = await this.withCloudUsageLimitCheck(() => + this.api.fetcher.fetch({ + method: "post", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/start/`, + overrides: { + body: JSON.stringify({ + pending_user_message: options?.pendingUserMessage, + pending_user_artifact_ids: options?.pendingUserArtifactIds, + }), + }, + }), + ); if (!response.ok) { throw new Error(`Failed to start task run: ${response.statusText}`); @@ -2777,6 +2803,37 @@ export class PostHogAPIClient { } } + /** + * Run a cloud-run request, re-throwing a backend 429 usage-limit error as a + * typed CloudUsageLimitError so the UI can show the upgrade prompt. + */ + private async withCloudUsageLimitCheck(fn: () => Promise): Promise { + try { + return await fn(); + } catch (error) { + const parsed = this.parseFetcherError(error); + if ( + parsed && + parsed.status === 429 && + parsed.body.code === "usage_limit_exceeded" + ) { + const limitType = parsed.body.limit_type; + throw new CloudUsageLimitError({ + limitType: + limitType === "burst" || limitType === "sustained" + ? limitType + : null, + resetAt: + typeof parsed.body.reset_at === "string" + ? parsed.body.reset_at + : null, + isPro: parsed.body.is_pro === true, + }); + } + throw error; + } + } + private throwSeatError(error: unknown): never { const parsed = this.parseFetcherError(error); diff --git a/apps/code/src/renderer/features/billing/preflightCloudUsage.test.ts b/apps/code/src/renderer/features/billing/preflightCloudUsage.test.ts new file mode 100644 index 0000000000..1cbd10c2a0 --- /dev/null +++ b/apps/code/src/renderer/features/billing/preflightCloudUsage.test.ts @@ -0,0 +1,151 @@ +import type { UsageOutput } from "@main/services/llm-gateway/schemas"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const refresh = vi.fn(); +const getLatest = vi.fn(); +const track = vi.fn(); + +vi.mock("@renderer/trpc/client", () => ({ + trpcClient: { + usageMonitor: { + refresh: { mutate: () => refresh() }, + getLatest: { query: () => getLatest() }, + }, + }, +})); + +vi.mock("@utils/analytics", () => ({ + track: (...args: unknown[]) => track(...args), +})); + +import { ANALYTICS_EVENTS } from "@shared/types/analytics"; +import { assertCloudUsageAvailable } from "./preflightCloudUsage"; +import { useUsageLimitStore } from "./stores/usageLimitStore"; + +function makeUsage( + overrides: Partial<{ + sustained: boolean; + burst: boolean; + isRateLimited: boolean; + isPro: boolean; + }> = {}, +): UsageOutput { + return { + product: "posthog_code", + user_id: 1, + sustained: { + used_percent: 50, + reset_at: "2026-05-01T13:00:00.000Z", + exceeded: overrides.sustained ?? false, + }, + burst: { + used_percent: 30, + reset_at: "2026-05-01T12:10:00.000Z", + exceeded: overrides.burst ?? false, + }, + is_rate_limited: overrides.isRateLimited ?? false, + is_pro: overrides.isPro ?? false, + }; +} + +interface Case { + name: string; + arrange: () => void; + available: boolean; + modal: { + isOpen: boolean; + bucket?: "burst" | "sustained" | null; + resetAt?: string | null; + isPro?: boolean | null; + }; + trackPayload?: { bucket: "burst" | "sustained" | null; is_pro: boolean }; +} + +const cases: Case[] = [ + { + name: "allows creation and shows no modal when under limit", + arrange: () => refresh.mockResolvedValue(makeUsage()), + available: true, + modal: { isOpen: false }, + }, + { + name: "blocks and shows the burst modal when the daily limit is exceeded", + arrange: () => + refresh.mockResolvedValue( + makeUsage({ burst: true, isRateLimited: true, isPro: true }), + ), + available: false, + modal: { + isOpen: true, + bucket: "burst", + resetAt: "2026-05-01T12:10:00.000Z", + isPro: true, + }, + trackPayload: { bucket: "burst", is_pro: true }, + }, + { + name: "falls back to the latest snapshot when refresh fails", + arrange: () => { + refresh.mockRejectedValue(new Error("network")); + getLatest.mockResolvedValue(makeUsage({ sustained: true })); + }, + available: false, + modal: { isOpen: true, bucket: "sustained" }, + trackPayload: { bucket: "sustained", is_pro: false }, + }, + { + name: "falls back to the monthly bucket when only is_rate_limited is set", + arrange: () => + refresh.mockResolvedValue(makeUsage({ isRateLimited: true })), + available: false, + modal: { + isOpen: true, + bucket: "sustained", + resetAt: "2026-05-01T13:00:00.000Z", + }, + trackPayload: { bucket: "sustained", is_pro: false }, + }, + { + name: "fails open (allows creation) when usage cannot be fetched", + arrange: () => { + refresh.mockRejectedValue(new Error("network")); + getLatest.mockRejectedValue(new Error("network")); + }, + available: true, + modal: { isOpen: false }, + }, +]; + +describe("assertCloudUsageAvailable", () => { + beforeEach(() => { + refresh.mockReset(); + getLatest.mockReset(); + track.mockReset(); + useUsageLimitStore.getState().hide(); + }); + + it.each(cases)( + "$name", + async ({ arrange, available, modal, trackPayload }) => { + arrange(); + + expect(await assertCloudUsageAvailable()).toBe(available); + + const state = useUsageLimitStore.getState(); + expect(state.isOpen).toBe(modal.isOpen); + if (modal.bucket !== undefined) expect(state.bucket).toBe(modal.bucket); + if (modal.resetAt !== undefined) + expect(state.resetAt).toBe(modal.resetAt); + if (modal.isPro !== undefined) expect(state.isPro).toBe(modal.isPro); + + if (trackPayload) { + expect(track).toHaveBeenCalledWith( + ANALYTICS_EVENTS.CLOUD_TASK_USAGE_BLOCKED, + trackPayload, + ); + } else { + expect(track).not.toHaveBeenCalled(); + } + }, + ); +}); diff --git a/apps/code/src/renderer/features/billing/preflightCloudUsage.ts b/apps/code/src/renderer/features/billing/preflightCloudUsage.ts new file mode 100644 index 0000000000..cea8ff8728 --- /dev/null +++ b/apps/code/src/renderer/features/billing/preflightCloudUsage.ts @@ -0,0 +1,62 @@ +import { + type UsageLimitBucket, + useUsageLimitStore, +} from "@features/billing/stores/usageLimitStore"; +import { isUsageExceeded } from "@features/billing/utils"; +import type { UsageOutput } from "@main/services/llm-gateway/schemas"; +import { trpcClient } from "@renderer/trpc/client"; +import { ANALYTICS_EVENTS } from "@shared/types/analytics"; +import { track } from "@utils/analytics"; +import { logger } from "@utils/logger"; + +const log = logger.scope("preflight-cloud-usage"); + +function usageLimitArgs(usage: UsageOutput): { + bucket: UsageLimitBucket; + resetAt: string; + isPro: boolean; +} { + // Prefer the bucket that's actually exceeded (burst/daily takes priority); if neither + // is flagged (is_rate_limited via a server-side valve), fall back to the monthly bucket + // so the modal still shows a title and reset time rather than a bare prompt. + const bucket: UsageLimitBucket = usage.burst.exceeded ? "burst" : "sustained"; + return { bucket, resetAt: usage[bucket].reset_at, isPro: usage.is_pro }; +} + +async function fetchUsageSnapshot(): Promise { + const fresh = await trpcClient.usageMonitor.refresh + .mutate() + .catch((error) => { + log.warn("Usage refresh failed; falling back to latest snapshot", { + error, + }); + return null; + }); + if (fresh) return fresh; + + return trpcClient.usageMonitor.getLatest.query().catch((error) => { + log.warn("Usage lookup failed; allowing cloud creation", { error }); + return null; + }); +} + +/** + * Pre-flight gate for cloud task creation. Returns false (and shows the upgrade + * modal) when the team is over its usage limit, so no cloud task/run is created. + * + * Best-effort: if usage can't be fetched, returns true (fail open) — a usage + * service hiccup must never block task creation. + */ +export async function assertCloudUsageAvailable(): Promise { + const usage = await fetchUsageSnapshot(); + if (usage && isUsageExceeded(usage)) { + const args = usageLimitArgs(usage); + track(ANALYTICS_EVENTS.CLOUD_TASK_USAGE_BLOCKED, { + bucket: args.bucket, + is_pro: usage.is_pro, + }); + useUsageLimitStore.getState().show(args); + return false; + } + return true; +} diff --git a/apps/code/src/renderer/features/inbox/hooks/useCreatePrReport.ts b/apps/code/src/renderer/features/inbox/hooks/useCreatePrReport.ts index bba623899d..fa046a6a7c 100644 --- a/apps/code/src/renderer/features/inbox/hooks/useCreatePrReport.ts +++ b/apps/code/src/renderer/features/inbox/hooks/useCreatePrReport.ts @@ -12,9 +12,10 @@ import { track } from "@utils/analytics"; import { logger } from "@utils/logger"; import { useCallback, useState } from "react"; import { toast as sonnerToast } from "sonner"; -import type { - TaskCreationInput, - TaskService, +import { + isUsageLimitResult, + type TaskCreationInput, + type TaskService, } from "../../task-detail/service/service"; import { buildCreatePrReportPrompt } from "../utils/buildCreatePrReportPrompt"; import { resolveDefaultModel } from "../utils/resolveDefaultModel"; @@ -136,15 +137,18 @@ export function useCreatePrReport({ }); } else { sonnerToast.dismiss(toastId); - toast.error("Failed to start PR task", { - description: result.error, - }); - log.error("Create PR task creation failed", { - failedStep: result.failedStep, - error: result.error, - reportId, - reportTitle, - }); + // Usage-limit blocks already show the upgrade modal; don't also toast an error. + if (!isUsageLimitResult(result)) { + toast.error("Failed to start PR task", { + description: result.error, + }); + log.error("Create PR task creation failed", { + failedStep: result.failedStep, + error: result.error, + reportId, + reportTitle, + }); + } } } catch (error) { sonnerToast.dismiss(toastId); diff --git a/apps/code/src/renderer/features/inbox/hooks/useDiscussReport.ts b/apps/code/src/renderer/features/inbox/hooks/useDiscussReport.ts index 2da49cd517..1a491a7018 100644 --- a/apps/code/src/renderer/features/inbox/hooks/useDiscussReport.ts +++ b/apps/code/src/renderer/features/inbox/hooks/useDiscussReport.ts @@ -12,9 +12,10 @@ import { track } from "@utils/analytics"; import { logger } from "@utils/logger"; import { useCallback, useState } from "react"; import { toast as sonnerToast } from "sonner"; -import type { - TaskCreationInput, - TaskService, +import { + isUsageLimitResult, + type TaskCreationInput, + type TaskService, } from "../../task-detail/service/service"; import { buildDiscussReportPrompt } from "../utils/buildDiscussReportPrompt"; import { resolveDefaultModel } from "../utils/resolveDefaultModel"; @@ -138,15 +139,18 @@ export function useDiscussReport({ }); } else { sonnerToast.dismiss(toastId); - toast.error("Failed to start discussion", { - description: result.error, - }); - log.error("Discuss task creation failed", { - failedStep: result.failedStep, - error: result.error, - reportId, - reportTitle, - }); + // Usage-limit blocks already show the upgrade modal; don't also toast an error. + if (!isUsageLimitResult(result)) { + toast.error("Failed to start discussion", { + description: result.error, + }); + log.error("Discuss task creation failed", { + failedStep: result.failedStep, + error: result.error, + reportId, + reportTitle, + }); + } } } catch (error) { sonnerToast.dismiss(toastId); diff --git a/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts b/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts index d0caad5dad..2ce9e77e9f 100644 --- a/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts +++ b/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts @@ -1,4 +1,6 @@ import { useAuthStateValue } from "@features/auth/hooks/authQueries"; +import { assertCloudUsageAvailable } from "@features/billing/preflightCloudUsage"; +import { useUsageLimitStore } from "@features/billing/stores/usageLimitStore"; import { buildCloudTaskDescription } from "@features/editor/utils/cloud-prompt"; import { useTaskInputHistoryStore } from "@features/message-editor/stores/taskInputHistoryStore"; import type { EditorHandle } from "@features/message-editor/types"; @@ -28,7 +30,11 @@ import { useQuery } from "@tanstack/react-query"; import { track } from "@utils/analytics"; import { logger } from "@utils/logger"; import { useCallback, useState } from "react"; -import type { TaskCreationInput, TaskService } from "../service/service"; +import { + isUsageLimitResult, + type TaskCreationInput, + type TaskService, +} from "../service/service"; const log = logger.scope("task-creation"); @@ -232,6 +238,11 @@ export function useTaskCreation({ const allowSubmit = contentOverride ? canSubmitBase : canSubmit; if (!allowSubmit) return false; + // Block over-limit cloud creation before the pending view so it doesn't flash. + if (workspaceMode === "cloud" && !(await assertCloudUsageAvailable())) { + return false; + } + setIsCreatingTask(true); const content = contentOverride ?? editor.getContent(); @@ -285,24 +296,29 @@ export function useTaskCreation({ } const taskService = get(RENDERER_TOKENS.TaskService); - const result = await taskService.createTask(input, (output) => { - invalidateTasks(output.task); - if (signalReportId) { - clearTaskInputReportAssociation(); - } - if (pendingTaskKey) { - pendingTaskPromptStoreApi.move(pendingTaskKey, output.task.id); - } - if (onTaskCreated) { - onTaskCreated(output.task); - } else { - void openTask(output.task); - } - useTourStore.getState().completeTour(createFirstTaskTour.id); - if (!pendingTaskKey && !contentOverride) { - editor.clear(); - } - }); + const result = await taskService.createTask( + input, + (output) => { + invalidateTasks(output.task); + if (signalReportId) { + clearTaskInputReportAssociation(); + } + if (pendingTaskKey) { + pendingTaskPromptStoreApi.move(pendingTaskKey, output.task.id); + } + if (onTaskCreated) { + onTaskCreated(output.task); + } else { + void openTask(output.task); + } + useTourStore.getState().completeTour(createFirstTaskTour.id); + if (!pendingTaskKey && !contentOverride) { + editor.clear(); + } + // Pre-flight already ran above for cloud; skip the service's duplicate check. + }, + { skipCloudUsagePreflight: true }, + ); if (result.success) { setAdditionalDirectoriesOverride(null); @@ -310,12 +326,18 @@ export function useTaskCreation({ } if (!result.success) { - const title = getErrorTitle(result.failedStep); - toast.error(title, { description: result.error }); - log.error("Task creation failed", { - failedStep: result.failedStep, - error: result.error, - }); + // Usage-limit blocks already show the upgrade modal; don't also toast an error. + if (isUsageLimitResult(result)) { + useUsageLimitStore.getState().show(); + log.warn("Cloud task creation blocked by usage limit"); + } else { + const title = getErrorTitle(result.failedStep); + toast.error(title, { description: result.error }); + log.error("Task creation failed", { + failedStep: result.failedStep, + error: result.error, + }); + } if (pendingTaskKey) { pendingTaskPromptStoreApi.clear(pendingTaskKey); openTaskInput({ initialPrompt: plainPromptText }); diff --git a/apps/code/src/renderer/features/task-detail/service/service.ts b/apps/code/src/renderer/features/task-detail/service/service.ts index 11adeb8711..0e155b72dc 100644 --- a/apps/code/src/renderer/features/task-detail/service/service.ts +++ b/apps/code/src/renderer/features/task-detail/service/service.ts @@ -1,9 +1,11 @@ import { getAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { assertCloudUsageAvailable } from "@features/billing/preflightCloudUsage"; import { useDraftStore } from "@features/message-editor/stores/draftStore"; import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { workspaceApi } from "@features/workspace/hooks/useWorkspace"; import type { Workspace } from "@main/services/workspace/schemas"; import type { SagaResult } from "@posthog/shared"; +import { CLOUD_USAGE_LIMIT_ERROR_MESSAGE } from "@renderer/api/posthogClient"; import { type TaskCreationInput, type TaskCreationOutput, @@ -20,6 +22,18 @@ const log = logger.scope("task-service"); export type CreateTaskResult = SagaResult; +/** + * True when a failed createTask was blocked by the usage limit. The upgrade modal is + * already shown in this case, so callers should suppress their own error toast. + */ +export function isUsageLimitResult(result: CreateTaskResult): boolean { + return ( + !result.success && + (result.failedStep === "usage_limit" || + result.error === CLOUD_USAGE_LIMIT_ERROR_MESSAGE) + ); +} + @injectable() export class TaskService { /** @@ -33,6 +47,7 @@ export class TaskService { public async createTask( input: TaskCreationInput, onTaskReady?: (output: TaskCreationOutput) => void, + options?: { skipCloudUsagePreflight?: boolean }, ): Promise { log.info("Creating task", { workspaceMode: input.workspaceMode, @@ -57,6 +72,20 @@ export class TaskService { }; } + // Backstop for callers that bypass useTaskCreation (e.g. inbox); the helper shows the modal. + // Callers that already pre-flighted pass skipCloudUsagePreflight to avoid a second fetch. + if ( + !options?.skipCloudUsagePreflight && + input.workspaceMode === "cloud" && + !(await assertCloudUsageAvailable()) + ) { + return { + success: false, + error: CLOUD_USAGE_LIMIT_ERROR_MESSAGE, + failedStep: "usage_limit", + }; + } + const saga = new TaskCreationSaga({ posthogClient, onTaskReady: onTaskReady diff --git a/apps/code/src/shared/types/analytics.ts b/apps/code/src/shared/types/analytics.ts index df312790a5..16afb44ebf 100644 --- a/apps/code/src/shared/types/analytics.ts +++ b/apps/code/src/shared/types/analytics.ts @@ -654,6 +654,11 @@ export interface UpgradePromptClickedProperties { surface: UpgradePromptClickedSurface; } +export interface CloudTaskUsageBlockedProperties { + bucket: "burst" | "sustained" | null; + is_pro: boolean; +} + export interface SubscriptionStartedProperties { plan_key: string; previous_plan_key?: string; @@ -788,6 +793,7 @@ export const ANALYTICS_EVENTS = { // Subscription events UPGRADE_PROMPT_SHOWN: "Upgrade prompt shown", UPGRADE_PROMPT_CLICKED: "Upgrade prompt clicked", + CLOUD_TASK_USAGE_BLOCKED: "Cloud task usage blocked", SUBSCRIPTION_STARTED: "Subscription started", SUBSCRIPTION_CANCELLED: "Subscription cancelled", } as const; @@ -910,6 +916,7 @@ export type EventPropertyMap = { // Subscription events [ANALYTICS_EVENTS.UPGRADE_PROMPT_SHOWN]: UpgradePromptShownProperties; [ANALYTICS_EVENTS.UPGRADE_PROMPT_CLICKED]: UpgradePromptClickedProperties; + [ANALYTICS_EVENTS.CLOUD_TASK_USAGE_BLOCKED]: CloudTaskUsageBlockedProperties; [ANALYTICS_EVENTS.SUBSCRIPTION_STARTED]: SubscriptionStartedProperties; [ANALYTICS_EVENTS.SUBSCRIPTION_CANCELLED]: SubscriptionCancelledProperties; };