1use std::collections::BTreeMap;
39use std::io::{Read, Write};
40use std::path::PathBuf;
41use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
42use std::sync::{Arc, LazyLock, Mutex, OnceLock};
43use std::time::Duration;
44
45use harn_vm::VmValue;
46
47use crate::error::HostlibError;
48use crate::json::vm_dict_to_json;
49use crate::process::{self as process_handle, ProcessHandle, ProcessKiller, SpawnSpec};
50use crate::tools::args::to_agent_path;
51use crate::tools::proc::{self, CaptureConfig, CommandStatus, EnvMode};
52
53static HANDLE_COUNTER: AtomicU64 = AtomicU64::new(1);
55
56struct CancelState {
58 cancelled: AtomicBool,
61 timed_out: AtomicBool,
65 process_cleanup: Mutex<Option<process_handle::ProcessCleanupReport>>,
67}
68
69#[derive(Default)]
70struct OutputState {
71 stdout: Vec<u8>,
72 stderr: Vec<u8>,
73}
74
75struct HandleEntry {
77 handle: Option<Box<dyn ProcessHandle>>,
79 killer: Arc<dyn ProcessKiller>,
81 session_id: String,
82 cancel_state: Arc<CancelState>,
84 completion_tx: Option<std::sync::mpsc::SyncSender<()>>,
88 result_tx: Option<std::sync::mpsc::SyncSender<VmValue>>,
91 snapshot_binding: Option<harn_vm::value::DictMap>,
93}
94
95#[derive(Default)]
96struct HandleStore {
97 entries: BTreeMap<String, HandleEntry>,
98}
99
100static HANDLE_STORE: LazyLock<Mutex<HandleStore>> =
101 LazyLock::new(|| Mutex::new(HandleStore::default()));
102
103pub struct LongRunningHandleInfo {
106 pub command_id: String,
108 pub handle_id: String,
110 pub started_at: String,
112 pub pid: u32,
114 pub process_group_id: Option<u32>,
116 pub command_display: String,
118 pub snapshot_binding: Option<harn_vm::value::DictMap>,
120}
121
122pub(crate) struct LongRunningSpawnOptions {
123 pub(crate) env_mode: EnvMode,
124 pub(crate) env_remove: Vec<String>,
125 pub(crate) capture: CaptureConfig,
126 pub(crate) session_id: String,
127 pub(crate) progress_interval: Option<Duration>,
128 pub(crate) progress_max_inline_bytes: usize,
129 pub(crate) snapshot_binding: Option<harn_vm::value::DictMap>,
130}
131
132struct WaiterContext {
133 command_id: String,
134 handle_id: String,
135 session_id: String,
136 started_at: String,
137 process_group_id: Option<u32>,
138 command_display: String,
139 progress_interval: Option<Duration>,
140 progress_max_inline_bytes: usize,
141 snapshot_binding: Option<harn_vm::value::DictMap>,
142}
143
144struct ProgressThreadContext {
145 command_id: String,
146 handle_id: String,
147 session_id: String,
148 started_at: String,
149 command_display: String,
150 process_group_id: Option<u32>,
151 output_path: PathBuf,
152 stdout_path: PathBuf,
153 stderr_path: PathBuf,
154 output_state: Arc<Mutex<OutputState>>,
155 cancel_state: Arc<CancelState>,
156 done: Arc<AtomicBool>,
157 started: std::time::Instant,
158 interval: Duration,
159 max_inline_bytes: usize,
160 snapshot_binding: Option<harn_vm::value::DictMap>,
161}
162
163impl LongRunningHandleInfo {
164 pub fn into_handle_response(self) -> VmValue {
166 let Self {
167 command_id,
168 handle_id,
169 started_at,
170 pid,
171 process_group_id,
172 command_display,
173 snapshot_binding,
174 } = self;
175 proc::running_response(
176 command_id,
177 handle_id,
178 pid,
179 process_group_id,
180 started_at,
181 command_display,
182 snapshot_binding.as_ref(),
183 )
184 }
185}
186
187pub fn spawn_long_running(
193 builtin: &'static str,
194 program: String,
195 args: Vec<String>,
196 cwd: Option<PathBuf>,
197 env: BTreeMap<String, String>,
198 session_id: String,
199) -> Result<LongRunningHandleInfo, HostlibError> {
200 spawn_long_running_with_options(
201 builtin,
202 program,
203 args,
204 cwd,
205 env,
206 LongRunningSpawnOptions {
207 env_mode: EnvMode::InheritClean,
208 env_remove: Vec::new(),
209 capture: CaptureConfig::default(),
210 session_id,
211 progress_interval: None,
212 progress_max_inline_bytes: CaptureConfig::default().max_inline_bytes,
213 snapshot_binding: None,
214 },
215 )
216}
217
218pub(crate) fn spawn_long_running_with_options(
219 builtin: &'static str,
220 program: String,
221 args: Vec<String>,
222 cwd: Option<PathBuf>,
223 env: BTreeMap<String, String>,
224 options: LongRunningSpawnOptions,
225) -> Result<LongRunningHandleInfo, HostlibError> {
226 let mut env = env;
227 proc::apply_toolchain_path(cwd.as_deref(), &mut env, options.env_mode);
228 let spec = SpawnSpec {
229 builtin,
230 program: program.clone(),
231 args: args.clone(),
232 cwd,
233 env,
234 env_remove: options.env_remove.clone(),
235 env_mode: options.env_mode,
236 use_stdin: false,
237 configure_process_group: true,
238 output_capture: process_handle::OutputCapture::Pipe,
239 };
240 let handle = process_handle::spawn_process(spec)
241 .map_err(|e| proc::process_error_to_hostlib(builtin, e))?;
242
243 let pid = handle.pid().unwrap_or(0);
244 let process_group_id = handle.process_group_id();
245 let killer = handle.killer();
246 let id = HANDLE_COUNTER.fetch_add(1, Ordering::SeqCst);
247 let handle_id = format!("hto-{:x}-{id}", std::process::id());
248 let command_id = proc::next_command_id();
249 let started_at = proc::now_rfc3339();
250 let _artifacts = proc::register_live_artifacts(&command_id, Some(&handle_id))?;
251
252 let mut all_argv = vec![program];
253 all_argv.extend(args.iter().cloned());
254 let command_display = all_argv.join(" ");
255
256 let cancel_state = Arc::new(CancelState {
257 cancelled: AtomicBool::new(false),
258 timed_out: AtomicBool::new(false),
259 process_cleanup: Mutex::new(None),
260 });
261
262 {
263 let mut store = HANDLE_STORE
264 .lock()
265 .expect("long-running handle store poisoned");
266 store.entries.insert(
267 handle_id.clone(),
268 HandleEntry {
269 handle: Some(handle),
270 killer,
271 session_id: options.session_id.clone(),
272 cancel_state: cancel_state.clone(),
273 completion_tx: None,
274 result_tx: None,
275 snapshot_binding: options.snapshot_binding.clone(),
276 },
277 );
278 }
279
280 let waiter_context = WaiterContext {
281 command_id: command_id.clone(),
282 handle_id: handle_id.clone(),
283 session_id: options.session_id,
284 started_at: started_at.clone(),
285 process_group_id,
286 command_display: command_display.clone(),
287 progress_interval: options.progress_interval,
288 progress_max_inline_bytes: options.progress_max_inline_bytes,
289 snapshot_binding: options.snapshot_binding.clone(),
290 };
291 let waiter_thread_name = waiter_context.handle_id.clone();
292 let capture = options.capture;
293 std::thread::Builder::new()
294 .name(format!("hto-waiter-{waiter_thread_name}"))
295 .spawn(move || {
296 waiter_thread(waiter_context, cancel_state, capture);
297 })
298 .map_err(|e| HostlibError::Backend {
299 builtin,
300 message: format!("failed to spawn waiter thread: {e}"),
301 })?;
302
303 Ok(LongRunningHandleInfo {
304 command_id,
305 handle_id,
306 started_at,
307 pid,
308 process_group_id,
309 command_display,
310 snapshot_binding: options.snapshot_binding,
311 })
312}
313
314fn waiter_thread(context: WaiterContext, cancel_state: Arc<CancelState>, capture: CaptureConfig) {
316 let waiter_start = std::time::Instant::now();
317
318 let mut handle = {
321 let mut store = HANDLE_STORE
322 .lock()
323 .expect("long-running handle store poisoned");
324 match store.entries.get_mut(&context.handle_id) {
325 Some(entry) => match entry.handle.take() {
326 Some(h) => h,
327 None => return, },
329 None => return, }
331 };
332
333 let output_state = Arc::new(Mutex::new(OutputState::default()));
334 let done = Arc::new(AtomicBool::new(false));
335 let planned = proc::planned_artifact_paths(&context.command_id);
336 if let Some(parent) = planned.output_path.parent() {
337 let _ = std::fs::create_dir_all(parent);
338 }
339 let _ = std::fs::File::create(&planned.stdout_path);
340 let _ = std::fs::File::create(&planned.stderr_path);
341 let combined_file = std::fs::File::create(&planned.output_path)
342 .ok()
343 .map(|file| Arc::new(Mutex::new(file)));
344
345 let stdout_thread = handle.take_stdout().map(|out| {
346 spawn_output_drain(
347 out,
348 output_state.clone(),
349 planned.stdout_path.clone(),
350 combined_file.clone(),
351 true,
352 )
353 });
354 let stderr_thread = handle.take_stderr().map(|err| {
355 spawn_output_drain(
356 err,
357 output_state.clone(),
358 planned.stderr_path.clone(),
359 combined_file.clone(),
360 false,
361 )
362 });
363
364 let progress_thread = context
365 .progress_interval
366 .filter(|interval| !interval.is_zero())
367 .map(|interval| {
368 spawn_progress_thread(ProgressThreadContext {
369 command_id: context.command_id.clone(),
370 handle_id: context.handle_id.clone(),
371 session_id: context.session_id.clone(),
372 started_at: context.started_at.clone(),
373 command_display: context.command_display.clone(),
374 process_group_id: context.process_group_id,
375 output_path: planned.output_path.clone(),
376 stdout_path: planned.stdout_path.clone(),
377 stderr_path: planned.stderr_path.clone(),
378 output_state: output_state.clone(),
379 cancel_state: cancel_state.clone(),
380 done: done.clone(),
381 started: waiter_start,
382 interval,
383 max_inline_bytes: context.progress_max_inline_bytes,
384 snapshot_binding: context.snapshot_binding.clone(),
385 })
386 });
387
388 let status = handle.wait().ok();
389
390 if let Some(thread) = stdout_thread {
391 let _ = thread.join();
392 }
393 if let Some(thread) = stderr_thread {
394 let _ = thread.join();
395 }
396 done.store(true, Ordering::Release);
397 drop(progress_thread);
398 let (stdout, stderr) = {
399 let state = output_state
400 .lock()
401 .unwrap_or_else(|poison| poison.into_inner());
402 (state.stdout.clone(), state.stderr.clone())
403 };
404
405 let (completion_tx, result_tx) = {
408 let mut store = HANDLE_STORE
409 .lock()
410 .expect("long-running handle store poisoned");
411 let entry = store
412 .entries
413 .remove(&context.handle_id)
414 .map(|mut e| (e.completion_tx.take(), e.result_tx.take()));
415 entry.unwrap_or((None, None))
416 };
417
418 let signal_done = move || {
419 if let Some(tx) = completion_tx {
420 let _ = tx.try_send(());
421 }
422 };
423
424 let cancelled = cancel_state.cancelled.load(Ordering::Acquire);
425 let timed_out = cancelled && cancel_state.timed_out.load(Ordering::Acquire);
426 let process_cleanup = cancel_state
427 .process_cleanup
428 .lock()
429 .unwrap_or_else(|poison| poison.into_inner())
430 .clone();
431
432 let (exit_code, signal_name) = match status {
433 Some(s) => decode_exit_status(s),
434 None => (-1, Some("SIGKILL".to_string())),
436 };
437 let command_status = if timed_out {
438 CommandStatus::TimedOut
439 } else if cancelled {
440 CommandStatus::Killed
441 } else {
442 CommandStatus::Completed
443 };
444 let duration = waiter_start.elapsed();
445 let duration_ms = duration.as_millis() as i64;
446 let artifacts = match proc::persist_artifacts(
447 &context.command_id,
448 &stdout,
449 &stderr,
450 Some(&context.handle_id),
451 ) {
452 Ok(artifacts) => artifacts,
453 Err(_) => return,
454 };
455 let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
456
457 let mut payload = serde_json::Map::new();
458 payload.insert(
459 "command_id".into(),
460 serde_json::Value::String(context.command_id.clone()),
461 );
462 payload.insert(
463 "status".into(),
464 serde_json::Value::String(command_status.as_str().to_string()),
465 );
466 payload.insert(
467 "handle_id".into(),
468 serde_json::Value::String(context.handle_id),
469 );
470 payload.insert(
471 "command_or_op_descriptor".into(),
472 serde_json::Value::String(context.command_display),
473 );
474 payload.insert(
475 "started_at".into(),
476 serde_json::Value::String(context.started_at),
477 );
478 payload.insert(
479 "ended_at".into(),
480 serde_json::Value::String(proc::now_rfc3339()),
481 );
482 payload.insert(
483 "duration_ms".into(),
484 serde_json::Value::Number(duration_ms.into()),
485 );
486 payload.insert(
487 "exit_code".into(),
488 serde_json::Value::Number(exit_code.into()),
489 );
490 payload.insert("timed_out".into(), serde_json::Value::Bool(timed_out));
491 payload.insert("stdout".into(), serde_json::Value::String(inline_stdout));
492 payload.insert("stderr".into(), serde_json::Value::String(inline_stderr));
493 payload.insert(
494 "output_path".into(),
495 serde_json::Value::String(to_agent_path(&artifacts.output_path)),
496 );
497 payload.insert(
498 "stdout_path".into(),
499 serde_json::Value::String(to_agent_path(&artifacts.stdout_path)),
500 );
501 payload.insert(
502 "stderr_path".into(),
503 serde_json::Value::String(to_agent_path(&artifacts.stderr_path)),
504 );
505 payload.insert(
506 "line_count".into(),
507 serde_json::Value::Number(artifacts.line_count.into()),
508 );
509 payload.insert(
510 "byte_count".into(),
511 serde_json::Value::Number(artifacts.byte_count.into()),
512 );
513 payload.insert(
514 "output_sha256".into(),
515 serde_json::Value::String(artifacts.output_sha256),
516 );
517 if let Some(pgid) = context.process_group_id {
518 payload.insert(
519 "process_group_id".into(),
520 serde_json::Value::Number((pgid as u64).into()),
521 );
522 }
523 if let Some(sig) = signal_name {
524 payload.insert("signal".into(), serde_json::Value::String(sig));
525 } else {
526 payload.insert("signal".into(), serde_json::Value::Null);
527 }
528 if let Some(snapshot_binding) = context.snapshot_binding.as_ref() {
529 payload.insert("snapshot_binding".into(), vm_dict_to_json(snapshot_binding));
530 }
531 if let Some(process_cleanup) = process_cleanup.as_ref() {
532 payload.insert(
533 "process_cleanup".into(),
534 proc::process_cleanup_to_json(process_cleanup),
535 );
536 }
537
538 if let Some(tx) = result_tx {
539 let value = serde_json::Value::Object(payload.clone());
540 let _ = tx.try_send(harn_vm::json_to_vm_value(&value));
541 }
542 if !cancelled {
543 let content = serde_json::to_string(&payload).unwrap_or_default();
544 harn_vm::orchestration::agent_inbox::push(
545 &context.session_id,
546 "tool_result",
547 &content,
548 "hostlib.long_running.exit",
549 );
550 }
551 signal_done();
552}
553
554fn spawn_output_drain(
555 mut reader: Box<dyn Read + Send>,
556 state: Arc<Mutex<OutputState>>,
557 path: std::path::PathBuf,
558 combined_file: Option<Arc<Mutex<std::fs::File>>>,
559 stdout: bool,
560) -> std::thread::JoinHandle<()> {
561 std::thread::spawn(move || {
562 let mut file = std::fs::File::create(path).ok();
563 let mut buf = [0_u8; 8192];
564 loop {
565 let read = match reader.read(&mut buf) {
566 Ok(0) => break,
567 Ok(read) => read,
568 Err(_) => break,
569 };
570 let chunk = &buf[..read];
571 if let Some(file) = file.as_mut() {
572 let _ = file.write_all(chunk);
573 }
574 if let Some(combined) = combined_file.as_ref() {
575 if let Ok(mut combined) = combined.lock() {
576 let _ = combined.write_all(chunk);
577 }
578 }
579 if let Ok(mut state) = state.lock() {
580 if stdout {
581 state.stdout.extend_from_slice(chunk);
582 } else {
583 state.stderr.extend_from_slice(chunk);
584 }
585 }
586 }
587 })
588}
589
590fn spawn_progress_thread(context: ProgressThreadContext) -> std::thread::JoinHandle<()> {
591 std::thread::spawn(move || {
592 while !context.done.load(Ordering::Acquire)
593 && !context.cancel_state.cancelled.load(Ordering::Acquire)
594 {
595 std::thread::sleep(context.interval);
596 if context.done.load(Ordering::Acquire)
597 || context.cancel_state.cancelled.load(Ordering::Acquire)
598 {
599 break;
600 }
601 let (stdout, stderr) = {
602 let state = context
603 .output_state
604 .lock()
605 .unwrap_or_else(|poison| poison.into_inner());
606 (state.stdout.clone(), state.stderr.clone())
607 };
608 let capture = CaptureConfig {
609 max_inline_bytes: context.max_inline_bytes,
610 ..CaptureConfig::default()
611 };
612 let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
613 let byte_count = stdout.len().saturating_add(stderr.len());
614 let mut payload = serde_json::json!({
615 "command_id": &context.command_id,
616 "handle_id": &context.handle_id,
617 "status": CommandStatus::Running.as_str(),
618 "command_or_op_descriptor": &context.command_display,
619 "started_at": &context.started_at,
620 "ended_at": null,
621 "duration_ms": context.started.elapsed().as_millis() as i64,
622 "exit_code": null,
623 "signal": null,
624 "stdout": inline_stdout,
625 "stderr": inline_stderr,
626 "output_path": to_agent_path(&context.output_path),
627 "stdout_path": to_agent_path(&context.stdout_path),
628 "stderr_path": to_agent_path(&context.stderr_path),
629 "byte_count": byte_count as i64,
630 "line_count": stdout.iter().chain(stderr.iter()).filter(|byte| **byte == b'\n').count() as i64,
631 "process_group_id": context.process_group_id,
632 });
633 if let (Some(object), Some(snapshot_binding)) =
634 (payload.as_object_mut(), context.snapshot_binding.as_ref())
635 {
636 object.insert(
637 "snapshot_binding".to_string(),
638 vm_dict_to_json(snapshot_binding),
639 );
640 }
641 harn_vm::orchestration::agent_inbox::push(
642 &context.session_id,
643 "tool_progress",
644 &payload.to_string(),
645 "hostlib.long_running.progress",
646 );
647 }
648 })
649}
650
651pub(crate) struct CancelOptions {
652 pub(crate) timed_out: bool,
653 pub(crate) wait_result: Option<Duration>,
654}
655
656pub(crate) struct CancelOutcome {
657 pub(crate) cancelled: bool,
658 pub(crate) result: Option<VmValue>,
659}
660
661pub fn cancel_handle(handle_id: &str) -> bool {
665 cancel_handle_with_options(
666 handle_id,
667 CancelOptions {
668 timed_out: false,
669 wait_result: None,
670 },
671 )
672 .cancelled
673}
674
675pub(crate) fn snapshot_binding_for_handle(handle_id: &str) -> Option<harn_vm::value::DictMap> {
676 let store = HANDLE_STORE
677 .lock()
678 .expect("long-running handle store poisoned");
679 store
680 .entries
681 .get(handle_id)
682 .and_then(|entry| entry.snapshot_binding.clone())
683}
684
685pub(crate) fn cancel_handle_with_options(handle_id: &str, options: CancelOptions) -> CancelOutcome {
686 let (killer, cancel_state, result_rx) = {
687 let mut store = HANDLE_STORE
688 .lock()
689 .expect("long-running handle store poisoned");
690 let Some(entry) = store.entries.get_mut(handle_id) else {
691 return CancelOutcome {
692 cancelled: false,
693 result: None,
694 };
695 };
696 if entry.cancel_state.cancelled.swap(true, Ordering::AcqRel) {
697 return CancelOutcome {
698 cancelled: false,
699 result: None,
700 };
701 }
702 entry
703 .cancel_state
704 .timed_out
705 .store(options.timed_out, Ordering::Release);
706 let result_rx = options.wait_result.map(|_| {
707 let (tx, rx) = std::sync::mpsc::sync_channel::<VmValue>(1);
708 entry.result_tx = Some(tx);
709 rx
710 });
711 (entry.killer.clone(), entry.cancel_state.clone(), result_rx)
712 };
713 do_kill(killer, cancel_state);
714 let result = match (options.wait_result, result_rx) {
715 (Some(timeout), Some(rx)) => rx.recv_timeout(timeout).ok(),
716 _ => None,
717 };
718 CancelOutcome {
719 cancelled: true,
720 result,
721 }
722}
723
724type SessionKillEntry = (Arc<dyn ProcessKiller>, Arc<CancelState>);
728
729pub fn cancel_session_handles(session_id: &str) {
732 let to_kill: Vec<SessionKillEntry> = {
733 let store = HANDLE_STORE
734 .lock()
735 .expect("long-running handle store poisoned");
736 let matching: Vec<String> = store
737 .entries
738 .iter()
739 .filter(|(_, e)| e.session_id == session_id)
740 .map(|(id, _)| id.clone())
741 .collect();
742 matching
743 .into_iter()
744 .filter_map(|id| {
745 let entry = store.entries.get(&id)?;
746 if entry.cancel_state.cancelled.swap(true, Ordering::AcqRel) {
747 return None;
748 }
749 entry.cancel_state.timed_out.store(false, Ordering::Release);
750 Some((entry.killer.clone(), entry.cancel_state.clone()))
751 })
752 .collect()
753 };
754 for (killer, cancel_state) in to_kill {
755 do_kill(killer, cancel_state);
756 }
757}
758
759fn do_kill(killer: Arc<dyn ProcessKiller>, cancel_state: Arc<CancelState>) {
762 let report = killer.kill();
765 {
766 let mut stored = cancel_state
767 .process_cleanup
768 .lock()
769 .unwrap_or_else(|poison| poison.into_inner());
770 match stored.as_mut() {
771 Some(existing) => existing.merge(report),
772 None => *stored = Some(report),
773 }
774 }
775 cancel_state.cancelled.store(true, Ordering::Release);
776}
777
778pub(crate) fn register_cleanup_hook() {
782 static REGISTERED: OnceLock<()> = OnceLock::new();
783 REGISTERED.get_or_init(|| {
784 let hook: Arc<dyn Fn(&str) + Send + Sync> = Arc::new(|session_id: &str| {
785 cancel_session_handles(session_id);
786 });
787 harn_vm::register_session_end_hook(hook);
788 });
789}
790
791fn decode_exit_status(status: process_handle::ExitStatus) -> (i32, Option<String>) {
792 if let Some(code) = status.code {
793 return (code, None);
794 }
795 if let Some(sig) = status.signal {
796 return (-1, Some(format!("SIG{sig}")));
797 }
798 (-1, None)
799}
800
801pub fn register_completion_notifier(handle_id: &str) -> Option<std::sync::mpsc::Receiver<()>> {
807 let (tx, rx) = std::sync::mpsc::sync_channel::<()>(1);
808 let mut store = HANDLE_STORE
809 .lock()
810 .expect("long-running handle store poisoned");
811 let entry = store.entries.get_mut(handle_id)?;
812 entry.completion_tx = Some(tx);
813 Some(rx)
814}