Repository: servo/servo
File: components/webdriver_server/lib.rs (Rust)

Code before:
        // Step 8.17 Remove an input source with input state and input id.
        self.session_mut()?.input_state_table.remove(&id);

        Ok(WebDriverResponse::Void)
    }

    fn take_screenshot(&self, rect: Option<Rect<f32, CSSPixel>>) -> WebDriverResult<String> {
        // Spec: Take screenshot after running the animation frame callbacks.
        let _ = self.handle_execute_async_script(JavascriptCommandParameters {
            script: "requestAnimationFrame(() => arguments[0]());".to_string(),
            args: None,
        });
        if rect.as_ref().is_some_and(Rect::is_empty) {
            return Err(WebDriverError::new(
                ErrorStatus::UnknownError,
                "The requested `rect` has zero width and/or height",
            ));
        }

        let webview_id = self.webview_id()?;
        let (sender, receiver) = crossbeam_channel::unbounded();
        self.send_message_to_embedder(WebDriverCommandMsg::TakeScreenshot(
            webview_id, rect, sender,
        ))?;

        let Ok(result) = receiver.recv_timeout(SCREENSHOT_TIMEOUT) else {
            return Err(WebDriverError::new(
                ErrorStatus::Timeout,
                "Timed out waiting to take screenshot. Test likely didn't finish.",
            ));
        };

        let image = result.map_err(|error| {
            WebDriverError::new(
                ErrorStatus::UnknownError,
                format!("Failed to take screenshot: {error:?}"),
            )
        })?;

        let mut png_data = Cursor::new(Vec::new());
        DynamicImage::ImageRgba8(image)
            .write_to(&mut png_data, ImageFormat::Png)
            .unwrap();

        Ok(base64::engine::general_purpose::STANDARD.encode(png_data.get_ref()))
    }

    fn handle_take_screenshot(&self) -> WebDriverResult<WebDriverResponse> {
        // Step 1. If session's current top-level browsing context is no longer open,
        // return error with error code no such window.
        let webview_id = self.webview_id()?;

Review comment (suggestion):
It is also possible that the error is not `RecvTimeoutError::Timeout`. Should use a match here to capture`Disconnected`.

Code after:
        // Step 8.17 Remove an input source with input state and input id.
        self.session_mut()?.input_state_table.remove(&id);

        Ok(WebDriverResponse::Void)
    }

    fn take_screenshot(&self, rect: Option<Rect<f32, CSSPixel>>) -> WebDriverResult<String> {
        // Spec: Take screenshot after running the animation frame callbacks.
        let _ = self.handle_execute_async_script(JavascriptCommandParameters {
            script: "requestAnimationFrame(() => arguments[0]());".to_string(),
            args: None,
        });
        if rect.as_ref().is_some_and(Rect::is_empty) {
            return Err(WebDriverError::new(
                ErrorStatus::UnknownError,
                "The requested `rect` has zero width and/or height",
            ));
        }

        let webview_id = self.webview_id()?;
        let (sender, receiver) = crossbeam_channel::unbounded();
        self.send_message_to_embedder(WebDriverCommandMsg::TakeScreenshot(
            webview_id, rect, sender,
        ))?;

        let result = match receiver.recv_timeout(SCREENSHOT_TIMEOUT) {
            Ok(result) => Ok(result),
            Err(RecvTimeoutError::Timeout) => Err(WebDriverError::new(
                ErrorStatus::Timeout,
                "Timed out waiting to take screenshot. Test likely didn't finish.",
            )),
            Err(RecvTimeoutError::Disconnected) => Err(WebDriverError::new(
                ErrorStatus::UnknownError,
                "Could not take screenshot because channel disconnected.",
            )),
        }?;

        let image = result.map_err(|error| {
            WebDriverError::new(
                ErrorStatus::UnknownError,
                format!("Failed to take screenshot: {error:?}"),
            )
        })?;

        let mut png_data = Cursor::new(Vec::new());
        DynamicImage::ImageRgba8(image)
            .write_to(&mut png_data, ImageFormat::Png)
            .unwrap();

        Ok(base64::engine::general_purpose::STANDARD.encode(png_data.get_ref()))
    }