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