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
77struct BuildAgentCommandParams<'a> {
78    binary_path: &'a PathBuf,
79    agent_name: &'a str,
80    port: u16,
81    profile_path: &'a str,
82    secrets: &'a Secrets,
83    config: &'a Config,
84    log_file: File,
85}
86
87fn build_agent_command(params: BuildAgentCommandParams<'_>) -> Command {
88    let BuildAgentCommandParams {
89        binary_path,
90        agent_name,
91        port,
92        profile_path,
93        secrets,
94        config,
95        log_file,
96    } = params;
97    let mut command = Command::new(binary_path);
98    for arg in CliPaths::agent_run_args() {
99        command.arg(arg);
100    }
101    command
102        .arg("--agent-name")
103        .arg(agent_name)
104        .arg("--port")
105        .arg(port.to_string())
106        .env_clear();
107    if let Ok(path) = std::env::var("PATH") {
108        command.env("PATH", path);
109    }
110    if let Ok(home) = std::env::var("HOME") {
111        command.env("HOME", home);
112    }
113    command
114        .env("SYSTEMPROMPT_PROFILE", profile_path)
115        .env("SYSTEMPROMPT_SUBPROCESS", "1")
116        .env("JWT_SECRET", &secrets.jwt_secret)
117        .env("DATABASE_URL", &secrets.database_url)
118        .env("AGENT_NAME", agent_name)
119        .env("AGENT_PORT", port.to_string())
120        .env("DATABASE_TYPE", &config.database_type)
121        .stdout(std::process::Stdio::null())
122        .stderr(std::process::Stdio::from(log_file))
123        .stdin(std::process::Stdio::null());
124
125    if let Ok(fly_app) = std::env::var("FLY_APP_NAME") {
126        command.env("FLY_APP_NAME", fly_app);
127    }
128
129    configure_secrets_env(&mut command, secrets);
130    command
131}
132
133pub fn spawn_detached(agent_name: &str, port: u16) -> OrchestrationResult<u32> {
134    let paths = AppPaths::get()
135        .map_err(|e| OrchestrationError::ProcessSpawnFailed(format!("Failed to get paths: {e}")))?;
136
137    let binary_path = paths.build().resolve_binary("systemprompt").map_err(|e| {
138        OrchestrationError::ProcessSpawnFailed(format!("Failed to find systemprompt binary: {e}"))
139    })?;
140
141    let config = Config::get().map_err(|e| {
142        OrchestrationError::ProcessSpawnFailed(format!("Failed to get config: {e}"))
143    })?;
144
145    let secrets = SecretsBootstrap::get().map_err(|e| {
146        OrchestrationError::ProcessSpawnFailed(format!("Failed to get secrets: {e}"))
147    })?;
148
149    let profile_path = ProfileBootstrap::get_path().map_err(|e| {
150        OrchestrationError::ProcessSpawnFailed(format!("Failed to get profile path: {e}"))
151    })?;
152
153    let log_file = prepare_agent_log_file(agent_name, &paths.system().logs())?;
154
155    let mut command = build_agent_command(BuildAgentCommandParams {
156        binary_path: &binary_path,
157        agent_name,
158        port,
159        profile_path,
160        secrets,
161        config,
162        log_file,
163    });
164
165    let child = command.spawn().map_err(|e| {
166        OrchestrationError::ProcessSpawnFailed(format!("Failed to spawn {agent_name}: {e}"))
167    })?;
168
169    let pid = child.id();
170    std::mem::forget(child);
171
172    if !verify_process_started(pid) {
173        return Err(OrchestrationError::ProcessSpawnFailed(format!(
174            "Agent {} (PID {}) died immediately after spawn",
175            agent_name, pid
176        )));
177    }
178
179    tracing::debug!(pid = %pid, agent_name = %agent_name, "Detached process spawned");
180    Ok(pid)
181}
182
183#[cfg(unix)]
184fn verify_process_started(pid: u32) -> bool {
185    use nix::sys::wait::{WaitPidFlag, WaitStatus, waitpid};
186    use nix::unistd::Pid;
187
188    match waitpid(Pid::from_raw(pid as i32), Some(WaitPidFlag::WNOHANG)) {
189        Ok(WaitStatus::StillAlive) => true,
190        Ok(_) => false,
191        Err(_) => process_exists(pid),
192    }
193}
194
195#[cfg(windows)]
196fn verify_process_started(pid: u32) -> bool {
197    process_exists(pid)
198}
199
200#[cfg(unix)]
201pub fn process_exists(pid: u32) -> bool {
202    use nix::sys::signal;
203    use nix::unistd::Pid;
204    signal::kill(Pid::from_raw(pid as i32), None).is_ok()
205}
206
207#[cfg(windows)]
208pub fn process_exists(pid: u32) -> bool {
209    match Command::new("tasklist")
210        .args(["/FI", &format!("PID eq {}", pid), "/NH"])
211        .output()
212    {
213        Ok(o) => {
214            let stdout = String::from_utf8_lossy(&o.stdout);
215            !stdout.contains("INFO: No tasks") && !stdout.trim().is_empty()
216        },
217        Err(e) => {
218            tracing::warn!(
219                pid = pid,
220                error = %e,
221                "failed to run `tasklist` while checking process; assuming dead",
222            );
223            false
224        },
225    }
226}
227
228#[cfg(unix)]
229pub fn terminate_process(pid: u32) -> Result<()> {
230    use nix::sys::signal::{self, Signal};
231    use nix::unistd::Pid;
232
233    signal::kill(Pid::from_raw(pid as i32), Signal::SIGTERM)
234        .with_context(|| format!("Failed to send SIGTERM to PID {pid}"))?;
235
236    Ok(())
237}
238
239#[cfg(windows)]
240pub fn terminate_process(pid: u32) -> Result<()> {
241    let output = Command::new("taskkill")
242        .args(["/PID", &pid.to_string()])
243        .output()
244        .with_context(|| format!("Failed to run taskkill for PID {pid}"))?;
245
246    if !output.status.success() {
247        anyhow::bail!("taskkill failed for PID {pid}");
248    }
249    Ok(())
250}
251
252#[cfg(unix)]
253pub fn force_kill_process(pid: u32) -> Result<()> {
254    use nix::sys::signal::{self, Signal};
255    use nix::unistd::Pid;
256
257    signal::kill(Pid::from_raw(pid as i32), Signal::SIGKILL)
258        .with_context(|| format!("Failed to send SIGKILL to PID {pid}"))?;
259
260    Ok(())
261}
262
263#[cfg(windows)]
264pub fn force_kill_process(pid: u32) -> Result<()> {
265    let output = Command::new("taskkill")
266        .args(["/PID", &pid.to_string(), "/F"])
267        .output()
268        .with_context(|| format!("Failed to force-kill PID {pid}"))?;
269
270    if !output.status.success() {
271        anyhow::bail!("taskkill /F failed for PID {pid}");
272    }
273    Ok(())
274}
275
276pub async fn terminate_gracefully(pid: u32, timeout_secs: u64) -> Result<()> {
277    if !process_exists(pid) {
278        return Ok(());
279    }
280
281    terminate_process(pid)?;
282
283    let check_interval = tokio::time::Duration::from_millis(100);
284    let max_checks = (timeout_secs * 1000) / 100;
285
286    for _ in 0..max_checks {
287        if !process_exists(pid) {
288            return Ok(());
289        }
290        tokio::time::sleep(check_interval).await;
291    }
292
293    force_kill_process(pid)?;
294
295    for _ in 0..50 {
296        if !process_exists(pid) {
297            return Ok(());
298        }
299        tokio::time::sleep(check_interval).await;
300    }
301
302    Err(anyhow::anyhow!(
303        "Failed to kill process {} even with SIGKILL",
304        pid
305    ))
306}
307
308pub fn kill_process(pid: u32) -> bool {
309    terminate_process(pid).is_ok()
310}
311
312pub fn is_port_in_use(port: u16) -> bool {
313    use std::net::TcpListener;
314    TcpListener::bind(format!("127.0.0.1:{port}")).is_err()
315}
316
317pub fn spawn_detached_process(agent_name: &str, port: u16) -> OrchestrationResult<u32> {
318    spawn_detached(agent_name, port)
319}
320
321pub fn validate_agent_binary() -> Result<()> {
322    let paths = AppPaths::get().map_err(|e| anyhow::anyhow!("{}", e))?;
323    let binary_path = paths.build().resolve_binary("systemprompt")?;
324
325    let metadata = fs::metadata(&binary_path)
326        .with_context(|| format!("Failed to get metadata for: {}", binary_path.display()))?;
327
328    if !metadata.is_file() {
329        return Err(anyhow::anyhow!(
330            "Agent binary is not a file: {}",
331            binary_path.display()
332        ));
333    }
334
335    #[cfg(unix)]
336    {
337        use std::os::unix::fs::PermissionsExt;
338        let permissions = metadata.permissions();
339        if permissions.mode() & 0o111 == 0 {
340            return Err(anyhow::anyhow!(
341                "Agent binary is not executable: {}",
342                binary_path.display()
343            ));
344        }
345    }
346
347    Ok(())
348}
systemprompt_agent/services/agent_orchestration/process.rs

systemprompt_agent/services/agent_orchestration/
process.rs
