diff --git a/nx-dev/feature-ai/src/lib/error-message.tsx b/nx-dev/feature-ai/src/lib/error-message.tsx
index a7432776a577d2..afe3d075d03835 100644
--- a/nx-dev/feature-ai/src/lib/error-message.tsx
+++ b/nx-dev/feature-ai/src/lib/error-message.tsx
@@ -1,9 +1,10 @@
+import { type JSX, memo } from 'react';
 import {
-  XCircleIcon,
   ExclamationTriangleIcon,
+  XCircleIcon,
 } from '@heroicons/react/24/outline';
 
-export function ErrorMessage({ error }: { error: any }): JSX.Element {
+function ErrorMessage({ error }: { error: any }): JSX.Element {
   try {
     if (error.message) {
       error = JSON.parse(error.message);
@@ -57,3 +58,6 @@ export function ErrorMessage({ error }: { error: any }): JSX.Element {
     );
   }
 }
+
+const MemoErrorMessage = memo(ErrorMessage);
+export { MemoErrorMessage as ErrorMessage };
diff --git a/nx-dev/feature-ai/src/lib/feed-container.tsx b/nx-dev/feature-ai/src/lib/feed-container.tsx
index b16b433a7f73ec..7db01fcbdbd908 100644
--- a/nx-dev/feature-ai/src/lib/feed-container.tsx
+++ b/nx-dev/feature-ai/src/lib/feed-container.tsx
@@ -1,11 +1,12 @@
 import { sendCustomEvent } from '@nx/nx-dev/feature-analytics';
-import { RefObject, useEffect, useRef, useState } from 'react';
+import { type JSX, RefObject, useEffect, useRef, useState } from 'react';
 import { ErrorMessage } from './error-message';
 import { Feed } from './feed/feed';
 import { LoadingState } from './loading-state';
 import { Prompt } from './prompt';
 import { getQueryFromUid, storeQueryForUid } from '@nx/nx-dev/util-ai';
 import { Message, useChat } from 'ai/react';
+import { cx } from '@nx/nx-dev/ui-primitives';
 
 const assistantWelcome: Message = {
   id: 'first-custom-message',
@@ -16,27 +17,41 @@ const assistantWelcome: Message = {
 
 export function FeedContainer(): JSX.Element {
   const [error, setError] = useState<Error | null>(null);
-  const [startedReply, setStartedReply] = useState(false);
+  const [hasReply, setHasReply] = useState(false);
 
   const feedContainer: RefObject<HTMLDivElement> | undefined = useRef(null);
-  const { messages, input, handleInputChange, handleSubmit, isLoading } =
-    useChat({
-      api: '/api/query-ai-handler',
-      onError: (error) => {
-        setError(error);
-      },
-      onResponse: (_response) => {
-        setStartedReply(true);
-        sendCustomEvent('ai_query', 'ai', 'query', undefined, {
-          query: input,
-        });
-        setError(null);
-      },
-      onFinish: (response: Message) => {
-        setStartedReply(false);
-        storeQueryForUid(response.id, input);
-      },
-    });
+  const {
+    messages,
+    setMessages,
+    input,
+    handleInputChange,
+    handleSubmit,
+    stop: handleStopGenerating,
+    reload: handleRegenerate,
+    isLoading,
+  } = useChat({
+    api: '/api/query-ai-handler',
+    onError: (error) => {
+      setError(error);
+    },
+    onResponse: (_response) => {
+      setHasReply(true);
+      sendCustomEvent('ai_query', 'ai', 'query', undefined, {
+        query: input,
+      });
+      setError(null);
+    },
+    onFinish: (response: Message) => {
+      setHasReply(false);
+      storeQueryForUid(response.id, input);
+    },
+  });
+
+  const handleNewChat = () => {
+    setMessages([]);
+    setError(null);
+    setHasReply(false);
+  };
 
   useEffect(() => {
     if (feedContainer.current) {
@@ -71,25 +86,32 @@ export function FeedContainer(): JSX.Element {
               <div
                 ref={feedContainer}
                 data-document="main"
-                className="relative"
+                className="relative pb-24"
               >
                 <Feed
                   activity={!!messages.length ? messages : [assistantWelcome]}
-                  handleFeedback={(statement, chatItemUid) =>
-                    handleFeedback(statement, chatItemUid)
-                  }
+                  onFeedback={handleFeedback}
                 />
 
                 {/* Change this message if it's loading but it's writing as well  */}
-                {isLoading && !startedReply && <LoadingState />}
+                {isLoading && !hasReply && <LoadingState />}
                 {error && <ErrorMessage error={error} />}
 
-                <div className="sticky bottom-0 left-0 right-0 w-full pt-6 pb-4 bg-gradient-to-t from-white via-white dark:from-slate-900 dark:via-slate-900">
+                <div
+                  className={cx(
+                    'fixed bottom-0 left0 right-0 w-full py-4 px-4 lg:py-6 lg:px-0',
+                    'bg-gradient-to-t from-white via-white/75 dark:from-slate-900 dark:via-slate-900/75'
+                  )}
+                >
                   <Prompt
-                    handleSubmit={handleSubmit}
-                    handleInputChange={handleInputChange}
+                    onSubmit={handleSubmit}
+                    onInputChange={handleInputChange}
+                    onNewChat={handleNewChat}
+                    onStopGenerating={handleStopGenerating}
+                    onRegenerate={handleRegenerate}
                     input={input}
-                    isDisabled={isLoading}
+                    isGenerating={isLoading}
+                    showNewChatCta={!isLoading && hasReply}
                   />
                 </div>
               </div>
diff --git a/nx-dev/feature-ai/src/lib/feed/feed-answer.tsx b/nx-dev/feature-ai/src/lib/feed/feed-answer.tsx
index 678fb96ea834ec..9ede178fee62c3 100644
--- a/nx-dev/feature-ai/src/lib/feed/feed-answer.tsx
+++ b/nx-dev/feature-ai/src/lib/feed/feed-answer.tsx
@@ -66,7 +66,7 @@ export function FeedAnswer({
           <ReactMarkdown children={content} />
         </div>
         {!isFirst && (
-          <div className="group text-xs flex-1 md:flex md:justify-end gap-4 md:items-center text-slate-400 hover:text-slate-500 transition">
+          <div className="group text-md flex-1 md:flex md:justify-end gap-4 md:items-center text-slate-400 hover:text-slate-500 transition">
             {feedbackStatement ? (
               <p className="italic group-hover:flex">
                 {feedbackStatement === 'good'
@@ -89,7 +89,7 @@ export function FeedAnswer({
                 title="Bad"
               >
                 <span className="sr-only">Bad answer</span>
-                <HandThumbDownIcon className="h-5 w-5" aria-hidden="true" />
+                <HandThumbDownIcon className="h-6 w-6" aria-hidden="true" />
               </button>
               <button
                 className={cx(
@@ -101,7 +101,7 @@ export function FeedAnswer({
                 title="Good"
               >
                 <span className="sr-only">Good answer</span>
-                <HandThumbUpIcon className="h-5 w-5" aria-hidden="true" />
+                <HandThumbUpIcon className="h-6 w-6" aria-hidden="true" />
               </button>
             </div>
           </div>
diff --git a/nx-dev/feature-ai/src/lib/feed/feed.tsx b/nx-dev/feature-ai/src/lib/feed/feed.tsx
index 8b3a00cb383d7c..aaa313a0070636 100644
--- a/nx-dev/feature-ai/src/lib/feed/feed.tsx
+++ b/nx-dev/feature-ai/src/lib/feed/feed.tsx
@@ -4,10 +4,10 @@ import { Message } from 'ai/react';
 
 export function Feed({
   activity,
-  handleFeedback,
+  onFeedback,
 }: {
   activity: Message[];
-  handleFeedback: (statement: 'bad' | 'good', chatItemUid: string) => void;
+  onFeedback: (statement: 'bad' | 'good', chatItemUid: string) => void;
 }) {
   return (
     <div className="flow-root my-12">
@@ -21,7 +21,7 @@ export function Feed({
               <FeedAnswer
                 content={activityItem.content}
                 feedbackButtonCallback={(statement) =>
-                  handleFeedback(statement, activityItem.id)
+                  onFeedback(statement, activityItem.id)
                 }
                 isFirst={activityItemIdx === 0}
               />
diff --git a/nx-dev/feature-ai/src/lib/prompt.tsx b/nx-dev/feature-ai/src/lib/prompt.tsx
index 64fb78e4c69e10..dd40dafb01cecd 100644
--- a/nx-dev/feature-ai/src/lib/prompt.tsx
+++ b/nx-dev/feature-ai/src/lib/prompt.tsx
@@ -1,40 +1,114 @@
-import { ChangeEvent, FormEvent, useEffect, useRef } from 'react';
-import { PaperAirplaneIcon } from '@heroicons/react/24/outline';
+import { ChangeEvent, FormEvent, useEffect, useRef, useState } from 'react';
+import {
+  ArrowPathIcon,
+  PaperAirplaneIcon,
+  PlusIcon,
+  StopIcon,
+} from '@heroicons/react/24/outline';
 import { Button } from '@nx/nx-dev/ui-common';
 import Textarea from 'react-textarea-autosize';
 import { ChatRequestOptions } from 'ai';
+import { cx } from '@nx/nx-dev/ui-primitives';
 
 export function Prompt({
-  isDisabled,
-  handleSubmit,
-  handleInputChange,
+  isGenerating,
+  showNewChatCta,
+  onSubmit,
+  onInputChange,
+  onNewChat,
+  onStopGenerating,
+  onRegenerate,
   input,
 }: {
-  isDisabled: boolean;
-  handleSubmit: (
+  isGenerating: boolean;
+  showNewChatCta: boolean;
+  onSubmit: (
     e: FormEvent<HTMLFormElement>,
     chatRequestOptions?: ChatRequestOptions | undefined
   ) => void;
-  handleInputChange: (
+  onInputChange: (
     e: ChangeEvent<HTMLTextAreaElement> | ChangeEvent<HTMLInputElement>
   ) => void;
+  onNewChat: () => void;
+  onStopGenerating: () => void;
+  onRegenerate: () => void;
   input: string;
 }) {
   const formRef = useRef<HTMLFormElement>(null);
   const inputRef = useRef<HTMLTextAreaElement>(null);
+  const [isStopped, setStopped] = useState(false);
 
   useEffect(() => {
-    if (inputRef.current) {
-      inputRef.current.focus();
-    }
-  }, []);
+    if (!isGenerating) inputRef.current?.focus();
+  }, [isGenerating]);
+  const handleSubmit = (event: FormEvent<HTMLFormElement>) => {
+    setStopped(false);
+    if (inputRef.current?.value.trim()) onSubmit(event);
+    else event.preventDefault();
+  };
+
+  const handleNewChat = () => {
+    onNewChat();
+    inputRef.current?.focus();
+  };
+
+  const handleStopGenerating = () => {
+    setStopped(true);
+    onStopGenerating();
+    inputRef.current?.focus();
+  };
+
+  const handleOnRegenerate = () => {
+    setStopped(false);
+    onRegenerate();
+  };
 
   return (
     <form
       ref={formRef}
       onSubmit={handleSubmit}
-      className="relative flex gap-2 max-w-2xl mx-auto py-0 px-2 shadow-lg rounded-md border border-slate-300 bg-white dark:border-slate-900 dark:bg-slate-700"
+      className="relative flex gap-2 max-w-3xl mx-auto py-0 px-2 shadow-lg rounded-md border border-slate-300 bg-white dark:border-slate-900 dark:bg-slate-700"
     >
+      <div
+        className={cx(
+          'absolute -top-full left-1/2 mt-1 -translate-x-1/2',
+          'flex gap-4'
+        )}
+      >
+        {isGenerating && (
+          <Button
+            variant="secondary"
+            size="small"
+            className={cx('bg-white dark:bg-slate-900')}
+            onClick={handleStopGenerating}
+          >
+            <StopIcon aria-hidden="true" className="h-5 w-5" />
+            <span className="text-base">Stop generating</span>
+          </Button>
+        )}
+        {showNewChatCta && (
+          <Button
+            variant="secondary"
+            size="small"
+            className={cx('bg-white dark:bg-slate-900')}
+            onClick={handleNewChat}
+          >
+            <PlusIcon aria-hidden="true" className="h-5 w-5" />
+            <span className="text-base">New Chat</span>
+          </Button>
+        )}
+        {isStopped && (
+          <Button
+            variant="secondary"
+            size="small"
+            className={cx('bg-white dark:bg-slate-900')}
+            onClick={handleOnRegenerate}
+          >
+            <ArrowPathIcon aria-hidden="true" className="h-5 w-5" />
+            <span className="text-base">Regenerated</span>
+          </Button>
+        )}
+      </div>
       <div className="overflow-y-auto w-full h-full max-h-[300px]">
         <Textarea
           onKeyDown={(event) => {
@@ -49,10 +123,10 @@ export function Prompt({
           }}
           ref={inputRef}
           value={input}
-          onChange={handleInputChange}
+          onChange={onInputChange}
           id="query-prompt"
           name="query"
-          disabled={isDisabled}
+          disabled={isGenerating}
           className="block w-full p-0 resize-none bg-transparent text-sm placeholder-slate-500 pl-2 py-[1.3rem] focus-within:outline-none focus:placeholder-slate-400 dark:focus:placeholder-slate-300 dark:text-white focus:outline-none focus:ring-0 border-none disabled:cursor-not-allowed"
           placeholder="How does caching work?"
           rows={1}
@@ -63,7 +137,7 @@ export function Prompt({
           variant="primary"
           size="small"
           type="submit"
-          disabled={isDisabled}
+          disabled={isGenerating}
           className="self-end w-12 h-12 disabled:cursor-not-allowed"
         >
           <div hidden className="sr-only">
diff --git a/nx-dev/ui-common/src/lib/button.tsx b/nx-dev/ui-common/src/lib/button.tsx
index 080f8dd8bee9dd..a7e9bb9f112ada 100644
--- a/nx-dev/ui-common/src/lib/button.tsx
+++ b/nx-dev/ui-common/src/lib/button.tsx
@@ -13,6 +13,7 @@ type AllowedSizes = 'large' | 'default' | 'small';
 interface ButtonProps {
   variant?: AllowedVariants;
   size?: AllowedSizes;
+  rounded?: 'full' | 'default';
   children: ReactNode | ReactNode[];
 }
 
@@ -45,12 +46,15 @@ function ButtonInner({
   children,
   variant = 'primary',
   size = 'default',
+  rounded = 'default',
 }: ButtonProps): JSX.Element {
   return (
     <>
       <span
         className={cx(
-          'flex h-full w-full items-center justify-center whitespace-nowrap rounded-md border border-transparent font-medium shadow-sm transition',
+          'flex h-full w-full items-center justify-center whitespace-nowrap',
+          rounded === 'full' ? 'rounded-full' : 'rounded-md',
+          'border border-transparent font-medium shadow-sm transition',
           variantStyles[variant],
           sizes[size]
         )}
@@ -69,11 +73,12 @@ export function Button({
   className = '',
   variant = 'primary',
   size = 'large',
+  rounded = 'default',
   ...props
 }: ButtonProps & JSX.IntrinsicElements['button']): JSX.Element {
   return (
     <button {...props} className={getLayoutClassName(className)}>
-      <ButtonInner variant={variant} size={size}>
+      <ButtonInner variant={variant} size={size} rounded={rounded}>
         {children}
       </ButtonInner>
     </button>