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

systemprompt_agent/services/agent_orchestration/
process.rs
