Webui - change setText command from parent window to also send the message.

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,10 +103,6 @@ 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);
   const messages: MessageDisplay[] = useMemo(() => {
@@ -132,7 +128,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 +137,10 @@ 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,17 @@ 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;

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,25 @@ 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>(
+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 +44,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(() => {