diff --git a/crates/pu-cli/src/commands/bench.rs b/crates/pu-cli/src/commands/bench.rs new file mode 100644 index 0000000..24055f8 --- /dev/null +++ b/crates/pu-cli/src/commands/bench.rs @@ -0,0 +1,80 @@ +use crate::client; +use crate::daemon_ctrl; +use crate::error::CliError; +use crate::output; +use pu_core::protocol::{Request, SuspendTarget}; +use std::path::Path; + +/// Suspend one or all agents ("bench" them). +/// +/// When `--all` is used, the invoking agent may be included in the suspend +/// results. A warning is emitted if this happens; use `pu play ` to resume. +pub async fn run_bench( + socket: &Path, + agent: Option, + all: bool, + json: bool, +) -> Result<(), CliError> { + let self_agent_id = std::env::var("PU_AGENT_ID").ok(); + + let target = if all { + SuspendTarget::All + } else if let Some(id) = agent { + // Self-protection: refuse to bench own agent + if let Some(ref self_id) = self_agent_id { + if &id == self_id { + return Err(CliError::Other("cannot bench self".into())); + } + } + SuspendTarget::Agent(id) + } else { + return Err(CliError::Other( + "bench target required — use or --all".into(), + )); + }; + + daemon_ctrl::ensure_daemon(socket).await?; + + let project_root = crate::commands::cwd_string()?; + let resp = client::send_request( + socket, + &Request::Suspend { + project_root, + target, + }, + ) + .await?; + let resp = output::check_response(resp, json)?; + + // Warn if self was included in bulk bench results + if let Some(ref self_id) = self_agent_id { + if let pu_core::protocol::Response::SuspendResult { ref suspended } = resp { + if suspended.contains(self_id) { + eprintln!( + "warning: agent {self_id} benched itself — use `pu play {self_id}` to resume", + ); + } + } + } + + output::print_response(&resp, json); + Ok(()) +} + +/// Resume a previously benched agent, putting it back in play. +pub async fn run_play(socket: &Path, agent_id: &str, json: bool) -> Result<(), CliError> { + daemon_ctrl::ensure_daemon(socket).await?; + + let project_root = crate::commands::cwd_string()?; + let resp = client::send_request( + socket, + &Request::Resume { + project_root, + agent_id: agent_id.to_string(), + }, + ) + .await?; + let resp = output::check_response(resp, json)?; + output::print_response(&resp, json); + Ok(()) +} diff --git a/crates/pu-cli/src/commands/mod.rs b/crates/pu-cli/src/commands/mod.rs index cb25e00..488301c 100644 --- a/crates/pu-cli/src/commands/mod.rs +++ b/crates/pu-cli/src/commands/mod.rs @@ -1,5 +1,6 @@ pub mod agent_def; pub mod attach; +pub mod bench; pub mod clean; pub mod diff; pub mod grid; diff --git a/crates/pu-cli/src/main.rs b/crates/pu-cli/src/main.rs index ee8111d..94f4b27 100644 --- a/crates/pu-cli/src/main.rs +++ b/crates/pu-cli/src/main.rs @@ -66,6 +66,29 @@ enum Commands { #[arg(long)] json: bool, }, + /// Bench (suspend) agents — pull them off the court + #[command(long_about = "Bench (suspend) agents — pull them off the court.\n\n\ + When using --all, the invoking agent may also be suspended.\n\ + Use `pu play ` to resume a benched agent.")] + Bench { + /// Agent ID to bench + #[arg(conflicts_with = "all")] + agent_id: Option, + /// Bench all active agents + #[arg(long)] + all: bool, + /// Output as JSON + #[arg(long)] + json: bool, + }, + /// Put a benched agent back in play (resume) + Play { + /// Agent ID to resume + agent_id: String, + /// Output as JSON + #[arg(long)] + json: bool, + }, /// Kill agents Kill { /// Kill specific agent @@ -500,6 +523,14 @@ async fn main() { ) .await } + Commands::Bench { + agent_id, + all, + json, + } => commands::bench::run_bench(&socket, agent_id, all, json).await, + Commands::Play { agent_id, json } => { + commands::bench::run_play(&socket, &agent_id, json).await + } Commands::Status { agent, json } => commands::status::run(&socket, agent, json).await, Commands::Kill { agent, diff --git a/crates/pu-cli/src/output.rs b/crates/pu-cli/src/output.rs index da59fa8..2eae2cc 100644 --- a/crates/pu-cli/src/output.rs +++ b/crates/pu-cli/src/output.rs @@ -23,7 +23,20 @@ pub fn check_response(resp: Response, json: bool) -> Result } } +/// Return a colored status string for display (delegates with `suspended = false`). fn status_colored(status: AgentStatus, exit_code: Option) -> String { + status_colored_with_suspended(status, exit_code, false) +} + +/// Return a colored status string, showing "benched" (yellow) for suspended alive agents. +fn status_colored_with_suspended( + status: AgentStatus, + exit_code: Option, + suspended: bool, +) -> String { + if suspended && status.is_alive() { + return "benched".yellow().to_string(); + } match status { AgentStatus::Streaming => "streaming".green().to_string(), AgentStatus::Waiting => "waiting".cyan().to_string(), @@ -95,7 +108,7 @@ pub fn print_response(response: &Response, json_mode: bool) { "{:<14} {:<16} {}", a.id.dimmed(), a.name, - status_colored(a.status, a.exit_code) + status_colored_with_suspended(a.status, a.exit_code, a.suspended) ); } } @@ -119,7 +132,7 @@ pub fn print_response(response: &Response, json_mode: bool) { " {:<14} {:<16} {}", a.id.dimmed(), a.name, - status_colored(a.status, a.exit_code), + status_colored_with_suspended(a.status, a.exit_code, a.suspended), ); } } @@ -130,7 +143,7 @@ pub fn print_response(response: &Response, json_mode: bool) { "{} {} {}", a.id.dimmed(), a.name.bold(), - status_colored(a.status, a.exit_code) + status_colored_with_suspended(a.status, a.exit_code, a.suspended) ); if let Some(pid) = a.pid { println!(" PID: {pid}"); @@ -149,11 +162,18 @@ pub fn print_response(response: &Response, json_mode: bool) { println!("Killed {} agent(s)", killed.len()); } Response::SuspendResult { suspended } => { - println!("Suspended {} agent(s)", suspended.len()); + if suspended.is_empty() { + println!("No agents to bench"); + } else { + println!("Benched {} agent(s)", suspended.len()); + for id in suspended { + println!(" {}", id.dimmed()); + } + } } Response::ResumeResult { agent_id, status } => { println!( - "Resumed agent {} ({})", + "Back in play: {} ({})", agent_id.bold(), status_colored(*status, None) ); @@ -701,6 +721,12 @@ mod tests { print_response(&resp, false); } + #[test] + fn given_empty_suspend_result_should_not_panic() { + let resp = Response::SuspendResult { suspended: vec![] }; + print_response(&resp, false); + } + #[test] fn given_resume_result_should_not_panic() { let resp = Response::ResumeResult { @@ -843,6 +869,33 @@ mod tests { assert!(s.contains("waiting")); } + // --- status_colored_with_suspended (bench) --- + + #[test] + fn given_suspended_streaming_should_show_benched() { + let s = status_colored_with_suspended(AgentStatus::Streaming, None, true); + assert!(s.contains("benched")); + } + + #[test] + fn given_suspended_waiting_should_show_benched() { + let s = status_colored_with_suspended(AgentStatus::Waiting, None, true); + assert!(s.contains("benched")); + } + + #[test] + fn given_suspended_broken_should_not_show_benched() { + let s = status_colored_with_suspended(AgentStatus::Broken, Some(0), true); + assert!(!s.contains("benched")); + assert!(s.contains("done")); + } + + #[test] + fn given_not_suspended_should_show_normal_status() { + let s = status_colored_with_suspended(AgentStatus::Streaming, None, false); + assert!(s.contains("streaming")); + } + // --- diff output --- #[test]