Store oauth token in DB + use it when doing API calls (#1885)

* store oauth token in DB

* store user token/oauth token in session

* use token when making API call if USE_USER_TOKEN is true in env

* lint

.env

@@ -7,9 +7,9 @@
 OPENAI_BASE_URL=https://router.huggingface.co/v1
 
 # Canonical auth token for any OpenAI-compatible provider
-OPENAI_API_KEY=#your provider API key (works for HF router, OpenAI, LM Studio, etc.)
-# Legacy alias (still supported): if set and OPENAI_API_KEY is empty, it will be used
-# HF_TOKEN=
+OPENAI_API_KEY=#your provider API key (works for HF router, OpenAI, LM Studio, etc.). 
+# When set to true, user token will be used for inference calls
+USE_USER_TOKEN=false
 
 ### MongoDB ###
 MONGODB_URL=#your mongodb URL here, use chat-ui-db image if you don't want to set this

src/app.d.ts

@@ -12,6 +12,7 @@ declare global {
 			sessionId: string;
 			user?: User & { logoutDisabled?: boolean };
 			isAdmin: boolean;
+			token?: string;
 		}
 
 		interface Error {

src/hooks.server.ts

@@ -128,6 +128,7 @@ export const handle: Handle = async ({ event, resolve }) => {
 
 	event.locals.user = auth.user || undefined;
 	event.locals.sessionId = auth.sessionId;
+	event.locals.token = auth.token;
 
 	event.locals.isAdmin =
 		event.locals.user?.isAdmin || adminTokenManager.isAdmin(event.locals.sessionId);

src/lib/server/apiToken.ts

@@ -0,0 +1,11 @@
+import { config } from "$lib/server/config";
+
+export function getApiToken(locals: App.Locals | undefined) {
+	if (config.USE_USER_TOKEN === "true") {
+		if (!locals?.token) {
+			throw new Error("User token not found");
+		}
+		return locals.token;
+	}
+	return config.OPENAI_API_KEY || config.HF_TOKEN;
+}

src/lib/server/auth.ts

@@ -18,6 +18,7 @@ import { ObjectId } from "mongodb";
 import type { Cookie } from "elysia";
 import { adminTokenManager } from "./adminToken";
 import type { User } from "$lib/types/User";
+import type { Session } from "$lib/types/Session";
 
 export interface OIDCSettings {
 	redirectURI: string;
@@ -79,6 +80,7 @@ export async function findUser(
 ): Promise<{
 	user: User | null;
 	invalidateSession: boolean;
+	oauth?: Session["oauth"];
 }> {
 	const session = await collections.sessions.findOne({ sessionId });
 
@@ -93,6 +95,7 @@ export async function findUser(
 	return {
 		user: await collections.users.findOne({ _id: session.userId }),
 		invalidateSession: false,
+		oauth: session.oauth,
 	};
 }
 export const authCondition = (locals: App.Locals) => {
@@ -283,6 +286,7 @@ export async function authenticateRequest(
 
 		return {
 			user: result.user ?? undefined,
+			token: result.oauth?.token?.value,
 			sessionId,
 			secretSessionId,
 			isAdmin: result.user?.isAdmin || adminTokenManager.isAdmin(sessionId),
@@ -308,6 +312,7 @@ export async function authenticateRequest(
 				return {
 					user,
 					sessionId,
+					token,
 					secretSessionId,
 					isAdmin: user.isAdmin || adminTokenManager.isAdmin(sessionId),
 				};
@@ -338,6 +343,7 @@ export async function authenticateRequest(
 				user,
 				sessionId,
 				secretSessionId,
+				token,
 				isAdmin: user.isAdmin || adminTokenManager.isAdmin(sessionId),
 			};
 		}

src/lib/server/endpoints/endpoints.ts

@@ -16,6 +16,7 @@ export interface EndpointParameters {
 	generateSettings?: Partial<Model["parameters"]>;
 	isMultimodal?: boolean;
 	conversationId?: ObjectId;
+	locals: App.Locals | undefined;
 }
 
 export type TextGenerationStreamOutputSimplified = TextGenerationStreamOutput & {

src/lib/server/generateFromDefaultEndpoint.ts

@@ -7,18 +7,20 @@ export async function* generateFromDefaultEndpoint({
 	preprompt,
 	generateSettings,
 	modelId,
+	locals,
 }: {
 	messages: EndpointMessage[];
 	preprompt?: string;
 	generateSettings?: Record<string, unknown>;
 	/** Optional: use this model instead of the default task model */
 	modelId?: string;
+	locals: App.Locals | undefined;
 }): AsyncGenerator<MessageUpdate, string, undefined> {
 	try {
 		// Choose endpoint based on provided modelId, else fall back to taskModel
 		const model = modelId ? (models.find((m) => m.id === modelId) ?? taskModel) : taskModel;
 		const endpoint = await model.getEndpoint();
-		const tokenStream = await endpoint({ messages, preprompt, generateSettings });
+		const tokenStream = await endpoint({ messages, preprompt, generateSettings, locals });
 
 		for await (const output of tokenStream) {
 			// if not generated_text is here it means the generation is not done

src/lib/server/models.ts

@@ -109,7 +109,7 @@ if (openaiBaseUrl) {
 		logger.info({ baseURL }, "[models] Using OpenAI-compatible base URL");
 
 		// Canonical auth token is OPENAI_API_KEY; keep HF_TOKEN as legacy alias
-		const authToken = config.OPENAI_API_KEY || config.HF_TOKEN || "";
+		const authToken = config.OPENAI_API_KEY || config.HF_TOKEN;
 
 		// Try unauthenticated request first (many model lists are public, e.g. HF router)
 		let response = await fetch(`${baseURL}/models`);

src/lib/server/router/arch.ts

@@ -3,6 +3,7 @@ import { logger } from "$lib/server/logger";
 import type { EndpointMessage } from "../endpoints/endpoints";
 import type { Route, RouteConfig } from "./types";
 import { getRoutes } from "./policy";
+import { getApiToken } from "$lib/server/apiToken";
 
 const DEFAULT_LAST_TURNS = 16;
 const PROMPT_TEMPLATE = `
@@ -68,7 +69,8 @@ function parseRouteName(text: string): string | undefined {
 
 export async function archSelectRoute(
 	messages: EndpointMessage[],
-	traceId?: string
+	traceId: string | undefined,
+	locals: App.Locals | undefined
 ): Promise<{ routeName: string }> {
 	const routes = await getRoutes();
 	const prompt = toRouterPrompt(messages, routes);
@@ -82,7 +84,7 @@ export async function archSelectRoute(
 	}
 
 	const headers: HeadersInit = {
-		Authorization: `Bearer ${config.OPENAI_API_KEY || config.HF_TOKEN}`,
+		Authorization: `Bearer ${getApiToken(locals)}`,
 		"Content-Type": "application/json",
 	};
 	const body = {

src/lib/server/router/endpoint.ts

@@ -10,6 +10,7 @@ import { config } from "$lib/server/config";
 import { logger } from "$lib/server/logger";
 import { archSelectRoute } from "./arch";
 import { getRoutes, resolveRouteModels } from "./policy";
+import { getApiToken } from "$lib/server/apiToken";
 
 const REASONING_BLOCK_REGEX = /<think>[\s\S]*?(?:<\/think>|$)/g;
 
@@ -72,7 +73,7 @@ export async function makeRouterEndpoint(routerModel: ProcessedModel): Promise<E
 			return endpoints.openai({
 				type: "openai",
 				baseURL: (config.OPENAI_BASE_URL || "https://router.huggingface.co/v1").replace(/\/$/, ""),
-				apiKey: config.OPENAI_API_KEY || config.HF_TOKEN || "sk-",
+				apiKey: getApiToken(params.locals),
 				model: modelForCall,
 				// Ensure streaming path is used
 				streamingSupported: true,
@@ -133,7 +134,7 @@ export async function makeRouterEndpoint(routerModel: ProcessedModel): Promise<E
 			}
 		}
 
-		const { routeName } = await archSelectRoute(sanitizedMessages);
+		const { routeName } = await archSelectRoute(sanitizedMessages, undefined, params.locals);
 
 		const fallbackModel = config.LLM_ROUTER_FALLBACK_MODEL || routerModel.id;
 		const { candidates } = resolveRouteModels(routeName, routes, fallbackModel);

src/lib/server/textGeneration/generate.ts

@@ -23,6 +23,7 @@ export async function* generate(
 		isContinue,
 		promptedAt,
 		forceMultimodal,
+		locals,
 	}: GenerateContext,
 	preprompt?: string
 ): AsyncIterable<MessageUpdate> {
@@ -57,6 +58,7 @@ export async function* generate(
 		// Allow user-level override to force multimodal
 		isMultimodal: (forceMultimodal ?? false) || model.multimodal,
 		conversationId: conv._id,
+		locals,
 	})) {
 		// Check if this output contains router metadata
 		if (
@@ -114,6 +116,7 @@ Do not use prefixes such as Response: or Answer: when answering to the user.`,
 							max_tokens: 1024,
 						},
 						modelId: model.id,
+						locals,
 					});
 					finalAnswer = summary;
 					yield {
@@ -224,7 +227,7 @@ Do not use prefixes such as Response: or Answer: when answering to the user.`,
 			) {
 				lastReasoningUpdate = new Date();
 				try {
-					generateSummaryOfReasoning(reasoningBuffer, model.id).then((summary) => {
+					generateSummaryOfReasoning(reasoningBuffer, model.id, locals).then((summary) => {
 						status = summary;
 					});
 				} catch (e) {

src/lib/server/textGeneration/index.ts

@@ -23,7 +23,7 @@ async function* keepAlive(done: AbortSignal): AsyncGenerator<MessageUpdate, unde
 export async function* textGeneration(ctx: TextGenerationContext) {
 	const done = new AbortController();
 
-	const titleGen = generateTitleForConversation(ctx.conv);
+	const titleGen = generateTitleForConversation(ctx.conv, ctx.locals);
 	const textGen = textGenerationWithoutTitle(ctx, done);
 	const keepAliveGen = keepAlive(done.signal);
 

src/lib/server/textGeneration/reasoning.ts

@@ -3,7 +3,8 @@ import { getReturnFromGenerator } from "$lib/utils/getReturnFromGenerator";
 
 export async function generateSummaryOfReasoning(
 	buffer: string,
-	modelId?: string
+	modelId: string | undefined,
+	locals: App.Locals | undefined
 ): Promise<string> {
 	let summary: string | undefined;
 
@@ -25,6 +26,7 @@ export async function generateSummaryOfReasoning(
 					max_tokens: 50,
 				},
 				modelId,
+				locals,
 			})
 		);
 	}

src/lib/server/textGeneration/title.ts

@@ -6,7 +6,8 @@ import type { Conversation } from "$lib/types/Conversation";
 import { getReturnFromGenerator } from "$lib/utils/getReturnFromGenerator";
 
 export async function* generateTitleForConversation(
-	conv: Conversation
+	conv: Conversation,
+	locals: App.Locals | undefined
 ): AsyncGenerator<MessageUpdate, undefined, undefined> {
 	try {
 		const userMessage = conv.messages.find((m) => m.from === "user");
@@ -15,7 +16,7 @@ export async function* generateTitleForConversation(
 
 		const prompt = userMessage.content;
 		const modelForTitle = config.TASK_MODEL?.trim() ? config.TASK_MODEL : conv.model;
-		const title = (await generateTitle(prompt, modelForTitle)) ?? "New Chat";
+		const title = (await generateTitle(prompt, modelForTitle, locals)) ?? "New Chat";
 
 		yield {
 			type: MessageUpdateType.Title,
@@ -26,7 +27,11 @@ export async function* generateTitleForConversation(
 	}
 }
 
-export async function generateTitle(prompt: string, modelId?: string) {
+async function generateTitle(
+	prompt: string,
+	modelId: string | undefined,
+	locals: App.Locals | undefined
+) {
 	if (config.LLM_SUMMARIZATION !== "true") {
 		// When summarization is disabled, use the first five words without adding emojis
 		return prompt.split(/\s+/g).slice(0, 5).join(" ");
@@ -48,6 +53,7 @@ Return ONLY the title text.`,
 				max_tokens: 30,
 			},
 			modelId,
+			locals,
 		})
 	)
 		.then((summary) => {

src/lib/server/textGeneration/types.ts

@@ -16,4 +16,5 @@ export interface TextGenerationContext {
 	username?: string;
 	/** Force-enable multimodal handling for endpoints that support it */
 	forceMultimodal?: boolean;
+	locals: App.Locals | undefined;
 }

src/lib/types/Session.ts

@@ -11,4 +11,12 @@ export interface Session extends Timestamps {
 	expiresAt: Date;
 	admin?: boolean;
 	coupledCookieHash?: string;
+
+	oauth?: {
+		token: {
+			value: string;
+			expiresAt: Date;
+		};
+		refreshToken?: string;
+	};
 }

src/routes/conversation/[id]/+server.ts

@@ -459,6 +459,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 							model.id
 						]
 					),
+					locals,
 				};
 				// run the text generation and send updates to the client
 				for await (const event of textGeneration(ctx)) await update(event);

src/routes/login/callback/+server.ts

@@ -52,7 +52,7 @@ export async function GET({ url, locals, cookies, request, getClientAddress }) {
 		throw error(403, "Invalid or expired CSRF token");
 	}
 
-	const { userData } = await getOIDCUserData(
+	const { userData, token } = await getOIDCUserData(
 		{ redirectURI: validatedToken.redirectUrl },
 		code,
 		iss
@@ -79,6 +79,7 @@ export async function GET({ url, locals, cookies, request, getClientAddress }) {
 
 	await updateUser({
 		userData,
+		token,
 		locals,
 		cookies,
 		userAgent: request.headers.get("user-agent") ?? undefined,

src/routes/login/callback/updateUser.spec.ts

@@ -6,6 +6,7 @@ import { ObjectId } from "mongodb";
 import { DEFAULT_SETTINGS } from "$lib/types/Settings";
 import { defaultModel } from "$lib/server/models";
 import { findUser } from "$lib/server/auth";
+import type { TokenSet } from "openid-client";
 
 const userData = {
 	preferred_username: "new-username",
@@ -21,6 +22,13 @@ const locals = {
 	isAdmin: false,
 };
 
+const token = {
+	access_token: "access_token",
+	refresh_token: "refresh_token",
+	expires_at: 1717334400,
+	expires_in: 3600,
+} as TokenSet;
+
 // @ts-expect-error SvelteKit cookies dumb mock
 const cookiesMock: Cookies = {
 	set: vi.fn(),
@@ -61,7 +69,7 @@ describe("login", () => {
 	it("should update user if existing", async () => {
 		await insertRandomUser();
 
-		await updateUser({ userData, locals, cookies: cookiesMock });
+		await updateUser({ userData, locals, cookies: cookiesMock, token });
 
 		const existingUser = await collections.users.findOne({ hfUserId: userData.sub });
 
@@ -75,7 +83,7 @@ describe("login", () => {
 
 		await insertRandomConversations(2);
 
-		await updateUser({ userData, locals, cookies: cookiesMock });
+		await updateUser({ userData, locals, cookies: cookiesMock, token });
 
 		const conversationCount = await collections.conversations.countDocuments({
 			userId: insertedId,
@@ -88,7 +96,7 @@ describe("login", () => {
 	});
 
 	it("should create default settings for new user", async () => {
-		await updateUser({ userData, locals, cookies: cookiesMock });
+		await updateUser({ userData, locals, cookies: cookiesMock, token });
 
 		const user = (await findUser(locals.sessionId)).user;
 
@@ -115,7 +123,7 @@ describe("login", () => {
 			shareConversationsWithModelAuthors: false,
 		});
 
-		await updateUser({ userData, locals, cookies: cookiesMock });
+		await updateUser({ userData, locals, cookies: cookiesMock, token });
 
 		const settings = await collections.settings.findOne({
 			_id: insertedId,

src/routes/login/callback/updateUser.ts

@@ -3,23 +3,24 @@ import { collections } from "$lib/server/database";
 import { ObjectId } from "mongodb";
 import { DEFAULT_SETTINGS } from "$lib/types/Settings";
 import { z } from "zod";
-import type { UserinfoResponse } from "openid-client";
+import type { UserinfoResponse, TokenSet } from "openid-client";
 import { error, type Cookies } from "@sveltejs/kit";
 import crypto from "crypto";
 import { sha256 } from "$lib/utils/sha256";
-import { addWeeks } from "date-fns";
+import { addWeeks, subMinutes } from "date-fns";
 import { OIDConfig } from "$lib/server/auth";
 import { config } from "$lib/server/config";
 import { logger } from "$lib/server/logger";
 
 export async function updateUser(params: {
 	userData: UserinfoResponse;
+	token: TokenSet;
 	locals: App.Locals;
 	cookies: Cookies;
 	userAgent?: string;
 	ip?: string;
 }) {
-	const { userData, locals, cookies, userAgent, ip } = params;
+	const { userData, token, locals, cookies, userAgent, ip } = params;
 
 	// Microsoft Entra v1 tokens do not provide preferred_username, instead the username is provided in the upn
 	// claim. See https://learn.microsoft.com/en-us/entra/identity-platform/access-token-claims-reference
@@ -122,6 +123,21 @@ export async function updateUser(params: {
 	// Get cookie hash if coupling is enabled
 	const coupledCookieHash = await getCoupledCookieHash({ type: "svelte", value: cookies });
 
+	// Prepare OAuth token data for session storage
+	const oauthData = token.access_token
+		? {
+				token: {
+					value: token.access_token,
+					expiresAt: token.expires_at
+						? subMinutes(new Date(token.expires_at * 1000), 1)
+						: token.expires_in
+							? subMinutes(new Date(Date.now() + token.expires_in * 1000), 1)
+							: addWeeks(new Date(), 2),
+				},
+				...(token.refresh_token ? { refreshToken: token.refresh_token } : {}),
+			}
+		: undefined;
+
 	if (existingUser) {
 		// update existing user if any
 		await collections.users.updateOne(
@@ -141,6 +157,7 @@ export async function updateUser(params: {
 			ip,
 			expiresAt: addWeeks(new Date(), 2),
 			...(coupledCookieHash ? { coupledCookieHash } : {}),
+			...(oauthData ? { oauth: oauthData } : {}),
 		});
 	} else {
 		// user doesn't exist yet, create a new one
@@ -169,6 +186,7 @@ export async function updateUser(params: {
 			ip,
 			expiresAt: addWeeks(new Date(), 2),
 			...(coupledCookieHash ? { coupledCookieHash } : {}),
+			...(oauthData ? { oauth: oauthData } : {}),
 		});
 
 		// move pre-existing settings to new user