platform/computer-use/src/desktop.rs

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

//! Select the desktop independently of the OS running the Agent.
#[cfg(feature = "native-desktop")]
use crate::session;
use crate::{Driver, protocol::Command};
use maka_agent_facts::{language::Language, tools::ToolError};
use maka_capabilities::computer::BrowserConnection;
use maka_capabilities::computer::Desktop;
use serde_json::Value;
use std::sync::Arc;
use std::{path::PathBuf, sync::OnceLock};
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;

pub(crate) mod worker;
pub use worker::bootstrap;

/// The embedding application owns distribution and version selection. Neither
/// the model nor plugin configuration can provide a downloader or executable.
pub type WindowsHelperResolver = fn() -> maka_agent_facts::tools::ToolFuture<PathBuf>;
static WINDOWS_HELPER: OnceLock<WindowsHelperResolver> = OnceLock::new();

pub fn register_windows_helper(resolver: WindowsHelperResolver) {
    assert!(
        WINDOWS_HELPER.set(resolver).is_ok(),
        "Windows helper resolver is configured once at application startup"
    );
}

pub(crate) fn local_platform() -> &'static str {
    match std::env::consts::OS {
        "macos" => "mac",
        other => other,
    }
}

fn is_wsl() -> bool {
    cfg!(target_os = "linux")
        && std::fs::read_to_string("/proc/sys/kernel/osrelease")
            .is_ok_and(|release| release.to_ascii_lowercase().contains("microsoft"))
}

fn use_windows(desktop: Desktop, platform: &str, wsl: bool) -> Result<bool, ToolError> {
    match desktop {
        Desktop::Local => Ok(false),
        Desktop::Auto => Ok(platform == "linux" && wsl),
        Desktop::Windows if platform == "windows" => Ok(false),
        Desktop::Windows if platform == "linux" && wsl => Ok(true),
        Desktop::Windows => Err(failed("Windows desktop requires Windows or WSL")),
    }
}

/// The selected desktop owns all targets, including browser connections. The
/// embedding Host still admits and journals every operation before invoking it.
pub struct Session {
    driver: Arc<Mutex<Driver>>,
    backend: Option<Backend>,
    closed: bool,
}
enum Backend {
    #[cfg(feature = "native-desktop")]
    Local(session::Session),
    Windows {
        id: Option<String>,
        browsers: Vec<BrowserConnection>,
    },
}
impl Session {
    pub fn new(driver: Arc<Mutex<Driver>>) -> Self {
        Self {
            driver,
            backend: None,
            closed: false,
        }
    }

    /// Snapshot display metadata from the same target owner used for execution.
    /// Remote/unknown bindings use a generic action, never a guessed window name.
    pub fn activity(&self, command: &Command) -> maka_agent_facts::display::Text {
        let target = match (&self.backend, command) {
            #[cfg(feature = "native-desktop")]
            (
                Some(Backend::Local(local)),
                Command::Observe { handle, .. } | Command::Action { handle, .. },
            ) => local.target_label(handle),
            _ => None,
        };
        crate::activity::command(command, target)
    }

    pub fn configure(
        &mut self,
        desktop: Desktop,
        browsers: Vec<BrowserConnection>,
        language: Language,
    ) -> Result<(), ToolError> {
        if self.closed {
            return Err(failed("Computer Use Session is closed"));
        }
        validate_connections(&browsers)?;
        let windows = use_windows(desktop, local_platform(), is_wsl())?;
        if let Some(backend) = &self.backend
            && windows != matches!(backend, Backend::Windows { .. })
        {
            return Err(failed(match language {
                Language::En => {
                    "desktop configuration changed; reset Cua before switching desktops"
                }
                Language::ZhCn => "桌面配置已改变；请先重置 Cua，再切换桌面",
                Language::ZhTw => "桌面設定已變更；請先重設 Cua，再切換桌面",
            }));
        }
        if self.backend.is_none() {
            self.backend = Some(if windows {
                Backend::Windows {
                    id: None,
                    browsers: Vec::new(),
                }
            } else {
                #[cfg(feature = "native-desktop")]
                {
                    Backend::Local(session::Session::new(self.driver.clone()))
                }
                #[cfg(not(feature = "native-desktop"))]
                return Err(native_unavailable());
            });
        }
        match self.backend.as_mut().unwrap() {
            #[cfg(feature = "native-desktop")]
            Backend::Local(local) => local.configure_browsers(browsers, language),
            Backend::Windows {
                browsers: current, ..
            } => {
                *current = browsers;
                Ok(())
            }
        }
    }

    pub fn configure_browsers(
        &mut self,
        browsers: Vec<BrowserConnection>,
        language: Language,
    ) -> Result<(), ToolError> {
        self.configure(Desktop::Auto, browsers, language)
    }

    pub fn platform(&self) -> &'static str {
        if matches!(self.backend, Some(Backend::Windows { .. })) {
            "windows"
        } else {
            local_platform()
        }
    }

    /// Prepare dependencies and the private channel outside the REPL deadline.
    /// This does not create a native Cua runtime or dispatch any input.
    pub async fn prepare(
        &mut self,
        cancellation: &CancellationToken,
        language: Language,
    ) -> Result<(), ToolError> {
        if cancellation.is_cancelled() {
            return Err(failed("Computer Use preparation cancelled"));
        }
        if self.closed {
            return Err(failed("Computer Use Session is closed"));
        }
        if self.backend.is_none() {
            self.configure(Desktop::Auto, Vec::new(), language)?;
        }
        if matches!(self.backend, Some(Backend::Windows { .. })) {
            let mut driver = tokio::select! {
                biased;
                _ = cancellation.cancelled() => return Err(failed("Computer Use preparation cancelled")),
                driver = self.driver.lock() => driver,
            };
            driver.windows.prepare(cancellation).await?;
        }
        Ok(())
    }

    pub async fn invoke(
        &mut self,
        command: Command,
        session: &str,
        cancellation: &CancellationToken,
        language: Language,
    ) -> Result<Value, ToolError> {
        if cancellation.is_cancelled() {
            return Err(failed("Computer Use cancelled before dispatch"));
        }
        if self.closed {
            return Err(failed("Computer Use Session is closed"));
        }
        if self.backend.is_none() {
            self.configure(Desktop::Auto, Vec::new(), language)?;
        }
        match self.backend.as_mut().unwrap() {
            #[cfg(feature = "native-desktop")]
            Backend::Local(local) => local.invoke(command, session, cancellation, language).await,
            Backend::Windows { id, browsers } => {
                if id.as_deref().is_some_and(|id| id != session) {
                    return Err(failed("Computer Use Session identity mismatch"));
                }
                id.get_or_insert_with(|| session.into());
                let mut driver = self.driver.lock().await;
                if cancellation.is_cancelled() {
                    return Err(failed("Computer Use cancelled before dispatch"));
                }
                driver
                    .windows
                    .call(
                        worker::Request::Invoke {
                            session: session.into(),
                            browsers: browsers.clone(),
                            command,
                            language,
                        },
                        cancellation,
                    )
                    .await
            }
        }
    }

    pub async fn close(&mut self) -> Result<(), ToolError> {
        self.closed = true;
        match &mut self.backend {
            #[cfg(feature = "native-desktop")]
            Some(Backend::Local(local)) => local.close().await,
            Some(Backend::Windows { id: Some(id), .. }) => {
                self.driver.lock().await.windows.release_session(id).await?;
                self.backend = None;
                Ok(())
            }
            _ => Ok(()),
        }
    }
}

pub(crate) fn validate_connections(connections: &[BrowserConnection]) -> Result<(), ToolError> {
    let mut ids = std::collections::HashSet::new();
    for connection in connections {
        let url = reqwest::Url::parse(&connection.endpoint).map_err(failed)?;
        let loopback = url.host_str().is_some_and(|host| {
            host == "localhost"
                || host
                    .trim_matches(['[', ']'])
                    .parse::<std::net::IpAddr>()
                    .is_ok_and(|ip| ip.is_loopback())
        });
        if !loopback
            || url.scheme() != "http"
            || !url.username().is_empty()
            || url.password().is_some()
            || url.query().is_some()
            || url.fragment().is_some()
            || url.path() != "/"
        {
            return Err(failed(
                "browser endpoint must be a loopback HTTP origin, without credentials, path or query",
            ));
        }
        if connection.id.is_empty() || !ids.insert(&connection.id) {
            return Err(failed("browser IDs must be nonempty and unique"));
        }
    }
    Ok(())
}

#[cfg(not(feature = "native-desktop"))]
fn native_unavailable() -> ToolError {
    failed(
        "Local desktop support is not included in this build; build Maka with the native-desktop feature, or use the Windows desktop from WSL",
    )
}

fn failed(error: impl std::fmt::Display) -> ToolError {
    ToolError::Failed(error.to_string())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn desktop_selection_never_falls_back_from_windows() {
        for platform in ["linux", "mac", "windows"] {
            for wsl in [false, true] {
                assert!(!use_windows(Desktop::Local, platform, wsl).unwrap());
                assert_eq!(
                    use_windows(Desktop::Auto, platform, wsl).unwrap(),
                    platform == "linux" && wsl
                );
                let result = use_windows(Desktop::Windows, platform, wsl);
                match platform {
                    "windows" => assert!(!result.unwrap()),
                    "linux" if wsl => assert!(result.unwrap()),
                    _ => assert!(
                        matches!(result, Err(ToolError::Failed(message)) if message == "Windows desktop requires Windows or WSL")
                    ),
                }
            }
        }
    }

    #[cfg(not(feature = "native-desktop"))]
    #[tokio::test]
    async fn bridge_build_rejects_local_configuration_and_keeps_windows_available() {
        let driver = Arc::new(Mutex::new(Driver::default()));
        let mut session = Session::new(driver.clone());
        let error = session
            .configure(Desktop::Local, Vec::new(), Language::En)
            .unwrap_err();
        assert_eq!(
            error.to_string(),
            "tool failed: Local desktop support is not included in this build; build Maka with the native-desktop feature, or use the Windows desktop from WSL"
        );
        // Auto may select the Windows worker only in WSL. Elsewhere it must
        // report the missing native capability during configuration too.
        if !is_wsl() {
            assert_eq!(
                session
                    .configure(Desktop::Auto, Vec::new(), Language::En)
                    .unwrap_err()
                    .to_string(),
                error.to_string()
            );
        } else {
            session
                .configure(Desktop::Auto, Vec::new(), Language::En)
                .unwrap();
            assert_eq!(session.platform(), "windows");
        }
        session.close().await.unwrap();
        driver.lock().await.close().await.unwrap();
    }
}