Gate premium chat models with Autumn credits

package.json

@@ -186,5 +186,10 @@
     "tw-animate-css": "^1.4.0",
     "typescript": "^6.0.2",
     "vitest": "^4.1.2"
+  },
+  "pnpm": {
+    "patchedDependencies": {
+      "autumn-js@1.2.9": "patches/autumn-js@1.2.9.patch"
+    }
   }
 }

patches/autumn-js@1.2.9.patch

@@ -0,0 +1,13 @@
+diff --git a/dist/better-auth/chunk-LRULDM7S.mjs b/dist/better-auth/chunk-LRULDM7S.mjs
+index 31334ee65452ccbf47388e1d4cd79ce7d1f908ce..4e24d0ed9013cd9968df667738d47aba98afd2e2 100644
+--- a/dist/better-auth/chunk-LRULDM7S.mjs
++++ b/dist/better-auth/chunk-LRULDM7S.mjs
+@@ -6,7 +6,7 @@ import {
+ } from "./chunk-GJAMWZNZ.mjs";
+
+ // src/better-auth/utils/createAutumnEndpoint.ts
+-import { createAuthEndpoint } from "better-auth/plugins";
++import { createAuthEndpoint } from "better-auth/api";
+ var getRouteConfig = (routeName) => {
+   const route = routeConfigs.find((r) => r.route === routeName);
+   if (!route) throw new Error(`Route not found: ${routeName}`);

pnpm-workspace.yaml

@@ -0,0 +1,2 @@
+patchedDependencies:
+  autumn-js@1.2.9: patches/autumn-js@1.2.9.patch

src/app/api/chat/route.ts

@@ -21,7 +21,11 @@ import {
 import { withServerObservability } from "@/lib/with-server-observability";
 import { normalizeLegacyToolMessages } from "@/lib/ai/legacy-tool-message-compat";
 import type { ReplySelection } from "@/lib/stores/ui-store";
-import { getDefaultChatModelId, resolveGatewayModelId } from "@/lib/ai/models";
+import {
+  getDefaultChatModelId,
+  isPremiumChatModel,
+  resolveGatewayModelId,
+} from "@/lib/ai/models";
 import {
   buildGatewayProviderOptions,
   createGatewayLanguageModel,
@@ -55,6 +59,17 @@ function getSelectedCardsContext(body: any): string {
   return body.selectedCardsContext || "";
 }
 
+type AutumnBillingAPI = {
+  check: (args: {
+    headers: Awaited<ReturnType<typeof headers>>;
+    body: { featureId: string };
+  }) => Promise<{ allowed: boolean }>;
+  track: (args: {
+    headers: Awaited<ReturnType<typeof headers>>;
+    body: { featureId: string; value: number };
+  }) => Promise<unknown>;
+};
+
 /**
  * Inject user-selected context (selected cards + reply quotes / workspace passages) into the last user message.
  * `custom` is body.metadata.custom from the composer's runConfig (replySelections only).
@@ -123,9 +138,11 @@ async function handlePOST(req: Request) {
   try {
     // FIX: Parallelize headers() and req.json() to eliminate waterfall
     const [headersObj, body] = await Promise.all([headers(), req.json()]);
+    const autumnBilling = auth.api as typeof auth.api & AutumnBillingAPI;
 
     // Get authenticated user ID
     const session = await auth.api.getSession({ headers: headersObj });
+    const isAnonymousUser = Boolean(session?.user?.isAnonymous);
     userId = session?.user?.id || null;
 
     const { messages = [] }: { messages?: UIMessage[] } = body;
@@ -188,9 +205,8 @@ async function handlePOST(req: Request) {
     // Get pre-formatted selected cards context from client (no DB fetch needed)
     const selectedCardsContext = getSelectedCardsContext(body);
 
-    const modelId = resolveGatewayModelId(
-      body.modelId || getDefaultChatModelId(),
-    );
+    const rawModelId = body.modelId || getDefaultChatModelId();
+    const modelId = resolveGatewayModelId(rawModelId);
 
     // Inject selected cards + reply selections into the last user message
     injectSelectionContext(
@@ -199,6 +215,32 @@ async function handlePOST(req: Request) {
       selectedCardsContext,
     );
 
+    if (userId && !isAnonymousUser && isPremiumChatModel(rawModelId)) {
+      try {
+        const checkResult = await autumnBilling.check({
+          headers: headersObj,
+          body: { featureId: "premium_message" },
+        });
+
+        if (!checkResult.allowed) {
+          return new Response(
+            JSON.stringify({
+              error: "credits_exhausted",
+              message:
+                "You've used all your premium AI messages for this month. Upgrade to Pro for unlimited access, or switch to another model.",
+              code: "CREDITS_EXHAUSTED",
+            }),
+            {
+              status: 403,
+              headers: { "Content-Type": "application/json" },
+            },
+          );
+        }
+      } catch (err) {
+        console.error("[billing] Autumn check failed, allowing request:", err);
+      }
+    }
+
     const posthogClient = getPostHogServerClient();
     const baseGatewayModel = createGatewayLanguageModel(modelId);
     const tracedModel = posthogClient
@@ -245,7 +287,7 @@ async function handlePOST(req: Request) {
         },
       },
       experimental_transform: smoothStream({ chunking: "word", delayInMs: 15 }),
-      onFinish: ({ usage, finishReason }) => {
+      onFinish: async ({ usage, finishReason }) => {
         const usageInfo = {
           inputTokens: usage?.inputTokens,
           outputTokens: usage?.outputTokens,
@@ -257,6 +299,17 @@ async function handlePOST(req: Request) {
         };
 
         logger.info("📊 [CHAT-API] Final Token Usage:", usageInfo);
+
+        if (userId && !isAnonymousUser && isPremiumChatModel(rawModelId)) {
+          try {
+            await autumnBilling.track({
+              headers: headersObj,
+              body: { featureId: "premium_message", value: 1 },
+            });
+          } catch (err) {
+            console.error("[billing] Autumn track failed:", err);
+          }
+        }
       },
       onStepFinish: (result) => {
         // stepType exists in runtime but may not be in type definitions

src/app/billing/layout.tsx

@@ -0,0 +1,52 @@
+"use client";
+
+import { useEffect } from "react";
+import { useRouter } from "next/navigation";
+import { AnonymousSessionHandler } from "@/components/layout/SessionHandler";
+import { Skeleton } from "@/components/ui/skeleton";
+import { useSession } from "@/lib/auth-client";
+
+function BillingAccessGate({ children }: { children: React.ReactNode }) {
+  const { data: session, isPending } = useSession();
+  const router = useRouter();
+
+  useEffect(() => {
+    if (!isPending && (!session || session.user?.isAnonymous)) {
+      router.replace("/auth/sign-up?redirect_url=%2Fbilling");
+    }
+  }, [isPending, router, session]);
+
+  if (isPending || !session || session.user?.isAnonymous) {
+    return (
+      <main className="mx-auto flex min-h-screen w-full max-w-4xl flex-col gap-8 px-4 py-12 md:px-6">
+        <div className="space-y-3 text-center">
+          <Skeleton className="mx-auto h-10 w-56" />
+          <Skeleton className="mx-auto h-4 w-72" />
+        </div>
+        <div className="space-y-4 rounded-2xl border bg-card p-6">
+          <Skeleton className="h-7 w-32" />
+          <Skeleton className="h-4 w-56" />
+          <Skeleton className="h-2 w-full" />
+        </div>
+        <div className="grid gap-4 md:grid-cols-2">
+          <Skeleton className="h-64 rounded-2xl" />
+          <Skeleton className="h-64 rounded-2xl" />
+        </div>
+      </main>
+    );
+  }
+
+  return <>{children}</>;
+}
+
+export default function BillingLayout({
+  children,
+}: {
+  children: React.ReactNode;
+}) {
+  return (
+    <AnonymousSessionHandler>
+      <BillingAccessGate>{children}</BillingAccessGate>
+    </AnonymousSessionHandler>
+  );
+}