step one

Signed-off-by: Andrei Gherghescu <8067229+andrei-ng@users.noreply.github.com>

Author: andrei-ng

plotly_static/src/lib.rs

@@ -640,6 +640,23 @@ impl StaticExporterBuilder {
     }
 }
 
+// Async builder for async-first exporter (added without reordering existing items)
+impl StaticExporterBuilder {
+    /// Build an async exporter for use within async contexts.
+    pub fn build_async(&self) -> Result<AsyncStaticExporter> {
+        let wd = self.create_webdriver()?;
+        Ok(AsyncStaticExporter {
+            webdriver_port: self.webdriver_port,
+            webdriver_url: self.webdriver_url.clone(),
+            webdriver: wd,
+            offline_mode: self.offline_mode,
+            pdf_export_timeout: self.pdf_export_timeout,
+            webdriver_browser_caps: self.webdriver_browser_caps.clone(),
+            webdriver_client: None,
+        })
+    }
+}
+
 /// Main struct for exporting Plotly plots to static images.
 ///
 /// This struct provides methods to convert Plotly JSON plots into various
@@ -1120,6 +1137,265 @@ impl StaticExporter {
     pub fn get_webdriver_diagnostics(&self) -> String {
         self.webdriver.get_diagnostics()
     }
+
+    /// Explicitly close the WebDriver session and stop the driver.
+    /// Prefer calling this in long-running applications to ensure deterministic cleanup.
+    pub fn close(&mut self) -> Result<()> {
+        if let Some(client) = self.webdriver_client.take() {
+            let runtime = self.runtime.clone();
+            runtime.block_on(async {
+                if let Err(e) = client.close().await {
+                    error!("Failed to close WebDriver client: {e}");
+                }
+            });
+        }
+        if let Err(e) = self.webdriver.stop() {
+            error!("Failed to stop WebDriver: {e}");
+        }
+        Ok(())
+    }
+}
+
+/// Async-first exporter for async contexts. Keeps existing sync API intact.
+pub struct AsyncStaticExporter {
+    /// WebDriver server port (default: 4444)
+    webdriver_port: u32,
+
+    /// WebDriver server base URL (default: "http://localhost")
+    webdriver_url: String,
+
+    /// WebDriver process manager for spawning and cleanup
+    webdriver: WebDriver,
+
+    /// Use bundled JS libraries instead of CDN
+    offline_mode: bool,
+
+    /// PDF export timeout in milliseconds
+    pdf_export_timeout: u32,
+
+    /// Browser command-line flags (e.g., "--headless", "--no-sandbox")
+    webdriver_browser_caps: Vec<String>,
+
+    /// Cached WebDriver client for session reuse
+    webdriver_client: Option<Client>,
+}
+
+impl AsyncStaticExporter {
+    pub async fn write_fig(
+        &mut self,
+        dst: &Path,
+        plot: &serde_json::Value,
+        format: ImageFormat,
+        width: usize,
+        height: usize,
+        scale: f64,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let mut dst = PathBuf::from(dst);
+        dst.set_extension(format.to_string());
+
+        let plot_data = PlotData {
+            format: format.clone(),
+            width,
+            height,
+            scale,
+            data: plot,
+        };
+
+        let image_data = self.export(plot_data).await?;
+        let data = match format {
+            ImageFormat::SVG => image_data.as_bytes().to_vec(),
+            _ => general_purpose::STANDARD.decode(image_data)?,
+        };
+        let mut file = File::create(dst.as_path())?;
+        file.write_all(&data)?;
+        file.flush()?;
+
+        Ok(())
+    }
+
+    pub async fn write_to_string(
+        &mut self,
+        plot: &serde_json::Value,
+        format: ImageFormat,
+        width: usize,
+        height: usize,
+        scale: f64,
+    ) -> Result<String, Box<dyn std::error::Error>> {
+        let plot_data = PlotData {
+            format,
+            width,
+            height,
+            scale,
+            data: plot,
+        };
+        let image_data = self.export(plot_data).await?;
+        Ok(image_data)
+    }
+
+    async fn export(&mut self, plot: PlotData<'_>) -> Result<String> {
+        let data = self.static_export(&plot).await?;
+        Ok(data)
+    }
+
+    async fn static_export(&mut self, plot: &PlotData<'_>) -> Result<String> {
+        let html_content = template::get_html_body(self.offline_mode);
+        self.extract(&html_content, plot)
+            .await
+            .with_context(|| "Failed to extract static image from browser session")
+    }
+
+    async fn extract(&mut self, html_content: &str, plot: &PlotData<'_>) -> Result<String> {
+        let caps = self.build_webdriver_caps()?;
+        debug!("Use WebDriver and headless browser to export static plot");
+        let webdriver_url = format!("{}:{}", self.webdriver_url, self.webdriver_port);
+
+        // Reuse existing client or create new one
+        let client = if let Some(ref client) = self.webdriver_client {
+            debug!("Reusing existing WebDriver session");
+            client.clone()
+        } else {
+            debug!("Creating new WebDriver session");
+            let new_client = ClientBuilder::native()
+                .capabilities(caps)
+                .connect(&webdriver_url)
+                .await
+                .with_context(|| "WebDriver session error")?;
+            self.webdriver_client = Some(new_client.clone());
+            new_client
+        };
+
+        // For offline mode, write HTML to file to avoid data URI size limits since JS
+        // libraries are embedded in the file
+        let url = if self.offline_mode {
+            let temp_file = template::to_file(html_content)
+                .with_context(|| "Failed to create temporary HTML file")?;
+            format!("file://{}", temp_file.to_string_lossy())
+        } else {
+            // For online mode, use data URI (smaller size since JS is loaded from CDN)
+            format!("data:text/html,{}", encode(html_content))
+        };
+
+        // Open the HTML
+        client.goto(&url).await?;
+
+        let (js_script, args) = match plot.format {
+            ImageFormat::PDF => {
+                // Always use SVG for PDF export
+                let args = vec![
+                    plot.data.clone(),
+                    ImageFormat::SVG.to_string().into(),
+                    plot.width.into(),
+                    plot.height.into(),
+                    plot.scale.into(),
+                ];
+
+                (pdf_export_js_script(self.pdf_export_timeout), args)
+            }
+            _ => {
+                let args = vec![
+                    plot.data.clone(),
+                    plot.format.to_string().into(),
+                    plot.width.into(),
+                    plot.height.into(),
+                    plot.scale.into(),
+                ];
+
+                (image_export_js_script(), args)
+            }
+        };
+
+        let data = client.execute_async(&js_script, args).await?;
+
+        let result = data.as_str().ok_or(anyhow!(
+            "Failed to execute Plotly.toImage in browser session"
+        ))?;
+
+        if let Some(err) = result.strip_prefix("ERROR:") {
+            return Err(anyhow!("JavaScript error during export: {err}"));
+        }
+
+        match plot.format {
+            ImageFormat::SVG => StaticExporter::extract_plain(result, &plot.format),
+            ImageFormat::PNG | ImageFormat::JPEG | ImageFormat::WEBP | ImageFormat::PDF => {
+                StaticExporter::extract_encoded(result, &plot.format)
+            }
+            #[allow(deprecated)]
+            ImageFormat::EPS => {
+                error!("EPS format is deprecated. Use SVG or PDF instead.");
+                StaticExporter::extract_encoded(result, &plot.format)
+            }
+        }
+    }
+
+    fn build_webdriver_caps(&self) -> Result<Capabilities> {
+        // Define browser capabilities (copied to avoid reordering existing code)
+        let mut caps = JsonMap::new();
+        let mut browser_opts = JsonMap::new();
+        let browser_args = self.webdriver_browser_caps.clone();
+
+        browser_opts.insert("args".to_string(), serde_json::json!(browser_args));
+
+        // Add Chrome binary capability if BROWSER_PATH is set
+        #[cfg(feature = "chromedriver")]
+        if let Ok(chrome_path) = std::env::var("BROWSER_PATH") {
+            browser_opts.insert("binary".to_string(), serde_json::json!(chrome_path));
+            debug!("Added Chrome binary capability: {chrome_path}");
+        }
+        // Add Firefox binary capability if BROWSER_PATH is set
+        #[cfg(feature = "geckodriver")]
+        if let Ok(firefox_path) = std::env::var("BROWSER_PATH") {
+            browser_opts.insert("binary".to_string(), serde_json::json!(firefox_path));
+            debug!("Added Firefox binary capability: {firefox_path}");
+        }
+
+        // Add Firefox-specific preferences for CI environments
+        #[cfg(feature = "geckodriver")]
+        {
+            let prefs = StaticExporter::get_firefox_ci_preferences();
+            browser_opts.insert("prefs".to_string(), serde_json::json!(prefs));
+            debug!("Added Firefox preferences for CI compatibility");
+        }
+
+        caps.insert(
+            "browserName".to_string(),
+            serde_json::json!(get_browser_name()),
+        );
+        caps.insert(
+            get_options_key().to_string(),
+            serde_json::json!(browser_opts),
+        );
+
+        debug!("WebDriver capabilities: {caps:?}");
+
+        Ok(caps)
+    }
+
+    /// Close the WebDriver session and stop the driver if it was spawned.
+    pub async fn close(&mut self) -> Result<()> {
+        if let Some(client) = self.webdriver_client.take() {
+            if let Err(e) = client.close().await {
+                error!("Failed to close WebDriver client: {e}");
+            }
+        }
+        if let Err(e) = self.webdriver.stop() {
+            error!("Failed to stop WebDriver: {e}");
+        }
+        Ok(())
+    }
+
+    /// Get diagnostic information about the underlying WebDriver process.
+    pub fn get_webdriver_diagnostics(&self) -> String {
+        self.webdriver.get_diagnostics()
+    }
+}
+
+impl Drop for AsyncStaticExporter {
+    fn drop(&mut self) {
+        // Best-effort sync cleanup only; do not block on async.
+        if let Err(e) = self.webdriver.stop() {
+            error!("Failed to stop WebDriver: {e}");
+        }
+    }
 }
 
 #[cfg(test)]