easier approach

Author: ngxson

tools/server/public/index.html.gz

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { CallbackGeneratedChunk, useAppContext } from '../utils/app.context';
 import ChatMessage from './ChatMessage';
 import { CanvasType, Message, PendingMessage } from '../utils/types';
@@ -103,6 +103,10 @@ 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(() => {
@@ -128,7 +132,7 @@ export default function ChatScreen() {
     scrollToBottom(true);
   };
 
-  const sendNewMessage = useCallback(async () => {
+  const sendNewMessage = async () => {
     const lastInpMsg = textarea.value();
     if (lastInpMsg.trim().length === 0 || isGenerating(currConvId ?? ''))
       return;
@@ -137,10 +141,6 @@ 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,17 +155,10 @@ 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]);
+  // for vscode context
+  textarea.refOnSubmit.current = sendNewMessage;
 
   const handleEditMessage = async (msg: Message, content: string) => {
     if (!viewingChat) return;

tools/server/webui/src/components/ChatScreen.tsx

@@ -37,6 +37,7 @@ export interface ChatTextareaApi {
   setValue: (value: string) => void;
   focus: () => void;
   ref: React.RefObject<HTMLTextAreaElement>;
+  refOnSubmit: React.MutableRefObject<(() => void) | null>; // Submit handler
   onInput: (event: React.FormEvent<HTMLTextAreaElement>) => void; // Input handler
 }
 
@@ -46,6 +47,7 @@ export interface ChatTextareaApi {
 export function useChatTextarea(initValue: string): ChatTextareaApi {
   const [savedInitValue, setSavedInitValue] = useState<string>(initValue);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
+  const onSubmitRef = useRef<(() => void) | null>(null);
 
   // Effect to set initial value and height on mount or when initValue changes
   useEffect(() => {
@@ -91,6 +93,7 @@ export function useChatTextarea(initValue: string): ChatTextareaApi {
       }
     },
     ref: textareaRef,
+    refOnSubmit: onSubmitRef,
     onInput: handleInput,
   };
 }

tools/server/webui/src/components/useChatTextarea.ts

@@ -1,4 +1,4 @@
-import { useEffect, useRef, useState } from 'react';
+import { useEffect, useState } from 'react';
 import { MessageExtraContext } from './types';
 import { ChatTextareaApi } from '../components/useChatTextarea.ts';
 
@@ -15,25 +15,14 @@ interface SetTextEvData {
  * window.postMessage({ command: 'setText', text: 'Spot the syntax error', context: 'def test()\n  return 123' }, '*');
  */
 
-export const useVSCodeContext = (
-  textarea: ChatTextareaApi,
-  onSend?: () => void
-) => {
-  const [extraContext, _setExtraContext] = useState<MessageExtraContext | null>(
+export const useVSCodeContext = (textarea: ChatTextareaApi) => {
+  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);
@@ -44,18 +33,13 @@ export const useVSCodeContext = (
           });
         }
         textarea.focus();
-        if (onSend && data?.text) {
-          // Use setTimeout to ensure state updates are processed
-          setTimeout(() => {
-            onSend();
-          }, 50);
-        }
+        textarea.refOnSubmit.current?.();
       }
     };
 
     window.addEventListener('message', handleMessage);
     return () => window.removeEventListener('message', handleMessage);
-  }, [textarea, onSend]);
+  }, [textarea]);
 
   // Add a keydown listener that sends the "escapePressed" message to the parent window
   useEffect(() => {