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, MutexGuard, OnceLock};
43use std::time::Duration;
44
45use harn_vm::VmDictExt;
46use harn_vm::VmValue;
47
48use crate::error::HostlibError;
49use crate::json::vm_dict_to_json;
50use crate::process::{
51 self as process_handle, OwnerDeathPolicy, ProcessHandle, ProcessKiller, SpawnSpec,
52};
53use crate::tools::args::to_agent_path;
54use crate::tools::proc::{self, CaptureConfig, CommandStatus, EnvMode};
55
56static HANDLE_COUNTER: AtomicU64 = AtomicU64::new(1);
58
59#[derive(Default)]
66struct CancelState {
67 state: Mutex<CancellationState>,
68}
69
70#[derive(Default)]
71struct CancellationState {
72 cancellation_requested: bool,
74 cancelled: bool,
76 timed_out: bool,
78 process_cleanup: Option<process_handle::ProcessCleanupReport>,
80 completed: bool,
83}
84
85#[derive(Clone, Debug)]
86struct CancellationSnapshot {
87 cancelled: bool,
88 timed_out: bool,
89 process_cleanup: Option<process_handle::ProcessCleanupReport>,
90}
91
92impl CancelState {
93 fn begin_cancellation(&self, timed_out: bool) -> Option<MutexGuard<'_, CancellationState>> {
94 let mut state = self
95 .state
96 .lock()
97 .unwrap_or_else(|poison| poison.into_inner());
98 if state.cancellation_requested || state.completed {
99 return None;
100 }
101 state.cancellation_requested = true;
102 state.timed_out = timed_out;
103 Some(state)
104 }
105
106 fn complete_wait(&self) -> CancellationSnapshot {
107 let mut state = self
108 .state
109 .lock()
110 .unwrap_or_else(|poison| poison.into_inner());
111 state.completed = true;
112 CancellationSnapshot {
113 cancelled: state.cancelled,
114 timed_out: state.timed_out,
115 process_cleanup: state.process_cleanup.clone(),
116 }
117 }
118
119 fn cancellation_published(&self) -> bool {
120 self.state
121 .lock()
122 .unwrap_or_else(|poison| poison.into_inner())
123 .cancelled
124 }
125}
126
127impl CancellationState {
128 fn record_cleanup(&mut self, report: process_handle::ProcessCleanupReport) {
129 match self.process_cleanup.as_mut() {
130 Some(existing) => existing.merge(report),
131 None => self.process_cleanup = Some(report),
132 }
133 }
134
135 fn publish_cancellation(&mut self) {
136 debug_assert!(self.cancellation_requested);
137 self.cancelled = true;
138 }
139}
140
141fn kill_and_publish(killer: &dyn ProcessKiller, cancellation: &mut CancellationState) {
142 let report = killer.kill();
143 cancellation.record_cleanup(report);
144 cancellation.publish_cancellation();
145}
146
147#[derive(Default)]
148pub(crate) struct OutputState {
149 pub(crate) stdout: Vec<u8>,
150 pub(crate) stderr: Vec<u8>,
151 pub(crate) combined: Vec<u8>,
152 pub(crate) terminal: Option<VmValue>,
153 last_output_at: Option<std::time::Instant>,
157}
158
159pub(crate) struct OutputFeed {
160 pub(crate) state: Mutex<OutputState>,
161 notify: tokio::sync::Notify,
162}
163
164impl Default for OutputFeed {
165 fn default() -> Self {
166 Self {
167 state: Mutex::new(OutputState::default()),
168 notify: tokio::sync::Notify::new(),
169 }
170 }
171}
172
173impl OutputFeed {
174 pub(crate) fn notified(&self) -> tokio::sync::futures::Notified<'_> {
175 self.notify.notified()
176 }
177}
178
179struct HandleEntry {
181 handle: Option<Box<dyn ProcessHandle>>,
183 killer: Arc<dyn ProcessKiller>,
185 session_id: String,
186 cancel_state: Arc<CancelState>,
188 output_feed: Arc<OutputFeed>,
189 completion_tx: Option<std::sync::mpsc::SyncSender<()>>,
193 result_txs: Vec<std::sync::mpsc::SyncSender<VmValue>>,
197 snapshot_binding: Option<harn_vm::value::DictMap>,
199 lease: LeaseTag,
201 command_display: String,
204 started_at: String,
206}
207
208#[derive(Default)]
209struct HandleStore {
210 entries: BTreeMap<String, HandleEntry>,
211}
212
213static HANDLE_STORE: LazyLock<Mutex<HandleStore>> =
214 LazyLock::new(|| Mutex::new(HandleStore::default()));
215
216type HandleNotifiers = (
217 Option<std::sync::mpsc::SyncSender<()>>,
218 Vec<std::sync::mpsc::SyncSender<VmValue>>,
219);
220
221fn take_handle_notifiers(handle_id: &str) -> HandleNotifiers {
222 let mut store = HANDLE_STORE
223 .lock()
224 .expect("long-running handle store poisoned");
225 store
226 .entries
227 .remove(handle_id)
228 .map(|mut entry| {
229 (
230 entry.completion_tx.take(),
231 std::mem::take(&mut entry.result_txs),
232 )
233 })
234 .unwrap_or((None, Vec::new()))
235}
236
237pub struct LongRunningHandleInfo {
240 pub command_id: String,
242 pub handle_id: String,
244 pub started_at: String,
246 pub pid: u32,
248 pub process_group_id: Option<u32>,
250 pub command_display: String,
252 pub snapshot_binding: Option<harn_vm::value::DictMap>,
254}
255
256const DEFAULT_PROGRESS_MAX_INTERVAL: Duration = Duration::from_secs(30);
263
264pub(crate) struct LongRunningSpawnOptions {
265 pub(crate) env_mode: EnvMode,
266 pub(crate) env_remove: Vec<String>,
267 pub(crate) capture: CaptureConfig,
268 pub(crate) session_id: String,
269 pub(crate) progress_interval: Option<Duration>,
270 pub(crate) progress_max_interval: Option<Duration>,
271 pub(crate) progress_max_inline_bytes: usize,
272 pub(crate) snapshot_binding: Option<harn_vm::value::DictMap>,
273 pub(crate) lease: LeaseTag,
279}
280
281#[derive(Clone, Copy, Debug, PartialEq, Eq)]
285pub(crate) enum LeaseTag {
286 Awaited,
287 Service,
288}
289
290impl LeaseTag {
291 fn as_str(self) -> &'static str {
292 match self {
293 LeaseTag::Awaited => "awaited",
294 LeaseTag::Service => "service",
295 }
296 }
297}
298
299struct WaiterContext {
300 command_id: String,
301 handle_id: String,
302 session_id: String,
303 started_at: String,
304 process_group_id: Option<u32>,
305 command_display: String,
306 progress_interval: Option<Duration>,
307 progress_max_interval: Option<Duration>,
308 progress_max_inline_bytes: usize,
309 snapshot_binding: Option<harn_vm::value::DictMap>,
310 output_feed: Arc<OutputFeed>,
311}
312
313struct ProgressThreadContext {
314 command_id: String,
315 handle_id: String,
316 session_id: String,
317 started_at: String,
318 command_display: String,
319 process_group_id: Option<u32>,
320 output_path: PathBuf,
321 stdout_path: PathBuf,
322 stderr_path: PathBuf,
323 output_feed: Arc<OutputFeed>,
324 cancel_state: Arc<CancelState>,
325 done: Arc<AtomicBool>,
326 started: std::time::Instant,
327 interval: Duration,
330 max_interval: Duration,
332 max_inline_bytes: usize,
333 snapshot_binding: Option<harn_vm::value::DictMap>,
334}
335
336impl LongRunningHandleInfo {
337 pub fn into_handle_response(self) -> VmValue {
339 let Self {
340 command_id,
341 handle_id,
342 started_at,
343 pid,
344 process_group_id,
345 command_display,
346 snapshot_binding,
347 } = self;
348 proc::running_response(
349 command_id,
350 handle_id,
351 pid,
352 process_group_id,
353 started_at,
354 command_display,
355 snapshot_binding.as_ref(),
356 )
357 }
358}
359
360pub fn spawn_long_running(
366 builtin: &'static str,
367 program: String,
368 args: Vec<String>,
369 cwd: Option<PathBuf>,
370 env: BTreeMap<String, String>,
371 session_id: String,
372) -> Result<LongRunningHandleInfo, HostlibError> {
373 spawn_long_running_with_options(
374 builtin,
375 program,
376 args,
377 cwd,
378 env,
379 LongRunningSpawnOptions {
380 env_mode: EnvMode::InheritClean,
381 env_remove: Vec::new(),
382 capture: CaptureConfig::default(),
383 session_id,
384 progress_interval: None,
385 progress_max_interval: None,
386 progress_max_inline_bytes: CaptureConfig::default().max_inline_bytes,
387 snapshot_binding: None,
388 lease: LeaseTag::Awaited,
389 },
390 )
391}
392
393pub(crate) fn spawn_long_running_with_options(
394 builtin: &'static str,
395 program: String,
396 args: Vec<String>,
397 cwd: Option<PathBuf>,
398 env: BTreeMap<String, String>,
399 options: LongRunningSpawnOptions,
400) -> Result<LongRunningHandleInfo, HostlibError> {
401 let mut env = env;
402 proc::apply_toolchain_path(cwd.as_deref(), &mut env, options.env_mode);
403 let spec = SpawnSpec {
404 builtin,
405 program: program.clone(),
406 args: args.clone(),
407 cwd,
408 env,
409 env_remove: options.env_remove.clone(),
410 env_mode: options.env_mode,
411 use_stdin: false,
412 configure_process_group: true,
413 owner_death: OwnerDeathPolicy::KillContainment,
414 output_capture: process_handle::OutputCapture::Pipe,
415 };
416 let handle = process_handle::spawn_process(spec)
417 .map_err(|e| proc::process_error_to_hostlib(builtin, e))?;
418
419 let pid = handle.pid().unwrap_or(0);
420 let process_group_id = handle.process_group_id();
421 let killer = handle.killer();
422 let id = HANDLE_COUNTER.fetch_add(1, Ordering::SeqCst);
423 let handle_id = format!("hto-{:x}-{id}", std::process::id());
424 let command_id = proc::next_command_id();
425 let started_at = proc::now_rfc3339();
426 let _artifacts = proc::register_live_artifacts(&command_id, Some(&handle_id))?;
427
428 let mut all_argv = vec![program];
429 all_argv.extend(args.iter().cloned());
430 let command_display = all_argv.join(" ");
431
432 let cancel_state = Arc::new(CancelState {
433 state: Mutex::new(CancellationState::default()),
434 });
435 let output_feed = Arc::new(OutputFeed::default());
436
437 {
438 let mut store = HANDLE_STORE
439 .lock()
440 .expect("long-running handle store poisoned");
441 store.entries.insert(
442 handle_id.clone(),
443 HandleEntry {
444 handle: Some(handle),
445 killer,
446 session_id: options.session_id.clone(),
447 cancel_state: cancel_state.clone(),
448 output_feed: output_feed.clone(),
449 completion_tx: None,
450 result_txs: Vec::new(),
451 snapshot_binding: options.snapshot_binding.clone(),
452 lease: options.lease,
453 command_display: command_display.clone(),
454 started_at: started_at.clone(),
455 },
456 );
457 }
458
459 let waiter_context = WaiterContext {
460 command_id: command_id.clone(),
461 handle_id: handle_id.clone(),
462 session_id: options.session_id,
463 started_at: started_at.clone(),
464 process_group_id,
465 command_display: command_display.clone(),
466 progress_interval: options.progress_interval,
467 progress_max_interval: options.progress_max_interval,
468 progress_max_inline_bytes: options.progress_max_inline_bytes,
469 snapshot_binding: options.snapshot_binding.clone(),
470 output_feed,
471 };
472 let waiter_thread_name = waiter_context.handle_id.clone();
473 let capture = options.capture;
474 std::thread::Builder::new()
475 .name(format!("hto-waiter-{waiter_thread_name}"))
476 .spawn(move || {
477 waiter_thread(waiter_context, cancel_state, capture);
478 })
479 .map_err(|e| HostlibError::Backend {
480 builtin,
481 message: format!("failed to spawn waiter thread: {e}"),
482 })?;
483
484 Ok(LongRunningHandleInfo {
485 command_id,
486 handle_id,
487 started_at,
488 pid,
489 process_group_id,
490 command_display,
491 snapshot_binding: options.snapshot_binding,
492 })
493}
494
495fn waiter_thread(context: WaiterContext, cancel_state: Arc<CancelState>, capture: CaptureConfig) {
497 let waiter_start = std::time::Instant::now();
498
499 let mut handle = {
502 let mut store = HANDLE_STORE
503 .lock()
504 .expect("long-running handle store poisoned");
505 match store.entries.get_mut(&context.handle_id) {
506 Some(entry) => match entry.handle.take() {
507 Some(h) => h,
508 None => return, },
510 None => return, }
512 };
513
514 let done = Arc::new(AtomicBool::new(false));
515 let planned = proc::planned_artifact_paths(&context.command_id);
516 if let Some(parent) = planned.output_path.parent() {
517 let _ = std::fs::create_dir_all(parent);
518 }
519 let _ = std::fs::File::create(&planned.stdout_path);
520 let _ = std::fs::File::create(&planned.stderr_path);
521 let combined_file = std::fs::File::create(&planned.output_path)
522 .ok()
523 .map(|file| Arc::new(Mutex::new(file)));
524
525 let stdout_thread = handle.take_stdout().map(|out| {
526 spawn_output_drain(
527 out,
528 context.output_feed.clone(),
529 planned.stdout_path.clone(),
530 combined_file.clone(),
531 true,
532 )
533 });
534 let stderr_thread = handle.take_stderr().map(|err| {
535 spawn_output_drain(
536 err,
537 context.output_feed.clone(),
538 planned.stderr_path.clone(),
539 combined_file.clone(),
540 false,
541 )
542 });
543
544 let progress_thread = context
545 .progress_interval
546 .filter(|interval| !interval.is_zero())
547 .map(|interval| {
548 let max_interval = context
552 .progress_max_interval
553 .filter(|cap| !cap.is_zero())
554 .unwrap_or(DEFAULT_PROGRESS_MAX_INTERVAL)
555 .max(interval);
556 spawn_progress_thread(ProgressThreadContext {
557 command_id: context.command_id.clone(),
558 handle_id: context.handle_id.clone(),
559 session_id: context.session_id.clone(),
560 started_at: context.started_at.clone(),
561 command_display: context.command_display.clone(),
562 process_group_id: context.process_group_id,
563 output_path: planned.output_path.clone(),
564 stdout_path: planned.stdout_path.clone(),
565 stderr_path: planned.stderr_path.clone(),
566 output_feed: context.output_feed.clone(),
567 cancel_state: cancel_state.clone(),
568 done: done.clone(),
569 started: waiter_start,
570 interval,
571 max_interval,
572 max_inline_bytes: context.progress_max_inline_bytes,
573 snapshot_binding: context.snapshot_binding.clone(),
574 })
575 });
576
577 let status = handle.wait().ok();
578
579 if let Some(thread) = stdout_thread {
580 let _ = thread.join();
581 }
582 if let Some(thread) = stderr_thread {
583 let _ = thread.join();
584 }
585 done.store(true, Ordering::Release);
586 drop(progress_thread);
587 let (stdout, stderr) = {
588 let state = context
589 .output_feed
590 .state
591 .lock()
592 .unwrap_or_else(|poison| poison.into_inner());
593 (state.stdout.clone(), state.stderr.clone())
594 };
595
596 let cancellation = cancel_state.complete_wait();
597 let cancelled = cancellation.cancelled;
598 let timed_out = cancelled && cancellation.timed_out;
599 let process_cleanup = cancellation.process_cleanup;
600
601 let (exit_code, signal_name) = match status {
602 Some(s) => decode_exit_status(s),
603 None => (-1, Some("SIGKILL".to_string())),
605 };
606 let command_status = if timed_out {
607 CommandStatus::TimedOut
608 } else if cancelled {
609 CommandStatus::Killed
610 } else {
611 CommandStatus::Completed
612 };
613 let duration = waiter_start.elapsed();
614 let duration_ms = duration.as_millis() as i64;
615 let artifacts = match proc::persist_artifacts(
616 &context.command_id,
617 &stdout,
618 &stderr,
619 Some(&context.handle_id),
620 ) {
621 Ok(artifacts) => artifacts,
622 Err(error) => {
623 tracing::warn!(
624 "long-running command {} could not persist artifacts: {error}; returning in-memory terminal metadata",
625 context.command_id
626 );
627 proc::summarize_artifacts(
628 &context.command_id,
629 &stdout,
630 &stderr,
631 Some(&context.handle_id),
632 )
633 }
634 };
635 let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
636
637 let mut payload = serde_json::Map::new();
638 payload.insert(
639 "command_id".into(),
640 serde_json::Value::String(context.command_id.clone()),
641 );
642 payload.insert(
643 "status".into(),
644 serde_json::Value::String(command_status.as_str().to_string()),
645 );
646 payload.insert(
647 "handle_id".into(),
648 serde_json::Value::String(context.handle_id.clone()),
649 );
650 payload.insert(
651 "command_or_op_descriptor".into(),
652 serde_json::Value::String(context.command_display),
653 );
654 payload.insert(
655 "started_at".into(),
656 serde_json::Value::String(context.started_at),
657 );
658 payload.insert(
659 "ended_at".into(),
660 serde_json::Value::String(proc::now_rfc3339()),
661 );
662 payload.insert(
663 "duration_ms".into(),
664 serde_json::Value::Number(duration_ms.into()),
665 );
666 payload.insert(
667 "exit_code".into(),
668 serde_json::Value::Number(exit_code.into()),
669 );
670 payload.insert("timed_out".into(), serde_json::Value::Bool(timed_out));
671 payload.insert("stdout".into(), serde_json::Value::String(inline_stdout));
672 payload.insert("stderr".into(), serde_json::Value::String(inline_stderr));
673 payload.insert(
674 "output_path".into(),
675 serde_json::Value::String(to_agent_path(&artifacts.output_path)),
676 );
677 payload.insert(
678 "stdout_path".into(),
679 serde_json::Value::String(to_agent_path(&artifacts.stdout_path)),
680 );
681 payload.insert(
682 "stderr_path".into(),
683 serde_json::Value::String(to_agent_path(&artifacts.stderr_path)),
684 );
685 payload.insert(
686 "line_count".into(),
687 serde_json::Value::Number(artifacts.line_count.into()),
688 );
689 payload.insert(
690 "byte_count".into(),
691 serde_json::Value::Number(artifacts.byte_count.into()),
692 );
693 payload.insert(
694 "output_sha256".into(),
695 serde_json::Value::String(artifacts.output_sha256),
696 );
697 if let Some(pgid) = context.process_group_id {
698 payload.insert(
699 "process_group_id".into(),
700 serde_json::Value::Number((pgid as u64).into()),
701 );
702 }
703 if let Some(sig) = signal_name {
704 payload.insert("signal".into(), serde_json::Value::String(sig));
705 } else {
706 payload.insert("signal".into(), serde_json::Value::Null);
707 }
708 if let Some(snapshot_binding) = context.snapshot_binding.as_ref() {
709 payload.insert("snapshot_binding".into(), vm_dict_to_json(snapshot_binding));
710 }
711 if let Some(process_cleanup) = process_cleanup.as_ref() {
712 payload.insert(
713 "process_cleanup".into(),
714 proc::process_cleanup_to_json(process_cleanup),
715 );
716 }
717
718 let result_value = harn_vm::json_to_vm_value(&serde_json::Value::Object(payload.clone()));
719 {
720 let mut state = context
721 .output_feed
722 .state
723 .lock()
724 .unwrap_or_else(|poison| poison.into_inner());
725 state.terminal = Some(result_value.clone());
726 }
727 context.output_feed.notify.notify_waiters();
728 if !cancelled {
729 let content = serde_json::to_string(&payload).unwrap_or_default();
730 harn_vm::orchestration::agent_inbox::push(
731 &context.session_id,
732 "tool_result",
733 &content,
734 "hostlib.long_running.exit",
735 );
736 }
737 let (completion_tx, result_txs) = take_handle_notifiers(&context.handle_id);
743 for tx in result_txs {
744 let _ = tx.try_send(result_value.clone());
745 }
746 if let Some(tx) = completion_tx {
747 let _ = tx.try_send(());
748 }
749}
750
751fn spawn_output_drain(
752 mut reader: Box<dyn Read + Send>,
753 output_feed: Arc<OutputFeed>,
754 path: std::path::PathBuf,
755 combined_file: Option<Arc<Mutex<std::fs::File>>>,
756 stdout: bool,
757) -> std::thread::JoinHandle<()> {
758 std::thread::spawn(move || {
759 let mut file = std::fs::File::create(path).ok();
760 let mut buf = [0_u8; 8192];
761 loop {
762 let read = match reader.read(&mut buf) {
763 Ok(0) => break,
764 Ok(read) => read,
765 Err(_) => break,
766 };
767 let chunk = &buf[..read];
768 if let Some(file) = file.as_mut() {
769 let _ = file.write_all(chunk);
770 }
771 if let Ok(mut state) = output_feed.state.lock() {
772 if let Some(combined) = combined_file.as_ref() {
773 if let Ok(mut combined) = combined.lock() {
774 let _ = combined.write_all(chunk);
775 }
776 }
777 if stdout {
778 state.stdout.extend_from_slice(chunk);
779 } else {
780 state.stderr.extend_from_slice(chunk);
781 }
782 state.combined.extend_from_slice(chunk);
783 state.last_output_at = Some(std::time::Instant::now());
784 }
785 output_feed.notify.notify_waiters();
786 }
787 })
788}
789
790fn next_progress_interval(current: Duration, max: Duration) -> Duration {
793 current.checked_mul(2).unwrap_or(max).min(max)
794}
795
796fn spawn_progress_thread(context: ProgressThreadContext) -> std::thread::JoinHandle<()> {
797 std::thread::spawn(move || {
798 let mut current = context.interval;
806 while !context.done.load(Ordering::Acquire)
807 && !context.cancel_state.cancellation_published()
808 {
809 std::thread::sleep(current);
810 if context.done.load(Ordering::Acquire) || context.cancel_state.cancellation_published()
811 {
812 break;
813 }
814 current = next_progress_interval(current, context.max_interval);
815 let (stdout, stderr, last_output_at) = {
816 let state = context
817 .output_feed
818 .state
819 .lock()
820 .unwrap_or_else(|poison| poison.into_inner());
821 (
822 state.stdout.clone(),
823 state.stderr.clone(),
824 state.last_output_at,
825 )
826 };
827 let capture = CaptureConfig {
828 max_inline_bytes: context.max_inline_bytes,
829 ..CaptureConfig::default()
830 };
831 let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
832 let byte_count = stdout.len().saturating_add(stderr.len());
833 let silence_ms = last_output_at
837 .map(|instant| instant.elapsed().as_millis() as i64)
838 .unwrap_or_else(|| context.started.elapsed().as_millis() as i64);
839 let mut payload = serde_json::json!({
840 "command_id": &context.command_id,
841 "handle_id": &context.handle_id,
842 "status": CommandStatus::Running.as_str(),
843 "command_or_op_descriptor": &context.command_display,
844 "started_at": &context.started_at,
845 "ended_at": null,
846 "duration_ms": context.started.elapsed().as_millis() as i64,
847 "exit_code": null,
848 "signal": null,
849 "stdout": inline_stdout,
850 "stderr": inline_stderr,
851 "output_path": to_agent_path(&context.output_path),
852 "stdout_path": to_agent_path(&context.stdout_path),
853 "stderr_path": to_agent_path(&context.stderr_path),
854 "byte_count": byte_count as i64,
855 "output_offset": byte_count as i64,
859 "stderr_byte_count": stderr.len() as i64,
863 "silence_ms": silence_ms,
864 "line_count": stdout.iter().chain(stderr.iter()).filter(|byte| **byte == b'\n').count() as i64,
865 "process_group_id": context.process_group_id,
866 });
867 if let (Some(object), Some(snapshot_binding)) =
868 (payload.as_object_mut(), context.snapshot_binding.as_ref())
869 {
870 object.insert(
871 "snapshot_binding".to_string(),
872 vm_dict_to_json(snapshot_binding),
873 );
874 }
875 harn_vm::orchestration::agent_inbox::push(
876 &context.session_id,
877 "tool_progress",
878 &payload.to_string(),
879 "hostlib.long_running.progress",
880 );
881 }
882 })
883}
884
885pub(crate) struct CancelOptions {
886 pub(crate) timed_out: bool,
887 pub(crate) wait_result: Option<Duration>,
888}
889
890pub(crate) struct CancelOutcome {
891 pub(crate) cancelled: bool,
892 pub(crate) result: Option<VmValue>,
893}
894
895pub fn cancel_handle(handle_id: &str) -> bool {
899 cancel_handle_with_options(
900 handle_id,
901 CancelOptions {
902 timed_out: false,
903 wait_result: None,
904 },
905 )
906 .cancelled
907}
908
909pub(crate) fn snapshot_binding_for_handle(handle_id: &str) -> Option<harn_vm::value::DictMap> {
910 let store = HANDLE_STORE
911 .lock()
912 .expect("long-running handle store poisoned");
913 store
914 .entries
915 .get(handle_id)
916 .and_then(|entry| entry.snapshot_binding.clone())
917}
918
919pub(crate) fn output_feed_for_handle(handle_id: &str) -> Option<Arc<OutputFeed>> {
920 HANDLE_STORE
921 .lock()
922 .expect("long-running handle store poisoned")
923 .entries
924 .get(handle_id)
925 .map(|entry| entry.output_feed.clone())
926}
927
928pub(crate) fn cancel_handle_with_options(handle_id: &str, options: CancelOptions) -> CancelOutcome {
929 let (killer, cancel_state, result_rx) = {
930 let mut store = HANDLE_STORE
931 .lock()
932 .expect("long-running handle store poisoned");
933 let Some((killer, cancel_state)) = store
934 .entries
935 .get(handle_id)
936 .map(|entry| (entry.killer.clone(), entry.cancel_state.clone()))
937 else {
938 return CancelOutcome {
939 cancelled: false,
940 result: None,
941 };
942 };
943 let result_rx = options.wait_result.map(|_| {
944 let (tx, rx) = std::sync::mpsc::sync_channel::<VmValue>(1);
945 store
946 .entries
947 .get_mut(handle_id)
948 .expect("handle entry disappeared while store was locked")
949 .result_txs
950 .push(tx);
951 rx
952 });
953 (killer, cancel_state, result_rx)
954 };
955 let cancellation = cancel_state.begin_cancellation(options.timed_out);
956 let Some(mut cancellation) = cancellation else {
957 return CancelOutcome {
958 cancelled: false,
959 result: match (options.wait_result, result_rx) {
960 (Some(timeout), Some(rx)) => rx.recv_timeout(timeout).ok(),
961 _ => None,
962 },
963 };
964 };
965 kill_and_publish(killer.as_ref(), &mut cancellation);
966 drop(cancellation);
967 let result = match (options.wait_result, result_rx) {
968 (Some(timeout), Some(rx)) => rx.recv_timeout(timeout).ok(),
969 _ => None,
970 };
971 CancelOutcome {
972 cancelled: true,
973 result,
974 }
975}
976
977pub(crate) fn wait_for_result(handle_id: &str, timeout: Duration) -> Option<VmValue> {
984 if timeout.is_zero() {
985 return None;
986 }
987 let rx = {
988 let mut store = HANDLE_STORE
989 .lock()
990 .expect("long-running handle store poisoned");
991 let entry = store.entries.get_mut(handle_id)?;
992 let (tx, rx) = std::sync::mpsc::sync_channel::<VmValue>(1);
993 entry.result_txs.push(tx);
994 rx
995 };
996 rx.recv_timeout(timeout).ok()
997}
998
999pub(crate) fn list_session_handles(session_id: &str) -> VmValue {
1005 let store = HANDLE_STORE
1006 .lock()
1007 .expect("long-running handle store poisoned");
1008 let handles: Vec<VmValue> = store
1009 .entries
1010 .iter()
1011 .filter(|(_id, entry)| entry.session_id == session_id)
1012 .map(|(id, entry)| {
1013 let mut row = harn_vm::value::DictMap::new();
1014 row.put_str("handle_id", id.clone());
1015 row.put_str("session_id", entry.session_id.clone());
1016 row.put_str("lease", entry.lease.as_str());
1017 row.put_str("command_or_op_descriptor", entry.command_display.clone());
1018 row.put_str("started_at", entry.started_at.clone());
1019 VmValue::dict(row)
1020 })
1021 .collect();
1022 let mut response = harn_vm::value::DictMap::new();
1023 response.insert(
1024 harn_vm::value::intern_key("handles"),
1025 VmValue::List(Arc::new(handles)),
1026 );
1027 VmValue::dict(response)
1028}
1029
1030type SessionKillEntry = (Arc<dyn ProcessKiller>, Arc<CancelState>);
1034
1035pub fn cancel_session_handles(session_id: &str) {
1038 let to_kill: Vec<SessionKillEntry> = {
1039 let store = HANDLE_STORE
1040 .lock()
1041 .expect("long-running handle store poisoned");
1042 let matching: Vec<String> = store
1043 .entries
1044 .iter()
1045 .filter(|(_, e)| e.session_id == session_id)
1046 .map(|(id, _)| id.clone())
1047 .collect();
1048 matching
1049 .into_iter()
1050 .filter_map(|id| {
1051 let entry = store.entries.get(&id)?;
1052 Some((entry.killer.clone(), entry.cancel_state.clone()))
1053 })
1054 .collect()
1055 };
1056 for (killer, cancel_state) in to_kill {
1057 if let Some(mut cancellation) = cancel_state.begin_cancellation(false) {
1058 kill_and_publish(killer.as_ref(), &mut cancellation);
1059 }
1060 }
1061}
1062
1063pub(crate) fn register_cleanup_hook() {
1067 static REGISTERED: OnceLock<harn_vm::SessionEndHookRegistration> = OnceLock::new();
1068 REGISTERED.get_or_init(|| {
1069 let hook: Arc<dyn Fn(&str) + Send + Sync> = Arc::new(|session_id: &str| {
1070 cancel_session_handles(session_id);
1071 });
1072 harn_vm::register_session_end_hook(hook)
1073 });
1074}
1075
1076fn decode_exit_status(status: process_handle::ExitStatus) -> (i32, Option<String>) {
1077 if let Some(code) = status.code {
1078 return (code, None);
1079 }
1080 if let Some(sig) = status.signal {
1081 return (-1, Some(format!("SIG{sig}")));
1082 }
1083 (-1, None)
1084}
1085
1086pub fn register_completion_notifier(handle_id: &str) -> Option<std::sync::mpsc::Receiver<()>> {
1092 let (tx, rx) = std::sync::mpsc::sync_channel::<()>(1);
1093 let mut store = HANDLE_STORE
1094 .lock()
1095 .expect("long-running handle store poisoned");
1096 let entry = store.entries.get_mut(handle_id)?;
1097 entry.completion_tx = Some(tx);
1098 Some(rx)
1099}
1100
1101pub fn register_result_notifier(handle_id: &str) -> Option<std::sync::mpsc::Receiver<VmValue>> {
1107 let (tx, rx) = std::sync::mpsc::sync_channel::<VmValue>(1);
1108 let mut store = HANDLE_STORE
1109 .lock()
1110 .expect("long-running handle store poisoned");
1111 let entry = store.entries.get_mut(handle_id)?;
1112 entry.result_txs.push(tx);
1113 Some(rx)
1114}
1115
1116#[cfg(test)]
1117mod tests {
1118 use std::sync::{mpsc, Arc};
1119 use std::time::Duration;
1120
1121 use super::{next_progress_interval, CancelState};
1122 use crate::process::ProcessCleanupReport;
1123
1124 #[test]
1125 fn progress_backoff_doubles_then_clamps_to_max() {
1126 let max = Duration::from_secs(30);
1127 assert_eq!(
1129 next_progress_interval(Duration::from_secs(2), max),
1130 Duration::from_secs(4)
1131 );
1132 assert_eq!(
1133 next_progress_interval(Duration::from_secs(8), max),
1134 Duration::from_secs(16)
1135 );
1136 assert_eq!(next_progress_interval(Duration::from_secs(16), max), max);
1138 assert_eq!(next_progress_interval(max, max), max);
1139 assert_eq!(next_progress_interval(Duration::MAX, max), max);
1141 }
1142
1143 #[test]
1144 fn terminal_snapshot_waits_for_cleanup_publication() {
1145 let state = Arc::new(CancelState::default());
1146 let mut publication = state
1147 .begin_cancellation(true)
1148 .expect("fresh handle should accept cancellation");
1149 let (started_tx, started_rx) = mpsc::sync_channel(1);
1150 let (snapshot_tx, snapshot_rx) = mpsc::sync_channel(1);
1151 let waiter_state = state.clone();
1152 let waiter = std::thread::spawn(move || {
1153 started_tx.send(()).expect("test waiter start receiver");
1154 snapshot_tx
1155 .send(waiter_state.complete_wait())
1156 .expect("test snapshot receiver");
1157 });
1158
1159 started_rx.recv().expect("test waiter did not start");
1160 assert!(matches!(
1161 snapshot_rx.try_recv(),
1162 Err(mpsc::TryRecvError::Empty)
1163 ));
1164
1165 publication.record_cleanup(ProcessCleanupReport::for_signal(Some(42), 9));
1166 publication.publish_cancellation();
1167 drop(publication);
1168
1169 let snapshot = snapshot_rx.recv().expect("waiter did not publish snapshot");
1170 waiter.join().expect("test waiter panicked");
1171 assert!(snapshot.cancelled);
1172 assert!(snapshot.timed_out);
1173 assert_eq!(
1174 snapshot
1175 .process_cleanup
1176 .expect("cleanup must publish with cancellation")
1177 .root_pid,
1178 Some(42)
1179 );
1180 assert!(
1181 state.begin_cancellation(false).is_none(),
1182 "a completed terminal result must not be retroactively cancelled"
1183 );
1184 }
1185}