diff --git a/.env.example b/.env.example
index c583d8bb..876d2cfb 100644
--- a/.env.example
+++ b/.env.example
@@ -5,6 +5,9 @@ NEXT_PUBLIC_APP_URL=http://localhost:3000
BETTER_AUTH_SECRET=your-better-auth-secret
BETTER_AUTH_URL=http://localhost:3000
+# Autumn Billing
+AUTUMN_SECRET_KEY=am_sk_test_...
+
# Google OAuth (for Better Auth)
GOOGLE_CLIENT_ID=...
GOOGLE_CLIENT_SECRET=...
diff --git a/package.json b/package.json
index afab3568..d746a39b 100644
--- a/package.json
+++ b/package.json
@@ -113,6 +113,7 @@
"ai": "^6.0.158",
"assemblyai": "^4.29.0",
"assistant-stream": "^0.3.10",
+ "autumn-js": "^1.2.9",
"better-auth": "^1.6.2",
"class-variance-authority": "^0.7.1",
"clsx": "^2.1.1",
@@ -185,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"
+ }
}
}
diff --git a/patches/autumn-js@1.2.9.patch b/patches/autumn-js@1.2.9.patch
new file mode 100644
index 00000000..9dcb01bd
--- /dev/null
+++ b/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}`);
diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml
new file mode 100644
index 00000000..c1aabe54
--- /dev/null
+++ b/pnpm-workspace.yaml
@@ -0,0 +1,2 @@
+patchedDependencies:
+ autumn-js@1.2.9: patches/autumn-js@1.2.9.patch
diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts
index 53e7dc4f..860acf1c 100644
--- a/src/app/api/chat/route.ts
+++ b/src/app/api/chat/route.ts
@@ -8,6 +8,7 @@ import {
wrapLanguageModel,
} from "ai";
import { devToolsMiddleware } from "@ai-sdk/devtools";
+import { Autumn } from "autumn-js";
import { withTracing } from "@posthog/ai";
import type { UIMessage } from "ai";
import { logger } from "@/lib/utils/logger";
@@ -21,13 +22,21 @@ 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,
getGatewayAttributionHeaders,
} from "@/lib/ai/gateway-provider-options";
+const autumn = new Autumn({
+ secretKey: process.env.AUTUMN_SECRET_KEY!,
+});
+
/**
* Extract workspaceId from system context or request body
*/
@@ -126,6 +135,7 @@ async function handlePOST(req: Request) {
// 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 +198,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 +208,33 @@ async function handlePOST(req: Request) {
selectedCardsContext,
);
+ if (userId && !isAnonymousUser && isPremiumChatModel(rawModelId)) {
+ try {
+ const checkResult = await autumn.check({
+ customerId: userId,
+ featureId: "premium_message",
+ requiredBalance: 1,
+ });
+
+ 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 +281,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 +293,18 @@ async function handlePOST(req: Request) {
};
logger.info("📊 [CHAT-API] Final Token Usage:", usageInfo);
+
+ if (userId && !isAnonymousUser && isPremiumChatModel(rawModelId)) {
+ try {
+ await autumn.track({
+ customerId: userId!,
+ 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
diff --git a/src/app/billing/layout.tsx b/src/app/billing/layout.tsx
new file mode 100644
index 00000000..dc429b89
--- /dev/null
+++ b/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 (
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ );
+ }
+
+ return <>{children}>;
+}
+
+export default function BillingLayout({
+ children,
+}: {
+ children: React.ReactNode;
+}) {
+ return (
+
+ {children}
+
+ );
+}
diff --git a/src/app/billing/page.tsx b/src/app/billing/page.tsx
new file mode 100644
index 00000000..90cf2ba9
--- /dev/null
+++ b/src/app/billing/page.tsx
@@ -0,0 +1,555 @@
+"use client";
+
+import { useCustomer, useListPlans } from "autumn-js/react";
+import { useState } from "react";
+import { Loader2, CreditCard, Sparkles, Zap } from "lucide-react";
+import { toast } from "sonner";
+
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Progress } from "@/components/ui/progress";
+import { Separator } from "@/components/ui/separator";
+import { Skeleton } from "@/components/ui/skeleton";
+
+type PremiumBalance = {
+ balance?: number;
+ limit?: number;
+ unlimited?: boolean;
+ remaining?: number;
+ granted?: number;
+};
+
+type CustomerSubscription = {
+ status?: string;
+ productId?: string;
+ planId?: string;
+ plan?: {
+ name?: string;
+ };
+};
+
+type Customer = {
+ features?: {
+ premium_message?: PremiumBalance;
+ };
+ balances?: {
+ premium_message?: PremiumBalance;
+ };
+ subscriptions?: CustomerSubscription[];
+};
+
+type PlanPrice = {
+ amount?: number;
+ interval?: string;
+ recurring?: {
+ interval?: string;
+ };
+};
+
+type Plan = {
+ id: string;
+ name?: string;
+ description?: string;
+ price?: PlanPrice | number | null;
+ prices?: PlanPrice[];
+ customerEligibility?: {
+ attachAction?: "activate" | "upgrade" | "downgrade" | "purchase" | "none";
+ status?: "active" | "scheduled" | string;
+ trialAvailable?: boolean;
+ };
+};
+
+type UseCustomerApi = ReturnType;
+
+function getAutumnCustomer(api: UseCustomerApi): Customer | null {
+ const customerApi = api as unknown as {
+ customer?: Customer | null;
+ data?: Customer | null;
+ };
+ return customerApi.customer ?? customerApi.data ?? null;
+}
+
+function getPremiumBalance(customer?: Customer | null) {
+ const premiumBalance =
+ customer?.features?.premium_message ?? customer?.balances?.premium_message;
+
+ if (!premiumBalance) return null;
+
+ return {
+ balance: premiumBalance.balance ?? premiumBalance.remaining ?? 0,
+ limit: premiumBalance.limit ?? premiumBalance.granted ?? 0,
+ unlimited: premiumBalance.unlimited ?? false,
+ };
+}
+
+async function attachPlan(api: UseCustomerApi, planId: string) {
+ const customerApi = api as unknown as {
+ attach?: (params: Record) => Promise;
+ };
+
+ if (!customerApi.attach) {
+ throw new Error("Attach is not available");
+ }
+
+ try {
+ return await customerApi.attach({ productId: planId });
+ } catch (error) {
+ return await customerApi.attach({ planId });
+ }
+}
+
+async function cancelPlan(api: UseCustomerApi, planId: string) {
+ const customerApi = api as unknown as {
+ cancel?: (params: Record) => Promise;
+ updateSubscription?: (params: Record) => Promise;
+ };
+
+ if (customerApi.cancel) {
+ return customerApi.cancel({ productId: planId });
+ }
+
+ if (customerApi.updateSubscription) {
+ return customerApi.updateSubscription({
+ planId,
+ cancelAction: "cancel_end_of_cycle",
+ });
+ }
+
+ throw new Error("Cancellation is not available");
+}
+
+async function openBillingPortal(api: UseCustomerApi, returnUrl: string) {
+ const customerApi = api as unknown as {
+ openBillingPortal?: (params: Record) => Promise;
+ openCustomerPortal?: (params: Record) => Promise;
+ };
+
+ if (customerApi.openBillingPortal) {
+ return customerApi.openBillingPortal({ returnUrl });
+ }
+
+ if (customerApi.openCustomerPortal) {
+ return customerApi.openCustomerPortal({ returnUrl });
+ }
+
+ throw new Error("Billing portal is not available");
+}
+
+function BillingSectionSkeleton() {
+ return (
+
+
+
+
+
+
+
+ );
+}
+
+function getPlanLabel(plan: Plan) {
+ const eligibility = plan.customerEligibility;
+
+ if (eligibility?.attachAction === "none") {
+ return eligibility.status === "scheduled" ? "Scheduled" : "Current Plan";
+ }
+
+ switch (eligibility?.attachAction) {
+ case "upgrade":
+ return "Upgrade";
+ case "downgrade":
+ return "Downgrade";
+ case "purchase":
+ return "Purchase";
+ case "activate":
+ return "Subscribe";
+ default:
+ return plan.price ? "Choose Plan" : "Get Started";
+ }
+}
+
+function getPlanPrice(plan: Plan) {
+ const priceCandidate =
+ Array.isArray(plan.prices) && plan.prices.length > 0
+ ? plan.prices[0]
+ : plan.price;
+
+ if (priceCandidate == null) {
+ return { amount: null as number | null, interval: "month" };
+ }
+
+ if (typeof priceCandidate === "number") {
+ return { amount: priceCandidate, interval: "month" };
+ }
+
+ return {
+ amount:
+ typeof priceCandidate.amount === "number" ? priceCandidate.amount : null,
+ interval:
+ priceCandidate.interval || priceCandidate.recurring?.interval || "month",
+ };
+}
+
+function formatPlanPrice(plan: Plan) {
+ const { amount, interval } = getPlanPrice(plan);
+
+ if (amount == null || amount <= 0) {
+ return "Free";
+ }
+
+ const normalizedAmount = amount;
+ return `$${normalizedAmount.toFixed(normalizedAmount % 1 === 0 ? 0 : 2)}/${interval === "year" ? "yr" : "mo"}`;
+}
+
+function UsageSection() {
+ const customerApi = useCustomer();
+ const customer = getAutumnCustomer(customerApi);
+ const isLoading = customerApi.isLoading;
+
+ if (isLoading) {
+ return ;
+ }
+
+ const premiumBalance = getPremiumBalance(customer);
+ const remaining = premiumBalance?.balance ?? 0;
+ const limit = premiumBalance?.limit ?? 0;
+ const used = limit > 0 ? Math.max(0, limit - remaining) : 0;
+ const progress = limit > 0 ? (used / limit) * 100 : 0;
+ const activeSubscription = customer?.subscriptions?.find(
+ (subscription) => subscription.status === "active",
+ );
+ const currentPlan =
+ activeSubscription?.plan?.name ||
+ activeSubscription?.productId ||
+ activeSubscription?.planId ||
+ (premiumBalance?.unlimited ? "Pro" : "Free");
+
+ return (
+
+
+
+
Your Plan
+
+ Track your premium AI usage and current subscription status.
+
+
+
+ {currentPlan}
+
+
+
+ {premiumBalance ? (
+ premiumBalance.unlimited ? (
+
+
+
+ Unlimited premium messages
+
+
+ You have unlimited access to premium chat models on your current
+ plan.
+
+
+ ) : (
+
+
+
+
Premium messages
+
+ {remaining} remaining out of {limit} this month
+
+
+
+ {limit > 0
+ ? `${Math.round((remaining / limit) * 100)}% left`
+ : "0% left"}
+
+
+
+
+ Only premium models consume credits. Basic models remain
+ unlimited.
+
+
+ )
+ ) : (
+
+ Usage data will appear here once your billing profile is loaded.
+
+ )}
+
+ );
+}
+
+function PricingTable() {
+ const {
+ data: plans,
+ isLoading,
+ error,
+ } = useListPlans() as {
+ data?: Plan[];
+ isLoading?: boolean;
+ error?: Error | null;
+ };
+ const customerApi = useCustomer();
+ const [submittingPlanId, setSubmittingPlanId] = useState(null);
+
+ const handleAttach = async (planId: string) => {
+ setSubmittingPlanId(planId);
+ try {
+ await attachPlan(customerApi, planId);
+ } catch (error) {
+ console.error("Failed to attach plan", error);
+ toast.error("Unable to start checkout", {
+ description: "Please try again in a moment.",
+ });
+ } finally {
+ setSubmittingPlanId(null);
+ }
+ };
+
+ if (isLoading) {
+ return (
+
+
+
Plans
+
+ Choose the plan that matches your usage.
+
+
+
+
+
+
+
+ );
+ }
+
+ if (error) {
+ return (
+
+ Plans
+
+ We couldn't load pricing right now. Please refresh and try again.
+
+
+ );
+ }
+
+ return (
+
+
+
Plans
+
+ Upgrade for unlimited premium messages, or stay on the free tier.
+
+
+
+
+ {plans?.map((plan) => {
+ const disabled = plan.customerEligibility?.attachAction === "none";
+ const isCurrent = plan.customerEligibility?.status === "active";
+ const isProcessing = submittingPlanId === plan.id;
+
+ return (
+
+
+
+
+
+ {plan.name || plan.id}
+
+
+ {plan.description ||
+ (plan.id === "pro"
+ ? "Unlimited premium messages and unrestricted access to all models."
+ : "100 premium AI messages per month with unlimited basic models.")}
+
+
+ {isCurrent ?
Active : null}
+
+
+ {formatPlanPrice(plan)}
+
+ {plan.customerEligibility?.trialAvailable ? (
+
+ Trial available
+
+ ) : null}
+
+
+
+ {plan.id === "pro" ? (
+ <>
+
+
+ Unlimited premium messages
+
+
+
+ Unlimited access to all chat models
+
+ >
+ ) : (
+ <>
+
+
+ 100 premium AI messages each month
+
+
+
+ Unlimited basic models
+
+ >
+ )}
+
+
+
+
+
+
+ );
+ })}
+
+
+ );
+}
+
+function SubscriptionManagement() {
+ const customerApi = useCustomer();
+ const customer = getAutumnCustomer(customerApi);
+ const [isCancelling, setIsCancelling] = useState(false);
+ const [isPortalLoading, setIsPortalLoading] = useState(false);
+
+ const premiumBalance = getPremiumBalance(customer);
+ const subscriptions = customer?.subscriptions || [];
+ const hasActivePaidPlan =
+ premiumBalance?.unlimited ||
+ subscriptions.some(
+ (subscription) =>
+ subscription.status === "active" &&
+ (subscription.productId === "pro" || subscription.planId === "pro"),
+ );
+
+ const handleCancel = async () => {
+ setIsCancelling(true);
+ try {
+ await cancelPlan(customerApi, "pro");
+ toast.success("Plan cancellation requested", {
+ description:
+ "Your Pro plan will remain active until the current billing period ends.",
+ });
+ } catch (error) {
+ console.error("Failed to cancel plan", error);
+ toast.error("Unable to cancel plan", {
+ description: "Please try again from the billing portal.",
+ });
+ } finally {
+ setIsCancelling(false);
+ }
+ };
+
+ const handleOpenPortal = async () => {
+ setIsPortalLoading(true);
+ try {
+ await openBillingPortal(customerApi, window.location.href);
+ } catch (error) {
+ console.error("Failed to open billing portal", error);
+ toast.error("Unable to open billing portal", {
+ description: "Please try again in a moment.",
+ });
+ } finally {
+ setIsPortalLoading(false);
+ }
+ };
+
+ return (
+
+
+
Subscription Management
+
+ Manage your subscription, payment method, and billing history.
+
+
+
+
+
+ {hasActivePaidPlan ? (
+
+
+
+
+ ) : (
+
+ You're currently on the free plan. Upgrade to Pro to manage
+ billing in Stripe.
+
+ )}
+
+ );
+}
+
+export default function BillingPage() {
+ return (
+
+
+
+ Plans & Billing
+
+
+ Manage your subscription, track premium message usage, and update
+ billing details.
+
+
+
+
+
+
+
+ );
+}
diff --git a/src/app/layout.tsx b/src/app/layout.tsx
index 40607c89..a912318a 100644
--- a/src/app/layout.tsx
+++ b/src/app/layout.tsx
@@ -1,5 +1,6 @@
import type { Metadata } from "next";
import Script from "next/script";
+import { AutumnProvider } from "autumn-js/react";
// Using system fonts instead of custom fonts
import { Providers } from "@/components/providers";
@@ -85,13 +86,23 @@ export default function RootLayout({
-
-
- {children}
-
+
+
+
+ {children}
+
+
diff --git a/src/components/assistant-ui/ModelPicker.tsx b/src/components/assistant-ui/ModelPicker.tsx
index c63be640..bb0e2517 100644
--- a/src/components/assistant-ui/ModelPicker.tsx
+++ b/src/components/assistant-ui/ModelPicker.tsx
@@ -125,37 +125,12 @@ function ModelPickerRow({
{model.ui.description}
-
-
-
- Speed
-
-
- {model.ui.speed}
-
+
+
+ Speed
-
-
- Cost
-
-
- {Array.from({ length: 3 }, (_, index) => {
- const isActive = index < model.ui.costLevel;
- return (
-
- $
-
- );
- })}
-
+
+ {model.ui.speed}
diff --git a/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx b/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx
index 8b76b7a8..ab57c1a4 100644
--- a/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx
+++ b/src/components/assistant-ui/WorkspaceRuntimeProvider.tsx
@@ -13,6 +13,7 @@ import {
import type { UIMessage } from "ai";
import { useMemo, useCallback, useRef } from "react";
import { AssistantAvailableProvider } from "@/contexts/AssistantAvailabilityContext";
+import { UpgradeDialog } from "@/components/billing/UpgradeDialog";
import { useUIStore } from "@/lib/stores/ui-store";
import { toast } from "sonner";
import { useWorkspaceState } from "@/hooks/workspace/use-workspace-state";
@@ -27,6 +28,12 @@ interface WorkspaceRuntimeProviderProps {
children: React.ReactNode;
}
+function UpgradeDialogConnected() {
+ const open = useUIStore((state) => state.showUpgradeDialog);
+ const setOpen = useUIStore((state) => state.setShowUpgradeDialog);
+ return
;
+}
+
function createWorkspaceChatRuntimeHook(
transport: AssistantChatTransport
,
onError: (error: Error) => void,
@@ -79,12 +86,7 @@ export function WorkspaceRuntimeProvider({
activePdfPageByItemId,
viewingItemIds,
);
- }, [
- workspaceState,
- contextCardIds,
- activePdfPageByItemId,
- viewingItemIds,
- ]);
+ }, [workspaceState, contextCardIds, activePdfPageByItemId, viewingItemIds]);
// Per AI SDK, transport `body` is `Resolvable