Skip to main content

agent_sim/cli/
mod.rs

1pub mod args;
2pub mod commands;
3mod env;
4pub mod error;
5pub mod output;
6mod recipe;
7
8use crate::cli::args::{CliArgs, CloseArgs, Command, WatchArgs};
9use crate::cli::commands::to_request;
10use crate::cli::error::CliError;
11use crate::config::load_config;
12use crate::connection::{send_env_request, send_request};
13use crate::daemon::lifecycle::{self, bootstrap_daemon};
14use crate::envd::lifecycle as env_lifecycle;
15use crate::load::resolve::resolve_standalone_load_spec;
16use crate::protocol::{
17    EnvAction, InstanceAction, Request, RequestAction, Response, ResponseData, SignalValueData,
18    WatchSampleData,
19};
20use std::path::Path;
21use std::process::ExitCode;
22use tokio::time::{Duration, sleep};
23use uuid::Uuid;
24
25pub async fn run_with_args(args: CliArgs) -> ExitCode {
26    match run_inner(args).await {
27        Ok(code) => code,
28        Err(err) => {
29            eprintln!("{err}");
30            ExitCode::from(1)
31        }
32    }
33}
34
35async fn run_inner(args: CliArgs) -> Result<ExitCode, CliError> {
36    let Some(command) = args.command.as_ref() else {
37        return Err(CliError::MissingCommand);
38    };
39
40    match command {
41        Command::Load(load) => run_load_command(&args, load).await,
42        Command::Watch(watch) => run_watch_command(&args, watch).await,
43        Command::Run(run) => recipe::run_recipe_command(&args, run).await,
44        Command::Close(close) if close.all || close.env.is_some() => run_close_command(close).await,
45        Command::Env(env) => env::run_env_command(&args, env).await,
46        _ => {
47            let request = to_request(&args)?;
48            let response = send_request(&args.instance, &request)
49                .await
50                .map_err(|e| CliError::CommandFailed(e.to_string()))?;
51            output::print_response(&response, args.json);
52            if response.success {
53                Ok(ExitCode::SUCCESS)
54            } else {
55                Ok(ExitCode::from(1))
56            }
57        }
58    }
59}
60
61async fn run_load_command(
62    args: &CliArgs,
63    load: &crate::cli::args::LoadArgs,
64) -> Result<ExitCode, CliError> {
65    let config = load_config(args.config.as_deref().map(Path::new))
66        .map_err(|err| CliError::CommandFailed(err.to_string()))?;
67    let config_base_dir = config.source_path.as_ref().and_then(|path| path.parent());
68    let load_spec = resolve_standalone_load_spec(
69        &config.file,
70        config_base_dir,
71        load.libpath.as_deref(),
72        &load.flash,
73        None,
74    )
75    .map_err(|err| CliError::CommandFailed(err.to_string()))?;
76
77    bootstrap_daemon(&args.instance, &load_spec)
78        .await
79        .map_err(|err| CliError::CommandFailed(err.to_string()))?;
80    let response = send_action(&args.instance, InstanceAction::Info).await?;
81    let ResponseData::ProjectInfo {
82        libpath,
83        signal_count,
84        ..
85    } = response
86        .data
87        .ok_or_else(|| CliError::CommandFailed("missing info response payload".to_string()))?
88    else {
89        return Err(CliError::CommandFailed(
90            "unexpected daemon response after load".to_string(),
91        ));
92    };
93    let response = Response::ok(
94        Uuid::new_v4(),
95        ResponseData::Loaded {
96            libpath,
97            signal_count,
98        },
99    );
100    output::print_response(&response, args.json);
101    Ok(ExitCode::SUCCESS)
102}
103
104async fn run_watch_command(args: &CliArgs, watch: &WatchArgs) -> Result<ExitCode, CliError> {
105    let count = watch.samples.unwrap_or(10).max(1);
106    let mut samples = Vec::with_capacity(count as usize);
107    for idx in 0..count {
108        let (tick, time_us, signal_value) =
109            fetch_signal_sample(&args.instance, &watch.selector).await?;
110        samples.push(WatchSampleData {
111            tick,
112            time_us,
113            signal: signal_value.name,
114            value: signal_value.value,
115        });
116        if idx + 1 < count {
117            sleep(Duration::from_millis(watch.interval_ms.max(1))).await;
118        }
119    }
120
121    let response = Response::ok(Uuid::new_v4(), ResponseData::WatchSamples { samples });
122    output::print_response(&response, args.json);
123    Ok(ExitCode::SUCCESS)
124}
125
126async fn run_close_command(close: &CloseArgs) -> Result<ExitCode, CliError> {
127    if let Some(env_name) = &close.env {
128        close_env_and_wait(env_name).await?;
129        return Ok(ExitCode::SUCCESS);
130    }
131
132    let env_targets = env_lifecycle::list_envs()
133        .await
134        .map_err(|err| CliError::CommandFailed(err.to_string()))?
135        .into_iter()
136        .filter(|(_, _, running)| *running)
137        .map(|(name, _, _)| name)
138        .collect::<Vec<_>>();
139    let mut close_errors = Vec::new();
140    for env_name in env_targets {
141        if let Err(err) = close_env_and_wait(&env_name).await {
142            close_errors.push(format!("env '{env_name}': {err}"));
143        }
144    }
145
146    let instances = lifecycle::list_instances()
147        .await
148        .map_err(|e| CliError::CommandFailed(e.to_string()))?;
149    let mut targets = instances
150        .into_iter()
151        .filter(|instance| instance.running)
152        .map(|instance| instance.name)
153        .collect::<Vec<_>>();
154    targets.sort();
155
156    for session_name in targets {
157        if send_action_success(&session_name, InstanceAction::Close)
158            .await
159            .is_err()
160            && let Some(pid) = lifecycle::read_pid(&session_name)
161        {
162            let _ = lifecycle::kill_pid(pid);
163        }
164    }
165    if !close_errors.is_empty() {
166        return Err(CliError::CommandFailed(format!(
167            "close --all completed with env shutdown errors: {}",
168            close_errors.join("; ")
169        )));
170    }
171    Ok(ExitCode::SUCCESS)
172}
173
174async fn close_env_and_wait(env_name: &str) -> Result<(), CliError> {
175    let request = Request {
176        id: Uuid::new_v4(),
177        action: RequestAction::Env(EnvAction::Close {
178            env: env_name.to_string(),
179        }),
180    };
181    let response = send_env_request(env_name, &request)
182        .await
183        .map_err(|err| CliError::CommandFailed(err.to_string()))?;
184    if !response.success {
185        return Err(CliError::CommandFailed(response_error(&response)));
186    }
187    let env_socket = env_lifecycle::socket_path(env_name);
188    let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
189    loop {
190        let running_instances = lifecycle::list_instances()
191            .await
192            .map_err(|err| CliError::CommandFailed(err.to_string()))?
193            .into_iter()
194            .any(|instance| instance.running && instance.env.as_deref() == Some(env_name));
195        if !env_socket.exists() && !running_instances {
196            return Ok(());
197        }
198        if tokio::time::Instant::now() >= deadline {
199            return Err(CliError::CommandFailed(format!(
200                "timed out waiting for env '{env_name}' to shut down"
201            )));
202        }
203        sleep(Duration::from_millis(100)).await;
204    }
205}
206
207pub(crate) async fn fetch_signal_sample(
208    session: &str,
209    selector: &str,
210) -> Result<(u64, u64, SignalValueData), CliError> {
211    let response = send_action(
212        session,
213        InstanceAction::Sample {
214            selectors: vec![selector.to_string()],
215        },
216    )
217    .await?;
218    if !response.success {
219        return Err(CliError::CommandFailed(response_error(&response)));
220    }
221    match response.data {
222        Some(ResponseData::SignalSample {
223            tick,
224            time_us,
225            mut values,
226        }) => {
227            let value = values.drain(..).next().ok_or_else(|| {
228                CliError::CommandFailed(format!("no values returned for '{selector}'"))
229            })?;
230            Ok((tick, time_us, value))
231        }
232        Some(other) => Err(CliError::CommandFailed(format!(
233            "unexpected sample response payload: {other:?}"
234        ))),
235        None => Err(CliError::CommandFailed(
236            "missing sample response payload".to_string(),
237        )),
238    }
239}
240
241pub(crate) async fn send_action_success(
242    session: &str,
243    action: InstanceAction,
244) -> Result<(), CliError> {
245    let response = send_action(session, action).await?;
246    if response.success {
247        Ok(())
248    } else {
249        Err(CliError::CommandFailed(response_error(&response)))
250    }
251}
252
253pub(crate) async fn send_action(
254    session: &str,
255    action: InstanceAction,
256) -> Result<Response, CliError> {
257    let request = Request {
258        id: Uuid::new_v4(),
259        action: RequestAction::Instance(action),
260    };
261    send_request(session, &request)
262        .await
263        .map_err(|e| CliError::CommandFailed(e.to_string()))
264}
265
266pub(crate) fn response_error(response: &Response) -> String {
267    response
268        .error
269        .clone()
270        .unwrap_or_else(|| "command failed".to_string())
271}