feat: Implement billing with seat management, plan/usage settings and onboarding flow (#1304)

Author: charlesvien

apps/code/src/main/db/repositories/auth-session-repository.ts

@@ -1,4 +1,4 @@
-import type { CloudRegion } from "@shared/types/oauth";
+import type { CloudRegion } from "@shared/types/regions";
 import { eq } from "drizzle-orm";
 import { inject, injectable } from "inversify";
 import { MAIN_TOKENS } from "../../di/tokens";

apps/code/src/main/services/auth/service.ts

@@ -1,9 +1,7 @@
-import {
-  getCloudUrlFromRegion,
-  OAUTH_SCOPE_VERSION,
-} from "@shared/constants/oauth";
-import type { CloudRegion } from "@shared/types/oauth";
+import { OAUTH_SCOPE_VERSION } from "@shared/constants/oauth";
+import type { CloudRegion } from "@shared/types/regions";
 import { type BackoffOptions, sleepWithBackoff } from "@shared/utils/backoff";
+import { getCloudUrlFromRegion } from "@shared/utils/urls";
 import { powerMonitor } from "electron";
 import { inject, injectable, postConstruct, preDestroy } from "inversify";
 import type { IAuthPreferenceRepository } from "../../db/repositories/auth-preference-repository";

apps/code/src/main/services/github-integration/service.ts

@@ -1,4 +1,4 @@
-import { getCloudUrlFromRegion } from "@shared/constants/oauth";
+import { getCloudUrlFromRegion } from "@shared/utils/urls";
 import { shell } from "electron";
 import { injectable } from "inversify";
 import { logger } from "../../utils/logger";

apps/code/src/main/services/linear-integration/service.ts

@@ -1,4 +1,4 @@
-import { getCloudUrlFromRegion } from "@shared/constants/oauth.js";
+import { getCloudUrlFromRegion } from "@shared/utils/urls.js";
 import { shell } from "electron";
 import { injectable } from "inversify";
 import { logger } from "../../utils/logger.js";

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

@@ -56,3 +56,20 @@ export interface AnthropicErrorResponse {
     code?: string;
   };
 }
+
+export const usageBucketSchema = z.object({
+  used_percent: z.number(),
+  resets_in_seconds: z.number(),
+  exceeded: z.boolean(),
+});
+
+export const usageOutput = z.object({
+  product: z.string(),
+  user_id: z.number(),
+  sustained: usageBucketSchema,
+  burst: usageBucketSchema,
+  is_rate_limited: z.boolean(),
+});
+
+export type UsageBucket = z.infer<typeof usageBucketSchema>;
+export type UsageOutput = z.infer<typeof usageOutput>;

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

@@ -1,15 +1,20 @@
-import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
+import {
+  getGatewayUsageUrl,
+  getLlmGatewayUrl,
+} from "@posthog/agent/posthog-api";
 import { net } from "electron";
 import { inject, injectable } from "inversify";
 import { MAIN_TOKENS } from "../../di/tokens";
 import { logger } from "../../utils/logger";
 import type { AuthService } from "../auth/service";
-import type {
-  AnthropicErrorResponse,
-  AnthropicMessagesRequest,
-  AnthropicMessagesResponse,
-  LlmMessage,
-  PromptOutput,
+import {
+  type AnthropicErrorResponse,
+  type AnthropicMessagesRequest,
+  type AnthropicMessagesResponse,
+  type LlmMessage,
+  type PromptOutput,
+  type UsageOutput,
+  usageOutput,
 } from "./schemas";
 
 const log = logger.scope("llm-gateway");
@@ -134,4 +139,27 @@ export class LlmGatewayService {
       },
     };
   }
+
+  async fetchUsage(): Promise<UsageOutput> {
+    const auth = await this.authService.getValidAccessToken();
+    const usageUrl = getGatewayUsageUrl(auth.apiHost);
+
+    log.debug("Fetching usage from gateway", { url: usageUrl });
+
+    const response = await this.authService.authenticatedFetch(
+      net.fetch,
+      usageUrl,
+    );
+
+    if (!response.ok) {
+      throw new LlmGatewayError(
+        `Failed to fetch usage: HTTP ${response.status}`,
+        "usage_error",
+        undefined,
+        response.status,
+      );
+    }
+
+    return usageOutput.parse(await response.json());
+  }
 }

apps/code/src/main/services/oauth/service.ts

@@ -2,10 +2,10 @@ import * as crypto from "node:crypto";
 import * as http from "node:http";
 import type { Socket } from "node:net";
 import {
-  getCloudUrlFromRegion,
   getOauthClientIdFromRegion,
   OAUTH_SCOPES,
 } from "@shared/constants/oauth";
+import { getCloudUrlFromRegion } from "@shared/utils/urls";
 import { shell } from "electron";
 import { inject, injectable } from "inversify";
 import { MAIN_TOKENS } from "../../di/tokens";

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

@@ -1,6 +1,10 @@
 import { container } from "../../di/container";
 import { MAIN_TOKENS } from "../../di/tokens";
-import { promptInput, promptOutput } from "../../services/llm-gateway/schemas";
+import {
+  promptInput,
+  promptOutput,
+  usageOutput,
+} from "../../services/llm-gateway/schemas";
 import type { LlmGatewayService } from "../../services/llm-gateway/service";
 import { publicProcedure, router } from "../trpc";
 
@@ -18,4 +22,8 @@ export const llmGatewayRouter = router({
         model: input.model,
       }),
     ),
+
+  usage: publicProcedure
+    .output(usageOutput)
+    .query(() => getService().fetchUsage()),
 });

apps/code/src/renderer/api/posthogClient.ts

@@ -1,5 +1,5 @@
 import { isSupportedReasoningEffort } from "@posthog/agent/adapters/reasoning-effort";
-import { type PermissionMode } from "@posthog/agent/execution-mode";
+import type { PermissionMode } from "@posthog/agent/execution-mode";
 import type {
   ActionabilityJudgmentArtefact,
   AvailableSuggestedReviewer,
@@ -24,11 +24,29 @@ import type {
   TaskRun,
 } from "@shared/types";
 import type { CloudRunSource, PrAuthorshipMode } from "@shared/types/cloud";
+import type { SeatData } from "@shared/types/seat";
+import { SEAT_PRODUCT_KEY } from "@shared/types/seat";
 import type { StoredLogEntry } from "@shared/types/session-events";
 import { logger } from "@utils/logger";
 import { buildApiFetcher } from "./fetcher";
 import { createApiClient, type Schemas } from "./generated";
 
+export class SeatSubscriptionRequiredError extends Error {
+  redirectUrl: string;
+  constructor(redirectUrl: string) {
+    super("Billing subscription required");
+    this.name = "SeatSubscriptionRequiredError";
+    this.redirectUrl = redirectUrl;
+  }
+}
+
+export class SeatPaymentFailedError extends Error {
+  constructor(message?: string) {
+    super(message ?? "Payment failed");
+    this.name = "SeatPaymentFailedError";
+  }
+}
+
 const log = logger.scope("posthog-client");
 
 export type McpRecommendedServer = Schemas.RecommendedServer;
@@ -1178,39 +1196,6 @@ export class PostHogAPIClient {
     return await response.json();
   }
 
-  /**
-   * Get billing information for a specific organization.
-   */
-  async getOrgBilling(orgId: string): Promise<{
-    has_active_subscription: boolean;
-    customer_id: string | null;
-  }> {
-    const url = new URL(
-      `${this.api.baseUrl}/api/organizations/${orgId}/billing/`,
-    );
-    const response = await this.api.fetcher.fetch({
-      method: "get",
-      url,
-      path: `/api/organizations/${orgId}/billing/`,
-    });
-
-    if (!response.ok) {
-      throw new Error(
-        `Failed to fetch organization billing: ${response.statusText}`,
-      );
-    }
-
-    const data = await response.json();
-    return {
-      has_active_subscription:
-        typeof data.has_active_subscription === "boolean"
-          ? data.has_active_subscription
-          : false,
-      customer_id:
-        typeof data.customer_id === "string" ? data.customer_id : null,
-    };
-  }
-
   async getSignalReports(
     params?: SignalReportsQueryParams,
   ): Promise<SignalReportsResponse> {
@@ -1741,6 +1726,145 @@ export class PostHogAPIClient {
     }
   }
 
+  async getMySeat(): Promise<SeatData | null> {
+    try {
+      const url = new URL(`${this.api.baseUrl}/api/seats/me/`);
+      url.searchParams.set("product_key", SEAT_PRODUCT_KEY);
+      const response = await this.api.fetcher.fetch({
+        method: "get",
+        url,
+        path: "/api/seats/me/",
+      });
+      return (await response.json()) as SeatData;
+    } catch (error) {
+      if (this.isFetcherStatusError(error, 404)) {
+        return null;
+      }
+      throw error;
+    }
+  }
+
+  async createSeat(planKey: string): Promise<SeatData> {
+    try {
+      const user = await this.getCurrentUser();
+      const distinctId = user.distinct_id;
+      if (!distinctId) {
+        throw new Error("Cannot create seat: user has no distinct_id");
+      }
+      const url = new URL(`${this.api.baseUrl}/api/seats/`);
+      const response = await this.api.fetcher.fetch({
+        method: "post",
+        url,
+        path: "/api/seats/",
+        overrides: {
+          body: JSON.stringify({
+            product_key: SEAT_PRODUCT_KEY,
+            plan_key: planKey,
+            user_distinct_id: distinctId,
+          }),
+        },
+      });
+      return (await response.json()) as SeatData;
+    } catch (error) {
+      this.throwSeatError(error);
+    }
+  }
+
+  async upgradeSeat(planKey: string): Promise<SeatData> {
+    try {
+      const url = new URL(`${this.api.baseUrl}/api/seats/me/`);
+      const response = await this.api.fetcher.fetch({
+        method: "patch",
+        url,
+        path: "/api/seats/me/",
+        overrides: {
+          body: JSON.stringify({
+            product_key: SEAT_PRODUCT_KEY,
+            plan_key: planKey,
+          }),
+        },
+      });
+      return (await response.json()) as SeatData;
+    } catch (error) {
+      this.throwSeatError(error);
+    }
+  }
+
+  async cancelSeat(): Promise<void> {
+    try {
+      const url = new URL(`${this.api.baseUrl}/api/seats/me/`);
+      url.searchParams.set("product_key", SEAT_PRODUCT_KEY);
+      await this.api.fetcher.fetch({
+        method: "delete",
+        url,
+        path: "/api/seats/me/",
+      });
+    } catch (error) {
+      if (this.isFetcherStatusError(error, 204)) {
+        return;
+      }
+      this.throwSeatError(error);
+    }
+  }
+
+  async reactivateSeat(): Promise<SeatData> {
+    try {
+      const url = new URL(`${this.api.baseUrl}/api/seats/me/reactivate/`);
+      const response = await this.api.fetcher.fetch({
+        method: "post",
+        url,
+        path: "/api/seats/me/reactivate/",
+        overrides: {
+          body: JSON.stringify({ product_key: SEAT_PRODUCT_KEY }),
+        },
+      });
+      return (await response.json()) as SeatData;
+    } catch (error) {
+      this.throwSeatError(error);
+    }
+  }
+
+  private isFetcherStatusError(error: unknown, status: number): boolean {
+    return error instanceof Error && error.message.includes(`[${status}]`);
+  }
+
+  private parseFetcherError(error: unknown): {
+    status: number;
+    body: Record<string, unknown>;
+  } | null {
+    if (!(error instanceof Error)) return null;
+    const match = error.message.match(/\[(\d+)\]\s*(.*)/);
+    if (!match) return null;
+    try {
+      return {
+        status: Number.parseInt(match[1], 10),
+        body: JSON.parse(match[2]) as Record<string, unknown>,
+      };
+    } catch {
+      return { status: Number.parseInt(match[1], 10), body: {} };
+    }
+  }
+
+  private throwSeatError(error: unknown): never {
+    const parsed = this.parseFetcherError(error);
+
+    if (parsed) {
+      if (
+        parsed.status === 400 &&
+        typeof parsed.body.redirect_url === "string"
+      ) {
+        throw new SeatSubscriptionRequiredError(parsed.body.redirect_url);
+      }
+      if (parsed.status === 402) {
+        const message =
+          typeof parsed.body.error === "string" ? parsed.body.error : undefined;
+        throw new SeatPaymentFailedError(message);
+      }
+    }
+
+    throw error;
+  }
+
   /**
    * Check if a feature flag is enabled for the current project.
    * Returns true if the flag exists and is active, false otherwise.

apps/code/src/renderer/features/auth/components/AuthScreen.tsx

@@ -9,8 +9,8 @@ import { Callout, Flex, Spinner, Text, Theme } from "@radix-ui/themes";
 import codeLogo from "@renderer/assets/images/code.svg";
 import logomark from "@renderer/assets/images/logomark.svg";
 import { trpcClient } from "@renderer/trpc/client";
-import { REGION_LABELS } from "@shared/constants/oauth";
-import type { CloudRegion } from "@shared/types/oauth";
+import type { CloudRegion } from "@shared/types/regions";
+import { REGION_LABELS } from "@shared/types/regions";
 import { RegionSelect } from "./RegionSelect";
 
 export const getErrorMessage = (error: unknown) => {