From 29a39369eb977bf6884b0d941e41b9f93e229cb1 Mon Sep 17 00:00:00 2001
From: gmega <giuliano.mega@gmail.com>
Date: Fri, 27 Jun 2025 08:23:28 -0300
Subject: [PATCH] feat: add no-timeout option to await; fix minor bugs

---
 experiments/k-node.sh     |  2 +-
 src/procmon.bash          |  6 ++++--
 test/test_experiment.bats |  1 -
 test/test_procmon.bats    | 34 +++++++++++++++++++++++++++++++++-
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/experiments/k-node.sh b/experiments/k-node.sh
index 08087a5..9e9ef39 100755
--- a/experiments/k-node.sh
+++ b/experiments/k-node.sh
@@ -56,7 +56,7 @@ for i in $(seq 1 "${repetitions}"); do
       handles+=("$result")
     done
 
-    await_all "${handles[@]}"
+    await_all "${handles[@]}" "Inf"
 
     cdx_log_timings_end
   done
diff --git a/src/procmon.bash b/src/procmon.bash
index a344632..0352eda 100644
--- a/src/procmon.bash
+++ b/src/procmon.bash
@@ -209,6 +209,8 @@ pm_list_descendants() {
 }
 
 pm_async() {
+  _pm_assert_state "running" || return 1
+
   proc_type=""
   command=()
   while [[ "$#" -gt 0 ]]; do
@@ -227,7 +229,7 @@ pm_async() {
   done
 
   (
-    set +e
+    set -e
     _pm_job_started "${BASHPID}" "$proc_type" "$@"
     trap '_pm_job_exited "${BASHPID}" "$proc_type" "killed" "$@"' TERM
     trap '_pm_job_exited "${BASHPID}" "$proc_type" "$?" "$@"' EXIT
@@ -239,7 +241,7 @@ pm_async() {
 await() {
   local pid=$1 timeout=${2:-30} start="${SECONDS}"
   while kill -0 "$pid" 2> /dev/null; do
-    if ((SECONDS - start > timeout)); then
+    if [ "$timeout" != 'Inf' ] && ((SECONDS - start > timeout)); then
       echoerr "Error: timeout waiting for process $pid to exit"
       return 1
     fi
diff --git a/test/test_experiment.bats b/test/test_experiment.bats
index fc71041..fcddae6 100644
--- a/test/test_experiment.bats
+++ b/test/test_experiment.bats
@@ -23,7 +23,6 @@ setup() {
   done
 
   assert [ "$found" = true ]
-
 }
 
 @test "should launch Codex nodes with metrics enabled when there is an experiment in scope" {
diff --git a/test/test_procmon.bats b/test/test_procmon.bats
index 8e6f3f9..f3b0e95 100644
--- a/test/test_procmon.bats
+++ b/test/test_procmon.bats
@@ -1,5 +1,5 @@
 #!/usr/bin/env bats
-# shellcheck disable=SC2128
+# shellcheck disable=SC2128,SC2317
 setup() {
   load test_helper/common_setup
   common_setup
@@ -10,6 +10,34 @@ setup() {
   pm_set_outputs "${TEST_OUTPUTS}/pm"
 }
 
+@test "should allow awaiting for a process" {
+  job() {
+    sleep 0.1
+  }
+
+  pm_start
+
+  pm_async job
+  pid=$result
+
+  assert await "$pid"
+  pm_stop
+}
+
+@test "should allow awaiting for a process up until a timeout" {
+  job() {
+    sleep 5000
+  }
+
+  pm_start
+
+  pm_async job
+  pid=$result
+
+  refute await "$pid" 1
+  pm_stop
+}
+
 @test "should kill processes recursively" {
   # Note that this is fragile. We need to structure
   # the process tree such that the parent does not exit
@@ -230,6 +258,8 @@ callback() {
 @test "should allow passing custom arguments to lifecycle callback" {
   pm_register_callback "sleepy" "callback"
 
+  pm_start
+
   pm_async sleep 0.1 -%- "sleepy" "arg1" "arg2"
   pid=$result
 
@@ -237,4 +267,6 @@ callback() {
 
   assert_equal "$(cat "${_pm_output}/${pid}-sleepy-start-args")" "arg1 arg2"
   assert_equal "$(cat "${_pm_output}/${pid}-sleepy-exit-args")" "arg1 arg2"
+
+  pm_stop
 }