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 477627a538..23f15cef83 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
@@ -1,6 +1,11 @@
 import React, { useEffect, useMemo, useRef } from "react"
 import clsx from "clsx"
-import { ArrowUpCircleSolid, LightBulb, LightBulbSolid } from "@medusajs/icons"
+import {
+  ArrowUpCircleSolid,
+  LightBulb,
+  LightBulbSolid,
+  StopCircleSolid,
+} from "@medusajs/icons"
 import {
   useAiAssistant,
   useAnalytics,
@@ -20,7 +25,7 @@ export const AiAssistantChatWindowInput = ({
 }: AiAssistantChatWindowInputProps) => {
   const { chatOpened, inputRef, loading, setChatOpened, isCaptchaLoaded } =
     useAiAssistant()
-  const { submitQuery, conversation } = useChat()
+  const { submitQuery, conversation, stopGeneration } = useChat()
   const { track } = useAnalytics()
   const { active, toggle } = useDeepThinking()
   const { isBrowser } = useIsBrowser()
@@ -43,6 +48,12 @@ export const AiAssistantChatWindowInput = ({
     overrideQuestion?: string
   ) => {
     e?.preventDefault()
+    if (loading) {
+      // stop the generation
+      stopGeneration()
+      return
+    }
+
     submitQuery(overrideQuestion || question)
     if (!conversation.length) {
       track({
@@ -206,10 +217,10 @@ export const AiAssistantChatWindowInput = ({
               "appearance-none p-0 text-medusa-fg-base disabled:text-medusa-fg-disabled",
               "transition-colors"
             )}
-            disabled={!question || loading || !isCaptchaLoaded}
+            disabled={!isCaptchaLoaded || (!question && !loading)}
             type="submit"
           >
-            <ArrowUpCircleSolid />
+            {loading ? <StopCircleSolid /> : <ArrowUpCircleSolid />}
           </button>
         </div>
       </form>
diff --git a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/index.tsx b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/index.tsx
index 53600823fe..dd1e816215 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/ChatWindow/index.tsx
@@ -53,6 +53,7 @@ export const AiAssistantChatWindow = () => {
             content: item.question,
             sources: item.sources,
             question_id: item.id,
+            isGenerationAborted: item.isGenerationAborted,
           }}
         />
         <AiAssistantThreadItem
@@ -61,6 +62,7 @@ export const AiAssistantChatWindow = () => {
             content: item.answer,
             sources: item.sources,
             question_id: item.id,
+            isGenerationAborted: item.isGenerationAborted,
           }}
         />
       </Fragment>
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 510689d5e5..4397c18c72 100644
--- a/www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx
+++ b/www/packages/docs-ui/src/components/AiAssistant/ThreadItem/index.tsx
@@ -25,6 +25,13 @@ export const AiAssistantThreadItem = ({ item }: AiAssistantThreadItemProps) => {
 
     return !item.question_id && item.content.length === 0
   }, [item, error])
+  if (
+    item.isGenerationAborted &&
+    item.type === "answer" &&
+    !item.content.length
+  ) {
+    return null
+  }
   return (
     <div
       className={clsx(
diff --git a/www/packages/docs-ui/src/providers/AiAssistant/index.tsx b/www/packages/docs-ui/src/providers/AiAssistant/index.tsx
index e993c587cf..6281e46389 100644
--- a/www/packages/docs-ui/src/providers/AiAssistant/index.tsx
+++ b/www/packages/docs-ui/src/providers/AiAssistant/index.tsx
@@ -32,6 +32,7 @@ export type AiAssistantThreadItem = {
   content: string
   question_id?: string | null
   sources?: Source[]
+  isGenerationAborted?: boolean
 }
 
 const AiAssistantContext = createContext<AiAssistantContextType | null>(null)