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 };
239 let handle = process_handle::spawn_process(spec)
240 .map_err(|e| proc::process_error_to_hostlib(builtin, e))?;
241
242 let pid = handle.pid().unwrap_or(0);
243 let process_group_id = handle.process_group_id();
244 let killer = handle.killer();
245 let id = HANDLE_COUNTER.fetch_add(1, Ordering::SeqCst);
246 let handle_id = format!("hto-{:x}-{id}", std::process::id());
247 let command_id = proc::next_command_id();
248 let started_at = proc::now_rfc3339();
249 let _artifacts = proc::register_live_artifacts(&command_id, Some(&handle_id))?;
250
251 let mut all_argv = vec![program];
252 all_argv.extend(args.iter().cloned());
253 let command_display = all_argv.join(" ");
254
255 let cancel_state = Arc::new(CancelState {
256 cancelled: AtomicBool::new(false),
257 timed_out: AtomicBool::new(false),
258 process_cleanup: Mutex::new(None),
259 });
260
261 {
262 let mut store = HANDLE_STORE
263 .lock()
264 .expect("long-running handle store poisoned");
265 store.entries.insert(
266 handle_id.clone(),
267 HandleEntry {
268 handle: Some(handle),
269 killer,
270 session_id: options.session_id.clone(),
271 cancel_state: cancel_state.clone(),
272 completion_tx: None,
273 result_tx: None,
274 snapshot_binding: options.snapshot_binding.clone(),
275 },
276 );
277 }
278
279 let waiter_context = WaiterContext {
280 command_id: command_id.clone(),
281 handle_id: handle_id.clone(),
282 session_id: options.session_id,
283 started_at: started_at.clone(),
284 process_group_id,
285 command_display: command_display.clone(),
286 progress_interval: options.progress_interval,
287 progress_max_inline_bytes: options.progress_max_inline_bytes,
288 snapshot_binding: options.snapshot_binding.clone(),
289 };
290 let waiter_thread_name = waiter_context.handle_id.clone();
291 let capture = options.capture;
292 std::thread::Builder::new()
293 .name(format!("hto-waiter-{waiter_thread_name}"))
294 .spawn(move || {
295 waiter_thread(waiter_context, cancel_state, capture);
296 })
297 .map_err(|e| HostlibError::Backend {
298 builtin,
299 message: format!("failed to spawn waiter thread: {e}"),
300 })?;
301
302 Ok(LongRunningHandleInfo {
303 command_id,
304 handle_id,
305 started_at,
306 pid,
307 process_group_id,
308 command_display,
309 snapshot_binding: options.snapshot_binding,
310 })
311}
312
313fn waiter_thread(context: WaiterContext, cancel_state: Arc<CancelState>, capture: CaptureConfig) {
315 let waiter_start = std::time::Instant::now();
316
317 let mut handle = {
320 let mut store = HANDLE_STORE
321 .lock()
322 .expect("long-running handle store poisoned");
323 match store.entries.get_mut(&context.handle_id) {
324 Some(entry) => match entry.handle.take() {
325 Some(h) => h,
326 None => return, },
328 None => return, }
330 };
331
332 let output_state = Arc::new(Mutex::new(OutputState::default()));
333 let done = Arc::new(AtomicBool::new(false));
334 let planned = proc::planned_artifact_paths(&context.command_id);
335 if let Some(parent) = planned.output_path.parent() {
336 let _ = std::fs::create_dir_all(parent);
337 }
338 let _ = std::fs::File::create(&planned.stdout_path);
339 let _ = std::fs::File::create(&planned.stderr_path);
340 let combined_file = std::fs::File::create(&planned.output_path)
341 .ok()
342 .map(|file| Arc::new(Mutex::new(file)));
343
344 let stdout_thread = handle.take_stdout().map(|out| {
345 spawn_output_drain(
346 out,
347 output_state.clone(),
348 planned.stdout_path.clone(),
349 combined_file.clone(),
350 true,
351 )
352 });
353 let stderr_thread = handle.take_stderr().map(|err| {
354 spawn_output_drain(
355 err,
356 output_state.clone(),
357 planned.stderr_path.clone(),
358 combined_file.clone(),
359 false,
360 )
361 });
362
363 let progress_thread = context
364 .progress_interval
365 .filter(|interval| !interval.is_zero())
366 .map(|interval| {
367 spawn_progress_thread(ProgressThreadContext {
368 command_id: context.command_id.clone(),
369 handle_id: context.handle_id.clone(),
370 session_id: context.session_id.clone(),
371 started_at: context.started_at.clone(),
372 command_display: context.command_display.clone(),
373 process_group_id: context.process_group_id,
374 output_path: planned.output_path.clone(),
375 stdout_path: planned.stdout_path.clone(),
376 stderr_path: planned.stderr_path.clone(),
377 output_state: output_state.clone(),
378 cancel_state: cancel_state.clone(),
379 done: done.clone(),
380 started: waiter_start,
381 interval,
382 max_inline_bytes: context.progress_max_inline_bytes,
383 snapshot_binding: context.snapshot_binding.clone(),
384 })
385 });
386
387 let status = handle.wait().ok();
388
389 if let Some(thread) = stdout_thread {
390 let _ = thread.join();
391 }
392 if let Some(thread) = stderr_thread {
393 let _ = thread.join();
394 }
395 done.store(true, Ordering::Release);
396 drop(progress_thread);
397 let (stdout, stderr) = {
398 let state = output_state
399 .lock()
400 .unwrap_or_else(|poison| poison.into_inner());
401 (state.stdout.clone(), state.stderr.clone())
402 };
403
404 let (completion_tx, result_tx) = {
407 let mut store = HANDLE_STORE
408 .lock()
409 .expect("long-running handle store poisoned");
410 let entry = store
411 .entries
412 .remove(&context.handle_id)
413 .map(|mut e| (e.completion_tx.take(), e.result_tx.take()));
414 entry.unwrap_or((None, None))
415 };
416
417 let signal_done = move || {
418 if let Some(tx) = completion_tx {
419 let _ = tx.try_send(());
420 }
421 };
422
423 let cancelled = cancel_state.cancelled.load(Ordering::Acquire);
424 let timed_out = cancelled && cancel_state.timed_out.load(Ordering::Acquire);
425 let process_cleanup = cancel_state
426 .process_cleanup
427 .lock()
428 .unwrap_or_else(|poison| poison.into_inner())
429 .clone();
430
431 let (exit_code, signal_name) = match status {
432 Some(s) => decode_exit_status(s),
433 None => (-1, Some("SIGKILL".to_string())),
435 };
436 let command_status = if timed_out {
437 CommandStatus::TimedOut
438 } else if cancelled {
439 CommandStatus::Killed
440 } else {
441 CommandStatus::Completed
442 };
443 let duration = waiter_start.elapsed();
444 let duration_ms = duration.as_millis() as i64;
445 let artifacts = match proc::persist_artifacts(
446 &context.command_id,
447 &stdout,
448 &stderr,
449 Some(&context.handle_id),
450 ) {
451 Ok(artifacts) => artifacts,
452 Err(_) => return,
453 };
454 let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
455
456 let mut payload = serde_json::Map::new();
457 payload.insert(
458 "command_id".into(),
459 serde_json::Value::String(context.command_id.clone()),
460 );
461 payload.insert(
462 "status".into(),
463 serde_json::Value::String(command_status.as_str().to_string()),
464 );
465 payload.insert(
466 "handle_id".into(),
467 serde_json::Value::String(context.handle_id),
468 );
469 payload.insert(
470 "command_or_op_descriptor".into(),
471 serde_json::Value::String(context.command_display),
472 );
473 payload.insert(
474 "started_at".into(),
475 serde_json::Value::String(context.started_at),
476 );
477 payload.insert(
478 "ended_at".into(),
479 serde_json::Value::String(proc::now_rfc3339()),
480 );
481 payload.insert(
482 "duration_ms".into(),
483 serde_json::Value::Number(duration_ms.into()),
484 );
485 payload.insert(
486 "exit_code".into(),
487 serde_json::Value::Number(exit_code.into()),
488 );
489 payload.insert("timed_out".into(), serde_json::Value::Bool(timed_out));
490 payload.insert("stdout".into(), serde_json::Value::String(inline_stdout));
491 payload.insert("stderr".into(), serde_json::Value::String(inline_stderr));
492 payload.insert(
493 "output_path".into(),
494 serde_json::Value::String(to_agent_path(&artifacts.output_path)),
495 );
496 payload.insert(
497 "stdout_path".into(),
498 serde_json::Value::String(to_agent_path(&artifacts.stdout_path)),
499 );
500 payload.insert(
501 "stderr_path".into(),
502 serde_json::Value::String(to_agent_path(&artifacts.stderr_path)),
503 );
504 payload.insert(
505 "line_count".into(),
506 serde_json::Value::Number(artifacts.line_count.into()),
507 );
508 payload.insert(
509 "byte_count".into(),
510 serde_json::Value::Number(artifacts.byte_count.into()),
511 );
512 payload.insert(
513 "output_sha256".into(),
514 serde_json::Value::String(artifacts.output_sha256),
515 );
516 if let Some(pgid) = context.process_group_id {
517 payload.insert(
518 "process_group_id".into(),
519 serde_json::Value::Number((pgid as u64).into()),
520 );
521 }
522 if let Some(sig) = signal_name {
523 payload.insert("signal".into(), serde_json::Value::String(sig));
524 } else {
525 payload.insert("signal".into(), serde_json::Value::Null);
526 }
527 if let Some(snapshot_binding) = context.snapshot_binding.as_ref() {
528 payload.insert("snapshot_binding".into(), vm_dict_to_json(snapshot_binding));
529 }
530 if let Some(process_cleanup) = process_cleanup.as_ref() {
531 payload.insert(
532 "process_cleanup".into(),
533 proc::process_cleanup_to_json(process_cleanup),
534 );
535 }
536
537 if let Some(tx) = result_tx {
538 let value = serde_json::Value::Object(payload.clone());
539 let _ = tx.try_send(harn_vm::json_to_vm_value(&value));
540 }
541 if !cancelled {
542 let content = serde_json::to_string(&payload).unwrap_or_default();
543 harn_vm::orchestration::agent_inbox::push(
544 &context.session_id,
545 "tool_result",
546 &content,
547 "hostlib.long_running.exit",
548 );
549 }
550 signal_done();
551}
552
553fn spawn_output_drain(
554 mut reader: Box<dyn Read + Send>,
555 state: Arc<Mutex<OutputState>>,
556 path: std::path::PathBuf,
557 combined_file: Option<Arc<Mutex<std::fs::File>>>,
558 stdout: bool,
559) -> std::thread::JoinHandle<()> {
560 std::thread::spawn(move || {
561 let mut file = std::fs::File::create(path).ok();
562 let mut buf = [0_u8; 8192];
563 loop {
564 let read = match reader.read(&mut buf) {
565 Ok(0) => break,
566 Ok(read) => read,
567 Err(_) => break,
568 };
569 let chunk = &buf[..read];
570 if let Some(file) = file.as_mut() {
571 let _ = file.write_all(chunk);
572 }
573 if let Some(combined) = combined_file.as_ref() {
574 if let Ok(mut combined) = combined.lock() {
575 let _ = combined.write_all(chunk);
576 }
577 }
578 if let Ok(mut state) = state.lock() {
579 if stdout {
580 state.stdout.extend_from_slice(chunk);
581 } else {
582 state.stderr.extend_from_slice(chunk);
583 }
584 }
585 }
586 })
587}
588
589fn spawn_progress_thread(context: ProgressThreadContext) -> std::thread::JoinHandle<()> {
590 std::thread::spawn(move || {
591 while !context.done.load(Ordering::Acquire)
592 && !context.cancel_state.cancelled.load(Ordering::Acquire)
593 {
594 std::thread::sleep(context.interval);
595 if context.done.load(Ordering::Acquire)
596 || context.cancel_state.cancelled.load(Ordering::Acquire)
597 {
598 break;
599 }
600 let (stdout, stderr) = {
601 let state = context
602 .output_state
603 .lock()
604 .unwrap_or_else(|poison| poison.into_inner());
605 (state.stdout.clone(), state.stderr.clone())
606 };
607 let capture = CaptureConfig {
608 max_inline_bytes: context.max_inline_bytes,
609 ..CaptureConfig::default()
610 };
611 let (inline_stdout, inline_stderr) = proc::inline_output(&stdout, &stderr, capture);
612 let byte_count = stdout.len().saturating_add(stderr.len());
613 let mut payload = serde_json::json!({
614 "command_id": &context.command_id,
615 "handle_id": &context.handle_id,
616 "status": CommandStatus::Running.as_str(),
617 "command_or_op_descriptor": &context.command_display,
618 "started_at": &context.started_at,
619 "ended_at": null,
620 "duration_ms": context.started.elapsed().as_millis() as i64,
621 "exit_code": null,
622 "signal": null,
623 "stdout": inline_stdout,
624 "stderr": inline_stderr,
625 "output_path": to_agent_path(&context.output_path),
626 "stdout_path": to_agent_path(&context.stdout_path),
627 "stderr_path": to_agent_path(&context.stderr_path),
628 "byte_count": byte_count as i64,
629 "line_count": stdout.iter().chain(stderr.iter()).filter(|byte| **byte == b'\n').count() as i64,
630 "process_group_id": context.process_group_id,
631 });
632 if let (Some(object), Some(snapshot_binding)) =
633 (payload.as_object_mut(), context.snapshot_binding.as_ref())
634 {
635 object.insert(
636 "snapshot_binding".to_string(),
637 vm_dict_to_json(snapshot_binding),
638 );
639 }
640 harn_vm::orchestration::agent_inbox::push(
641 &context.session_id,
642 "tool_progress",
643 &payload.to_string(),
644 "hostlib.long_running.progress",
645 );
646 }
647 })
648}
649
650pub(crate) struct CancelOptions {
651 pub(crate) timed_out: bool,
652 pub(crate) wait_result: Option<Duration>,
653}
654
655pub(crate) struct CancelOutcome {
656 pub(crate) cancelled: bool,
657 pub(crate) result: Option<VmValue>,
658}
659
660pub fn cancel_handle(handle_id: &str) -> bool {
664 cancel_handle_with_options(
665 handle_id,
666 CancelOptions {
667 timed_out: false,
668 wait_result: None,
669 },
670 )
671 .cancelled
672}
673
674pub(crate) fn snapshot_binding_for_handle(handle_id: &str) -> Option<harn_vm::value::DictMap> {
675 let store = HANDLE_STORE
676 .lock()
677 .expect("long-running handle store poisoned");
678 store
679 .entries
680 .get(handle_id)
681 .and_then(|entry| entry.snapshot_binding.clone())
682}
683
684pub(crate) fn cancel_handle_with_options(handle_id: &str, options: CancelOptions) -> CancelOutcome {
685 let (killer, cancel_state, result_rx) = {
686 let mut store = HANDLE_STORE
687 .lock()
688 .expect("long-running handle store poisoned");
689 let Some(entry) = store.entries.get_mut(handle_id) else {
690 return CancelOutcome {
691 cancelled: false,
692 result: None,
693 };
694 };
695 if entry.cancel_state.cancelled.swap(true, Ordering::AcqRel) {
696 return CancelOutcome {
697 cancelled: false,
698 result: None,
699 };
700 }
701 entry
702 .cancel_state
703 .timed_out
704 .store(options.timed_out, Ordering::Release);
705 let result_rx = options.wait_result.map(|_| {
706 let (tx, rx) = std::sync::mpsc::sync_channel::<VmValue>(1);
707 entry.result_tx = Some(tx);
708 rx
709 });
710 (entry.killer.clone(), entry.cancel_state.clone(), result_rx)
711 };
712 do_kill(killer, cancel_state);
713 let result = match (options.wait_result, result_rx) {
714 (Some(timeout), Some(rx)) => rx.recv_timeout(timeout).ok(),
715 _ => None,
716 };
717 CancelOutcome {
718 cancelled: true,
719 result,
720 }
721}
722
723type SessionKillEntry = (Arc<dyn ProcessKiller>, Arc<CancelState>);
727
728pub fn cancel_session_handles(session_id: &str) {
731 let to_kill: Vec<SessionKillEntry> = {
732 let store = HANDLE_STORE
733 .lock()
734 .expect("long-running handle store poisoned");
735 let matching: Vec<String> = store
736 .entries
737 .iter()
738 .filter(|(_, e)| e.session_id == session_id)
739 .map(|(id, _)| id.clone())
740 .collect();
741 matching
742 .into_iter()
743 .filter_map(|id| {
744 let entry = store.entries.get(&id)?;
745 if entry.cancel_state.cancelled.swap(true, Ordering::AcqRel) {
746 return None;
747 }
748 entry.cancel_state.timed_out.store(false, Ordering::Release);
749 Some((entry.killer.clone(), entry.cancel_state.clone()))
750 })
751 .collect()
752 };
753 for (killer, cancel_state) in to_kill {
754 do_kill(killer, cancel_state);
755 }
756}
757
758fn do_kill(killer: Arc<dyn ProcessKiller>, cancel_state: Arc<CancelState>) {
761 let report = killer.kill();
764 {
765 let mut stored = cancel_state
766 .process_cleanup
767 .lock()
768 .unwrap_or_else(|poison| poison.into_inner());
769 match stored.as_mut() {
770 Some(existing) => existing.merge(report),
771 None => *stored = Some(report),
772 }
773 }
774 cancel_state.cancelled.store(true, Ordering::Release);
775}
776
777pub(crate) fn register_cleanup_hook() {
781 static REGISTERED: OnceLock<()> = OnceLock::new();
782 REGISTERED.get_or_init(|| {
783 let hook: Arc<dyn Fn(&str) + Send + Sync> = Arc::new(|session_id: &str| {
784 cancel_session_handles(session_id);
785 });
786 harn_vm::register_session_end_hook(hook);
787 });
788}
789
790fn decode_exit_status(status: process_handle::ExitStatus) -> (i32, Option<String>) {
791 if let Some(code) = status.code {
792 return (code, None);
793 }
794 if let Some(sig) = status.signal {
795 return (-1, Some(format!("SIG{sig}")));
796 }
797 (-1, None)
798}
799
800pub fn register_completion_notifier(handle_id: &str) -> Option<std::sync::mpsc::Receiver<()>> {
806 let (tx, rx) = std::sync::mpsc::sync_channel::<()>(1);
807 let mut store = HANDLE_STORE
808 .lock()
809 .expect("long-running handle store poisoned");
810 let entry = store.entries.get_mut(handle_id)?;
811 entry.completion_tx = Some(tx);
812 Some(rx)
813}