fix: support setting visibility on initial chat creation (#975)

app/(chat)/api/chat/route.ts

@@ -49,7 +49,8 @@ export async function POST(request: Request) {
   }
 
   try {
-    const { id, message, selectedChatModel } = requestBody;
+    const { id, message, selectedChatModel, selectedVisibilityType } =
+      requestBody;
 
     const session = await auth();
 
@@ -80,7 +81,12 @@ export async function POST(request: Request) {
         message,
       });
 
-      await saveChat({ id, userId: session.user.id, title });
+      await saveChat({
+        id,
+        userId: session.user.id,
+        title,
+        visibility: selectedVisibilityType,
+      });
     } else {
       if (chat.userId !== session.user.id) {
         return new Response('Forbidden', { status: 403 });
@@ -236,7 +242,7 @@ export async function GET(request: Request) {
     return new Response('Not found', { status: 404 });
   }
 
-  if (chat.userId !== session.user.id) {
+  if (chat.visibility === 'private' && chat.userId !== session.user.id) {
     return new Response('Forbidden', { status: 403 });
   }
 

app/(chat)/api/chat/schema.ts

@@ -24,6 +24,7 @@ export const postRequestBodySchema = z.object({
       .optional(),
   }),
   selectedChatModel: z.enum(['chat-model', 'chat-model-reasoning']),
+  selectedVisibilityType: z.enum(['public', 'private']),
 });
 
 export type PostRequestBody = z.infer<typeof postRequestBodySchema>;

app/(chat)/chat/[id]/page.tsx

@@ -60,8 +60,8 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
         <Chat
           id={chat.id}
           initialMessages={convertToUIMessages(messagesFromDb)}
-          selectedChatModel={DEFAULT_CHAT_MODEL}
+          initialChatModel={DEFAULT_CHAT_MODEL}
-          selectedVisibilityType={chat.visibility}
+          initialVisibilityType={chat.visibility}
           isReadonly={session?.user?.id !== chat.userId}
           session={session}
           autoResume={true}
@@ -76,8 +76,8 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
       <Chat
         id={chat.id}
         initialMessages={convertToUIMessages(messagesFromDb)}
-        selectedChatModel={chatModelFromCookie.value}
+        initialChatModel={chatModelFromCookie.value}
-        selectedVisibilityType={chat.visibility}
+        initialVisibilityType={chat.visibility}
         isReadonly={session?.user?.id !== chat.userId}
         session={session}
         autoResume={true}

app/(chat)/page.tsx

@@ -26,8 +26,8 @@ export default async function Page() {
           key={id}
           id={id}
           initialMessages={[]}
-          selectedChatModel={DEFAULT_CHAT_MODEL}
+          initialChatModel={DEFAULT_CHAT_MODEL}
-          selectedVisibilityType="private"
+          initialVisibilityType="private"
           isReadonly={false}
           session={session}
           autoResume={false}
@@ -43,8 +43,8 @@ export default async function Page() {
         key={id}
         id={id}
         initialMessages={[]}
-        selectedChatModel={modelIdFromCookie.value}
+        initialChatModel={modelIdFromCookie.value}
-        selectedVisibilityType="private"
+        initialVisibilityType="private"
         isReadonly={false}
         session={session}
         autoResume={false}

components/artifact.tsx

@@ -27,6 +27,7 @@ import { sheetArtifact } from '@/artifacts/sheet/client';
 import { textArtifact } from '@/artifacts/text/client';
 import equal from 'fast-deep-equal';
 import type { UseChatHelpers } from '@ai-sdk/react';
+import type { VisibilityType } from './visibility-selector';
 
 export const artifactDefinitions = [
   textArtifact,
@@ -66,6 +67,7 @@ function PureArtifact({
   reload,
   votes,
   isReadonly,
+  selectedVisibilityType,
 }: {
   chatId: string;
   input: string;
@@ -81,6 +83,7 @@ function PureArtifact({
   handleSubmit: UseChatHelpers['handleSubmit'];
   reload: UseChatHelpers['reload'];
   isReadonly: boolean;
+  selectedVisibilityType: VisibilityType;
 }) {
   const { artifact, setArtifact, metadata, setMetadata } = useArtifact();
 
@@ -335,6 +338,7 @@ function PureArtifact({
                     append={append}
                     className="bg-background dark:bg-muted"
                     setMessages={setMessages}
+                    selectedVisibilityType={selectedVisibilityType}
                   />
                 </form>
               </div>
@@ -503,6 +507,8 @@ export const Artifact = memo(PureArtifact, (prevProps, nextProps) => {
   if (!equal(prevProps.votes, nextProps.votes)) return false;
   if (prevProps.input !== nextProps.input) return false;
   if (!equal(prevProps.messages, nextProps.messages.length)) return false;
+  if (prevProps.selectedVisibilityType !== nextProps.selectedVisibilityType)
+    return false;
 
   return true;
 });

components/chat.tsx

@@ -17,26 +17,32 @@ import { getChatHistoryPaginationKey } from './sidebar-history';
 import { toast } from './toast';
 import type { Session } from 'next-auth';
 import { useSearchParams } from 'next/navigation';
+import { useChatVisibility } from '@/hooks/use-chat-visibility';
 
 export function Chat({
   id,
   initialMessages,
-  selectedChatModel,
+  initialChatModel,
-  selectedVisibilityType,
+  initialVisibilityType,
   isReadonly,
   session,
   autoResume,
 }: {
   id: string;
   initialMessages: Array<UIMessage>;
-  selectedChatModel: string;
+  initialChatModel: string;
-  selectedVisibilityType: VisibilityType;
+  initialVisibilityType: VisibilityType;
   isReadonly: boolean;
   session: Session;
   autoResume: boolean;
 }) {
   const { mutate } = useSWRConfig();
 
+  const { visibilityType } = useChatVisibility({
+    chatId: id,
+    initialVisibilityType,
+  });
+
   const {
     messages,
     setMessages,
@@ -57,7 +63,8 @@ export function Chat({
     experimental_prepareRequestBody: (body) => ({
       id,
       message: body.messages.at(-1),
-      selectedChatModel,
+      selectedChatModel: initialChatModel,
+      selectedVisibilityType: visibilityType,
     }),
     onFinish: () => {
       mutate(unstable_serialize(getChatHistoryPaginationKey));
@@ -109,8 +116,8 @@ export function Chat({
       <div className="flex flex-col min-w-0 h-dvh bg-background">
         <ChatHeader
           chatId={id}
-          selectedModelId={selectedChatModel}
+          selectedModelId={initialChatModel}
-          selectedVisibilityType={selectedVisibilityType}
+          selectedVisibilityType={initialVisibilityType}
           isReadonly={isReadonly}
           session={session}
         />
@@ -140,6 +147,7 @@ export function Chat({
               messages={messages}
               setMessages={setMessages}
               append={append}
+              selectedVisibilityType={visibilityType}
             />
           )}
         </form>
@@ -160,6 +168,7 @@ export function Chat({
         reload={reload}
         votes={votes}
         isReadonly={isReadonly}
+        selectedVisibilityType={visibilityType}
       />
     </>
   );

components/multimodal-input.tsx

@@ -26,6 +26,7 @@ import type { UseChatHelpers } from '@ai-sdk/react';
 import { AnimatePresence, motion } from 'framer-motion';
 import { ArrowDown } from 'lucide-react';
 import { useScrollToBottom } from '@/hooks/use-scroll-to-bottom';
+import type { VisibilityType } from './visibility-selector';
 
 function PureMultimodalInput({
   chatId,
@@ -40,6 +41,7 @@ function PureMultimodalInput({
   append,
   handleSubmit,
   className,
+  selectedVisibilityType,
 }: {
   chatId: string;
   input: UseChatHelpers['input'];
@@ -53,6 +55,7 @@ function PureMultimodalInput({
   append: UseChatHelpers['append'];
   handleSubmit: UseChatHelpers['handleSubmit'];
   className?: string;
+  selectedVisibilityType: VisibilityType;
 }) {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { width } = useWindowSize();
@@ -220,7 +223,11 @@ function PureMultimodalInput({
       {messages.length === 0 &&
         attachments.length === 0 &&
         uploadQueue.length === 0 && (
-          <SuggestedActions append={append} chatId={chatId} />
+          <SuggestedActions
+            append={append}
+            chatId={chatId}
+            selectedVisibilityType={selectedVisibilityType}
+          />
         )}
 
       <input
@@ -309,6 +316,8 @@ export const MultimodalInput = memo(
     if (prevProps.input !== nextProps.input) return false;
     if (prevProps.status !== nextProps.status) return false;
     if (!equal(prevProps.attachments, nextProps.attachments)) return false;
+    if (prevProps.selectedVisibilityType !== nextProps.selectedVisibilityType)
+      return false;
 
     return true;
   },

components/sidebar-history-item.tsx

@@ -1,4 +1,4 @@
-import { Chat } from '@/lib/db/schema';
+import type { Chat } from '@/lib/db/schema';
 import {
   SidebarMenuAction,
   SidebarMenuButton,
@@ -39,7 +39,7 @@ const PureChatItem = ({
 }) => {
   const { visibilityType, setVisibilityType } = useChatVisibility({
     chatId: chat.id,
-    initialVisibility: chat.visibility,
+    initialVisibilityType: chat.visibility,
   });
 
   return (

components/suggested-actions.tsx

@@ -3,14 +3,20 @@
 import { motion } from 'framer-motion';
 import { Button } from './ui/button';
 import { memo } from 'react';
-import { UseChatHelpers } from '@ai-sdk/react';
+import type { UseChatHelpers } from '@ai-sdk/react';
+import type { VisibilityType } from './visibility-selector';
 
 interface SuggestedActionsProps {
   chatId: string;
   append: UseChatHelpers['append'];
+  selectedVisibilityType: VisibilityType;
 }
 
-function PureSuggestedActions({ chatId, append }: SuggestedActionsProps) {
+function PureSuggestedActions({
+  chatId,
+  append,
+  selectedVisibilityType,
+}: SuggestedActionsProps) {
   const suggestedActions = [
     {
       title: 'What are the advantages',
@@ -71,4 +77,13 @@ function PureSuggestedActions({ chatId, append }: SuggestedActionsProps) {
   );
 }
 
-export const SuggestedActions = memo(PureSuggestedActions, () => true);
+export const SuggestedActions = memo(
+  PureSuggestedActions,
+  (prevProps, nextProps) => {
+    if (prevProps.chatId !== nextProps.chatId) return false;
+    if (prevProps.selectedVisibilityType !== nextProps.selectedVisibilityType)
+      return false;
+
+    return true;
+  },
+);

components/visibility-selector.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import { ReactNode, useMemo, useState } from 'react';
+import { type ReactNode, useMemo, useState } from 'react';
 import { Button } from '@/components/ui/button';
 import {
   DropdownMenu,
@@ -9,7 +9,6 @@ import {
   DropdownMenuTrigger,
 } from '@/components/ui/dropdown-menu';
 import { cn } from '@/lib/utils';
-
 import {
   CheckCircleFillIcon,
   ChevronDownIcon,
@@ -52,7 +51,7 @@ export function VisibilitySelector({
 
   const { visibilityType, setVisibilityType } = useChatVisibility({
     chatId,
-    initialVisibility: selectedVisibilityType,
+    initialVisibilityType: selectedVisibilityType,
   });
 
   const selectedVisibility = useMemo(
@@ -70,6 +69,7 @@ export function VisibilitySelector({
         )}
       >
         <Button
+          data-testid="visibility-selector"
           variant="outline"
           className="hidden md:flex md:px-2 md:h-[34px]"
         >
@@ -82,6 +82,7 @@ export function VisibilitySelector({
       <DropdownMenuContent align="start" className="min-w-[300px]">
         {visibilities.map((visibility) => (
           <DropdownMenuItem
+            data-testid={`visibility-selector-item-${visibility.id}`}
             key={visibility.id}
             onSelect={() => {
               setVisibilityType(visibility.id);

hooks/use-chat-visibility.ts

@@ -12,10 +12,10 @@ import type { VisibilityType } from '@/components/visibility-selector';
 
 export function useChatVisibility({
   chatId,
-  initialVisibility,
+  initialVisibilityType,
 }: {
   chatId: string;
-  initialVisibility: VisibilityType;
+  initialVisibilityType: VisibilityType;
 }) {
   const { mutate, cache } = useSWRConfig();
   const history: ChatHistory = cache.get('/api/history')?.data;
@@ -24,7 +24,7 @@ export function useChatVisibility({
     `${chatId}-visibility`,
     null,
     {
-      fallbackData: initialVisibility,
+      fallbackData: initialVisibilityType,
     },
   );
 

lib/db/queries.ts

@@ -31,6 +31,7 @@ import {
 import type { ArtifactKind } from '@/components/artifact';
 import { generateUUID } from '../utils';
 import { generateHashedPassword } from './utils';
+import type { VisibilityType } from '@/components/visibility-selector';
 
 // Optionally, if not using email/pass login, you can
 // use the Drizzle adapter for Auth.js / NextAuth
@@ -79,10 +80,12 @@ export async function saveChat({
   id,
   userId,
   title,
+  visibility,
 }: {
   id: string;
   userId: string;
   title: string;
+  visibility: VisibilityType;
 }) {
   try {
     return await db.insert(chat).values({
@@ -90,6 +93,7 @@ export async function saveChat({
       createdAt: new Date(),
       userId,
       title,
+      visibility,
     });
   } catch (error) {
     console.error('Failed to save chat in database');

package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-chatbot",
-  "version": "3.0.16",
+  "version": "3.0.17",
   "private": true,
   "scripts": {
     "dev": "next dev --turbo",

tests/pages/chat.ts

@@ -115,6 +115,23 @@ export class ChatPage {
     expect(await this.getSelectedModel()).toBe(chatModel.name);
   }
 
+  public async getSelectedVisibility() {
+    const visibilityId = await this.page
+      .getByTestId('visibility-selector')
+      .innerText();
+    return visibilityId;
+  }
+
+  public async chooseVisibilityFromSelector(
+    chatVisibility: 'public' | 'private',
+  ) {
+    await this.page.getByTestId('visibility-selector').click();
+    await this.page
+      .getByTestId(`visibility-selector-item-${chatVisibility}`)
+      .click();
+    expect(await this.getSelectedVisibility()).toBe(chatVisibility);
+  }
+
   async getRecentAssistantMessage() {
     const messageElements = await this.page
       .getByTestId('message-assistant')

tests/routes/chat.test.ts

@@ -26,6 +26,7 @@ test.describe
           id: chatId,
           message: TEST_PROMPTS.SKY.MESSAGE,
           selectedChatModel: 'chat-model',
+          selectedVisibilityType: 'private',
         },
       });
       expect(response.status()).toBe(200);
@@ -49,6 +50,7 @@ test.describe
           id: chatId,
           message: TEST_PROMPTS.GRASS.MESSAGE,
           selectedChatModel: 'chat-model',
+          selectedVisibilityType: 'private',
         },
       });
       expect(response.status()).toBe(403);
@@ -109,6 +111,7 @@ test.describe
             createdAt: new Date().toISOString(),
           },
           selectedChatModel: 'chat-model',
+          selectedVisibilityType: 'private',
         },
       });
 
@@ -162,6 +165,7 @@ test.describe
             createdAt: new Date().toISOString(),
           },
           selectedChatModel: 'chat-model',
+          selectedVisibilityType: 'private',
         },
       });
 
@@ -190,7 +194,7 @@ test.describe
       expect(secondResponseContent).toEqual('');
     });
 
-    test('Babbage cannot resume chat generation that belongs to Ada', async ({
+    test('Babbage cannot resume a private chat generation that belongs to Ada', async ({
       adaContext,
       babbageContext,
     }) => {
@@ -212,6 +216,7 @@ test.describe
             createdAt: new Date().toISOString(),
           },
           selectedChatModel: 'chat-model',
+          selectedVisibilityType: 'private',
         },
       });
 
@@ -234,4 +239,57 @@ test.describe
       expect(firstStatusCode).toBe(200);
       expect(secondStatusCode).toBe(403);
     });
+
+    test('Babbage can resume a public chat generation that belongs to Ada', async ({
+      adaContext,
+      babbageContext,
+    }) => {
+      const chatId = generateUUID();
+
+      const firstRequest = adaContext.request.post('/api/chat', {
+        data: {
+          id: chatId,
+          message: {
+            id: generateUUID(),
+            role: 'user',
+            content: 'Help me write an essay about Silicon Valley',
+            parts: [
+              {
+                type: 'text',
+                text: 'Help me write an essay about Silicon Valley',
+              },
+            ],
+            createdAt: new Date().toISOString(),
+          },
+          selectedChatModel: 'chat-model',
+          selectedVisibilityType: 'public',
+        },
+      });
+
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+
+      const secondRequest = babbageContext.request.get(
+        `/api/chat?chatId=${chatId}`,
+      );
+
+      const [firstResponse, secondResponse] = await Promise.all([
+        firstRequest,
+        secondRequest,
+      ]);
+
+      const [firstStatusCode, secondStatusCode] = await Promise.all([
+        firstResponse.status(),
+        secondResponse.status(),
+      ]);
+
+      expect(firstStatusCode).toBe(200);
+      expect(secondStatusCode).toBe(200);
+
+      const [firstResponseContent, secondResponseContent] = await Promise.all([
+        firstResponse.text(),
+        secondResponse.text(),
+      ]);
+
+      expect(firstResponseContent).toEqual(secondResponseContent);
+    });
   });