diff --git a/www/apps/api-reference/app/globals.css b/www/apps/api-reference/app/globals.css
index 3f3d74d08e..2ec89e36d4 100644
--- a/www/apps/api-reference/app/globals.css
+++ b/www/apps/api-reference/app/globals.css
@@ -55,6 +55,10 @@
   body[data-modal="opened"] {
     @apply !overflow-hidden;
   }
+
+  .text-wrap {
+    text-wrap: wrap;
+  }
 }
 
 .grecaptcha-badge {
diff --git a/www/apps/book/app/globals.css b/www/apps/book/app/globals.css
index 41d7c98507..4a75d63add 100644
--- a/www/apps/book/app/globals.css
+++ b/www/apps/book/app/globals.css
@@ -32,6 +32,10 @@
   body[data-modal="opened"] {
     @apply !overflow-hidden;
   }
+
+  .text-wrap {
+    text-wrap: wrap;
+  }
 }
 
 .grecaptcha-badge {
diff --git a/www/apps/resources/app/globals.css b/www/apps/resources/app/globals.css
index 336288f82d..a5b9e80f86 100644
--- a/www/apps/resources/app/globals.css
+++ b/www/apps/resources/app/globals.css
@@ -35,6 +35,10 @@
   body[data-modal="opened"] {
     @apply !overflow-hidden;
   }
+
+  .text-wrap {
+    text-wrap: wrap;
+  }
 }
 
 .grecaptcha-badge {
diff --git a/www/apps/ui/src/styles/globals.css b/www/apps/ui/src/styles/globals.css
index ac4daf4dc5..3c92254990 100644
--- a/www/apps/ui/src/styles/globals.css
+++ b/www/apps/ui/src/styles/globals.css
@@ -40,6 +40,10 @@
   body[data-modal="opened"] {
     @apply !overflow-hidden text-ui-fg-base;
   }
+
+  .text-wrap {
+    text-wrap: wrap;
+  }
 }
 
 .grecaptcha-badge {
diff --git a/www/apps/user-guide/app/globals.css b/www/apps/user-guide/app/globals.css
index 03dee4656a..1d4eb3275c 100644
--- a/www/apps/user-guide/app/globals.css
+++ b/www/apps/user-guide/app/globals.css
@@ -28,6 +28,10 @@
   body[data-modal="opened"] {
     @apply !overflow-hidden;
   }
+
+  .text-wrap {
+    text-wrap: wrap;
+  }
 }
 
 .grecaptcha-badge {
diff --git a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Footer/index.tsx b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Footer/index.tsx
index f5ec39c8aa..34a2d79372 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Footer/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Footer/index.tsx
@@ -12,7 +12,7 @@ export const AiAssistantChatWindowFooter = () => {
       )}
     >
       <span className="text-medusa-fg-muted">Chat is cleared on refresh</span>
-      <span className="h-full w-px bg-medusa-border-base"></span>
+      <span className="h-docs_0.75 w-px bg-medusa-border-base"></span>
       <div className="flex items-center gap-docs_0.5">
         <span className="text-medusa-fg-subtle">Line break</span>
         <div className="flex items-center gap-[5px]">
diff --git a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Header/index.tsx b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Header/index.tsx
index ac5fc37376..46f2e45d4f 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Header/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Header/index.tsx
@@ -23,7 +23,9 @@ export const AiAssistantChatWindowHeader = () => {
         <Tooltip
           tooltipChildren={
             <>
-              This site is protected by reCAPTCHA and the{" "}
+              This site is protected by reCAPTCHA and
+              <br />
+              the{" "}
               <Link href="https://policies.google.com/privacy">
                 Google Privacy Policy
               </Link>{" "}
@@ -32,6 +34,7 @@ export const AiAssistantChatWindowHeader = () => {
             </>
           }
           clickable={true}
+          tooltipClassName={"!text-compact-small-plus"}
         >
           <ShieldCheck className="text-medusa-fg-muted" />
         </Tooltip>
diff --git a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Input/index.tsx b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Input/index.tsx
index 7e1e6d6daa..acb005ab61 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Input/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/Input/index.tsx
@@ -2,6 +2,7 @@ import React, { useEffect, useRef } from "react"
 import clsx from "clsx"
 import { useAiAssistantChat } from "../../../../providers/AiAssistant/Chat"
 import { ArrowUpCircleSolid } from "@medusajs/icons"
+import { useAiAssistant } from "../../../../providers"
 
 export const AiAssistantChatWindowInput = () => {
   const {
@@ -12,6 +13,7 @@ export const AiAssistantChatWindowInput = () => {
     loading,
     getThreadItems,
   } = useAiAssistantChat()
+  const { chatOpened } = useAiAssistant()
   const formRef = useRef<HTMLFormElement | null>(null)
 
   const onSubmit = (e?: React.FormEvent<HTMLFormElement>) => {
@@ -33,7 +35,17 @@ export const AiAssistantChatWindowInput = () => {
       return
     }
     if (e.shiftKey) {
-      setQuestion((prev) => `${prev}\n`)
+      const { selectionStart, selectionEnd } = e.currentTarget
+      setQuestion(
+        (prev) =>
+          `${prev.substring(0, selectionStart)}\n${prev.substring(selectionEnd)}`
+      )
+      setTimeout(() => {
+        if (inputRef.current) {
+          inputRef.current.selectionStart = inputRef.current.selectionEnd =
+            selectionStart + 1
+        }
+      }, 0)
     } else {
       onSubmit()
     }
@@ -52,6 +64,7 @@ export const AiAssistantChatWindowInput = () => {
 
   useEffect(() => {
     adjustTextareaHeight()
+    inputRef.current?.focus()
   }, [question])
 
   const handleTouch = (e: React.TouchEvent<HTMLTextAreaElement>) => {
@@ -61,6 +74,19 @@ export const AiAssistantChatWindowInput = () => {
     })
   }
 
+  useEffect(() => {
+    if (!chatOpened || !inputRef.current) {
+      return
+    }
+
+    const isCursorAtEnd =
+      inputRef.current.selectionStart === inputRef.current.value.length
+
+    if (isCursorAtEnd) {
+      inputRef.current.scrollTop = inputRef.current.scrollHeight
+    }
+  }, [chatOpened, inputRef.current])
+
   return (
     <div
       className={clsx(
diff --git a/www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx b/www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx
index 76888b135e..e62c646f32 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx
@@ -1,6 +1,13 @@
 import clsx from "clsx"
 import React from "react"
-import { AiAssistantIcon, DotsLoading, MarkdownContent } from "@/components"
+import {
+  AiAssistantIcon,
+  CodeMdx,
+  CodeMdxProps,
+  DotsLoading,
+  MarkdownContent,
+  MDXComponents,
+} from "@/components"
 import { AiAssistantThreadItemActions } from "./Actions"
 import { AiAssistantThread } from "../../../providers/AiAssistant/Chat"
 
@@ -27,7 +34,7 @@ export const AiAssistantThreadItem = ({ item }: AiAssistantThreadItemProps) => {
           "txt-small text-medusa-fg-base",
           item.type === "question" && [
             "rounded-docs_xl bg-medusa-tag-neutral-bg",
-            "px-docs_0.75 py-docs_0.5",
+            "px-docs_0.75 py-docs_0.5 max-w-full",
           ],
           item.type !== "question" && "flex-1",
           item.type === "answer" && "text-pretty flex-1 max-w-[calc(100%-20px)]"
@@ -38,6 +45,39 @@ export const AiAssistantThreadItem = ({ item }: AiAssistantThreadItemProps) => {
             className="[&>*:last-child]:mb-0"
             allowedElements={["br", "p", "code", "pre"]}
             unwrapDisallowed={true}
+            components={{
+              ...MDXComponents,
+              code: (props: CodeMdxProps) => {
+                return (
+                  <CodeMdx
+                    {...props}
+                    noCopy
+                    noReport
+                    forceNoTitle
+                    noAskAi
+                    inlineCodeProps={{
+                      ...props.inlineCodeProps,
+                      className: "!text-wrap !break-words",
+                      variant: "grey-bg",
+                    }}
+                    collapsibleLines="11"
+                    codeBlockProps={{
+                      className: clsx(
+                        "rounded-docs_lg p-[5px]",
+                        props.className
+                      ),
+                      wrapperClassName: "rounded-docs_lg",
+                      innerClassName: "border rounded-docs_lg",
+                      overrideColors: {
+                        bg: "bg-medusa-contrast-bg-subtle",
+                        innerBg: "bg-medusa-contrast-bg-subtle",
+                        innerBorder: "border-medusa-contrast-border-bot",
+                      },
+                    }}
+                  />
+                )
+              },
+            }}
           >
             {item.content}
           </MarkdownContent>
@@ -45,7 +85,15 @@ export const AiAssistantThreadItem = ({ item }: AiAssistantThreadItemProps) => {
         {item.type === "answer" && (
           <div className="flex flex-col gap-docs_0.75">
             {!item.question_id && item.content.length === 0 && <DotsLoading />}
-            <MarkdownContent className="[&>*:last-child]:mb-0">
+            <MarkdownContent
+              className="[&>*:last-child]:mb-0"
+              components={{
+                ...MDXComponents,
+                code: (props: CodeMdxProps) => {
+                  return <CodeMdx {...props} noReport noAskAi />
+                },
+              }}
+            >
               {item.content}
             </MarkdownContent>
             {item.question_id && <AiAssistantThreadItemActions item={item} />}
diff --git a/www/packages/docs-ui/src/components/AiAssistant/TriggerButton/index.tsx b/www/packages/docs-ui/src/components/AiAssistant/TriggerButton/index.tsx
index e74d1a3310..3eb23014e2 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/TriggerButton/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/TriggerButton/index.tsx
@@ -1,6 +1,6 @@
 "use client"
 
-import React, { useMemo, useState } from "react"
+import React from "react"
 import { Button } from "../../Button"
 import { Tooltip } from "../../Tooltip"
 import { Kbd } from "../../Kbd"
@@ -9,17 +9,13 @@ import { useAiAssistant, useSearch, useSiteConfig } from "../../../providers"
 import { useKeyboardShortcut } from "../../../hooks"
 import Image from "next/image"
 
-const AI_ASSISTANT_ICON = "/images/ai-assistent-luminosity.png"
+// const AI_ASSISTANT_ICON = "/images/ai-assistent-luminosity.png"
 const AI_ASSISTANT_ICON_ACTIVE = "/images/ai-assistent.png"
 
 export const AiAssistantTriggerButton = () => {
-  const [hovered, setHovered] = useState(false)
   const { config } = useSiteConfig()
-  const { chatOpened, setChatOpened } = useAiAssistant()
+  const { setChatOpened } = useAiAssistant()
   const { setIsOpen } = useSearch()
-  const isActive = useMemo(() => {
-    return hovered || chatOpened
-  }, [hovered, chatOpened])
   const osShortcut = getOsShortcut()
 
   useKeyboardShortcut({
@@ -35,36 +31,29 @@ export const AiAssistantTriggerButton = () => {
   return (
     <Tooltip
       render={() => (
-        <span className="flex gap-[6px] items-center">
-          <span className="text-compact-x-small-plus text-medusa-fg-base">
-            Ask AI
-          </span>
-          <span className="flex gap-[5px] items-center">
-            <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
-              {osShortcut}
-            </Kbd>
-            <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
-              i
-            </Kbd>
-          </span>
+        <span className="flex gap-[5px] items-center">
+          <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
+            {osShortcut}
+          </Kbd>
+          <Kbd className="bg-medusa-bg-field-component border-medusa-border-strong w-[18px] h-[18px] inline-block">
+            i
+          </Kbd>
         </span>
       )}
     >
       <Button
         variant="transparent-clear"
-        className="!p-[6.5px]"
-        onMouseOver={() => setHovered(true)}
-        onMouseOut={() => setHovered(false)}
-        onTouchStart={() => setHovered(true)}
-        onTouchEnd={() => setHovered(false)}
         onClick={() => setChatOpened((prev) => !prev)}
       >
         <Image
-          src={`${config.basePath}${isActive ? AI_ASSISTANT_ICON_ACTIVE : AI_ASSISTANT_ICON}`}
+          src={`${config.basePath}${AI_ASSISTANT_ICON_ACTIVE}`}
           width={15}
           height={15}
           alt="AI Assistant"
         />
+        <span className="hidden md:inline-block text-medusa-fg-subtle">
+          Ask AI
+        </span>
       </Button>
     </Tooltip>
   )
diff --git a/www/packages/docs-ui/src/components/CodeBlock/Actions/AskAi/index.tsx b/www/packages/docs-ui/src/components/CodeBlock/Actions/AskAi/index.tsx
new file mode 100644
index 0000000000..a796d74610
--- /dev/null
+++ b/www/packages/docs-ui/src/components/CodeBlock/Actions/AskAi/index.tsx
@@ -0,0 +1,59 @@
+"use client"
+
+import React from "react"
+import { useAiAssistant, useSiteConfig } from "../../../../providers"
+import { useAiAssistantChat } from "../../../../providers/AiAssistant/Chat"
+import clsx from "clsx"
+import { Tooltip } from "../../../Tooltip"
+import Image from "next/image"
+
+export type CodeBlockCopyActionProps = {
+  source: string
+  inHeader: boolean
+}
+
+export const CodeBlockAskAiAction = ({
+  source,
+  inHeader,
+}: CodeBlockCopyActionProps) => {
+  const { setChatOpened } = useAiAssistant()
+  const { setQuestion, loading } = useAiAssistantChat()
+  const { config } = useSiteConfig()
+
+  const handleClick = () => {
+    if (loading) {
+      return
+    }
+    setQuestion(`\`\`\`tsx\n${source.trim()}\n\`\`\`\n\nExplain the code above`)
+    setChatOpened(true)
+  }
+
+  return (
+    <Tooltip
+      text="Ask AI"
+      tooltipClassName="font-base"
+      className={clsx("group")}
+      innerClassName={clsx(
+        inHeader && "flex",
+        "h-fit rounded-docs_sm",
+        "group-hover:bg-medusa-contrast-bg-base-hover group-focus:bg-medusa-contrast-bg-base-hover"
+      )}
+    >
+      <span
+        className={clsx(
+          !inHeader && "p-[6px]",
+          inHeader && "p-[4.5px]",
+          "cursor-pointer"
+        )}
+        onClick={handleClick}
+      >
+        <Image
+          src={`${config.basePath}/images/ai-assistent-luminosity.png`}
+          width={15}
+          height={15}
+          alt="Ask AI"
+        />
+      </span>
+    </Tooltip>
+  )
+}
diff --git a/www/packages/docs-ui/src/components/CodeBlock/Actions/Copy/index.tsx b/www/packages/docs-ui/src/components/CodeBlock/Actions/Copy/index.tsx
index 792ea70d5c..935b42301c 100644
--- a/www/packages/docs-ui/src/components/CodeBlock/Actions/Copy/index.tsx
+++ b/www/packages/docs-ui/src/components/CodeBlock/Actions/Copy/index.tsx
@@ -1,3 +1,5 @@
+"use client"
+
 import React, { useEffect, useState } from "react"
 import { CopyButton } from "../../../.."
 import clsx from "clsx"
diff --git a/www/packages/docs-ui/src/components/CodeBlock/Actions/index.tsx b/www/packages/docs-ui/src/components/CodeBlock/Actions/index.tsx
index 8e63eabcaa..f040cf96ff 100644
--- a/www/packages/docs-ui/src/components/CodeBlock/Actions/index.tsx
+++ b/www/packages/docs-ui/src/components/CodeBlock/Actions/index.tsx
@@ -6,6 +6,7 @@ import { Link, Tooltip } from "@/components"
 import { ExclamationCircle, PlaySolid } from "@medusajs/icons"
 import { GITHUB_ISSUES_LINK } from "@/constants"
 import { CodeBlockCopyAction } from "./Copy"
+import { CodeBlockAskAiAction } from "./AskAi"
 
 export type CodeBlockActionsProps = {
   source: string
@@ -74,6 +75,7 @@ export const CodeBlockActions = ({
           ]
         )}
       >
+        <CodeBlockAskAiAction source={source} inHeader={inHeader} />
         {canShowApiTesting && (
           <Tooltip
             text="Test API"
diff --git a/www/packages/docs-ui/src/components/CodeBlock/index.tsx b/www/packages/docs-ui/src/components/CodeBlock/index.tsx
index 02acc99c7b..92ea3f8292 100644
--- a/www/packages/docs-ui/src/components/CodeBlock/index.tsx
+++ b/www/packages/docs-ui/src/components/CodeBlock/index.tsx
@@ -38,9 +38,13 @@ export type CodeBlockMetaFields = {
   noCopy?: boolean
   noReport?: boolean
   noLineNumbers?: boolean
+  noAskAi?: boolean
   collapsibleLines?: string
   expandButtonLabel?: string
   isTerminal?: boolean
+  forceNoTitle?: boolean
+  collapsed?: boolean
+  wrapperClassName?: string
 } & CodeBlockHeaderMeta
 
 export type CodeBlockStyle = "loud" | "subtle" | "inline"
@@ -48,14 +52,11 @@ export type CodeBlockStyle = "loud" | "subtle" | "inline"
 export type CodeBlockProps = {
   source: string
   lang?: string
-  wrapperClassName?: string
   innerClassName?: string
   className?: string
-  collapsed?: boolean
   blockStyle?: CodeBlockStyle
   children?: React.ReactNode
   style?: React.HTMLAttributes<HTMLDivElement>["style"]
-  forceNoTitle?: boolean
   animateTokenHighlights?: boolean
   overrideColors?: {
     bg?: string
@@ -277,6 +278,7 @@ export const CodeBlock = ({
     getCollapsedLinesElm,
     getNonCollapsedLinesElm,
     type: collapsibleType,
+    isCollapsible,
     ...collapsibleResult
   } = useCollapsibleCodeLines({
     collapsibleLinesStr: collapsibleLines,
@@ -462,7 +464,7 @@ export const CodeBlock = ({
                       isSingleLine={tokens.length <= 1}
                     />
                   )}
-                {collapsibleType === "end" && (
+                {collapsibleType === "end" && isCollapsible(tokens) && (
                   <>
                     <CodeBlockCollapsibleFade
                       type={collapsibleType}
diff --git a/www/packages/docs-ui/src/components/CodeMdx/index.tsx b/www/packages/docs-ui/src/components/CodeMdx/index.tsx
index 9a72cc2bab..83d644c223 100644
--- a/www/packages/docs-ui/src/components/CodeMdx/index.tsx
+++ b/www/packages/docs-ui/src/components/CodeMdx/index.tsx
@@ -2,7 +2,9 @@ import React from "react"
 import {
   CodeBlock,
   CodeBlockMetaFields,
+  CodeBlockProps,
   InlineCode,
+  InlineCodeProps,
   MermaidDiagram,
 } from "@/components"
 import { Npm2YarnCode } from "../Npm2YarnCode"
@@ -10,13 +12,21 @@ import { Npm2YarnCode } from "../Npm2YarnCode"
 export type CodeMdxProps = {
   className?: string
   children?: React.ReactNode
+  inlineCodeProps?: Partial<InlineCodeProps>
+  codeBlockProps?: Partial<CodeBlockProps>
 } & CodeBlockMetaFields
 
 // due to how mdx handles code blocks
 // it is required that a code block specify a language
 // to be considered a block. Otherwise, it will be
 // considered as inline code
-export const CodeMdx = ({ className, children, ...rest }: CodeMdxProps) => {
+export const CodeMdx = ({
+  className,
+  children,
+  inlineCodeProps = {},
+  codeBlockProps = {},
+  ...rest
+}: CodeMdxProps) => {
   if (!children) {
     return <></>
   }
@@ -33,8 +43,15 @@ export const CodeMdx = ({ className, children, ...rest }: CodeMdxProps) => {
     } else if (match[1] === "mermaid") {
       return <MermaidDiagram diagramContent={codeContent} />
     }
-    return <CodeBlock source={codeContent} lang={match[1]} {...rest} />
+    return (
+      <CodeBlock
+        source={codeContent}
+        lang={match[1]}
+        {...codeBlockProps}
+        {...rest}
+      />
+    )
   }
 
-  return <InlineCode>{codeContent}</InlineCode>
+  return <InlineCode {...inlineCodeProps}>{codeContent}</InlineCode>
 }
diff --git a/www/packages/docs-ui/src/components/InlineCode/index.tsx b/www/packages/docs-ui/src/components/InlineCode/index.tsx
index 9d514954fb..cd84b9e593 100644
--- a/www/packages/docs-ui/src/components/InlineCode/index.tsx
+++ b/www/packages/docs-ui/src/components/InlineCode/index.tsx
@@ -4,9 +4,14 @@ import React from "react"
 import clsx from "clsx"
 import { CopyButton } from "@/components"
 
-export type InlineCodeProps = React.ComponentProps<"code">
+export type InlineCodeProps = React.ComponentProps<"code"> & {
+  variant?: "default" | "grey-bg"
+}
 
-export const InlineCode = (props: InlineCodeProps) => {
+export const InlineCode = ({
+  variant = "default",
+  ...props
+}: InlineCodeProps) => {
   return (
     <CopyButton
       text={props.children as string}
@@ -18,11 +23,18 @@ export const InlineCode = (props: InlineCodeProps) => {
       <code
         {...props}
         className={clsx(
-          "border-medusa-tag-neutral-border border",
-          "text-medusa-tag-neutral-text",
-          "bg-medusa-tag-neutral-bg font-monospace text-code-label rounded-docs_sm py-0 px-[6px]",
-          "group-active:bg-medusa-bg-subtle-pressed group-focus:bg-medusa-bg-subtle-pressed",
-          "group-hover:bg-medusa-tag-neutral-bg-hover",
+          "text-medusa-tag-neutral-text border",
+          "font-monospace text-code-label rounded-docs_sm py-0 px-[5px]",
+          variant === "default" && [
+            "bg-medusa-tag-neutral-bg group-hover:bg-medusa-tag-neutral-bg-hover",
+            "group-active:bg-medusa-bg-subtle-pressed group-focus:bg-medusa-bg-subtle-pressed",
+            "border-medusa-tag-neutral-border",
+          ],
+          variant === "grey-bg" && [
+            "bg-medusa-bg-switch-off group-hover:bg-medusa-bg-switch-off-hover",
+            "group-active:bg-medusa-bg-switch-off-hover group-focus:bg-medusa-switch-off-hover",
+            "border-medusa-border-strong",
+          ],
           props.className
         )}
       />
diff --git a/www/packages/docs-ui/src/hooks/use-ai-assistant-chat-navigation/index.ts b/www/packages/docs-ui/src/hooks/use-ai-assistant-chat-navigation/index.ts
index f7d411f5be..6dda96c8f6 100644
--- a/www/packages/docs-ui/src/hooks/use-ai-assistant-chat-navigation/index.ts
+++ b/www/packages/docs-ui/src/hooks/use-ai-assistant-chat-navigation/index.ts
@@ -7,6 +7,7 @@ import {
   useKeyboardShortcut,
   type useKeyboardShortcutOptions,
 } from "../use-keyboard-shortcut"
+import { useAiAssistantChat } from "../../providers/AiAssistant/Chat"
 
 export type UseAiAssistantChatNavigationProps = {
   getChatWindowElm: () => HTMLElement | null
@@ -23,6 +24,7 @@ export const useAiAssistantChatNavigation = ({
 }: UseAiAssistantChatNavigationProps) => {
   const shortcutKeys = useMemo(() => ["ArrowUp", "ArrowDown", "Enter"], [])
   const { chatOpened } = useAiAssistant()
+  const { question } = useAiAssistantChat()
 
   const handleKeyAction = (e: KeyboardEvent) => {
     const chatElm = getChatWindowElm()
@@ -30,7 +32,8 @@ export const useAiAssistantChatNavigation = ({
       !chatOpened ||
       e.metaKey ||
       e.ctrlKey ||
-      !chatElm?.contains(document.activeElement)
+      !chatElm?.contains(document.activeElement) ||
+      (question.length && question.includes("\n"))
     ) {
       return
     }
@@ -135,6 +138,7 @@ export const useAiAssistantChatNavigation = ({
     shortcutKeys: shortcutKeys,
     checkEditing: false,
     isLoading: false,
+    preventDefault: false,
     action: handleKeyAction,
     ...keyboardProps,
   })
diff --git a/www/packages/docs-ui/src/hooks/use-collapsible-code-lines/index.tsx b/www/packages/docs-ui/src/hooks/use-collapsible-code-lines/index.tsx
index c2335c1e19..9bd86aa786 100644
--- a/www/packages/docs-ui/src/hooks/use-collapsible-code-lines/index.tsx
+++ b/www/packages/docs-ui/src/hooks/use-collapsible-code-lines/index.tsx
@@ -34,7 +34,7 @@ export const useCollapsibleCodeLines = ({
   const collapsedRange:
     | {
         start: number
-        end: number
+        end: number | undefined
       }
     | undefined = useMemo(() => {
     if (!collapsibleLinesStr) {
@@ -46,8 +46,10 @@ export const useCollapsibleCodeLines = ({
       .map((lineNumber) => parseInt(lineNumber))
 
     if (
-      splitCollapsedLines.length !== 2 ||
-      (splitCollapsedLines[0] !== 1 && splitCollapsedLines[1] < 2)
+      !splitCollapsedLines.length ||
+      (splitCollapsedLines.length >= 2 &&
+        splitCollapsedLines[0] !== 1 &&
+        splitCollapsedLines[1] < 2)
     ) {
       return
     }
@@ -58,6 +60,13 @@ export const useCollapsibleCodeLines = ({
     }
   }, [collapsibleLinesStr])
 
+  const isCollapsible = useCallback(
+    (tokens: Token[][]) => {
+      return collapsedRange && collapsedRange.start < tokens.length
+    },
+    [collapsedRange]
+  )
+
   const type: CollapsedCodeLinesPosition | undefined = useMemo(() => {
     if (!collapsedRange) {
       return undefined
@@ -81,7 +90,7 @@ export const useCollapsibleCodeLines = ({
       tokens: Token[][]
       highlightProps: HighlightProps
     }) => {
-      if (!collapsedRange || !type) {
+      if (!collapsedRange || !type || !isCollapsible(tokens)) {
         return <></>
       }
 
@@ -90,7 +99,9 @@ export const useCollapsibleCodeLines = ({
 
       const lines = tokens.slice(
         startIndex,
-        Math.min(collapsedRange.end, tokens.length)
+        collapsedRange.end
+          ? Math.min(collapsedRange.end, tokens.length)
+          : tokens.length
       )
 
       return (
@@ -99,7 +110,7 @@ export const useCollapsibleCodeLines = ({
         </CodeBlockCollapsibleLines>
       )
     },
-    [collapsedRange, collapsibleHookResult]
+    [collapsedRange, collapsibleHookResult, isCollapsible, type]
   )
 
   const getNonCollapsedLinesElm = useCallback(
@@ -110,13 +121,13 @@ export const useCollapsibleCodeLines = ({
       tokens: Token[][]
       highlightProps: HighlightProps
     }) => {
-      if (!collapsedRange) {
+      if (!collapsedRange || !isCollapsible(tokens)) {
         return getLines(tokens, highlightProps)
       }
 
       const isCollapseBeginning = collapsedRange.start === 1
       const lines = tokens.slice(
-        isCollapseBeginning ? collapsedRange.end : 0,
+        isCollapseBeginning ? collapsedRange.end || tokens.length : 0,
         isCollapseBeginning ? undefined : collapsedRange.start
       )
 
@@ -126,13 +137,14 @@ export const useCollapsibleCodeLines = ({
         isCollapseBeginning ? collapsedRange.end : 0
       )
     },
-    [collapsedRange, collapsibleHookResult]
+    [collapsedRange, collapsibleHookResult, isCollapsible]
   )
 
   return {
     getCollapsedLinesElm,
     getNonCollapsedLinesElm,
     type,
+    isCollapsible,
     ...collapsibleHookResult,
   }
 }
diff --git a/www/packages/docs-ui/src/providers/AiAssistant/Chat/index.tsx b/www/packages/docs-ui/src/providers/AiAssistant/Chat/index.tsx
index 18e5b1d160..93c60a40c5 100644
--- a/www/packages/docs-ui/src/providers/AiAssistant/Chat/index.tsx
+++ b/www/packages/docs-ui/src/providers/AiAssistant/Chat/index.tsx
@@ -260,45 +260,47 @@ export const AiAssistantChatProvider = ({
 
   useEffect(() => {
     if (
-      !loading &&
-      answer.length &&
-      thread[lastAnswerIndex]?.content !== answer
+      loading ||
+      !answer.length ||
+      thread[lastAnswerIndex]?.content === answer
     ) {
-      const uniqueAnswerSources = answerSources
-        .filter(
-          (source, index) =>
-            answerSources.findIndex(
-              (s) => s.source_url === source.source_url
-            ) === index
-        )
-        .map((source) => {
-          const separatorIndex = source.title.indexOf("|")
-          return {
-            ...source,
-            title:
-              separatorIndex !== -1
-                ? source.title.slice(0, separatorIndex)
-                : source.title,
-          }
-        })
-      setThread((prevThread) => [
-        ...prevThread,
-        {
-          type: "answer",
-          content: answer,
-          question_id: identifiers?.question_answer_id,
-          order: getNewOrder(prevThread),
-          sources:
-            uniqueAnswerSources.length > 3
-              ? uniqueAnswerSources.slice(0, 3)
-              : uniqueAnswerSources,
-        },
-      ])
-      setAnswer("")
-      setAnswerSources([])
-      setMessagesCount((prev) => prev + 1)
-      inputRef.current?.focus()
+      return
     }
+
+    const uniqueAnswerSources = answerSources
+      .filter(
+        (source, index) =>
+          answerSources.findIndex((s) => s.source_url === source.source_url) ===
+          index
+      )
+      .map((source) => {
+        const separatorIndex = source.title.indexOf("|")
+        return {
+          ...source,
+          title:
+            separatorIndex !== -1
+              ? source.title.slice(0, separatorIndex)
+              : source.title,
+        }
+      })
+    setThread((prevThread) => [
+      ...prevThread,
+      {
+        type: "answer",
+        content: answer,
+        question_id: identifiers?.question_answer_id,
+        order: getNewOrder(prevThread),
+        sources:
+          uniqueAnswerSources.length > 3
+            ? uniqueAnswerSources.slice(0, 3)
+            : uniqueAnswerSources,
+      },
+    ])
+    setAnswer("")
+    setAnswerSources([])
+    setMessagesCount((prev) => prev + 1)
+    inputRef.current?.focus()
+    scrollToBottom()
   }, [loading, answer, thread, lastAnswerIndex, inputRef.current])
 
   useResizeObserver(contentRef as React.RefObject<HTMLDivElement>, () => {
diff --git a/www/packages/tailwind/base.tailwind.config.js b/www/packages/tailwind/base.tailwind.config.js
index 5408845386..26bbc8a411 100644
--- a/www/packages/tailwind/base.tailwind.config.js
+++ b/www/packages/tailwind/base.tailwind.config.js
@@ -508,8 +508,8 @@ module.exports = {
         "code-label": [
           "12px",
           {
-            lineHeight: "20px",
-            fontWeight: "400",
+            lineHeight: "15px",
+            fontWeight: "500",
           },
         ],
         "code-body": [
diff --git a/www/packages/tailwind/tailwind.config.js b/www/packages/tailwind/tailwind.config.js
index 044fcfeab6..b0e22a9368 100644
--- a/www/packages/tailwind/tailwind.config.js
+++ b/www/packages/tailwind/tailwind.config.js
@@ -1,12 +1,11 @@
 import path from "path"
 import coreConfig from "./modified.tailwind.config"
 
-// Get two levels up from require.resolve("@medusajs/ui")
 const root = path.join(require.resolve("docs-ui"), "../..")
-const uiPath = path.join(root, "**/*.{js,ts,jsx,tsx,mdx}")
+const files = path.join(root, "**/*.{js,ts,jsx,tsx,mdx}")
 
 /** @type {import('tailwindcss').Config} */
 module.exports = {
   ...coreConfig,
-  content: [uiPath],
+  content: [files],
 }