From ecdc63d5377317adfd367fe3005daaba2b544910 Mon Sep 17 00:00:00 2001 From: junhsss Date: Sat, 12 Sep 2026 05:50:42 +0900 Subject: [PATCH] feat(computer): add stop, start and restart --- src/api/computers.rs | 54 ++++++++++++++++++++++++ src/commands/computer/mod.rs | 73 ++++++++++++++++++++++++++++++-- tests/computer_lifecycle.rs | 82 ++++++++++++++++++++++++++++++++++++ 3 files changed, 205 insertions(+), 4 deletions(-) create mode 100644 tests/computer_lifecycle.rs diff --git a/src/api/computers.rs b/src/api/computers.rs index a476e0b..7f8d525 100644 --- a/src/api/computers.rs +++ b/src/api/computers.rs @@ -164,6 +164,60 @@ impl SteelClient { .await } + pub async fn stop_computer( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + id: &str, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::POST, + &format!("{}/stop", computer_path(id)), + None, + auth, + ) + .await + } + + pub async fn start_computer( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + id: &str, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::POST, + &format!("{}/start", computer_path(id)), + None, + auth, + ) + .await + } + + pub async fn restart_computer( + &self, + base_url: &str, + mode: ApiMode, + auth: &Auth, + id: &str, + ) -> Result { + self.request( + base_url, + mode, + reqwest::Method::POST, + &format!("{}/restart", computer_path(id)), + None, + auth, + ) + .await + } + pub async fn exec_computer( &self, base_url: &str, diff --git a/src/commands/computer/mod.rs b/src/commands/computer/mod.rs index 8580123..ee7b234 100644 --- a/src/commands/computer/mod.rs +++ b/src/commands/computer/mod.rs @@ -36,7 +36,16 @@ pub enum Command { Pause(IdArgs), /// Resume a paused computer - Resume(ResumeArgs), + Resume(LifecycleArgs), + + /// Stop a computer but keep its disk + Stop(LifecycleArgs), + + /// Start a stopped computer + Start(LifecycleArgs), + + /// Reboot a running computer + Restart(LifecycleArgs), /// Remember a computer as the default for other commands Use(UseArgs), @@ -63,6 +72,9 @@ impl Command { Self::Delete(_) => "delete", Self::Pause(_) => "pause", Self::Resume(_) => "resume", + Self::Stop(_) => "stop", + Self::Start(_) => "start", + Self::Restart(_) => "restart", Self::Use(_) => "use", Self::Exec(_) => "exec", Self::Ssh(_) => "ssh", @@ -110,11 +122,11 @@ pub struct IdArgs { } #[derive(Parser)] -pub struct ResumeArgs { +pub struct LifecycleArgs { /// Computer ID (defaults to STEEL_COMPUTER_ID or `steel computer use`) pub computer_id: Option, - /// Wait until the computer is running + /// Wait until the computer reaches its new state #[arg(long)] pub wait: bool, } @@ -151,6 +163,9 @@ pub async fn run(command: Command) -> Result<()> { Command::Delete(args) => run_delete(args).await, Command::Pause(args) => run_pause(args).await, Command::Resume(args) => run_resume(args).await, + Command::Stop(args) => run_stop(args).await, + Command::Start(args) => run_start(args).await, + Command::Restart(args) => run_restart(args).await, Command::Use(args) => run_use(args), Command::Exec(args) => exec::run(args).await, Command::Ssh(args) => ssh::run(args).await, @@ -287,7 +302,7 @@ async fn run_pause(args: IdArgs) -> Result<()> { Ok(()) } -async fn run_resume(args: ResumeArgs) -> Result<()> { +async fn run_resume(args: LifecycleArgs) -> Result<()> { let id = resolve_computer_id(args.computer_id.as_deref())?; let (mode, base_url, auth) = api::resolve_with_auth(); let client = SteelClient::new()?; @@ -305,6 +320,56 @@ async fn run_resume(args: ResumeArgs) -> Result<()> { Ok(()) } +#[derive(Clone, Copy)] +enum Lifecycle { + Stop, + Start, + Restart, +} + +impl Lifecycle { + const fn settled(self) -> &'static str { + match self { + Self::Stop => "stopped", + Self::Start | Self::Restart => "running", + } + } +} + +async fn run_stop(args: LifecycleArgs) -> Result<()> { + run_lifecycle(args, Lifecycle::Stop).await +} + +async fn run_start(args: LifecycleArgs) -> Result<()> { + run_lifecycle(args, Lifecycle::Start).await +} + +async fn run_restart(args: LifecycleArgs) -> Result<()> { + run_lifecycle(args, Lifecycle::Restart).await +} + +async fn run_lifecycle(args: LifecycleArgs, verb: Lifecycle) -> 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 sent = match verb { + Lifecycle::Stop => client.stop_computer(&base_url, mode, &auth, &id).await?, + Lifecycle::Start => client.start_computer(&base_url, mode, &auth, &id).await?, + Lifecycle::Restart => client.restart_computer(&base_url, mode, &auth, &id).await?, + }; + let data = if args.wait { + wait_for(&client, &base_url, mode, &auth, &id, verb.settled()).await? + } else { + sent + }; + if output::is_json() { + output::success_data(data); + } else { + println!("{id} is {}.", status_of(&data)); + } + Ok(()) +} + async fn run_quota() -> Result<()> { let (mode, base_url, auth) = api::resolve_with_auth(); let client = SteelClient::new()?; diff --git a/tests/computer_lifecycle.rs b/tests/computer_lifecycle.rs new file mode 100644 index 0000000..7d76499 --- /dev/null +++ b/tests/computer_lifecycle.rs @@ -0,0 +1,82 @@ +//! End-to-end tests for `steel computer stop|start|restart` 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") +} + +const ID: &str = "cmp_00x1492gqevd1nvs7vya5py3d0m33"; + +async fn mount(server: &MockServer, verb: &str, status: &str) { + Mock::given(method("POST")) + .and(path(format!("/v1/computers/{ID}/{verb}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": ID, + "status": status + }))) + .expect(1) + .mount(server) + .await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn stop_reports_the_settled_status() { + let server = MockServer::start().await; + mount(&server, "stop", "stopped").await; + + let output = run_steel(&server, &["computer", "stop", ID]).await; + + assert!(output.status.success()); + let text = String::from_utf8_lossy(&output.stdout).to_string(); + assert!(text.contains(&format!("{ID} is stopped."))); +} + +#[tokio::test(flavor = "multi_thread")] +async fn start_reports_the_settled_status() { + let server = MockServer::start().await; + mount(&server, "start", "running").await; + + let output = run_steel(&server, &["computer", "start", ID]).await; + + assert!(output.status.success()); + let text = String::from_utf8_lossy(&output.stdout).to_string(); + assert!(text.contains(&format!("{ID} is running."))); +} + +#[tokio::test(flavor = "multi_thread")] +async fn restart_calls_restart_and_not_start() { + let server = MockServer::start().await; + mount(&server, "restart", "running").await; + Mock::given(method("POST")) + .and(path(format!("/v1/computers/{ID}/start"))) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&server) + .await; + + let output = run_steel(&server, &["computer", "restart", ID]).await; + + assert!(output.status.success()); + let text = String::from_utf8_lossy(&output.stdout).to_string(); + assert!(text.contains(&format!("{ID} is running."))); +}