From e2ec22f9abfd7fc1c44022b563b3bf2214cf820f Mon Sep 17 00:00:00 2001
From: igardev <ivailo.gardev@akros.ch>
Date: Mon, 5 May 2025 07:29:14 +0300
Subject: [PATCH] setText command from parent window for llama-vscode now sends
 the message automatically.

---
 .../webui/src/components/ChatScreen.tsx       | 19 ++++++++++-----
 tools/server/webui/src/utils/llama-vscode.ts  | 24 ++++++++++++++-----
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/tools/server/webui/src/components/ChatScreen.tsx b/tools/server/webui/src/components/ChatScreen.tsx
index 29ab5ea64f..dc111b262a 100644
--- a/tools/server/webui/src/components/ChatScreen.tsx
+++ b/tools/server/webui/src/components/ChatScreen.tsx
@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useState } from 'react';
+import {useCallback, useEffect, useMemo, useRef, useState} from 'react';
 import { CallbackGeneratedChunk, useAppContext } from '../utils/app.context';
 import ChatMessage from './ChatMessage';
 import { CanvasType, Message, PendingMessage } from '../utils/types';
@@ -103,9 +103,7 @@ export default function ChatScreen() {
 
   const textarea: ChatTextareaApi = useChatTextarea(prefilledMsg.content());
 
-  const { extraContext, clearExtraContext } = useVSCodeContext(textarea);
-  // TODO: improve this when we have "upload file" feature
-  const currExtra: Message['extra'] = extraContext ? [extraContext] : undefined;
+
 
   // keep track of leaf node for rendering
   const [currNodeId, setCurrNodeId] = useState<number>(-1);
@@ -132,7 +130,7 @@ export default function ChatScreen() {
     scrollToBottom(true);
   };
 
-  const sendNewMessage = async () => {
+  const sendNewMessage = useCallback(async () => {
     const lastInpMsg = textarea.value();
     if (lastInpMsg.trim().length === 0 || isGenerating(currConvId ?? ''))
       return;
@@ -141,6 +139,8 @@ export default function ChatScreen() {
     setCurrNodeId(-1);
     // get the last message node
     const lastMsgNodeId = messages.at(-1)?.msg.id ?? null;
+    // TODO: improve this when we have "upload file" feature
+    const currExtra = extraContextRef.current ? [extraContextRef.current] : undefined;
     if (
       !(await sendMessage(
         currConvId,
@@ -155,7 +155,14 @@ export default function ChatScreen() {
     }
     // OK
     clearExtraContext();
-  };
+  }, [textarea, currConvId, isGenerating, messages, sendMessage, onChunk]);
+
+  const { extraContext, clearExtraContext } = useVSCodeContext(textarea,sendNewMessage);
+  const extraContextRef = useRef(extraContext);
+
+  useEffect(() => {
+    extraContextRef.current = extraContext;
+  }, [extraContext]);
 
   const handleEditMessage = async (msg: Message, content: string) => {
     if (!viewingChat) return;
diff --git a/tools/server/webui/src/utils/llama-vscode.ts b/tools/server/webui/src/utils/llama-vscode.ts
index c45b0d3973..0e200cce5a 100644
--- a/tools/server/webui/src/utils/llama-vscode.ts
+++ b/tools/server/webui/src/utils/llama-vscode.ts
@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react';
+import {useEffect, useRef, useState} from 'react';
 import { MessageExtraContext } from './types';
 import { ChatTextareaApi } from '../components/useChatTextarea.ts';
 
@@ -15,14 +15,20 @@ interface SetTextEvData {
  * window.postMessage({ command: 'setText', text: 'Spot the syntax error', context: 'def test()\n  return 123' }, '*');
  */
 
-export const useVSCodeContext = (textarea: ChatTextareaApi) => {
-  const [extraContext, setExtraContext] = useState<MessageExtraContext | null>(
-    null
-  );
+export const useVSCodeContext = (textarea: ChatTextareaApi, onSend?: () => void) => {
+    const [extraContext, _setExtraContext] = useState<MessageExtraContext | null>(null);
 
+    // Use ref to store the latest value
+    const extraContextRef = useRef(extraContext);
+
+    const setExtraContext = (value: MessageExtraContext | null) => {
+        extraContextRef.current = value;
+        _setExtraContext(value);
+    };
   // Accept setText message from a parent window and set inputMsg and extraContext
   useEffect(() => {
     const handleMessage = (event: MessageEvent) => {
+      if (event.source !== window.parent) return;
       if (event.data?.command === 'setText') {
         const data: SetTextEvData = event.data;
         textarea.setValue(data?.text);
@@ -33,12 +39,18 @@ export const useVSCodeContext = (textarea: ChatTextareaApi) => {
           });
         }
         textarea.focus();
+        if (onSend && data?.text) {
+          // Use setTimeout to ensure state updates are processed
+          setTimeout(() => {
+              onSend()
+          }, 50);
+        }
       }
     };
 
     window.addEventListener('message', handleMessage);
     return () => window.removeEventListener('message', handleMessage);
-  }, [textarea]);
+  },[textarea, onSend]);
 
   // Add a keydown listener that sends the "escapePressed" message to the parent window
   useEffect(() => {