Wire up gateway usage endpoint in Plan & Usage settings

charlesvien · charlesvien · commit cefd39b9893e · 2026-04-08T01:06:14.000-07:00
diff --git a/apps/code/src/main/services/llm-gateway/schemas.ts b/apps/code/src/main/services/llm-gateway/schemas.ts
@@ -56,3 +56,21 @@ export interface AnthropicErrorResponse {
     code?: string;
   };
 }
+
+export const usageBucketSchema = z.object({
+  used_usd: z.number(),
+  limit_usd: z.number(),
+  remaining_usd: 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 UsageOutput = z.infer<typeof usageOutput>;
diff --git a/apps/code/src/main/services/llm-gateway/service.ts b/apps/code/src/main/services/llm-gateway/service.ts
@@ -1,4 +1,7 @@
-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";
@@ -10,6 +13,7 @@ import type {
   AnthropicMessagesResponse,
   LlmMessage,
   PromptOutput,
+  UsageOutput,
 } from "./schemas";
 
 const log = logger.scope("llm-gateway");
@@ -134,4 +138,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 (await response.json()) as UsageOutput;
+  }
 }
diff --git a/apps/code/src/main/trpc/routers/llm-gateway.ts b/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()),
 });
diff --git a/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx b/apps/code/src/renderer/features/settings/components/sections/PlanUsageSettings.tsx
@@ -15,8 +15,63 @@ import {
   Spinner,
   Text,
 } from "@radix-ui/themes";
+import { trpcClient } from "@renderer/trpc/client";
+import { logger } from "@utils/logger";
 import { getPostHogUrl } from "@utils/urls";
-import { useState } from "react";
+import { useEffect, useState } from "react";
+
+const log = logger.scope("plan-usage");
+
+interface UsageBucket {
+  used_usd: number;
+  limit_usd: number;
+  remaining_usd: number;
+  resets_in_seconds: number;
+  exceeded: boolean;
+}
+
+interface UsageData {
+  sustained: UsageBucket;
+  burst: UsageBucket;
+  is_rate_limited: boolean;
+}
+
+function formatUsd(amount: number): string {
+  return `$${amount.toFixed(2)}`;
+}
+
+function formatResetTime(seconds: number): string {
+  const days = Math.ceil(seconds / 86400);
+  if (days === 1) return "1 day";
+  return `${days} days`;
+}
+
+function useUsage() {
+  const [usage, setUsage] = useState<UsageData | null>(null);
+  const [isLoading, setIsLoading] = useState(true);
+
+  useEffect(() => {
+    let cancelled = false;
+
+    trpcClient.llmGateway.usage
+      .query()
+      .then((data) => {
+        if (!cancelled) setUsage(data);
+      })
+      .catch((error) => {
+        log.warn("Failed to fetch usage", error);
+      })
+      .finally(() => {
+        if (!cancelled) setIsLoading(false);
+      });
+
+    return () => {
+      cancelled = true;
+    };
+  }, []);
+
+  return { usage, isLoading };
+}
 
 export function PlanUsageSettings() {
   const {
@@ -31,6 +86,7 @@ export function PlanUsageSettings() {
   const { upgradeToPro, cancelSeat, reactivateSeat, clearError } =
     useSeatStore();
   const [showUpgradeDialog, setShowUpgradeDialog] = useState(false);
+  const { usage, isLoading: usageLoading } = useUsage();
 
   const formattedActiveUntil = activeUntil
     ? activeUntil.toLocaleDateString(undefined, {
@@ -181,59 +237,30 @@ export function PlanUsageSettings() {
         <Text size="2" weight="medium" style={{ color: "var(--gray-9)" }}>
           Usage
         </Text>
-        {isPro ? (
+        {usageLoading ? (
           <Flex
-            direction="column"
-            gap="3"
+            align="center"
+            justify="center"
             p="4"
             style={{
-              border: "1px solid var(--accent-7)",
+              border: "1px solid var(--gray-5)",
               borderRadius: "var(--radius-3)",
             }}
           >
-            <Flex align="center" justify="between">
-              <Text size="2" weight="medium">
-                Token usage
-              </Text>
-              <Text
-                size="2"
-                weight="medium"
-                style={{ color: "var(--accent-9)" }}
-              >
-                Unlimited
-              </Text>
-            </Flex>
-            <div
-              style={{
-                position: "relative",
-                height: 8,
-                borderRadius: 4,
-                overflow: "hidden",
-                background: "var(--gray-4)",
-              }}
-            >
-              <div
-                style={{
-                  position: "absolute",
-                  inset: 0,
-                  borderRadius: 4,
-                  background:
-                    "linear-gradient(90deg, var(--amber-9), var(--orange-9), var(--amber-8), var(--orange-10), var(--amber-9))",
-                  backgroundSize: "300% 100%",
-                  animation: "lava-flow 4s linear infinite",
-                  boxShadow: "0 0 10px var(--orange-8)",
-                }}
-              />
-              <style>{`
-                @keyframes lava-flow {
-                  0% { background-position: 300% 50%; }
-                  100% { background-position: 0% 50%; }
-                }
-              `}</style>
-            </div>
-            <Text size="1" style={{ color: "var(--gray-9)" }}>
-              Unlimited tokens included with Pro (go crazy)
-            </Text>
+            <Spinner size="2" />
+          </Flex>
+        ) : usage ? (
+          <Flex direction="column" gap="3">
+            <UsageMeter
+              label="Sustained"
+              bucket={usage.sustained}
+              color={usage.sustained.exceeded ? "red" : undefined}
+            />
+            <UsageMeter
+              label="Burst"
+              bucket={usage.burst}
+              color={usage.burst.exceeded ? "red" : undefined}
+            />
           </Flex>
         ) : (
           <Flex
@@ -245,17 +272,8 @@ export function PlanUsageSettings() {
               borderRadius: "var(--radius-3)",
             }}
           >
-            <Flex align="center" justify="between">
-              <Text size="2" weight="medium">
-                Token usage
-              </Text>
-              <Text size="2" weight="medium">
-                0%
-              </Text>
-            </Flex>
-            <Progress value={0} size="2" />
-            <Text size="1" style={{ color: "var(--gray-9)" }}>
-              0 tokens used this period
+            <Text size="2" color="gray">
+              Unable to load usage data
             </Text>
           </Flex>
         )}
@@ -349,6 +367,52 @@ export function PlanUsageSettings() {
   );
 }
 
+interface UsageMeterProps {
+  label: string;
+  bucket: UsageBucket;
+  color?: "red";
+}
+
+function UsageMeter({ label, bucket, color }: UsageMeterProps) {
+  const percentage =
+    bucket.limit_usd > 0
+      ? Math.min(100, (bucket.used_usd / bucket.limit_usd) * 100)
+      : 0;
+
+  const borderColor = color === "red" ? "var(--red-7)" : "var(--gray-5)";
+
+  return (
+    <Flex
+      direction="column"
+      gap="3"
+      p="4"
+      style={{
+        border: `1px solid ${borderColor}`,
+        borderRadius: "var(--radius-3)",
+      }}
+    >
+      <Flex align="center" justify="between">
+        <Text size="2" weight="medium">
+          {label}
+        </Text>
+        <Text size="2" weight="medium">
+          {formatUsd(bucket.used_usd)} / {formatUsd(bucket.limit_usd)}
+        </Text>
+      </Flex>
+      <Progress
+        value={percentage}
+        size="2"
+        color={color === "red" ? "red" : undefined}
+      />
+      <Text size="1" style={{ color: "var(--gray-9)" }}>
+        {bucket.exceeded
+          ? "Limit exceeded"
+          : `${formatUsd(bucket.remaining_usd)} remaining \u00b7 resets in ${formatResetTime(bucket.resets_in_seconds)}`}
+      </Text>
+    </Flex>
+  );
+}
+
 interface PlanCardProps {
   name: string;
   price: string;
diff --git a/packages/agent/src/posthog-api.ts b/packages/agent/src/posthog-api.ts
@@ -7,9 +7,9 @@ import type {
   TaskRun,
   TaskRunArtifact,
 } from "./types";
-import { getLlmGatewayUrl } from "./utils/gateway";
+import { getGatewayUsageUrl, getLlmGatewayUrl } from "./utils/gateway";
 
-export { getLlmGatewayUrl };
+export { getGatewayUsageUrl, getLlmGatewayUrl };
 
 const DEFAULT_USER_AGENT = `posthog/agent.hog.dev; version: ${packageJson.version}`;
 
diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts
@@ -1,23 +1,31 @@
 export type GatewayProduct = "posthog_code" | "background_agents";
 
-export function getLlmGatewayUrl(
-  posthogHost: string,
-  product: GatewayProduct = "posthog_code",
-): string {
+function getGatewayBaseUrl(posthogHost: string): string {
   const url = new URL(posthogHost);
   const hostname = url.hostname;
 
-  // Local development (normalize 127.0.0.1 to localhost)
   if (hostname === "localhost" || hostname === "127.0.0.1") {
-    return `${url.protocol}//localhost:3308/${product}`;
+    return `${url.protocol}//localhost:3308`;
   }
 
-  // Docker containers accessing host
   if (hostname === "host.docker.internal") {
-    return `${url.protocol}//host.docker.internal:3308/${product}`;
+    return `${url.protocol}//host.docker.internal:3308`;
   }
 
-  // Production - extract region from hostname, default to US
   const region = hostname.match(/^(us|eu)\.posthog\.com$/)?.[1] ?? "us";
-  return `https://gateway.${region}.posthog.com/${product}`;
+  return `https://gateway.${region}.posthog.com`;
+}
+
+export function getLlmGatewayUrl(
+  posthogHost: string,
+  product: GatewayProduct = "posthog_code",
+): string {
+  return `${getGatewayBaseUrl(posthogHost)}/${product}`;
+}
+
+export function getGatewayUsageUrl(
+  posthogHost: string,
+  product: GatewayProduct = "posthog_code",
+): string {
+  return `${getGatewayBaseUrl(posthogHost)}/v1/usage/${product}`;
 }
