[Feature] Get token to Zero in a more robust way

.env.sample

@@ -2,8 +2,8 @@ NEXT_PUBLIC_ZERO_SERVER='http://localhost:4848'
 ZERO_UPSTREAM_DB="postgresql://user:password@127.0.0.1:5432/zstart"
 ZERO_CVR_DB="postgresql://user:password@127.0.0.1:5432/zstart_cvr"
 ZERO_CHANGE_DB="postgresql://user:password@127.0.0.1:5432/zstart_cdb"
-ZERO_AUTH_SECRET="secretkey"
 ZERO_REPLICA_FILE="/tmp/zstart_replica.db"
+ZERO_AUTH_JWKS_URL="http://localhost:3000/api/auth/jwks"
 
 AUTH_SECRET="" # Added by `npx auth`. Read more: https://cli.authjs.dev
 
@@ -12,3 +12,6 @@ AUTH_DRIZZLE_URL=$ZERO_UPSTREAM_DB
 AUTH_STRAVA_ID=
 AUTH_STRAVA_SECRET=
 AUTH_STRAVA_REDIRECT_URL="http://localhost:3000"
+
+BETTER_AUTH_SECRET=
+BETTER_AUTH_URL="http://localhost:3000"

app/(app)/layout.tsx

@@ -23,9 +23,7 @@ export default async function Layout({
 			<Header />
 
 			<ClientOnly fallback={<div>Loading...</div>}>
-				<ZeroProvider userID={session.user.id} token={session.token}>
-					{children}
-				</ZeroProvider>
+				<ZeroProvider userID={session.user.id}>{children}</ZeroProvider>
 			</ClientOnly>
 		</>
 	);

components/single-user.tsx

@@ -15,11 +15,15 @@ export const User = ({ id }: { id: string }) => {
 		const { value } = e.target;
 
 		z.mutate.user.update({
-			id: user.id,
+			id: user?.id,
 			name: value,
 		});
 	};
 
+	if (!user) {
+		return <div>Loading</div>;
+	}
+
 	return (
 		<div>
 			<button
@@ -41,9 +45,9 @@ export const User = ({ id }: { id: string }) => {
 			) : (
 				<div>
 					<div>User: {user?.name}</div>
-					<div>User Email: {user?.email}</div>
+					{/* <div>User Email: {user?.email}</div> */}
 					<div>ID: {user?.id}</div>
-					<div>
+					{/* <div>
 						User Image:
 						{user?.image && (
 							<img
@@ -54,8 +58,8 @@ export const User = ({ id }: { id: string }) => {
 								alt="user avatar"
 							/>
 						)}
-					</div>
-					<div>Provider: {user?.provider?.provider}</div>
+					</div> */}
+					<div>Provider: {user?.provider?.providerId}</div>
 				</div>
 			)}
 		</div>

components/user-list.tsx

@@ -13,7 +13,7 @@ export function UserList() {
 			{users.length > 0 ? (
 				<ul className="list-disc list-inside">
 					{users.map((u) => {
-						const name = `${u.name} - ${u?.provider?.provider ?? "no provider provided"}`;
+						const name = `${u.name} - ${u?.provider?.providerId ?? "no provider provided"}`;
 
 						return (
 							<li key={u.id}>

components/zero.tsx

@@ -12,21 +12,35 @@ import { type Schema, schema } from "../schema";
 export function ZeroProvider({
 	children,
 	userID,
-	token,
 }: {
 	children: ReactNode;
 	userID: string;
-	token: string;
 }) {
 	const z = useMemo(() => {
+		const jwtStorageKey = `jwt-${userID}`;
+
 		return new Zero({
 			userID,
-			auth: token,
+			auth: async (error) => {
+				if (error === "invalid-token") {
+					sessionStorage.removeItem(jwtStorageKey);
+				}
+				let token = sessionStorage.getItem(jwtStorageKey);
+				if (!token) {
+					if (!userID) return undefined;
+					const response = await fetch("/api/auth/token");
+					const data = await response.json();
+					token = data.token;
+					if (!token) throw new Error("No token found");
+					sessionStorage.setItem(jwtStorageKey, token);
+				}
+				return token ?? undefined;
+			},
 			server: process.env.NEXT_PUBLIC_ZERO_SERVER,
 			schema,
 			kvStore: "mem",
 		});
-	}, [userID, token]);
+	}, [userID]);
 
 	return <ZeroProviderBase zero={z}>{children}</ZeroProviderBase>;
 }

db/schema.ts

@@ -1,4 +1,3 @@
-import type { AdapterAccountType } from "@auth/core/adapters";
 import { relations, sql } from "drizzle-orm";
 import {
 	boolean,
@@ -231,3 +230,10 @@ export const tasksRelations = relations(tasks, ({ one }) => ({
 		references: [user.id],
 	}),
 }));
+
+export const jwks = pgTable("jwks", {
+	id: text("id").primaryKey(),
+	publicKey: text("public_key").notNull(),
+	privateKey: text("private_key").notNull(),
+	createdAt: timestamp("created_at").notNull(),
+});

drizzle-zero.config.ts

@@ -12,8 +12,6 @@ export default drizzleZeroConfig(drizzleSchema, {
 	tables: {
 		providers: true,
 		account: true,
-		// this can be set to false
-		// e.g. users: false,
 		activities: true,
 		user: true,
 		todos: {
@@ -24,6 +22,7 @@ export default drizzleZeroConfig(drizzleSchema, {
 			assignedToId: true,
 			timestamp: true,
 		},
+		jwks: false,
 		tasks: {
 			id: true,
 			name: true,

drizzle/0000_chubby_absorbing_man.sql → drizzle/0000_nappy_tarot.sql

@@ -28,6 +28,13 @@ CREATE TABLE "activity" (
 	"updatedAt" text DEFAULT (CURRENT_TIMESTAMP) NOT NULL
 );
 --> statement-breakpoint
+CREATE TABLE "jwks" (
+	"id" text PRIMARY KEY NOT NULL,
+	"public_key" text NOT NULL,
+	"private_key" text NOT NULL,
+	"created_at" timestamp NOT NULL
+);
+--> statement-breakpoint
 CREATE TABLE "session" (
 	"id" text PRIMARY KEY NOT NULL,
 	"expires_at" timestamp NOT NULL,

drizzle/meta/0000_snapshot.json

@@ -1,5 +1,5 @@
 {
-  "id": "e6e6d81b-76f5-4b8d-911a-987c3ca1f36e",
+  "id": "9004671c-1eb4-4f01-9e8f-349feb8c0522",
   "prevId": "00000000-0000-0000-0000-000000000000",
   "version": "7",
   "dialect": "postgresql",
@@ -190,6 +190,43 @@
       "checkConstraints": {},
       "isRLSEnabled": false
     },
+    "public.jwks": {
+      "name": "jwks",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "text",
+          "primaryKey": true,
+          "notNull": true
+        },
+        "public_key": {
+          "name": "public_key",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "private_key": {
+          "name": "private_key",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {},
+      "policies": {},
+      "checkConstraints": {},
+      "isRLSEnabled": false
+    },
     "public.session": {
       "name": "session",
       "schema": "",

drizzle/meta/_journal.json

@@ -5,8 +5,8 @@
     {
       "idx": 0,
       "version": "7",
-      "when": 1760804593389,
-      "tag": "0000_chubby_absorbing_man",
+      "when": 1760974090046,
+      "tag": "0000_nappy_tarot",
       "breakpoints": true
     }
   ]

lib/auth-client.ts

@@ -3,7 +3,6 @@ import { createAuthClient } from "better-auth/react";
 
 export const authClient = createAuthClient({
 	plugins: [genericOAuthClient()],
-	// baseURL: "http://localhost:3000",
 });
 
 export const { signIn, signOut, useSession } = authClient;

lib/auth.ts

@@ -1,8 +1,7 @@
 import { betterAuth } from "better-auth";
 import { drizzleAdapter } from "better-auth/adapters/drizzle";
-import { customSession, genericOAuth } from "better-auth/plugins";
+import { customSession, genericOAuth, jwt } from "better-auth/plugins";
 import { db } from "@/db"; // your drizzle instance
-import { getNewToken } from "../app/actions";
 
 export const auth = betterAuth({
 	database: drizzleAdapter(db, {
@@ -42,7 +41,6 @@ export const auth = betterAuth({
 						const data = await fetch("https://www.strava.com/api/v3/athlete", {
 							headers: { Authorization: `Bearer ${tokens.accessToken}` },
 						}).then((data) => data.json());
-						console.log({ data });
 						return {
 							...data,
 							email: data.username, // This is required to fix the error, even though it makes no sense b/c it isn't an email
@@ -58,21 +56,13 @@ export const auth = betterAuth({
 				where: (account, { eq }) => eq(account.userId, user.id),
 			});
 
-			const getJWTToken = await getNewToken(user.id);
-
-			// console.log({ account });
-			// const roles = findUserRoles(session.session.userId);
 			return {
-				// roles,
-				user: {
-					...user,
-					newField: "newField",
-				},
+				user,
 				session,
 				account,
-				token: getJWTToken,
 			};
 		}),
+		jwt(),
 	],
 	baseURL: process.env.BETTER_AUTH_URL,
 });