lazyBoy/crates/sandbox/src/docker.rs

use async_trait::async_trait;
use base64::Engine;
use lazyboy_contracts::ComputerCapabilities;
use lazyboy_contracts::{
    ActiveWindow, ComputerObservation, CursorPosition, DEFAULT_SCREEN_HEIGHT, DEFAULT_SCREEN_WIDTH,
    SandboxKind,
};
use lazyboy_control::{
    ActionRequest, ActionResult, AdapterContext, BrowserRequest, CdpPage, CommandRequest,
    CommandResult, ComputerRef, EnsureScreenRequest, EnsureScreenResult, FileEntry,
    ProvisionRequest, RecordingRequest, RecordingResult, RecordingSession, SandboxError,
    SandboxProvider, ScreenSession, image_dimensions, observation_from_png,
    observation_with_elements, parse_ui_elements,
};
use reqwest::Client;
use serde_json::Value;

pub struct DockerSandbox {
    client: Client,
    base_url: String,
    token: String,
}

impl DockerSandbox {
    pub fn new(base_url: String, token: String) -> Self {
        Self {
            client: Client::new(),
            base_url: base_url.trim_end_matches('/').to_string(),
            token,
        }
    }

    fn headers(&self, context: &AdapterContext) -> reqwest::header::HeaderMap {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert(
            reqwest::header::AUTHORIZATION,
            format!("Bearer {}", self.token).parse().unwrap(),
        );
        headers.insert("x-lazyboy-space-id", context.space_id.parse().unwrap());
        if let Some(bot_id) = &context.bot_id {
            headers.insert("x-lazyboy-bot-id", bot_id.parse().unwrap());
        }
        if let Some(display) = &context.display
            && let Ok(value) = display.parse()
        {
            headers.insert("x-lazyboy-display", value);
        }
        if let Some(profile) = &context.profile_path
            && let Ok(value) = profile.parse()
        {
            headers.insert("x-lazyboy-profile", value);
        }
        if let Some(slot) = context.screen_slot
            && let Ok(value) = slot.to_string().parse()
        {
            headers.insert("x-lazyboy-screen-slot", value);
        }
        headers
    }

    fn url(&self, path: &str) -> String {
        format!("{}{path}", self.base_url)
    }

    async fn post_lifecycle(
        &self,
        computer: &ComputerRef,
        context: &AdapterContext,
        action: &str,
    ) -> Result<(), SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/{action}", computer.id)))
            .headers(self.headers(context))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        if response.status().is_success() {
            Ok(())
        } else {
            Err(SandboxError::message(format!(
                "{action} failed: {}",
                response.status()
            )))
        }
    }
}

#[async_trait]
impl SandboxProvider for DockerSandbox {
    async fn provision(
        &self,
        request: ProvisionRequest,
        context: &AdapterContext,
    ) -> Result<ComputerRef, SandboxError> {
        let response = self
            .client
            .post(self.url("/computers"))
            .headers(self.headers(context))
            .json(&serde_json::json!({
                "homeKey": request.home_key,
                "homePath": request.home_path,
                "spaceId": context.space_id,
            }))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(SandboxError::message(format!(
                "provision failed: {status} {body}"
            )));
        }
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let id = body
            .get("id")
            .and_then(Value::as_str)
            .ok_or_else(|| SandboxError::message("missing computer id"))?
            .to_string();
        Ok(ComputerRef {
            id: id.clone(),
            home_key: request.home_key,
            kind: SandboxKind::Docker,
            provider_ref: id,
            fresh: body.get("resumed").and_then(Value::as_bool) != Some(true),
        })
    }

    async fn capabilities(
        &self,
        _computer: &ComputerRef,
        _context: &AdapterContext,
    ) -> Result<ComputerCapabilities, SandboxError> {
        Ok(ComputerCapabilities { multi_screen: true })
    }

    async fn ensure_screen(
        &self,
        computer: &ComputerRef,
        request: EnsureScreenRequest,
        context: &AdapterContext,
    ) -> Result<EnsureScreenResult, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/screens", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(SandboxError::message(format!(
                "ensure screen failed: {status} {body}"
            )));
        }
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        Ok(EnsureScreenResult {
            slot: body
                .get("slot")
                .and_then(Value::as_u64)
                .unwrap_or(request.slot as u64) as u32,
            display: body
                .get("display")
                .and_then(Value::as_str)
                .unwrap_or(":1")
                .to_string(),
            view_port: body
                .get("viewPort")
                .or_else(|| body.get("view_port"))
                .and_then(Value::as_u64)
                .unwrap_or(6080) as u16,
        })
    }

    async fn execute(
        &self,
        computer: &ComputerRef,
        request: CommandRequest,
        context: &AdapterContext,
    ) -> Result<CommandResult, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/exec", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(SandboxError::message(format!(
                "exec failed: {status} {body}"
            )));
        }
        response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))
    }

    async fn observe(
        &self,
        computer: &ComputerRef,
        context: &AdapterContext,
    ) -> Result<ComputerObservation, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/observe", computer.id)))
            .headers(self.headers(context))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        decode_observation(&body)
    }

    async fn act(
        &self,
        computer: &ComputerRef,
        request: ActionRequest,
        context: &AdapterContext,
    ) -> Result<ActionResult, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/act", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        Ok(ActionResult {
            completed: body.get("completed").and_then(Value::as_u64).unwrap_or(0) as usize,
            observation: if body.get("png_base64").is_some() {
                Some(decode_observation(&body)?)
            } else {
                None
            },
        })
    }

    async fn browser(
        &self,
        computer: &ComputerRef,
        request: BrowserRequest,
        context: &AdapterContext,
    ) -> Result<CdpPage, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/browser", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        serde_json::from_value(body).map_err(|error| SandboxError::message(error.to_string()))
    }

    async fn start_recording(
        &self,
        computer: &ComputerRef,
        request: RecordingRequest,
        context: &AdapterContext,
    ) -> Result<RecordingSession, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/recording/start", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        serde_json::from_value(body).map_err(|error| SandboxError::message(error.to_string()))
    }

    async fn stop_recording(
        &self,
        computer: &ComputerRef,
        request: RecordingRequest,
        context: &AdapterContext,
    ) -> Result<(), SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/recording/stop", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        if response.status().is_success() {
            Ok(())
        } else {
            Err(SandboxError::message(format!(
                "stop recording failed: {}",
                response.status()
            )))
        }
    }

    async fn collect_recording(
        &self,
        computer: &ComputerRef,
        request: RecordingRequest,
        context: &AdapterContext,
    ) -> Result<RecordingResult, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/recording/collect", computer.id)))
            .headers(self.headers(context))
            .json(&request)
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        serde_json::from_value(body).map_err(|error| SandboxError::message(error.to_string()))
    }

    async fn connect_screen(
        &self,
        computer: &ComputerRef,
        interactive: bool,
        context: &AdapterContext,
    ) -> Result<ScreenSession, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/screen-mode", computer.id)))
            .headers(self.headers(context))
            .json(&serde_json::json!({
                "interactive": interactive,
                "slot": context.screen_slot.unwrap_or(0),
            }))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        Ok(ScreenSession {
            url: body
                .get("screenUrl")
                .and_then(Value::as_str)
                .map(str::to_string),
            interactive,
        })
    }

    async fn list_files(
        &self,
        computer: &ComputerRef,
        path: &str,
        context: &AdapterContext,
    ) -> Result<Vec<FileEntry>, SandboxError> {
        let response = self
            .client
            .get(self.url(&format!("/computers/{}/files", computer.id)))
            .headers(self.headers(context))
            .query(&[("path", path)])
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let Some(items) = body.as_array() else {
            return Ok(Vec::new());
        };
        Ok(items
            .iter()
            .filter_map(|item| {
                Some(FileEntry {
                    path: item.get("path")?.as_str()?.to_string(),
                    kind: item.get("kind")?.as_str()?.to_string(),
                    size: item.get("size")?.as_u64().unwrap_or(0),
                })
            })
            .collect())
    }

    async fn read_file(
        &self,
        computer: &ComputerRef,
        path: &str,
        context: &AdapterContext,
    ) -> Result<Vec<u8>, SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/read", computer.id)))
            .headers(self.headers(context))
            .json(&serde_json::json!({ "path": path }))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        let body: Value = response
            .json()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        Ok(body
            .get("content")
            .and_then(Value::as_str)
            .unwrap_or_default()
            .as_bytes()
            .to_vec())
    }

    async fn write_file(
        &self,
        computer: &ComputerRef,
        path: &str,
        content: &[u8],
        context: &AdapterContext,
    ) -> Result<(), SandboxError> {
        let response = self
            .client
            .post(self.url(&format!("/computers/{}/files", computer.id)))
            .headers(self.headers(context))
            .json(&serde_json::json!({
                "path": path,
                "content": String::from_utf8_lossy(content),
            }))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        if response.status().is_success() {
            Ok(())
        } else {
            Err(SandboxError::message(format!(
                "write failed: {}",
                response.status()
            )))
        }
    }

    async fn suspend(
        &self,
        computer: &ComputerRef,
        context: &AdapterContext,
    ) -> Result<(), SandboxError> {
        self.post_lifecycle(computer, context, "pause").await
    }

    async fn resume(
        &self,
        request: ProvisionRequest,
        context: &AdapterContext,
    ) -> Result<ComputerRef, SandboxError> {
        if let Some(provider_ref) = request.provider_ref.clone() {
            let computer = ComputerRef {
                id: provider_ref.clone(),
                home_key: request.home_key.clone(),
                kind: SandboxKind::Docker,
                provider_ref,
                fresh: false,
            };
            if self
                .post_lifecycle(&computer, context, "unpause")
                .await
                .is_ok()
            {
                return Ok(computer);
            }
        }
        self.provision(request, context).await
    }

    async fn stop(
        &self,
        computer: &ComputerRef,
        context: &AdapterContext,
    ) -> Result<(), SandboxError> {
        self.post_lifecycle(computer, context, "stop").await
    }

    async fn destroy(
        &self,
        computer: &ComputerRef,
        context: &AdapterContext,
    ) -> Result<(), SandboxError> {
        self.client
            .delete(self.url(&format!("/computers/{}", computer.id)))
            .headers(self.headers(context))
            .send()
            .await
            .map_err(|error| SandboxError::message(error.to_string()))?;
        Ok(())
    }
}

fn decode_observation(body: &Value) -> Result<ComputerObservation, SandboxError> {
    let encoded = body
        .get("png_base64")
        .and_then(Value::as_str)
        .ok_or_else(|| SandboxError::message("missing png"))?;
    let png = base64::engine::general_purpose::STANDARD
        .decode(encoded)
        .map_err(|error| SandboxError::message(error.to_string()))?;
    let cursor = body.get("cursor").and_then(|cursor| {
        Some(CursorPosition {
            x: cursor.get("x")?.as_i64()? as i32,
            y: cursor.get("y")?.as_i64()? as i32,
        })
    });
    let window = body
        .get("activeWindow")
        .or_else(|| body.get("active_window"))
        .and_then(|window| {
            let id = window.get("id")?.as_str()?.to_string();
            if id.is_empty() {
                return None;
            }
            Some(ActiveWindow {
                id,
                title: window
                    .get("title")
                    .and_then(Value::as_str)
                    .filter(|title| !title.is_empty())
                    .map(str::to_string),
            })
        });
    let elements = body
        .get("elements")
        .map(|value| parse_ui_elements(&value.to_string()))
        .unwrap_or_default();
    // The control endpoint does not put the mode on the wire, so the frame
    // itself decides: a desktop that is not 1280x800 must not be reported to
    // the model as 1280x800.
    let (width, height) =
        image_dimensions(&png).unwrap_or((DEFAULT_SCREEN_WIDTH, DEFAULT_SCREEN_HEIGHT));
    let mut observation = observation_with_elements(
        observation_from_png(png, width, height, cursor, window),
        elements,
    );
    observation.native_observation_complete = body
        .get("native_observation_complete")
        .and_then(Value::as_bool)
        .unwrap_or(false);
    Ok(observation)
}

#[cfg(test)]
mod tests {
    use super::decode_observation;
    use lazyboy_contracts::{DEFAULT_SCREEN_HEIGHT, DEFAULT_SCREEN_WIDTH};
    use serde_json::json;

    /// A 3x2 solid-red PNG: small enough to inline, real enough to decode.
    const FRAME_3X2: &str = "iVBORw0KGgoAAAANSUhEUgAAAAMAAAACCAIAAAASFvFNAAAAEElEQVR4nGP4z8AAQQxwFgBB\
0gX7h/C5SAAAAABJRU5ErkJggg==";

    #[test]
    fn the_mode_is_read_from_the_frame() {
        let observation = decode_observation(&json!({ "png_base64": FRAME_3X2 }))
            .expect("a decodable frame makes an observation");
        assert_eq!(observation.width, 3);
        assert_eq!(observation.height, 2);
    }

    #[test]
    fn an_undecodable_frame_falls_back_to_the_default_mode() {
        let observation =
            decode_observation(&json!({ "png_base64": "AAAAAAAAAAA=" })).expect("base64 decodes");
        assert_eq!(observation.width, DEFAULT_SCREEN_WIDTH);
        assert_eq!(observation.height, DEFAULT_SCREEN_HEIGHT);
    }
}