Better first load

packages/mdx/pages/new-chat.tsx

@@ -5,7 +5,7 @@ import theme from "@code-hike/lighter/themes/github-dark.json"
 import { Message } from "../src/chat/types"
 
 export default function Page() {
-  const [progress, setProgress] = React.useState(0)
+  const [progress, setProgress] = React.useState(max)
   const messages = steps[progress]
   const isStreaming =
     messages[messages.length - 1]?.isStreaming
@@ -20,7 +20,7 @@ export default function Page() {
   // console.log(progress)
   // console.log({ messages })
   // console.log(messages[messages.length - 1]?.content)
-  // console.log({ conversation })
+  console.log({ conversation })
 
   React.useEffect(() => {
     // focus input range
@@ -132,6 +132,13 @@ const messages = [
 ~~~js foo.js
 console.log("this is foo")
 console.log("this too")
+console.log("this too")
+console.log("this too")
+console.log("this too")
+~~~
+
+~~~py bar.py
+print("this is foo")
 ~~~
 
 That is bar.

packages/mdx/src/chat/chat.tsx

@@ -55,6 +55,11 @@ function StaticChat({
         entry.type === "answer" ? (
           <React.Fragment key={i}>
             <Code
+              key={
+                entry.files?.find(
+                  f => f.name === entry.activeFile
+                )?.text
+              }
               files={entry.files}
               activeFile={entry.activeFile}
             />
@@ -192,8 +197,16 @@ const Code = React.memo(
 
     if (!files || !files.length) return null
 
+    // console.log("render code", files, activeFile)
+
+    const raw = files.find(f => f.name === active)?.raw
+    // console.log(
+    //   raw,
+    //   files.find(f => f.name === active)
+    // )
     return (
       <InnerCode
+        key={raw?.toString()}
         codeConfig={{
           showCopyButton: true,
           showExpandButton: true,
@@ -216,7 +229,9 @@ const Code = React.memo(
       prevProps.activeFile === nextProps.activeFile &&
       prevProps.files?.length === nextProps.files?.length &&
       prevProps.files?.every(
-        (f, i) => f.text === nextProps.files[i]?.text
+        (f, i) =>
+          f.text === nextProps.files[i]?.text &&
+          f.raw === nextProps.files[i]?.raw
       )
     )
   }

packages/mdx/src/chat/highlight-code.ts

@@ -2,7 +2,7 @@ import {
   highlightSync,
   RawTheme,
 } from "@code-hike/lighter/dist/browser.esm.mjs"
-import { isLangLoaded } from "./laguages"
+import { isLangLoaded, isLangSupported } from "./languages"
 import { EntryCodeFile, FileInfo } from "./types"
 
 import { diffLines } from "diff"
@@ -31,15 +31,21 @@ export function highlightFile(
         lang: fileInfo.lang,
       },
       text: "",
+      raw: true,
       focus: "",
       annotations: [],
     }
   }
+
   const text = getStreamingCode(
     fileInfo.text,
     prevFile?.text
   )
-  const result = highlightSync(text, fileInfo.lang, theme)
+  const result =
+    isLangSupported(fileInfo.lang) &&
+    isLangLoaded(fileInfo.lang)
+      ? highlightSync(text, fileInfo.lang, theme)
+      : highlightSync(text, "text", theme)
   const lines = result.lines.map(line => ({
     tokens: line.map(token => ({
       content: token.content,
@@ -55,6 +61,7 @@ export function highlightFile(
     },
     text,
     focus: getDiffFocus(prevFile?.text || "", text),
+    raw: !isLangLoaded(fileInfo.lang),
     annotations: [],
   }
 }

packages/mdx/src/chat/laguages.ts → packages/mdx/src/chat/languages.ts

@@ -1,4 +1,7 @@
-import { preload } from "@code-hike/lighter/dist/browser.esm.mjs"
+import {
+  preload,
+  LANG_NAMES,
+} from "@code-hike/lighter/dist/browser.esm.mjs"
 
 const cache = new Map<
   string,
@@ -11,7 +14,15 @@ export function getLoadedLanguages() {
     .map(([lang]) => lang)
 }
 
+export function isLangSupported(lang: string) {
+  return LANG_NAMES.includes(lang)
+}
+
 export function isLangLoaded(lang: string) {
+  if (!isLangSupported(lang)) {
+    return true
+  }
+
   // if is server side return false
   if (typeof window === "undefined") {
     return false

packages/mdx/src/chat/types.ts

@@ -26,4 +26,7 @@ export type FileInfo = {
   open: boolean
 }
 
-export type EntryCodeFile = CodeFile & { text: string }
+export type EntryCodeFile = CodeFile & {
+  text: string
+  raw?: boolean
+}

packages/mdx/src/chat/use-conversation.tsx

@@ -9,15 +9,12 @@ import {
   Message,
 } from "./types"
 import React from "react"
-import {
-  highlightSync,
-  RawTheme,
-} from "@code-hike/lighter/dist/browser.esm.mjs"
+import { RawTheme } from "@code-hike/lighter/dist/browser.esm.mjs"
 import {
   getLoadedLanguages,
   isLangLoaded,
   loadLang,
-} from "./laguages"
+} from "./languages"
 import { highlightFile } from "./highlight-code"
 
 export function useConversation(
@@ -30,8 +27,14 @@ export function useConversation(
     oldConversation: [] as ConversationEntry[],
   })
 
+  const [loadingLangs, setLoadingLangs] = React.useState(
+    [] as string[]
+  )
+
+  console.log(loadingLangs)
+
   const { oldConversation } = ref.current
-  const newConversation = getNewConversation(
+  let newConversation = getNewConversation(
     [],
     messages,
     oldConversation,
@@ -40,6 +43,44 @@ export function useConversation(
     theme
   )
 
+  const missingLangs = [] as string[]
+  for (const entry of newConversation) {
+    if (entry.type === "answer") {
+      for (const file of entry.files) {
+        if (
+          !isLangLoaded(file.code.lang) &&
+          !missingLangs.includes(file.code.lang) &&
+          !loadingLangs.includes(file.code.lang)
+        ) {
+          missingLangs.push(file.code.lang)
+        }
+      }
+    }
+  }
+
+  React.useEffect(() => {
+    if (missingLangs.length === 0) return
+    missingLangs.forEach(lang => {
+      loadLang(lang).then(() => {
+        setLoadingLangs(prev =>
+          prev.includes(lang) ? prev : [...prev, lang]
+        )
+      })
+    })
+  }, [missingLangs.join(",")])
+
+  // newConversation.forEach(entry => {
+  //   if (entry.type === "answer") {
+  //     entry.files = entry.files?.filter(
+  //       f => !missingLangs.includes(f.code.lang)
+  //     )
+  //   }
+  // })
+
+  if (!isStreaming && missingLangs.length > 0) {
+    newConversation = []
+  }
+
   ref.current.oldConversation = newConversation
 
   // console.log(messages, newConversation)
@@ -54,10 +95,6 @@ function getNewConversation(
   onReply: (reply: string) => void,
   theme: RawTheme
 ): ConversationEntry[] {
-  const [loadedLangs, setLoadedLangs] = React.useState(
-    getLoadedLanguages
-  )
-
   const headCount = Math.max(oldConversation.length - 2, 0)
   const conversation = oldConversation.slice(0, headCount)
 
@@ -83,29 +120,6 @@ function getNewConversation(
     }
   }
 
-  const missingLangs = [] as string[]
-  for (const entry of conversation) {
-    if (entry.type === "answer") {
-      for (const file of entry.files) {
-        if (
-          !loadedLangs.includes(file.code.lang) &&
-          !missingLangs.includes(file.code.lang)
-        ) {
-          missingLangs.push(file.code.lang)
-        }
-      }
-    }
-  }
-
-  // todo - should be effect
-  missingLangs.forEach(lang => {
-    loadLang(lang).then(() => {
-      setLoadedLangs(prev =>
-        prev.includes(lang) ? prev : [...prev, lang]
-      )
-    })
-  })
-
   return conversation
 }