Skip to main content

vv_agent/runtime/
background_sessions.rs

1mod listeners;
2mod options;
3mod session;
4mod subscription;
5#[cfg(test)]
6mod tests;
7
8use std::collections::BTreeMap;
9use std::path::PathBuf;
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::sync::{Mutex, OnceLock};
12use std::thread;
13use std::time::Duration;
14
15use serde_json::{json, Value};
16
17use crate::runtime::processes::start_captured_process_with_env;
18
19use listeners::notify_background_listeners;
20use session::BackgroundSession;
21
22pub use listeners::BackgroundSessionListener;
23pub use options::{BackgroundSessionAdoptOptions, BackgroundSessionStartOptions};
24pub use subscription::BackgroundSessionSubscription;
25
26static MANAGER: OnceLock<BackgroundSessionManager> = OnceLock::new();
27
28pub fn background_session_manager() -> &'static BackgroundSessionManager {
29    MANAGER.get_or_init(BackgroundSessionManager::default)
30}
31
32#[derive(Default)]
33pub struct BackgroundSessionManager {
34    sessions: Mutex<BTreeMap<String, BackgroundSession>>,
35    next_id: AtomicU64,
36    next_listener_id: AtomicU64,
37}
38
39impl BackgroundSessionManager {
40    pub fn start(
41        &self,
42        command: impl Into<String>,
43        cwd: impl Into<PathBuf>,
44        timeout_seconds: u64,
45        options: BackgroundSessionStartOptions,
46    ) -> Result<String, String> {
47        let command = command.into();
48        let cwd = cwd.into();
49        let prepared = super::shell::prepare_shell_execution(
50            &command,
51            options.auto_confirm,
52            options.stdin.as_deref(),
53            options.shell.as_deref(),
54            options.windows_shell_priority.as_deref(),
55        )?;
56        let started = start_captured_process_with_env(
57            &prepared.command,
58            &cwd,
59            prepared.stdin.as_deref(),
60            options.env.as_ref(),
61        )
62        .map_err(|error| error.to_string())?;
63        Ok(self.adopt_running_process(
64            command,
65            cwd,
66            timeout_seconds,
67            started.child,
68            started.output_path,
69            prepared.shell,
70        ))
71    }
72
73    pub fn adopt_running_process(
74        &self,
75        command: impl Into<String>,
76        cwd: impl Into<PathBuf>,
77        timeout_seconds: u64,
78        child: std::process::Child,
79        output_path: PathBuf,
80        shell: Option<String>,
81    ) -> String {
82        let mut options =
83            BackgroundSessionAdoptOptions::new(command, cwd, timeout_seconds, child, output_path);
84        options.shell = shell;
85        self.adopt_running_process_with_options(options)
86    }
87
88    pub fn adopt_running_process_with_options(
89        &self,
90        options: BackgroundSessionAdoptOptions,
91    ) -> String {
92        let id = self.next_id.fetch_add(1, Ordering::Relaxed) + 1;
93        let session_id = format!("bg_{id:012x}");
94        let session = BackgroundSession::from_adopt_options(session_id.clone(), options);
95        self.sessions
96            .lock()
97            .expect("background session manager poisoned")
98            .insert(session_id.clone(), session);
99        self.start_watch_thread(session_id.clone());
100        session_id
101    }
102
103    pub fn subscribe(
104        &'static self,
105        session_id: &str,
106        listener: BackgroundSessionListener,
107    ) -> BackgroundSessionSubscription {
108        let mut snapshot = None;
109        let listener_id = self.next_listener_id.fetch_add(1, Ordering::Relaxed) + 1;
110        {
111            let mut sessions = self
112                .sessions
113                .lock()
114                .expect("background session manager poisoned");
115            let Some(session) = sessions.get_mut(session_id) else {
116                return BackgroundSessionSubscription::noop();
117            };
118            if session.is_terminal() {
119                snapshot = Some(session.snapshot());
120            } else {
121                session.add_listener(listener_id, listener.clone());
122            }
123        }
124        if let Some(payload) = snapshot {
125            listener(&payload);
126            return BackgroundSessionSubscription::noop();
127        }
128        BackgroundSessionSubscription::new(session_id.to_string(), listener_id, self)
129    }
130
131    fn unsubscribe(&self, session_id: &str, listener_id: u64) {
132        let mut sessions = self
133            .sessions
134            .lock()
135            .expect("background session manager poisoned");
136        if let Some(session) = sessions.get_mut(session_id) {
137            session.remove_listener(listener_id);
138        }
139    }
140
141    fn start_watch_thread(&self, session_id: String) {
142        let thread_name = format!("vv-agent-bg-{session_id}");
143        let _ = thread::Builder::new()
144            .name(thread_name)
145            .spawn(move || loop {
146                thread::sleep(Duration::from_millis(200));
147                let payload = background_session_manager().check(&session_id);
148                let status = payload
149                    .get("status")
150                    .and_then(Value::as_str)
151                    .unwrap_or("missing");
152                if status != "running" {
153                    break;
154                }
155            });
156    }
157
158    pub fn check(&self, session_id: &str) -> Value {
159        let mut sessions = self
160            .sessions
161            .lock()
162            .expect("background session manager poisoned");
163        let Some(session) = sessions.get_mut(session_id) else {
164            return json!({
165                "status": "missing",
166                "session_id": session_id,
167                "error": "Background session not found",
168            });
169        };
170
171        if session.is_terminal() {
172            return session.snapshot();
173        }
174
175        let elapsed = session.elapsed();
176        if session.timed_out(elapsed) {
177            session.finalize_timeout();
178            let payload = session.snapshot();
179            let terminal_listeners = session.take_listeners();
180            drop(sessions);
181            notify_background_listeners(terminal_listeners, &payload);
182            return payload;
183        }
184
185        match session.try_wait() {
186            Ok(Some(exit_code)) => {
187                session.finalize_completed(exit_code);
188                let payload = session.snapshot();
189                let terminal_listeners = session.take_listeners();
190                drop(sessions);
191                notify_background_listeners(terminal_listeners, &payload);
192                payload
193            }
194            Ok(None) => session.running_snapshot(elapsed),
195            Err(error) => {
196                session.finalize_failed_with_output(-1, error.to_string());
197                let payload = session.snapshot();
198                let terminal_listeners = session.take_listeners();
199                drop(sessions);
200                notify_background_listeners(terminal_listeners, &payload);
201                payload
202            }
203        }
204    }
205}