Let users access the API using their HF token (#1216)

.env

@@ -143,6 +143,7 @@ PUBLIC_APP_DISCLAIMER_MESSAGE="Disclaimer: AI is an area of active research with
 LLM_SUMMARIZATION=true
 
 EXPOSE_API=true
+USE_HF_TOKEN_IN_API=false
 # PUBLIC_APP_NAME=HuggingChat
 # PUBLIC_APP_ASSETS=huggingchat
 # PUBLIC_APP_COLOR=yellow

src/hooks.server.ts

@@ -110,8 +110,8 @@ export const handle: Handle = async ({ event, resolve }) => {
 		? event.request.headers.get(env.TRUSTED_EMAIL_HEADER)
 		: null;
 
-	let secretSessionId: string;
-	let sessionId: string;
+	let secretSessionId: string | null = null;
+	let sessionId: string | null = null;
 
 	if (email) {
 		secretSessionId = sessionId = await sha256(email);
@@ -136,8 +136,63 @@ export const handle: Handle = async ({ event, resolve }) => {
 		if (user) {
 			event.locals.user = user;
 		}
-	} else {
-		// if the user doesn't have any cookie, we generate one for him
+	} else if (event.url.pathname.startsWith(`${base}/api/`) && env.USE_HF_TOKEN_IN_API === "true") {
+		// if the request goes to the API and no user is available in the header
+		// check if a bearer token is available in the Authorization header
+
+		const authorization = event.request.headers.get("Authorization");
+
+		if (authorization && authorization.startsWith("Bearer ")) {
+			const token = authorization.slice(7);
+
+			const hash = await sha256(token);
+
+			sessionId = secretSessionId = hash;
+
+			// check if the hash is in the DB and get the user
+			// else check against https://huggingface.co/api/whoami-v2
+
+			const cacheHit = await collections.tokenCaches.findOne({ tokenHash: hash });
+
+			if (cacheHit) {
+				const user = await collections.users.findOne({ hfUserId: cacheHit.userId });
+
+				if (!user) {
+					return errorResponse(500, "User not found");
+				}
+
+				event.locals.user = user;
+			} else {
+				const response = await fetch("https://huggingface.co/api/whoami-v2", {
+					headers: {
+						Authorization: `Bearer ${token}`,
+					},
+				});
+
+				if (!response.ok) {
+					return errorResponse(401, "Unauthorized");
+				}
+
+				const data = await response.json();
+				const user = await collections.users.findOne({ hfUserId: data.id });
+
+				if (!user) {
+					return errorResponse(500, "User not found");
+				}
+
+				await collections.tokenCaches.insertOne({
+					tokenHash: hash,
+					userId: data.id,
+					createdAt: new Date(),
+					updatedAt: new Date(),
+				});
+
+				event.locals.user = user;
+			}
+		}
+	}
+
+	if (!sessionId || !secretSessionId) {
 		secretSessionId = crypto.randomUUID();
 		sessionId = await sha256(secretSessionId);
 

src/lib/server/database.ts

@@ -17,6 +17,7 @@ import type { CommunityToolDB } from "$lib/types/Tool";
 
 import { logger } from "$lib/server/logger";
 import { building } from "$app/environment";
+import type { TokenCache } from "$lib/types/TokenCache";
 import { onExit } from "./exitHandler";
 
 export const CONVERSATION_STATS_COLLECTION = "conversations.stats";
@@ -85,6 +86,7 @@ export class Database {
 		const bucket = new GridFSBucket(db, { bucketName: "files" });
 		const migrationResults = db.collection<MigrationResult>("migrationResults");
 		const semaphores = db.collection<Semaphore>("semaphores");
+		const tokenCaches = db.collection<TokenCache>("tokens");
 		const tools = db.collection<CommunityToolDB>("tools");
 
 		return {
@@ -102,6 +104,7 @@ export class Database {
 			bucket,
 			migrationResults,
 			semaphores,
+			tokenCaches,
 			tools,
 		};
 	}
@@ -124,6 +127,7 @@ export class Database {
 			sessions,
 			messageEvents,
 			semaphores,
+			tokenCaches,
 			tools,
 		} = this.getCollections();
 
@@ -217,7 +221,10 @@ export class Database {
 		semaphores
 			.createIndex({ createdAt: 1 }, { expireAfterSeconds: 60 })
 			.catch((e) => logger.error(e));
-
+		tokenCaches
+			.createIndex({ createdAt: 1 }, { expireAfterSeconds: 5 * 60 })
+			.catch((e) => logger.error(e));
+		tokenCaches.createIndex({ tokenHash: 1 }).catch((e) => logger.error(e));
 		tools.createIndex({ createdById: 1, userCount: -1 }).catch((e) => logger.error(e));
 		tools.createIndex({ userCount: 1 }).catch((e) => logger.error(e));
 		tools.createIndex({ last24HoursCount: 1 }).catch((e) => logger.error(e));

src/lib/types/TokenCache.ts

@@ -0,0 +1,6 @@
+import type { Timestamps } from "./Timestamps";
+
+export interface TokenCache extends Timestamps {
+	tokenHash: string; // sha256 of the bearer token
+	userId: string; // the matching hf user id
+}