diff --git a/packages/bolt/app/components/chat/BaseChat.tsx b/packages/bolt/app/components/chat/BaseChat.tsx index b3820e1..887593f 100644 --- a/packages/bolt/app/components/chat/BaseChat.tsx +++ b/packages/bolt/app/components/chat/BaseChat.tsx @@ -9,6 +9,7 @@ import { Messages } from './Messages.client'; import { SendButton } from './SendButton.client'; import styles from './BaseChat.module.scss'; +import { useLoaderData } from '@remix-run/react'; interface BaseChatProps { textareaRef?: React.RefObject | undefined; @@ -58,6 +59,7 @@ export const BaseChat = React.forwardRef( ref, ) => { const TEXTAREA_MAX_HEIGHT = chatStarted ? 400 : 200; + const { avatar } = useLoaderData<{ avatar?: string }>(); return (
( className="flex flex-col w-full flex-1 max-w-chat px-4 pb-6 mx-auto z-1" messages={messages} isStreaming={isStreaming} + avatar={avatar} /> ) : null; }} diff --git a/packages/bolt/app/components/chat/Messages.client.tsx b/packages/bolt/app/components/chat/Messages.client.tsx index 18cfefc..6d04542 100644 --- a/packages/bolt/app/components/chat/Messages.client.tsx +++ b/packages/bolt/app/components/chat/Messages.client.tsx @@ -9,10 +9,11 @@ interface MessagesProps { className?: string; isStreaming?: boolean; messages?: Message[]; + avatar?: string; } export const Messages = React.forwardRef((props: MessagesProps, ref) => { - const { id, isStreaming = false, messages = [] } = props; + const { id, isStreaming = false, messages = [], avatar } = props; return (
@@ -34,12 +35,8 @@ export const Messages = React.forwardRef((props: })} > {isUserMessage && ( -
-
+
+ {avatar ? :
}
)}
diff --git a/packages/bolt/app/lib/.server/auth.ts b/packages/bolt/app/lib/.server/auth.ts new file mode 100644 index 0000000..1ab24cd --- /dev/null +++ b/packages/bolt/app/lib/.server/auth.ts @@ -0,0 +1,41 @@ +import { json, redirect, type LoaderFunctionArgs, type TypedResponse } from '@remix-run/cloudflare'; +import { isAuthenticated, type Session } from './sessions'; + +type RequestArgs = Pick; + +export async function loadWithAuth( + args: T, + handler: (args: T, session: Session) => Promise, +) { + return handleWithAuth(args, handler, (response) => redirect('/login', response)); +} + +export async function actionWithAuth( + args: T, + handler: (args: T, session: Session) => Promise, +) { + return await handleWithAuth(args, handler, (response) => json({}, { status: 401, ...response })); +} + +async function handleWithAuth( + args: T, + handler: (args: T, session: Session) => Promise, + fallback: (partial: ResponseInit) => R, +) { + const { request, context } = args; + const { session, response } = await isAuthenticated(request, context.cloudflare.env); + + if (session == null && !import.meta.env.VITE_DISABLE_AUTH) { + return fallback(response); + } + + const handlerResponse = await handler(args, session || {}); + + if (response) { + for (const [key, value] of Object.entries(response.headers)) { + handlerResponse.headers.append(key, value); + } + } + + return handlerResponse; +} diff --git a/packages/bolt/app/lib/.server/login.ts b/packages/bolt/app/lib/.server/login.ts deleted file mode 100644 index 5c84c5c..0000000 --- a/packages/bolt/app/lib/.server/login.ts +++ /dev/null @@ -1,34 +0,0 @@ -import { json, redirect, type LoaderFunctionArgs } from '@remix-run/cloudflare'; -import { isAuthenticated } from './sessions'; - -type RequestArgs = Pick; - -export async function handleAuthRequest(args: T, body: object = {}) { - const { request, context } = args; - const { authenticated, response } = await isAuthenticated(request, context.cloudflare.env); - - if (authenticated || import.meta.env.VITE_DISABLE_AUTH) { - return json(body, response); - } - - return redirect('/login', response); -} - -export async function handleWithAuth(args: T, handler: (args: T) => Promise) { - const { request, context } = args; - const { authenticated, response } = await isAuthenticated(request, context.cloudflare.env); - - if (authenticated || import.meta.env.VITE_DISABLE_AUTH) { - const handlerResponse = await handler(args); - - if (response) { - for (const [key, value] of Object.entries(response.headers)) { - handlerResponse.headers.append(key, value); - } - } - - return handlerResponse; - } - - return json({}, { status: 401 }); -} diff --git a/packages/bolt/app/lib/.server/sessions.ts b/packages/bolt/app/lib/.server/sessions.ts index f735b62..2986619 100644 --- a/packages/bolt/app/lib/.server/sessions.ts +++ b/packages/bolt/app/lib/.server/sessions.ts @@ -1,4 +1,4 @@ -import { createCookieSessionStorage, redirect } from '@remix-run/cloudflare'; +import { createCookieSessionStorage, redirect, type Session as RemixSession } from '@remix-run/cloudflare'; import { decodeJwt } from 'jose'; import { CLIENT_ID, CLIENT_ORIGIN } from '~/lib/constants'; import { request as doRequest } from '~/lib/fetch'; @@ -13,30 +13,43 @@ const TOKEN_KEY = 't'; const EXPIRES_KEY = 'e'; const USER_ID_KEY = 'u'; const SEGMENT_KEY = 's'; +const AVATAR_KEY = 'a'; +const ENCRYPTED_KEY = 'd'; -interface SessionData { +interface PrivateSession { [TOKEN_KEY]: string; [EXPIRES_KEY]: number; [USER_ID_KEY]?: string; [SEGMENT_KEY]?: string; } +interface PublicSession { + [ENCRYPTED_KEY]: string; + [AVATAR_KEY]?: string; +} + +export interface Session { + userId?: string; + segmentWriteKey?: string; + avatar?: string; +} + export async function isAuthenticated(request: Request, env: Env) { const { session, sessionStorage } = await getSession(request, env); - const sessionData: SessionData | null = await decryptSessionData(env, session.get('d')); + const sessionData: PrivateSession | null = await decryptSessionData(env, session.get(ENCRYPTED_KEY)); const header = async (cookie: Promise) => ({ headers: { 'Set-Cookie': await cookie } }); const destroy = () => header(sessionStorage.destroySession(session)); if (sessionData?.[TOKEN_KEY] == null) { - return { authenticated: false as const, response: await destroy() }; + return { session: null, response: await destroy() }; } const expiresAt = sessionData[EXPIRES_KEY] ?? 0; if (Date.now() < expiresAt) { - return { authenticated: true as const }; + return { session: getSessionData(session, sessionData) }; } logger.debug('Renewing token'); @@ -56,11 +69,14 @@ export async function isAuthenticated(request: Request, env: Env) { const newSessionData = { ...sessionData, [EXPIRES_KEY]: expiresAt }; const encryptedData = await encryptSessionData(env, newSessionData); - session.set('d', encryptedData); + session.set(ENCRYPTED_KEY, encryptedData); - return { authenticated: true as const, response: await header(sessionStorage.commitSession(session)) }; + return { + session: getSessionData(session, newSessionData), + response: await header(sessionStorage.commitSession(session)), + }; } else { - return { authenticated: false as const, response: await destroy() }; + return { session: null, response: await destroy() }; } } @@ -74,7 +90,7 @@ export async function createUserSession( const expiresAt = cookieExpiration(tokens.expires_in, tokens.created_at); - const sessionData: SessionData = { + const sessionData: PrivateSession = { [TOKEN_KEY]: tokens.refresh, [EXPIRES_KEY]: expiresAt, [USER_ID_KEY]: identity?.userId ?? undefined, @@ -82,7 +98,8 @@ export async function createUserSession( }; const encryptedData = await encryptSessionData(env, sessionData); - session.set('d', encryptedData); + session.set(ENCRYPTED_KEY, encryptedData); + session.set(AVATAR_KEY, identity?.avatar); return { headers: { @@ -94,7 +111,7 @@ export async function createUserSession( } function getSessionStorage(cloudflareEnv: Env) { - return createCookieSessionStorage<{ d: string }>({ + return createCookieSessionStorage({ cookie: { name: '__session', httpOnly: true, @@ -108,7 +125,7 @@ function getSessionStorage(cloudflareEnv: Env) { export async function logout(request: Request, env: Env) { const { session, sessionStorage } = await getSession(request, env); - const sessionData = await decryptSessionData(env, session.get('d')); + const sessionData = await decryptSessionData(env, session.get(ENCRYPTED_KEY)); if (sessionData) { revokeToken(sessionData[TOKEN_KEY]); @@ -127,14 +144,11 @@ export function validateAccessToken(access: string) { return jwtPayload.bolt === true; } -export async function getSessionData(request: Request, env: Env) { - const { session } = await getSession(request, env); - - const decrypted = await decryptSessionData(env, session.get('d')); - +function getSessionData(session: RemixSession, data: PrivateSession): Session { return { - userId: decrypted?.[USER_ID_KEY], - segmentWriteKey: decrypted?.[SEGMENT_KEY], + userId: data?.[USER_ID_KEY], + segmentWriteKey: data?.[SEGMENT_KEY], + avatar: session.get(AVATAR_KEY), }; } @@ -212,12 +226,12 @@ function urlParams(data: Record) { async function decryptSessionData(env: Env, encryptedData?: string) { const decryptedData = encryptedData ? await decrypt(payloadSecret(env), encryptedData) : undefined; - const sessionData: SessionData | null = JSON.parse(decryptedData ?? 'null'); + const sessionData: PrivateSession | null = JSON.parse(decryptedData ?? 'null'); return sessionData; } -async function encryptSessionData(env: Env, sessionData: SessionData) { +async function encryptSessionData(env: Env, sessionData: PrivateSession) { return await encrypt(payloadSecret(env), JSON.stringify(sessionData)); } diff --git a/packages/bolt/app/lib/analytics.ts b/packages/bolt/app/lib/analytics.ts index bf1bcdc..883ebde 100644 --- a/packages/bolt/app/lib/analytics.ts +++ b/packages/bolt/app/lib/analytics.ts @@ -7,6 +7,7 @@ export interface Identity { userId?: string | null; guestId?: string | null; segmentWriteKey?: string | null; + avatar?: string; } const MESSAGE_PREFIX = 'Bolt'; diff --git a/packages/bolt/app/routes/_index.tsx b/packages/bolt/app/routes/_index.tsx index e9a1f64..dbffde7 100644 --- a/packages/bolt/app/routes/_index.tsx +++ b/packages/bolt/app/routes/_index.tsx @@ -1,16 +1,16 @@ -import { type LoaderFunctionArgs, type MetaFunction } from '@remix-run/cloudflare'; +import { json, type LoaderFunctionArgs, type MetaFunction } from '@remix-run/cloudflare'; import { ClientOnly } from 'remix-utils/client-only'; import { BaseChat } from '~/components/chat/BaseChat'; import { Chat } from '~/components/chat/Chat.client'; import { Header } from '~/components/header/Header'; -import { handleAuthRequest } from '~/lib/.server/login'; +import { loadWithAuth } from '~/lib/.server/auth'; export const meta: MetaFunction = () => { return [{ title: 'Bolt' }, { name: 'description', content: 'Talk with Bolt, an AI assistant from StackBlitz' }]; }; export async function loader(args: LoaderFunctionArgs) { - return handleAuthRequest(args); + return loadWithAuth(args, async (_args, session) => json({ avatar: session.avatar })); } export default function Index() { diff --git a/packages/bolt/app/routes/api.analytics.ts b/packages/bolt/app/routes/api.analytics.ts index 9bc2718..da61715 100644 --- a/packages/bolt/app/routes/api.analytics.ts +++ b/packages/bolt/app/routes/api.analytics.ts @@ -1,12 +1,11 @@ import { json, type ActionFunctionArgs } from '@remix-run/cloudflare'; -import { handleWithAuth } from '~/lib/.server/login'; -import { getSessionData } from '~/lib/.server/sessions'; +import { actionWithAuth } from '~/lib/.server/auth'; +import type { Session } from '~/lib/.server/sessions'; import { sendEventInternal, type AnalyticsEvent } from '~/lib/analytics'; -async function analyticsAction({ request, context }: ActionFunctionArgs) { +async function analyticsAction({ request }: ActionFunctionArgs, session: Session) { const event: AnalyticsEvent = await request.json(); - const sessionData = await getSessionData(request, context.cloudflare.env); - const { success, error } = await sendEventInternal(sessionData, event); + const { success, error } = await sendEventInternal(session, event); if (!success) { return json({ error }, { status: 500 }); @@ -16,5 +15,5 @@ async function analyticsAction({ request, context }: ActionFunctionArgs) { } export async function action(args: ActionFunctionArgs) { - return handleWithAuth(args, analyticsAction); + return actionWithAuth(args, analyticsAction); } diff --git a/packages/bolt/app/routes/api.chat.ts b/packages/bolt/app/routes/api.chat.ts index 8175366..747d1e5 100644 --- a/packages/bolt/app/routes/api.chat.ts +++ b/packages/bolt/app/routes/api.chat.ts @@ -1,17 +1,17 @@ import { type ActionFunctionArgs } from '@remix-run/cloudflare'; +import { actionWithAuth } from '~/lib/.server/auth'; import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '~/lib/.server/llm/constants'; import { CONTINUE_PROMPT } from '~/lib/.server/llm/prompts'; import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text'; import SwitchableStream from '~/lib/.server/llm/switchable-stream'; -import { handleWithAuth } from '~/lib/.server/login'; -import { getSessionData } from '~/lib/.server/sessions'; +import type { Session } from '~/lib/.server/sessions'; import { AnalyticsAction, AnalyticsTrackEvent, sendEventInternal } from '~/lib/analytics'; export async function action(args: ActionFunctionArgs) { - return handleWithAuth(args, chatAction); + return actionWithAuth(args, chatAction); } -async function chatAction({ context, request }: ActionFunctionArgs) { +async function chatAction({ context, request }: ActionFunctionArgs, session: Session) { const { messages } = await request.json<{ messages: Messages }>(); const stream = new SwitchableStream(); @@ -21,9 +21,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) { toolChoice: 'none', onFinish: async ({ text: content, finishReason, usage }) => { if (finishReason !== 'length') { - const sessionData = await getSessionData(request, context.cloudflare.env); - - await sendEventInternal(sessionData, { + await sendEventInternal(session, { action: AnalyticsAction.Track, payload: { event: AnalyticsTrackEvent.MessageComplete, diff --git a/packages/bolt/app/routes/api.enhancer.ts b/packages/bolt/app/routes/api.enhancer.ts index 4bc1201..a3ee8b7 100644 --- a/packages/bolt/app/routes/api.enhancer.ts +++ b/packages/bolt/app/routes/api.enhancer.ts @@ -1,14 +1,14 @@ import { type ActionFunctionArgs } from '@remix-run/cloudflare'; import { StreamingTextResponse, parseStreamPart } from 'ai'; +import { actionWithAuth } from '~/lib/.server/auth'; import { streamText } from '~/lib/.server/llm/stream-text'; -import { handleWithAuth } from '~/lib/.server/login'; import { stripIndents } from '~/utils/stripIndent'; const encoder = new TextEncoder(); const decoder = new TextDecoder(); export async function action(args: ActionFunctionArgs) { - return handleWithAuth(args, enhancerAction); + return actionWithAuth(args, enhancerAction); } async function enhancerAction({ context, request }: ActionFunctionArgs) { diff --git a/packages/bolt/app/routes/chat.$id.tsx b/packages/bolt/app/routes/chat.$id.tsx index e0ed0af..a1315c9 100644 --- a/packages/bolt/app/routes/chat.$id.tsx +++ b/packages/bolt/app/routes/chat.$id.tsx @@ -1,9 +1,9 @@ -import type { LoaderFunctionArgs } from '@remix-run/cloudflare'; +import { json, type LoaderFunctionArgs } from '@remix-run/cloudflare'; import { default as IndexRoute } from './_index'; -import { handleAuthRequest } from '~/lib/.server/login'; +import { loadWithAuth } from '~/lib/.server/auth'; export async function loader(args: LoaderFunctionArgs) { - return handleAuthRequest(args, { id: args.params.id }); + return loadWithAuth(args, async (_args, session) => json({ id: args.params.id, avatar: session.avatar })); } export default IndexRoute; diff --git a/packages/bolt/app/routes/login.tsx b/packages/bolt/app/routes/login.tsx index caabcfe..c715337 100644 --- a/packages/bolt/app/routes/login.tsx +++ b/packages/bolt/app/routes/login.tsx @@ -16,9 +16,9 @@ import { auth, type AuthAPI } from '~/lib/webcontainer/auth.client'; import { logger } from '~/utils/logger'; export async function loader({ request, context }: LoaderFunctionArgs) { - const { authenticated, response } = await isAuthenticated(request, context.cloudflare.env); + const { session, response } = await isAuthenticated(request, context.cloudflare.env); - if (authenticated) { + if (session != null) { return redirect('/', response); }