mcp iframe updates

daemon/daemon.go

@@ -462,8 +462,17 @@ func (d *Daemon) handleCommand(w http.ResponseWriter, r *http.Request) {
 	case "switch-iframe":
 		tabID := cmd.Params["tab"]
 		selector := cmd.Params["selector"]
+		timeoutStr := cmd.Params["timeout"]
 
-		err := d.switchToIframe(tabID, selector)
+		// Parse timeout (default to 5 seconds if not specified)
+		timeout := 5
+		if timeoutStr != "" {
+			if parsedTimeout, err := strconv.Atoi(timeoutStr); err == nil && parsedTimeout > 0 {
+				timeout = parsedTimeout
+			}
+		}
+
+		err := d.switchToIframe(tabID, selector, timeout)
 		if err != nil {
 			response = Response{Success: false, Error: err.Error()}
 		} else {
@@ -1678,10 +1687,13 @@ func (d *Daemon) takeScreenshot(tabID, outputPath string, fullPage bool, timeout
 }
 
 // switchToIframe switches the context to an iframe for subsequent commands
-func (d *Daemon) switchToIframe(tabID, selector string) error {
+func (d *Daemon) switchToIframe(tabID, selector string, timeout int) error {
+	d.debugLog("Switching to iframe: selector=%s, tab=%s, timeout=%d", selector, tabID, timeout)
+
 	// Get the main page first (not iframe context)
 	actualTabID, err := d.getTabID(tabID)
 	if err != nil {
+		d.debugLog("Failed to get tab ID: %v", err)
 		return err
 	}
 
@@ -1691,49 +1703,155 @@ func (d *Daemon) switchToIframe(tabID, selector string) error {
 	// Get the main page (bypass iframe context)
 	mainPage, exists := d.tabs[actualTabID]
 	if !exists {
+		d.debugLog("Tab %s not in cache, trying to find it", actualTabID)
 		// Try to find it
 		mainPage, err = d.findPageByID(actualTabID)
 		if err != nil {
+			d.debugLog("Failed to find tab %s: %v", actualTabID, err)
 			return err
 		}
 		if mainPage == nil {
+			d.debugLog("Tab %s not found", actualTabID)
 			return fmt.Errorf("tab not found: %s", actualTabID)
 		}
 		d.tabs[actualTabID] = mainPage
 	}
 
-	// Find the iframe element
-	iframeElement, err := mainPage.Element(selector)
+	d.debugLog("Found main page for tab %s, looking for iframe element", actualTabID)
+
+	// Find the iframe element with timeout
+	var iframeElement *rod.Element
+	if timeout > 0 {
+		// Use timeout context for finding the iframe element
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
+		defer cancel()
+
+		// Create a channel to signal completion
+		done := make(chan struct {
+			element *rod.Element
+			err     error
+		}, 1)
+
+		// Execute the element search in a goroutine
+		go func() {
+			defer func() {
+				if r := recover(); r != nil {
+					done <- struct {
+						element *rod.Element
+						err     error
+					}{nil, fmt.Errorf("iframe element search panicked: %v", r)}
+				}
+			}()
+
+			element, err := mainPage.Timeout(time.Duration(timeout) * time.Second).Element(selector)
+			done <- struct {
+				element *rod.Element
+				err     error
+			}{element, err}
+		}()
+
+		// Wait for either completion or timeout
+		select {
+		case result := <-done:
+			iframeElement = result.element
+			err = result.err
+		case <-ctx.Done():
+			d.debugLog("Iframe element search timed out after %d seconds", timeout)
+			return fmt.Errorf("failed to find iframe element (timeout after %ds): %s", timeout, selector)
+		}
+	} else {
+		// No timeout
+		iframeElement, err = mainPage.Element(selector)
+	}
+
 	if err != nil {
+		d.debugLog("Failed to find iframe element: %v", err)
 		return fmt.Errorf("failed to find iframe element: %w", err)
 	}
 
-	// Get the iframe's page context
-	iframePage, err := iframeElement.Frame()
+	d.debugLog("Found iframe element, getting frame context")
+
+	// Get the iframe's page context with timeout
+	var iframePage *rod.Page
+	if timeout > 0 {
+		// Use timeout context for getting the frame
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
+		defer cancel()
+
+		// Create a channel to signal completion
+		done := make(chan struct {
+			page *rod.Page
+			err  error
+		}, 1)
+
+		// Execute the frame access in a goroutine
+		go func() {
+			defer func() {
+				if r := recover(); r != nil {
+					done <- struct {
+						page *rod.Page
+						err  error
+					}{nil, fmt.Errorf("iframe frame access panicked: %v", r)}
+				}
+			}()
+
+			page, err := iframeElement.Frame()
+			done <- struct {
+				page *rod.Page
+				err  error
+			}{page, err}
+		}()
+
+		// Wait for either completion or timeout
+		select {
+		case result := <-done:
+			iframePage = result.page
+			err = result.err
+		case <-ctx.Done():
+			d.debugLog("Iframe frame access timed out after %d seconds", timeout)
+			return fmt.Errorf("failed to get iframe context (timeout after %ds)", timeout)
+		}
+	} else {
+		// No timeout
+		iframePage, err = iframeElement.Frame()
+	}
+
 	if err != nil {
+		d.debugLog("Failed to get iframe context: %v", err)
 		return fmt.Errorf("failed to get iframe context: %w", err)
 	}
 
 	// Store the iframe page context
 	d.iframePages[actualTabID] = iframePage
+	d.debugLog("Successfully switched to iframe context for tab %s", actualTabID)
 
 	return nil
 }
 
 // switchToMain switches back to the main page context
 func (d *Daemon) switchToMain(tabID string) error {
+	d.debugLog("Switching back to main context: tab=%s", tabID)
+
 	// Get the tab ID to use (may be the current tab)
 	actualTabID, err := d.getTabID(tabID)
 	if err != nil {
+		d.debugLog("Failed to get tab ID: %v", err)
 		return err
 	}
 
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
-	// Remove the iframe context for this tab
-	delete(d.iframePages, actualTabID)
+	// Check if there was an iframe context to remove
+	if _, exists := d.iframePages[actualTabID]; exists {
+		d.debugLog("Removing iframe context for tab %s", actualTabID)
+		// Remove the iframe context for this tab
+		delete(d.iframePages, actualTabID)
+	} else {
+		d.debugLog("No iframe context found for tab %s", actualTabID)
+	}
 
+	d.debugLog("Successfully switched back to main context for tab %s", actualTabID)
 	return nil
 }