Merge pull request #85 from PathOnAI/search-algorithm-cleanup

fix the bug to terminate the correct browserbase session when click end button

Author: TataKKKL

visual-tree-search-app/components/MessageLogPanel.tsx

@@ -38,6 +38,7 @@ interface Message {
 interface MessageLogPanelProps {
   messages: Message[];
   messagesEndRef?: React.RefObject<HTMLDivElement | null>;
+  onSessionIdChange?: (sessionId: string) => void;
 }
 
 interface ParsedMessage {
@@ -62,20 +63,35 @@ interface ParsedMessage {
   server_info?: {
     hostname?: string;
   };
+  session_id?: string;
 }
 
 interface PathStep {
   natural_language_description: string;
   action: string;
 }
 
-const MessageLogPanel: React.FC<MessageLogPanelProps> = ({ messages, messagesEndRef }) => {
+const MessageLogPanel: React.FC<MessageLogPanelProps> = ({ messages, messagesEndRef, onSessionIdChange }) => {
   useEffect(() => {
     if (messagesEndRef?.current) {
       messagesEndRef.current.scrollIntoView({ behavior: 'smooth' });
     }
   }, [messages, messagesEndRef]);
 
+  useEffect(() => {
+    if (!messages.length || !onSessionIdChange) return;
+    
+    const latestMessage = messages[messages.length - 1];
+    try {
+      const data = JSON.parse(latestMessage.content);
+      if (data.type === 'browser_setup' && data.status === 'success' && data.session_id) {
+        onSessionIdChange(data.session_id);
+      }
+    } catch {
+      // Not JSON or doesn't contain session info, ignore
+    }
+  }, [messages, onSessionIdChange]);
+
   const getCardStyle = (type: string) => {
     switch (type) {
       // System Status Messages
@@ -536,7 +552,7 @@ const MessageLogPanel: React.FC<MessageLogPanelProps> = ({ messages, messagesEnd
                   </div>
                 </div>
               </div>
-            </div>
+          </div>
           );
         })}
         {messagesEndRef && <div ref={messagesEndRef} />}

visual-tree-search-app/components/MessageLogPanelLATS.tsx

@@ -43,6 +43,7 @@ interface Message {
 interface MessageLogPanelProps {
   messages: Message[];
   messagesEndRef?: React.RefObject<HTMLDivElement | null>;
+  onSessionIdChange?: (sessionId: string) => void;
 }
 
 interface ParsedMessage {
@@ -75,20 +76,35 @@ interface ParsedMessage {
   step?: number;
   step_name?: string;
   iteration?: number;
+  session_id?: string;
 }
 
 interface PathStep {
   natural_language_description: string;
   action: string;
 }
 
-const MessageLogPanelLATS: React.FC<MessageLogPanelProps> = ({ messages, messagesEndRef }) => {
+const MessageLogPanelLATS: React.FC<MessageLogPanelProps> = ({ messages, messagesEndRef, onSessionIdChange }) => {
   useEffect(() => {
     if (messagesEndRef?.current) {
       messagesEndRef.current.scrollIntoView({ behavior: 'smooth' });
     }
   }, [messages, messagesEndRef]);
 
+  useEffect(() => {
+    if (!messages.length || !onSessionIdChange) return;
+    
+    const latestMessage = messages[messages.length - 1];
+    try {
+      const data = JSON.parse(latestMessage.content);
+      if (data.type === 'browser_setup' && data.status === 'success' && data.session_id) {
+        onSessionIdChange(data.session_id);
+      }
+    } catch {
+      // Not JSON or doesn't contain session info, ignore
+    }
+  }, [messages, onSessionIdChange]);
+
   const getCardStyle = (type: string) => {
     switch (type) {
       // System Status Messages
@@ -408,7 +424,8 @@ const MessageLogPanelLATS: React.FC<MessageLogPanelProps> = ({ messages, message
 
   const parseMessage = (content: string): ParsedMessage => {
     try {
-      return JSON.parse(content);
+      const data = JSON.parse(content);
+      return data;
     } catch {
       return { content };
     }

visual-tree-search-app/pages/LATSAgent.tsx

@@ -89,16 +89,7 @@ const LATSAgent = () => {
         wsRef.current.onopen = () => {
           logMessage('Connected to LATS WebSocket server');
           setConnected(true);
-          
-        //   const request = {
-        //     type: "start_search",
-        //     agent_type: "LATSAgent",
-        //     starting_url: searchParams.startingUrl,
-        //     goal: searchParams.goal,
-        //     max_depth: searchParams.maxDepth,
-        //     num_simulations: searchParams.num_simulations
-        //   };
-
+        
           const request = {
             type: "start_search",
             agent_type: "LATSAgent",
@@ -214,9 +205,18 @@ const LATSAgent = () => {
           <LATSVisual messages={messages} />
         </div>
 
+        {/* In LATSAgent.tsx */}
         <MessageLogPanelLATS
           messages={messages}
           messagesEndRef={messagesEndRef}
+          onSessionIdChange={(newSessionId) => {
+            // Handle session ID change
+            if (sessionId && sessionId !== newSessionId) {
+              // Terminate old session if needed
+              // ...terminate code...
+            }
+            setSessionId(newSessionId);
+          }}
         />
       </div>
     </div>

visual-tree-search-app/pages/SimpleSearchAgent.tsx

@@ -201,9 +201,18 @@ const SimpleSearchAgent = () => {
           <SimpleSearchVisual messages={messages} />
         </div>
 
+        {/* In LATSAgent.tsx */}
         <MessageLogPanel
           messages={messages}
           messagesEndRef={messagesEndRef}
+          onSessionIdChange={(newSessionId) => {
+            // Handle session ID change
+            if (sessionId && sessionId !== newSessionId) {
+              // Terminate old session if needed
+              // ...terminate code...
+            }
+            setSessionId(newSessionId);
+          }}
         />
       </div>
     </div>