From fae169bb80c0ee1d5bc1c2255698c0049c3dbc8b Mon Sep 17 00:00:00 2001 From: junhsss Date: Sat, 12 Sep 2026 05:59:32 +0900 Subject: [PATCH] feat(computer): add --idle-timeout and --env --- src/api/checkpoints.rs | 7 +++++- src/api/computers.rs | 23 ++++++++++++++++++ src/commands/checkpoint.rs | 5 ++++ src/commands/computer/mod.rs | 39 ++++++++++++++++++++++++++++++ tests/computer_lifecycle.rs | 46 ++++++++++++++++++++++++++++++++++++ 5 files changed, 119 insertions(+), 1 deletion(-) diff --git a/src/api/checkpoints.rs b/src/api/checkpoints.rs index 526423f..bd1e49a 100644 --- a/src/api/checkpoints.rs +++ b/src/api/checkpoints.rs @@ -13,6 +13,7 @@ pub fn checkpoint_path(id: &str) -> String { pub struct RestoreCheckpoint { pub timeout_seconds: Option, pub auto_pause: Option, + pub idle_timeout_seconds: Option, } impl RestoreCheckpoint { @@ -24,6 +25,9 @@ impl RestoreCheckpoint { if let Some(auto_pause) = self.auto_pause { body["autoPause"] = json!(auto_pause); } + if let Some(idle_timeout) = self.idle_timeout_seconds { + body["idleTimeoutSeconds"] = json!(idle_timeout); + } body } } @@ -140,10 +144,11 @@ mod tests { let request = RestoreCheckpoint { timeout_seconds: Some(120), auto_pause: Some(true), + idle_timeout_seconds: Some(60), }; assert_eq!( request.body(), - json!({ "timeoutSeconds": 120, "autoPause": true }) + json!({ "timeoutSeconds": 120, "autoPause": true, "idleTimeoutSeconds": 60 }) ); } } diff --git a/src/api/computers.rs b/src/api/computers.rs index 7f8d525..7374a4e 100644 --- a/src/api/computers.rs +++ b/src/api/computers.rs @@ -1,3 +1,5 @@ +use std::collections::BTreeMap; + use serde_json::{Value, json}; use crate::api::client::{ApiError, SteelClient}; @@ -15,6 +17,8 @@ pub struct CreateComputer { pub memory_mib: Option, pub timeout_seconds: Option, pub auto_pause: Option, + pub idle_timeout_seconds: Option, + pub env: BTreeMap, } impl CreateComputer { @@ -35,6 +39,12 @@ impl CreateComputer { if let Some(auto_pause) = self.auto_pause { body["autoPause"] = json!(auto_pause); } + if let Some(idle_timeout) = self.idle_timeout_seconds { + body["idleTimeoutSeconds"] = json!(idle_timeout); + } + if !self.env.is_empty() { + body["env"] = json!(self.env); + } body } } @@ -258,6 +268,8 @@ mod tests { memory_mib: Some(4096), timeout_seconds: Some(600), auto_pause: Some(true), + idle_timeout_seconds: Some(90), + env: BTreeMap::from([("TOKEN".into(), "abc".into())]), }; assert_eq!( full.body(), @@ -267,10 +279,21 @@ mod tests { "memoryMib": 4096, "timeoutSeconds": 600, "autoPause": true, + "idleTimeoutSeconds": 90, + "env": { "TOKEN": "abc" }, }) ); } + #[test] + fn zero_idle_timeout_is_sent_and_empty_env_is_not() { + let request = CreateComputer { + idle_timeout_seconds: Some(0), + ..Default::default() + }; + assert_eq!(request.body(), json!({ "idleTimeoutSeconds": 0 })); + } + #[test] fn computer_paths_escape_ids() { assert_eq!( diff --git a/src/commands/checkpoint.rs b/src/commands/checkpoint.rs index 4deb700..4339229 100644 --- a/src/commands/checkpoint.rs +++ b/src/commands/checkpoint.rs @@ -53,6 +53,10 @@ pub struct RestoreArgs { #[arg(long = "auto-pause")] pub auto_pause: bool, + /// Pause after this many seconds without incoming traffic (0 disables it) + #[arg(long = "idle-timeout", value_name = "SECONDS")] + pub idle_timeout_seconds: Option, + /// Wait until the computer is running #[arg(long)] pub wait: bool, @@ -117,6 +121,7 @@ async fn run_restore(args: RestoreArgs) -> Result<()> { let request = RestoreCheckpoint { timeout_seconds: args.timeout_seconds, auto_pause: args.auto_pause.then_some(true), + idle_timeout_seconds: args.idle_timeout_seconds, }; let restored = client .restore_checkpoint(&base_url, mode, &auth, &args.checkpoint_id, &request) diff --git a/src/commands/computer/mod.rs b/src/commands/computer/mod.rs index ee7b234..e948d7a 100644 --- a/src/commands/computer/mod.rs +++ b/src/commands/computer/mod.rs @@ -2,6 +2,7 @@ pub mod exec; pub mod ssh; pub mod wsio; +use std::collections::BTreeMap; use std::time::{Duration, Instant}; use anyhow::{Context, Result, bail}; @@ -106,6 +107,14 @@ pub struct CreateArgs { #[arg(long = "auto-pause")] pub auto_pause: bool, + /// Pause after this many seconds without incoming traffic (0 disables it) + #[arg(long = "idle-timeout", value_name = "SECONDS")] + pub idle_timeout_seconds: Option, + + /// Environment variable for the computer, repeatable + #[arg(long = "env", value_name = "KEY=VALUE")] + pub env: Vec, + /// Wait until the computer is running #[arg(long)] pub wait: bool, @@ -203,6 +212,22 @@ pub fn remember_computer(id: Option<&str>) -> Result<()> { settings::write_config(&config).context("Failed to save the default computer") } +pub fn parse_env(pairs: &[String]) -> Result> { + let mut env = BTreeMap::new(); + for pair in pairs { + let Some((name, value)) = pair.split_once('=') else { + bail!("--env wants KEY=VALUE, got {pair:?}."); + }; + if name.is_empty() { + bail!("--env wants a name before the '=', got {pair:?}."); + } + if env.insert(name.to_string(), value.to_string()).is_some() { + bail!("--env {name} was given twice."); + } + } + Ok(env) +} + pub fn status_of(computer: &Value) -> &str { computer["status"].as_str().unwrap_or("unknown") } @@ -223,6 +248,8 @@ async fn run_create(args: CreateArgs) -> Result<()> { memory_mib: args.memory_mib, timeout_seconds: args.timeout_seconds, auto_pause: args.auto_pause.then_some(true), + idle_timeout_seconds: args.idle_timeout_seconds, + env: parse_env(&args.env)?, }; let created = client .create_computer(&base_url, mode, &auth, &request) @@ -529,6 +556,18 @@ fn print_computers(data: &Value) { mod tests { use super::*; + #[test] + fn parse_env_reads_pairs_and_refuses_bad_ones() { + assert!(parse_env(&[]).unwrap().is_empty()); + let parsed = parse_env(&["A=1".into(), "B=x=y".into(), "C=".into()]).unwrap(); + assert_eq!(parsed["A"], "1"); + assert_eq!(parsed["B"], "x=y"); + assert_eq!(parsed["C"], ""); + assert!(parse_env(&["NOPE".into()]).is_err()); + assert!(parse_env(&["=1".into()]).is_err()); + assert!(parse_env(&["A=1".into(), "A=2".into()]).is_err()); + } + #[test] fn explicit_id_wins_and_blank_ids_are_ignored() { assert_eq!( diff --git a/tests/computer_lifecycle.rs b/tests/computer_lifecycle.rs index 7d76499..57dbf88 100644 --- a/tests/computer_lifecycle.rs +++ b/tests/computer_lifecycle.rs @@ -80,3 +80,49 @@ async fn restart_calls_restart_and_not_start() { let text = String::from_utf8_lossy(&output.stdout).to_string(); assert!(text.contains(&format!("{ID} is running."))); } + +#[tokio::test(flavor = "multi_thread")] +async fn create_sends_the_idle_timeout_and_env() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/computers")) + .and(wiremock::matchers::body_partial_json(json!({ + "idleTimeoutSeconds": 120, + "env": { "A": "1", "TOKEN": "x=y" } + }))) + .respond_with(ResponseTemplate::new(201).set_body_json(json!({ + "id": ID, + "status": "running" + }))) + .expect(1) + .mount(&server) + .await; + + let output = run_steel( + &server, + &[ + "computer", + "create", + "--idle-timeout", + "120", + "--env", + "A=1", + "--env", + "TOKEN=x=y", + ], + ) + .await; + + assert!(output.status.success()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn create_refuses_an_env_without_an_equals() { + let server = MockServer::start().await; + + let output = run_steel(&server, &["computer", "create", "--env", "NOPE"]).await; + + assert!(!output.status.success()); + let text = String::from_utf8_lossy(&output.stderr).to_string(); + assert!(text.contains("KEY=VALUE")); +}