From b3e0173651c9f557b062bf10b1362e3fe1cfc902 Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Wed, 15 Apr 2026 23:42:28 +0000
Subject: [PATCH 1/8] Install Autumn SDK and integrate with Better Auth
---
.env.example | 3 +++
package.json | 1 +
src/app/layout.tsx | 24 +++++++++++++++++-------
src/lib/auth-client.ts | 25 ++++++++++++++-----------
src/lib/auth.ts | 2 ++
5 files changed, 37 insertions(+), 18 deletions(-)
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..5931345c 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",
diff --git a/src/app/layout.tsx b/src/app/layout.tsx
index 40607c89..6f634e89 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,22 @@ export default function RootLayout({
-
-
- {children}
-
+
+
+
+ {children}
+
+
diff --git a/src/lib/auth-client.ts b/src/lib/auth-client.ts
index 075fbe82..050e0778 100644
--- a/src/lib/auth-client.ts
+++ b/src/lib/auth-client.ts
@@ -1,17 +1,20 @@
import { createAuthClient } from "better-auth/react";
import { anonymousClient } from "better-auth/client/plugins";
+import type { AutumnPlugin } from "autumn-js/better-auth";
+
+const autumnClient = () => ({
+ id: "autumn" as const,
+ $InferServerPlugin: {} as AutumnPlugin,
+});
export const authClient = createAuthClient({
- baseURL: process.env.NEXT_PUBLIC_BETTER_AUTH_URL ||
- process.env.NEXT_PUBLIC_APP_URL ||
- (typeof window !== "undefined" ? window.location.origin : "http://localhost:3000"),
- plugins: [anonymousClient()],
+ baseURL:
+ process.env.NEXT_PUBLIC_BETTER_AUTH_URL ||
+ process.env.NEXT_PUBLIC_APP_URL ||
+ (typeof window !== "undefined"
+ ? window.location.origin
+ : "http://localhost:3000"),
+ plugins: [anonymousClient(), autumnClient()],
});
-export const {
- signIn,
- signUp,
- signOut,
- useSession,
- getSession
-} = authClient;
+export const { signIn, signUp, signOut, useSession, getSession } = authClient;
diff --git a/src/lib/auth.ts b/src/lib/auth.ts
index 775f266e..f24a1939 100644
--- a/src/lib/auth.ts
+++ b/src/lib/auth.ts
@@ -2,6 +2,7 @@ import { betterAuth } from "better-auth/minimal";
import { drizzleAdapter } from "better-auth/adapters/drizzle";
import { nextCookies } from "better-auth/next-js";
import { anonymous } from "better-auth/plugins";
+import { autumn } from "autumn-js/better-auth";
import { db } from "@/lib/db/client";
// Determine the base URL - prioritize explicit config, then Vercel URL, then localhost
@@ -177,6 +178,7 @@ export const auth = betterAuth({
}
},
}),
+ autumn(),
// Automatically set cookies in server actions
// Make sure this is the last plugin in the array
nextCookies(),
From a7fc39fdf7952ceef49e4bc6ee78606c83a10e16 Mon Sep 17 00:00:00 2001
From: urjitc <135136842+urjitc@users.noreply.github.com>
Date: Wed, 15 Apr 2026 23:53:34 +0000
Subject: [PATCH 2/8] Enforce Autumn credits on premium chat models
---
src/app/api/chat/route.ts | 61 +++++++++++++++++--
src/components/assistant-ui/ModelPicker.tsx | 35 ++---------
.../assistant-ui/WorkspaceRuntimeProvider.tsx | 24 +++++---
src/components/billing/UpgradeDialog.tsx | 44 +++++++++++++
src/lib/ai/models.ts | 12 ++++
src/lib/stores/ui-store.ts | 5 ++
6 files changed, 140 insertions(+), 41 deletions(-)
create mode 100644 src/components/billing/UpgradeDialog.tsx
diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts
index 53e7dc4f..aaa98d7f 100644
--- a/src/app/api/chat/route.ts
+++ b/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>;
+ body: { featureId: string };
+ }) => Promise<{ allowed: boolean }>;
+ track: (args: {
+ headers: Awaited>;
+ body: { featureId: string; value: number };
+ }) => Promise;
+};
+
/**
* 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
@@ -257,6 +299,17 @@ async function handlePOST(req: Request) {
};
logger.info("📊 [CHAT-API] Final Token Usage:", usageInfo);
+
+ if (userId && !isAnonymousUser && isPremiumChatModel(rawModelId)) {
+ autumnBilling
+ .track({
+ headers: headersObj,
+ body: { featureId: "premium_message", value: 1 },
+ })
+ .catch((err: unknown) => {
+ console.error("[billing] Autumn track failed:", err);
+ });
+ }
},
onStepFinish: (result) => {
// stepType exists in runtime but may not be in type definitions
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