From c4012a910fb42840969b350fa16c98c7cbd843db Mon Sep 17 00:00:00 2001
From: Akshita Goyal <36129505+gakshita@users.noreply.github.com>
Date: Fri, 22 Aug 2025 17:35:05 +0530
Subject: [PATCH] [PAI-708] fix: Pi workspace_id + scroll issue #3970

---
 .../pi-chat/conversation/ai-message.tsx       | 57 +++++++++----------
 .../pi-chat/conversation/messages.tsx         | 18 +++---
 .../pi-chat/conversation/reasoning.tsx        |  6 +-
 .../pi-chat/conversation/thinking.tsx         |  4 +-
 apps/web/ee/store/pi-chat/pi-chat.ts          |  6 +-
 5 files changed, 41 insertions(+), 50 deletions(-)

diff --git a/apps/web/ee/components/pi-chat/conversation/ai-message.tsx b/apps/web/ee/components/pi-chat/conversation/ai-message.tsx
index be91ecc201..7c2f319db7 100644
--- a/apps/web/ee/components/pi-chat/conversation/ai-message.tsx
+++ b/apps/web/ee/components/pi-chat/conversation/ai-message.tsx
@@ -20,9 +20,10 @@ type TProps = {
   isPiThinking?: boolean;
   isLoading?: boolean;
   feedback?: EFeedback;
+  isLatest?: boolean;
 };
 export const AiMessage = observer((props: TProps) => {
-  const { message = "", reasoning, isPiThinking = false, id, isLoading = false, feedback } = props;
+  const { message = "", reasoning, isPiThinking = false, id, isLoading = false, feedback, isLatest } = props;
   // store
   const { workspaceSlug } = useParams();
   const { sendFeedback, activeChatId, isPiTyping } = usePiChat();
@@ -65,42 +66,36 @@ export const AiMessage = observer((props: TProps) => {
       </div>
       <div className="flex flex-col text-base break-words w-full">
         {/* Message */}
-        {!isPiThinking && !isLoading && (
-          <div className="flex flex-col gap-4">
-            {reasoning && <ReasoningBlock reasoning={reasoning} isReasoning={!message} />}
-            <Markdown
-              remarkPlugins={[remarkGfm]}
-              className="pi-chat-root [&>*:first-child]:mt-0"
-              components={{
-                a: ({ children, href }) => (
-                  <a href={href || ""} target="_blank" rel="noopener noreferrer">
-                    {children}
-                  </a>
-                ),
-                table: ({ children }) => (
-                  <div className="overflow-x-auto w-full my-4 border-custom-border-200">
-                    <table className="min-w-full border-collapse">{children}</table>
-                  </div>
-                ),
-                th: ({ children }) => <th className="px-2 py-3 border-custom-border-200">{children}</th>,
-                td: ({ children }) => <td className="px-2 py-3 border-custom-border-200">{children}</td>,
-              }}
-            >
-              {message}
-            </Markdown>
-          </div>
-        )}
-
-        {/* Typing */}
-        {isPiThinking && <Thinking />}
-
+        <div className="flex flex-col gap-4">
+          {!isLoading && <ReasoningBlock reasoning={reasoning} showLoading={isPiThinking && isLatest} />}
+          <Markdown
+            remarkPlugins={[remarkGfm]}
+            className="pi-chat-root [&>*:first-child]:mt-0"
+            components={{
+              a: ({ children, href }) => (
+                <a href={href || ""} target="_blank" rel="noopener noreferrer">
+                  {children}
+                </a>
+              ),
+              table: ({ children }) => (
+                <div className="overflow-x-auto w-full my-4 border-custom-border-200">
+                  <table className="min-w-full border-collapse">{children}</table>
+                </div>
+              ),
+              th: ({ children }) => <th className="px-2 py-3 border-custom-border-200">{children}</th>,
+              td: ({ children }) => <td className="px-2 py-3 border-custom-border-200">{children}</td>,
+            }}
+          >
+            {message}
+          </Markdown>
+        </div>
         {isLoading && (
           <Loader>
             <Loader.Item width="50px" height="42px" />
           </Loader>
         )}
         {/* Action bar */}
-        {!isPiTyping && !isPiThinking && !isLoading && (
+        {message && (
           <div className="flex gap-4 mt-6">
             {/* Copy */}
             <Tooltip tooltipContent="Copy to clipboard" position="bottom" className="mb-4">
diff --git a/apps/web/ee/components/pi-chat/conversation/messages.tsx b/apps/web/ee/components/pi-chat/conversation/messages.tsx
index 6796d0820b..845b8a5b98 100644
--- a/apps/web/ee/components/pi-chat/conversation/messages.tsx
+++ b/apps/web/ee/components/pi-chat/conversation/messages.tsx
@@ -81,18 +81,16 @@ export const Messages = observer((props: TProps) => {
       {activeChat?.dialogue?.map((message: TDialogue, index: number) => (
         <div key={index} className="space-y-4">
           <MyMessage message={message.query} currentUser={currentUser} id={index.toString()} />
-          {(message.answer || message.reasoning) && (
-            <AiMessage
-              message={message.answer}
-              reasoning={message.reasoning}
-              id={index.toString()}
-              feedback={message.feedback}
-            />
-          )}
+          <AiMessage
+            message={message.answer}
+            reasoning={message.reasoning}
+            id={index.toString()}
+            feedback={message.feedback}
+            isPiThinking={isPiThinking}
+            isLatest={index === activeChat?.dialogue.length - 1}
+          />
         </div>
       ))}
-      {/* Typing */}
-      {isPiThinking && <AiMessage isPiThinking={isPiThinking} id={""} />}
 
       {/* Loading */}
       {isLoading && <AiMessage isLoading={isLoading} id={""} />}
diff --git a/apps/web/ee/components/pi-chat/conversation/reasoning.tsx b/apps/web/ee/components/pi-chat/conversation/reasoning.tsx
index 1cbc10a158..26cd6700f4 100644
--- a/apps/web/ee/components/pi-chat/conversation/reasoning.tsx
+++ b/apps/web/ee/components/pi-chat/conversation/reasoning.tsx
@@ -6,11 +6,11 @@ import { Thinking } from "./thinking";
 
 type TProps = {
   reasoning: string | undefined;
-  isReasoning: boolean;
+  showLoading?: boolean;
 };
 
 export const ReasoningBlock = (props: TProps) => {
-  const { reasoning, isReasoning } = props;
+  const { reasoning, showLoading } = props;
   const [isOpen, setIsOpen] = useState(false);
 
   return (
@@ -21,7 +21,7 @@ export const ReasoningBlock = (props: TProps) => {
         className="w-fit hover:bg-custom-background-90 rounded-full px-4 py-2 transition-all duration-500 ease-in-out border border-custom-border-200 hover:border-transparent"
       >
         <div className="flex gap-2 items-center">
-          {isReasoning ? (
+          {showLoading ? (
             <Thinking />
           ) : (
             <div className="text-base text-custom-text-200 font-medium">{isOpen ? "Hide" : "Show"} thinking</div>
diff --git a/apps/web/ee/components/pi-chat/conversation/thinking.tsx b/apps/web/ee/components/pi-chat/conversation/thinking.tsx
index b75fb12280..8556a20999 100644
--- a/apps/web/ee/components/pi-chat/conversation/thinking.tsx
+++ b/apps/web/ee/components/pi-chat/conversation/thinking.tsx
@@ -1,8 +1,8 @@
 import Typing from "./typing";
 
-export const Thinking = () => (
+export const Thinking = ({ text = "Thinking" }: { text?: string }) => (
   <div className="flex">
-    <span className="text-base">Thinking &nbsp;</span>
+    <span className="text-base">{text} &nbsp;</span>
     <Typing />
   </div>
 );
diff --git a/apps/web/ee/store/pi-chat/pi-chat.ts b/apps/web/ee/store/pi-chat/pi-chat.ts
index e39045fc39..1414872b0c 100644
--- a/apps/web/ee/store/pi-chat/pi-chat.ts
+++ b/apps/web/ee/store/pi-chat/pi-chat.ts
@@ -204,12 +204,12 @@ export class PiChatStore implements IPiChatStore {
     let payload: TInitPayload = {
       workspace_in_context: focus.isInWorkspaceContext,
       is_project_chat: isProjectChat,
+      workspace_id: workspaceId,
     };
     if (focus.isInWorkspaceContext) {
       payload = {
         ...payload,
         [focus.entityType]: focus.entityIdentifier,
-        workspace_id: workspaceId,
       };
     }
 
@@ -252,6 +252,7 @@ export class PiChatStore implements IPiChatStore {
     // 🔹 Handles `delta` chunks
     eventSource.addEventListener("delta", (event: MessageEvent) => {
       try {
+        if (this.isPiThinkingMap[chatId]) this.isPiThinkingMap[chatId] = false;
         const data = JSON.parse(event.data);
         callback("answer", data.chunk);
       } catch (e) {
@@ -268,9 +269,6 @@ export class PiChatStore implements IPiChatStore {
         console.error("Reasoning parse error", e);
       }
     });
-    eventSource.onopen = () => {
-      if (this.isPiThinkingMap[chatId]) this.isPiThinkingMap[chatId] = false;
-    };
 
     // 🔹 Handles done event
     eventSource.addEventListener("done", async () => {