Skip to main content

systemprompt_agent/services/agent_orchestration/
process.rs

1use anyhow::{Context, Result};
2use std::fs::{self, File};
3use std::path::{Path, PathBuf};
4use std::process::Command;
5use systemprompt_models::{
6    AppPaths, CliPaths, Config, ProfileBootstrap, Secrets, SecretsBootstrap,
7};
8
9use crate::services::agent_orchestration::{OrchestrationError, OrchestrationResult};
10
11const MAX_LOG_SIZE: u64 = 10 * 1024 * 1024;
12
13fn rotate_log_if_needed(log_path: &Path) -> Result<()> {
14    if let Ok(metadata) = fs::metadata(log_path) {
15        if metadata.len() > MAX_LOG_SIZE {
16            let backup_path = log_path.with_extension("log.old");
17            fs::rename(log_path, &backup_path)?;
18        }
19    }
20    Ok(())
21}
22
23fn prepare_agent_log_file(agent_name: &str, log_dir: &Path) -> OrchestrationResult<File> {
24    if let Err(e) = fs::create_dir_all(log_dir) {
25        tracing::error!(
26            error = %e,
27            path = %log_dir.display(),
28            "Failed to create agent log directory - agent may fail to start"
29        );
30    }
31
32    let log_file_path = log_dir.join(format!("agent-{}.log", agent_name));
33    if let Err(e) = rotate_log_if_needed(&log_file_path) {
34        tracing::warn!(
35            error = %e,
36            path = %log_file_path.display(),
37            "Failed to rotate agent log file"
38        );
39    }
40
41    fs::OpenOptions::new()
42        .create(true)
43        .append(true)
44        .open(&log_file_path)
45        .map_err(|e| {
46            OrchestrationError::ProcessSpawnFailed(format!(
47                "Failed to create log file {}: {}",
48                log_file_path.display(),
49                e
50            ))
51        })
52}
53
54fn configure_secrets_env(command: &mut Command, secrets: &Secrets) {
55    if let Some(ref key) = secrets.gemini {
56        command.env("GEMINI_API_KEY", key);
57    }
58    if let Some(ref key) = secrets.anthropic {
59        command.env("ANTHROPIC_API_KEY", key);
60    }
61    if let Some(ref key) = secrets.openai {
62        command.env("OPENAI_API_KEY", key);
63    }
64    if let Some(ref key) = secrets.github {
65        command.env("GITHUB_TOKEN", key);
66    }
67
68    if !secrets.custom.is_empty() {
69        let uppercase_keys = secrets.custom_env_var_names();
70        command.env("SYSTEMPROMPT_CUSTOM_SECRETS", uppercase_keys.join(","));
71        for (env_name, value) in secrets.custom_env_vars() {
72            command.env(env_name, value);
73        }
74    }
75}
76
77fn build_agent_command(
78    binary_path: &PathBuf,
79    agent_name: &str,
80    port: u16,
81    profile_path: &str,
82    secrets: &Secrets,
83    config: &Config,
84    log_file: File,
85) -> Command {
86    let mut command = Command::new(binary_path);
87    for arg in CliPaths::agent_run_args() {
88        command.arg(arg);
89    }
90    command
91        .arg("--agent-name")
92        .arg(agent_name)
93        .arg("--port")
94        .arg(port.to_string())
95        .envs(std::env::vars())
96        .env("SYSTEMPROMPT_PROFILE", profile_path)
97        .env("SYSTEMPROMPT_SUBPROCESS", "1")
98        .env("JWT_SECRET", &secrets.jwt_secret)
99        .env("DATABASE_URL", &secrets.database_url)
100        .env("AGENT_NAME", agent_name)
101        .env("AGENT_PORT", port.to_string())
102        .env("DATABASE_TYPE", &config.database_type)
103        .stdout(std::process::Stdio::null())
104        .stderr(std::process::Stdio::from(log_file))
105        .stdin(std::process::Stdio::null());
106
107    configure_secrets_env(&mut command, secrets);
108    command
109}
110
111pub async fn spawn_detached(agent_name: &str, port: u16) -> OrchestrationResult<u32> {
112    let paths = AppPaths::get()
113        .map_err(|e| OrchestrationError::ProcessSpawnFailed(format!("Failed to get paths: {e}")))?;
114
115    let binary_path = paths.build().resolve_binary("systemprompt").map_err(|e| {
116        OrchestrationError::ProcessSpawnFailed(format!("Failed to find systemprompt binary: {e}"))
117    })?;
118
119    let config = Config::get().map_err(|e| {
120        OrchestrationError::ProcessSpawnFailed(format!("Failed to get config: {e}"))
121    })?;
122
123    let secrets = SecretsBootstrap::get().map_err(|e| {
124        OrchestrationError::ProcessSpawnFailed(format!("Failed to get secrets: {e}"))
125    })?;
126
127    let profile_path = ProfileBootstrap::get_path().map_err(|e| {
128        OrchestrationError::ProcessSpawnFailed(format!("Failed to get profile path: {e}"))
129    })?;
130
131    let log_file = prepare_agent_log_file(agent_name, &paths.system().logs())?;
132
133    let mut command = build_agent_command(
134        &binary_path,
135        agent_name,
136        port,
137        profile_path,
138        &secrets,
139        &config,
140        log_file,
141    );
142
143    let child = command.spawn().map_err(|e| {
144        OrchestrationError::ProcessSpawnFailed(format!("Failed to spawn {agent_name}: {e}"))
145    })?;
146
147    let pid = child.id();
148    std::mem::forget(child);
149
150    if !verify_process_started(pid) {
151        return Err(OrchestrationError::ProcessSpawnFailed(format!(
152            "Agent {} (PID {}) died immediately after spawn",
153            agent_name, pid
154        )));
155    }
156
157    tracing::debug!(pid = %pid, agent_name = %agent_name, "Detached process spawned");
158    Ok(pid)
159}
160
161#[cfg(unix)]
162fn verify_process_started(pid: u32) -> bool {
163    use nix::sys::wait::{waitpid, WaitPidFlag, WaitStatus};
164    use nix::unistd::Pid;
165
166    match waitpid(Pid::from_raw(pid as i32), Some(WaitPidFlag::WNOHANG)) {
167        Ok(WaitStatus::StillAlive) => true,
168        Ok(_) => false,
169        Err(_) => process_exists(pid),
170    }
171}
172
173#[cfg(windows)]
174fn verify_process_started(pid: u32) -> bool {
175    process_exists(pid)
176}
177
178#[cfg(unix)]
179pub fn process_exists(pid: u32) -> bool {
180    use nix::sys::signal;
181    use nix::unistd::Pid;
182    signal::kill(Pid::from_raw(pid as i32), None).is_ok()
183}
184
185#[cfg(windows)]
186pub fn process_exists(pid: u32) -> bool {
187    Command::new("tasklist")
188        .args(["/FI", &format!("PID eq {}", pid), "/NH"])
189        .output()
190        .map(|o| {
191            let stdout = String::from_utf8_lossy(&o.stdout);
192            !stdout.contains("INFO: No tasks") && !stdout.trim().is_empty()
193        })
194        .unwrap_or(false)
195}
196
197#[cfg(unix)]
198pub fn terminate_process(pid: u32) -> Result<()> {
199    use nix::sys::signal::{self, Signal};
200    use nix::unistd::Pid;
201
202    signal::kill(Pid::from_raw(pid as i32), Signal::SIGTERM)
203        .with_context(|| format!("Failed to send SIGTERM to PID {pid}"))?;
204
205    Ok(())
206}
207
208#[cfg(windows)]
209pub fn terminate_process(pid: u32) -> Result<()> {
210    let output = Command::new("taskkill")
211        .args(["/PID", &pid.to_string()])
212        .output()
213        .with_context(|| format!("Failed to run taskkill for PID {pid}"))?;
214
215    if !output.status.success() {
216        anyhow::bail!("taskkill failed for PID {pid}");
217    }
218    Ok(())
219}
220
221#[cfg(unix)]
222pub fn force_kill_process(pid: u32) -> Result<()> {
223    use nix::sys::signal::{self, Signal};
224    use nix::unistd::Pid;
225
226    signal::kill(Pid::from_raw(pid as i32), Signal::SIGKILL)
227        .with_context(|| format!("Failed to send SIGKILL to PID {pid}"))?;
228
229    Ok(())
230}
231
232#[cfg(windows)]
233pub fn force_kill_process(pid: u32) -> Result<()> {
234    let output = Command::new("taskkill")
235        .args(["/PID", &pid.to_string(), "/F"])
236        .output()
237        .with_context(|| format!("Failed to force-kill PID {pid}"))?;
238
239    if !output.status.success() {
240        anyhow::bail!("taskkill /F failed for PID {pid}");
241    }
242    Ok(())
243}
244
245pub async fn terminate_gracefully(pid: u32, timeout_secs: u64) -> Result<()> {
246    if !process_exists(pid) {
247        return Ok(());
248    }
249
250    terminate_process(pid)?;
251
252    let check_interval = tokio::time::Duration::from_millis(100);
253    let max_checks = (timeout_secs * 1000) / 100;
254
255    for _ in 0..max_checks {
256        if !process_exists(pid) {
257            return Ok(());
258        }
259        tokio::time::sleep(check_interval).await;
260    }
261
262    force_kill_process(pid)?;
263
264    for _ in 0..50 {
265        if !process_exists(pid) {
266            return Ok(());
267        }
268        tokio::time::sleep(check_interval).await;
269    }
270
271    Err(anyhow::anyhow!(
272        "Failed to kill process {} even with SIGKILL",
273        pid
274    ))
275}
276
277pub fn kill_process(pid: u32) -> bool {
278    terminate_process(pid).is_ok()
279}
280
281pub fn is_port_in_use(port: u16) -> bool {
282    use std::net::TcpListener;
283    TcpListener::bind(format!("127.0.0.1:{port}")).is_err()
284}
285
286pub async fn spawn_detached_process(agent_name: &str, port: u16) -> OrchestrationResult<u32> {
287    spawn_detached(agent_name, port).await
288}
289
290pub fn validate_agent_binary() -> Result<()> {
291    let paths = AppPaths::get().map_err(|e| anyhow::anyhow!("{}", e))?;
292    let binary_path = paths.build().resolve_binary("systemprompt")?;
293
294    let metadata = fs::metadata(&binary_path)
295        .with_context(|| format!("Failed to get metadata for: {}", binary_path.display()))?;
296
297    if !metadata.is_file() {
298        return Err(anyhow::anyhow!(
299            "Agent binary is not a file: {}",
300            binary_path.display()
301        ));
302    }
303
304    #[cfg(unix)]
305    {
306        use std::os::unix::fs::PermissionsExt;
307        let permissions = metadata.permissions();
308        if permissions.mode() & 0o111 == 0 {
309            return Err(anyhow::anyhow!(
310                "Agent binary is not executable: {}",
311                binary_path.display()
312            ));
313        }
314    }
315
316    Ok(())
317}