Merge pull request #144 -- fix: chat mode tab selection and context handling

felarof99 · web-flow · commit 5a0bc02fb0ff · 2025-10-22T15:12:50.000-07:00
fix: chat mode tab selection and context handling
diff --git a/src/lib/agent/ChatAgent.prompt.ts b/src/lib/agent/ChatAgent.prompt.ts
@@ -17,88 +17,91 @@ interface ExtractedPageContext {
  * This is added ONCE at the beginning of a fresh conversation
  */
 export function generateSystemPrompt(): string {
-  return `You are a helpful AI assistant that can answer questions about web pages.
+  return `You are a helpful AI assistant for answering questions and providing information.
 
 ## Your Capabilities
-- You can analyze and understand web page content
-- You can answer questions based on the information provided
-- You have access to screenshot_tool for visual information
-- You have access to scroll_tool to navigate content
+- Answer questions naturally and conversationally
+- Analyze and understand web page content when it's available
+- Use screenshot_tool to view visual elements when needed
+- Use scroll_tool to navigate through page content when needed
 
-## Important: Browser State
-- The current web page content is provided in <browser-state> tags
-- Always refer to the content within <browser-state> tags when answering questions about the page
-- This browser state is automatically updated when tabs change
+## Operating Modes
+You operate in two distinct modes based on whether the user has selected page context:
 
-## Instructions
+**Mode 1: Page Context Selected**
+- When page content is provided, you must ONLY answer from that content
+- If the answer isn't in the provided page(s), politely say the information is not available on this page
+- Do NOT use general knowledge when page context is active
+
+**Mode 2: No Page Context (General Mode)**
+- When no page content is provided, freely use your general knowledge
+- Answer any question to the best of your ability
+
+## General Instructions
 1. Be concise and direct in your responses
-2. Answer based on the page content within <browser-state> tags
-3. Use tools only when necessary for answering the question
-4. Focus on providing accurate, helpful answers
+2. Never mention internal technical details like tags, data structures, or system implementation
+3. Use tools only when they would genuinely help answer the question better
+4. Speak naturally as if you're having a conversation with a person
 
-You're in Q&A mode. Provide direct answers without planning or task management.`
+You're in Q&A mode. Provide helpful, accurate answers in a natural conversational tone.`
 }
 
 /**
  * Generate page context message to be added as assistant message
  * This contains the actual page content extracted from tabs
  */
 export function generatePageContextMessage(pageContext: ExtractedPageContext, isUpdate: boolean = false): string {
-  const prefix = isUpdate 
-    ? "I've detected that the tabs have changed. Here's the updated page content:"
-    : "I've extracted the content from the current page(s). Here's what I found:"
+  // Handle case where user explicitly removed all tabs (no page context)
+  // Return empty string - ChatAgent will remove browser state entirely
+  if (pageContext.tabs.length === 0) {
+    return ''
+  }
 
+  // No verbose announcements - just provide the content cleanly
   if (pageContext.isSingleTab) {
-    return generateSingleTabContext(pageContext.tabs[0], prefix)
+    return generateSingleTabContext(pageContext.tabs[0])
   } else {
-    return generateMultiTabContext(pageContext.tabs, prefix)
+    return generateMultiTabContext(pageContext.tabs)
   }
 }
 
 /**
  * Generate context message for single tab
  */
-function generateSingleTabContext(tab: ExtractedPageContext['tabs'][0], prefix: string): string {
-  return `${prefix}
-
-**Page: ${tab.title}**
+function generateSingleTabContext(tab: ExtractedPageContext['tabs'][0]): string {
+  return `<browser-state>
+Page: ${tab.title}
 URL: ${tab.url}
 
-## Content:
-${tab.text}`
+${tab.text}
+</browser-state>
+
+IMPORTANT: The user has selected this page as context. You must ONLY answer questions based on the content above. If a question cannot be answered using this page's content, politely inform the user that the information is not available on this page. Do not use your general knowledge to answer questions unrelated to this page.`
 }
 
 /**
  * Generate context message for multiple tabs
  */
-function generateMultiTabContext(tabs: ExtractedPageContext['tabs'], prefix: string): string {
+function generateMultiTabContext(tabs: ExtractedPageContext['tabs']): string {
   const tabSections = tabs.map((tab, index) => `
-**Tab ${index + 1}: ${tab.title}**
+Tab ${index + 1}: ${tab.title}
 URL: ${tab.url}
 
 ${tab.text}`).join('\n\n---\n')
 
-  return `${prefix}
+  return `<browser-state>
+${tabSections}
+</browser-state>
 
-I'm analyzing ${tabs.length} tabs:
-
-${tabSections}`
+IMPORTANT: The user has selected these ${tabs.length} pages as context. You must ONLY answer questions based on the content above. If a question cannot be answered using these pages' content, politely inform the user that the information is not available on the selected pages. Do not use your general knowledge to answer questions unrelated to these pages.`
 }
 
 /**
- * Generate task prompt that wraps the user's query
- * This tells the LLM to refer to the BrowserState content
+ * Generate task prompt for the user's query
+ * Simply returns the query without wrapper - the agent naturally uses available context
  */
 export function generateTaskPrompt(query: string, contextJustExtracted: boolean): string {
-  if (contextJustExtracted) {
-    // Context was just extracted and added above
-    return `Based on the page content in the <browser-state> tags above, please answer the following question:
-
-"${query}"`
-  } else {
-    // Context already exists from previous extraction
-    return `Using the page content from the <browser-state> tags, please answer:
-
-"${query}"`
-  }
+  // Return the query directly without any wrapper
+  // The agent will naturally reference browser-state content when relevant
+  return query
 }
diff --git a/src/lib/agent/ChatAgent.ts b/src/lib/agent/ChatAgent.ts
@@ -1,5 +1,5 @@
 import { ExecutionContext } from '@/lib/runtime/ExecutionContext'
-import { MessageManager } from '@/lib/runtime/MessageManager'
+import { MessageManager, LLMMessageType } from '@/lib/runtime/MessageManager'
 import { ToolManager } from '@/lib/tools/ToolManager'
 import { ScrollTool, ScreenshotTool } from '@/lib/tools'
 import { generateSystemPrompt, generatePageContextMessage, generateTaskPrompt } from './ChatAgent.prompt'
@@ -103,18 +103,25 @@ export class ChatAgent {
         if (isFreshConversation) {
           // Fresh conversation: Clear and add simple system prompt once
           this.messageManager.clear()
-          
+
           // Simple system prompt - just sets the role
           const systemPrompt = generateSystemPrompt()
           this.messageManager.addSystem(systemPrompt)
-          
-          // Add page context as browser state message
+
+          // Add page context as browser state message (only if tabs exist)
           const contextMessage = generatePageContextMessage(pageContext, false)
-          this.messageManager.addBrowserState(contextMessage)
+          if (contextMessage) {
+            this.messageManager.addBrowserState(contextMessage)
+          }
         } else {
           // Tabs changed: replace browser state to remove old page content
           const contextMessage = generatePageContextMessage(pageContext, true)
-          this.messageManager.addBrowserState(contextMessage)
+          if (contextMessage) {
+            this.messageManager.addBrowserState(contextMessage)
+          } else {
+            // No tabs - remove browser state entirely to enable general knowledge mode
+            this.messageManager.removeMessagesByType(LLMMessageType.BROWSER_STATE)
+          }
         }
         
         // Update tracked tab IDs
@@ -169,26 +176,22 @@ export class ChatAgent {
    */
   private async _getCurrentTabIds(): Promise<Set<number>> {
     const selectedTabIds = this.executionContext.getSelectedTabIds()
-    
+
+    // If null or empty array, user deselected all tabs - return empty Set
+    if (!selectedTabIds || selectedTabIds.length === 0) {
+      return new Set()
+    }
+
     // Check if user has explicitly selected multiple tabs (using "@" selector)
-    // If only 1 tab or null, it's likely just the default current tab from NxtScape
-    const hasExplicitSelection = selectedTabIds && selectedTabIds.length > 1
-    
+    const hasExplicitSelection = selectedTabIds.length > 1
+
     if (hasExplicitSelection) {
       // User explicitly selected multiple tabs - use those
       return new Set(selectedTabIds)
     }
-    
-    // No explicit multi-tab selection - get the ACTUAL current active tab
-    // This ensures we detect tab changes even when user switches tabs between queries
-    try {
-      const currentPage = await this.executionContext.browserContext.getCurrentPage()
-      return new Set([currentPage.tabId])
-    } catch (error) {
-      // Fallback to ExecutionContext if getCurrentPage fails
-      Logging.log('ChatAgent', `Failed to get current page, using ExecutionContext: ${error}`, 'warning')
-      return new Set(selectedTabIds || [])
-    }
+
+    // Single tab selected - use it directly (don't auto-detect current page)
+    return new Set([selectedTabIds[0]])
   }
   
   /**
@@ -214,13 +217,22 @@ export class ChatAgent {
   private async _extractPageContext(): Promise<ExtractedPageContext> {
     // Get selected tab IDs from execution context
     const selectedTabIds = this.executionContext.getSelectedTabIds()
+
+    // If explicitly null, user removed all tabs - return empty context
+    if (selectedTabIds === null) {
+      return {
+        tabs: [],
+        isSingleTab: false
+      }
+    }
+
     const hasUserSelectedTabs = Boolean(selectedTabIds && selectedTabIds.length > 0)
-    
+
     // Get browser pages
     const pages = await this.executionContext.browserContext.getPages(
       hasUserSelectedTabs && selectedTabIds ? selectedTabIds : undefined
     )
-    
+
     if (pages.length === 0) {
       throw new Error('No tabs available for context extraction')
     }
diff --git a/src/lib/execution/Execution.ts b/src/lib/execution/Execution.ts
@@ -167,7 +167,8 @@ export class Execution {
       });
 
       // Set selected tab IDs for context
-      executionContext.setSelectedTabIds(this.options.tabIds || []);
+      // null means "explicitly no tabs", undefined/[] means "use default current tab"
+      executionContext.setSelectedTabIds(this.options.tabIds === null ? null : (this.options.tabIds || []));
       executionContext.startExecution(this.options.tabId || 0);
 
       // Evals2: start a session and attach parent span to context
diff --git a/src/lib/runtime/ExecutionContext.ts b/src/lib/runtime/ExecutionContext.ts
@@ -183,7 +183,7 @@ export class ExecutionContext {
     return this._chatMode;
   }
 
-  public setSelectedTabIds(tabIds: number[]): void {
+  public setSelectedTabIds(tabIds: number[] | null): void {
     this.selectedTabIds = tabIds;
   }
 
diff --git a/src/sidepanel/components/ChatInput.tsx b/src/sidepanel/components/ChatInput.tsx
@@ -30,6 +30,7 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
   const [showTabSelector, setShowTabSelector] = useState(false)
   const [showSlashPalette, setShowSlashPalette] = useState(false)
   const textareaRef = useRef<HTMLTextAreaElement>(null)
+  const chatModeInitializedRef = useRef<boolean>(false)
   const [providerOk, setProviderOk] = useState<boolean>(true)
   const [historyIndex, setHistoryIndex] = useState<number>(-1)
   const [draftBeforeHistory, setDraftBeforeHistory] = useState<string>('')
@@ -38,7 +39,7 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
   const messages = useChatStore(state => state.messages)
   const { chatMode } = useSettingsStore()
   const { sendMessage, addMessageListener, removeMessageListener, connected: portConnected } = useSidePanelPortMessaging()
-  const { getContextTabs, toggleTabSelection, clearSelectedTabs } = useTabsStore()
+  const { getContextTabs, toggleTabSelection, clearSelectedTabs, fetchOpenTabs } = useTabsStore()
   const { agents, loadAgents } = useAgentsStore()
   
   // Load agents from Chrome storage on mount
@@ -102,15 +103,65 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
       setInput(e.detail)
       textareaRef.current?.focus()
     }
-    
+
     window.addEventListener('setInputValue', handleSetInput as EventListener)
     return () => {
       window.removeEventListener('setInputValue', handleSetInput as EventListener)
     }
   }, [])
-  
 
-  
+  // Fetch tabs on mount to ensure currentTabId is set
+  useEffect(() => {
+    fetchOpenTabs(undefined, true)
+  }, [fetchOpenTabs])
+
+  // Auto-select current tab when entering chat mode
+  useEffect(() => {
+    if (chatMode && !chatModeInitializedRef.current) {
+      chatModeInitializedRef.current = true
+      fetchOpenTabs(undefined, true)
+
+      const timer = setTimeout(() => {
+        const state = useTabsStore.getState()
+        // Only auto-select if current tab hasn't been explicitly removed
+        if (state.currentTabId !== null && !state.selectedTabs.includes(state.currentTabId) && !state.isCurrentTabRemoved) {
+          toggleTabSelection(state.currentTabId)
+        }
+      }, 150)
+
+      return () => clearTimeout(timer)
+    } else if (!chatMode) {
+      chatModeInitializedRef.current = false
+      // Clear selections when exiting chat mode
+      clearSelectedTabs()
+    }
+  }, [chatMode, fetchOpenTabs, toggleTabSelection, clearSelectedTabs])
+
+  // Auto-select new tab when user switches tabs in chat mode
+  useEffect(() => {
+    if (!chatMode) return
+
+    const handleTabActivated = (activeInfo: chrome.tabs.TabActiveInfo) => {
+      fetchOpenTabs(undefined, true)
+
+      setTimeout(() => {
+        const state = useTabsStore.getState()
+        const newCurrentTabId = state.currentTabId
+        const currentSelectedTabs = state.selectedTabs
+        const isCurrentTabRemoved = state.isCurrentTabRemoved
+
+        // Only auto-select if current tab hasn't been explicitly removed by user
+        if (newCurrentTabId !== null && !currentSelectedTabs.includes(newCurrentTabId) && !isCurrentTabRemoved) {
+          toggleTabSelection(newCurrentTabId)
+        }
+      }, 150)
+    }
+
+    chrome.tabs.onActivated.addListener(handleTabActivated)
+    return () => chrome.tabs.onActivated.removeListener(handleTabActivated)
+  }, [chatMode, fetchOpenTabs, toggleTabSelection])
+
+
   const submitTask = (query: string) => {
     if (!query.trim()) return
     
@@ -127,7 +178,11 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
     
     // Get selected tab IDs from tabsStore
     const contextTabs = getContextTabs()
-    const tabIds = contextTabs.length > 0 ? contextTabs.map(tab => tab.id) : undefined
+    // In chat mode, if no tabs are selected, explicitly send null to prevent using current page
+    // In agent mode, undefined means "use current active tab"
+    const tabIds = contextTabs.length > 0
+      ? contextTabs.map(tab => tab.id)
+      : (chatMode ? null : undefined)
     
     // Send to background
     setProcessing(true)
@@ -138,11 +193,16 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
       chatMode  // Include chat mode setting
     })
     
-    // Clear input and selected tabs
+    // Clear input and history
     setInput('')
     setHistoryIndex(-1)
     setDraftBeforeHistory('')
-    clearSelectedTabs()
+
+    // Only clear selected tabs in agent mode (not in chat mode)
+    if (!chatMode) {
+      clearSelectedTabs()
+    }
+
     setShowTabSelector(false)
   }
   
@@ -298,8 +358,8 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
           </Button>
         </div> */}
         
-        {/* Selected tabs chips */}
-        {selectedContextTabs.length > 0 && (
+        {/* Selected tabs chips - only show in chat mode */}
+        {chatMode && selectedContextTabs.length > 0 && (
           <div className="px-2 mb-1">
             <div className="flex items-center gap-2 overflow-x-auto no-scrollbar py-1">
               {selectedContextTabs.map(tab => (
@@ -320,7 +380,7 @@ export function ChatInput({ isConnected, isProcessing }: ChatInputProps) {
                   </span>
                   <button
                     type="button"
-                    className="ml-1 inline-flex items-center justify-center w-4 h-4 rounded-full hover:bg-foreground/10 text-xs text-muted-foreground hover:text-foreground"
+                    className="ml-1 inline-flex items-center justify-center w-5 h-5 rounded-full hover:bg-foreground/10 text-sm text-muted-foreground hover:text-foreground"
                     aria-label={`Remove ${tab.title} from selection`}
                     onClick={() => handleRemoveSelectedTab(tab.id)}
                   >
diff --git a/src/sidepanel/stores/tabsStore.ts b/src/sidepanel/stores/tabsStore.ts

Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,7 @@ export class ExecutionContext {`
`183`	`183`	`return this._chatMode;`
`184`	`184`	`}`
`185`	`185`
`186`		`- public setSelectedTabIds(tabIds: number[]): void {`
	`186`	`+ public setSelectedTabIds(tabIds: number[] \| null): void {`
`187`	`187`	`this.selectedTabIds = tabIds;`
`188`	`188`	`}`
`189`	`189`