533 lines
17 KiB
Rust
533 lines
17 KiB
Rust
|
mod browser;
|
||
|
|
mod client;
|
||
|
|
mod record;
|
||
|
|
mod translate;
|
||
|
|
|
||
|
|
use std::collections::HashMap;
|
||
|
|
use std::path::PathBuf;
|
||
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||
|
|
|
||
|
|
use async_trait::async_trait;
|
||
|
|
use lazyboy_contracts::{ActiveWindow, ComputerObservation, CursorPosition, UiElement};
|
||
|
|
use serde_json::{Value, json};
|
||
|
|
use tokio::time::{Duration, sleep};
|
||
|
|
|
||
|
|
use crate::controller::{
|
||
|
|
ComputerController, ComputerDriver, ControlContext, ControlError, ControllerHealth,
|
||
|
|
};
|
||
|
|
use crate::legacy::apply_action as legacy_apply_action;
|
||
|
|
use crate::process::spawn_detached;
|
||
|
|
use crate::{
|
||
|
|
ActionRequest, ActionResult, BrowserRequest, CdpPage, RecordingRequest, RecordingResult,
|
||
|
|
RecordingSession, action_pause_ms, normalize_display, observation_from_png,
|
||
|
|
observation_with_elements, teach_trajectory_dir,
|
||
|
|
};
|
||
|
|
use client::first_array_of_objects;
|
||
|
|
|
||
|
|
pub use client::CuaClient;
|
||
|
|
pub use translate::{TranslatedAction, translate_action};
|
||
|
|
|
||
|
|
#[derive(Debug, Default)]
|
||
|
|
pub struct CuaController {
|
||
|
|
client: CuaClient,
|
||
|
|
browser: tokio::sync::Mutex<HashMap<String, browser::BrowserBind>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[async_trait]
|
||
|
|
impl ComputerController for CuaController {
|
||
|
|
fn backend(&self) -> ComputerDriver {
|
||
|
|
ComputerDriver::Cua
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn health(&self, ctx: &ControlContext) -> Result<ControllerHealth, ControlError> {
|
||
|
|
let version = self.client.version().await.ok();
|
||
|
|
match self.client.status(&ctx.display).await {
|
||
|
|
Ok(status) => Ok(ControllerHealth {
|
||
|
|
backend: ComputerDriver::Cua.as_str().to_string(),
|
||
|
|
version,
|
||
|
|
healthy: true,
|
||
|
|
degraded: false,
|
||
|
|
details: vec![status.trim().to_string()],
|
||
|
|
}),
|
||
|
|
Err(error) => Ok(ControllerHealth {
|
||
|
|
backend: ComputerDriver::Cua.as_str().to_string(),
|
||
|
|
version,
|
||
|
|
healthy: false,
|
||
|
|
degraded: true,
|
||
|
|
details: vec![error.to_string()],
|
||
|
|
}),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn observe(&self, ctx: &ControlContext) -> Result<ComputerObservation, ControlError> {
|
||
|
|
self.browser
|
||
|
|
.lock()
|
||
|
|
.await
|
||
|
|
.remove(normalize_display(&ctx.display));
|
||
|
|
self.observe_display(&ctx.display).await
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn act(
|
||
|
|
&self,
|
||
|
|
request: &ActionRequest,
|
||
|
|
ctx: &ControlContext,
|
||
|
|
) -> Result<ActionResult, ControlError> {
|
||
|
|
let display = ctx.display.as_str();
|
||
|
|
let profile = ctx.profile_path.as_deref();
|
||
|
|
let mut completed = 0usize;
|
||
|
|
for action in &request.actions {
|
||
|
|
match translate_action(action, display, profile) {
|
||
|
|
Ok(translated) => self.dispatch(display, translated).await?,
|
||
|
|
Err(ControlError::Unsupported) => {
|
||
|
|
legacy_apply_action(display, profile, action).await?;
|
||
|
|
}
|
||
|
|
Err(error) => return Err(error),
|
||
|
|
}
|
||
|
|
let pause = action_pause_ms(action);
|
||
|
|
if pause > 0 {
|
||
|
|
sleep(Duration::from_millis(pause)).await;
|
||
|
|
}
|
||
|
|
completed += 1;
|
||
|
|
}
|
||
|
|
if request.settle_ms > 0 {
|
||
|
|
sleep(Duration::from_millis(u64::from(request.settle_ms))).await;
|
||
|
|
}
|
||
|
|
let observation = if request.observe {
|
||
|
|
Some(self.observe_display(display).await?)
|
||
|
|
} else {
|
||
|
|
None
|
||
|
|
};
|
||
|
|
Ok(ActionResult {
|
||
|
|
completed,
|
||
|
|
observation,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn browser(
|
||
|
|
&self,
|
||
|
|
request: &BrowserRequest,
|
||
|
|
ctx: &ControlContext,
|
||
|
|
) -> Result<CdpPage, ControlError> {
|
||
|
|
let key = normalize_display(&ctx.display).to_string();
|
||
|
|
let mut windows = self.windows(&ctx.display).await.unwrap_or_default();
|
||
|
|
let mut cache = self.browser.lock().await;
|
||
|
|
for attempt in 0..2 {
|
||
|
|
let mut bind = cache.remove(&key);
|
||
|
|
let had_bind = bind.is_some();
|
||
|
|
match browser::run(
|
||
|
|
&self.client,
|
||
|
|
&ctx.display,
|
||
|
|
ctx.profile_path.as_deref(),
|
||
|
|
request,
|
||
|
|
&windows,
|
||
|
|
&mut bind,
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
{
|
||
|
|
Err(ControlError::BrowserUnavailable) if !had_bind => {
|
||
|
|
return Ok(CdpPage {
|
||
|
|
ok: false,
|
||
|
|
error: Some("cdp unavailable".into()),
|
||
|
|
..CdpPage::default()
|
||
|
|
});
|
||
|
|
}
|
||
|
|
Err(error)
|
||
|
|
if attempt == 0
|
||
|
|
&& had_bind
|
||
|
|
&& matches!(
|
||
|
|
error,
|
||
|
|
ControlError::BrowserUnavailable
|
||
|
|
| ControlError::StaleReference
|
||
|
|
| ControlError::TargetNotFound
|
||
|
|
) =>
|
||
|
|
{
|
||
|
|
windows = self.windows(&ctx.display).await.unwrap_or_default();
|
||
|
|
}
|
||
|
|
other => {
|
||
|
|
if let Some(current) = bind {
|
||
|
|
cache.insert(key, current);
|
||
|
|
}
|
||
|
|
return map_browser_unavailable(other);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
map_browser_unavailable(Err(ControlError::BrowserUnavailable))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn start_recording(
|
||
|
|
&self,
|
||
|
|
request: &RecordingRequest,
|
||
|
|
ctx: &ControlContext,
|
||
|
|
) -> Result<RecordingSession, ControlError> {
|
||
|
|
if request.skill_id.trim().is_empty() {
|
||
|
|
return Err(ControlError::InvalidAction(
|
||
|
|
"recording needs a skill id".into(),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
let output_dir = teach_trajectory_dir(&request.skill_id);
|
||
|
|
let _ = tokio::fs::create_dir_all(&output_dir).await;
|
||
|
|
let _ = self
|
||
|
|
.client
|
||
|
|
.call(&ctx.display, "stop_recording", &json!({}), &[])
|
||
|
|
.await;
|
||
|
|
if let Err(error) = self
|
||
|
|
.client
|
||
|
|
.call(
|
||
|
|
&ctx.display,
|
||
|
|
"start_recording",
|
||
|
|
&json!({ "output_dir": output_dir, "record_video": false }),
|
||
|
|
&[],
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
{
|
||
|
|
tracing::warn!(error = %error, "cua start_recording failed");
|
||
|
|
}
|
||
|
|
if let Err(error) = crate::legacy::start_cdp_recorder(
|
||
|
|
&ctx.display,
|
||
|
|
ctx.profile_path.as_deref(),
|
||
|
|
&request.skill_id,
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
{
|
||
|
|
tracing::warn!(error = %error, "cdp recorder start failed");
|
||
|
|
}
|
||
|
|
Ok(RecordingSession {
|
||
|
|
skill_id: request.skill_id.clone(),
|
||
|
|
output_dir,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn stop_recording(
|
||
|
|
&self,
|
||
|
|
request: &RecordingRequest,
|
||
|
|
ctx: &ControlContext,
|
||
|
|
) -> Result<(), ControlError> {
|
||
|
|
let _ = self
|
||
|
|
.client
|
||
|
|
.call(&ctx.display, "stop_recording", &json!({}), &[])
|
||
|
|
.await;
|
||
|
|
crate::legacy::stop_cdp_recorder(&request.skill_id).await
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn collect_recording(
|
||
|
|
&self,
|
||
|
|
request: &RecordingRequest,
|
||
|
|
_ctx: &ControlContext,
|
||
|
|
) -> Result<RecordingResult, ControlError> {
|
||
|
|
let mut events = crate::legacy::collect_cdp_events(&request.skill_id).await;
|
||
|
|
let dir = teach_trajectory_dir(&request.skill_id);
|
||
|
|
events.extend(record::events_from_dir(std::path::Path::new(&dir)));
|
||
|
|
events.sort_by_key(|event| event.get("at").and_then(Value::as_i64).unwrap_or(0));
|
||
|
|
let _ = tokio::fs::remove_dir_all(&dir).await;
|
||
|
|
Ok(RecordingResult { events })
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl CuaController {
|
||
|
|
async fn observe_display(&self, display: &str) -> Result<ComputerObservation, ControlError> {
|
||
|
|
let png_path = observe_png_path(display);
|
||
|
|
let _ = tokio::fs::remove_file(&png_path).await;
|
||
|
|
self.client
|
||
|
|
.call(
|
||
|
|
display,
|
||
|
|
"get_desktop_state",
|
||
|
|
&json!({ "screenshot_out_file": png_path.to_string_lossy() }),
|
||
|
|
&["--screenshot-out-file", &png_path.to_string_lossy()],
|
||
|
|
)
|
||
|
|
.await?;
|
||
|
|
let png = tokio::fs::read(&png_path)
|
||
|
|
.await
|
||
|
|
.map_err(|_| ControlError::Internal("screenshot failed".into()))?;
|
||
|
|
let _ = tokio::fs::remove_file(&png_path).await;
|
||
|
|
if png.is_empty() {
|
||
|
|
return Err(ControlError::Internal("screenshot failed".into()));
|
||
|
|
}
|
||
|
|
let (width, height) = png_dimensions(&png);
|
||
|
|
let cursor = self.cursor(display).await;
|
||
|
|
let windows = self.windows(display).await.unwrap_or_default();
|
||
|
|
let active = windows
|
||
|
|
.iter()
|
||
|
|
.max_by_key(|window| window.z)
|
||
|
|
.map(|window| ActiveWindow {
|
||
|
|
id: window.id.to_string(),
|
||
|
|
title: Some(window.title.clone()).filter(|title| !title.is_empty()),
|
||
|
|
});
|
||
|
|
let elements: Vec<UiElement> = windows
|
||
|
|
.into_iter()
|
||
|
|
.enumerate()
|
||
|
|
.map(|(index, window)| UiElement {
|
||
|
|
id: (index + 1) as u32,
|
||
|
|
title: window.title.chars().take(80).collect(),
|
||
|
|
x: window.x,
|
||
|
|
y: window.y,
|
||
|
|
w: window.w,
|
||
|
|
h: window.h,
|
||
|
|
selector: None,
|
||
|
|
kind: Some("window".into()),
|
||
|
|
role: None,
|
||
|
|
})
|
||
|
|
.collect();
|
||
|
|
Ok(observation_with_elements(
|
||
|
|
observation_from_png(png, width, height, cursor, active),
|
||
|
|
elements,
|
||
|
|
))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn cursor(&self, display: &str) -> Option<CursorPosition> {
|
||
|
|
let value = self
|
||
|
|
.client
|
||
|
|
.call(display, "get_cursor_position", &json!({}), &[])
|
||
|
|
.await
|
||
|
|
.ok()?;
|
||
|
|
cursor_from_value(&value)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn windows(&self, display: &str) -> Result<Vec<ListedWindow>, ControlError> {
|
||
|
|
Self::list_windows_now(&self.client, display).await
|
||
|
|
}
|
||
|
|
|
||
|
|
pub(crate) async fn list_windows_now(
|
||
|
|
client: &CuaClient,
|
||
|
|
display: &str,
|
||
|
|
) -> Result<Vec<ListedWindow>, ControlError> {
|
||
|
|
let value = client
|
||
|
|
.call(
|
||
|
|
display,
|
||
|
|
"list_windows",
|
||
|
|
&json!({ "on_screen_only": true }),
|
||
|
|
&[],
|
||
|
|
)
|
||
|
|
.await?;
|
||
|
|
Ok(parse_listed_windows(&value))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn dispatch(
|
||
|
|
&self,
|
||
|
|
display: &str,
|
||
|
|
translated: TranslatedAction,
|
||
|
|
) -> Result<(), ControlError> {
|
||
|
|
match translated {
|
||
|
|
TranslatedAction::Sleep { ms } => {
|
||
|
|
sleep(Duration::from_millis(ms)).await;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
TranslatedAction::LegacyArgv { argv, detached } => {
|
||
|
|
if detached {
|
||
|
|
spawn_detached(&argv).await.map_err(ControlError::internal)
|
||
|
|
} else {
|
||
|
|
crate::process::run_output(&argv)
|
||
|
|
.await
|
||
|
|
.map_err(ControlError::internal)
|
||
|
|
.and_then(|output| {
|
||
|
|
if output.status.success() {
|
||
|
|
Ok(())
|
||
|
|
} else {
|
||
|
|
Err(ControlError::internal(String::from_utf8_lossy(
|
||
|
|
&output.stderr,
|
||
|
|
)))
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
TranslatedAction::FocusTitle { title } => self.focus_title(display, &title).await,
|
||
|
|
TranslatedAction::Cua { tool, payload } => {
|
||
|
|
let payload = self.with_window_target(display, tool, payload).await?;
|
||
|
|
self.client.call(display, tool, &payload, &[]).await?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn with_window_target(
|
||
|
|
&self,
|
||
|
|
display: &str,
|
||
|
|
tool: &str,
|
||
|
|
mut payload: Value,
|
||
|
|
) -> Result<Value, ControlError> {
|
||
|
|
if tool != "mouse_button_down" && tool != "mouse_button_up" {
|
||
|
|
return Ok(payload);
|
||
|
|
}
|
||
|
|
if payload.get("pid").is_some() && payload.get("window_id").is_some() {
|
||
|
|
return Ok(payload);
|
||
|
|
}
|
||
|
|
let window = self
|
||
|
|
.windows(display)
|
||
|
|
.await?
|
||
|
|
.into_iter()
|
||
|
|
.max_by_key(|window| window.z)
|
||
|
|
.ok_or(ControlError::TargetNotFound)?;
|
||
|
|
if let Some(object) = payload.as_object_mut() {
|
||
|
|
object.insert("pid".into(), json!(window.pid));
|
||
|
|
object.insert("window_id".into(), json!(window.id));
|
||
|
|
}
|
||
|
|
Ok(payload)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn focus_title(&self, display: &str, title: &str) -> Result<(), ControlError> {
|
||
|
|
let needle = title.to_ascii_lowercase();
|
||
|
|
let windows = self.windows(display).await?;
|
||
|
|
let window = windows
|
||
|
|
.into_iter()
|
||
|
|
.find(|window| window.title.to_ascii_lowercase().contains(&needle))
|
||
|
|
.ok_or(ControlError::TargetNotFound)?;
|
||
|
|
self.client
|
||
|
|
.call(
|
||
|
|
display,
|
||
|
|
"bring_to_front",
|
||
|
|
&json!({ "pid": window.pid, "window_id": window.id }),
|
||
|
|
&[],
|
||
|
|
)
|
||
|
|
.await?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Debug, Clone)]
|
||
|
|
pub(crate) struct ListedWindow {
|
||
|
|
pub(crate) id: u64,
|
||
|
|
pub(crate) pid: u64,
|
||
|
|
pub(crate) title: String,
|
||
|
|
pub(crate) app_name: String,
|
||
|
|
pub(crate) x: u32,
|
||
|
|
pub(crate) y: u32,
|
||
|
|
pub(crate) w: u32,
|
||
|
|
pub(crate) h: u32,
|
||
|
|
pub(crate) z: i64,
|
||
|
|
}
|
||
|
|
|
||
|
|
fn map_browser_unavailable(result: Result<CdpPage, ControlError>) -> Result<CdpPage, ControlError> {
|
||
|
|
match result {
|
||
|
|
Err(ControlError::BrowserUnavailable) => Ok(CdpPage {
|
||
|
|
ok: false,
|
||
|
|
error: Some("cdp unavailable".into()),
|
||
|
|
..CdpPage::default()
|
||
|
|
}),
|
||
|
|
other => other,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn cursor_from_value(value: &Value) -> Option<CursorPosition> {
|
||
|
|
let mut nodes = Vec::new();
|
||
|
|
client::walk(value, &mut nodes);
|
||
|
|
for node in nodes {
|
||
|
|
if let (Some(x), Some(y)) = (
|
||
|
|
node.get("x").and_then(Value::as_i64),
|
||
|
|
node.get("y").and_then(Value::as_i64),
|
||
|
|
) {
|
||
|
|
return Some(CursorPosition {
|
||
|
|
x: x as i32,
|
||
|
|
y: y as i32,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
None
|
||
|
|
}
|
||
|
|
|
||
|
|
fn parse_listed_windows(value: &Value) -> Vec<ListedWindow> {
|
||
|
|
first_array_of_objects(value, "window_id")
|
||
|
|
.into_iter()
|
||
|
|
.filter_map(listed_window)
|
||
|
|
.filter(|window| !ignored_window(window))
|
||
|
|
.collect()
|
||
|
|
}
|
||
|
|
|
||
|
|
fn listed_window(value: &Value) -> Option<ListedWindow> {
|
||
|
|
let bounds = value.get("bounds");
|
||
|
|
let x = number(value, "x").or_else(|| bounds.and_then(|bounds| number(bounds, "x")))?;
|
||
|
|
let y = number(value, "y").or_else(|| bounds.and_then(|bounds| number(bounds, "y")))?;
|
||
|
|
let w = number(value, "width")
|
||
|
|
.or_else(|| number(value, "w"))
|
||
|
|
.or_else(|| bounds.and_then(|bounds| number(bounds, "width")))?;
|
||
|
|
let h = number(value, "height")
|
||
|
|
.or_else(|| number(value, "h"))
|
||
|
|
.or_else(|| bounds.and_then(|bounds| number(bounds, "height")))?;
|
||
|
|
if w < 32 || h < 16 {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
Some(ListedWindow {
|
||
|
|
id: number(value, "window_id")?,
|
||
|
|
pid: number(value, "pid").unwrap_or(0),
|
||
|
|
title: value
|
||
|
|
.get("title")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or("")
|
||
|
|
.to_string(),
|
||
|
|
app_name: value
|
||
|
|
.get("app_name")
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.unwrap_or("")
|
||
|
|
.to_string(),
|
||
|
|
x: x as u32,
|
||
|
|
y: y as u32,
|
||
|
|
w: w as u32,
|
||
|
|
h: h as u32,
|
||
|
|
z: number(value, "z_index").unwrap_or(0) as i64,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
fn ignored_window(window: &ListedWindow) -> bool {
|
||
|
|
let title = window.title.to_ascii_lowercase();
|
||
|
|
title.is_empty() || title == "desktop" || title == "xfce4-panel" || title.contains("xfdesktop")
|
||
|
|
}
|
||
|
|
|
||
|
|
fn number(value: &Value, key: &str) -> Option<u64> {
|
||
|
|
value.get(key).and_then(Value::as_u64).or_else(|| {
|
||
|
|
value
|
||
|
|
.get(key)
|
||
|
|
.and_then(Value::as_f64)
|
||
|
|
.filter(|n| *n >= 0.0)
|
||
|
|
.map(|n| n as u64)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
fn png_dimensions(bytes: &[u8]) -> (u32, u32) {
|
||
|
|
image::load_from_memory(bytes)
|
||
|
|
.map(|image| (image.width(), image.height()))
|
||
|
|
.unwrap_or((1280, 800))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn observe_png_path(display: &str) -> PathBuf {
|
||
|
|
let nanos = SystemTime::now()
|
||
|
|
.duration_since(UNIX_EPOCH)
|
||
|
|
.map(|duration| duration.as_nanos())
|
||
|
|
.unwrap_or(0);
|
||
|
|
PathBuf::from(format!(
|
||
|
|
"/tmp/lazyboy/cua-obs-{}-{nanos}.png",
|
||
|
|
display.trim_start_matches(':')
|
||
|
|
))
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn window_list_skips_panel_and_numbers_from_one() {
|
||
|
|
let raw = json!([
|
||
|
|
{
|
||
|
|
"window_id": 1,
|
||
|
|
"pid": 8,
|
||
|
|
"title": "xfce4-panel",
|
||
|
|
"x": 0,
|
||
|
|
"y": 759,
|
||
|
|
"width": 1280,
|
||
|
|
"height": 41,
|
||
|
|
"z_index": 3
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"window_id": 9,
|
||
|
|
"pid": 20,
|
||
|
|
"title": "終端機",
|
||
|
|
"app_name": "xfce4-terminal",
|
||
|
|
"bounds": { "x": 53, "y": 55, "width": 753, "height": 699 },
|
||
|
|
"z_index": 1
|
||
|
|
}
|
||
|
|
]);
|
||
|
|
let windows = parse_listed_windows(&raw);
|
||
|
|
assert_eq!(windows.len(), 1);
|
||
|
|
assert_eq!(windows[0].title, "終端機");
|
||
|
|
assert_eq!(windows[0].app_name, "xfce4-terminal");
|
||
|
|
assert_eq!(windows[0].w, 753);
|
||
|
|
}
|
||
|
|
}
|