feat: add avatar (#47)

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<HTMLTextAreaElement> | undefined;
@@ -58,6 +59,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
     ref,
   ) => {
     const TEXTAREA_MAX_HEIGHT = chatStarted ? 400 : 200;
+    const { avatar } = useLoaderData<{ avatar?: string }>();
 
     return (
       <div
@@ -94,6 +96,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                       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;
                 }}

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<HTMLDivElement, MessagesProps>((props: MessagesProps, ref) => {
-  const { id, isStreaming = false, messages = [] } = props;
+  const { id, isStreaming = false, messages = [], avatar } = props;
 
   return (
     <div id={id} ref={ref} className={props.className}>
@@ -34,12 +35,8 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
                 })}
               >
                 {isUserMessage && (
-                  <div
-                    className={classNames(
-                      'flex items-center justify-center min-w-[34px] min-h-[34px] bg-white text-gray-600 rounded-full p-1 self-start',
-                    )}
-                  >
-                    <div className="i-ph:user-fill text-xl"></div>
+                  <div className="flex items-center justify-center min-w-[34px] min-h-[34px] bg-white text-gray-600 rounded-full p-1 self-start">
+                    {avatar ? <img className="w-6 h-6" src={avatar} /> : <div className="i-ph:user-fill text-xl"></div>}
                   </div>
                 )}
                 <div className="grid grid-col-1 w-full">

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<LoaderFunctionArgs, 'request' | 'context'>;
+
+export async function loadWithAuth<T extends RequestArgs>(
+  args: T,
+  handler: (args: T, session: Session) => Promise<Response>,
+) {
+  return handleWithAuth(args, handler, (response) => redirect('/login', response));
+}
+
+export async function actionWithAuth<T extends RequestArgs>(
+  args: T,
+  handler: (args: T, session: Session) => Promise<TypedResponse>,
+) {
+  return await handleWithAuth(args, handler, (response) => json({}, { status: 401, ...response }));
+}
+
+async function handleWithAuth<T extends RequestArgs, R extends TypedResponse>(
+  args: T,
+  handler: (args: T, session: Session) => Promise<R>,
+  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;
+}

packages/bolt/app/lib/.server/login.ts

@@ -1,34 +0,0 @@
-import { json, redirect, type LoaderFunctionArgs } from '@remix-run/cloudflare';
-import { isAuthenticated } from './sessions';
-
-type RequestArgs = Pick<LoaderFunctionArgs, 'request' | 'context'>;
-
-export async function handleAuthRequest<T extends RequestArgs>(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<T extends RequestArgs>(args: T, handler: (args: T) => Promise<Response>) {
-  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 });
-}

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<string>) => ({ 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<PublicSession>({
     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<PublicSession>, 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<string, string>) {
 
 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));
 }
 

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';

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() {

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);
 }

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,

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) {

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;

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);
   }