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