From 6d96bcd71aa115da776b12af22e89c4b78c1ca5f Mon Sep 17 00:00:00 2001 From: junhsss Date: Fri, 11 Sep 2026 00:42:07 +0900 Subject: [PATCH 1/3] feat(computer): add checkpoint commands --- docs/cli-reference.md | 96 ++++++++++++++++ docs/references/steel-cli.md | 6 + src/api/checkpoints.rs | 149 ++++++++++++++++++++++++ src/api/mod.rs | 1 + src/commands/checkpoint.rs | 215 +++++++++++++++++++++++++++++++++++ src/commands/computer/mod.rs | 57 +++++++++- src/commands/mod.rs | 9 ++ tests/cli-spec.json | 10 ++ tests/computer_checkpoint.rs | 194 +++++++++++++++++++++++++++++++ 9 files changed, 731 insertions(+), 6 deletions(-) create mode 100644 src/api/checkpoints.rs create mode 100644 src/commands/checkpoint.rs create mode 100644 tests/computer_checkpoint.rs diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 7130505..fc5c846 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -101,6 +101,12 @@ Steel CLI - browser automation for AI agents. This file is generated from `steel - [steel computer use](#steel-computer-use) - [steel computer exec](#steel-computer-exec) - [steel computer ssh](#steel-computer-ssh) +- [steel computer checkpoint](#steel-computer-checkpoint) +- [steel checkpoint](#steel-checkpoint) +- [steel checkpoint list](#steel-checkpoint-list) +- [steel checkpoint get](#steel-checkpoint-get) +- [steel checkpoint delete](#steel-checkpoint-delete) +- [steel checkpoint restore](#steel-checkpoint-restore) - [steel init](#steel-init) - [steel login](#steel-login) - [steel logout](#steel-logout) @@ -1550,6 +1556,7 @@ steel computer - `use`: Remember a computer as the default for other commands - `exec`: Run one command in a computer - `ssh`: Open an SSH session to a computer +- `checkpoint`: Save a computer as a checkpoint ## steel computer create @@ -1700,6 +1707,95 @@ steel computer ssh Example: `steel computer ssh ` Streaming: websocket `/v1/computers/{id}/ssh` +## steel computer checkpoint + +Save a computer as a checkpoint + +### Usage + +```bash +steel computer checkpoint +``` + +### Parameters + +- `computer_id` (string, optional): Computer ID (defaults to STEEL_COMPUTER_ID or `steel computer use`) +- `--name` (string, optional): Name for the checkpoint +- `--wait` (boolean, optional): Wait until the checkpoint is ready + +## steel checkpoint + +Computer checkpoints: list, restore, delete + +### Usage + +```bash +steel checkpoint +``` + +### Subcommands + +- `list`: List checkpoints +- `get`: Get one checkpoint +- `delete`: Delete a checkpoint +- `restore`: Start a new computer from a checkpoint + +## steel checkpoint list + +List checkpoints + +### Usage + +```bash +steel checkpoint list +``` + +## steel checkpoint get + +Get one checkpoint + +### Usage + +```bash +steel checkpoint get +``` + +### Parameters + +- `checkpoint_id` (string, required): Checkpoint ID + +## steel checkpoint delete + +Delete a checkpoint + +### Usage + +```bash +steel checkpoint delete +``` + +### Parameters + +- `checkpoint_id` (string, required): Checkpoint ID + +## steel checkpoint restore + +Start a new computer from a checkpoint + +### Usage + +```bash +steel checkpoint restore +``` + +### Parameters + +- `checkpoint_id` (string, required): Checkpoint ID +- `--timeout` (string, optional): Stop the computer after this many seconds of running time +- `--auto-pause` (boolean, optional): Pause instead of stopping when the timeout is reached +- `--wait` (boolean, optional): Wait until the computer is running +- `--use` (boolean, optional): Make the new computer the default for other commands + ## steel init One-command onboarding: login + verify + install agent skills diff --git a/docs/references/steel-cli.md b/docs/references/steel-cli.md index e834517..3c34f27 100644 --- a/docs/references/steel-cli.md +++ b/docs/references/steel-cli.md @@ -34,6 +34,12 @@ For generated flags and argument schemas, use [../cli-reference.md](../cli-refer - `steel computer use`: remember a default computer so other commands need no id. - `steel computer exec -- `: run one command; output streams and the exit code is returned. - `steel computer ssh`: open an SSH shell, or run a command over SSH with `-- `. +- `steel computer checkpoint`: save a computer's disk and memory as a checkpoint. + +### Checkpoint Commands + +- `steel checkpoint list`, `steel checkpoint get`, `steel checkpoint delete`: inspect and remove checkpoints. +- `steel checkpoint restore`: start a new computer from a ready checkpoint. ### Credentials Commands diff --git a/src/api/checkpoints.rs b/src/api/checkpoints.rs new file mode 100644 index 0000000..526423f --- /dev/null +++ b/src/api/checkpoints.rs @@ -0,0 +1,149 @@ +use serde_json::{Value, json}; + +use crate::api::client::{ApiError, SteelClient}; +use crate::api::computers::computer_path; +use crate::config::auth::Auth; +use crate::config::settings::ApiMode; + +pub fn checkpoint_path(id: &str) -> String { + format!("/checkpoints/{}", urlencoding::encode(id)) +} + +#[derive(Debug, Default, Clone)] +pub struct RestoreCheckpoint { + pub timeout_seconds: Option, + pub auto_pause: Option, +} + +impl RestoreCheckpoint { + pub fn body(&self) -> Value { + let mut body = json!({}); + if let Some(timeout) = self.timeout_seconds { + body["timeoutSeconds"] = json!(timeout); + } + if let Some(auto_pause) = self.auto_pause { + body["autoPause"] = json!(auto_pause); + } + body + } +} + +impl SteelClient { + pub async fn create_checkpoint( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + computer_id: &str, + name: Option<&str>, + ) -> Result { + let mut body = json!({}); + if let Some(name) = name { + body["name"] = json!(name); + } + self.request( + base_url, + mode, + reqwest::Method::POST, + &format!("{}/checkpoints", computer_path(computer_id)), + Some(body), + auth, + ) + .await + } + + pub async fn list_checkpoints( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::GET, + "/checkpoints", + None, + auth, + ) + .await + } + + pub async fn get_checkpoint( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + id: &str, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::GET, + &checkpoint_path(id), + None, + auth, + ) + .await + } + + pub async fn delete_checkpoint( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + id: &str, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::DELETE, + &checkpoint_path(id), + None, + auth, + ) + .await + } + + pub async fn restore_checkpoint( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + id: &str, + request: &RestoreCheckpoint, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::POST, + &format!("{}/computers", checkpoint_path(id)), + Some(request.body()), + auth, + ) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn paths_are_encoded() { + assert_eq!(checkpoint_path("ckpt a"), "/checkpoints/ckpt%20a"); + } + + #[test] + fn restore_body_only_carries_given_options() { + assert_eq!(RestoreCheckpoint::default().body(), json!({})); + let request = RestoreCheckpoint { + timeout_seconds: Some(120), + auto_pause: Some(true), + }; + assert_eq!( + request.body(), + json!({ "timeoutSeconds": 120, "autoPause": true }) + ); + } +} diff --git a/src/api/mod.rs b/src/api/mod.rs index 4126f12..8c7a11b 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,3 +1,4 @@ +pub mod checkpoints; pub mod client; pub mod computers; pub mod generated; diff --git a/src/commands/checkpoint.rs b/src/commands/checkpoint.rs new file mode 100644 index 0000000..4deb700 --- /dev/null +++ b/src/commands/checkpoint.rs @@ -0,0 +1,215 @@ +use anyhow::{Result, bail}; +use clap::{Parser, Subcommand}; +use serde_json::Value; + +use crate::api::checkpoints::RestoreCheckpoint; +use crate::api::client::SteelClient; +use crate::commands::computer; +use crate::status; +use crate::util::{api, output}; + +#[derive(Subcommand)] +pub enum Command { + /// List checkpoints + List, + + /// Get one checkpoint + Get(IdArgs), + + /// Delete a checkpoint + Delete(IdArgs), + + /// Start a new computer from a checkpoint + Restore(RestoreArgs), +} + +impl Command { + pub const fn telemetry_name(&self) -> &'static str { + match self { + Self::List => "list", + Self::Get(_) => "get", + Self::Delete(_) => "delete", + Self::Restore(_) => "restore", + } + } +} + +#[derive(Parser)] +pub struct IdArgs { + /// Checkpoint ID + pub checkpoint_id: String, +} + +#[derive(Parser)] +pub struct RestoreArgs { + /// Checkpoint ID + pub checkpoint_id: String, + + /// Stop the computer after this many seconds of running time + #[arg(long = "timeout", value_name = "SECONDS")] + pub timeout_seconds: Option, + + /// Pause instead of stopping when the timeout is reached + #[arg(long = "auto-pause")] + pub auto_pause: bool, + + /// Wait until the computer is running + #[arg(long)] + pub wait: bool, + + /// Make the new computer the default for other commands + #[arg(long = "use")] + pub use_as_default: bool, +} + +pub async fn run(command: Command) -> Result<()> { + match command { + Command::List => run_list().await, + Command::Get(args) => run_get(args).await, + Command::Delete(args) => run_delete(args).await, + Command::Restore(args) => run_restore(args).await, + } +} + +pub fn status_of(checkpoint: &Value) -> &str { + checkpoint["status"].as_str().unwrap_or("unknown") +} + +async fn run_list() -> Result<()> { + let (mode, base_url, auth) = api::resolve_with_auth(); + let client = SteelClient::new()?; + let data = client.list_checkpoints(&base_url, mode, &auth).await?; + if output::is_json() { + output::success_data(data); + } else { + print_checkpoints(&data); + } + Ok(()) +} + +async fn run_get(args: IdArgs) -> Result<()> { + let (mode, base_url, auth) = api::resolve_with_auth(); + let client = SteelClient::new()?; + let data = client + .get_checkpoint(&base_url, mode, &auth, &args.checkpoint_id) + .await?; + output::success_data(data); + Ok(()) +} + +async fn run_delete(args: IdArgs) -> Result<()> { + let (mode, base_url, auth) = api::resolve_with_auth(); + let client = SteelClient::new()?; + let data = client + .delete_checkpoint(&base_url, mode, &auth, &args.checkpoint_id) + .await?; + if output::is_json() { + output::success_data(data); + } else { + println!("Deleted {}.", args.checkpoint_id); + } + Ok(()) +} + +async fn run_restore(args: RestoreArgs) -> Result<()> { + let (mode, base_url, auth) = api::resolve_with_auth(); + let client = SteelClient::new()?; + let request = RestoreCheckpoint { + timeout_seconds: args.timeout_seconds, + auto_pause: args.auto_pause.then_some(true), + }; + let restored = client + .restore_checkpoint(&base_url, mode, &auth, &args.checkpoint_id, &request) + .await?; + let id = computer::id_of(&restored)?; + let computer = if args.wait { + status!("Restored {id}, waiting until it is running."); + computer::wait_for(&client, &base_url, mode, &auth, &id, "running").await? + } else { + restored + }; + if args.use_as_default { + computer::remember_computer(Some(&id))?; + } + if output::is_json() { + output::success_data(computer); + } else { + println!("{id} is {}.", computer::status_of(&computer)); + if args.use_as_default { + println!("It is now the default computer."); + } + } + Ok(()) +} + +pub async fn wait_until_ready( + client: &SteelClient, + base_url: &str, + mode: crate::config::settings::ApiMode, + auth: &crate::config::auth::Auth, + id: &str, +) -> Result { + let started = std::time::Instant::now(); + loop { + let checkpoint = client.get_checkpoint(base_url, mode, auth, id).await?; + match status_of(&checkpoint) { + "ready" => return Ok(checkpoint), + "creating" => {} + status => bail!("{id} is {status}, it will not become ready."), + } + if started.elapsed() > computer::WAIT_TIMEOUT { + bail!( + "{id} is still creating after {}s.", + computer::WAIT_TIMEOUT.as_secs() + ); + } + tokio::time::sleep(computer::WAIT_POLL_INTERVAL).await; + } +} + +fn print_checkpoints(data: &Value) { + let checkpoints = data["checkpoints"].as_array().cloned().unwrap_or_default(); + if checkpoints.is_empty() { + println!("No checkpoints."); + return; + } + let rows: Vec<[String; 5]> = checkpoints + .iter() + .map(|checkpoint| { + [ + checkpoint["id"].as_str().unwrap_or("").to_string(), + status_of(checkpoint).to_string(), + checkpoint["name"].as_str().unwrap_or("-").to_string(), + checkpoint["computerId"].as_str().unwrap_or("-").to_string(), + match checkpoint["sizeBytes"].as_u64() { + Some(bytes) => format!("{:.1}G", bytes as f64 / 1_073_741_824.0), + None => "-".to_string(), + }, + ] + }) + .collect(); + let header = ["ID", "STATUS", "NAME", "COMPUTER", "SIZE"]; + let widths: Vec = (0..header.len()) + .map(|column| { + rows.iter() + .map(|row| row[column].len()) + .chain(std::iter::once(header[column].len())) + .max() + .unwrap_or(0) + }) + .collect(); + let line = |cells: [&str; 5]| { + cells + .iter() + .enumerate() + .map(|(column, cell)| format!("{cell:>() + .join(" ") + .trim_end() + .to_string() + }; + println!("{}", line(header)); + for row in &rows { + println!("{}", line([&row[0], &row[1], &row[2], &row[3], &row[4]])); + } +} diff --git a/src/commands/computer/mod.rs b/src/commands/computer/mod.rs index 3641892..eeb2db8 100644 --- a/src/commands/computer/mod.rs +++ b/src/commands/computer/mod.rs @@ -10,12 +10,13 @@ use serde_json::Value; use crate::api::client::SteelClient; use crate::api::computers::CreateComputer; +use crate::commands::checkpoint; use crate::config::settings::{self, ComputerConfig}; use crate::status; use crate::util::{api, output}; -const WAIT_POLL_INTERVAL: Duration = Duration::from_secs(1); -const WAIT_TIMEOUT: Duration = Duration::from_secs(180); +pub const WAIT_POLL_INTERVAL: Duration = Duration::from_secs(1); +pub const WAIT_TIMEOUT: Duration = Duration::from_secs(180); #[derive(Subcommand)] pub enum Command { @@ -45,6 +46,9 @@ pub enum Command { /// Open an SSH session to a computer Ssh(ssh::Args), + + /// Save a computer as a checkpoint + Checkpoint(CheckpointArgs), } impl Command { @@ -59,6 +63,7 @@ impl Command { Self::Use(_) => "use", Self::Exec(_) => "exec", Self::Ssh(_) => "ssh", + Self::Checkpoint(_) => "checkpoint", } } } @@ -118,6 +123,20 @@ pub struct ResumeArgs { pub wait: bool, } +#[derive(Parser)] +pub struct CheckpointArgs { + /// Computer ID (defaults to STEEL_COMPUTER_ID or `steel computer use`) + pub computer_id: Option, + + /// Name for the checkpoint + #[arg(long)] + pub name: Option, + + /// Wait until the checkpoint is ready + #[arg(long)] + pub wait: bool, +} + #[derive(Parser)] pub struct UseArgs { /// Computer ID to remember @@ -139,6 +158,7 @@ pub async fn run(command: Command) -> Result<()> { Command::Use(args) => run_use(args), Command::Exec(args) => exec::run(args).await, Command::Ssh(args) => ssh::run(args).await, + Command::Checkpoint(args) => run_checkpoint(args).await, } } @@ -163,7 +183,7 @@ pub fn choose_computer_id( bail!("No computer given. Pass an id, set STEEL_COMPUTER_ID, or run `steel computer use `.") } -fn remember_computer(id: Option<&str>) -> Result<()> { +pub fn remember_computer(id: Option<&str>) -> Result<()> { let mut config = settings::read_config().unwrap_or_default(); config.computer = id.map(|id| ComputerConfig { default_id: Some(id.to_string()), @@ -171,11 +191,11 @@ fn remember_computer(id: Option<&str>) -> Result<()> { settings::write_config(&config).context("Failed to save the default computer") } -fn status_of(computer: &Value) -> &str { +pub fn status_of(computer: &Value) -> &str { computer["status"].as_str().unwrap_or("unknown") } -fn id_of(computer: &Value) -> Result { +pub fn id_of(computer: &Value) -> Result { computer["id"] .as_str() .map(str::to_string) @@ -290,6 +310,31 @@ async fn run_resume(args: ResumeArgs) -> Result<()> { Ok(()) } +async fn run_checkpoint(args: CheckpointArgs) -> Result<()> { + let id = resolve_computer_id(args.computer_id.as_deref())?; + let (mode, base_url, auth) = api::resolve_with_auth(); + let client = SteelClient::new()?; + let created = client + .create_checkpoint(&base_url, mode, &auth, &id, args.name.as_deref()) + .await?; + let checkpoint_id = created["id"] + .as_str() + .map(str::to_string) + .context("the API returned a checkpoint without an id")?; + let checkpoint = if args.wait { + status!("Saving {id} as {checkpoint_id}, this takes a while."); + checkpoint::wait_until_ready(&client, &base_url, mode, &auth, &checkpoint_id).await? + } else { + created + }; + if output::is_json() { + output::success_data(checkpoint); + } else { + println!("{checkpoint_id} is {}.", checkpoint::status_of(&checkpoint)); + } + Ok(()) +} + fn run_use(args: UseArgs) -> Result<()> { if args.clear { remember_computer(None)?; @@ -317,7 +362,7 @@ fn run_use(args: UseArgs) -> Result<()> { Ok(()) } -async fn wait_for( +pub async fn wait_for( client: &SteelClient, base_url: &str, mode: crate::config::settings::ApiMode, diff --git a/src/commands/mod.rs b/src/commands/mod.rs index f0dfcc0..f450a35 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -1,5 +1,6 @@ pub mod browser; pub mod cache; +pub mod checkpoint; pub mod completion; pub mod computer; pub mod config; @@ -297,6 +298,12 @@ pub enum Command { command: computer::Command, }, + /// Computer checkpoints: list, restore, delete + Checkpoint { + #[command(subcommand)] + command: checkpoint::Command, + }, + /// One-command onboarding: login + verify + install agent skills Init(init::Args), @@ -359,6 +366,7 @@ fn telemetry_command_path(command: &Command) -> Option { Command::Browser(args) => Some(format!("browser.{}", args.command.telemetry_name())), Command::Sessions { command } => Some(format!("sessions.{}", command.telemetry_name())), Command::Computer { command } => Some(format!("computer.{}", command.telemetry_name())), + Command::Checkpoint { command } => Some(format!("checkpoint.{}", command.telemetry_name())), Command::Init(_) => Some("init".to_string()), Command::Login(_) => Some("login".to_string()), Command::Logout(_) => Some("logout".to_string()), @@ -407,6 +415,7 @@ pub async fn run(cli: Cli) -> anyhow::Result<()> { Command::Browser(args) => browser::run(args).await, Command::Sessions { command } => sessions::run(command).await, Command::Computer { command } => computer::run(command).await, + Command::Checkpoint { command } => checkpoint::run(command).await, Command::Init(args) => init::run(args).await, Command::Login(args) => login::run(args).await, Command::Logout(args) => logout::run(args).await, diff --git a/tests/cli-spec.json b/tests/cli-spec.json index 4fe647e..b2e50ca 100644 --- a/tests/cli-spec.json +++ b/tests/cli-spec.json @@ -18,9 +18,19 @@ ] }, { "name": "cache" }, + { + "name": "checkpoint", + "subcommands": [ + { "name": "delete" }, + { "name": "get" }, + { "name": "list" }, + { "name": "restore" } + ] + }, { "name": "computer", "subcommands": [ + { "name": "checkpoint" }, { "name": "create" }, { "name": "delete" }, { "name": "exec" }, diff --git a/tests/computer_checkpoint.rs b/tests/computer_checkpoint.rs new file mode 100644 index 0000000..9b0f25a --- /dev/null +++ b/tests/computer_checkpoint.rs @@ -0,0 +1,194 @@ +//! End-to-end tests for the checkpoint commands against a fake API host. +//! +//! The real `steel` binary runs against a wiremock server, so the request +//! shapes, the waiting loops and the printed tables are exercised as a user +//! sees them. + +use std::process::{Command, Output}; + +use serde_json::json; +use wiremock::matchers::{body_partial_json, header, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const COMPUTER: &str = "cmp_0000123456789abcdefghjkmnpqrs"; +const CHECKPOINT: &str = "ckpt_0000123456789abcdefghjkmn"; + +async fn run_steel(server: &MockServer, args: &[&str]) -> Output { + let tmp = tempfile::tempdir().expect("temp dir"); + let mut cmd = Command::new(env!("CARGO_BIN_EXE_steel")); + cmd.env("STEEL_CONFIG_DIR", tmp.path()); + cmd.env("STEEL_API_URL", format!("{}/v1", server.uri())); + cmd.env("STEEL_API_KEY", "ste-test-key"); + cmd.env("STEEL_TELEMETRY_DISABLED", "1"); + cmd.env("STEEL_FORCE_TTY", "1"); + cmd.env_remove("STEEL_COMPUTER_ID"); + cmd.arg("--no-update-check"); + cmd.args(args); + tokio::task::spawn_blocking(move || { + let output = cmd.output().expect("failed to execute steel binary"); + drop(tmp); + output + }) + .await + .expect("steel process") +} + +fn stdout(output: &Output) -> String { + String::from_utf8_lossy(&output.stdout).to_string() +} + +fn checkpoint(status: &str) -> serde_json::Value { + json!({ + "id": CHECKPOINT, + "computerId": COMPUTER, + "name": "ready", + "status": status, + "sizeBytes": 2_147_483_648u64, + "statusChangedAt": "2026-09-11T00:00:00Z" + }) +} + +#[tokio::test(flavor = "multi_thread")] +async fn checkpoint_create_waits_until_the_snapshot_is_ready() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/v1/computers/{COMPUTER}/checkpoints"))) + .and(header("steel-api-key", "ste-test-key")) + .and(body_partial_json(json!({ "name": "ready" }))) + .respond_with(ResponseTemplate::new(201).set_body_json(checkpoint("creating"))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("/v1/checkpoints/{CHECKPOINT}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(checkpoint("creating"))) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("/v1/checkpoints/{CHECKPOINT}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(checkpoint("ready"))) + .expect(1) + .mount(&server) + .await; + + let output = run_steel( + &server, + &[ + "computer", + "checkpoint", + COMPUTER, + "--name", + "ready", + "--wait", + ], + ) + .await; + + assert!(output.status.success()); + assert!(stdout(&output).contains(&format!("{CHECKPOINT} is ready."))); +} + +#[tokio::test(flavor = "multi_thread")] +async fn checkpoint_create_stops_when_the_snapshot_fails() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/v1/computers/{COMPUTER}/checkpoints"))) + .respond_with(ResponseTemplate::new(201).set_body_json(checkpoint("creating"))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("/v1/checkpoints/{CHECKPOINT}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(checkpoint("failed"))) + .expect(1) + .mount(&server) + .await; + + let output = run_steel(&server, &["computer", "checkpoint", COMPUTER, "--wait"]).await; + + assert!(!output.status.success()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn checkpoint_list_prints_a_table() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/checkpoints")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({ "checkpoints": [checkpoint("ready")] })), + ) + .expect(1) + .mount(&server) + .await; + + let output = run_steel(&server, &["checkpoint", "list"]).await; + + assert!(output.status.success()); + let text = stdout(&output); + assert!(text.contains("ID")); + assert!(text.contains(CHECKPOINT)); + assert!(text.contains("ready")); + assert!(text.contains("2.0G")); +} + +#[tokio::test(flavor = "multi_thread")] +async fn checkpoint_restore_starts_a_computer_and_can_wait_for_it() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/v1/checkpoints/{CHECKPOINT}/computers"))) + .and(body_partial_json(json!({ "timeoutSeconds": 600 }))) + .respond_with(ResponseTemplate::new(201).set_body_json(json!({ + "id": COMPUTER, + "status": "creating" + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("/v1/computers/{COMPUTER}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": COMPUTER, + "status": "running" + }))) + .expect(1) + .mount(&server) + .await; + + let output = run_steel( + &server, + &[ + "checkpoint", + "restore", + CHECKPOINT, + "--timeout", + "600", + "--wait", + "--use", + ], + ) + .await; + + assert!(output.status.success()); + let text = stdout(&output); + assert!(text.contains(&format!("{COMPUTER} is running."))); + assert!(text.contains("default computer")); +} + +#[tokio::test(flavor = "multi_thread")] +async fn checkpoint_delete_reports_the_id() { + let server = MockServer::start().await; + Mock::given(method("DELETE")) + .and(path(format!("/v1/checkpoints/{CHECKPOINT}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(checkpoint("deleting"))) + .expect(1) + .mount(&server) + .await; + + let output = run_steel(&server, &["checkpoint", "delete", CHECKPOINT]).await; + + assert!(output.status.success()); + assert!(stdout(&output).contains(&format!("Deleted {CHECKPOINT}."))); +} From 0cf8d70e597872a78e865a219d96dd11b24e32cc Mon Sep 17 00:00:00 2001 From: junhsss Date: Fri, 11 Sep 2026 00:44:15 +0900 Subject: [PATCH 2/3] feat(computer): add a quota command --- docs/cli-reference.md | 12 +++++++++ docs/references/steel-cli.md | 1 + src/api/computers.rs | 17 ++++++++++++ src/commands/computer/mod.rs | 32 ++++++++++++++++++++++ tests/cli-spec.json | 1 + tests/computer_quota.rs | 52 ++++++++++++++++++++++++++++++++++++ 6 files changed, 115 insertions(+) create mode 100644 tests/computer_quota.rs diff --git a/docs/cli-reference.md b/docs/cli-reference.md index fc5c846..e315107 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -102,6 +102,7 @@ Steel CLI - browser automation for AI agents. This file is generated from `steel - [steel computer exec](#steel-computer-exec) - [steel computer ssh](#steel-computer-ssh) - [steel computer checkpoint](#steel-computer-checkpoint) +- [steel computer quota](#steel-computer-quota) - [steel checkpoint](#steel-checkpoint) - [steel checkpoint list](#steel-checkpoint-list) - [steel checkpoint get](#steel-checkpoint-get) @@ -1557,6 +1558,7 @@ steel computer - `exec`: Run one command in a computer - `ssh`: Open an SSH session to a computer - `checkpoint`: Save a computer as a checkpoint +- `quota`: Show how many computers and checkpoints you can have ## steel computer create @@ -1723,6 +1725,16 @@ steel computer checkpoint - `--name` (string, optional): Name for the checkpoint - `--wait` (boolean, optional): Wait until the checkpoint is ready +## steel computer quota + +Show how many computers and checkpoints you can have + +### Usage + +```bash +steel computer quota +``` + ## steel checkpoint Computer checkpoints: list, restore, delete diff --git a/docs/references/steel-cli.md b/docs/references/steel-cli.md index 3c34f27..d41d89e 100644 --- a/docs/references/steel-cli.md +++ b/docs/references/steel-cli.md @@ -35,6 +35,7 @@ For generated flags and argument schemas, use [../cli-reference.md](../cli-refer - `steel computer exec -- `: run one command; output streams and the exit code is returned. - `steel computer ssh`: open an SSH shell, or run a command over SSH with `-- `. - `steel computer checkpoint`: save a computer's disk and memory as a checkpoint. +- `steel computer quota`: show how many computers and checkpoints the account can have. ### Checkpoint Commands diff --git a/src/api/computers.rs b/src/api/computers.rs index e365a86..240a316 100644 --- a/src/api/computers.rs +++ b/src/api/computers.rs @@ -62,6 +62,23 @@ impl SteelClient { .await } + pub async fn get_computer_quota( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::GET, + "/computers/quota", + None, + auth, + ) + .await + } + pub async fn get_computer( &self, base_url: &str, diff --git a/src/commands/computer/mod.rs b/src/commands/computer/mod.rs index eeb2db8..2e0e8f7 100644 --- a/src/commands/computer/mod.rs +++ b/src/commands/computer/mod.rs @@ -49,6 +49,9 @@ pub enum Command { /// Save a computer as a checkpoint Checkpoint(CheckpointArgs), + + /// Show how many computers and checkpoints you can have + Quota, } impl Command { @@ -64,6 +67,7 @@ impl Command { Self::Exec(_) => "exec", Self::Ssh(_) => "ssh", Self::Checkpoint(_) => "checkpoint", + Self::Quota => "quota", } } } @@ -159,6 +163,7 @@ pub async fn run(command: Command) -> Result<()> { Command::Exec(args) => exec::run(args).await, Command::Ssh(args) => ssh::run(args).await, Command::Checkpoint(args) => run_checkpoint(args).await, + Command::Quota => run_quota().await, } } @@ -310,6 +315,33 @@ async fn run_resume(args: ResumeArgs) -> Result<()> { Ok(()) } +async fn run_quota() -> Result<()> { + let (mode, base_url, auth) = api::resolve_with_auth(); + let client = SteelClient::new()?; + let data = client.get_computer_quota(&base_url, mode, &auth).await?; + if output::is_json() { + output::success_data(data); + } else { + let count = |key: &str| data[key].as_u64().unwrap_or(0); + println!( + "Computers {}/{}", + count("computerCount"), + count("computerLimit") + ); + println!( + "Running {}/{}", + count("runningCount"), + count("runningLimit") + ); + println!( + "Checkpoints {}/{}", + count("checkpointCount"), + count("checkpointLimit") + ); + } + Ok(()) +} + async fn run_checkpoint(args: CheckpointArgs) -> Result<()> { let id = resolve_computer_id(args.computer_id.as_deref())?; let (mode, base_url, auth) = api::resolve_with_auth(); diff --git a/tests/cli-spec.json b/tests/cli-spec.json index b2e50ca..83a0506 100644 --- a/tests/cli-spec.json +++ b/tests/cli-spec.json @@ -37,6 +37,7 @@ { "name": "get" }, { "name": "list" }, { "name": "pause" }, + { "name": "quota" }, { "name": "resume" }, { "name": "ssh" }, { "name": "use" } diff --git a/tests/computer_quota.rs b/tests/computer_quota.rs new file mode 100644 index 0000000..21fd971 --- /dev/null +++ b/tests/computer_quota.rs @@ -0,0 +1,52 @@ +//! End-to-end test for `steel computer quota` against a fake API host. + +use std::process::{Command, Output}; + +use serde_json::json; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +async fn run_steel(server: &MockServer, args: &[&str]) -> Output { + let tmp = tempfile::tempdir().expect("temp dir"); + let mut cmd = Command::new(env!("CARGO_BIN_EXE_steel")); + cmd.env("STEEL_CONFIG_DIR", tmp.path()); + cmd.env("STEEL_API_URL", format!("{}/v1", server.uri())); + cmd.env("STEEL_API_KEY", "ste-test-key"); + cmd.env("STEEL_TELEMETRY_DISABLED", "1"); + cmd.env("STEEL_FORCE_TTY", "1"); + cmd.arg("--no-update-check"); + cmd.args(args); + tokio::task::spawn_blocking(move || { + let output = cmd.output().expect("failed to execute steel binary"); + drop(tmp); + output + }) + .await + .expect("steel process") +} + +#[tokio::test(flavor = "multi_thread")] +async fn quota_prints_usage_against_each_limit() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/computers/quota")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "computerCount": 2, + "computerLimit": 5, + "runningCount": 1, + "runningLimit": 3, + "checkpointCount": 0, + "checkpointLimit": 2 + }))) + .expect(1) + .mount(&server) + .await; + + let output = run_steel(&server, &["computer", "quota"]).await; + + assert!(output.status.success()); + let text = String::from_utf8_lossy(&output.stdout).to_string(); + assert!(text.contains("Computers 2/5")); + assert!(text.contains("Running 1/3")); + assert!(text.contains("Checkpoints 0/2")); +} From 6bea5a74ba4b2c749013a051d095f5262abccf21 Mon Sep 17 00:00:00 2001 From: junhsss Date: Fri, 11 Sep 2026 00:44:15 +0900 Subject: [PATCH 3/3] chore: bump version --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cb7b274..7cc8609 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3762,7 +3762,7 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "steel-cli" -version = "0.5.0-preview.1" +version = "0.5.0-preview.2" dependencies = [ "aes 0.8.4", "aes-gcm 0.10.3", diff --git a/Cargo.toml b/Cargo.toml index 48f8e7b..359fbdd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "steel-cli" -version = "0.5.0-preview.1" +version = "0.5.0-preview.2" edition = "2024" description = "Steel CLI - Browser automation for AI agents" license = "MIT"