Improved reactivity in Fragments

app/page.tsx

@@ -19,7 +19,13 @@ import { ExecutionResult } from '@/lib/types'
 import { DeepPartial } from 'ai'
 import { experimental_useObject as useObject } from 'ai/react'
 import { usePostHog } from 'posthog-js/react'
-import { SetStateAction, useEffect, useState } from 'react'
+import {
+  SetStateAction,
+  useEffect,
+  useState,
+  useMemo,
+  useCallback,
+} from 'react'
 import { useLocalStorage } from 'usehooks-ts'
 
 export default function Home() {
@@ -40,6 +46,7 @@ export default function Home() {
   const [result, setResult] = useState<ExecutionResult>()
   const [messages, setMessages] = useState<Message[]>([])
   const [fragment, setFragment] = useState<DeepPartial<FragmentSchema>>()
+  const [isSideBarOpen, setIsSideBarOpen] = useState(false)
   const [currentTab, setCurrentTab] = useState<'code' | 'fragment'>('code')
   const [isPreviewLoading, setIsPreviewLoading] = useState(false)
   const [isAuthDialogOpen, setAuthDialog] = useState(false)
@@ -52,26 +59,54 @@ export default function Home() {
     process.env.NEXT_PUBLIC_USE_MORPH_APPLY === 'true',
   )
 
-  const filteredModels = modelsList.models.filter((model) => {
-    if (process.env.NEXT_PUBLIC_HIDE_LOCAL_MODELS) {
-      return model.providerId !== 'ollama'
-    }
-    return true
-  })
+  const filteredModels = useMemo(
+    () =>
+      modelsList.models.filter((model) => {
+        if (process.env.NEXT_PUBLIC_HIDE_LOCAL_MODELS) {
+          return model.providerId !== 'ollama'
+        }
+        return true
+      }),
+    [],
+  )
+
+  const currentModel = useMemo(
+    () => filteredModels.find((model) => model.id === languageModel.model),
+    [filteredModels, languageModel.model],
+  )
 
-  const currentModel = filteredModels.find(
-    (model) => model.id === languageModel.model,
+  const lastMessage = useMemo(() => messages[messages.length - 1], [messages])
+
+  const setLastMessage = useCallback((updates: Partial<Message>) => {
+    setMessages((previousMessages) => {
+      const lastMessage = previousMessages[previousMessages.length - 1]
+      if (!lastMessage) return previousMessages
+
+      return [
+        ...previousMessages.slice(0, -1),
+        { ...lastMessage, ...updates },
+      ]
+    })
+  }, [])
+
+  const currentTemplate = useMemo(
+    () =>
+      selectedTemplate === 'auto'
+        ? templates
+        : { [selectedTemplate]: templates[selectedTemplate] },
+    [selectedTemplate],
   )
-  const currentTemplate =
-    selectedTemplate === 'auto'
-      ? templates
-      : { [selectedTemplate]: templates[selectedTemplate] }
-  const lastMessage = messages[messages.length - 1]
 
   // Determine which API to use based on morph toggle and existing fragment
-  const shouldUseMorph =
-    useMorphApply && fragment && fragment.code && fragment.file_path
-  const apiEndpoint = shouldUseMorph ? '/api/morph-chat' : '/api/chat'
+  const shouldUseMorph = useMemo(
+    () => useMorphApply && fragment,
+    [useMorphApply, fragment],
+  )
+
+  const apiEndpoint = useMemo(
+    () => (shouldUseMorph ? '/api/morph-chat' : '/api/chat'),
+    [shouldUseMorph],
+  )
 
   const { object, submit, isLoading, stop, error } = useObject({
     api: apiEndpoint,
@@ -109,7 +144,7 @@ export default function Home() {
 
         setResult(result)
         setCurrentPreview({ fragment, result })
-        setMessage({ result })
+        setLastMessage({ result })
         setCurrentTab('fragment')
         setIsPreviewLoading(false)
       }
@@ -130,10 +165,8 @@ export default function Home() {
           content,
           object,
         })
-      }
-
-      if (lastMessage && lastMessage.role === 'assistant') {
-        setMessage({
+      } else if (lastMessage && lastMessage.role === 'assistant') {
+        setLastMessage({
           content,
           object,
         })
@@ -145,18 +178,6 @@ export default function Home() {
     if (error) stop()
   }, [error])
 
-  function setMessage(message: Partial<Message>, index?: number) {
-    setMessages((previousMessages) => {
-      const updatedMessages = [...previousMessages]
-      updatedMessages[index ?? previousMessages.length - 1] = {
-        ...previousMessages[index ?? previousMessages.length - 1],
-        ...message,
-      }
-
-      return updatedMessages
-    })
-  }
-
   async function handleSubmitAuth(e: React.FormEvent<HTMLFormElement>) {
     e.preventDefault()
 
@@ -189,12 +210,13 @@ export default function Home() {
       template: currentTemplate,
       model: currentModel,
       config: languageModel,
-      ...(shouldUseMorph && fragment ? { currentFragment: fragment } : {}),
+      ...(shouldUseMorph ? { currentFragment: fragment } : {}),
     })
 
     setChatInput('')
     setFiles([])
     setCurrentTab('code')
+    setIsSideBarOpen(true)
 
     posthog.capture('chat_submit', {
       template: selectedTemplate,
@@ -210,7 +232,7 @@ export default function Home() {
       template: currentTemplate,
       model: currentModel,
       config: languageModel,
-      ...(shouldUseMorph && fragment ? { currentFragment: fragment } : {}),
+      ...(shouldUseMorph ? { currentFragment: fragment } : {}),
     })
   }
 
@@ -266,6 +288,7 @@ export default function Home() {
   }) {
     setFragment(preview.fragment)
     setResult(preview.result)
+    setIsSideBarOpen(true)
   }
 
   function handleUndo() {
@@ -283,9 +306,11 @@ export default function Home() {
           supabase={supabase}
         />
       )}
-      <div className="grid w-full md:grid-cols-2">
+      <div
+        className={`grid w-full ${isSideBarOpen ? 'md:grid-cols-2' : 'md:grid-cols-1'}`}
+      >
         <div
-          className={`flex flex-col w-full max-h-full max-w-[800px] mx-auto px-4 overflow-auto ${fragment ? 'col-span-1' : 'col-span-2'}`}
+          className={`flex flex-col w-full max-h-full max-w-[800px] mx-auto px-4 overflow-auto`}
         >
           <NavBar
             session={session}
@@ -343,7 +368,8 @@ export default function Home() {
           isPreviewLoading={isPreviewLoading}
           fragment={fragment}
           result={result as ExecutionResult}
-          onClose={() => setFragment(undefined)}
+          isSideBarOpen={isSideBarOpen}
+          setIsSideBarOpen={setIsSideBarOpen}
         />
       </div>
     </main>

components/preview.tsx

@@ -24,7 +24,8 @@ export function Preview({
   isPreviewLoading,
   fragment,
   result,
-  onClose,
+  isSideBarOpen,
+  setIsSideBarOpen,
 }: {
   teamID: string | undefined
   accessToken: string | undefined
@@ -34,9 +35,10 @@ export function Preview({
   isPreviewLoading: boolean
   fragment?: DeepPartial<FragmentSchema>
   result?: ExecutionResult
-  onClose: () => void
+  isSideBarOpen: boolean
+  setIsSideBarOpen: Dispatch<SetStateAction<boolean>>
 }) {
-  if (!fragment) {
+  if (!isSideBarOpen) {
     return null
   }
 
@@ -59,7 +61,7 @@ export function Preview({
                   variant="ghost"
                   size="icon"
                   className="text-muted-foreground"
-                  onClick={onClose}
+                  onClick={() => setIsSideBarOpen(false)}
                 >
                   <ChevronsRight className="h-5 w-5" />
                 </Button>